Migrate to langgraph

Dockerfile

@@ -45,6 +45,8 @@ 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 curl -fsSL https://ollama.com/install.sh | sh
 RUN ollama pull ${OLLAMA_MODEL}
 
 COPY --chown=1000 src ${HOME}/app

requirements.txt

@@ -27,3 +27,4 @@ azure-search-documents==11.6.0b1
 azure-identity==1.16.0
 loguru==0.7.2
 openai==1.30.1
+langgraph==0.0.50

src/agents/__init__.py

@@ -1,29 +1,60 @@
+from typing import Dict
+import langchain_community
+from langchain.agents import AgentExecutor
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
 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 functions import tools, get_openai_tools
 from config import config
 
-llm = ChatOllama(model = config.ollama_model, temperature = 0.55)
+from prompts.prompt import (
+    rag_agent_prompt, 
+    doc_grader_agent_prompt, 
+    router_agent_prompt, 
+    hallucination_grader_agent_prompt, 
+    rephrase_agent_prompt,
+    output_agent_prompt,
+)
 
 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)
-
+def get_agents(llm: langchain_community.chat_models.ChatOllama) -> Dict[str, AgentExecutor]:
+    """Returns all available agents."""
+    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)
+    
+    rephrase_agent = rephrase_agent_prompt | llm | StrOutputParser()
+    router_agent = router_agent_prompt | llm | JsonOutputParser()
+    hallucination_grader_agent = hallucination_grader_agent_prompt | llm | JsonOutputParser()
+    doc_grader_agent = doc_grader_agent_prompt | llm | JsonOutputParser()
+    output_agent = output_agent_prompt | llm | StrOutputParser()
+    
+    return {
+        "function_agent": functions_agent_executor,
+        "doc_grader_agent": doc_grader_agent,
+        "hallucination_grader_agent": hallucination_grader_agent,
+        "router_agent": router_agent,
+        "rephrase_agent": rephrase_agent,
+        "output_agent": output_agent,
+    }
+    
 
 if __name__ == "__main__":
+    from langchain.memory import ChatMessageHistory
+    history = ChatMessageHistory()
+    
+    llm = ChatOllama(model = config.ollama_model, temperature = 0.55)
+    function_agent = get_agents(llm)["function_agent"]
+    
     while True:
         try:
             inp = input("User:")
             if inp == "/bye":
                 break
 
-            response = functions_agent_executor.invoke({"input": inp, "chat_history": history, "tools" : tools_dict})
+            response = function_agent.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'])

src/agents/functions_agent/base.py

@@ -45,4 +45,4 @@ def create_functions_agent(
        | llm
        | FunctionsAgentOutputParser()
    )
-   return agent
+   return agent

src/app.py

@@ -1,71 +1,100 @@
-import os
 from time import time
+from pprint import pprint
 import huggingface_hub
 import streamlit as st
+from typing import Literal, Dict
+from typing_extensions import TypedDict
+import langchain
+from langgraph.graph import END, StateGraph
+from langchain_community.chat_models import ChatOllama
+from logger import logger
+
 from config import config
-from functioncall import ModelInference
+from agents import get_agents, tools_dict
+
+
+class GraphState(TypedDict):
+    """Represents the state of the graph."""
+    question: str
+    rephrased_question: str
+    function_agent_output: str
+    generation: str
 
 
 @st.cache_resource(show_spinner="Loading model..")
-def init_llm():
+def init_agents() -> dict[str, langchain.agents.AgentExecutor]:
     huggingface_hub.login(token=config.hf_token, new_session=False)
