Update agent.py

agent.py

@@ -850,13 +850,13 @@ tool_node = ToolNode(tools)
 from langgraph.graph import StateGraph
 from typing import TypedDict
 
-class InputSchema(TypedDict):
-    question: str
+# Define the state schema
+class AgentState(TypedDict):
+    input: str
+    result: str
 
-class OutputSchema(TypedDict):
-    answer: str
-
-builder = StateGraph(input_schema=InputSchema, output_schema=OutputSchema)
+# Create the state graph builder with state_schema
+builder = StateGraph(AgentState)
 
 
 builder.add_node("call_llm", call_llm)