Update agent.py

agent.py

@@ -10,28 +10,16 @@ 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: 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}"})
-    
+    def assistant(state: MessagesState):
         response = chat_with_tools.invoke(context + state["messages"])
         return {
-            "messages": [response],
-            "wiki_table": table,
+            "messages": [response]
         }
 
-    def enhancer(state: AgentState):
+    def enhancer(state: MessagesState):
         sys_msg = """You are a reasoning assistant that can answer complex questions using external tools. You must follow a structured reasoning format and always output a final answer.
                             Always follow this step-by-step structure:
                             
@@ -50,12 +38,11 @@ def build_agent():
                             - If a comma-separated list: apply the above rules to each element.
                             - Only include the final answer, do not include any reasoning steps."""
         return {
-                "messages": [{"role": "system", "content": sys_msg}] + state["messages"],
-                "wiki_table": None
+                "messages": [{"role": "system", "content": sys_msg}] + state["messages"]
             }
 
     ## The graph
-    builder = StateGraph(AgentState)
+    builder = StateGraph(MessagesState)
     
     # Define nodes: these do the work
     builder.add_node("enhancer", enhancer)