472 lines
17 KiB
Python
472 lines
17 KiB
Python
"""Multi-agent orchestration endpoints.
|
||
|
||
Wraps the legacy ``newpython/OUT.py`` demo into the Gateway. Provides:
|
||
|
||
* ``POST /api/multi-agent/init`` — create a coordinator agent and per-sub-agent
|
||
threads. Returns the ``{agent_name: thread_id}`` map the client passes back
|
||
on subsequent calls.
|
||
* ``POST /api/multi-agent/run/stream`` — drive a *leader* or *special* agent
|
||
via SSE. The leader dispatches work to sub-agents through the
|
||
``agent_orchestration`` skill; sub-agents execute and broadcast their result.
|
||
|
||
Internal calls loop back to the same Gateway over HTTP, reusing the caller's
|
||
``Authorization`` header — so the demo's hard-coded ``get_token("ji")`` is gone.
|
||
The model used by both the coordinator agent and every streamed run is taken
|
||
from the request's ``model`` field.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import re
|
||
from typing import Any, AsyncIterator
|
||
|
||
import httpx
|
||
from fastapi import APIRouter, Body, HTTPException, Request
|
||
from fastapi.responses import StreamingResponse
|
||
|
||
logger = logging.getLogger(__name__)
|
||
router = APIRouter(prefix="/api/multi-agent", tags=["multi-agent"])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Loopback helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _loopback_base(request: Request) -> str:
|
||
return f"{request.url.scheme}://{request.url.netloc}"
|
||
|
||
|
||
def _auth_headers(request: Request) -> dict[str, str]:
|
||
token = request.headers.get("authorization")
|
||
if not token:
|
||
raise HTTPException(status_code=401, detail="Missing Authorization header")
|
||
return {"Authorization": token, "Content-Type": "application/json"}
|
||
|
||
|
||
async def _create_thread(client: httpx.AsyncClient, base: str, headers: dict[str, str]) -> str:
|
||
resp = await client.post(f"{base}/api/threads", json={}, headers=headers)
|
||
resp.raise_for_status()
|
||
return resp.json()["thread_id"]
|
||
|
||
|
||
async def _get_agent(client: httpx.AsyncClient, base: str, headers: dict[str, str], agent_id: str) -> dict[str, Any]:
|
||
resp = await client.get(f"{base}/api/agents/{agent_id}", headers=headers)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
|
||
async def _create_agent(
|
||
client: httpx.AsyncClient,
|
||
base: str,
|
||
headers: dict[str, str],
|
||
*,
|
||
name: str,
|
||
description: str,
|
||
model: str | None,
|
||
skills: list[str] | None,
|
||
soul: str,
|
||
agent_id: str | None = None,
|
||
) -> tuple[bool, dict[str, Any] | str]:
|
||
payload: dict[str, Any] = {"name": name, "description": description, "soul": soul}
|
||
if agent_id is not None:
|
||
payload["id"] = agent_id
|
||
if model is not None:
|
||
payload["model"] = model
|
||
if skills is not None:
|
||
payload["skills"] = skills
|
||
resp = await client.post(f"{base}/api/agents", json=payload, headers=headers)
|
||
if resp.status_code == 201:
|
||
return True, resp.json()
|
||
return False, resp.text
|
||
|
||
|
||
async def _append_thread_message(
|
||
client: httpx.AsyncClient,
|
||
base: str,
|
||
headers: dict[str, str],
|
||
thread_id: str,
|
||
content: str,
|
||
) -> None:
|
||
state_url = f"{base}/api/threads/{thread_id}/state"
|
||
get_resp = await client.get(state_url, headers=headers)
|
||
if get_resp.status_code != 200:
|
||
logger.warning("get state %s failed: %s %s", thread_id, get_resp.status_code, get_resp.text[:200])
|
||
return
|
||
state = get_resp.json()
|
||
values = state.setdefault("values", {})
|
||
messages = values.setdefault("messages", [])
|
||
if not isinstance(messages, list):
|
||
messages = []
|
||
values["messages"] = messages
|
||
messages.append({"type": "human", "content": content})
|
||
post_resp = await client.post(state_url, json=state, headers=headers)
|
||
if post_resp.status_code >= 400:
|
||
logger.warning("update state %s failed: %s %s", thread_id, post_resp.status_code, post_resp.text[:200])
|
||
|
||
|
||
async def _broadcast(
|
||
client: httpx.AsyncClient,
|
||
base: str,
|
||
headers: dict[str, str],
|
||
message: str,
|
||
agent_threads: dict[str, str],
|
||
skip: list[str],
|
||
) -> None:
|
||
targets = [tid for name, tid in agent_threads.items() if name not in skip]
|
||
if not targets:
|
||
return
|
||
await asyncio.gather(
|
||
*(_append_thread_message(client, base, headers, tid, message) for tid in targets),
|
||
return_exceptions=True,
|
||
)
|
||
|
||
|
||
_TOOL_CALL_PATTERN = re.compile(r'\.py\s+"([^"]*)"\s+"([^"]*)"')
|
||
|
||
|
||
def _extract_name_task(text: str) -> tuple[str, str]:
|
||
match = _TOOL_CALL_PATTERN.search(text)
|
||
if match:
|
||
return match.group(1), match.group(2)
|
||
return "", ""
|
||
|
||
|
||
def _flatten_content(content: Any) -> str:
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
parts: list[str] = []
|
||
for item in content:
|
||
if isinstance(item, dict):
|
||
parts.append(str(item.get("text", "")))
|
||
else:
|
||
parts.append(str(item))
|
||
return "".join(parts)
|
||
return str(content)
|
||
|
||
|
||
def _run_payload(
|
||
*,
|
||
agent_name: str,
|
||
model_name: str,
|
||
thread_id: str,
|
||
new_message: str,
|
||
excluded_tools: list[str],
|
||
skill_stop_names: list[str],
|
||
) -> dict[str, Any]:
|
||
return {
|
||
"input": {
|
||
"messages": [
|
||
{
|
||
"type": "human",
|
||
"content": [{"type": "text", "text": new_message}],
|
||
"additional_kwargs": {},
|
||
}
|
||
]
|
||
},
|
||
"config": {"recursion_limit": 1000},
|
||
"context": {
|
||
"agent_name": agent_name,
|
||
"model_name": model_name,
|
||
"mode": "pro",
|
||
"reasoning_effort": "medium",
|
||
"thinking_enabled": True,
|
||
"is_plan_mode": False,
|
||
"subagent_enabled": False,
|
||
"thread_id": thread_id,
|
||
},
|
||
"stream_mode": ["messages-tuple", "values"],
|
||
"stream_subgraphs": True,
|
||
"stream_resumable": True,
|
||
"assistant_id": "lead_agent",
|
||
"on_disconnect": "continue",
|
||
"excluded_tools": excluded_tools,
|
||
"skill_stop_names": skill_stop_names,
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Init endpoint
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@router.post("/init")
|
||
async def init_multi_agent(request: Request, payload: dict = Body(...)) -> dict[str, Any]:
|
||
agent_names_in = payload.get("agent_names") or []
|
||
agent_names: list[str] = [str(name).lower() for name in agent_names_in]
|
||
main_agent_name: str = str(payload.get("main_agent_name") or "main-agent").lower()
|
||
model: str | None = payload.get("model")
|
||
|
||
if not agent_names:
|
||
raise HTTPException(status_code=400, detail="agent_names is required")
|
||
|
||
base = _loopback_base(request)
|
||
headers = _auth_headers(request)
|
||
|
||
async with httpx.AsyncClient(timeout=None) as client:
|
||
# Optimization A: fetch every sub-agent's metadata and create its
|
||
# thread concurrently. With 6 seats this turns ~12 sequential
|
||
# loopback round-trips into one wave of 12 parallel ones.
|
||
async def _prepare_sub_agent(sub_id: str) -> tuple[str, str, str]:
|
||
info = await _get_agent(client, base, headers, sub_id)
|
||
sub_name = str(info["name"]).lower()
|
||
thread_id = await _create_thread(client, base, headers)
|
||
return sub_name, thread_id, info.get("description", "") or ""
|
||
|
||
prepared = await asyncio.gather(*(_prepare_sub_agent(sid) for sid in agent_names))
|
||
|
||
agent_threads: dict[str, str] = {sub_name: thread_id for sub_name, thread_id, _ in prepared}
|
||
descriptions: list[str] = [
|
||
f"agent_name:{sub_name},{description}" for sub_name, _, description in prepared
|
||
]
|
||
|
||
soul = (
|
||
"#核心定位：\n"
|
||
"用途：根据用户的任务目标,调用其他智能体干活,只负责协调汇总,"
|
||
"当你觉得已有信息支持回答用户问题时候,停止调度,并输出最后的结果\n"
|
||
"#调度方法：\n"
|
||
"通过使用agent_orchestration技能调用其他智能体干活,使用技能时传入agent_name和需要做的活\n"
|
||
"#可调度智能体：\n"
|
||
f"{descriptions}"
|
||
)
|
||
|
||
agent_threads[main_agent_name] = await _create_thread(client, base, headers)
|
||
ok, result = await _create_agent(
|
||
client,
|
||
base,
|
||
headers,
|
||
name=main_agent_name,
|
||
description="你是总控控制协调智能体,负责协调其他智能体干活",
|
||
model=model,
|
||
skills=["agent_orchestration"],
|
||
soul=soul,
|
||
agent_id=main_agent_name,
|
||
)
|
||
if not ok and isinstance(result, str) and "exists" not in result:
|
||
raise HTTPException(status_code=500, detail=f"create coordinator agent failed: {result}")
|
||
|
||
# Optimization C: the sub-agent SOULs already declare "你只在被点到时
|
||
# 发言,由总控协调智能体派活", so the old intro broadcast (one human
|
||
# message per thread, 6 × GET-state + 6 × POST-state) is redundant.
|
||
# Skipping it removes the slowest stage of init entirely.
|
||
|
||
return {
|
||
"status": "success",
|
||
"main_agent_name": main_agent_name,
|
||
"thread_ids": agent_threads,
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Stream endpoint
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def _stream_upstream(
|
||
client: httpx.AsyncClient,
|
||
base: str,
|
||
headers: dict[str, str],
|
||
thread_id: str,
|
||
body: dict[str, Any],
|
||
) -> AsyncIterator[tuple[str | None, list[str] | None]]:
|
||
"""Yield (raw_line, None) for each SSE line, then ``(None, captured_data_lines)`` once."""
|
||
|
||
url = f"{base}/api/threads/{thread_id}/runs/stream"
|
||
captured: list[str] = []
|
||
async with client.stream("POST", url, json=body, headers=headers) as resp:
|
||
if resp.status_code != 200:
|
||
body_text = (await resp.aread()).decode("utf-8", "replace")[:1000]
|
||
raise HTTPException(
|
||
status_code=502,
|
||
detail=f"upstream {url} returned HTTP {resp.status_code}: {body_text}",
|
||
)
|
||
async for line in resp.aiter_lines():
|
||
if not line:
|
||
continue
|
||
yield line, None
|
||
if line.startswith("data:"):
|
||
captured.append(line)
|
||
yield None, captured
|
||
|
||
|
||
def _parse_last_message(captured: list[str]) -> dict[str, Any]:
|
||
if len(captured) < 2:
|
||
raise HTTPException(
|
||
status_code=502,
|
||
detail=f"upstream produced no usable data lines (got {len(captured)})",
|
||
)
|
||
raw = captured[-2][5:].strip()
|
||
try:
|
||
parsed = json.loads(raw)
|
||
except json.JSONDecodeError as exc:
|
||
raise HTTPException(status_code=502, detail=f"failed to parse upstream frame: {exc}") from exc
|
||
messages = parsed.get("messages") or []
|
||
if not messages:
|
||
return {}
|
||
return messages[-1] or {}
|
||
|
||
|
||
async def _leader_run(
|
||
client: httpx.AsyncClient,
|
||
base: str,
|
||
headers: dict[str, str],
|
||
agent_threads: dict[str, str],
|
||
agent_name: str,
|
||
new_message: str,
|
||
skill_stop_names: list[str],
|
||
model_name: str,
|
||
) -> AsyncIterator[str | dict[str, Any]]:
|
||
await _broadcast(client, base, headers, "总控智能体收到用户消息:" + new_message, agent_threads, [agent_name])
|
||
thread_id = agent_threads[agent_name]
|
||
body = _run_payload(
|
||
agent_name=agent_name,
|
||
model_name=model_name,
|
||
thread_id=thread_id,
|
||
new_message=new_message,
|
||
excluded_tools=["web_search"],
|
||
skill_stop_names=skill_stop_names,
|
||
)
|
||
|
||
captured: list[str] = []
|
||
async for line, done in _stream_upstream(client, base, headers, thread_id, body):
|
||
if line is not None:
|
||
yield line
|
||
elif done is not None:
|
||
captured = done
|
||
|
||
last_msg = _parse_last_message(captured)
|
||
leader_text = _flatten_content(last_msg.get("content", ""))
|
||
tool_calls = last_msg.get("tool_calls") or []
|
||
if not tool_calls:
|
||
yield {"status": [], "content": leader_text, "agent_name": agent_name}
|
||
return
|
||
|
||
dispatched: list[list[str]] = []
|
||
for tool_call in tool_calls:
|
||
try:
|
||
sub_name, task_text = _extract_name_task(str(tool_call.get("args", "")))
|
||
if not sub_name:
|
||
continue
|
||
await _broadcast(
|
||
client,
|
||
base,
|
||
headers,
|
||
f"总控智能体给{sub_name}派活,内容为:" + task_text,
|
||
agent_threads,
|
||
[agent_name, sub_name],
|
||
)
|
||
dispatched.append([sub_name, task_text])
|
||
except Exception as exc:
|
||
logger.warning("dispatch failed: %s", exc)
|
||
yield {"status": dispatched, "content": leader_text, "agent_name": agent_name}
|
||
|
||
|
||
async def _special_run(
|
||
client: httpx.AsyncClient,
|
||
base: str,
|
||
headers: dict[str, str],
|
||
agent_threads: dict[str, str],
|
||
agent_name: str,
|
||
new_message: str,
|
||
skill_stop_names: list[str],
|
||
model_name: str,
|
||
) -> AsyncIterator[str | dict[str, Any]]:
|
||
thread_id = agent_threads[agent_name]
|
||
body = _run_payload(
|
||
agent_name=agent_name,
|
||
model_name=model_name,
|
||
thread_id=thread_id,
|
||
new_message=new_message,
|
||
excluded_tools=["web_search", "ask_clarification", "present_files", "view_image"],
|
||
skill_stop_names=skill_stop_names,
|
||
)
|
||
|
||
captured: list[str] = []
|
||
async for line, done in _stream_upstream(client, base, headers, thread_id, body):
|
||
if line is not None:
|
||
yield line
|
||
elif done is not None:
|
||
captured = done
|
||
|
||
last_msg = _parse_last_message(captured)
|
||
content = _flatten_content(last_msg.get("content", ""))
|
||
await _broadcast(
|
||
client,
|
||
base,
|
||
headers,
|
||
f"子智能体{agent_name}完成{new_message}工作,交付内容为:" + content,
|
||
agent_threads,
|
||
[agent_name],
|
||
)
|
||
yield {
|
||
"status": f"子智能体{agent_name}完成{new_message}工作",
|
||
"content": content,
|
||
"agent_name": agent_name,
|
||
}
|
||
|
||
|
||
@router.post("/run/stream")
|
||
async def run_multi_agent_stream(request: Request, payload: dict = Body(...)) -> StreamingResponse:
|
||
agent_type: str = str(payload.get("agent_type") or "leader").lower()
|
||
agent_threads: dict[str, str] = payload.get("d_agent_thread_id") or {}
|
||
agent_name: str = str(payload.get("agent_name") or "").lower()
|
||
new_message: str = str(payload.get("new_message") or "")
|
||
skill_stop_names: list[str] = payload.get("skill_stop_names") or []
|
||
model_name: str = str(payload.get("model") or "deepseek-chat")
|
||
|
||
if not agent_name or agent_name not in agent_threads:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"agent_name '{agent_name}' is missing from d_agent_thread_id",
|
||
)
|
||
|
||
base = _loopback_base(request)
|
||
headers = _auth_headers(request)
|
||
|
||
async def stream_generator() -> AsyncIterator[str]:
|
||
async with httpx.AsyncClient(timeout=None) as client:
|
||
try:
|
||
if agent_type == "leader":
|
||
gen = _leader_run(
|
||
client, base, headers, agent_threads, agent_name, new_message, skill_stop_names, model_name,
|
||
)
|
||
else:
|
||
gen = _special_run(
|
||
client, base, headers, agent_threads, agent_name, new_message, skill_stop_names, model_name,
|
||
)
|
||
async for chunk in gen:
|
||
if isinstance(chunk, dict):
|
||
yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
|
||
else:
|
||
yield chunk if chunk.endswith("\n") else chunk + "\n"
|
||
except HTTPException as exc:
|
||
logger.warning("multi-agent stream HTTPException: %s %s", exc.status_code, exc.detail)
|
||
err_frame = {
|
||
"status": "error",
|
||
"content": f"[{exc.status_code}] {exc.detail}",
|
||
"agent_name": agent_name,
|
||
}
|
||
yield f"data: {json.dumps(err_frame, ensure_ascii=False)}\n\n"
|
||
except Exception as exc:
|
||
logger.exception("multi-agent stream unexpected failure")
|
||
err_frame = {
|
||
"status": "error",
|
||
"content": f"unexpected: {exc!r}",
|
||
"agent_name": agent_name,
|
||
}
|
||
yield f"data: {json.dumps(err_frame, ensure_ascii=False)}\n\n"
|
||
|
||
return StreamingResponse(
|
||
stream_generator(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"X-Accel-Buffering": "no",
|
||
"Connection": "keep-alive",
|
||
},
|
||
)
|