249 lines
10 KiB
Python
249 lines
10 KiB
Python
"""LangChain Agent service with tool calling capabilities."""
|
|
|
|
import asyncio
|
|
from typing import List, Dict, Any, Optional, AsyncGenerator
|
|
from langgraph.prebuilt import create_react_agent
|
|
from langchain_core.messages import HumanMessage, AIMessage
|
|
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
|
|
from loguru import logger
|
|
|
|
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):
|
|
self.settings = get_settings()
|
|
|
|
async def initialize(self, session=None):
|
|
self.tool_registry = ToolRegistry()
|
|
self.config = AgentConfig()
|
|
self.session = session
|
|
self.config_service = AgentConfigService(session) if session else None
|
|
self._initialize_tools()
|
|
await 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()}")
|
|
|
|
async def _load_config(self):
|
|
"""Load configuration from database if available."""
|
|
if self.config_service:
|
|
try:
|
|
config_dict = await 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)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load config from database, using defaults: {str(e)}")
|
|
|
|
def _get_enabled_tools(self) -> List[Any]:
|
|
"""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:
|
|
enabled_tools.append(tool)
|
|
logger.debug(f"Enabled tool: {tool_name}")
|
|
else:
|
|
logger.warning(f"Tool not found: {tool_name}")
|
|
|
|
return enabled_tools
|
|
|
|
async def _create_agent_executor(self) -> Any:
|
|
"""Create LangChain agent executor."""
|
|
from ...core.new_agent import new_llm
|
|
llm = await new_llm()
|
|
tools = self._get_enabled_tools()
|
|
return create_react_agent(llm, tools, prompt=self.config.system_message)
|
|
|
|
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
|
|
agent = await self._create_agent_executor()
|
|
|
|
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"]))
|
|
messages = langchain_history + [HumanMessage(content=message)]
|
|
result = await agent.ainvoke({"messages": messages})
|
|
msgs = result.get("messages") or []
|
|
last = msgs[-1] if msgs else None
|
|
response = (getattr(last, "content", None) or str(last)) if last else str(result)
|
|
logger.info(f"Agent response generated successfully")
|
|
return {
|
|
"response": response,
|
|
"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
|
|
agent = await self._create_agent_executor()
|
|
|
|
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"]))
|
|
messages = langchain_history + [HumanMessage(content=message)]
|
|
yield {"type": "status", "content": "🤖 开始分析您的请求...", "done": False}
|
|
await asyncio.sleep(0.2)
|
|
result = await agent.ainvoke({"messages": messages})
|
|
msgs = result.get("messages") or []
|
|
last = msgs[-1] if msgs else None
|
|
response_content = (getattr(last, "content", None) or str(last)) if last else str(result)
|
|
|
|
# Yield the final response in chunks to simulate streaming
|
|
words = response_content.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")
|
|
|
|
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}")
|
|
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
|
|
_global_agent_service: Optional[AgentService] = None
|
|
|
|
|
|
async def get_agent_service(session=None) -> AgentService:
|
|
"""Get global agent service instance."""
|
|
global _global_agent_service
|
|
if _global_agent_service is None:
|
|
_global_agent_service = AgentService()
|
|
await _global_agent_service.initialize(session)
|
|
elif session and session != _global_agent_service.session:
|
|
# Update with database session if not already set
|
|
_global_agent_service.session = session
|
|
_global_agent_service.config_service = AgentConfigService(session)
|
|
_global_agent_service._load_config()
|
|
return _global_agent_service
|