Update agent.py

agent.py

@@ -847,7 +847,17 @@ def call_llm(state):
 tool_node = ToolNode(tools)
 
 
-builder = StateGraph()
+from langgraph.graph import StateGraph
+from typing import TypedDict
+
+class InputSchema(TypedDict):
+    question: str
+
+class OutputSchema(TypedDict):
+    answer: str
+
+builder = StateGraph(input_schema=InputSchema, output_schema=OutputSchema)
+
 
 builder.add_node("call_llm", call_llm)
 builder.add_node("call_tool", tool_node)