imes/src/agents/expert_agents.py

"""专家Agent模块 - 处理不同类型的问题"""

from abc import ABC, abstractmethod
from typing import Dict, Any, List
from loguru import logger
from openai import OpenAI

from ..models.schemas import QuestionType, AgentContext, RAGResult, ToolCallResult


class BaseExpertAgent(ABC):
    """专家Agent基类"""
    
    def __init__(self, openai_client: OpenAI):
        self.client = openai_client
    
    @abstractmethod
    async def need_rag(self, user_question: str) -> bool:
        """判断是否需要RAG查询"""
        pass
    
    @abstractmethod
    async def need_tool_call(self, user_question: str, context: AgentContext) -> bool:
        """判断是否需要工具调用"""
        pass
    
    @abstractmethod
    async def generate_response(self, context: AgentContext) -> str:
        """生成最终回复"""
        pass


class PageNavigationAgent(BaseExpertAgent):
    """页面跳转引导专家Agent"""
    
    async def need_rag(self, user_question: str) -> bool:
        """页面跳转通常需要查询页面路径信息"""
        return True
    
    async def need_tool_call(self, user_question: str, context: AgentContext) -> bool:
        """页面跳转需要调用跳转工具"""
        return True
    
    async def generate_response(self, context: AgentContext) -> str:
        """生成页面跳转引导回复"""
        try:
            prompt = f"""
你是MES系统的页面跳转助手。用户想要执行某个操作，请根据以下信息生成友好的引导回复：

用户问题：{context.user_question}

RAG查询结果：
{self._format_rag_results(context.rag_results)}

工具调用结果：
{self._format_tool_results(context.tool_call_results)}

请生成一个友好、明确的回复，告诉用户如何完成他们想要的操作。
            """
            
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=300
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            logger.error(f"页面跳转Agent生成回复失败: {e}")
            return "抱歉，页面跳转引导暂时无法使用，请稍后再试。"
    
    def _format_rag_results(self, rag_results: List[RAGResult]) -> str:
        """格式化RAG结果"""
        if not rag_results:
            return "无相关信息"
        
        formatted = []
        for result in rag_results:
            formatted.append(f"- {result.content} (来源: {result.source})")
        return "\n".join(formatted)
    
    def _format_tool_results(self, tool_results: List[ToolCallResult]) -> str:
        """格式化工具调用结果"""
        if not tool_results:
            return "无工具调用"
        
        formatted = []
        for result in tool_results:
            status = "成功" if result.success else f"失败: {result.error_message}"
            formatted.append(f"- 工具 {result.tool_name}: {status}")
        return "\n".join(formatted)


class SystemGuideAgent(BaseExpertAgent):
    """系统使用引导专家Agent"""
    
    async def need_rag(self, user_question: str) -> bool:
        """系统引导总是需要查询知识库"""
        return True
    
    async def need_tool_call(self, user_question: str, context: AgentContext) -> bool:
        """系统引导通常不需要工具调用"""
        return False
    
    async def generate_response(self, context: AgentContext) -> str:
        """生成系统使用引导回复"""
        try:
            prompt = f"""
你是MES系统的使用引导专家。请根据以下信息为用户提供详细的使用指导：

用户问题：{context.user_question}

知识库信息：
{self._format_rag_results(context.rag_results)}

请生成一个详细、易懂的使用指导，包含具体步骤。
            """
            
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.5,
                max_tokens=500
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            logger.error(f"系统引导Agent生成回复失败: {e}")
            return "抱歉，系统使用引导暂时无法使用，请稍后再试。"
    
    def _format_rag_results(self, rag_results: List[RAGResult]) -> str:
        """格式化RAG结果"""
        if not rag_results:
            return "无相关文档"
        
        formatted = []
        for result in rag_results:
            formatted.append(f"- {result.content}")
        return "\n".join(formatted)


class ProductionQAAgent(BaseExpertAgent):
    """生产QA专家Agent"""
    
    async def need_rag(self, user_question: str) -> bool:
        """生产QA可能需要RAG查询历史数据"""
        return True
    
    async def need_tool_call(self, user_question: str, context: AgentContext) -> bool:
        """生产QA可能需要查询数据库"""
        return True
    
    async def generate_response(self, context: AgentContext) -> str:
        """生成生产QA回复"""
        try:
            prompt = f"""
你是MES系统的生产数据分析专家。请根据以下信息回答用户关于生产情况的问题：

用户问题：{context.user_question}

RAG查询结果：
{self._format_rag_results(context.rag_results)}

数据库查询结果：
{self._format_tool_results(context.tool_call_results)}

请生成一个数据驱动的专业回复。
            """
            
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=400
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            logger.error(f"生产QA Agent生成回复失败: {e}")
            return "抱歉，生产数据查询暂时无法使用，请稍后再试。"
    
    def _format_rag_results(self, rag_results: List[RAGResult]) -> str:
        """格式化RAG结果"""
        if not rag_results:
            return "无历史数据"
        
        formatted = []
        for result in rag_results:
            formatted.append(f"- {result.content}")
        return "\n".join(formatted)
    
    def _format_tool_results(self, tool_results: List[ToolCallResult]) -> str:
        """格式化工具调用结果"""
        if not tool_results:
            return "无数据库查询结果"
        
        formatted = []
        for result in tool_results:
            if result.success:
                formatted.append(f"- {result.tool_name}: {result.result}")
            else:
                formatted.append(f"- {result.tool_name}: 查询失败 - {result.error_message}")
        return "\n".join(formatted)


class ChatAgent(BaseExpertAgent):
    """闲聊专家Agent"""
    
    async def need_rag(self, user_question: str) -> bool:
        """闲聊通常不需要RAG"""
        return False
    
    async def need_tool_call(self, user_question: str, context: AgentContext) -> bool:
        """闲聊通常不需要工具调用"""
        return False
    
    async def generate_response(self, context: AgentContext) -> str:
        """生成闲聊回复"""
        try:
            prompt = f"""
你是一个友好的MES系统助手。用户在进行日常闲聊，请以友好、自然的方式回复：

用户问题：{context.user_question}

请生成一个友好、自然的回复，同时可以适当引导用户使用MES系统的功能。
            """
            
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.8,
                max_tokens=200
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            logger.error(f"闲聊Agent生成回复失败: {e}")
            return "你好！我是MES系统助手，很高兴为您服务。有什么我可以帮助您的吗？"


class ExpertAgentFactory:
    """专家Agent工厂类"""
    
    @staticmethod
    def create_agent(question_type: QuestionType, openai_client: OpenAI) -> BaseExpertAgent:
        """根据问题类型创建对应的专家Agent"""
        agent_mapping = {
            QuestionType.PAGE_NAVIGATION: PageNavigationAgent,
            QuestionType.SYSTEM_GUIDE: SystemGuideAgent,
            QuestionType.PRODUCTION_QA: ProductionQAAgent,
            QuestionType.CHAT: ChatAgent,
        }
        
        agent_class = agent_mapping.get(question_type, ChatAgent)
        return agent_class(openai_client)