99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
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("generate_sql", self._generate_sql)
|
|
workflow.add_node("agent", self._agent_node)
|
|
workflow.add_node("tools", self.tool_node)
|
|
|
|
workflow.add_edge("normalize_input", "generate_sql")
|
|
workflow.add_edge("generate_sql", "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 _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]
|
|
|
|
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")
|
|
}
|