Integrate working langchain ollama model

Dockerfile

@@ -45,6 +45,7 @@ RUN pyenv install ${PYTHON_VERSION} && \
 COPY --chown=1000 ./requirements.txt /tmp/requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /tmp/requirements.txt && \
     pip install flash-attn --no-build-isolation
+RUN ollama pull ${OLLAMA_MODEL}
 
 COPY --chown=1000 src ${HOME}/app
 

src/agents/__init__.py

@@ -0,0 +1,34 @@
+from langchain_community.chat_models import ChatOllama
+from prompts.prompt import rag_agent_prompt
+from agents.functions_agent.base import create_functions_agent
+from langchain.agents import AgentExecutor
+from langchain.memory import ChatMessageHistory
+from functions import get_openai_functions, tools, get_openai_tools
+from config import config
+
+llm = ChatOllama(model = config.ollama_model, temperature = 0.55)
+
+tools_dict = get_openai_tools()
+
+history = ChatMessageHistory()
+
+functions_agent = create_functions_agent(llm=llm, prompt=rag_agent_prompt)
+functions_agent_executor = AgentExecutor(agent=functions_agent, tools=tools, verbose=True, return_intermediate_steps=True)
+
+
+if __name__ == "__main__":
+    while True:
+        try:
+            inp = input("User:")
+            if inp == "/bye":
+                break
+
+            response = functions_agent_executor.invoke({"input": inp, "chat_history": history, "tools" : tools_dict})
+            response['output'] = response['output'].replace("<|im_end|>", "")
+            history.add_user_message(inp)
+            history.add_ai_message(response['output'])
+
+            print(response['output'])
+        except Exception as e:
+            print(e)
+            continue

src/agents/format_scratchpad/functions.py

@@ -0,0 +1,63 @@
+import json
+from typing import List, Sequence, Tuple
+
+from langchain_core.agents import AgentAction, AgentActionMessageLog
+from langchain_core.messages import AIMessage, BaseMessage, FunctionMessage
+
+def _convert_agent_action_to_messages(
+   agent_action: AgentAction, observation: str
+) -> List[BaseMessage]:
+   """Convert an agent action to a message.
+   This code is used to reconstruct the original AI message from the agent action.
+   Args:
+       agent_action: Agent action to convert.
+   Returns:
+       AIMessage that corresponds to the original tool invocation.
+   """
+
+   if isinstance(agent_action, AgentActionMessageLog):
+       return list(agent_action.message_log) + [f"<tool_response>\n{_create_function_message(agent_action, observation)}\n</tool_response>"]
+   else:
+       return [AIMessage(content=agent_action.log)]
+
+def _create_function_message(
+   agent_action: AgentAction, observation: str
+) -> str:
+   """Convert agent action and observation into a function message.
+   Args:
+       agent_action: the tool invocation request from the agent
+       observation: the result of the tool invocation
+   Returns:
+       FunctionMessage that corresponds to the original tool invocation
+   """
+
+   if not isinstance(observation, str):
+       try:
+           content = json.dumps(observation, ensure_ascii=False)
+       except Exception:
+           content = str(observation)
+   else:
+       content = observation
+   tool_response = {
+       "name": agent_action.tool,
+       "content": content,
+   }
+   return json.dumps(tool_response)
+
+def format_to_function_messages(
+   intermediate_steps: Sequence[Tuple[AgentAction, str]],
+) -> List[BaseMessage]:
+   """Convert (AgentAction, tool output) tuples into FunctionMessages.
+   Args:
+       intermediate_steps: Steps the LLM has taken to date, along with observations
+   Returns:
+       list of messages to send to the LLM for the next prediction
+   """
+
+   messages = []
+   for agent_action, observation in intermediate_steps:
+       messages.extend(_convert_agent_action_to_messages(agent_action, observation))
+   return messages
+
+# Backwards compatibility
+format_to_functions = format_to_function_messages

src/agents/functions_agent/base.py

