86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
from typing import Dict, Any, List, Optional
|
|
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
|
|
from langchain_core.tools import BaseTool
|
|
from langgraph.graph import StateGraph, END
|
|
from langgraph.prebuilt import ToolNode
|
|
from .base_agent import BaseAgent, AgentState
|
|
from tools.calculator import CalculatorTool
|
|
from tools.web_search import WebSearchTool
|
|
|
|
|
|
class ToolAgent(BaseAgent):
|
|
"""Agent that can use tools to accomplish tasks"""
|
|
|
|
def __init__(self, model_name: str = None, tools: List[BaseTool] = None):
|
|
# Initialize with default tools if none provided
|
|
if tools is None:
|
|
tools = [CalculatorTool(), WebSearchTool()]
|
|
|
|
self.tools = tools
|
|
self.tool_node = ToolNode(tools)
|
|
super().__init__(model_name)
|
|
|
|
def _build_graph(self) -> StateGraph:
|
|
"""Build tool-using graph"""
|
|
workflow = StateGraph(AgentState)
|
|
|
|
# Add nodes
|
|
workflow.add_node("agent", self._agent_node)
|
|
workflow.add_node("tools", self.tool_node)
|
|
|
|
# Define edges
|
|
workflow.add_edge("tools", "agent")
|
|
|
|
# Conditional routing
|
|
workflow.add_conditional_edges(
|
|
"agent",
|
|
self._should_use_tools,
|
|
{
|
|
"tools": "tools",
|
|
"end": END,
|
|
}
|
|
)
|
|
|
|
# Set entry point
|
|
workflow.set_entry_point("agent")
|
|
|
|
return workflow.compile()
|
|
|
|
def _agent_node(self, state: AgentState) -> AgentState:
|
|
"""Agent node that decides whether to use tools"""
|
|
# Bind tools to the model
|
|
model_with_tools = self.model.bind_tools(self.tools)
|
|
|
|
# Get the last message
|
|
if state.messages:
|
|
response = model_with_tools.invoke(state.messages)
|
|
state.messages.append(response)
|
|
|
|
return state
|
|
|
|
def _should_use_tools(self, state: AgentState) -> str:
|
|
"""Determine if tools should be used"""
|
|
last_message = state.messages[-1]
|
|
|
|
# If the last message has tool calls, route to tools
|
|
if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
|
|
return "tools"
|
|
|
|
# Otherwise, end the workflow
|
|
return "end"
|
|
|
|
def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
|
|
"""Run the tool-using agent"""
|
|
initial_state = AgentState(
|
|
messages=[HumanMessage(content=user_input)],
|
|
context=kwargs
|
|
)
|
|
|
|
result = self.graph.invoke(initial_state)
|
|
|
|
return {
|
|
"messages": result.messages,
|
|
"context": result.context,
|
|
"tools_used": [tool.name for tool in self.tools],
|
|
"final_step": result.current_step
|
|
} |