Merge branch 'simple-rag'

.gitignore

@@ -0,0 +1,7 @@
+.env
+
+# Python
+__pycache__/
+
+# vLLM
+inference_logs/

app.py

@@ -1,43 +1,43 @@
 import os
 import huggingface_hub
 import streamlit as st
-from vllm import LLM, SamplingParams
+from config import config
+from utils import get_assistant_message
+from functioncall import ModelInference
+from prompter import PromptManager
 
-sys_msg = """#Context:
-You are an expert financial advisor named IRAI. You have a comprehensive understanding of finance and investing with experience and expertise in all areas of finance. 
-#Objective:
-Please answer questions as best as possible given your current knowledge. You do not have access to up-to-date current market data. Try to demonstrate analytical depth and showcase ability to integrate complex data into practical advice, but answer the question directly.
-#Style and tone:
-Answer in a friendly and engaging manner representing a top female investment professional working at a leading investment bank.
-#Audience:
-The questions will be asked by top executives and managers of successful startups. Assume the audience is composed of 40 year old males with high wealth and income, high risk appetite with high threshold for volatility.
-#Response:
-Direct answer to question, concise yet insightful."""
 
 @st.cache_resource(show_spinner="Loading model..")
 def init_llm():
-    huggingface_hub.login(token=os.getenv("HF_TOKEN"))
-    llm = LLM(model="InvestmentResearchAI/LLM-ADE-dev")
-    tok = llm.get_tokenizer()
-    tok.eos_token = '<|im_end|>' # Override to use turns
+    huggingface_hub.login(token=config.hf_token, new_session=False)
+    llm = ModelInference(chat_template=config.chat_template)
     return llm
 
 def get_response(prompt):
     try:
+        return llm.generate_function_call(
+            prompt, 
+            config.chat_template, 
+            config.num_fewshot, 
+            config.max_depth
+        )
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+    
+def get_output(context, user_input):
+    try:
+        prompt_schema = llm.prompter.read_yaml_file("prompt_assets/output_sys_prompt.yml")
+        sys_prompt = llm.prompter.format_yaml_prompt(prompt_schema, dict()) + \
+            f"Information:\n{context}"
         convo = [
-            {"role": "system", "content": sys_msg},
-            {"role": "user", "content": prompt},
+            {"role": "system", "content": sys_prompt},
+            {"role": "user", "content": user_input},
         ]
-        llm = init_llm()
-        prompts = [llm.get_tokenizer().apply_chat_template(convo, tokenize=False)]
-        sampling_params = SamplingParams(temperature=0.3, top_p=0.95, max_tokens=500, stop_token_ids=[128009])
-        outputs = llm.generate(prompts, sampling_params)
-        for output in outputs:
-            return output.outputs[0].text
+        response = llm.run_inference(convo)
+        return get_assistant_message(response, config.chat_template, llm.tokenizer.eos_token)
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-
 def main():
     st.title("LLM-ADE 9B Demo")
     
@@ -46,13 +46,21 @@ def main():
     if st.button("Generate"):
         if input_text:
             with st.spinner('Generating response...'):
-                response_text = get_response(input_text)
-                st.write(response_text)
+                agent_resp = get_response(input_text)
+                st.write(get_output(agent_resp, input_text))
         else:
             st.warning("Please enter some text to generate a response.")
 
 llm = init_llm()
 
-if __name__ == "__main__":
-    main()
+def main_headless():
+    while True:
+       input_text = input("Enter your text here: ")
+       agent_resp = get_response(input_text)
+       print('\033[94m' + get_output(agent_resp, input_text) + '\033[0m')
 
+if __name__ == "__main__":
+    if config.headless:
+        main_headless()
+    else:
+        main()

config.py

@@ -0,0 +1,14 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+class Config(BaseSettings):
+    hf_token: str = Field(...)
+    model_path: str = Field("InvestmentResearchAI/LLM-ADE-dev")
+    headless: bool = Field(False, description="Run in headless mode.")
+    
+    chat_template: str = Field("chatml", description="Chat template for prompt formatting")
+    num_fewshot: int | None = Field(None, description="Option to use json mode examples")
+    load_in_4bit: str = Field("False", description="Option to load in 4bit with bitsandbytes")
+    max_depth: int = Field(5, description="Maximum number of recursive iteration")
+
+config = Config(_env_file=".env")

