Migrate to loguru

.gitattributes

@@ -1,3 +1,4 @@
+*.log
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -5,3 +5,5 @@ __pycache__/
 
 # vLLM
 inference_logs/
+
+logs/*

requirements.txt

@@ -25,4 +25,4 @@ langchain==0.1.9
 accelerate==0.27.2
 azure-search-documents==11.6.0b1
 azure-identity==1.16.0
-
+loguru==0.7.2

src/app.py

@@ -3,7 +3,6 @@ from time import time
 import huggingface_hub
 import streamlit as st
 from config import config
-from utils import get_assistant_message
 from functioncall import ModelInference
 
 

src/functioncall.py

@@ -3,21 +3,18 @@ import torch
 import json
 from config import config
 from typing import List, Dict
+from logger import logger
 
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig
-)
+from transformers import AutoTokenizer
 
 import functions
 from prompter import PromptManager
 from validator import validate_function_call_schema
 from langchain_community.chat_models import ChatOllama
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
 
 from utils import (
-    inference_logger,
-    get_assistant_message,
     get_chat_template,
     validate_and_extract_tool_calls
 )
@@ -26,8 +23,9 @@ class ModelInference:
     def __init__(self, chat_template: str):
         self.prompter = PromptManager()
         
-        self.model = ChatOllama(model=config.ollama_model, 
-                                temperature=0.0, format='json')
+        self.model = ChatOllama(model=config.ollama_model, temperature=0.0, format='json')
+        template = PromptTemplate(template="""<|im_start|>system\nYou are a function calling AI model. You are provided with function signatures within <tools></tools> XML tags. You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions. Here are the available tools: <tools> {"type": "function", "function": {"name": "get_stock_fundamentals", "description": "get_stock_fundamentals(symbol: str) -> dict - Get fundamental data for a given stock symbol using yfinance API.\\n\\n    Args:\\n        symbol (str): The stock symbol.\\n\\n    Returns:\\n        dict: A dictionary containing fundamental data.\\n            Keys:\\n                - \'symbol\': The stock symbol.\\n                - \'company_name\': The long name of the company.\\n                - \'sector\': The sector to which the company belongs.\\n                - \'industry\': The industry to which the company belongs.\\n                - \'market_cap\': The market capitalization of the company.\\n                - \'pe_ratio\': The forward price-to-earnings ratio.\\n                - \'pb_ratio\': The price-to-book ratio.\\n                - \'dividend_yield\': The dividend yield.\\n                - \'eps\': The trailing earnings per share.\\n                - \'beta\': The beta value of the stock.\\n                - \'52_week_high\': The 52-week high price of the stock.\\n                - \'52_week_low\': The 52-week low price of the stock.", "parameters": {"type": "object", "properties": {"symbol": {"type": "string"}}, "required": ["symbol"]}}}  </tools> Use the following pydantic model json schema for each tool call you will make: {"properties": {"arguments": {"title": "Arguments", "type": "object"}, "name": {"title": "Name", "type": "string"}}, "required": ["arguments", "name"], "title": "FunctionCall", "type": "object"} For each function call return a json object with function name and arguments within <tool_call></tool_call> XML tags as follows:\n<tool_call>\n{"arguments": <args-dict>, "name": <function-name>}\n</tool_call><|im_end|>\n""", input_variables=["question"])
+        chain = template | self.model | StrOutputParser()
         
         self.tokenizer = AutoTokenizer.from_pretrained(config.hf_model, trust_remote_code=True)
         self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -37,19 +35,22 @@ class ModelInference:
             print("No chat template defined, getting chat_template...")
             self.tokenizer.chat_template = get_chat_template(chat_template)
 
+        logger.info(f"Model loaded: {self.model}")
 
     def process_completion_and_validate(self, completion, chat_template):
         if completion:
+            # completion = f"<tool_call>\n{completion}\n</tool_call>"
+            breakpoint()
             validation, tool_calls, error_message = validate_and_extract_tool_calls(completion)
 
             if validation:
-                inference_logger.info(f"parsed tool calls:\n{json.dumps(tool_calls, indent=2)}")
+                logger.info(f"parsed tool calls:\n{json.dumps(tool_calls, indent=2)}")
                 return tool_calls, completion, error_message
             else:
                 tool_calls = None
                 return tool_calls, completion, error_message
         else:
-            inference_logger.warning("Assistant message is None")
+            logger.warning("Assistant message is None")
             raise ValueError("Assistant message is None")
         
     def execute_function_call(self, tool_call):
@@ -58,7 +59,7 @@ class ModelInference:
         function_to_call = getattr(functions, function_name, None)
         function_args = tool_call.get("arguments", {})
 
-        inference_logger.info(f"Invoking function call {function_name} ...")
+        logger.info(f"Invoking function call {function_name} ...")
         function_response = function_to_call(*function_args.values())
         results_dict = f'{{"name": "{function_name}", "content": {function_response}}}'
         return results_dict
@@ -88,8 +89,9 @@ class ModelInference:
                 prompt.append({"role": "assistant", "content": assistant_message})
 
                 tool_message = f"Agent iteration {depth} to assist with user query: {query}\n"
+                logger.info(f"Found tool calls: {tool_calls}")
                 if tool_calls:
-                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+                    logger.info(f"Assistant Message:\n{assistant_message}")
 
                     for tool_call in tool_calls:
                         validation, message = validate_function_call_schema(tool_call, tools)
@@ -97,12 +99,12 @@ class ModelInference:
                             try:
                                 function_response = self.execute_function_call(tool_call)
                                 tool_message += f"<tool_response>\n{function_response}\n</tool_response>\n"
-                                inference_logger.info(f"Here's the response from the function call: {tool_call.get('name')}\n{function_response}")
+                                logger.info(f"Here's the response from the function call: {tool_call.get('name')}\n{function_response}")
                             except Exception as e:
-                                inference_logger.info(f"Could not execute function: {e}")
+                                logger.info(f"Could not execute function: {e}")
                                 tool_message += f"<tool_response>\nThere was an error when executing the function: {tool_call.get('name')}\nHere's the error traceback: {e}\nPlease call this function again with correct arguments within XML tags <tool_call></tool_call>\n</tool_response>\n"
                         else:
-                            inference_logger.info(message)
+                            logger.info(message)
                             tool_message += f"<tool_response>\nThere was an error validating function call against function signature: {tool_call.get('name')}\nHere's the error traceback: {message}\nPlease call this function again with correct arguments within XML tags <tool_call></tool_call>\n</tool_response>\n"
                     prompt.append({"role": "tool", "content": tool_message})
 
@@ -116,7 +118,7 @@ class ModelInference:
                     completion = self.run_inference(prompt)
                     return recursive_loop(prompt, completion, depth)
                 elif error_message:
-                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+                    logger.info(f"Assistant Message:\n{assistant_message}")
                     tool_message += f"<tool_response>\nThere was an error parsing function calls\n Here's the error stack trace: {error_message}\nPlease call the function again with correct syntax<tool_response>"
                     prompt.append({"role": "tool", "content": tool_message})
 
@@ -128,11 +130,11 @@ class ModelInference:
                     completion = self.run_inference(prompt)
                     return recursive_loop(prompt, completion, depth)
                 else:
-                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+                    logger.info(f"Assistant Message:\n{assistant_message}")
                     return assistant_message
 
             return recursive_loop(prompt, completion, depth)
 
         except Exception as e:
-            inference_logger.error(f"Exception occurred: {e}")
+            logger.error(f"Exception occurred: {e}")
             raise e

src/functions.py

@@ -8,7 +8,7 @@ from datetime import datetime
 
 from typing import List
 from bs4 import BeautifulSoup
-from utils import inference_logger
+from logger import logger
 from langchain.tools import tool
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from config import config
@@ -69,13 +69,13 @@ def google_search_and_scrape(query: str) -> dict:
     params = {'q': query, 'num': num_results}
     headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/94.0.4606.61 Safari/537.3'}
     
-    inference_logger.info(f"Performing google search with query: {query}\nplease wait...")
+    logger.info(f"Performing google search with query: {query}\nplease wait...")
     response = requests.get(url, params=params, headers=headers)
     soup = BeautifulSoup(response.text, 'html.parser')
     urls = [result.find('a')['href'] for result in soup.find_all('div', class_='tF2Cxc')]
     
-    inference_logger.info(f"Scraping text from urls, please wait...") 
-    [inference_logger.info(url) for url in urls]
+    logger.info(f"Scraping text from urls, please wait...") 
+    [logger.info(url) for url in urls]
     with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
         futures = [executor.submit(lambda url: (url, requests.get(url, headers=headers).text if isinstance(url, str) else None), url) for url in urls[:num_results] if isinstance(url, str)]
         results = []
@@ -196,25 +196,6 @@ def get_key_financial_ratios(symbol: str) -> dict:
         print(f"Error fetching key financial ratios for {symbol}: {e}")
         return {}
 
-@tool
-def get_analyst_recommendations(symbol: str) -> pd.DataFrame:
-    """
-    Get analyst recommendations for a given stock symbol.
-
-    Args:
-    symbol (str): The stock symbol.
-
-    Returns:
-    pd.DataFrame: DataFrame containing analyst recommendations.
-    """
-    try:
-        stock = yf.Ticker(symbol)
-        recommendations = stock.recommendations
-        return recommendations
-    except Exception as e:
-        print(f"Error fetching analyst recommendations for {symbol}: {e}")
-        return pd.DataFrame()
-
 @tool
 def get_dividend_data(symbol: str) -> pd.DataFrame:
     """
@@ -245,6 +226,7 @@ def get_company_news(symbol: str) -> pd.DataFrame:
     Returns:
     pd.DataFrame: DataFrame containing company news and press releases.
     """
+    config.status.update(label=":newspaper: Getting news")
     try:
         news = yf.Ticker(symbol).news
         return news
@@ -293,7 +275,7 @@ def get_openai_tools() -> List[dict]:
         get_analysis,  
         # google_search_and_scrape,
         get_current_stock_price,
-        # get_company_news,
+        get_company_news,
         # get_company_profile,
         # get_stock_fundamentals,
         # get_financial_statements,

src/logger.py

@@ -0,0 +1,13 @@
+import sys
+from time import time
+from loguru import logger
+from pathlib import Path
+
+log_dir = Path("logs")
+log_dir.mkdir(exist_ok=True)
+
+logger.remove()  # Remove the default logger configuration
+
+# Configure the logger to write logs to both files and stdout
+logger.add(sys.stdout, format="{time} - {file} - {line} - {message}", backtrace=True)
+logger.add(log_dir / f"{time()}.log", format="{time} - {file} - {line} - {message}", backtrace=True)

src/utils.py

@@ -5,6 +5,7 @@ import json
 import logging
 import datetime
 import xml.etree.ElementTree as ET
+from logger import logger
 
 from logging.handlers import RotatingFileHandler
 
@@ -27,9 +28,6 @@ file_handler.setLevel(logging.INFO)
 formatter = logging.Formatter("%(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s", datefmt="%Y-%m-%d:%H:%M:%S")
 file_handler.setFormatter(formatter)
 
-inference_logger = logging.getLogger("function-calling-inference")
-inference_logger.addHandler(file_handler)
-
 def get_fewshot_examples(num_fewshot):
     """return a list of few shot examples"""
     example_path = os.path.join(script_dir, 'prompt_assets', 'few_shot.json')
@@ -45,7 +43,7 @@ def get_chat_template(chat_template):
 
     if not os.path.exists(template_path):
         print
-        inference_logger.error(f"Template file not found: {chat_template}")
+        logger.error(f"Template file not found: {chat_template}")
         return None
     try:
         with open(template_path, 'r') as file:
@@ -55,31 +53,6 @@ def get_chat_template(chat_template):
         print(f"Error loading template: {e}")
         return None
 
-def get_assistant_message(completion, chat_template, eos_token):
-    """define and match pattern to find the assistant message"""
-    completion = completion.strip()
-
-    if chat_template == "zephyr":
-        assistant_pattern = re.compile(r'<\|assistant\|>((?:(?!<\|assistant\|>).)*)$', re.DOTALL)
-    elif chat_template == "chatml":
-        assistant_pattern = re.compile(r'<\|im_start\|>\s*assistant((?:(?!<\|im_start\|>\s*assistant).)*)$', re.DOTALL)
-
-    elif chat_template == "vicuna":
-        assistant_pattern = re.compile(r'ASSISTANT:\s*((?:(?!ASSISTANT:).)*)$', re.DOTALL)
-    else:
-        raise NotImplementedError(f"Handling for chat_template '{chat_template}' is not implemented.")
-    
-    assistant_match = assistant_pattern.search(completion)
-    if assistant_match:
-        assistant_content = assistant_match.group(1).strip()
-        if chat_template == "vicuna":
-            eos_token = f"</s>{eos_token}"
-        return assistant_content.replace(eos_token, "")
-    else:
-        assistant_content = None
-        inference_logger.info("No match found for the assistant pattern")
-        return assistant_content
-
 def validate_and_extract_tool_calls(assistant_content):
     validation_result = False
     tool_calls = []
@@ -108,11 +81,11 @@ def validate_and_extract_tool_calls(assistant_content):
                                         f"- JSON Decode Error: {json_err}\n"\
                                         f"- Fallback Syntax/Value Error: {eval_err}\n"\
                                         f"- Problematic JSON text: {json_text}"
-                        inference_logger.error(error_message)
+                        logger.error(error_message)
                         continue
             except Exception as e:
                 error_message = f"Cannot strip text: {e}"
-                inference_logger.error(error_message)
+                logger.error(error_message)
 
             if json_data is not None:
                 tool_calls.append(json_data)
@@ -120,7 +93,7 @@ def validate_and_extract_tool_calls(assistant_content):
 
     except ET.ParseError as err:
         error_message = f"XML Parse Error: {err}"
-        inference_logger.error(f"XML Parse Error: {err}")
+        logger.error(f"XML Parse Error: {err}")
 
     # Return default values if no valid data is extracted
     return validation_result, tool_calls, error_message

src/validator.py

@@ -2,9 +2,11 @@ import ast
 import json
 from jsonschema import validate
 from pydantic import ValidationError
-from utils import inference_logger, extract_json_from_markdown
+from logger import logger
+from utils import extract_json_from_markdown
 from schema import FunctionCall, FunctionSignature
 
+
 def validate_function_call_schema(call, signatures):
     try:
         call_data = FunctionCall(**call)
@@ -16,18 +18,26 @@ def validate_function_call_schema(call, signatures):
             signature_data = FunctionSignature(**signature)
             if signature_data.function.name == call_data.name:
                 # Validate types in function arguments
-                for arg_name, arg_schema in signature_data.function.parameters.get('properties', {}).items():
+                for arg_name, arg_schema in signature_data.function.parameters.get(
+                    "properties", {}
+                ).items():
                     if arg_name in call_data.arguments:
                         call_arg_value = call_data.arguments[arg_name]
                         if call_arg_value:
                             try:
-                                validate_argument_type(arg_name, call_arg_value, arg_schema)
+                                validate_argument_type(
+                                    arg_name, call_arg_value, arg_schema
+                                )
                             except Exception as arg_validation_error:
                                 return False, str(arg_validation_error)
 
                 # Check if all required arguments are present
-                required_arguments = signature_data.function.parameters.get('required', [])
-                result, missing_arguments = check_required_arguments(call_data.arguments, required_arguments)
+                required_arguments = signature_data.function.parameters.get(
+                    "required", []
+                )
+                result, missing_arguments = check_required_arguments(
+                    call_data.arguments, required_arguments
+                )
                 if not result:
                     return False, f"Missing required arguments: {missing_arguments}"
 
@@ -39,21 +49,24 @@ def validate_function_call_schema(call, signatures):
     # No matching function signature found
     return False, f"No matching function signature found for function: {call_data.name}"
 
+
 def check_required_arguments(call_arguments, required_arguments):
     missing_arguments = [arg for arg in required_arguments if arg not in call_arguments]
     return not bool(missing_arguments), missing_arguments
 
+
 def validate_enum_value(arg_name, arg_value, enum_values):
     if arg_value not in enum_values:
         raise Exception(
             f"Invalid value '{arg_value}' for parameter {arg_name}. Expected one of {', '.join(map(str, enum_values))}"
         )
 
+
 def validate_argument_type(arg_name, arg_value, arg_schema):
-    arg_type = arg_schema.get('type', None)
+    arg_type = arg_schema.get("type", None)
     if arg_type:
-        if arg_type == 'string' and 'enum' in arg_schema:
-            enum_values = arg_schema['enum']
+        if arg_type == "string" and "enum" in arg_schema:
+            enum_values = arg_schema["enum"]
             if None not in enum_values and enum_values != []:
                 try:
                     validate_enum_value(arg_name, arg_value, enum_values)
@@ -63,20 +76,24 @@ def validate_argument_type(arg_name, arg_value, arg_schema):
 
         python_type = get_python_type(arg_type)
         if not isinstance(arg_value, python_type):
-            raise Exception(f"Type mismatch for parameter {arg_name}. Expected: {arg_type}, Got: {type(arg_value)}")
+            raise Exception(
+                f"Type mismatch for parameter {arg_name}. Expected: {arg_type}, Got: {type(arg_value)}"
+            )
+
 
 def get_python_type(json_type):
     type_mapping = {
-        'string': str,
-        'number': (int, float),
-        'integer': int,
-        'boolean': bool,
-        'array': list,
-        'object': dict,
-        'null': type(None),
+        "string": str,
+        "number": (int, float),
+        "integer": int,
+        "boolean": bool,
+        "array": list,
+        "object": dict,
+        "null": type(None),
     }
     return type_mapping[json_type]
 
+
 def validate_json_data(json_object, json_schema):
     valid = False
     error_message = None
@@ -95,13 +112,13 @@ def validate_json_data(json_object, json_schema):
                     result_json = extract_json_from_markdown(json_object)
                 except Exception as e:
                     error_message = f"JSON decoding error: {e}"
-                    inference_logger.info(f"Validation failed for JSON data: {error_message}")
+                    logger.info(f"Validation failed for JSON data: {error_message}")
                     return valid, result_json, error_message
 
         # Return early if both json.loads and ast.literal_eval fail
         if result_json is None:
             error_message = "Failed to decode JSON data"
-            inference_logger.info(f"Validation failed for JSON data: {error_message}")
+            logger.info(f"Validation failed for JSON data: {error_message}")
             return valid, result_json, error_message
 
         # Validate each item in the list against schema if it's a list
@@ -109,7 +126,7 @@ def validate_json_data(json_object, json_schema):
             for index, item in enumerate(result_json):
                 try:
                     validate(instance=item, schema=json_schema)
-                    inference_logger.info(f"Item {index+1} is valid against the schema.")
+                    logger.info(f"Item {index+1} is valid against the schema.")
                 except ValidationError as e:
                     error_message = f"Validation failed for item {index+1}: {e}"
                     break
@@ -125,8 +142,8 @@ def validate_json_data(json_object, json_schema):
 
     if error_message is None:
         valid = True
-        inference_logger.info("JSON data is valid against the schema.")
+        logger.info("JSON data is valid against the schema.")
     else:
-        inference_logger.info(f"Validation failed for JSON data: {error_message}")
+        logger.info(f"Validation failed for JSON data: {error_message}")
 
     return valid, result_json, error_message