@@ -0,0 +1,48 @@
+from typing import Sequence
+
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts.chat import ChatPromptTemplate
+from langchain_core.runnables import Runnable, RunnablePassthrough
+from langchain_core.tools import BaseTool
+
+from agents.format_scratchpad.functions import (
+   format_to_function_messages,
+)
+from agents.output_parsers.functions import (
+   FunctionsAgentOutputParser,
+)
+
+def create_functions_agent(
+   llm: BaseLanguageModel, prompt: ChatPromptTemplate
+) -> Runnable:
+   """Create an agent that uses function calling.
+
+   Args:
+       llm: LLM to use as the agent. Should work with Nous Hermes function calling,
+           so either be an Nous Hermes based model that supports that or a wrapper of
+           a different model that adds in equivalent support.
+       prompt: The prompt to use. See Prompt section below for more.
+
+   Returns:
+       A Runnable sequence representing an agent. It takes as input all the same input
+       variables as the prompt passed in does. It returns as output either an
+       AgentAction or AgentFinish.
+   """
+   if "agent_scratchpad" not in (
+       prompt.input_variables + list(prompt.partial_variables)
+   ):
+       raise ValueError(
+           "Prompt must have input variable `agent_scratchpad`, but wasn't found."
+           f"Found {prompt.input_variables} instead."
+       )
+   agent = (
+       RunnablePassthrough.assign(
+           agent_scratchpad=lambda x: format_to_function_messages(
+               x["intermediate_steps"]
+           )
+       )
+       | prompt
+       | llm
+       | FunctionsAgentOutputParser()
+   )
+   return agent

src/agents/output_parsers/functions.py

@@ -0,0 +1,77 @@
+import re
+from json import JSONDecodeError
+from typing import List, Union
+
+from langchain_core.agents import AgentAction, AgentActionMessageLog, AgentFinish
+from langchain_core.exceptions import OutputParserException
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+)
+from langchain_core.outputs import ChatGeneration, Generation
+
+from langchain.agents.agent import AgentOutputParser
+from agents.output_parsers.utils import parse_tool_call, check_tool_call
+import ast
+
+class FunctionsAgentOutputParser(AgentOutputParser):
+    """Parses a message into agent action/finish.
+
+    Is meant to be used with a model with Nous Hermes 2 Pro as the base, as it relies on the specific
+    function_call parameter from Nous Research to convey what tools to use.
+
+    If a function_call parameter is passed, then that is used to get
+    the tool and tool input.
+
+    If one is not passed, then the AIMessage is assumed to be the final output.
+    It was add a 
+    """
+
+    @property
+    def _type(self) -> str:
+        return "functions-agent"
+
+    @staticmethod
+    def _parse_ai_message(message: BaseMessage):
+        """Parse an AI message."""
+        if not isinstance(message, AIMessage):
+            raise TypeError(f"Expected an AI message got {type(message)}")
+
+        actions = []
+
+        pattern = re.compile(r"<tool_call>(.*?)</tool_call>", re.DOTALL)
+        try:
+            tool_calls = [parse_tool_call(t.strip()) for t in pattern.findall(message.content)]
+        except:
+            raise OutputParserException(
+                f"Could not parse tool calls from message content: {message.content}. Please ensure that the tool calls are valid JSON."
+            )
+
+        if not tool_calls:
+            return AgentFinish(
+                return_values={"output": message.content}, log=str(message.content)
+            )
+
+        for tool_call in tool_calls:
+            tool_name, tool_input = check_tool_call(tool_call)
+            content_msg = f"\n{message.content}\n" if message.content else "\n"
+            log = f"\nInvoking: `{tool_name}` with `{tool_input}`\n{content_msg}\n"
+            actions.append(AgentActionMessageLog(
+                tool=tool_name,
+                tool_input=tool_input,
+                log=log,
+                message_log=[message],
+            ))
+        
+        return actions
+
+    def parse_result(
+        self, result: List[Generation], *, partial: bool = False
+    ) -> Union[AgentAction, AgentFinish]:
+        if not isinstance(result[0], ChatGeneration):
+            raise ValueError("This output parser only works on ChatGeneration output")
+        message = result[0].message
+        return self._parse_ai_message(message)
+
+    def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
+        raise ValueError("Can only parse messages")

src/agents/output_parsers/utils.py

