92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
from typing import Dict, Any, List, Optional, cast
|
|
from langchain_core.messages import HumanMessage
|
|
from langchain_core.tools import BaseTool
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.prebuilt import ToolNode
|
|
|
|
from agent.core.base_agent import BaseAgent
|
|
from agent.core.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) -> Any:
|
|
"""构建可使用工具的图"""
|
|
workflow = StateGraph(cast(Any, AgentState))
|
|
|
|
self._add_shared_sql_nodes(workflow)
|
|
workflow.add_node("agent", cast(Any, self._agent_node))
|
|
workflow.add_node("tools", cast(Any, self.tool_node))
|
|
|
|
workflow.set_entry_point("process_input")
|
|
self._add_shared_sql_edges(workflow, start_node="process_input", end_node="agent")
|
|
workflow.add_edge("tools", "agent")
|
|
|
|
workflow.add_conditional_edges(
|
|
"agent",
|
|
self._should_use_tools,
|
|
{
|
|
"tools": "tools",
|
|
"end": END,
|
|
}
|
|
)
|
|
|
|
return cast(Any, 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)
|
|
final_state = self._coerce_state(initial_state, result)
|
|
|
|
return {
|
|
"messages": final_state.messages,
|
|
"context": final_state.sync_context(),
|
|
"tools_used": [tool.name for tool in self.tools],
|
|
"final_step": final_state.current_step,
|
|
}
|