hxf/backend/th_agenter/services/agent/agent_service.py

"""LangChain Agent service with tool calling capabilities."""

import asyncio
from typing import List, Dict, Any, Optional, AsyncGenerator
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
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, db_session=None):
        self.settings = get_settings()
        self.tool_registry = ToolRegistry()
        self.config = AgentConfig()
        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[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
    
    def _create_agent_executor(self) -> Any:
        """Create LangChain agent executor."""
        # 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}"),
        ])
        
        # Create agent using new LangChain 1.0+ API
        agent = create_agent(
            llm=llm,
            tools=tools,
            prompt=prompt
        )
        
        return agent
    
    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 = 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 agent.ainvoke({
                "input": message,
                "chat_history": langchain_history
            })
            
            logger.info(f"Agent response generated successfully")
            
            return {
                "response": result["output"] if isinstance(result, dict) and "output" in result else str(result),
                "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 = 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)
            
            # Generate response
            result = await agent.ainvoke({
                "input": message,
                "chat_history": langchain_history
            })
            
            response_content = result["output"] if isinstance(result, dict) and "output" in result 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
_agent_service: Optional[AgentService] = None


def get_agent_service(db_session=None) -> AgentService:
    """Get global agent service instance."""
    global _agent_service
    if _agent_service is None:
        _agent_service = AgentService(db_session)
    elif db_session 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