@@ -0,0 +1,64 @@
+from langchain_core.utils.function_calling import convert_to_openai_function
+from functions import tools
+import re
+import ast
+
+def parse_args(args: str):
+    args = args.strip()
+    args = args.replace("true", "True")
+    args = args.replace("false", "False")
+    args = args.replace("null", "None")
+    args = args.replace("\"", "\"\"\"")
+    i = 0
+    while args[i] != "\"" and args[i] != "\'" and i < len(args) - 1:
+        i += 1
+    args = args[i:]
+    if args[-4:] != "True" and args[-5:] != "False":
+        i = len(args) - 1
+        while args[i] != "\"" and args[i] != "\'" and i > 0:
+            i -= 1
+        args = args[:i + 1]
+    print(args)
+    return ast.literal_eval("{" + args + "}")
+
+def parse_tool_call(call: str):
+    call = call.strip()
+    name: bool = "\"name\": " in call or "\'name\':" in call
+    args: bool = "\"arguments\": " in call or "\'arguments\':" in call
+    if not name:
+        print({"arguments": {}, "name": "missing_function_call"})
+        return {"arguments": {}, "name": "missing_function_call"}
+    if not args:
+        pattern = re.compile(r"\"name\": \"(.*?)\"|\'name\': \'(.*?)\'", re.DOTALL)
+        match = pattern.findall(call)
+        for n in match:
+            if isinstance(n, tuple):
+                n = n[0]
+            print({"arguments": {}, "name": n})
+            return {"arguments": {}, "name": n}
+    args_pattern = re.compile(r"\"arguments\": {(.*?)}|\'arguments\': {(.*?)}", re.DOTALL)
+    args_match = args_pattern.findall(call)
+    for a in args_match:
+        print(a, "\n")
+        print(a[0])
+        args = parse_args(a[0])
+    name_pattern = re.compile(r"\"name\": \"(.*?)\"", re.DOTALL)
+    name_match = name_pattern.findall(call)
+    for n in name_match:
+        if isinstance(n, tuple):
+            n = n[0]
+        print({"arguments": args, "name": n})
+        return {"arguments": args, "name": n}
+
+    
+def check_tool_call(call: dict):
+    global tools
+    tools = [convert_to_openai_function(t) for t in tools]
+    if call["name"] not in [t["name"] for t in tools]:
+        return "handle_tools_error", {"error": {"error": {"name": call["name"]}}}
+    tool = next((t for t in tools if t["name"] == call["name"]), None)
+
+    if set(list(tool["parameters"]["properties"])) != set(list(call["arguments"])):
+        print({"tool_response": {"error": {"expected": list(tool["parameters"]["properties"]), "received": list(call["arguments"])}, "name": call["name"]}})
+        return "handle_tools_error", {"error": {"error": {"expected": list(tool["parameters"]["properties"]), "received": list(call["arguments"])}, "name": call["name"]}}
+    return call["name"], call["arguments"]

src/functioncall.py

@@ -11,6 +11,7 @@ import functions
 from prompter import PromptManager
 from validator import validate_function_call_schema
 from langchain_community.chat_models import ChatOllama
+from langchain_community.llms import Ollama
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 
@@ -23,7 +24,7 @@ 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 = Ollama(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()
         
@@ -69,6 +70,7 @@ class ModelInference:
             add_generation_prompt=True,
             tokenize=False,
         )
+        inputs = inputs.replace("<|begin_of_text|>", "") # Something wrong with the chat template, hotfix
         completion = self.model.invoke(inputs, format='json')
         return completion.content
 
@@ -84,7 +86,6 @@ class ModelInference:
 
             def recursive_loop(prompt, completion, depth):
                 nonlocal max_depth
-                breakpoint()
                 tool_calls, assistant_message, error_message = self.process_completion_and_validate(completion, chat_template)
                 prompt.append({"role": "assistant", "content": assistant_message})
 

src/functions.py

@@ -11,7 +11,7 @@ from bs4 import BeautifulSoup
 from logger import logger
 from openai import AzureOpenAI
 from langchain.tools import tool
-from langchain_core.utils.function_calling import convert_to_openai_tool
+from langchain_core.utils.function_calling import convert_to_openai_tool, convert_to_openai_function
 from config import config
 
 from azure.core.credentials import AzureKeyCredential
@@ -281,20 +281,24 @@ def get_company_profile(symbol: str) -> dict:
         print(f"Error fetching company profile for {symbol}: {e}")
         return {}
 
+tools = [
+    get_analysis,  
+    # 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
+]
+
 def get_openai_tools() -> List[dict]:
-    functions = [
-        get_analysis,  
-        # 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
+    tools_ = [convert_to_openai_tool(f) for f in tools]
+    return tools_
+
+def get_openai_functions() -> List[str]:
+    functions = [convert_to_openai_function(f) for f in tools]
+    return functions

src/prompts/prompt.py

@@ -0,0 +1,17 @@
+import os
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder
+import yaml
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+with open(f"{current_dir}/rag_template.yaml", "r") as yaml_file:
+    templates = yaml.safe_load(yaml_file)
+
+# RAG Agent
+sys_msg_template: str = templates["sys_msg"]
+human_msg_template: str = templates["human_msg"]
+rag_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+    HumanMessagePromptTemplate.from_template(human_msg_template),
+    MessagesPlaceholder(variable_name = "agent_scratchpad")
+])

src/prompts/rag_template.yaml

@@ -0,0 +1,12 @@
+sys_msg: "
+You 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>
+{tools}
+</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:
+<tool_call>
+{{\"arguments\": <args-dict>, \"name\": <function-name>}}
+</tool_call>"
+human_msg: "
+  {input}"