init

agent/tool.py

@@ -0,0 +1,91 @@
+from typing import Dict, Any, List, Optional
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode
+
+from .graph import BaseAgent
+from .state import AgentState
+from tools.calculator import CalculatorTool
+from tools.web_search import WebSearchTool
+from tools.rest_api_tool import RestApiTool
+from tools.sr_api_tool import SrApiQueryTool
+
+
+class ToolAgent(BaseAgent):
+    """可使用工具完成任务的代理"""
+
+    def __init__(self, model_section: Optional[str] = None, tools: List[BaseTool] = None):
+        if tools is None:
+            tools = [CalculatorTool(), WebSearchTool(), RestApiTool(), SrApiQueryTool()]
+
+        self.tools = tools
+        self.tool_node = ToolNode(tools)
+        super().__init__(model_section)
+
+    def _build_graph(self) -> StateGraph:
+        """构建可使用工具的图"""
+        workflow = StateGraph(AgentState)
+
+        workflow.add_node("normalize_input", self._normalize_input)
+        workflow.add_node("agent", self._agent_node)
+        workflow.add_node("tools", self.tool_node)
+
+        workflow.add_edge("normalize_input", "agent")
+        workflow.add_edge("tools", "agent")
+
+        workflow.add_conditional_edges(
+            "agent",
+            self._should_use_tools,
+            {
+                "tools": "tools",
+                "end": END,
+            }
+        )
+
+        workflow.set_entry_point("normalize_input")
+
+        return workflow.compile()
+
+    def _agent_node(self, state: AgentState) -> AgentState:
+        """决定是否调用工具的代理节点"""
+        model_with_tools = self.model.bind_tools(self.tools)
+
+        if state.messages:
+            try:
+                response = model_with_tools.invoke(state.messages)
+            except Exception as e:
+                error_text = str(e)
+                if "tool choice" in error_text and "auto" in error_text:
+                    fallback_model = self.model.bind_tools(self.tools, tool_choice="none")
+                    response = fallback_model.invoke(state.messages)
+                else:
+                    raise
+            state.messages.append(response)
+
+        return state
+
+    def _should_use_tools(self, state: AgentState) -> str:
+        """判断是否需要使用工具"""
+        last_message = state.messages[-1]
+
+        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
+            return "tools"
+
+        return "end"
+
+    def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
+        """运行工具型代理"""
+        initial_state = AgentState(
+            messages=[HumanMessage(content=user_input)],
+            context=kwargs
+        )
+
+        result = self.graph.invoke(initial_state)
+
+        return {
+            "messages": result.get("messages", []),
+            "context": result.get("context", {}),
+            "tools_used": [tool.name for tool in self.tools],
+            "final_step": result.get("current_step", "unknown")
+        }