Update agent.py

agent.py

@@ -9,20 +9,17 @@ from typing import TypedDict
 
 tools = [add, substract, multiply, divide, web_search]
 
-class AgentState(TypedDict):
-    messages: str
-
 def build_agent():
     llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
     chat_with_tools = llm.bind_tools(tools)
 
-    def assistant(state: AgentState):
+    def assistant(state: MessagesState):
         return {
             "messages": [chat_with_tools.invoke(state["messages"])],
         }
 
     ## The graph
-    builder = StateGraph(AgentState)
+    builder = StateGraph(MessagesState)
     
     # Define nodes: these do the work
     builder.add_node("assistant", assistant)