functioncall.py

@@ -0,0 +1,163 @@
+import argparse
+import torch
+import json
+from config import config
+from typing import List, Dict
+from vllm import LLM, SamplingParams
+
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig
+)
+
+import functions
+from prompter import PromptManager
+from validator import validate_function_call_schema
+
+from utils import (
+    inference_logger,
+    get_assistant_message,
+    get_chat_template,
+    validate_and_extract_tool_calls
+)
+
+class ModelInference:
+    def __init__(self, chat_template: str, load_in_4bit: bool = False):
+        self.prompter = PromptManager()
+        self.bnb_config = None
+
+        if load_in_4bit == "True": # Never use this
+            self.bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+            )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            config.model_path,
+            trust_remote_code=True,
+            return_dict=True,
+            quantization_config=self.bnb_config,
+            torch_dtype=torch.float16,
+            attn_implementation="flash_attention_2",
+            device_map="auto",
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(config.model_path, trust_remote_code=True)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "left"
+
+        if self.tokenizer.chat_template is None:
+            print("No chat template defined, getting chat_template...")
+            self.tokenizer.chat_template = get_chat_template(chat_template)
+        
+        inference_logger.info(self.model.config)
+        inference_logger.info(self.model.generation_config)
+        inference_logger.info(self.tokenizer.special_tokens_map)
+
+    def process_completion_and_validate(self, completion, chat_template):
+
+        assistant_message = get_assistant_message(completion, chat_template, self.tokenizer.eos_token)
+
+        if assistant_message:
+            validation, tool_calls, error_message = validate_and_extract_tool_calls(assistant_message)
+
+            if validation:
+                inference_logger.info(f"parsed tool calls:\n{json.dumps(tool_calls, indent=2)}")
+                return tool_calls, assistant_message, error_message
+            else:
+                tool_calls = None
+                return tool_calls, assistant_message, error_message
+        else:
+            inference_logger.warning("Assistant message is None")
+            raise ValueError("Assistant message is None")
+        
+    def execute_function_call(self, tool_call):
+        function_name = tool_call.get("name")
+        function_to_call = getattr(functions, function_name, None)
+        function_args = tool_call.get("arguments", {})
+
+        inference_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
+    
+    def run_inference(self, prompt: List[Dict[str, str]]):
+        inputs = self.tokenizer.apply_chat_template(
+            prompt,
+            add_generation_prompt=True,
+            return_tensors='pt'
+        )
+
+        tokens = self.model.generate(
+            inputs.to(self.model.device),
+            max_new_tokens=1500,
+            temperature=0.8,
+            repetition_penalty=1.1,
+            do_sample=True,
+            eos_token_id=self.tokenizer.eos_token_id
+        )
+        completion = self.tokenizer.decode(tokens[0], skip_special_tokens=False, clean_up_tokenization_space=True)
+        return completion
+
+    def generate_function_call(self, query, chat_template, num_fewshot, max_depth=5):
+        try:
+            depth = 0
+            user_message = f"{query}\nThis is the first turn and you don't have <tool_results> to analyze yet"
+            chat = [{"role": "user", "content": user_message}]
+            tools = functions.get_openai_tools()
+            prompt = self.prompter.generate_prompt(chat, tools, num_fewshot)
+            completion = self.run_inference(prompt)
+
+            def recursive_loop(prompt, completion, depth):
+                nonlocal max_depth
+                tool_calls, assistant_message, error_message = self.process_completion_and_validate(completion, chat_template)
+                prompt.append({"role": "assistant", "content": assistant_message})
+
+                tool_message = f"Agent iteration {depth} to assist with user query: {query}\n"
+                if tool_calls:
+                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+
+                    for tool_call in tool_calls:
+                        validation, message = validate_function_call_schema(tool_call, tools)
+                        if validation:
+                            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}")
+                            except Exception as e:
+                                inference_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)
+                            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})
+
+                    depth += 1
+                    if depth >= max_depth:
+                        print(f"Maximum recursion depth reached ({max_depth}). Stopping recursion.")
+                        return
+
+                    completion = self.run_inference(prompt)
+                    return recursive_loop(prompt, completion, depth)
+                elif error_message:
+                    inference_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})
+
+                    depth += 1
+                    if depth >= max_depth:
+                        print(f"Maximum recursion depth reached ({max_depth}). Stopping recursion.")
+                        return
+
+                    completion = self.run_inference(prompt)
+                    return recursive_loop(prompt, completion, depth)
+                else:
+                    inference_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}")
+            raise e

