281 lines
11 KiB
Python
281 lines
11 KiB
Python
|
"""
|
|||
|
|
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"},
|
|||
|
|
)
|