Use messsage state from langgraph

app.py

@@ -9,12 +9,11 @@ import pandas as pd
 
 from duckduckgo_search import DDGS
 
-from langgraph.graph import StateGraph, START
+from langgraph.graph import MessagesState, StateGraph, START
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode, tools_condition
 
 from langchain_openai import ChatOpenAI
-# from langchain_community.tools import DuckDuckGoSearchRun
 from langchain_core.messages import SystemMessage, HumanMessage, AnyMessage
 from langchain_core.tools import tool
 
@@ -22,9 +21,13 @@ from langchain_core.tools import tool
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-class State(TypedDict):
+# class State(TypedDict):
+#     question: str
+#     messages: Annotated[List[AnyMessage], add_messages]
+
+
+class State(MessagesState):
     question: str
-    messages: Annotated[List[AnyMessage], add_messages]
 
 
 class BasicAgent:
@@ -85,10 +88,8 @@ class BasicAgent:
 
     def assistant(self, state: State):
         print("Assistant invoked. State:", state)
-
         response = self.model_with_tools.invoke(state["messages"])
         print("Assistant response:", response)
-
         return {
             "messages": [response]
         }