Update agent.py

agent.py

@@ -5,20 +5,33 @@ from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_core.messages import SystemMessage, HumanMessage
 from tools import *
-from typing import TypedDict
+from typing import TypedDict, Optional
 
 tools = [add, multiply, divide, subtract, search_wikipedia]
 
 def build_agent():
+
+    class AgentState(TypedDict):
+        messages: list  # chat history
+        wiki_table: Optional[str] 
+    
     llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
     chat_with_tools = llm.bind_tools(tools)
 
-    def assistant(state: MessagesState):
+    def assistant(state: AgentState):
+        table = state.get("wiki_table")
+        context = []
+    
+        if table:
+            context.append({"role": "system", "content": f"Here is a Wikipedia table you can use:\n\n{table}"})
+    
+        response = chat_with_tools.invoke(context + state["messages"])
         return {
-            "messages": [chat_with_tools.invoke(state["messages"])],
+            "messages": [response],
+            "wiki_table": table,
         }
 
-    def enhancer(state: MessagesState):
+    def enhancer(state: AgentState):
         sys_msg = """
                             You are a helpful assistant tasked with answering questions using a set of tools.  You can reason through problems and take actions using available tools to solve complex tasks.
                             Follow this format strictly:
@@ -35,12 +48,13 @@ def build_agent():
                             If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
                             Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
                         """
-        return {
-            "messages": [sys_msg] + state["messages"]
-        }
+        context = []
+        context.append({"role": "system", "content": sys_msg)
+        
+        return context + state["messages"]
 
     ## The graph
-    builder = StateGraph(MessagesState)
+    builder = StateGraph(AgentState)
     
     # Define nodes: these do the work
     builder.add_node("enhancer", enhancer)