X-Agents/agent/app/agent/core/supervisor.py

"""
Supervisor - 多智能体调度器
"""
import asyncio
from typing import List, Dict, Any
from app.agent.core.agent import AgentCore


class Supervisor:
    """多智能体调度器"""

    def __init__(self, supervisor_agent: AgentCore, members: List[AgentCore], strategy: str = "parallel"):
        """
        初始化调度器

        Args:
            supervisor_agent: 主智能体
            members: 子智能体列表
            strategy: 调度策略 (parallel/sequential)
        """
        self.supervisor = supervisor_agent
        self.members = members
        self.strategy = strategy

    async def run(self, task: str, user_id: int, session_id: str) -> Dict[str, Any]:
        """
        执行多智能体协作

        Args:
            task: 用户任务
            user_id: 用户 ID
            session_id: 会话 ID

        Returns:
            Dict: 包含主响应和子任务结果
        """
        # 1. 任务分解
        subtasks = await self._decompose_task(task)

        # 2. 分配任务
        if self.strategy == "parallel":
            results = await self._dispatch_parallel(subtasks, user_id, session_id)
        else:
            results = await self._dispatch_sequential(subtasks, user_id, session_id)

        # 3. 汇总结果
        final_result = await self._aggregate(results)

        return {
            "main_response": final_result,
            "subtask_results": results,
            "strategy": self.strategy
        }

    async def _decompose_task(self, task: str) -> List[Dict[str, str]]:
        """任务分解"""
        # 调用 LLM 分解任务
        prompt = f"""分解以下任务为子任务，返回 JSON 数组格式:
任务: {task}

返回格式示例:
[{{"task": "子任务1描述", "agent_type": "适合的智能体类型"}}, {{"task": "子任务2描述", "agent_type": "适合的智能体类型"}}]

请直接返回 JSON 数组，不要其他内容。"""

        response = await self.supervisor.llm.generate(prompt, [])

        try:
            import json
            # 尝试解析 JSON
            subtasks = json.loads(response)
            return subtasks
        except:
            # 解析失败，创建默认子任务
            return [{"task": task, "agent_type": "general"}]

    async def _dispatch_parallel(self, subtasks: List[Dict], user_id: int, session_id: str) -> List[Dict]:
        """并行分发任务"""
        tasks = []
        for i, subtask in enumerate(subtasks):
            if i < len(self.members):
                member = self.members[i]
            else:
                # 如果子任务多于成员，使用轮询
                member = self.members[i % len(self.members)]

            tasks.append(member.run(subtask['task'], user_id, session_id))

        results = await asyncio.gather(*tasks, return_exceptions=True)

        # 处理结果
        formatted_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                formatted_results.append({
                    "task": subtasks[i]['task'],
                    "success": False,
                    "error": str(result)
                })
            else:
                formatted_results.append({
                    "task": subtasks[i]['task'],
                    "success": True,
                    "content": result.content,
                    "tool_calls": result.tool_calls
                })

        return formatted_results

    async def _dispatch_sequential(self, subtasks: List[Dict], user_id: int, session_id: str) -> List[Dict]:
        """顺序分发任务"""
        results = []
        context = ""

        for i, subtask in enumerate(subtasks):
            # 选择子智能体
            if i < len(self.members):
                member = self.members[i]
            else:
                member = self.members[i % len(self.members)]

            # 传递前一个结果作为上下文
            enhanced_task = f"{context}\n\n当前任务: {subtask['task']}" if context else subtask['task']

            result = await member.run(enhanced_task, user_id, session_id)

            results.append({
                "task": subtask['task'],
                "success": True,
                "content": result.content,
                "tool_calls": result.tool_calls
            })

            # 累加上下文
            context += f"\n\n=== 任务: {subtask['task']} ===\n{result.content}"

        return results

    async def _aggregate(self, results: List[Dict]) -> str:
        """汇总结果"""
        # 过滤成功的结果
        success_results = [r for r in results if r.get('success')]

        if not success_results:
            return "所有子任务执行失败"

        if len(success_results) == 1:
            return success_results[0].get('content', '')

        # 调用 LLM 汇总
        summary_prompt = "请汇总以下所有任务的结果，生成一个完整的回复:\n\n"
        for i, result in enumerate(success_results, 1):
            summary_prompt += f"=== 任务 {i}: {result.get('task', '')} ===\n{result.get('content', '')}\n\n"

        final_response = await self.supervisor.llm.generate(summary_prompt, [])
        return final_response