functions.py

@@ -0,0 +1,262 @@
+import re
+import inspect
+import requests
+import pandas as pd
+import yfinance as yf
+import concurrent.futures
+
+from typing import List
+from bs4 import BeautifulSoup
+from utils import inference_logger
+from langchain.tools import tool
+from langchain_core.utils.function_calling import convert_to_openai_tool
+
+@tool
+def google_search_and_scrape(query: str) -> dict:
+    """
+    Performs a Google search for the given query, retrieves the top search result URLs,
+    and scrapes the text content and table data from those pages in parallel.
+
+    Args:
+        query (str): The search query.
+    Returns:
+        list: A list of dictionaries containing the URL, text content, and table data for each scraped page.
+    """
+    num_results = 2
+    url = 'https://www.google.com/search'
+    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...")
+    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]
+    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 = []
+        for future in concurrent.futures.as_completed(futures):
+            url, html = future.result()
+            soup = BeautifulSoup(html, 'html.parser')
+            paragraphs = [p.text.strip() for p in soup.find_all('p') if p.text.strip()]
+            text_content = ' '.join(paragraphs)
+            text_content = re.sub(r'\s+', ' ', text_content)
+            table_data = [[cell.get_text(strip=True) for cell in row.find_all('td')] for table in soup.find_all('table') for row in table.find_all('tr')]
+            if text_content or table_data:
+                results.append({'url': url, 'content': text_content, 'tables': table_data})
+    return results
+
+@tool
+def get_current_stock_price(symbol: str) -> float:
+  """
+  Get the current stock price for a given symbol.
+
+  Args:
+    symbol (str): The stock symbol.
+
+  Returns:
+    float: The current stock price, or None if an error occurs.
+  """
+  try:
+    stock = yf.Ticker(symbol)
+    # Use "regularMarketPrice" for regular market hours, or "currentPrice" for pre/post market
+    current_price = stock.info.get("regularMarketPrice", stock.info.get("currentPrice"))
+    return current_price if current_price else None
+  except Exception as e:
+    print(f"Error fetching current price for {symbol}: {e}")
+    return None
+
+@tool
+def get_stock_fundamentals(symbol: str) -> dict:
+    """
+    Get fundamental data for a given stock symbol using yfinance API.
+
+    Args:
+        symbol (str): The stock symbol.
+
+    Returns:
+        dict: A dictionary containing fundamental data.
+            Keys:
+                - 'symbol': The stock symbol.
+                - 'company_name': The long name of the company.
+                - 'sector': The sector to which the company belongs.
+                - 'industry': The industry to which the company belongs.
+                - 'market_cap': The market capitalization of the company.
+                - 'pe_ratio': The forward price-to-earnings ratio.
+                - 'pb_ratio': The price-to-book ratio.
+                - 'dividend_yield': The dividend yield.
+                - 'eps': The trailing earnings per share.
+                - 'beta': The beta value of the stock.
+                - '52_week_high': The 52-week high price of the stock.
+                - '52_week_low': The 52-week low price of the stock.
+    """
+    try:
+        stock = yf.Ticker(symbol)
+        info = stock.info
+        fundamentals = {
+            'symbol': symbol,
+            'company_name': info.get('longName', ''),
+            'sector': info.get('sector', ''),
+            'industry': info.get('industry', ''),
+            'market_cap': info.get('marketCap', None),
+            'pe_ratio': info.get('forwardPE', None),
+            'pb_ratio': info.get('priceToBook', None),
+            'dividend_yield': info.get('dividendYield', None),
+            'eps': info.get('trailingEps', None),
+            'beta': info.get('beta', None),
+            '52_week_high': info.get('fiftyTwoWeekHigh', None),
+            '52_week_low': info.get('fiftyTwoWeekLow', None)
+        }
+        return fundamentals
+    except Exception as e:
+        print(f"Error getting fundamentals for {symbol}: {e}")
+        return {}
+
+@tool
+def get_financial_statements(symbol: str) -> dict:
+    """
+    Get financial statements for a given stock symbol.
+
+    Args:
+    symbol (str): The stock symbol.
+
+    Returns:
+    dict: Dictionary containing financial statements (income statement, balance sheet, cash flow statement).
+    """
+    try:
+        stock = yf.Ticker(symbol)
+        financials = stock.financials
+        return financials
+    except Exception as e:
+        print(f"Error fetching financial statements for {symbol}: {e}")
+        return {}
+
+@tool
+def get_key_financial_ratios(symbol: str) -> dict:
+    """
+    Get key financial ratios for a given stock symbol.
+
+    Args:
+    symbol (str): The stock symbol.
+
+    Returns:
+    dict: Dictionary containing key financial ratios.
+    """
+    try:
+        stock = yf.Ticker(symbol)
+        key_ratios = stock.info
+        return key_ratios
+    except Exception as e:
+        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:
+    """
+    Get dividend data for a given stock symbol.
+
+    Args:
+    symbol (str): The stock symbol.
+
+    Returns:
+    pd.DataFrame: DataFrame containing dividend data.
+    """
+    try:
+        stock = yf.Ticker(symbol)
+        dividends = stock.dividends
+        return dividends
+    except Exception as e:
+        print(f"Error fetching dividend data for {symbol}: {e}")
+        return pd.DataFrame()
+
+@tool
+def get_company_news(symbol: str) -> pd.DataFrame:
+    """
+    Get company news and press releases for a given stock symbol.
+
+    Args:
+    symbol (str): The stock symbol.
+
+    Returns:
+    pd.DataFrame: DataFrame containing company news and press releases.
+    """
+    try:
+        news = yf.Ticker(symbol).news
+        return news
+    except Exception as e:
+        print(f"Error fetching company news for {symbol}: {e}")
+        return pd.DataFrame()
+
+@tool
+def get_technical_indicators(symbol: str) -> pd.DataFrame:
+    """
+    Get technical indicators for a given stock symbol.
+
+    Args:
+    symbol (str): The stock symbol.
+
+    Returns:
+    pd.DataFrame: DataFrame containing technical indicators.
+    """
+    try:
+        indicators = yf.Ticker(symbol).history(period="max")
+        return indicators
+    except Exception as e:
+        print(f"Error fetching technical indicators for {symbol}: {e}")
+        return pd.DataFrame()
+
+@tool
+def get_company_profile(symbol: str) -> dict:
+    """
+    Get company profile and overview for a given stock symbol.
+
+    Args:
+    symbol (str): The stock symbol.
+
+    Returns:
+    dict: Dictionary containing company profile and overview.
+    """
+    try:
+        profile = yf.Ticker(symbol).info
+        return profile
+    except Exception as e:
+        print(f"Error fetching company profile for {symbol}: {e}")
+        return {}
+
+def get_openai_tools() -> List[dict]:
+    functions = [
+        google_search_and_scrape,
+        get_current_stock_price,
+        get_company_news,
+        get_company_profile,
+        get_stock_fundamentals,
+        get_financial_statements,
+        get_key_financial_ratios,
+        get_analyst_recommendations,
+        get_dividend_data,
+        get_technical_indicators
+    ]
+
+    tools = [convert_to_openai_tool(f) for f in functions]
+    return tools

