""" Agent 蜂群 API — Leader/Teammate 并行协作 POST /api/v1/swarm/run {"message": "帮我做三件事: ...", "mode": "parallel", "max_teammates": 5} → Leader 分解任务 → Teammates 并行执行 → Leader 汇总 参考 Claude Code src/tools/AgentTool/ + forkSubagent.ts """ from __future__ import annotations import logging from typing import Any, Dict, List, Optional from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from app.core.database import get_db from sqlalchemy.orm import Session from app.api.auth import get_current_user from app.models.user import User from app.models.agent import Agent from app.agent_runtime.swarm import ( SwarmRuntime, SwarmConfig, SwarmMode, SwarmResult, SwarmTask, create_swarm, ) from app.agent_runtime.schemas import AgentConfig, AgentLLMConfig, AgentToolConfig logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/v1/swarm", tags=["swarm"]) # ──────────────────────────── 请求/响应模型 ──────────────────────────── class SwarmRunRequest(BaseModel): message: str = Field(..., description="用户输入") mode: str = Field(default="parallel", description="蜂群模式: parallel | pipeline | debate") max_teammates: int = Field(default=5, ge=1, le=20) leader_model: Optional[str] = Field(default=None, description="Leader 模型") teammate_model: Optional[str] = Field(default=None, description="Teammate 模型") mailbox_enabled: bool = Field(default=True, description="启用 Agent 间消息传递") agent_ids: Optional[List[str]] = Field(default=None, description="指定的 Agent ID 列表（作为 Teammates）") retry_failed: bool = Field(default=True, description="失败任务是否重试") class SwarmTaskItem(BaseModel): id: str description: str assigned_agent_id: Optional[str] = None status: str result: Optional[str] = None error: Optional[str] = None iterations_used: int = 0 tool_calls_made: int = 0 duration_ms: int = 0 class SwarmTeammateItem(BaseModel): agent_id: str agent_name: str task_id: str success: bool output: str duration_ms: int iterations_used: int tool_calls_made: int error: Optional[str] = None class MailboxMessageItem(BaseModel): id: str from_: str = Field(alias="from") to: str content: str timestamp: float class SwarmRunResponse(BaseModel): success: bool final_answer: str mode: str tasks: List[SwarmTaskItem] = Field(default_factory=list) teammate_results: List[SwarmTeammateItem] = Field(default_factory=list) mailbox_messages: List[Dict[str, Any]] = Field(default_factory=list) total_duration_ms: int = 0 total_iterations: int = 0 total_tool_calls: int = 0 error: Optional[str] = None # ──────────────────────────── 端点 ──────────────────────────── @router.post("/run", response_model=SwarmRunResponse) async def swarm_run( req: SwarmRunRequest, current_user: User = Depends(get_current_user), db: Session = Depends(get_db), ): """运行 Agent 蜂群 — Leader 分解任务 → Teammates 并行执行 → 汇总。 支持三种模式: - parallel: 所有子任务并发执行（无依赖） - pipeline: 按依赖顺序执行 - debate: 多个 Agent 独立回答后汇总 Teammates 来源（优先级）: 1. agent_ids 参数指定 → 从数据库加载 Agent 配置 2. 自动生成 → 使用 teammate_model 创建轻量 Teammate """ uid = current_user.id # 解析模式 mode = SwarmMode.PARALLEL if req.mode == "pipeline": mode = SwarmMode.PIPELINE elif req.mode == "debate": mode = SwarmMode.DEBATE # 构建 SwarmConfig config = SwarmConfig( mode=mode, max_teammates=req.max_teammates, leader_model=req.leader_model or "deepseek-v4-pro", teammate_model=req.teammate_model or "deepseek-v4-flash", mailbox_enabled=req.mailbox_enabled, retry_failed=req.retry_failed, ) # 加载指定的 Agent 作为 Teammates teammate_configs: List[AgentConfig] = [] if req.agent_ids: for aid in req.agent_ids: agent = db.query(Agent).filter(Agent.id == aid).first() if agent: wc = agent.workflow_config or {} nodes = wc.get("nodes", []) agent_node_cfg = {} for node in nodes: if node.get("type") in ("agent", "llm", "template"): agent_node_cfg = node.get("data") or {} break teammate_configs.append(AgentConfig( name=agent.name, system_prompt=agent_node_cfg.get("system_prompt") or agent.description or "你是一个有用的AI助手。", llm=AgentLLMConfig( model=agent_node_cfg.get("model", req.teammate_model or "deepseek-v4-flash"), provider=agent_node_cfg.get("provider", "deepseek"), temperature=float(agent_node_cfg.get("temperature", 0.7)), max_iterations=int(agent_node_cfg.get("max_iterations", 10)), ), tools=AgentToolConfig( include_tools=agent_node_cfg.get("tools", []), ), user_id=uid, )) # 构建 Leader 配置 leader_config = AgentConfig( name="SwarmLeader", system_prompt="你是一个AI任务协调者。将复杂问题分解为子任务，协调多个AI Agent并行处理，并汇总结果。", llm=AgentLLMConfig(model=config.leader_model, temperature=0.3, max_iterations=10), user_id=uid, ) # 创建并运行 Swarm swarm = SwarmRuntime( config=config, leader_config=leader_config, teammate_configs=teammate_configs, ) result = await swarm.run(req.message) return SwarmRunResponse( success=result.success, final_answer=result.final_answer, mode=result.mode.value, tasks=[ SwarmTaskItem( id=t.id, description=t.description, assigned_agent_id=t.assigned_agent_id, status=t.status.value, result=t.result[:500] if t.result else None, error=t.error, iterations_used=t.iterations_used, tool_calls_made=t.tool_calls_made, duration_ms=t.duration_ms, ) for t in result.tasks ], teammate_results=[ SwarmTeammateItem( agent_id=tr["agent_id"], agent_name=tr["agent_name"], task_id=tr["task_id"], success=tr["success"], output=tr["output"][:500], duration_ms=tr["duration_ms"], iterations_used=tr["iterations_used"], tool_calls_made=tr["tool_calls_made"], error=tr.get("error"), ) for tr in result.teammate_results ], mailbox_messages=result.mailbox_messages, total_duration_ms=result.total_duration_ms, total_iterations=result.total_iterations, total_tool_calls=result.total_tool_calls, error=result.error, ) @router.post("/run/stream") async def swarm_run_stream( req: SwarmRunRequest, current_user: User = Depends(get_current_user), db: Session = Depends(get_db), ): """运行 Agent 蜂群（流式 SSE）— 实时推送任务分解、执行进度、汇总结果。""" import json as _json async def _stream(): uid = current_user.id mode = {"parallel": SwarmMode.PARALLEL, "pipeline": SwarmMode.PIPELINE, "debate": SwarmMode.DEBATE}.get(req.mode, SwarmMode.PARALLEL) config = SwarmConfig( mode=mode, max_teammates=req.max_teammates, leader_model=req.leader_model or "deepseek-v4-pro", teammate_model=req.teammate_model or "deepseek-v4-flash", mailbox_enabled=req.mailbox_enabled, retry_failed=req.retry_failed, ) # Load specified agents teammate_configs = [] if req.agent_ids: for aid in req.agent_ids: agent = db.query(Agent).filter(Agent.id == aid).first() if agent: wc = agent.workflow_config or {} agent_node_cfg = {} for node in wc.get("nodes", []): if node.get("type") in ("agent", "llm", "template"): agent_node_cfg = node.get("data") or {} break teammate_configs.append(AgentConfig( name=agent.name, system_prompt=agent_node_cfg.get("system_prompt") or agent.description or "你是一个有用的AI助手。", llm=AgentLLMConfig( model=agent_node_cfg.get("model", req.teammate_model or "deepseek-v4-flash"), provider=agent_node_cfg.get("provider", "deepseek"), temperature=float(agent_node_cfg.get("temperature", 0.7)), max_iterations=int(agent_node_cfg.get("max_iterations", 10)), ), tools=AgentToolConfig(include_tools=agent_node_cfg.get("tools", [])), user_id=uid, )) leader_config = AgentConfig( name="SwarmLeader", system_prompt="你是一个AI任务协调者。", llm=AgentLLMConfig(model=config.leader_model, temperature=0.3, max_iterations=10), user_id=uid, ) swarm = SwarmRuntime(config=config, leader_config=leader_config, teammate_configs=teammate_configs) yield f"data: {_json.dumps({'type': 'swarm_start', 'mode': req.mode, 'max_teammates': req.max_teammates}, ensure_ascii=False)}\n\n" result = await swarm.run(req.message) yield f"data: {_json.dumps({'type': 'swarm_done', 'success': result.success, 'final_answer': result.final_answer, 'total_duration_ms': result.total_duration_ms, 'total_iterations': result.total_iterations, 'total_tool_calls': result.total_tool_calls}, ensure_ascii=False)}\n\n" return StreamingResponse( _stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, )