468 lines
19 KiB
Python
468 lines
19 KiB
Python
|
"""LangChain Agent service with tool calling capabilities."""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from typing import List, Dict, Any, Optional, AsyncGenerator
|
||
|
|
from langchain_core.tools import BaseTool as LangChainBaseTool
|
||
|
|
from langchain_openai import ChatOpenAI
|
||
|
|
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
||
|
|
from langchain_core.messages import HumanMessage, AIMessage
|
||
|
|
|
||
|
|
from ...utils.logger import get_logger
|
||
|
|
logger = get_logger("agent_service")
|
||
|
|
|
||
|
|
# Try to import langchain_classic with exception handling
|
||
|
|
try:
|
||
|
|
from langchain_classic.agents import AgentExecutor
|
||
|
|
from langchain_classic.agents.tool_calling_agent.base import create_tool_calling_agent
|
||
|
|
LANGCHAIN_CLASSIC_AVAILABLE = True
|
||
|
|
except ImportError:
|
||
|
|
logger.warning("langchain_classic not available. Agent functionality will be disabled.")
|
||
|
|
AgentExecutor = None
|
||
|
|
create_tool_calling_agent = None
|
||
|
|
LANGCHAIN_CLASSIC_AVAILABLE = False
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
|
||
|
|
from .base import BaseTool, ToolRegistry, ToolResult
|
||
|
|
from th_agenter.services.tools import WeatherQueryTool, TavilySearchTool, DateTimeTool
|
||
|
|
from ..postgresql_tool_manager import get_postgresql_tool
|
||
|
|
from ..mysql_tool_manager import get_mysql_tool
|
||
|
|
from ...core.config import get_settings
|
||
|
|
from ..agent_config import AgentConfigService
|
||
|
|
|
||
|
|
|
||
|
|
class LangChainToolWrapper(LangChainBaseTool):
|
||
|
|
"""Wrapper to convert our BaseTool to LangChain tool."""
|
||
|
|
|
||
|
|
name: str = Field(...)
|
||
|
|
description: str = Field(...)
|
||
|
|
base_tool: BaseTool = Field(...)
|
||
|
|
|
||
|
|
def __init__(self, base_tool: BaseTool, **kwargs):
|
||
|
|
super().__init__(
|
||
|
|
name=base_tool.get_name(),
|
||
|
|
description=base_tool.get_description(),
|
||
|
|
base_tool=base_tool,
|
||
|
|
**kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
def _run(self, *args, **kwargs) -> str:
|
||
|
|
"""Synchronous run method."""
|
||
|
|
# Handle both positional and keyword arguments
|
||
|
|
if args:
|
||
|
|
# If positional arguments are provided, convert them to kwargs
|
||
|
|
# based on the tool's parameter names
|
||
|
|
params = self.base_tool.get_parameters()
|
||
|
|
for i, arg in enumerate(args):
|
||
|
|
if i < len(params):
|
||
|
|
kwargs[params[i].name] = arg
|
||
|
|
|
||
|
|
# Run async method in sync context
|
||
|
|
loop = asyncio.new_event_loop()
|
||
|
|
asyncio.set_event_loop(loop)
|
||
|
|
try:
|
||
|
|
result = loop.run_until_complete(self.base_tool.execute(**kwargs))
|
||
|
|
return self._format_result(result)
|
||
|
|
finally:
|
||
|
|
loop.close()
|
||
|
|
|
||
|
|
async def _arun(self, *args, **kwargs) -> str:
|
||
|
|
"""Asynchronous run method."""
|
||
|
|
# Handle both positional and keyword arguments
|
||
|
|
if args:
|
||
|
|
# If positional arguments are provided, convert them to kwargs
|
||
|
|
# based on the tool's parameter names
|
||
|
|
params = self.base_tool.get_parameters()
|
||
|
|
for i, arg in enumerate(args):
|
||
|
|
if i < len(params):
|
||
|
|
kwargs[params[i].name] = arg
|
||
|
|
|
||
|
|
result = await self.base_tool.execute(**kwargs)
|
||
|
|
return self._format_result(result)
|
||
|
|
|
||
|
|
def _format_result(self, result: ToolResult) -> str:
|
||
|
|
"""Format tool result for LangChain."""
|
||
|
|
if result.success:
|
||
|
|
if isinstance(result.result, dict) and "summary" in result.result:
|
||
|
|
return result.result["summary"]
|
||
|
|
return str(result.result)
|
||
|
|
else:
|
||
|
|
return f"Error: {result.error}"
|
||
|
|
|
||
|
|
|
||
|
|
class AgentConfig(BaseModel):
|
||
|
|
"""Agent configuration."""
|
||
|
|
enabled_tools: List[str] = Field(default_factory=lambda: [
|
||
|
|
"calculator", "weather", "search", "datetime", "file", "generate_image", "postgresql_mcp", "mysql_mcp"
|
||
|
|
])
|
||
|
|
max_iterations: int = Field(default=10)
|
||
|
|
temperature: float = Field(default=0.1)
|
||
|
|
system_message: str = Field(
|
||
|
|
default="You are a helpful AI assistant with access to various tools. "
|
||
|
|
"Use the available tools to help answer user questions accurately. "
|
||
|
|
"Always explain your reasoning and the tools you're using."
|
||
|
|
)
|
||
|
|
verbose: bool = Field(default=True)
|
||
|
|
|
||
|
|
|
||
|
|
class AgentService:
|
||
|
|
"""LangChain Agent service with tool calling capabilities."""
|
||
|
|
|
||
|
|
def __init__(self, db_session=None):
|
||
|
|
self.settings = get_settings()
|
||
|
|
self.tool_registry = ToolRegistry()
|
||
|
|
self.config = AgentConfig()
|
||
|
|
self.agent_executor: Optional[AgentExecutor] = None
|
||
|
|
self.db_session = db_session
|
||
|
|
self.config_service = AgentConfigService(db_session) if db_session else None
|
||
|
|
self._initialize_tools()
|
||
|
|
self._load_config()
|
||
|
|
|
||
|
|
def _initialize_tools(self):
|
||
|
|
"""Initialize and register all available tools."""
|
||
|
|
tools = [
|
||
|
|
WeatherQueryTool(),
|
||
|
|
TavilySearchTool(),
|
||
|
|
DateTimeTool(),
|
||
|
|
get_postgresql_tool(), # 使用单例PostgreSQL MCP工具
|
||
|
|
get_mysql_tool() # 使用单例MySQL MCP工具
|
||
|
|
]
|
||
|
|
|
||
|
|
for tool in tools:
|
||
|
|
self.tool_registry.register(tool)
|
||
|
|
logger.info(f"Registered tool: {tool.get_name()}")
|
||
|
|
|
||
|
|
def _load_config(self):
|
||
|
|
"""Load configuration from database if available."""
|
||
|
|
if self.config_service:
|
||
|
|
try:
|
||
|
|
config_dict = self.config_service.get_config_dict()
|
||
|
|
# Update config with database values
|
||
|
|
for key, value in config_dict.items():
|
||
|
|
if hasattr(self.config, key):
|
||
|
|
setattr(self.config, key, value)
|
||
|
|
logger.info("Loaded agent configuration from database")
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning(f"Failed to load config from database, using defaults: {str(e)}")
|
||
|
|
|
||
|
|
def _get_enabled_tools(self) -> List[LangChainToolWrapper]:
|
||
|
|
"""Get list of enabled LangChain tools."""
|
||
|
|
enabled_tools = []
|
||
|
|
|
||
|
|
for tool_name in self.config.enabled_tools:
|
||
|
|
tool = self.tool_registry.get_tool(tool_name)
|
||
|
|
if tool:
|
||
|
|
langchain_tool = LangChainToolWrapper(base_tool=tool)
|
||
|
|
enabled_tools.append(langchain_tool)
|
||
|
|
logger.debug(f"Enabled tool: {tool_name}")
|
||
|
|
else:
|
||
|
|
logger.warning(f"Tool not found: {tool_name}")
|
||
|
|
|
||
|
|
return enabled_tools
|
||
|
|
|
||
|
|
def _create_agent_executor(self) -> AgentExecutor:
|
||
|
|
"""Create LangChain agent executor."""
|
||
|
|
if not LANGCHAIN_CLASSIC_AVAILABLE:
|
||
|
|
raise ValueError("Agent functionality is disabled because langchain_classic is not available.")
|
||
|
|
|
||
|
|
# Get LLM configuration
|
||
|
|
from ...core.llm import create_llm
|
||
|
|
llm = create_llm()
|
||
|
|
|
||
|
|
# Get enabled tools
|
||
|
|
tools = self._get_enabled_tools()
|
||
|
|
|
||
|
|
# Create prompt template
|
||
|
|
prompt = ChatPromptTemplate.from_messages([
|
||
|
|
("system", self.config.system_message),
|
||
|
|
MessagesPlaceholder(variable_name="chat_history"),
|
||
|
|
("human", "{input}"),
|
||
|
|
MessagesPlaceholder(variable_name="agent_scratchpad")
|
||
|
|
])
|
||
|
|
|
||
|
|
# Create agent
|
||
|
|
agent = create_tool_calling_agent(llm, tools, prompt)
|
||
|
|
|
||
|
|
# Create agent executor
|
||
|
|
agent_executor = AgentExecutor(
|
||
|
|
agent=agent,
|
||
|
|
tools=tools,
|
||
|
|
max_iterations=self.config.max_iterations,
|
||
|
|
verbose=self.config.verbose,
|
||
|
|
return_intermediate_steps=True
|
||
|
|
)
|
||
|
|
|
||
|
|
return agent_executor
|
||
|
|
|
||
|
|
async def chat(self, message: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
|
||
|
|
"""Process chat message with agent."""
|
||
|
|
try:
|
||
|
|
logger.info(f"Processing agent chat message: {message[:100]}...")
|
||
|
|
|
||
|
|
# Create agent executor if not exists
|
||
|
|
if not self.agent_executor:
|
||
|
|
self.agent_executor = self._create_agent_executor()
|
||
|
|
|
||
|
|
# Convert chat history to LangChain format
|
||
|
|
langchain_history = []
|
||
|
|
if chat_history:
|
||
|
|
for msg in chat_history:
|
||
|
|
if msg["role"] == "user":
|
||
|
|
langchain_history.append(HumanMessage(content=msg["content"]))
|
||
|
|
elif msg["role"] == "assistant":
|
||
|
|
langchain_history.append(AIMessage(content=msg["content"]))
|
||
|
|
|
||
|
|
# Execute agent
|
||
|
|
result = await self.agent_executor.ainvoke({
|
||
|
|
"input": message,
|
||
|
|
"chat_history": langchain_history
|
||
|
|
})
|
||
|
|
|
||
|
|
# Extract response and intermediate steps
|
||
|
|
response = result["output"]
|
||
|
|
intermediate_steps = result.get("intermediate_steps", [])
|
||
|
|
|
||
|
|
# Format tool calls for response
|
||
|
|
tool_calls = []
|
||
|
|
for step in intermediate_steps:
|
||
|
|
if len(step) >= 2:
|
||
|
|
action, observation = step[0], step[1]
|
||
|
|
tool_calls.append({
|
||
|
|
"tool": action.tool,
|
||
|
|
"input": action.tool_input,
|
||
|
|
"output": observation
|
||
|
|
})
|
||
|
|
|
||
|
|
logger.info(f"Agent response generated successfully with {len(tool_calls)} tool calls")
|
||
|
|
|
||
|
|
return {
|
||
|
|
"response": response,
|
||
|
|
"tool_calls": tool_calls,
|
||
|
|
"success": True
|
||
|
|
}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Agent chat error: {str(e)}", exc_info=True)
|
||
|
|
return {
|
||
|
|
"response": f"Sorry, I encountered an error: {str(e)}",
|
||
|
|
"tool_calls": [],
|
||
|
|
"success": False,
|
||
|
|
"error": str(e)
|
||
|
|
}
|
||
|
|
|
||
|
|
async def chat_stream(self, message: str, chat_history: Optional[List[Dict[str, str]]] = None) -> AsyncGenerator[Dict[str, Any], None]:
|
||
|
|
"""Process chat message with agent (streaming)."""
|
||
|
|
tool_calls = [] # Initialize tool_calls at the beginning
|
||
|
|
try:
|
||
|
|
logger.info(f"Processing agent chat stream: {message[:100]}...")
|
||
|
|
|
||
|
|
# Create agent executor if not exists
|
||
|
|
if not self.agent_executor:
|
||
|
|
self.agent_executor = self._create_agent_executor()
|
||
|
|
|
||
|
|
# Convert chat history to LangChain format
|
||
|
|
langchain_history = []
|
||
|
|
if chat_history:
|
||
|
|
for msg in chat_history:
|
||
|
|
if msg["role"] == "user":
|
||
|
|
langchain_history.append(HumanMessage(content=msg["content"]))
|
||
|
|
elif msg["role"] == "assistant":
|
||
|
|
langchain_history.append(AIMessage(content=msg["content"]))
|
||
|
|
|
||
|
|
# Yield initial status
|
||
|
|
yield {
|
||
|
|
"type": "status",
|
||
|
|
"content": "🤖 开始分析您的请求...",
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.2)
|
||
|
|
|
||
|
|
# Use astream_events for real streaming (if available) or fallback to simulation
|
||
|
|
try:
|
||
|
|
# Try to use streaming events if available
|
||
|
|
async for event in self.agent_executor.astream_events(
|
||
|
|
{"input": message, "chat_history": langchain_history},
|
||
|
|
version="v1"
|
||
|
|
):
|
||
|
|
if event["event"] == "on_tool_start":
|
||
|
|
tool_name = event["name"]
|
||
|
|
yield {
|
||
|
|
"type": "tool_start",
|
||
|
|
"content": f"🔧 正在使用工具: {tool_name}",
|
||
|
|
"tool_name": tool_name,
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.1)
|
||
|
|
|
||
|
|
elif event["event"] == "on_tool_end":
|
||
|
|
tool_name = event["name"]
|
||
|
|
yield {
|
||
|
|
"type": "tool_end",
|
||
|
|
"content": f"✅ 工具 {tool_name} 执行完成",
|
||
|
|
"tool_name": tool_name,
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.1)
|
||
|
|
|
||
|
|
elif event["event"] == "on_chat_model_stream":
|
||
|
|
chunk = event["data"]["chunk"]
|
||
|
|
if hasattr(chunk, 'content') and chunk.content:
|
||
|
|
yield {
|
||
|
|
"type": "content",
|
||
|
|
"content": chunk.content,
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.05)
|
||
|
|
|
||
|
|
except Exception as stream_error:
|
||
|
|
logger.warning(f"Streaming events not available, falling back to simulation: {stream_error}")
|
||
|
|
|
||
|
|
# Fallback: Execute agent and simulate streaming
|
||
|
|
result = await self.agent_executor.ainvoke({
|
||
|
|
"input": message,
|
||
|
|
"chat_history": langchain_history
|
||
|
|
})
|
||
|
|
|
||
|
|
# Extract response and intermediate steps
|
||
|
|
response = result["output"]
|
||
|
|
intermediate_steps = result.get("intermediate_steps", [])
|
||
|
|
|
||
|
|
# Yield tool execution steps
|
||
|
|
tool_calls = []
|
||
|
|
for i, step in enumerate(intermediate_steps):
|
||
|
|
if len(step) >= 2:
|
||
|
|
action, observation = step[0], step[1]
|
||
|
|
tool_calls.append({
|
||
|
|
"tool": action.tool,
|
||
|
|
"input": action.tool_input,
|
||
|
|
"output": observation
|
||
|
|
})
|
||
|
|
|
||
|
|
# Yield tool execution status
|
||
|
|
yield {
|
||
|
|
"type": "tool",
|
||
|
|
"content": f"🔧 使用工具 {action.tool}: {str(action.tool_input)[:100]}...",
|
||
|
|
"tool_name": action.tool,
|
||
|
|
"tool_input": action.tool_input,
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.3)
|
||
|
|
|
||
|
|
yield {
|
||
|
|
"type": "tool_result",
|
||
|
|
"content": f"✅ 工具结果: {str(observation)[:200]}...",
|
||
|
|
"tool_name": action.tool,
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.2)
|
||
|
|
|
||
|
|
# Yield thinking status
|
||
|
|
yield {
|
||
|
|
"type": "thinking",
|
||
|
|
"content": "🤔 正在整理回答...",
|
||
|
|
"done": False
|
||
|
|
}
|
||
|
|
await asyncio.sleep(0.3)
|
||
|
|
|
||
|
|
# Yield the final response in chunks to simulate streaming
|
||
|
|
words = response.split()
|
||
|
|
current_content = ""
|
||
|
|
|
||
|
|
for i, word in enumerate(words):
|
||
|
|
current_content += word + " "
|
||
|
|
|
||
|
|
# Yield every 2-3 words or at the end
|
||
|
|
if (i + 1) % 2 == 0 or i == len(words) - 1:
|
||
|
|
yield {
|
||
|
|
"type": "response",
|
||
|
|
"content": current_content.strip(),
|
||
|
|
"tool_calls": tool_calls if i == len(words) - 1 else [],
|
||
|
|
"done": i == len(words) - 1
|
||
|
|
}
|
||
|
|
|
||
|
|
# Small delay to simulate typing
|
||
|
|
if i < len(words) - 1:
|
||
|
|
await asyncio.sleep(0.05)
|
||
|
|
|
||
|
|
logger.info(f"Agent stream response completed with {len(tool_calls)} tool calls")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Agent chat stream error: {str(e)}", exc_info=True)
|
||
|
|
yield {
|
||
|
|
"type": "error",
|
||
|
|
"content": f"Sorry, I encountered an error: {str(e)}",
|
||
|
|
"done": True
|
||
|
|
}
|
||
|
|
|
||
|
|
def update_config(self, config: Dict[str, Any]):
|
||
|
|
"""Update agent configuration."""
|
||
|
|
try:
|
||
|
|
# Update configuration
|
||
|
|
for key, value in config.items():
|
||
|
|
if hasattr(self.config, key):
|
||
|
|
setattr(self.config, key, value)
|
||
|
|
logger.info(f"Updated agent config: {key} = {value}")
|
||
|
|
|
||
|
|
# Reset agent executor to apply new config
|
||
|
|
self.agent_executor = None
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error updating agent config: {str(e)}", exc_info=True)
|
||
|
|
raise
|
||
|
|
|
||
|
|
def load_config_from_db(self, config_id: Optional[int] = None):
|
||
|
|
"""Load configuration from database."""
|
||
|
|
if not self.config_service:
|
||
|
|
logger.warning("No database session available for loading config")
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
config_dict = self.config_service.get_config_dict(config_id)
|
||
|
|
self.update_config(config_dict)
|
||
|
|
logger.info(f"Loaded configuration from database (ID: {config_id})")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error loading config from database: {str(e)}")
|
||
|
|
raise
|
||
|
|
|
||
|
|
def get_available_tools(self) -> List[Dict[str, Any]]:
|
||
|
|
"""Get list of available tools."""
|
||
|
|
tools = []
|
||
|
|
for tool_name, tool in self.tool_registry._tools.items():
|
||
|
|
tools.append({
|
||
|
|
"name": tool.get_name(),
|
||
|
|
"description": tool.get_description(),
|
||
|
|
"parameters": [{
|
||
|
|
"name": param.name,
|
||
|
|
"type": param.type.value,
|
||
|
|
"description": param.description,
|
||
|
|
"required": param.required,
|
||
|
|
"default": param.default,
|
||
|
|
"enum": param.enum
|
||
|
|
} for param in tool.get_parameters()],
|
||
|
|
"enabled": tool_name in self.config.enabled_tools
|
||
|
|
})
|
||
|
|
return tools
|
||
|
|
|
||
|
|
def get_config(self) -> Dict[str, Any]:
|
||
|
|
"""Get current agent configuration."""
|
||
|
|
return self.config.dict()
|
||
|
|
|
||
|
|
|
||
|
|
# Global agent service instance
|
||
|
|
_agent_service: Optional[AgentService] = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_agent_service(db_session=None) -> Optional[AgentService]:
|
||
|
|
"""Get global agent service instance."""
|
||
|
|
global _agent_service
|
||
|
|
if _agent_service is None:
|
||
|
|
try:
|
||
|
|
_agent_service = AgentService(db_session)
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning(f"Failed to initialize AgentService: {str(e)}. Agent functionality will be disabled.")
|
||
|
|
_agent_service = None
|
||
|
|
elif db_session and _agent_service and not _agent_service.db_session:
|
||
|
|
# Update with database session if not already set
|
||
|
|
_agent_service.db_session = db_session
|
||
|
|
_agent_service.config_service = AgentConfigService(db_session)
|
||
|
|
_agent_service._load_config()
|
||
|
|
return _agent_service
|