prompt_assets/few_shot.json

@@ -0,0 +1,8 @@
+[
+    {
+        "example": "```\nSYSTEM: You are a helpful assistant who has access to functions. Use them if required\n<tools>[\n {\n \"name\": \"calculate_distance\",\n \"description\": \"Calculate the distance between two locations\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"origin\": {\n \"type\": \"string\",\n \"description\": \"The starting location\"\n },\n \"destination\": {\n \"type\": \"string\",\n \"description\": \"The destination location\"\n },\n \"mode\": {\n \"type\": \"string\",\n \"description\": \"The mode of transportation\"\n }\n },\n \"required\": [\n \"origin\",\n \"destination\",\n \"mode\"\n ]\n }\n },\n {\n \"name\": \"generate_password\",\n \"description\": \"Generate a random password\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"length\": {\n \"type\": \"integer\",\n \"description\": \"The length of the password\"\n }\n },\n \"required\": [\n \"length\"\n ]\n }\n }\n]\n\n</tools>\nUSER: Hi, I need to know the distance from New York to Los Angeles by car.\nASSISTANT:\n<tool_call>\n{\"arguments\": {\"origin\": \"New York\",\n \"destination\": \"Los Angeles\", \"mode\": \"car\"}, \"name\": \"calculate_distance\"}\n</tool_call>\n```\n"
+    },
+    {
+        "example": "```\nSYSTEM: You are a helpful assistant with access to functions. Use them if required\n<tools>[\n {\n \"name\": \"calculate_distance\",\n \"description\": \"Calculate the distance between two locations\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"origin\": {\n \"type\": \"string\",\n \"description\": \"The starting location\"\n },\n \"destination\": {\n \"type\": \"string\",\n \"description\": \"The destination location\"\n },\n \"mode\": {\n \"type\": \"string\",\n \"description\": \"The mode of transportation\"\n }\n },\n \"required\": [\n \"origin\",\n \"destination\",\n \"mode\"\n ]\n }\n },\n {\n \"name\": \"generate_password\",\n \"description\": \"Generate a random password\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"length\": {\n \"type\": \"integer\",\n \"description\": \"The length of the password\"\n }\n },\n \"required\": [\n \"length\"\n ]\n }\n }\n]\n\n</tools>\nUSER: Can you help me generate a random password with a length of 8 characters?\nASSISTANT:\n<tool_call>\n{\"arguments\": {\"length\": 8}, \"name\": \"generate_password\"}\n</tool_call>\n```"
+    }
+]

