X-Agents/agent/app/agent/multi/integration.py

"""
多智能体系统 - 与现有系统集成
"""
import logging
from typing import Optional
from app.llm.factory import LLMFactory
from app.agent.tools.registry import ToolRegistry
from app.agent.memory.session import SessionManager

from .types import create_initial_state
from .graph import create_multi_agent_graph

logger = logging.getLogger(__name__)


class MultiAgentSystem:
    """多智能体系统 - 集成现有组件"""

    def __init__(
        self,
        llm_provider: str = "openai",
        openai_api_key: Optional[str] = None,
        anthropic_api_key: Optional[str] = None,
        max_iterations: int = 3,
        max_tasks: int = 10
    ):
        """
        初始化多智能体系统

        Args:
            llm_provider: LLM 提供商
            openai_api_key: OpenAI API Key
            anthropic_api_key: Anthropic API Key
            max_iterations: 最大迭代次数
            max_tasks: 最大任务数
        """
        # 初始化 LLM Factory
        self.llm_factory = LLMFactory(
            provider=llm_provider,
            openai_api_key=openai_api_key,
            anthropic_api_key=anthropic_api_key
        )

        # 初始化 Tool Registry
        self.tool_registry = ToolRegistry()
        self._register_default_tools()

        # 初始化 Session Manager
        self.session_manager = SessionManager()

        # 配置
        self.max_iterations = max_iterations
        self.max_tasks = max_tasks

        # 图实例（延迟初始化）
        self._graph = None

    def _register_default_tools(self):
        """注册默认工具"""
        try:
            from app.agent.tools.impl import search, calculator, time_tool

            # 安全工具
            self.tool_registry.register(
                name="search",
                func=search.search_web,
                description="Search the web for information",
                security_level="safe"
            )

            self.tool_registry.register(
                name="calculator",
                func=calculator.calculate,
                description="Perform mathematical calculations",
                security_level="safe"
            )

            self.tool_registry.register(
                name="get_current_time",
                func=time_tool.get_current_time,
                description="Get current date and time",
                security_level="safe"
            )

            # 执行代码工具
            try:
                from app.agent.tools.impl import sandbox
                self.tool_registry.register(
                    name="execute_code",
                    func=sandbox.sandbox.execute,
                    description="Execute code in sandbox",
                    security_level="review",
                    require_approval=True
                )
            except ImportError:
                pass

        except ImportError as e:
            logger.warning(f"Failed to import default tools: {e}")

    @property
    def graph(self):
        """获取或创建 LangGraph"""
        if self._graph is None:
            llm = self.llm_factory.get_llm()
            self._graph = create_multi_agent_graph(
                llm=llm,
                tool_registry=self.tool_registry,
                max_iterations=self.max_iterations,
                max_tasks=self.max_tasks
            )
        return self._graph

    async def execute(self, task: str, session_id: str = None) -> dict:
        """
        执行多 Agent 任务

        Args:
            task: 任务描述
            session_id: 会话 ID（可选）

        Returns:
            dict: 执行结果
        """
        # 创建初始状态
        initial_state = create_initial_state(task, session_id)

        try:
            # 执行图
            result = await self.graph.ainvoke(initial_state)

            # 保存到 session
            if session_id:
                self.session_manager.add_message(session_id, "user", task)
                self.session_manager.add_message(
                    session_id,
                    "assistant",
                    result.get("final_output", "")
                )

            return {
                "success": result.get("status") != "failed",
                "output": result.get("final_output", ""),
                "status": result.get("status", "unknown"),
                "task_plan": result.get("task_plan", []),
                "results": result.get("results", {})
            }

        except Exception as e:
            logger.error(f"Multi-agent execution failed: {e}")
            return {
                "success": False,
                "output": f"执行失败: {str(e)}",
                "status": "failed",
                "error": str(e)
            }

    async def execute_simple(self, task: str, session_id: str = None) -> dict:
        """
        执行简单任务（不使用 Supervisor）

        Args:
            task: 任务描述
            session_id: 会话 ID（可选）

        Returns:
            dict: 执行结果
        """
        from .graph import create_simple_graph

        # 创建简单图
        llm = self.llm_factory.get_llm()
        simple_graph = create_simple_graph(llm, self.tool_registry)

        # 创建初始状态
        initial_state = create_initial_state(task, session_id)

        try:
            # 执行图
            result = await simple_graph.ainvoke(initial_state)

            return {
                "success": True,
                "output": result.get("final_output", ""),
                "status": result.get("status", "completed")
            }

        except Exception as e:
            logger.error(f"Simple execution failed: {e}")
            return {
                "success": False,
                "output": f"执行失败: {str(e)}",
                "status": "failed",
                "error": str(e)
            }

    def list_tools(self) -> list:
        """列出所有可用工具"""
        return self.tool_registry.list_tools()


# 全局实例
_global_system: Optional[MultiAgentSystem] = None


def get_multi_agent_system(
    llm_provider: str = "openai",
    openai_api_key: str = None,
    anthropic_api_key: str = None,
    **kwargs
) -> MultiAgentSystem:
    """获取全局多智能体系统实例"""
    global _global_system

    if _global_system is None:
        _global_system = MultiAgentSystem(
            llm_provider=llm_provider,
            openai_api_key=openai_api_key,
            anthropic_api_key=anthropic_api_key,
            **kwargs
        )

    return _global_system