init

agent/tool.py

@@ -28,10 +28,12 @@ class ToolAgent(BaseAgent):
         workflow = StateGraph(AgentState)
 
         workflow.add_node("normalize_input", self._normalize_input)
+        workflow.add_node("generate_sql", self._generate_sql)
         workflow.add_node("agent", self._agent_node)
         workflow.add_node("tools", self.tool_node)
 
-        workflow.add_edge("normalize_input", "agent")
+        workflow.add_edge("normalize_input", "generate_sql")
+        workflow.add_edge("generate_sql", "agent")
         workflow.add_edge("tools", "agent")
 
         workflow.add_conditional_edges(
@@ -65,6 +67,11 @@ class ToolAgent(BaseAgent):
 
         return state
 
+    def _generate_sql(self, state: AgentState) -> AgentState:
+        """生成 SQL"""
+        from . import nodes
+        return nodes.generate_sql(state, self.model)
+
     def _should_use_tools(self, state: AgentState) -> str:
         """判断是否需要使用工具"""
         last_message = state.messages[-1]