prompt_assets/output_sys_prompt.yml

@@ -0,0 +1,10 @@
+Role: |
+  You are an expert financial advisor named IRAI. 
+  You have a comprehensive understanding of finance and investing with experience and expertise in all areas of finance.
+  You can use information given to you, but do not mention function calls.
+Objective: |
+  Answer questions accurately and truthfully given your current knowledge. Answer the question directly.
+Instructions: |
+  The questions will be asked by top technology executives and CFO of large fintech companies and successful startups.
+  Answer in a friendly and engaging manner representing a top female investment professional working at a leading investment bank.
+  Give a direct answer to question, concise yet insightful.

prompt_assets/sys_prompt.yml

@@ -0,0 +1,38 @@
+Role: |
+  You are a function calling AI agent with self-recursion.
+  You can call only one function at a time and analyse data you get from function response.
+  You are provided with function signatures within <tools></tools> XML tags.
+  The current date is: {date}.
+Objective: |
+  You may use agentic frameworks for reasoning and planning to help with user query.
+  Please call a function and wait for function results to be provided to you in the next iteration.
+  Don't make assumptions about what values to plug into function arguments.
+  Once you have called a function, results will be fed back to you within <tool_response></tool_response> XML tags.
+  Don't make assumptions about tool results if <tool_response> XML tags are not present since function hasn't been executed yet.
+  Analyze the data once you get the results and call another function.
+  At each iteration please continue adding the your analysis to previous summary.
+  Your final response should directly answer the user query with an anlysis or summary of the results of function calls.
+Tools: |
+  Here are the available tools:
+  <tools> {tools} </tools>
+  If the provided function signatures doesn't have the function you must call, you may write executable python code in markdown syntax and call code_interpreter() function as follows:
+  <tool_call>
+  {{"arguments": {{"code_markdown": <python-code>, "name": "code_interpreter"}}}}
+  </tool_call>
+  Make sure that the json object above with code markdown block is parseable with json.loads() and the XML block with XML ElementTree.
+Examples: |
+  Here are some example usage of functions:
+  {examples}
+Schema: |
+  Use the following pydantic model json schema for each tool call you will make:
+  {schema}
+Instructions: |
+  At the very first turn you don't have <tool_results> so you shouldn't not make up the results.
+  Please keep a running summary with analysis of previous function results and summaries from previous iterations.
+  Do not stop calling functions until the task has been accomplished or you've reached max iteration of 10.
+  Calling multiple functions at once can overload the system and increase cost so call one function at a time please.
+  If you plan to continue with analysis, always call another function.
+  For each function call return a valid json object (using doulbe quotes) with function name and arguments within <tool_call></tool_call> XML tags as follows:
+  <tool_call>
+  {{"arguments": <args-dict>, "name": <function-name>}}
+  </tool_call>

