"""协作团队的实时性：谁先做完谁先报到，调度者在想什么也要有事件。 以前成员的 agent.step.end 要等整轮 gather 结束才按派活顺序统一发出——并行时先做完 的成员一直显示"进行中"，直到最慢的那个做完；矩阵声称能回答的"谁还在跑"其实答 不了。调度者那次结构化模型调用前后没有任何事件，一个 74 秒的团队节点里约 66 秒 是调度者在决策，界面上是个黑盒。 用量也一样：agent 节点、成员、调度者的模型调用都不发 llm.end，界面只能等终态 才知道花了多少 token。补上之后必须守住另一头：run.finished 的 usage 是后端累计 的权威值，事件只是把同一笔账逐次摊开，加起来要对得上，不能记两遍。 """ from __future__ import annotations import asyncio import json import pytest from langchain_core.messages import AIMessage, AIMessageChunk from langchain_core.outputs import ChatGenerationChunk from sqlalchemy import select from app.db.base import SessionLocal from app.db.models import Run, RunEvent from app.engine.runner import run_manager SLOW, FAST = 0.6, 0.05 ROUTE_SECONDS = 0.15 @pytest.fixture(autouse=True) async def engine_up(): await run_manager.setup() yield await run_manager.shutdown() def _team_graph() -> dict: agents = [ {"name": n, "description": f"{n} 的职责", "system": f"你是 {n}", "max_steps": 1} for n in ("researcher", "analyst") ] return { "nodes": [ {"id": "in", "type": "input", "data": {"config": {}}}, {"id": "team", "type": "supervisor", "data": {"label": "协作团队", "config": { "goal": "比较两种向量库", "agents": agents, "max_rounds": 3, "max_parallel": 2, }}}, {"id": "out", "type": "output", "data": {"config": {}}}, ], "edges": [{"source": "in", "target": "team"}, {"source": "team", "target": "out"}], } def _usage(messages, reply: str) -> dict: text = " ".join(str(getattr(m, "content", m)) for m in messages) return {"input_tokens": 10 + len(text) // 50, "output_tokens": 3 + len(reply) // 20, "total_tokens": 0} def _script(monkeypatch, decisions: list[dict]) -> None: """调度者按剧本派活；成员 researcher 慢、analyst 快；每次调用都报 token。 调度者走 with_structured_output(include_raw=True)：真实供应商会连原始消息 一起给回来，用量就在那条消息上。mock 默认只给解析结果，这里补齐这层行为。 """ from app.providers import mock_model turns = {"n": 0} def _decide(self, messages): if self.response_format: i = turns["n"] turns["n"] += 1 payload = decisions[i] if i < len(decisions) else {"assignments": [], "done": True} return AIMessage(content=json.dumps(payload, ensure_ascii=False)) return AIMessage(content="（干完了）") async def _agenerate(self, messages, stop=None, run_manager=None, **kw): text = " ".join(str(getattr(m, "content", "")) for m in messages) if self.response_format: await asyncio.sleep(ROUTE_SECONDS) else: await asyncio.sleep(FAST if "快活" in text else SLOW) result = self._generate(messages, stop, **kw) msg = result.generations[0].message msg.usage_metadata = _usage(messages, str(msg.content)) return result def _structured(self, schema, *, include_raw=False, **kw): bound = self.model_copy(update={"response_format": schema}) async def parse(msg): parsed = json.loads(str(msg.content)) return {"raw": msg, "parsed": parsed, "parsing_error": None} if include_raw else parsed return bound | parse monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide) monkeypatch.setattr(mock_model.MockChatModel, "_agenerate", _agenerate) monkeypatch.setattr(mock_model.MockChatModel, "with_structured_output", _structured) async def _run(graph: dict) -> tuple[list[RunEvent], Run]: run = await run_manager.start(graph=graph, input_payload={"question": "比较"}) for _ in range(400): await asyncio.sleep(0.05) async with SessionLocal() as session: row = await session.get(Run, run.id) if row.status in ("succeeded", "failed", "cancelled"): break else: pytest.fail("运行没有在 20 秒内结束") async with SessionLocal() as session: events = list((await session.execute( select(RunEvent).where(RunEvent.run_id == run.id).order_by(RunEvent.seq) )).scalars()) return events, row PARALLEL_ROUND = [{ "assignments": [ {"agent": "researcher", "instruction": "慢活"}, {"agent": "analyst", "instruction": "快活"}, ], "reason": "两边互不相干", }, {"assignments": [], "done": True, "reason": "够了"}] async def test_members_report_done_in_the_order_they_actually_finish(monkeypatch): _script(monkeypatch, PARALLEL_ROUND) events, run = await _run(_team_graph()) assert run.status == "succeeded", run.error ends = [e for e in events if e.type == "agent.step.end"] assert [e.data["agent"] for e in ends] == ["analyst", "researcher"], \ "快的那个先做完，就该先报到——而不是按派活顺序等整轮结束再一起报" start = next(e for e in events if e.type == "agent.step.start") fast_end = ends[0] # 快的那个结束事件的时刻要贴着它自己做完的时刻，而不是被慢的那个拖到整轮结束 assert fast_end.ts - start.ts < (SLOW + FAST) / 2, \ f"analyst {FAST}s 就做完了，结束事件却在 {fast_end.ts - start.ts:.2f}s 才出现" # 字段不变：界面照旧按这些字段画矩阵 for e in ends: assert {"agent", "duration_ms", "round", "parallel", "preview"} <= set(e.data) async def test_the_coordinator_is_not_a_black_box(monkeypatch): _script(monkeypatch, PARALLEL_ROUND) events, run = await _run(_team_graph()) assert run.status == "succeeded", run.error starts = [e for e in events if e.type == "agent.route.start"] ends = [e for e in events if e.type == "agent.route.end"] assert [e.data["round"] for e in starts] == [0, 1] assert [e.data["round"] for e in ends] == [0, 1] assert all(e.node_id == "team" for e in starts + ends) first, last = ends assert first.data["agents"] == ["researcher", "analyst"] assert first.data["parallel"] == 2 assert first.data["done"] is False assert first.data["reason"] == "两边互不相干" assert last.data["done"] is True and last.data["agents"] == [] assert last.data["reason"] == "够了" for e in ends: assert e.data["duration_ms"] >= ROUTE_SECONDS * 1000 * 0.8, e.data # 先"开始想"，再"想好了"，然后才派人出去 seq = {e.type + str(e.data.get("round", "")): e.seq for e in events} first_dispatch = next(e.seq for e in events if e.type == "agent.step.start") assert seq["agent.route.start0"] < seq["agent.route.end0"] < first_dispatch async def test_every_model_call_reports_usage_and_the_sum_matches_the_run(monkeypatch): _script(monkeypatch, PARALLEL_ROUND) events, run = await _run(_team_graph()) assert run.status == "succeeded", run.error calls = [e for e in events if e.type == "llm.end" and e.node_id == "team"] by_agent: dict[str, list[dict]] = {} for e in calls: by_agent.setdefault(e.data["agent"], []).append(e.data) assert set(by_agent) == {"调度者", "researcher", "analyst"}, by_agent.keys() assert len(by_agent["调度者"]) == 2, "调度了两次，就是两次模型调用" for d in (d for items in by_agent.values() for d in items): assert {"model", "input_tokens", "output_tokens", "cost_usd"} <= set(d), d assert d["input_tokens"] > 0 and d["output_tokens"] > 0, d finished = next(e for e in events if e.type == "run.finished") usage = finished.data["usage"] # 权威值是后端累计；逐次事件加起来必须等于它——多一笔少一笔都说明记重或漏记了 assert usage["input_tokens"] == sum(e.data["input_tokens"] for e in calls) assert usage["output_tokens"] == sum(e.data["output_tokens"] for e in calls) assert usage["calls"] == len(calls) assert run.usage["input_tokens"] == usage["input_tokens"] async def test_agent_node_reports_each_model_call(monkeypatch): """agent 节点以前只有 llm.start 没有 llm.end：用量要等整个节点结束才知道。""" from app.providers import mock_model async def _astream(self, messages, stop=None, run_manager=None, **kw): reply = "查完了，结论如下。" await asyncio.sleep(0.01) yield ChatGenerationChunk(message=AIMessageChunk( content=reply, usage_metadata=_usage(messages, reply))) monkeypatch.setattr(mock_model.MockChatModel, "_astream", _astream) graph = { "nodes": [ {"id": "in", "type": "input", "data": {"config": {}}}, {"id": "ag", "type": "agent", "data": {"label": "分析员", "config": { "prompt": "{{ input.question }}", "tools": [], "max_steps": 3}}}, {"id": "out", "type": "output", "data": {"config": {}}}, ], "edges": [{"source": "in", "target": "ag"}, {"source": "ag", "target": "out"}], } events, run = await _run(graph) assert run.status == "succeeded", run.error starts = [e for e in events if e.type == "llm.start" and e.node_id == "ag"] ends = [e for e in events if e.type == "llm.end" and e.node_id == "ag"] assert len(ends) == len(starts) >= 1 assert all(e.data["agent"] == "分析员" for e in ends) usage = next(e for e in events if e.type == "run.finished").data["usage"] assert usage["input_tokens"] == sum(e.data["input_tokens"] for e in ends) > 0 assert usage["output_tokens"] == sum(e.data["output_tokens"] for e in ends) > 0