assist/src/agent_assistant_new.py

# -*- coding: utf-8 -*-
"""
增强版TSP助手 - 集成Agent功能
重构版本：模块化设计，降低代码复杂度
"""

import logging
import asyncio
from typing import Dict, Any, List, Optional
from datetime import datetime

from src.agent.agent_assistant_core import TSPAgentAssistantCore
from src.agent.agent_message_handler import AgentMessageHandler
from src.agent.agent_sample_actions import AgentSampleActions

logger = logging.getLogger(__name__)

class TSPAgentAssistant(TSPAgentAssistantCore):
    """TSP Agent助手 - 重构版本"""
    
    def __init__(self, llm_config=None):
        # 初始化核心功能
        super().__init__(llm_config)
        
        # 初始化消息处理器
        self.message_handler = AgentMessageHandler(self)
        
        # 初始化示例动作处理器
        self.sample_actions = AgentSampleActions(self)
        
        logger.info("TSP Agent助手初始化完成（重构版本）")
    
    # ==================== 消息处理功能 ====================
    
    async def process_message_agent(self, message: str, user_id: str = "admin", 
                                  work_order_id: Optional[int] = None, 
                                  enable_proactive: bool = True) -> Dict[str, Any]:
        """使用Agent处理消息"""
        return await self.message_handler.process_message_agent(
            message, user_id, work_order_id, enable_proactive
        )
    
    async def process_conversation_agent(self, conversation_data: Dict[str, Any]) -> Dict[str, Any]:
        """使用Agent处理对话"""
        return await self.message_handler.process_conversation_agent(conversation_data)
    
    async def process_workorder_agent(self, workorder_data: Dict[str, Any]) -> Dict[str, Any]:
        """使用Agent处理工单"""
        return await self.message_handler.process_workorder_agent(workorder_data)
    
    async def process_alert_agent(self, alert_data: Dict[str, Any]) -> Dict[str, Any]:
        """使用Agent处理预警"""
        return await self.message_handler.process_alert_agent(alert_data)
    
    # ==================== 建议功能 ====================
    
    def get_conversation_suggestions(self, context: Dict[str, Any]) -> List[str]:
        """获取对话建议"""
        return self.message_handler.get_conversation_suggestions(context)
    
    def get_workorder_suggestions(self, workorder_data: Dict[str, Any]) -> List[str]:
        """获取工单建议"""
        return self.message_handler.get_workorder_suggestions(workorder_data)
    
    def get_alert_suggestions(self, alert_data: Dict[str, Any]) -> List[str]:
        """获取预警建议"""
        return self.message_handler.get_alert_suggestions(alert_data)
    
    # ==================== 示例动作功能 ====================
    
    async def trigger_sample_actions(self) -> Dict[str, Any]:
        """触发示例动作"""
        return await self.sample_actions.trigger_sample_actions()
    
    async def run_performance_test(self) -> Dict[str, Any]:
        """运行性能测试"""
        return await self.sample_actions.run_performance_test()
    
    # ==================== 兼容性方法 ====================
    
    def get_agent_status(self) -> Dict[str, Any]:
        """获取Agent状态（兼容性方法）"""
        return super().get_agent_status()
    
    def toggle_agent_mode(self, enabled: bool) -> bool:
        """切换Agent模式（兼容性方法）"""
        return super().toggle_agent_mode(enabled)
    
    def start_proactive_monitoring(self) -> bool:
        """启动主动监控（兼容性方法）"""
        return super().start_proactive_monitoring()
    
    def stop_proactive_monitoring(self) -> bool:
        """停止主动监控（兼容性方法）"""
        return super().stop_proactive_monitoring()
    
    def run_proactive_monitoring(self) -> Dict[str, Any]:
        """运行主动监控检查（兼容性方法）"""
        return super().run_proactive_monitoring()
    
    def run_intelligent_analysis(self) -> Dict[str, Any]:
        """运行智能分析（兼容性方法）"""
        return super().run_intelligent_analysis()
    
    def get_action_history(self, limit: int = 50) -> List[Dict[str, Any]]:
        """获取动作执行历史（兼容性方法）"""
        return super().get_action_history(limit)
    
    def clear_execution_history(self) -> Dict[str, Any]:
        """清空执行历史（兼容性方法）"""
        return super().clear_execution_history()
    
    def get_llm_usage_stats(self) -> Dict[str, Any]:
        """获取LLM使用统计（兼容性方法）"""
        return super().get_llm_usage_stats()
    
    # ==================== 高级功能 ====================
    
    async def comprehensive_analysis(self) -> Dict[str, Any]:
        """综合分析 - 结合多个模块的分析结果"""
        try:
            # 运行智能分析
            intelligent_analysis = self.run_intelligent_analysis()
            
            # 运行主动监控
            proactive_monitoring = self.run_proactive_monitoring()
            
            # 运行性能测试
            performance_test = await self.run_performance_test()
            
            # 综合结果
            comprehensive_result = {
                "timestamp": self.execution_history[-1]["timestamp"] if self.execution_history else None,
                "intelligent_analysis": intelligent_analysis,
                "proactive_monitoring": proactive_monitoring,
                "performance_test": performance_test,
                "overall_status": self._determine_overall_status(
                    intelligent_analysis, proactive_monitoring, performance_test
                )
            }
            
            # 记录综合分析
            self._record_execution("comprehensive_analysis", comprehensive_result)
            
            return comprehensive_result
            
        except Exception as e:
            logger.error(f"综合分析失败: {e}")
            return {"error": str(e)}
    
    def _determine_overall_status(self, intelligent_analysis: Dict, 
                                proactive_monitoring: Dict, 
                                performance_test: Dict) -> str:
        """确定整体状态"""
        try:
            # 检查各个模块的状态
            statuses = []
            
            if intelligent_analysis.get("success"):
                statuses.append("intelligent_analysis_ok")
            else:
                statuses.append("intelligent_analysis_error")
            
            if proactive_monitoring.get("success"):
                statuses.append("proactive_monitoring_ok")
            else:
                statuses.append("proactive_monitoring_error")
            
            if performance_test.get("success"):
                statuses.append("performance_test_ok")
            else:
                statuses.append("performance_test_error")
            
            # 根据状态确定整体状态
            if all("ok" in status for status in statuses):
                return "excellent"
            elif any("error" in status for status in statuses):
                return "needs_attention"
            else:
                return "good"
                
        except Exception:
            return "unknown"
    
    async def batch_process_requests(self, requests: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量处理请求"""
        try:
            results = []
            
            for request in requests:
                request_type = request.get("type", "message")
                
                if request_type == "message":
                    result = await self.process_message_agent(
                        request.get("message", ""),
                        request.get("user_id", "admin"),
                        request.get("work_order_id"),
                        request.get("enable_proactive", True)
                    )
                elif request_type == "conversation":
                    result = await self.process_conversation_agent(request)
                elif request_type == "workorder":
                    result = await self.process_workorder_agent(request)
                elif request_type == "alert":
                    result = await self.process_alert_agent(request)
                else:
                    result = {"error": f"未知请求类型: {request_type}"}
                
                results.append(result)
            
            # 记录批量处理
            self._record_execution("batch_process", {
                "request_count": len(requests),
                "results": results
            })
            
            return results
            
        except Exception as e:
            logger.error(f"批量处理请求失败: {e}")
            return [{"error": str(e)} for _ in requests]
    
    def get_system_summary(self) -> Dict[str, Any]:
        """获取系统摘要"""
        try:
            # 获取各种状态信息
            agent_status = self.get_agent_status()
            system_health = self.get_system_health()
            workorders_status = self._check_workorders_status()
            
            # 计算摘要指标
            summary = {
                "timestamp": datetime.now().isoformat(),
                "agent_status": agent_status,
                "system_health": system_health,
                "workorders_status": workorders_status,
                "execution_history_count": len(self.execution_history),
                "llm_usage_stats": self.get_llm_usage_stats(),
                "overall_health_score": system_health.get("health_score", 0)
            }
            
            return summary
            
        except Exception as e:
            logger.error(f"获取系统摘要失败: {e}")
            return {"error": str(e)}
    
    def export_agent_data(self) -> Dict[str, Any]:
        """导出Agent数据"""
        try:
            export_data = {
                "export_timestamp": datetime.now().isoformat(),
                "agent_status": self.get_agent_status(),
                "execution_history": self.execution_history,
                "llm_usage_stats": self.get_llm_usage_stats(),
                "system_summary": self.get_system_summary()
            }
            
            return {
                "success": True,
                "data": export_data,
                "message": "Agent数据导出成功"
            }
            
        except Exception as e:
            logger.error(f"导出Agent数据失败: {e}")
            return {
                "success": False,
                "error": str(e)
            }
    
    def import_agent_data(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """导入Agent数据"""
        try:
            # 验证数据格式
            if not isinstance(data, dict):
                raise ValueError("数据格式不正确")
            
            # 导入执行历史
            if "execution_history" in data:
                self.execution_history = data["execution_history"]
            
            # 其他数据的导入逻辑...
            
            return {
                "success": True,
                "message": "Agent数据导入成功"
            }
            
        except Exception as e:
            logger.error(f"导入Agent数据失败: {e}")
            return {
                "success": False,
                "error": str(e)
            }

# 测试函数
async def main():
    """测试函数"""
    print("🚀 TSP Agent助手测试")
    
    # 创建Agent助手实例
    agent_assistant = TSPAgentAssistant()
    
    # 测试基本功能
    status = agent_assistant.get_agent_status()
    print("Agent状态:", status)
    
    # 测试消息处理
    result = await agent_assistant.process_message_agent("你好，请帮我分析系统状态")
    print("消息处理结果:", result)
    
    # 测试示例动作
    sample_result = await agent_assistant.trigger_sample_actions()
    print("示例动作结果:", sample_result)
    
    # 测试综合分析
    analysis_result = await agent_assistant.comprehensive_analysis()
    print("综合分析结果:", analysis_result)

if __name__ == "__main__":
    asyncio.run(main())