prompter.py

@@ -0,0 +1,76 @@
+import datetime
+from pydantic import BaseModel
+from typing import Dict
+from schema import FunctionCall
+from utils import (
+    get_fewshot_examples
+)
+import yaml
+import json
+import os
+
+class PromptSchema(BaseModel):
+    Role: str
+    Objective: str
+    Tools: str
+    Examples: str
+    Schema: str
+    Instructions: str 
+
+class PromptManager:
+    def __init__(self):
+        self.script_dir = os.path.dirname(os.path.abspath(__file__))
+        
+    def format_yaml_prompt(self, prompt_schema: PromptSchema, variables: Dict) -> str:
+        formatted_prompt = ""
+        for field, value in prompt_schema.dict().items():
+            if field == "Examples" and variables.get("examples") is None:
+                continue
+            formatted_value = value.format(**variables)
+            if field == "Instructions":
+                formatted_prompt += f"{formatted_value}"
+            else:
+                formatted_value = formatted_value.replace("\n", " ")
+                formatted_prompt += f"{formatted_value}"
+        return formatted_prompt
+
+    def read_yaml_file(self, file_path: str) -> PromptSchema:
+        with open(file_path, 'r') as file:
+            yaml_content = yaml.safe_load(file)
+        
+        prompt_schema = PromptSchema(
+            Role=yaml_content.get('Role', ''),
+            Objective=yaml_content.get('Objective', ''),
+            Tools=yaml_content.get('Tools', ''),
+            Examples=yaml_content.get('Examples', ''),
+            Schema=yaml_content.get('Schema', ''),
+            Instructions=yaml_content.get('Instructions', ''),
+        )
+        return prompt_schema
+    
+    def generate_prompt(self, user_prompt, tools, num_fewshot=None):
+        prompt_path = os.path.join(self.script_dir, 'prompt_assets', 'sys_prompt.yml')
+        prompt_schema = self.read_yaml_file(prompt_path)
+
+        if num_fewshot is not None:
+            examples = get_fewshot_examples(num_fewshot)
+        else:
+            examples = None
+
+        schema_json = json.loads(FunctionCall.schema_json())
+
+        variables = {
+            "date": datetime.date.today(),
+            "tools": tools,
+            "examples": examples,
+            "schema": schema_json
+        }
+        sys_prompt = self.format_yaml_prompt(prompt_schema, variables)
+
+        prompt = [
+                {'content': sys_prompt, 'role': 'system'}
+            ]
+        prompt.extend(user_prompt)
+        return prompt
+        
+        

requirements.txt