-    llm = ModelInference(chat_template=config.chat_template)
-    return llm
-
-
-def function_agent(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 output_agent(context, user_input):
-    """Takes the output of the RAG and generates a final response."""
-    try:
-        config.status.update(label=":bulb: Preparing answer..")
-        script_dir = os.path.dirname(os.path.abspath(__file__))
-        prompt_path = os.path.join(script_dir, "prompt_assets", "output_sys_prompt.yml")
-        prompt_schema = llm.prompter.read_yaml_file(prompt_path)
-        sys_prompt = (
-            llm.prompter.format_yaml_prompt(prompt_schema, dict())
-            + f"Information:\n{context}"
-        )
-        convo = [
-            {"role": "system", "content": sys_prompt},
-            {"role": "user", "content": user_input},
-        ]
-        response = llm.run_inference(convo)
-        return response
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
-
-def query_agent(prompt):
-    """Modifies the prompt and runs inference on it."""
-    try:
-        config.status.update(label=":brain: Starting inference..")
-        script_dir = os.path.dirname(os.path.abspath(__file__))
-        prompt_path = os.path.join(script_dir, "prompt_assets", "output_sys_prompt.yml")
-        prompt_schema = llm.prompter.read_yaml_file(prompt_path)
-        sys_prompt = llm.prompter.format_yaml_prompt(prompt_schema, dict())
-        convo = [
-            {"role": "system", "content": sys_prompt},
-            {"role": "user", "content": prompt},
-        ]
-        response = llm.run_inference(convo)
-        return response
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
-
-
-def get_response(input_text: str):
-    """This is the main function that generates the final response."""
-    agent_resp = function_agent(input_text)
-    output = output_agent(agent_resp, input_text)
-    return output
-
+    llm = ChatOllama(model = config.ollama_model, temperature = 0.8)
+    return get_agents(llm)
+
+
+# Nodes -----------------------------------------------------------------------
+
+def question_node(state: GraphState) -> Dict[str, str]:
+    """
+    Generate a question for the function agent.
+    """
+    logger.info("Generating question for function agent")
+    config.status.update(label=":question: Breaking down question")
+    question = state["question"]
+    logger.info(f"Original question: {question}")
+    rephrased_question = agents["rephrase_agent"].invoke({"question": question})
+    logger.info(f"Rephrased question: {rephrased_question}")
+    return {"rephrased_question": rephrased_question}
+
+def function_agent_node(state: GraphState) -> Literal["finished"]:
+    """
+    Call the function agent
+    """
+    logger.info("Calling function agent")
+    question = state["rephrased_question"]
+    response = agents["function_agent"].invoke({"input": question, "tools": tools_dict}).get("output")
+    config.status.update(label=":brain: Analysing data..")
+    logger.info(f"Function agent output: {response}")
+    return {"function_agent_output": response}
+
+def output_node(state: GraphState) -> Dict[str, str]:
+    """
+    Generate the final output
+    """
+    logger.info("Generating output")
+    config.status.update(label=":bulb: Preparing response..")
+    generation = agents["output_agent"].invoke({"context": state["function_agent_output"],
+                                                "question": state["rephrased_question"]})
+    return {"generation": generation}
+
+# Conditional Edge ------------------------------------------------------------
+
+def route_question(state: GraphState) -> Literal["vectorstore", "websearch"]:
+    """
+    Route quesition to web search or RAG
+    """
+    logger.info("Routing question")
+    config.state.update(label=":chart_with_upwards_trend: Routing question")
+    question = state["question"]
+    logger.info(f"Question: {question}")
+    source = agents["router_agent"].invoke({"question": question})
+    logger.info(source)
+    logger.info(source["datasource"])
+    if source["datasource"] == "vectorstore":
+        return "vectorstore"
+    elif source["datasource"] == "websearch":
+        return "websearch"
+
+
+# Graph -----------------------------------------------------------------------
+
+workflow = StateGraph(GraphState)
+workflow.add_node("question_rephrase", question_node)
+workflow.add_node("function_agent", function_agent_node)
+workflow.add_node("output_node", output_node)
+
+workflow.set_entry_point("question_rephrase")
+workflow.add_edge("question_rephrase", "function_agent")
+workflow.add_edge("function_agent", "output_node")
+workflow.set_finish_point("output_node")
+
+flow = workflow.compile()
 
 def main():
     st.title("LLM-ADE 9B Demo")
@@ -76,7 +105,10 @@ def main():
         if input_text:
             with st.status("Generating response...") as status:
                 config.status = status
-                st.write(get_response(input_text))
+                for output in flow.stream({"question": input_text}):
+                    for key, value in output.items():
+                        pprint(f"Finished running: {key}")
+                st.write(value["generation"])
                 config.status.update(label="Finished!", state="complete", expanded=True)
         else:
             st.warning("Please enter some text to generate a response.")
@@ -84,17 +116,18 @@ def main():
 
 def main_headless(prompt: str):
     start = time()
-    print("\033[94m" + get_response(prompt) + "\033[0m")
+    for output in flow.stream({"question": prompt}):
+        for key, value in output.items():
+            pprint(f"Finished running: {key}")
+    print("\033[94m" + value["generation"] + "\033[0m")
     print(f"Time taken: {time() - start:.2f}s\n" + "-" * 20)
 
 
-llm = init_llm()
-
+agents = init_agents()
 
 if __name__ == "__main__":
     if config.headless:
         import fire
-
         fire.Fire(main_headless)
     else:
         main()

src/functions.py

@@ -230,6 +230,7 @@ def get_dividend_data(symbol: str) -> pd.DataFrame:
 def get_company_news(symbol: str) -> pd.DataFrame:
     """
     Get company news and press releases for a given stock symbol.
+    If you use this, you must use the articles to backup your analysis.
 
     Args:
     symbol (str): The stock symbol.
@@ -285,7 +286,7 @@ tools = [
     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/prompts/answer_grader_template.yaml

@@ -0,0 +1,10 @@
+sys_msg: "
+You are a grader assessing whether an 
+answer is useful to resolve a question. Give a binary score 'yes' or 'no' to indicate whether the answer is 
+useful to resolve a question. Provide the binary score as a JSON with a single key 'score' and no preamble or explanation."
+human_msg: "
+Here is the answer:
+\n ------- \n
+{generation}
+\n ------- \n
+Here is the question: {question}"

src/prompts/doc_grader_template.yaml

@@ -0,0 +1,11 @@
+sys_msg: "
+You are a grader assessing relevance 
+of a retrieved document to a user question. If the document contains keywords related to the user question, 
+grade it as relevant. It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
+Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question. \n
+Provide the binary score as a JSON with a single key 'score' and no premable or explanation."
+human_msg: "
+Here is the retrieved document: \n\n {document}
+
+Here is the user question: {question} \n 
+"

src/prompts/hallucination_grader_template.yaml

@@ -0,0 +1,6 @@
+sys_msg: "
+You are a grader assessing whether 
+an answer is grounded in / supported by a set of facts. Give a binary 'yes' or 'no' score to indicate 
+whether the answer is grounded in / supported by a set of facts. Provide the binary score as a JSON with a 
+single key 'score' and no preamble or explanation.
+"

src/prompts/output_agent_template.yaml

@@ -0,0 +1,30 @@
+sys_msg: "
+#ROLE:
+You are a financial expert named IRAI with experience and expertise in stocks and cryptocurrency.
+You have a comprehensive understanding of finance, investing, and quantiative analysis; you are a thought leader in these fields.
+You will be fielding questions about finance from top executives and managers of successful startups.
+#OBJECTIVE:
+Answer questions as best as accurately as possible given your current knowledge and all available information.
+Your answers should demonstrate expert insight and indepth analysis. 
+Use analysis, information, and data from the assistant message to form your answers.
+#INSTRUCTIONS:
+Try to incorprate the numerical data from the function calls to support your analysis, do not state any other numerical data in your answer.
+Use all the available information to answer the question, including the data and information from the function calls.
+Do not mention any function calls, such as \"get_analysis\", \"get_current_stock_price\", or \"get_key_financial_ratios\" in your answer.
+Give a direct answer to question, concise yet insightful. Your answer should aim to impress the user with your insight.
+Do not give additional instructions related to seeking financial advisor or professional in your answer.
+Do not ask any followup questions in your answer.
+Do not ask for any additional information in your answer.
+Do not mention Morgan Stanley in your answer.
+#Style and tone:
+Please answer in a friendly and engaging manner representing a top female crypto expert.
+Do not be formal, answer as you would in a conversation with a friend.
+Keep your response short unless the question requires a longer response.
+#Audience:
+You are talking to a high net worth individual with high risk tolerance and lots of investing experience - they do not need traditional, generic financial advice.
+\n ------- \n
+Your thoughts are as follows:
+{context}
+\n ------- \n"
+human_msg: "
+My question: {question}"

src/prompts/prompt.py

@@ -4,14 +4,90 @@ from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTem
 import yaml
 
 current_dir = os.path.dirname(os.path.abspath(__file__))
+
+
+# Function Agent
 with open(f"{current_dir}/rag_template.yaml", "r") as yaml_file:
-    templates = yaml.safe_load(yaml_file)
+    rag_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = rag_templates["sys_msg"]
+human_msg_template: str = rag_templates["human_msg"]
 
-# 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")
 ])
+
+
+# Question Rephrase Agent
+with open(f"{current_dir}/question_rephrase_template.yaml", "r") as yaml_file:
+    rephrase_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = rephrase_templates["sys_msg"]
+human_msg_template: str = rephrase_templates["human_msg"]
+
+rephrase_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+    HumanMessagePromptTemplate.from_template(human_msg_template),
+])
+
+
+# Router Agent
+with open(f"{current_dir}/router_template.yaml", "r") as yaml_file:
+    router_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = router_templates["sys_msg"]
+
+router_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+])
+
+
+# Document Grader Agent
+with open(f"{current_dir}/doc_grader_template.yaml", "r") as yaml_file:
+    grader_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = grader_templates["sys_msg"]
+human_msg_template: str = grader_templates["human_msg"]
+
+doc_grader_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+    HumanMessagePromptTemplate.from_template(human_msg_template),
+])
+
+
+# Hallucination Grader Agent
+with open(f"{current_dir}/hallucination_grader_template.yaml", "r") as yaml_file:
+    hallucination_grader_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = hallucination_grader_templates["sys_msg"]
+
+hallucination_grader_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+])
+
+
+# Answer Grader Agent
+with open(f"{current_dir}/answer_grader_template.yaml", "r") as yaml_file:
+    answer_grader_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = answer_grader_templates["sys_msg"]
+human_msg_template: str = answer_grader_templates["human_msg"]
+
+answer_grader_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+    HumanMessagePromptTemplate.from_template(human_msg_template),
+])
+
+# Output Agent
+with open(f"{current_dir}/output_agent_template.yaml", "r") as yaml_file:
+    output_agent_templates = yaml.safe_load(yaml_file)
+
+sys_msg_template: str = output_agent_templates["sys_msg"]
+human_msg_template: str = output_agent_templates["human_msg"]
+
+output_agent_prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(sys_msg_template),
+    HumanMessagePromptTemplate.from_template(human_msg_template),
+])

src/prompts/question_rephrase_template.yaml

@@ -0,0 +1,17 @@
+sys_msg: "
+#ROLE:
+You are a crypto expert named IRAI with experience and expertise in cryptocurrency.
+You are taking questions from a group of high net worth individuals with high risk tolerance and lots of investing experience.
+You have an intern hired to assist you in finding information regarding crypto and digital assets.
+#OBJECTIVE:
+Identify and infer what information is required to answer the question
+Your rephrased question is directed at the intern to help them understand the user's question better.
+#INSTRUCTIONS:
+Rephrase the question to include more context and specifics such that the intern can answer the user's question.
+Rephrasing is necessary when the question is too vague or broad, you will need to infer context.
+If the question has no context, use your expertise to create an intelligent question regarding the subject in the original question.
+If rephrasing is unncecessary, return the original question.
+Never ask for additional information in your rephrased question.
+The rephrased question should be concise and clear. Do not be conversational. Ask informally."
+human_msg: "
+{question}"

src/prompts/router_template.yaml

@@ -0,0 +1,7 @@
+sys_msg: "
+You are an expert at routing a 
+user question to a vectorstore or web search. Use the vectorstore for questions on LLM  agents, 
+prompt engineering, and adversarial attacks. You do not need to be stringent with the keywords 
+in the question related to these topics. Otherwise, use web-search. Give a binary choice 'web_search' 
+or 'vectorstore' based on the question. Return the a JSON with a single key 'datasource' and 
+no premable or explanation. Question to route: {question}:"