Update my_agent.py

my_agent.py

@@ -9,6 +9,10 @@ from langchain_ollama import ChatOllama
 from langchain.agents.middleware.types import AgentState
 from langchain.messages import HumanMessage, AIMessage, SystemMessage
 
+
+from prompts import system_prompt, qa_system_prompt
+from my_tools import wiki_search, arxiv_search, web_search, visit_webpage, translate_to_english
+
 hf_token = os.getenv("HF_TOKEN")
 
 class AgentResponseState(AgentState):
@@ -34,13 +38,22 @@ class BasicAgent:
         #     # debug=True,
         # )
         tools = [
-            DuckDuckGoSearchRun(),
+            wiki_search,
+            arxiv_search,
+            web_search,
+            visit_webpage,
+            translate_to_english,
         ]
 
-        builder = StateGraph(MessagesState)
+        builder = StateGraph(GraphMessagesState)
 
-        model = create_agent(llm, tools)
+        model = create_agent(
+            llm,
+            tools,
+            system_prompt=system_prompt,
+        )
         builder.add_node("assistant", model)
+        builder.add_node("assistant_qa", self.call_qa)
         builder.add_node("tools", ToolNode(tools))
 
         # Define edges: these determine how the control flow moves
@@ -50,8 +63,20 @@ class BasicAgent:
             # If the latest message requires a tool, route to tools
             # Otherwise, provide a direct response
             tools_condition,
+            {
+                "tools": "tools",
+                END: "assistant_qa",
+            },
         )
         builder.add_edge("tools", "assistant")
+        builder.add_conditional_edges(
+            "assistant_qa",
+            tools_condition,
+            {
+                "tools": "tools",
+                END: END,
+            },
+        )
         self.agent = builder.compile()
 
         print("BasicAgent initialized.")
@@ -64,18 +89,56 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
+    def call_qa(self, graph_state: GraphMessagesState) -> str:
+        # print(f"Calling LLM QA for question: {graph_state['question']}")
+        # print(type(graph_state["messages"]))
+        # print(graph_state["messages"])
+
+        # parsed_messages = [
+        #     {"role": m.type, "content": m.content} for m in graph_state["messages"]
+        # ]
+        parsed_messages = [
+            SystemMessage(content=qa_system_prompt)
+        ]
+        parsed_messages.extend(graph_state["messages"][1:])
+        parsed_messages.append(HumanMessage(content=f"Question: {graph_state['question']}"))
+        print(f"\n\n\n parsed_messages => {parsed_messages}")
+
+        # response = self.llm_qa.invoke(
+        #     {
+        #         "messages": [
+        #             *parsed_messages,
+        #             {
+        #                 "role": "human",
+        #                 "content": graph_state["question"],
+        #             },
+        #         ]
+        #     },
+        #     {"callbacks": [langfuse_handler]},
+        # )
+        response = self.llm_qa.invoke(
+            parsed_messages,
+            {"callbacks": [langfuse_handler]},
+        )
+        print(f"LLAMA 2 -> QA Agent raw response: {response}")
+        return response.model_dump()
+
     def generate_answer(self, question: str) -> str:
         response = self.agent.invoke(
             {
                 "messages": [
                     {
-                        "role": "user",
+                        "role": "system",
+                        "content": system_prompt,
+                    },
+                    {
+                        "role": "human",
                         "content": question,
-                    }
-                ]
-            }
+                    },
+                ],
+                "question": question,
+            },
+            {"callbacks": [langfuse_handler]},
         )
         print(f"Agent raw response: {response}")
-        print(f"response.content => {response['messages']}")
-        print(f"AI response => {response['messages'][-1].content}")
-        return response['messages'][-1].content
+        return response["messages"][-1].content