@@ -1,6 +1,131 @@
-streamlit
-transformers
-torch
-vllm
+aiohttp==3.9.5
+aioprometheus==23.12.0
+aiosignal==1.3.1
+altair==5.3.0
+annotated-types==0.6.0
+anyio==4.3.0
+appdirs==1.4.4
+async-timeout==4.0.3
+attrs==23.2.0
+beautifulsoup4==4.12.3
+blinker==1.8.2
+cachetools==5.3.3
+certifi==2024.2.2
+charset-normalizer==3.3.2
+click==8.1.7
+dataclasses-json==0.6.5
+dnspython==2.6.1
+email_validator==2.1.1
+exceptiongroup==1.2.1
+fastapi==0.111.0
+fastapi-cli==0.0.3
+filelock==3.14.0
+frozendict==2.4.4
+frozenlist==1.4.1
+fsspec==2024.3.1
+gitdb==4.0.11
+GitPython==3.1.43
+greenlet==3.0.3
+h11==0.14.0
+html5lib==1.1
+httpcore==1.0.5
+httptools==0.6.1
+httpx==0.27.0
+huggingface-hub==0.23.0
+idna==3.7
+Jinja2==3.1.4
+jsonpatch==1.33
+jsonpointer==2.4
+jsonschema==4.22.0
+jsonschema-specifications==2023.12.1
+langchain==0.1.17
+langchain-community==0.0.37
+langchain-core==0.1.52
+langchain-text-splitters==0.0.1
+langsmith==0.1.54
+lxml==5.2.1
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+marshmallow==3.21.2
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.0.8
+multidict==6.0.5
+multitasking==0.0.11
+mypy-extensions==1.0.0
+networkx==3.3
+ninja==1.11.1.1
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.1.105
+orjson==3.10.3
+packaging==23.2
+pandas==2.2.2
+peewee==3.17.3
+pillow==10.3.0
+protobuf==4.25.3
+psutil==5.9.8
+pyarrow==16.0.0
+pydantic==2.7.1
+pydantic-settings==2.2.1
+pydantic_core==2.18.2
+pydeck==0.9.0
+Pygments==2.18.0
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.1
+quantile-python==1.1
+ray==2.20.0
+referencing==0.35.1
+regex==2024.4.28
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.1
+safetensors==0.4.3
+sentencepiece==0.2.0
+shellingham==1.5.4
+six==1.16.0
+smmap==5.0.1
+sniffio==1.3.1
+soupsieve==2.5
+SQLAlchemy==2.0.30
+starlette==0.37.2
+streamlit==1.34.0
+sympy==1.12
+tenacity==8.3.0
+tokenizers==0.19.1
+toml==0.10.2
+toolz==0.12.1
+torch==2.1.1
+tornado==6.4
+tqdm==4.66.4
+transformers==4.40.2
+triton==2.1.0
+typer==0.12.3
+typing-inspect==0.9.0
+typing_extensions==4.11.0
+tzdata==2024.1
+ujson==5.9.0
+urllib3==2.2.1
+uvicorn==0.29.0
+uvloop==0.19.0
+vllm==0.2.5
+watchdog==4.0.0
+watchfiles==0.21.0
+webencodings==0.5.1
+websockets==12.0
 xformers==0.0.23
-
+yarl==1.9.4
+yfinance==0.2.38

schema.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel
+from typing import List, Dict, Literal, Optional
+
+class FunctionCall(BaseModel):
+    arguments: dict
+    """
+    The arguments to call the function with, as generated by the model in JSON
+    format. Note that the model does not always generate valid JSON, and may
+    hallucinate parameters not defined by your function schema. Validate the
+    arguments in your code before calling your function.
+    """
+
+    name: str
+    """The name of the function to call."""
+
+class FunctionDefinition(BaseModel):
+    name: str
+    description: Optional[str] = None
+    parameters: Optional[Dict[str, object]] = None
+
+class FunctionSignature(BaseModel):
+    function: FunctionDefinition
+    type: Literal["function"]

utils.py

