Update agent.py

agent.py

@@ -1,5 +1,5 @@
 from typing import Dict, Any, Optional
-from langgraph.graph import Graph
+from langgraph.graph import StateGraph, START, END
 from langgraph.prebuilt import ToolNode
 from langgraph.checkpoint.sqlite import SqliteSaver
 from ai_tools import AITools
@@ -14,7 +14,7 @@ class CustomAgent:
         tool_node = ToolNode([self._route_to_tool])
         
         # 构建工作流
-        workflow = Graph()
+        workflow = StateGraph(State)
         workflow.add_node("tools", tool_node)
         workflow.add_node("process", self._process_result)