Update app.py

app.py

@@ -5,7 +5,7 @@ import inspect
 import pandas as pd
 from tools import *
 from typing import TypedDict, Annotated
-from langgraph.graph import START, StateGraph
+from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode, tools_condition
 from langchain_core.messages import AnyMessage, HumanMessage
@@ -33,12 +33,8 @@ tools = [
     search_tool
 ]
 chat_with_tools = chat.bind_tools(tools)
-
-# Generate the AgentState and Agent graph
-class AgentState(TypedDict):
-    messages: Annotated[list[AnyMessage], add_messages]
     
-def assistant(state: AgentState):
+def assistant(state: MessagesState):
     return {
         "messages": [chat_with_tools.invoke(state["messages"])],
     }   
@@ -59,7 +55,7 @@ class BasicAgent:
         return response
         
     def _build_graph(self):
-        graph = StateGraph(AgentState)
+        graph = StateGraph(MessagesState)
         graph.add_node("assistant", assistant)
         graph.add_node("tools", ToolNode(tools))
         graph.add_edge(START, "assistant")