@@ -0,0 +1,149 @@
+import ast
+import os
+import re
+import json
+import logging
+import datetime
+import xml.etree.ElementTree as ET
+
+from logging.handlers import RotatingFileHandler
+
+logging.basicConfig(
+    format="%(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s",
+    datefmt="%Y-%m-%d:%H:%M:%S",
+    level=logging.INFO,
+)
+script_dir = os.path.dirname(os.path.abspath(__file__))
+now = datetime.datetime.now()
+log_folder = os.path.join(script_dir, "inference_logs")
+os.makedirs(log_folder, exist_ok=True)
+log_file_path = os.path.join(
+    log_folder, f"function-calling-inference_{now.strftime('%Y-%m-%d_%H-%M-%S')}.log"
+)
+# Use RotatingFileHandler from the logging.handlers module
+file_handler = RotatingFileHandler(log_file_path, maxBytes=0, backupCount=0)
+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')
+    with open(example_path, 'r') as file:
+        examples = json.load(file)  # Use json.load with the file object, not the file path
+    if num_fewshot > len(examples):
+        raise ValueError(f"Not enough examples (got {num_fewshot}, but there are only {len(examples)} examples).")
+    return examples[:num_fewshot]
+
+def get_chat_template(chat_template):
+    """read chat template from jinja file"""
+    template_path = os.path.join(script_dir, 'chat_templates', f"{chat_template}.j2")
+
+    if not os.path.exists(template_path):
+        print
+        inference_logger.error(f"Template file not found: {chat_template}")
+        return None
+    try:
+        with open(template_path, 'r') as file:
+            template = file.read()
+        return template
+    except Exception as e:
+        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 = []
+    error_message = None
+
+    try:
+        # wrap content in root element
+        xml_root_element = f"<root>{assistant_content}</root>"
+        root = ET.fromstring(xml_root_element)
+
+        # extract JSON data
+        for element in root.findall(".//tool_call"):
+            json_data = None
+            try:
+                json_text = element.text.strip()
+
+                try:
+                    # Prioritize json.loads for better error handling
+                    json_data = json.loads(json_text)
+                except json.JSONDecodeError as json_err:
+                    try:
+                        # Fallback to ast.literal_eval if json.loads fails
+                        json_data = ast.literal_eval(json_text)
+                    except (SyntaxError, ValueError) as eval_err:
+                        error_message = f"JSON parsing failed with both json.loads and ast.literal_eval:\n"\
+                                        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)
+                        continue
+            except Exception as e:
+                error_message = f"Cannot strip text: {e}"
+                inference_logger.error(error_message)
+
+            if json_data is not None:
+                tool_calls.append(json_data)
+                validation_result = True
+
+    except ET.ParseError as err:
+        error_message = f"XML Parse Error: {err}"
+        inference_logger.error(f"XML Parse Error: {err}")
+
+    # Return default values if no valid data is extracted
+    return validation_result, tool_calls, error_message
+
+def extract_json_from_markdown(text):
+    """
+    Extracts the JSON string from the given text using a regular expression pattern.
+    
+    Args:
+        text (str): The input text containing the JSON string.
+        
+    Returns:
+        dict: The JSON data loaded from the extracted string, or None if the JSON string is not found.
+    """
+    json_pattern = r'```json\r?\n(.*?)\r?\n```'
+    match = re.search(json_pattern, text, re.DOTALL)
+    if match:
+        json_string = match.group(1)
+        try:
+            data = json.loads(json_string)
+            return data
+        except json.JSONDecodeError as e:
+            print(f"Error decoding JSON string: {e}")
+    else:
+        print("JSON string not found in the text.")
+    return None

validator.py

@@ -0,0 +1,132 @@
+import ast
+import json
+from jsonschema import validate
+from pydantic import ValidationError
+from utils import inference_logger, extract_json_from_markdown
+from schema import FunctionCall, FunctionSignature
+
+def validate_function_call_schema(call, signatures):
+    try:
+        call_data = FunctionCall(**call)
+    except ValidationError as e:
+        return False, str(e)
+
+    for signature in signatures:
+        try:
+            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():
+                    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)
+                            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)
+                if not result:
+                    return False, f"Missing required arguments: {missing_arguments}"
+
+                return True, None
+        except Exception as e:
+            # Handle validation errors for the function signature
+            return False, str(e)
+
+    # 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)
+    if arg_type:
+        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)
+                except Exception as e:
+                    # Propagate the validation error message
+                    raise Exception(f"Error validating function call: {e}")
+
+        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)}")
+
+def get_python_type(json_type):
+    type_mapping = {
+        '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
+    result_json = None
+
+    try:
+        # Attempt to load JSON using json.loads
+        try:
+            result_json = json.loads(json_object)
+        except json.decoder.JSONDecodeError:
+            # If json.loads fails, try ast.literal_eval
+            try:
+                result_json = ast.literal_eval(json_object)
+            except (SyntaxError, ValueError) as e:
+                try:
+                    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}")
+                    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}")
+            return valid, result_json, error_message
+
+        # Validate each item in the list against schema if it's a list
+        if isinstance(result_json, list):
+            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.")
+                except ValidationError as e:
+                    error_message = f"Validation failed for item {index+1}: {e}"
+                    break
+        else:
+            # Default to validation without list
+            try:
+                validate(instance=result_json, schema=json_schema)
+            except ValidationError as e:
+                error_message = f"Validation failed: {e}"
+
+    except Exception as e:
+        error_message = f"Error occurred: {e}"
+
+    if error_message is None:
+        valid = True
+        inference_logger.info("JSON data is valid against the schema.")
+    else:
+        inference_logger.info(f"Validation failed for JSON data: {error_message}")
+
+    return valid, result_json, error_message