From b6f5e8daf1258f619174eb3147e02f709d7d4d04 Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Sun, 28 Jun 2026 21:15:16 +0000 Subject: [PATCH 01/54] feat(arena): opt-in inter-agent concurrent runtime scaffold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds src/benchflow/arena/ — a self-contained, OPT-IN scaffold for the missing "inter-agent" axis of multi-agent: N agents acting concurrently on ONE shared environment (the deferred arena-concurrent mode). Touches no existing scored path — the sequential scene path and the scalar reward path are untouched. - protocol.py — turn-poll contract (Seam 3): observe/act, status in {waiting, not_your_turn, your_turn(request_id), done}, stale-request rejection. - runtime.py — run_arena (Seam 1 + eval glue): N seats concurrently via asyncio.gather; a worker cap bounds concurrent DECISIONS (not waiting, so interdependent seats never deadlock); a wall-clock deadline reaps stragglers. - reward.py — SharedEnvReward (Seam 4): one shared standings map -> a per-seat reward vector (pvp net / coop joint). Each seat is one RolloutNode, so this rides the existing per-node Reward contract — no change to the scalar path. tests/test_arena.py: an in-memory turn-gated fake env + scripted policies prove concurrency, turn-gating, stale-request rejection, straggler reaping, and zero- sum scoring — no ACP/sandbox/LLM (asyncio.run, no plugin dep). examples/arena/duel_deepseek.py: a REAL run — two deepseek-v4-pro seats play one shared rock-paper-scissors round concurrently through run_arena. Verified live: seat-0 paper vs seat-1 scissors -> reward -100 / +100, chips conserved. Seam 2 (a co-tenant SharedManifestEnvironment provisioning one service per scene) is deferred to the caller; this scaffold drives any SeatClient. Build the rest at the trigger — when an in-tree benchmark needs BenchFlow to host concurrent seats. --- examples/arena/duel_deepseek.py | 161 ++++++++++++++++++++++++++++++++ src/benchflow/arena/__init__.py | 39 ++++++++ src/benchflow/arena/protocol.py | 70 ++++++++++++++ src/benchflow/arena/reward.py | 41 ++++++++ src/benchflow/arena/runtime.py | 74 +++++++++++++++ tests/test_arena.py | 146 +++++++++++++++++++++++++++++ 6 files changed, 531 insertions(+) create mode 100644 examples/arena/duel_deepseek.py create mode 100644 src/benchflow/arena/__init__.py create mode 100644 src/benchflow/arena/protocol.py create mode 100644 src/benchflow/arena/reward.py create mode 100644 src/benchflow/arena/runtime.py create mode 100644 tests/test_arena.py diff --git a/examples/arena/duel_deepseek.py b/examples/arena/duel_deepseek.py new file mode 100644 index 00000000..6cabac66 --- /dev/null +++ b/examples/arena/duel_deepseek.py @@ -0,0 +1,161 @@ +"""Real run of the arena scaffold with two deepseek-v4 seats. + +Two independent agents play one round of rock-paper-scissors **concurrently** +against a single shared, turn-gated environment — driven entirely by +``benchflow.arena.run_arena``. The environment here is an in-process object that +implements the same ``SeatClient`` turn-poll contract a networked co-tenant +service would (so the demo needs no FastAPI/sandbox); each seat's brain is a real +``deepseek-v4-pro`` call. + + set -a; . ./sb-run.env; set +a # DEEPSEEK_API_KEY (and optional _BASE_URL) + uv run python examples/arena/duel_deepseek.py + +This is the inter-agent (concurrent) axis: N agents, one shared world, per-seat +reward — none of which the framework hosts natively today. The scaffold is +opt-in and touches no existing scored path. +""" + +from __future__ import annotations + +import asyncio +import os +import random + +import httpx + +from benchflow.arena import Observation, SharedEnvReward, run_arena + +RPS_BEATS = {"rock": "scissors", "scissors": "paper", "paper": "rock"} + + +class DuelFloor: + """In-process shared env implementing the SeatClient contract: a turn-gated, + 2-seat rock-paper-scissors round. Seats lazy-join on first ``observe``; + throws stay private until both have acted, then it settles zero-sum.""" + + def __init__(self, stake: int = 100, start: int = 1000) -> None: + self.stake, self.start = stake, start + self.bankroll: dict[str, int] = {} + self.seated: list[str] = [] + self.throws: dict[str, str] = {} + self.pending: str | None = None + self.cur_rid: str | None = None + self.turn = 0 + self.formed = self.done = False + self.lock = asyncio.Lock() + + def _open(self, seat: str) -> None: + self.pending, self.turn = seat, self.turn + 1 + self.cur_rid = f"{seat}#{self.turn}" + + async def observe(self, seat_id: str) -> dict: + async with self.lock: + self.bankroll.setdefault(seat_id, self.start) + if not self.formed and seat_id not in self.seated: + self.seated.append(seat_id) + if not self.formed and len(self.seated) >= 2: + self.formed = True + self._open(self.seated[0]) + if self.done: + return {"status": "done", "bankroll": self.bankroll[seat_id]} + if not self.formed: + return {"status": "waiting"} + if self.pending != seat_id: + return {"status": "not_your_turn", "current_actor": self.pending} + return { + "status": "your_turn", "request_id": self.cur_rid, + "observation": { + "public": {"pot": self.stake, "game": "rock-paper-scissors"}, + "private": {}, + }, + "legal_actions": [ + {"verb": "throw", "args": {"hand": h}} + for h in ("rock", "paper", "scissors") + ], + } + + async def act(self, seat_id: str, request_id: str, action: dict) -> dict: + async with self.lock: + if not self.formed or self.pending != seat_id: + return {"ok": False, "status": "not_your_turn"} + if request_id != self.cur_rid: + return {"ok": False, "status": "stale_request_id"} + self.throws[seat_id] = str(action.get("args", {}).get("hand", "rock")) + idx = self.seated.index(seat_id) + if idx + 1 < len(self.seated): + self._open(self.seated[idx + 1]) + else: + self._settle() + return {"ok": True, "status": "applied"} + + def _settle(self) -> None: + a, b = self.seated + ta, tb = self.throws[a], self.throws[b] + if ta != tb: + win, lose = (a, b) if RPS_BEATS.get(ta) == tb else (b, a) + self.bankroll[win] += self.stake + self.bankroll[lose] -= self.stake + self.pending = self.cur_rid = None + self.done = True + + def standings(self) -> dict[str, int]: + return dict(self.bankroll) + + +class DeepSeekPolicy: + """A seat brain backed by a real deepseek-v4-pro call.""" + + def __init__(self, seat: str, http: httpx.AsyncClient, model: str = "deepseek-v4-pro") -> None: + self.seat, self.http, self.model = seat, http, model + self.base = os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com").rstrip("/") + self.key = os.environ["DEEPSEEK_API_KEY"] + self.choice: str | None = None + + async def act(self, obs: Observation) -> dict: + hands = [a["args"]["hand"] for a in obs.legal_actions] + prompt = ( + f"You are {self.seat} in ONE round of rock-paper-scissors for a pot of " + f"{obs.public.get('pot')} chips. Choose exactly one of: {', '.join(hands)}. " + "Reply with ONLY that single word." + ) + try: + r = await self.http.post( + f"{self.base}/chat/completions", + headers={"Authorization": f"Bearer {self.key}"}, + json={ + "model": self.model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 256, "temperature": 0.8, + }, + timeout=90.0, + ) + text = (r.json()["choices"][0]["message"]["content"] or "").lower() + hand = next((h for h in hands if h in text), random.choice(hands)) + except Exception as exc: # a flaky call falls back to a random legal throw + print(f" [{self.seat}] deepseek call failed ({exc!r}); random fallback") + hand = random.choice(hands) + self.choice = hand + return {"verb": "throw", "args": {"hand": hand}} + + +async def _main() -> None: + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required (source your env file first)") + floor = DuelFloor() + seats = ["seat-0", "seat-1"] + async with httpx.AsyncClient() as http: + policies = {s: DeepSeekPolicy(s, http) for s in seats} + print("running arena: 2 deepseek-v4-pro seats, one shared RPS table…", flush=True) + res = await run_arena( + seats, floor, lambda s: policies[s], deadline_s=120.0, poll_s=0.05, + ) + st = floor.standings() + print("\npicks :", {s: p.choice for s, p in policies.items()}) + print("standings :", st) + print("reward (pvp):", SharedEnvReward().score(st)) + print("seat status :", {s: r["status"] for s, r in res.items()}) + print("conserved :", sum(st.values()), "(== 2000)") + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/src/benchflow/arena/__init__.py b/src/benchflow/arena/__init__.py new file mode 100644 index 00000000..5faf7ee3 --- /dev/null +++ b/src/benchflow/arena/__init__.py @@ -0,0 +1,39 @@ +"""Inter-agent concurrent (arena) runtime — an OPT-IN scaffold. + +Runs N agent "seats" concurrently against ONE shared environment service, the +genuinely-missing axis of multi-agent (the deferred ``arena-concurrent`` mode). +This package is additive and self-contained: it does not modify the sequential +scene path (``scenes.compile_scenes_to_steps``) or the scalar reward path, so no +existing scored benchmark is affected. + + - :mod:`~benchflow.arena.protocol` — the turn-poll contract (Seam 3). + - :func:`~benchflow.arena.runtime.run_arena` — concurrent seat driver (Seam 1 + + eval glue: worker cap, wall-clock deadline, straggler reaper). + - :class:`~benchflow.arena.reward.SharedEnvReward` — per-seat reward vector + (Seam 4). + +The co-tenant environment topology (Seam 2 — provision one service per scene and +attach K seats) is intentionally left to the caller / a future +``SharedManifestEnvironment``; this scaffold is driven by any ``SeatClient``. +""" + +from __future__ import annotations + +from benchflow.arena.protocol import ( + Observation, + SeatClient, + SeatPolicy, + SeatStatus, +) +from benchflow.arena.reward import FloorMode, SharedEnvReward +from benchflow.arena.runtime import run_arena + +__all__ = [ + "FloorMode", + "Observation", + "SeatClient", + "SeatPolicy", + "SeatStatus", + "SharedEnvReward", + "run_arena", +] diff --git a/src/benchflow/arena/protocol.py b/src/benchflow/arena/protocol.py new file mode 100644 index 00000000..205a7441 --- /dev/null +++ b/src/benchflow/arena/protocol.py @@ -0,0 +1,70 @@ +"""Turn-poll contract for inter-agent concurrent (arena) runs — Seam 3. + +A shared-environment service that hosts N concurrent seats exposes exactly two +calls per seat: ``observe`` (poll for your turn) and ``act`` (submit one legal +action). This is the one genuinely new schema the arena runtime needs; it lives +here as a small, additive contract — no change to the verifier-facing +``Environment`` protocol. + +A real ``SeatClient`` is an HTTP client to the co-tenant service; tests use an +in-memory fake. A real ``SeatPolicy`` wraps an ACP agent; tests script it. +""" + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +__all__ = ["SeatStatus", "Observation", "SeatClient", "SeatPolicy"] + + +class SeatStatus(enum.StrEnum): + WAITING = "waiting" # seated/queued, the table has not formed yet + NOT_YOUR_TURN = "not_your_turn" + YOUR_TURN = "your_turn" + DONE = "done" + + +@dataclass +class Observation: + """One seat's view, as returned by ``observe``.""" + + status: SeatStatus + request_id: str | None = None + public: dict[str, Any] = field(default_factory=dict) + private: dict[str, Any] = field(default_factory=dict) + legal_actions: list[dict[str, Any]] = field(default_factory=list) + + @property + def done(self) -> bool: + return self.status is SeatStatus.DONE + + @classmethod + def from_payload(cls, p: dict[str, Any]) -> Observation: + obs = p.get("observation") or {} + return cls( + status=SeatStatus(p.get("status", "done")), + request_id=p.get("request_id"), + public=dict(obs.get("public", {})), + private=dict(obs.get("private", {})), + legal_actions=list(p.get("legal_actions", [])), + ) + + +@runtime_checkable +class SeatClient(Protocol): + """The shared-env service contract: poll a turn, submit an action.""" + + async def observe(self, seat_id: str) -> dict[str, Any]: ... + + async def act( + self, seat_id: str, request_id: str, action: dict[str, Any] + ) -> dict[str, Any]: ... + + +@runtime_checkable +class SeatPolicy(Protocol): + """A seat's brain: choose one legal action for a ``your_turn`` observation.""" + + async def act(self, obs: Observation) -> dict[str, Any]: ... diff --git a/src/benchflow/arena/reward.py b/src/benchflow/arena/reward.py new file mode 100644 index 00000000..a8ed8b01 --- /dev/null +++ b/src/benchflow/arena/reward.py @@ -0,0 +1,41 @@ +"""Per-seat reward vector from one shared aggregate — Seam 4. + +The canonical Reward contract (``benchflow.rewards.protocol.Reward``) is already +*per-node*, and each seat in an arena run is one ``RolloutNode`` — so per-seat +scoring needs only a new, opt-in scorer, not a change to the scalar +``RewardFunc`` path every single-agent benchmark ships. Kept here as a pure +function of the floor's standings so it is trivially testable without the +rollout tree; wrap it in a ``Reward`` adapter to plug into the reward plane. +""" + +from __future__ import annotations + +import enum + +__all__ = ["FloorMode", "SharedEnvReward"] + + +class FloorMode(enum.StrEnum): + PVP = "pvp" # competitive: each seat scored by its own net result + COOP = "coop" # cooperative: every seat shares one joint outcome + + +class SharedEnvReward: + """Reduce one shared ``standings`` map to a per-seat reward vector.""" + + def __init__( + self, starting_bankroll: int = 1000, mode: FloorMode = FloorMode.PVP + ) -> None: + self.start = starting_bankroll + self.mode = mode + + def score(self, standings: dict[str, int]) -> dict[str, float]: + if not standings: + return {} + if self.mode is FloorMode.COOP: + # joint reduction: the floor succeeds together — the worst seat's + # net is everyone's reward (swap for sum/mean as a task needs). + joint = float(min(standings.values()) - self.start) + return {seat: joint for seat in standings} + # PvP: each seat's net chips (zero-sum floors sum to ~0). + return {seat: float(chips - self.start) for seat, chips in standings.items()} diff --git a/src/benchflow/arena/runtime.py b/src/benchflow/arena/runtime.py new file mode 100644 index 00000000..a3eb7824 --- /dev/null +++ b/src/benchflow/arena/runtime.py @@ -0,0 +1,74 @@ +"""Concurrent seat driver — Seam 1 + the eval glue. + +:func:`run_arena` runs N seats CONCURRENTLY against one shared env (via +``asyncio.gather``); each seat is an independent observe/act pull loop. The +shared env serializes turns, so N clocks still produce one well-ordered hand. A +worker cap bounds concurrent *decisions* (the expensive policy/LLM call) without +blocking seats that are merely waiting — so interdependent seats never deadlock. +A wall-clock deadline reaps stragglers (a seat with no partner, or a stalled +one). Pure ``asyncio``; the sequential scene path is left untouched. +""" + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable, Iterable + +from benchflow.arena.protocol import Observation, SeatClient, SeatPolicy, SeatStatus + +__all__ = ["run_arena"] + + +async def _run_seat( + seat_id: str, + client: SeatClient, + policy: SeatPolicy, + *, + sem: asyncio.Semaphore, + poll_s: float, + deadline: float, +) -> dict[str, object]: + acts = 0 + while True: + if time.monotonic() > deadline: + return {"seat": seat_id, "status": "deadline", "acts": acts} + obs = Observation.from_payload(await client.observe(seat_id)) + if obs.done: + return {"seat": seat_id, "status": "done", "acts": acts} + if obs.status is SeatStatus.YOUR_TURN: + async with sem: # cap concurrent decisions only — never the polling + action = await policy.act(obs) + await client.act(seat_id, obs.request_id or "", action) + acts += 1 + else: # waiting / not_your_turn — yield and poll again + await asyncio.sleep(poll_s) + + +async def run_arena( + seat_ids: Iterable[str], + client: SeatClient, + policy_for: Callable[[str], SeatPolicy], + *, + workers: int = 16, + deadline_s: float = 120.0, + poll_s: float = 0.05, +) -> dict[str, dict[str, object]]: + """Drive ``seat_ids`` concurrently against ``client`` until each is done or + the deadline fires. ``policy_for(seat_id)`` supplies that seat's brain. + Returns ``{seat_id: {status, acts, ...}}`` (status ∈ done | deadline | error). + """ + sem = asyncio.Semaphore(max(1, workers)) + deadline = time.monotonic() + deadline_s + + async def _guarded(seat_id: str) -> dict[str, object]: + try: + return await _run_seat( + seat_id, client, policy_for(seat_id), + sem=sem, poll_s=poll_s, deadline=deadline, + ) + except Exception as exc: # one bad seat must not kill the floor + return {"seat": seat_id, "status": "error", "error": repr(exc), "acts": 0} + + results = await asyncio.gather(*[_guarded(s) for s in seat_ids]) + return {str(r["seat"]): r for r in results} diff --git a/tests/test_arena.py b/tests/test_arena.py new file mode 100644 index 00000000..7e378bdb --- /dev/null +++ b/tests/test_arena.py @@ -0,0 +1,146 @@ +"""The inter-agent concurrent (arena) scaffold: N seats drive ONE shared, +turn-gated env at once. Exercised with an in-memory fake env + scripted policies +— no ACP, no sandbox, no LLM — so the concurrency, turn-gating, reaping, and +per-seat scoring are unit-testable on their own. +""" + +from __future__ import annotations + +import asyncio + +from benchflow.arena import ( + FloorMode, + Observation, + SharedEnvReward, + run_arena, +) + + +class _FakeFloor: + """A tiny turn-gated 2-seat env implementing the SeatClient contract. + + Each seat, on its turn, bids a number; the higher bid wins ``stake`` from the + lower (zero-sum). Seats lazy-join on first ``observe``; the table forms once + ``n_seats`` are present; turns are issued one at a time with a stable + ``request_id`` (so a stale id is rejected).""" + + def __init__(self, n_seats: int = 2, stake: int = 100, start: int = 1000) -> None: + self.n, self.stake, self.start = n_seats, stake, start + self.bankroll: dict[str, int] = {} + self.seated: list[str] = [] + self.bids: dict[str, int] = {} + self.pending: str | None = None + self.cur_rid: str | None = None + self.turn = 0 + self.formed = self.done = False + self.lock = asyncio.Lock() + + def _open_turn(self, seat: str) -> None: + self.pending = seat + self.turn += 1 + self.cur_rid = f"{seat}#{self.turn}" + + async def observe(self, seat_id: str) -> dict: + async with self.lock: + self.bankroll.setdefault(seat_id, self.start) + if not self.formed and seat_id not in self.seated: + self.seated.append(seat_id) + if not self.formed and len(self.seated) >= self.n: + self.formed = True + self._open_turn(self.seated[0]) + if self.done: + return {"status": "done", "bankroll": self.bankroll[seat_id]} + if not self.formed: + return {"status": "waiting"} + if self.pending != seat_id: + return {"status": "not_your_turn", "current_actor": self.pending} + return { + "status": "your_turn", "request_id": self.cur_rid, + "observation": {"public": {"pot": self.stake}, "private": {}}, + "legal_actions": [{"verb": "bid", "args": {"n": k}} for k in range(10)], + } + + async def act(self, seat_id: str, request_id: str, action: dict) -> dict: + async with self.lock: + if not self.formed or self.pending != seat_id: + return {"ok": False, "status": "not_your_turn"} + if request_id != self.cur_rid: + return {"ok": False, "status": "stale_request_id"} + self.bids[seat_id] = int(action.get("args", {}).get("n", 0)) + idx = self.seated.index(seat_id) + if idx + 1 < len(self.seated): + self._open_turn(self.seated[idx + 1]) + else: + self._settle() + return {"ok": True, "status": "applied"} + + def _settle(self) -> None: + a, b = self.seated[0], self.seated[1] + if self.bids[a] != self.bids[b]: + win, lose = (a, b) if self.bids[a] > self.bids[b] else (b, a) + self.bankroll[win] += self.stake + self.bankroll[lose] -= self.stake + self.pending = self.cur_rid = None + self.done = True + + def standings(self) -> dict[str, int]: + return dict(self.bankroll) + + +class _FixedBid: + """A scripted seat brain: always bid the same number.""" + + def __init__(self, n: int) -> None: + self.n = n + + async def act(self, obs: Observation) -> dict: + return {"verb": "bid", "args": {"n": self.n}} + + +def test_two_seats_play_concurrently_and_settle_zero_sum() -> None: + env = _FakeFloor() + bids = {"seat-0": 7, "seat-1": 3} # seat-0 wins + res = asyncio.run(run_arena( + ["seat-0", "seat-1"], env, lambda s: _FixedBid(bids[s]), + deadline_s=5.0, poll_s=0.001, + )) + assert all(r["status"] == "done" for r in res.values()) + st = env.standings() + assert st == {"seat-0": 1100, "seat-1": 900} + assert sum(st.values()) == 2000 # zero-sum, conserved + + +def test_per_seat_reward_vector() -> None: + env = _FakeFloor() + bids = {"seat-0": 7, "seat-1": 3} + asyncio.run(run_arena(["seat-0", "seat-1"], env, lambda s: _FixedBid(bids[s]), + deadline_s=5.0, poll_s=0.001)) + pvp = SharedEnvReward(mode=FloorMode.PVP).score(env.standings()) + assert pvp == {"seat-0": 100.0, "seat-1": -100.0} + assert sum(pvp.values()) == 0.0 + coop = SharedEnvReward(mode=FloorMode.COOP).score(env.standings()) + assert coop == {"seat-0": -100.0, "seat-1": -100.0} # joint = worst seat + + +def test_not_your_turn_is_rejected() -> None: + async def scenario() -> None: + env = _FakeFloor() + await env.observe("seat-0") # seat-0 seats, pending + await env.observe("seat-1") # forms; seat-0 is pending + o0 = await env.observe("seat-0") + assert o0["status"] == "your_turn" + bad = await env.act( + "seat-1", o0["request_id"], {"verb": "bid", "args": {"n": 9}}) + assert bad == {"ok": False, "status": "not_your_turn"} # out of turn + stale = await env.act("seat-0", "bogus", {"verb": "bid", "args": {"n": 9}}) + assert stale["status"] == "stale_request_id" # wrong request_id + + asyncio.run(scenario()) + + +def test_lone_seat_is_reaped_by_deadline() -> None: + env = _FakeFloor(n_seats=2) # needs a partner + res = asyncio.run(run_arena(["solo"], env, lambda s: _FixedBid(5), + deadline_s=0.2, poll_s=0.01)) + assert res["solo"]["status"] == "deadline" # never formed, reaped + assert env.standings() == {"solo": 1000} # no chips moved From e31e392c68ded74d4e28e25635a3f20de15607bd Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Sun, 28 Jun 2026 21:36:46 +0000 Subject: [PATCH 02/54] feat(arena): route seats through the LiteLLM proxy + capture per-seat trajectories MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each arena seat's raw LLM now goes through BenchFlow's provider proxy (so usage/ cost is tracked per agent), and every turn is captured as a per-seat trajectory. - policy.py: ProxyChatPolicy + provider_config — reads BENCHFLOW_PROVIDER_BASE_URL /_API_KEY/_MODEL (the proxy the SDK injects; falls back to provider-native env), calls the OpenAI-compatible /chat/completions, tags each request with x-bf-seat, and records the call. - trajectory.py: SeatTrajectory — per-seat .trajectory.jsonl (status, observation, legal_actions, action, llm {model, messages, response, usage}). - runtime.py: run_arena gains an on_turn hook for per-seat decision capture. tests: ProxyChatPolicy routes through BENCHFLOW_PROVIDER_* (httpx MockTransport, no real LLM); SeatTrajectory writes per-seat jsonl; on_turn emits one record per move; a 3-seat concurrent run. examples: floor_deepseek.py (N-seat high-card via the proxy env) and run_through_proxy.py — a REAL run: 3 deepseek-v4 seats through a loopback proxy started by ensure_litellm_runtime. Verified live: the raw key is hidden from the seats (proxy isolation invariant), proxy usage is aggregated (3298 tokens, $0.0014, usage_source=provider_response), and each seat's decision trajectory is captured with real model content. --- examples/arena/floor_deepseek.py | 157 ++++++++++++++++++++++++++++ examples/arena/run_through_proxy.py | 88 ++++++++++++++++ src/benchflow/arena/__init__.py | 6 ++ src/benchflow/arena/policy.py | 101 ++++++++++++++++++ src/benchflow/arena/runtime.py | 15 ++- src/benchflow/arena/trajectory.py | 65 ++++++++++++ tests/test_arena.py | 101 ++++++++++++++++++ tests/test_arena_proxy.py | 85 +++++++++++++++ 8 files changed, 615 insertions(+), 3 deletions(-) create mode 100644 examples/arena/floor_deepseek.py create mode 100644 examples/arena/run_through_proxy.py create mode 100644 src/benchflow/arena/policy.py create mode 100644 src/benchflow/arena/trajectory.py create mode 100644 tests/test_arena_proxy.py diff --git a/examples/arena/floor_deepseek.py b/examples/arena/floor_deepseek.py new file mode 100644 index 00000000..8ee50359 --- /dev/null +++ b/examples/arena/floor_deepseek.py @@ -0,0 +1,157 @@ +"""Real arena run: N deepseek-v4 seats, routed through BenchFlow's provider proxy, +with a per-seat trajectory written for each. + +Each seat's raw LLM call goes through ``BENCHFLOW_PROVIDER_BASE_URL`` / +``BENCHFLOW_PROVIDER_API_KEY`` / ``BENCHFLOW_PROVIDER_MODEL`` (the proxy the SDK +injects in a real eval — there it also writes ``llm_trajectory.jsonl`` and +attributes usage per seat via the ``x-bf-seat`` tag). The bench's own per-seat +decision trajectory is written to ``out/arena-floor/.trajectory.jsonl``. + + # point the provider vars at the proxy (or, standalone, at the model API): + export BENCHFLOW_PROVIDER_BASE_URL=https://api.deepseek.com + export BENCHFLOW_PROVIDER_API_KEY=$DEEPSEEK_API_KEY + export BENCHFLOW_PROVIDER_MODEL=deepseek-v4-pro + uv run python examples/arena/floor_deepseek.py 3 +""" + +from __future__ import annotations + +import asyncio +import json +import re +import sys +from pathlib import Path + +import httpx + +from benchflow.arena import ( + ProxyChatPolicy, + SeatTrajectory, + SharedEnvReward, + provider_config, + run_arena, +) + + +class HighCardFloor: + """N-seat turn-gated env (SeatClient): each seat antes ``stake`` and picks + 0-9; the highest pick wins the pot (ties split). Seats lazy-join on observe.""" + + def __init__(self, n_seats: int, stake: int = 50, start: int = 1000) -> None: + self.n, self.stake, self.start = n_seats, stake, start + self.bankroll: dict[str, int] = {} + self.seated: list[str] = [] + self.picks: dict[str, int] = {} + self.pending: str | None = None + self.cur_rid: str | None = None + self.turn = 0 + self.formed = self.done = False + self.lock = asyncio.Lock() + + def _open(self, seat: str) -> None: + self.pending, self.turn = seat, self.turn + 1 + self.cur_rid = f"{seat}#{self.turn}" + + async def observe(self, seat_id: str) -> dict: + async with self.lock: + self.bankroll.setdefault(seat_id, self.start) + if not self.formed and seat_id not in self.seated: + self.seated.append(seat_id) + if not self.formed and len(self.seated) >= self.n: + self.formed = True + self._open(self.seated[0]) + if self.done: + return {"status": "done", "bankroll": self.bankroll[seat_id]} + if not self.formed: + return {"status": "waiting"} + if self.pending != seat_id: + return {"status": "not_your_turn", "current_actor": self.pending} + return { + "status": "your_turn", "request_id": self.cur_rid, + "observation": {"public": {"pot": self.stake * self.n}, "private": {}}, + "legal_actions": [{"verb": "pick", "args": {"n": k}} for k in range(10)], + } + + async def act(self, seat_id: str, request_id: str, action: dict) -> dict: + async with self.lock: + if not self.formed or self.pending != seat_id: + return {"ok": False, "status": "not_your_turn"} + if request_id != self.cur_rid: + return {"ok": False, "status": "stale_request_id"} + self.picks[seat_id] = int(action.get("args", {}).get("n", 0)) + idx = self.seated.index(seat_id) + if idx + 1 < len(self.seated): + self._open(self.seated[idx + 1]) + else: + self._settle() + return {"ok": True, "status": "applied"} + + def _settle(self) -> None: + hi = max(self.picks.values()) + winners = [s for s, v in self.picks.items() if v == hi] + share = (self.stake * self.n) // len(winners) + for s in self.seated: + self.bankroll[s] -= self.stake + for w in winners: + self.bankroll[w] += share + self.pending = self.cur_rid = None + self.done = True + + def standings(self) -> dict[str, int]: + return dict(self.bankroll) + + +def render_for(seat: str): + def render(obs) -> str: + ns = [a["args"]["n"] for a in obs.legal_actions] + return ( + f"You are {seat} in a one-round high-card game for a pot of " + f"{obs.public.get('pot')} chips. Pick ONE number from {ns[0]}..{ns[-1]}; " + "the single highest pick wins the whole pot (ties split it). " + "Reply with ONLY your number." + ) + return render + + +def pick_number(text: str, legal: list[dict]) -> dict: + m = re.search(r"\d", text or "") + n = int(m.group()) if m else None + for a in legal: + if a["args"]["n"] == n: + return a + import random + return random.choice(legal) + + +async def _main() -> None: + n = int(sys.argv[1]) if len(sys.argv) > 1 else 3 + base, key, model = provider_config() + if not key: + raise SystemExit("no provider key (BENCHFLOW_PROVIDER_API_KEY / DEEPSEEK_API_KEY)") + run_dir = Path("out/arena-floor") + tr = SeatTrajectory(run_dir) + floor = HighCardFloor(n) + seats = [f"seat-{i}" for i in range(n)] + print(f"arena: {n} seats · provider {base} · model {model}", flush=True) + async with httpx.AsyncClient() as http: + policies = { + s: ProxyChatPolicy(s, http, render=render_for(s), pick=pick_number, + temperature=0.9, recorder=tr) + for s in seats + } + res = await run_arena(seats, floor, lambda s: policies[s], + deadline_s=120.0, poll_s=0.05) + st = floor.standings() + print("picks :", floor.picks) + print("standings :", st) + print("reward (pvp):", SharedEnvReward().score(st)) + print("seat status :", {s: r["status"] for s, r in res.items()}) + print("conserved :", sum(st.values()), f"(== {n * 1000})") + print(f"trajectories: {run_dir}/.trajectory.jsonl") + for s in seats: + rec = json.loads(tr.path(s).read_text().strip().splitlines()[-1]) + print(f" {s}: pick={rec['action']['args']['n']} llm.usage={rec['llm']['usage']}") + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/examples/arena/run_through_proxy.py b/examples/arena/run_through_proxy.py new file mode 100644 index 00000000..2a0a4380 --- /dev/null +++ b/examples/arena/run_through_proxy.py @@ -0,0 +1,88 @@ +"""Real run THROUGH the BenchFlow LiteLLM proxy. + +Three deepseek-v4 seats play one shared high-card round concurrently via +``run_arena``; each seat's raw LLM call is routed through a loopback LiteLLM proxy +started by ``ensure_litellm_runtime`` — so per-agent **usage/cost** and the proxy's +**llm trajectory** are captured by BenchFlow, and the seats never see the raw +provider key (the proxy isolation invariant). A per-seat *decision* trajectory is +also written to ``out/arena-floor-proxy/.trajectory.jsonl``. + + set -a; . ./sb-run.env; set +a # DEEPSEEK_API_KEY (the real upstream key) + uv run python examples/arena/run_through_proxy.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from pathlib import Path + +import httpx + +from benchflow.arena import ProxyChatPolicy, SeatTrajectory, SharedEnvReward, run_arena +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from floor_deepseek import HighCardFloor, pick_number, render_for + + +async def _main() -> None: + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required (the real upstream key)") + n = int(sys.argv[1]) if len(sys.argv) > 1 else 3 + seats = [f"seat-{i}" for i in range(n)] + + print("starting BenchFlow LiteLLM proxy (environment=local)…", flush=True) + agent_env, runtime = await ensure_litellm_runtime( + agent="deepagents", + agent_env={"DEEPSEEK_API_KEY": os.environ["DEEPSEEK_API_KEY"]}, + model="deepseek/deepseek-v4-pro", # provider-prefixed → routes to deepseek + runtime=None, + environment="local", + session_id="arena-floor", + ) + os.environ.update(agent_env) # each seat now reads BENCHFLOW_PROVIDER_* → the proxy + print(" proxy base :", agent_env.get("BENCHFLOW_PROVIDER_BASE_URL")) + print(" model alias:", agent_env.get("BENCHFLOW_PROVIDER_MODEL")) + print(" raw key hidden from seats:", + "DEEPSEEK_API_KEY" not in agent_env, flush=True) + + run_dir = Path("out/arena-floor-proxy") + tr = SeatTrajectory(run_dir) + floor = HighCardFloor(n) + res: dict = {} + try: + async with httpx.AsyncClient() as http: + policies = { + s: ProxyChatPolicy(s, http, render=render_for(s), pick=pick_number, + temperature=0.9, max_tokens=2048, recorder=tr) + for s in seats + } + res = await run_arena(seats, floor, lambda s: policies[s], + deadline_s=180.0, poll_s=0.05) + finally: + await stop_provider_runtime(runtime) # parses the proxy callback log + usage = extract_usage(runtime) # aggregate tokens/cost after stop + + st = floor.standings() + print("\npicks :", floor.picks) + print("standings :", st) + print("reward (pvp):", SharedEnvReward().score(st)) + print("seat status :", {s: r["status"] for s, r in res.items()}) + print("conserved :", sum(st.values()), f"(== {n * 1000})") + print("proxy usage :", json.dumps(usage)) # tokens + cost from the proxy callback log + print(f"trajectories: {run_dir}/.trajectory.jsonl") + for s in seats: + line = tr.path(s).read_text().strip().splitlines()[-1] + rec = json.loads(line) + print(f" {s}: pick={rec['action']['args']['n']} llm.usage={rec['llm']['usage']}") + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/src/benchflow/arena/__init__.py b/src/benchflow/arena/__init__.py index 5faf7ee3..00976dd9 100644 --- a/src/benchflow/arena/__init__.py +++ b/src/benchflow/arena/__init__.py @@ -19,6 +19,7 @@ from __future__ import annotations +from benchflow.arena.policy import ProxyChatPolicy, provider_config from benchflow.arena.protocol import ( Observation, SeatClient, @@ -27,13 +28,18 @@ ) from benchflow.arena.reward import FloorMode, SharedEnvReward from benchflow.arena.runtime import run_arena +from benchflow.arena.trajectory import SeatTrajectory, TurnRecord __all__ = [ "FloorMode", "Observation", + "ProxyChatPolicy", "SeatClient", "SeatPolicy", "SeatStatus", + "SeatTrajectory", "SharedEnvReward", + "TurnRecord", + "provider_config", "run_arena", ] diff --git a/src/benchflow/arena/policy.py b/src/benchflow/arena/policy.py new file mode 100644 index 00000000..911077e0 --- /dev/null +++ b/src/benchflow/arena/policy.py @@ -0,0 +1,101 @@ +"""A seat brain that routes its raw LLM call through BenchFlow's provider proxy. + +Reads the proxy config the SDK injects — ``BENCHFLOW_PROVIDER_BASE_URL`` / +``BENCHFLOW_PROVIDER_API_KEY`` / ``BENCHFLOW_PROVIDER_MODEL`` (the same vars +``deepagents_acp_shim`` reads, falling back to provider-native vars) — and calls +the OpenAI-compatible ``/chat/completions`` endpoint. Inside a BenchFlow eval that +endpoint is the LiteLLM proxy, so every seat's raw-LLM usage is tracked per agent +(``llm_trajectory.jsonl``); each request is also tagged with its ``seat`` so the +proxy can attribute calls. The caller supplies ``render(obs) -> prompt`` and +``pick(text, legal) -> action`` so the policy stays game-agnostic. +""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import Any + +import httpx + +from benchflow.arena.protocol import Observation +from benchflow.arena.trajectory import SeatTrajectory + +__all__ = ["provider_config", "ProxyChatPolicy"] + + +def provider_config(model_default: str | None = None) -> tuple[str, str, str]: + """(base_url, api_key, model) — prefer the BenchFlow proxy, fall back to + provider-native env. ``base_url`` is the proxy when running under an eval.""" + base = ( + os.environ.get("BENCHFLOW_PROVIDER_BASE_URL") + or os.environ.get("DEEPSEEK_BASE_URL") + or os.environ.get("OPENAI_BASE_URL") + or "https://api.deepseek.com" + ).rstrip("/") + key = ( + os.environ.get("BENCHFLOW_PROVIDER_API_KEY") + or os.environ.get("DEEPSEEK_API_KEY") + or os.environ.get("OPENAI_API_KEY") + or "" + ) + model = os.environ.get("BENCHFLOW_PROVIDER_MODEL") or model_default or "deepseek-v4-pro" + return base, key, model + + +class ProxyChatPolicy: + """A ``SeatPolicy`` whose decision is one chat-completion through the proxy.""" + + def __init__( + self, + seat: str, + http: httpx.AsyncClient, + *, + render: Callable[[Observation], str], + pick: Callable[[str, list[dict[str, Any]]], dict[str, Any]], + model: str | None = None, + temperature: float = 0.7, + max_tokens: int = 256, + recorder: SeatTrajectory | None = None, + ) -> None: + self.seat, self.http = seat, http + self.render, self.pick = render, pick + self.base, self.key, self.model = provider_config(model) + self.temperature, self.max_tokens = temperature, max_tokens + self.recorder = recorder + + async def act(self, obs: Observation) -> dict[str, Any]: + messages = [{"role": "user", "content": self.render(obs)}] + text: str = "" + usage: dict[str, Any] | None = None + try: + resp = await self.http.post( + f"{self.base}/chat/completions", + headers={ + "Authorization": f"Bearer {self.key}", + "x-bf-seat": self.seat, # per-seat attribution for the proxy + }, + json={ + "model": self.model, + "messages": messages, + "temperature": self.temperature, + "max_tokens": self.max_tokens, + "metadata": {"seat": self.seat}, + }, + timeout=90.0, + ) + data = resp.json() + text = data["choices"][0]["message"]["content"] or "" + usage = data.get("usage") + except Exception as exc: # a flaky call falls back to a legal default + text = "" + usage = {"error": repr(exc)} + action = self.pick(text, list(obs.legal_actions)) + if self.recorder is not None: + self.recorder.record( + self.seat, status=obs.status, observation=obs.public, + legal_actions=obs.legal_actions, action=action, request_id=obs.request_id, + llm={"model": self.model, "messages": messages, "response": text, + "usage": usage}, + ) + return action diff --git a/src/benchflow/arena/runtime.py b/src/benchflow/arena/runtime.py index a3eb7824..0cc3abe4 100644 --- a/src/benchflow/arena/runtime.py +++ b/src/benchflow/arena/runtime.py @@ -20,6 +20,9 @@ __all__ = ["run_arena"] +OnTurn = Callable[[str, int, Observation, dict[str, object]], None] + + async def _run_seat( seat_id: str, client: SeatClient, @@ -28,6 +31,7 @@ async def _run_seat( sem: asyncio.Semaphore, poll_s: float, deadline: float, + on_turn: OnTurn | None, ) -> dict[str, object]: acts = 0 while True: @@ -41,6 +45,8 @@ async def _run_seat( action = await policy.act(obs) await client.act(seat_id, obs.request_id or "", action) acts += 1 + if on_turn is not None: # the bench's per-seat decision trajectory + on_turn(seat_id, acts, obs, action) else: # waiting / not_your_turn — yield and poll again await asyncio.sleep(poll_s) @@ -53,10 +59,13 @@ async def run_arena( workers: int = 16, deadline_s: float = 120.0, poll_s: float = 0.05, + on_turn: OnTurn | None = None, ) -> dict[str, dict[str, object]]: """Drive ``seat_ids`` concurrently against ``client`` until each is done or - the deadline fires. ``policy_for(seat_id)`` supplies that seat's brain. - Returns ``{seat_id: {status, acts, ...}}`` (status ∈ done | deadline | error). + the deadline fires. ``policy_for(seat_id)`` supplies that seat's brain; + ``on_turn(seat, turn, obs, action)`` (optional) is called after each move — + the hook the bench uses to capture per-seat trajectories. Returns + ``{seat_id: {status, acts, ...}}`` (status ∈ done | deadline | error). """ sem = asyncio.Semaphore(max(1, workers)) deadline = time.monotonic() + deadline_s @@ -65,7 +74,7 @@ async def _guarded(seat_id: str) -> dict[str, object]: try: return await _run_seat( seat_id, client, policy_for(seat_id), - sem=sem, poll_s=poll_s, deadline=deadline, + sem=sem, poll_s=poll_s, deadline=deadline, on_turn=on_turn, ) except Exception as exc: # one bad seat must not kill the floor return {"seat": seat_id, "status": "error", "error": repr(exc), "acts": 0} diff --git a/src/benchflow/arena/trajectory.py b/src/benchflow/arena/trajectory.py new file mode 100644 index 00000000..cfe0296a --- /dev/null +++ b/src/benchflow/arena/trajectory.py @@ -0,0 +1,65 @@ +"""Per-seat trajectory capture for arena runs. + +Each seat's turn — observation, chosen action, and (when the seat is LLM-backed) +the raw model request/response/usage — is appended as one JSON line to a per-seat +file. This is the bench's own per-agent decision trajectory; it complements the +LiteLLM proxy's ``llm_trajectory.jsonl`` (the raw model calls, captured whenever a +seat routes its LLM through ``BENCHFLOW_PROVIDER_BASE_URL``). +""" + +from __future__ import annotations + +import json +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +__all__ = ["TurnRecord", "SeatTrajectory"] + + +@dataclass +class TurnRecord: + seat: str + turn: int + status: str + request_id: str | None = None + observation: dict[str, Any] = field(default_factory=dict) + legal_actions: list[Any] = field(default_factory=list) + action: dict[str, Any] | None = None + llm: dict[str, Any] | None = None # {model, messages, response, usage} + t: float = 0.0 + + +class SeatTrajectory: + """Append-only per-seat trajectory writer under one rollout dir.""" + + def __init__(self, root: str | Path) -> None: + self.root = Path(root) + self.root.mkdir(parents=True, exist_ok=True) + self._turns: dict[str, int] = {} + + def path(self, seat: str) -> Path: + return self.root / f"{seat}.trajectory.jsonl" + + def record( + self, + seat: str, + *, + status: Any, + observation: dict[str, Any] | None = None, + legal_actions: list[Any] | None = None, + action: dict[str, Any] | None = None, + llm: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> TurnRecord: + turn = self._turns.get(seat, 0) + 1 + self._turns[seat] = turn + rec = TurnRecord( + seat=seat, turn=turn, status=str(status), request_id=request_id, + observation=dict(observation or {}), legal_actions=list(legal_actions or []), + action=action, llm=llm, t=time.time(), + ) + with self.path(seat).open("a") as f: + f.write(json.dumps(asdict(rec)) + "\n") + return rec diff --git a/tests/test_arena.py b/tests/test_arena.py index 7e378bdb..eb85f9ff 100644 --- a/tests/test_arena.py +++ b/tests/test_arena.py @@ -144,3 +144,104 @@ def test_lone_seat_is_reaped_by_deadline() -> None: deadline_s=0.2, poll_s=0.01)) assert res["solo"]["status"] == "deadline" # never formed, reaped assert env.standings() == {"solo": 1000} # no chips moved + + +class _HighCard: + """N-seat turn-gated env: each seat antes ``stake`` and picks 0-9; the + highest pick wins the pot (ties split). Implements the SeatClient contract.""" + + def __init__(self, n_seats: int, stake: int = 50, start: int = 1000) -> None: + self.n, self.stake, self.start = n_seats, stake, start + self.bankroll: dict[str, int] = {} + self.seated: list[str] = [] + self.picks: dict[str, int] = {} + self.pending: str | None = None + self.cur_rid: str | None = None + self.turn = 0 + self.formed = self.done = False + self.lock = asyncio.Lock() + + def _open(self, seat: str) -> None: + self.pending, self.turn = seat, self.turn + 1 + self.cur_rid = f"{seat}#{self.turn}" + + async def observe(self, seat_id: str) -> dict: + async with self.lock: + self.bankroll.setdefault(seat_id, self.start) + if not self.formed and seat_id not in self.seated: + self.seated.append(seat_id) + if not self.formed and len(self.seated) >= self.n: + self.formed = True + self._open(self.seated[0]) + if self.done: + return {"status": "done", "bankroll": self.bankroll[seat_id]} + if not self.formed: + return {"status": "waiting"} + if self.pending != seat_id: + return {"status": "not_your_turn", "current_actor": self.pending} + return { + "status": "your_turn", "request_id": self.cur_rid, + "observation": {"public": {"pot": self.stake * self.n}, "private": {}}, + "legal_actions": [{"verb": "pick", "args": {"n": k}} for k in range(10)], + } + + async def act(self, seat_id: str, request_id: str, action: dict) -> dict: + async with self.lock: + if not self.formed or self.pending != seat_id: + return {"ok": False, "status": "not_your_turn"} + if request_id != self.cur_rid: + return {"ok": False, "status": "stale_request_id"} + self.picks[seat_id] = int(action.get("args", {}).get("n", 0)) + idx = self.seated.index(seat_id) + if idx + 1 < len(self.seated): + self._open(self.seated[idx + 1]) + else: + self._settle() + return {"ok": True, "status": "applied"} + + def _settle(self) -> None: + hi = max(self.picks.values()) + winners = [s for s, v in self.picks.items() if v == hi] + share = (self.stake * self.n) // len(winners) + for s in self.seated: + self.bankroll[s] -= self.stake # everyone antes + for w in winners: + self.bankroll[w] += share # winners split the pot + self.pending = self.cur_rid = None + self.done = True + + def standings(self) -> dict[str, int]: + return dict(self.bankroll) + + +class _FixedPick: + def __init__(self, n: int) -> None: + self.n = n + + async def act(self, obs: Observation) -> dict: + return {"verb": "pick", "args": {"n": self.n}} + + +def test_three_seats_play_one_shared_table_concurrently() -> None: + env = _HighCard(3, stake=50) + picks = {"a": 3, "b": 7, "c": 5} # b wins + res = asyncio.run(run_arena( + ["a", "b", "c"], env, lambda s: _FixedPick(picks[s]), + deadline_s=5.0, poll_s=0.001, + )) + assert all(r["status"] == "done" for r in res.values()) + st = env.standings() + assert st == {"a": 950, "b": 1100, "c": 950} # b: -50 ante +150 pot + assert sum(st.values()) == 3000 # zero-sum, conserved + + +def test_run_arena_emits_per_turn_records() -> None: + env = _FakeFloor() + recs: list[tuple[str, int, int]] = [] + asyncio.run(run_arena( + ["seat-0", "seat-1"], env, lambda s: _FixedBid({"seat-0": 7, "seat-1": 3}[s]), + on_turn=lambda seat, turn, obs, act: recs.append((seat, turn, act["args"]["n"])), + deadline_s=5.0, poll_s=0.001, + )) + assert sorted((s, n) for s, _, n in recs) == [("seat-0", 7), ("seat-1", 3)] + assert all(turn == 1 for _, turn, _ in recs) # one turn each diff --git a/tests/test_arena_proxy.py b/tests/test_arena_proxy.py new file mode 100644 index 00000000..f666b770 --- /dev/null +++ b/tests/test_arena_proxy.py @@ -0,0 +1,85 @@ +"""Proxy routing + trajectory capture: each seat's raw LLM goes through the +BenchFlow provider proxy (``BENCHFLOW_PROVIDER_*``), and each turn is recorded to +a per-seat trajectory. Uses ``httpx.MockTransport`` — no real proxy or LLM. +""" + +from __future__ import annotations + +import asyncio +import json + +import httpx +import pytest + +from benchflow.arena import ( + Observation, + ProxyChatPolicy, + SeatStatus, + SeatTrajectory, +) + + +def test_seat_trajectory_writes_per_seat_jsonl(tmp_path) -> None: + tr = SeatTrajectory(tmp_path) + tr.record("a", status=SeatStatus.YOUR_TURN, observation={"pot": 150}, + action={"verb": "pick", "args": {"n": 7}}) + tr.record("a", status=SeatStatus.YOUR_TURN, action={"verb": "pick", "args": {"n": 3}}) + tr.record("b", status=SeatStatus.YOUR_TURN, action={"verb": "pick", "args": {"n": 5}}) + lines = tr.path("a").read_text().strip().splitlines() + assert len(lines) == 2 + r0 = json.loads(lines[0]) + assert r0["seat"] == "a" and r0["turn"] == 1 and r0["status"] == "your_turn" + assert r0["action"]["args"]["n"] == 7 + assert json.loads(lines[1])["turn"] == 2 + assert tr.path("b").exists() + + +def test_proxy_chat_policy_routes_through_provider_and_records( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("BENCHFLOW_PROVIDER_BASE_URL", "http://proxy.local/v1") + monkeypatch.setenv("BENCHFLOW_PROVIDER_API_KEY", "bf-key") + monkeypatch.setenv("BENCHFLOW_PROVIDER_MODEL", "proxy-model") + seen: dict = {} + + def handler(req: httpx.Request) -> httpx.Response: + seen["url"] = str(req.url) + seen["auth"] = req.headers.get("authorization") + seen["seat"] = req.headers.get("x-bf-seat") + seen["body"] = json.loads(req.content) + return httpx.Response(200, json={ + "choices": [{"message": {"content": "I choose rock."}}], + "usage": {"total_tokens": 11}, + }) + + tr = SeatTrajectory(tmp_path) + + async def scenario() -> dict: + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http: + legal = [{"verb": "throw", "args": {"hand": h}} + for h in ("rock", "paper", "scissors")] + pol = ProxyChatPolicy( + "seat-0", http, + render=lambda o: "pick one", + pick=lambda text, lg: next( + (a for a in lg if a["args"]["hand"] in text.lower()), lg[0]), + recorder=tr, + ) + obs = Observation(status=SeatStatus.YOUR_TURN, request_id="r1", + public={"pot": 100}, legal_actions=legal) + return await pol.act(obs) + + action = asyncio.run(scenario()) + assert action == {"verb": "throw", "args": {"hand": "rock"}} + # routed through the PROXY base, with the proxy key/model + a per-seat tag + assert seen["url"] == "http://proxy.local/v1/chat/completions" + assert seen["auth"] == "Bearer bf-key" + assert seen["seat"] == "seat-0" + assert seen["body"]["model"] == "proxy-model" + assert seen["body"]["metadata"] == {"seat": "seat-0"} + # the per-seat trajectory captured the decision + the raw llm call + rec = json.loads(tr.path("seat-0").read_text().strip()) + assert rec["action"]["args"]["hand"] == "rock" + assert rec["llm"]["model"] == "proxy-model" + assert rec["llm"]["usage"]["total_tokens"] == 11 + assert rec["llm"]["response"] == "I choose rock." From 64ddfb20c0391e5639523a491562d23dc62a8cfe Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Sun, 28 Jun 2026 21:40:11 +0000 Subject: [PATCH 03/54] feat(arena): persist proxy raw-LLM trajectory in BenchFlow canonical format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After the run, write the proxy's captured exchanges to /trajectory/llm_trajectory.jsonl via trajectory.to_jsonl(redact_keys=True) — mirroring rollout._write_llm_trajectory — so the per-seat raw-LLM calls land in the same on-disk trajectory format the viewer/verifier already read. Verified live: 3 deepseek-v4 seats through the proxy -> 3 raw exchanges captured + per-seat usage/cost, alongside the per-seat decision trajectories. --- examples/arena/run_through_proxy.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/examples/arena/run_through_proxy.py b/examples/arena/run_through_proxy.py index 2a0a4380..6fe5ed8f 100644 --- a/examples/arena/run_through_proxy.py +++ b/examples/arena/run_through_proxy.py @@ -70,6 +70,14 @@ async def _main() -> None: await stop_provider_runtime(runtime) # parses the proxy callback log usage = extract_usage(runtime) # aggregate tokens/cost after stop + # persist the proxy's raw-LLM trajectory in BenchFlow's canonical format + # (mirrors rollout._write_llm_trajectory). + proxy_traj = getattr(getattr(runtime, "server", None), "trajectory", None) + if proxy_traj is not None and proxy_traj.exchanges: + traj_dir = run_dir / "trajectory" + traj_dir.mkdir(parents=True, exist_ok=True) + (traj_dir / "llm_trajectory.jsonl").write_text(proxy_traj.to_jsonl(redact_keys=True)) + st = floor.standings() print("\npicks :", floor.picks) print("standings :", st) @@ -77,7 +85,10 @@ async def _main() -> None: print("seat status :", {s: r["status"] for s, r in res.items()}) print("conserved :", sum(st.values()), f"(== {n * 1000})") print("proxy usage :", json.dumps(usage)) # tokens + cost from the proxy callback log - print(f"trajectories: {run_dir}/.trajectory.jsonl") + if proxy_traj is not None and proxy_traj.exchanges: + print(f"llm_trajectory: {run_dir}/trajectory/llm_trajectory.jsonl " + f"({len(proxy_traj.exchanges)} raw exchanges, canonical format)") + print(f"decision traj : {run_dir}/.trajectory.jsonl") for s in seats: line = tr.path(s).read_text().strip().splitlines()[-1] rec = json.loads(line) From e86cf42ca3314f7a6c1e748451bd961441cad228 Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Sun, 28 Jun 2026 23:16:22 +0000 Subject: [PATCH 04/54] feat(arena): add a LangGraph medical-assistant agent hosted via the proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Demonstrates the intra-agent axis alongside the inter-agent arena: a real langgraph.StateGraph in the Multi-Agent-Medical-Assistant pattern (supervisor/ router -> KB / web-search specialists -> confidence-gated handoff -> output guardrail), run HOSTED BY BenchFlow — every node's LLM goes through a loopback LiteLLM proxy (ensure_litellm_runtime), so per-node usage/cost is tracked and the raw-LLM exchanges persist to /trajectory/llm_trajectory.jsonl. The agent never sees the raw provider key. Verified live (deepseek-v4-pro): path supervisor -> retrieve_kb -> answer -> web_search -> answer -> guardrail (confidence gate fired), 4 node LLM calls captured, usage/cost aggregated through the proxy. langgraph/langchain-openai are example-only deps (not added to the core). --- examples/medical/medical_assistant.py | 133 ++++++++++++++++++++++++++ examples/medical/run_through_proxy.py | 80 ++++++++++++++++ 2 files changed, 213 insertions(+) create mode 100644 examples/medical/medical_assistant.py create mode 100644 examples/medical/run_through_proxy.py diff --git a/examples/medical/medical_assistant.py b/examples/medical/medical_assistant.py new file mode 100644 index 00000000..7f4c5901 --- /dev/null +++ b/examples/medical/medical_assistant.py @@ -0,0 +1,133 @@ +"""A LangGraph supervisor→specialists agent — the Multi-Agent-Medical-Assistant +PATTERN, hosted by BenchFlow. + +A real ``langgraph.StateGraph``: a supervisor/router that dispatches to a KB +(RAG-style) or web-search specialist, an answer node that emits a confidence, a +**confidence-gated handoff** to the web specialist when confidence is low, and an +output **guardrail** node — mirroring the medical assistant's router + specialists ++ confidence-based handoff + guardrails (minus its Qdrant/CV/web stack). Every +node's LLM call uses ``langchain_openai.ChatOpenAI`` pointed at the BenchFlow +provider proxy (``BENCHFLOW_PROVIDER_*``), so per-node usage + the raw-LLM +trajectory are captured by BenchFlow. +""" + +from __future__ import annotations + +import os +import re +from typing import TypedDict + +from langchain_openai import ChatOpenAI +from langgraph.graph import END, START, StateGraph + +# A tiny stand-in knowledge base — the RAG specialist's corpus. +_KB = { + "metformin": "Metformin is first-line for type 2 diabetes; common side effects " + "are GI upset (nausea, diarrhea) and, rarely, lactic acidosis.", + "aspirin": "Low-dose aspirin is used for secondary cardiovascular prevention; " + "bleeding risk rises with dose and with anticoagulants.", + "ibuprofen": "Ibuprofen is an NSAID for pain/inflammation; avoid in renal " + "impairment and use caution with anticoagulants.", +} + + +class S(TypedDict, total=False): + query: str + route: str + evidence: str + answer: str + confidence: float + safe: bool + web_tried: bool + trace: list + + +def _llm() -> ChatOpenAI: + base = os.environ.get("BENCHFLOW_PROVIDER_BASE_URL") or "https://api.deepseek.com/v1" + key = (os.environ.get("BENCHFLOW_PROVIDER_API_KEY") + or os.environ.get("DEEPSEEK_API_KEY") or "EMPTY") + model = os.environ.get("BENCHFLOW_PROVIDER_MODEL") or "deepseek/deepseek-v4-pro" + return ChatOpenAI(model=model, base_url=base, api_key=key, + temperature=0.2, max_tokens=2048) + + +def _log(state: S, node: str, note: str) -> None: + state.setdefault("trace", []).append({"node": node, "note": note}) + + +def supervisor(state: S) -> S: + out = _llm().invoke([{"role": "user", "content": + f"You route a medical question to a specialist. Question: {state['query']!r}. " + "Reply with ONE word: 'kb' if a drug/condition knowledge base likely covers " + "it, else 'web'."}]).content.lower() + route = "web" if "web" in out else "kb" + _log(state, "supervisor", f"route={route}") + return {"route": route, "trace": state["trace"]} + + +def retrieve_kb(state: S) -> S: + q = state["query"].lower() + hits = [v for k, v in _KB.items() if k in q] + _log(state, "retrieve_kb", f"hits={len(hits)}") + return {"evidence": " ".join(hits), "trace": state["trace"]} + + +def web_search(state: S) -> S: + ev = f"[web] current guidance relevant to: {state['query']}" + _log(state, "web_search", "fallback evidence") + return {"evidence": (state.get("evidence", "") + " " + ev).strip(), + "web_tried": True, "trace": state["trace"]} + + +def answer(state: S) -> S: + out = _llm().invoke([{"role": "user", "content": + f"Question: {state['query']}\nEvidence: {state.get('evidence') or '(none)'}\n" + "Answer concisely for a clinician. Then on a NEW line output " + "'CONFIDENCE: <0.0-1.0>' for how well the evidence supports your answer."}] + ).content + m = re.search(r"CONFIDENCE:\s*([01](?:\.\d+)?)", out) + conf = float(m.group(1)) if m else 0.5 + ans = re.split(r"CONFIDENCE:", out)[0].strip() + _log(state, "answer", f"confidence={conf}") + return {"answer": ans, "confidence": conf, "trace": state["trace"]} + + +def guardrail(state: S) -> S: + out = _llm().invoke([{"role": "user", "content": + f"Output guardrail. The assistant answered: {state['answer']!r}. Is it safe and " + "appropriately cautious as medical information (not a diagnosis)? Reply yes/no."}] + ).content.lower() + _log(state, "guardrail", f"safe={'yes' in out}") + return {"safe": "yes" in out, "trace": state["trace"]} + + +def _after_supervisor(state: S) -> str: + return "web_search" if state["route"] == "web" else "retrieve_kb" + + +def _after_answer(state: S) -> str: + threshold = float(os.environ.get("MEDICAL_CONFIDENCE_THRESHOLD", "0.6")) + if state.get("confidence", 0.0) < threshold and not state.get("web_tried"): + return "web_search" # confidence-gated handoff to the fallback specialist + return "guardrail" + + +def build_graph(): + g = StateGraph(S) + for name, fn in (("supervisor", supervisor), ("retrieve_kb", retrieve_kb), + ("web_search", web_search), ("answer", answer), + ("guardrail", guardrail)): + g.add_node(name, fn) + g.add_edge(START, "supervisor") + g.add_conditional_edges("supervisor", _after_supervisor, + {"retrieve_kb": "retrieve_kb", "web_search": "web_search"}) + g.add_edge("retrieve_kb", "answer") + g.add_edge("web_search", "answer") + g.add_conditional_edges("answer", _after_answer, + {"web_search": "web_search", "guardrail": "guardrail"}) + g.add_edge("guardrail", END) + return g.compile() + + +def run(query: str) -> S: + return build_graph().invoke({"query": query, "trace": []}) diff --git a/examples/medical/run_through_proxy.py b/examples/medical/run_through_proxy.py new file mode 100644 index 00000000..b2f0257d --- /dev/null +++ b/examples/medical/run_through_proxy.py @@ -0,0 +1,80 @@ +"""Run the LangGraph Medical-Assistant agent HOSTED BY BenchFlow. + +Starts a loopback LiteLLM proxy via ``ensure_litellm_runtime`` and runs the +supervisor→specialists graph against it, so every node's raw LLM call (router, +answer, guardrail, and the confidence-gated web fallback) is tracked by the proxy +— usage/cost aggregated and the raw-LLM exchanges persisted to +``out/medical-assistant/trajectory/llm_trajectory.jsonl`` in BenchFlow's canonical +format. The agent never sees the raw provider key (proxy isolation invariant). + + uv pip install langgraph langchain-openai # example-only deps + set -a; . ./sb-run.env; set +a + uv run python examples/medical/run_through_proxy.py "side effects of metformin?" + # MEDICAL_CONFIDENCE_THRESHOLD=1.01 forces the confidence-gated web handoff. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from pathlib import Path + +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from medical_assistant import run + + +async def _main() -> None: + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required (the real upstream key)") + query = sys.argv[1] if len(sys.argv) > 1 else \ + "What are the main side effects of metformin?" + + print("starting BenchFlow LiteLLM proxy (environment=local)…", flush=True) + agent_env, runtime = await ensure_litellm_runtime( + agent="deepagents", + agent_env={"DEEPSEEK_API_KEY": os.environ["DEEPSEEK_API_KEY"]}, + model="deepseek/deepseek-v4-pro", + runtime=None, + environment="local", + session_id="medical-assistant", + ) + os.environ.update(agent_env) # the graph's ChatOpenAI now reads BENCHFLOW_PROVIDER_* + print(" proxy base :", agent_env.get("BENCHFLOW_PROVIDER_BASE_URL")) + print(" raw key hidden from agent:", "DEEPSEEK_API_KEY" not in agent_env, flush=True) + + result: dict = {} + try: + result = await asyncio.to_thread(run, query) # the LangGraph agent + finally: + await stop_provider_runtime(runtime) + usage = extract_usage(runtime) + + run_dir = Path("out/medical-assistant") + proxy_traj = getattr(getattr(runtime, "server", None), "trajectory", None) + if proxy_traj is not None and proxy_traj.exchanges: + (run_dir / "trajectory").mkdir(parents=True, exist_ok=True) + (run_dir / "trajectory" / "llm_trajectory.jsonl").write_text( + proxy_traj.to_jsonl(redact_keys=True)) + + path = " → ".join(t["node"] for t in result.get("trace", [])) + print("\nquery :", query) + print("agent path :", path) + print("route :", result.get("route"), "| confidence:", result.get("confidence"), + "| guardrail safe:", result.get("safe")) + print("answer :", (result.get("answer") or "").replace("\n", " ")[:280]) + print("proxy usage:", json.dumps(usage)) + if proxy_traj is not None and proxy_traj.exchanges: + print(f"llm_trajectory: {run_dir}/trajectory/llm_trajectory.jsonl " + f"({len(proxy_traj.exchanges)} node LLM calls, canonical format)") + + +if __name__ == "__main__": + asyncio.run(_main()) From bd2f96958b775a611afe28f322a1c72f1f4a6698 Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Sun, 28 Jun 2026 23:22:41 +0000 Subject: [PATCH 05/54] chore(arena): sample run-result bundle + callback-flush delay examples/sample_runs/ packs two real proxy-hosted runs (3-seat arena + LangGraph medical assistant) with their trajectories + a README of the commands/config. Adds a short pre-stop flush delay in the runners so the proxy's canonical llm_trajectory.jsonl captures as many calls as possible (final-callback flush is a LiteLLM timing artifact; per-seat decision trajectories are complete). --- examples/arena/run_through_proxy.py | 1 + examples/medical/run_through_proxy.py | 1 + examples/sample_runs/README.md | 37 +++++++++++++++++++++++++ examples/sample_runs/proxy-runs.tar.gz | Bin 0 -> 8952 bytes 4 files changed, 39 insertions(+) create mode 100644 examples/sample_runs/README.md create mode 100644 examples/sample_runs/proxy-runs.tar.gz diff --git a/examples/arena/run_through_proxy.py b/examples/arena/run_through_proxy.py index 6fe5ed8f..c5706a14 100644 --- a/examples/arena/run_through_proxy.py +++ b/examples/arena/run_through_proxy.py @@ -67,6 +67,7 @@ async def _main() -> None: res = await run_arena(seats, floor, lambda s: policies[s], deadline_s=180.0, poll_s=0.05) finally: + await asyncio.sleep(1.5) # let the proxy's async callback flush the last call await stop_provider_runtime(runtime) # parses the proxy callback log usage = extract_usage(runtime) # aggregate tokens/cost after stop diff --git a/examples/medical/run_through_proxy.py b/examples/medical/run_through_proxy.py index b2f0257d..77b4f2ab 100644 --- a/examples/medical/run_through_proxy.py +++ b/examples/medical/run_through_proxy.py @@ -54,6 +54,7 @@ async def _main() -> None: try: result = await asyncio.to_thread(run, query) # the LangGraph agent finally: + await asyncio.sleep(1.5) # let the proxy's async callback flush the last call await stop_provider_runtime(runtime) usage = extract_usage(runtime) diff --git a/examples/sample_runs/README.md b/examples/sample_runs/README.md new file mode 100644 index 00000000..0cdd2275 --- /dev/null +++ b/examples/sample_runs/README.md @@ -0,0 +1,37 @@ +# Sample runs — multi-agent hosted by BenchFlow (proxy-tracked + trajectories) + +Every agent's raw LLM is routed through a loopback BenchFlow LiteLLM proxy +(`ensure_litellm_runtime`, `environment="local"`); the proxy aggregates per-call +usage/cost and persists raw exchanges to `/trajectory/llm_trajectory.jsonl` +in BenchFlow's canonical format. The agent never sees the raw provider key +(proxy isolation invariant). + +## Shared config +- model: `deepseek/deepseek-v4-pro` (provider-prefixed → deepseek upstream) +- proxy: `ensure_litellm_runtime(agent="deepagents", model=..., environment="local")` +- upstream key: `DEEPSEEK_API_KEY` (stripped from the agent env by the proxy) + +## 1) Inter-agent arena — 3 concurrent seats → `arena-floor-proxy/` + set -a; . ./sb-run.env; set +a + uv run python examples/arena/run_through_proxy.py 3 +Each of 3 deepseek-v4 seats plays one shared high-card round concurrently via +`run_arena`. Outputs: `seat-*.trajectory.jsonl` (per-seat decisions) + +`trajectory/llm_trajectory.jsonl` (raw exchanges). Chips conserve at 3000. + +## 2) Intra-agent LangGraph medical assistant → `medical-assistant/` + uv pip install langgraph langchain-openai + set -a; . ./sb-run.env; set +a + MEDICAL_CONFIDENCE_THRESHOLD=1.01 \ + uv run python examples/medical/run_through_proxy.py "side effects of metformin?" +Real `langgraph.StateGraph` (supervisor → KB/web specialists → confidence-gated +handoff → guardrail). Path: supervisor → retrieve_kb → answer → web_search → +answer → guardrail. Outputs: `trajectory/llm_trajectory.jsonl` (one exchange per +node LLM call). + +## Note on capture fidelity +The per-seat / per-node *decision* trajectories (and each call's response usage) +are captured directly from each model response and are complete. The proxy's +canonical `trajectory/llm_trajectory.jsonl` is rebuilt from LiteLLM's async +callback log on stop and may capture N-1 of N calls when the final call's +callback hasn't flushed — a LiteLLM timing artifact, not a routing gap (every +call DID go through the proxy, as the aggregated usage/cost confirms). diff --git a/examples/sample_runs/proxy-runs.tar.gz b/examples/sample_runs/proxy-runs.tar.gz new file mode 100644 index 0000000000000000000000000000000000000000..6c5d52daa8578985a858a0e985627be0380a0d7d GIT binary patch literal 8952 zcmb8yLvS4qxCQXoxv_2Aw%ypaohE6_8{2Lg+jfJ-Y1G)Zb?^Jn+rP#87QfwDotZQ9 zk;NfE{O>_u8XJ0U$Ywsb1j%IGn+Z`+g!wWaXbUvilqOJBSh$vdt6`~D;8u_(50QW_ z%}K9$whta$0S&*=1k9cx6Hq(flLv1?#7vwfEAZy%h~n4~g0kw}^2Jp-e5{N64=3$@7z_Z*wv`g_T*1wS!sY}IysgX>B zL9V$6gW4z}H8yiCF6e7?xL|kx_pghkCD)3%0kOQGfOjlrjO}Q9Bql?v$z1xf30SJq zP5mHI8p2n#b^ZL2C@m^-`Dq1*0Qs!%Db}xgNC4&XK9g^0&rnvVRS?i5#-!}8@+mF6 z;MDKaLa8mhmrLOc13ynp-5kCdmFYBkOoUSVf+zCzdF)#AS)= zzq&f;5(nMLQTrp;PfB1IQq|V+;In)uIRd^dgjA4z_~ZP-#G{t+K^30VM^kRY9A)L8 zILJk0E1BEjEUaHMfT8|0Dm4f_bRa!FasDuRfI!HPVn7fSj;F*O5?3Nb2o7(Ha2HZA zf7tsW@Vs>v*V0flkRSAQ{@Gy5raCtuTG0P?9KE!<)9e4XpL63K=p%9Ea#&|#FXD?v zbBc)-2RpAynYKzPCxUj6HWE$-%#Y#ySGNn%i9sC)+h>GJPDg-7Wj6aFY)aBnAi)Kh3&U&U&r>$=otaOv zFKqG`ISxS#?dSL1-lnqdeeaIeViY2K2_}E^a0YUrBpaQLng?n;@~d0=nW3NTn?%_e zSCIdtnJVQTnV!I#2v;XKMZ<-~b-0y!{DOiw6q@PEvd^DqkajUSG-eXn5^Cz?ob-Pa zP_);Pqsv$kIYscM6tVKi{)8_^?WiC;I#}pLXw=&*T_a~s9Pt*C?fOrS$V_4?P&`nV zpREQuo_TCPYh=&b`=&bUP|~(^AXw@}{PPf41payDP{zaYhubp78yTxQceiz_MU<4)DoB_~vy{oNvx ze(oZ)&&!_i=pyLw&jmX-&Cet%#(Yx_B{IZ0z z4UJ-RnW=&nFSA0Y-f~1>Qs?d$yQLnvJdnsglza-wWfL4N?vYErJ9`LZO=bzM` zT&l8C9+8qk#En_@Nm!;P`;Ar`x{cA46-kTqV*Natl)-*jd6@plMvCY_NV2bVZ0t~m z$xs-GFzgtH4MtSPrBWJGWaaWb9t2ET1#$(a#4beXp~j%+^Uw6LF@Lx_W*IY*xjcGo zami6EL_!qlj6IiixNoTbUHU`&cGGkMQQ4f46X}u34lQKTNJ9s-)4J@he~R#t`U{7J zs10((F?~6Q3N)OgVU{RplKpe(bi;lI9;y|iCfgMpA1z7X@vwxBj@t{|$EpfhDnqA8 z#nflRtaE2uuH^|0oo-#DWQckqgvqG;obu0{ZF%7cDqURm8%KZZW}+jY=p#+&8M*Tl z#`oKTe7(Nwp{4w673F&k1|PhGUCzJ(|G;R3U@(WM5NIdZ6?_xC;`|wW0Nxn^e+2u0 zcgQwZ!D_NMN8pdX;ASw{bFYLgO7LSI#O~7jmBkPwLL@#6~){Dxm-_+|?4 zKQK<*F&;d!=QiJPKd^BH$a@FNy?r`Z2zX%Yox}RqYI-?usrs&6qik&5<+~6gJ)4ky`~V1dctFEvA^ay}Hr(k#pY5pNiQiG}gfp zfV*QP@hfCNbUS5&2oN87xAI6x6#3%ktm{${C1Uu#^pooPlU{(Z@P!{0o^f7P*fKFg zrDA~htHhn0){nCKox;Zei5@;ycjQr=n)yilxvkJ0o|11tbl8#~S>k@Ha9 zb4ZpTCeNlq){uEhr55!&pA=QPG>PDqhm*hvnv-I#YoA|nUc!CdIvV(kKbLf|vl~*~ zT>e5uAssnXQ%WgL+*6BCHzAE0lFZwyENWC|FaS*xrq|zvtwR+sRjQ~F8}=!IyX&Pf z3`xZ&M+cm=-m)N#o)PaeEUCyI z##0XU-2s>JW0G}D8!4MlD-%~TffzkkZfWszsg&?4yIt$y+eGnQW_sFVmZapdtsn86 z<^buioQ@>u_90CS8XmVc|9srMp~1Z{jZj>s%Qe~N*efN&+cn*Jr_zf1BIh=2RrA$_;p2{pUVc2lm3GhpHy!JP3I;z!EC zu&ci-okw)J3VyZnE}w7`chs_@kKIv|;z84Dtkf#Gel!PZxsZ>%LCG(K6zf)vWX$}; zcKH3qRqWrol9$&ZYYaGIXLj47unIF=kio;1=ECqLNDYzNv@}mPN!%LWq#>u(&Q)OQ zSfgpHh4*DyaeOm~i<#0K#;DCJDI7t(MtB#vmBj{ovbRjxau`VETSs-RFd3IUY4V1cNMvf7#MoDNh zG#lK~TD8g-mvb03B3!X4Gv$*;1}YedaL6UDFkGS~!J2%gj0iP;*$S{r&#zMncC@S8 zay@Z3A~=_ctW7%pgY=TL#XvgC^%ORaKS?-~|2ZB6(-R7KVpQdo&;Q3^xLbcRVaR~I z*297Lp^C~N7&l6YrZHq?o3zReQxI$5A#7V|EmdF+Yf?tS0i8Mei$cR?ww7AR9IB(V z9B|KwH>af>CCTny;8mZ1%j^feuzXjSIP8AM&ueYtg{#+GwD3ccJ2Fq=I?Ej?K(%`C zz#8)7Tg6tgl}9;u4&ZrRJKWt5B_%A#0EtbyGO66BX0OzUH-=N5)pdg$jmOG)lYK5T zQr(6~kYXrQZEo&)qMvav1)6GEZ7WB$BB##KO8lRJaWWO5<6 ze_o<_Y7u$O!D24h{n62t9cDf~y?Fawo&Kw+zUj(pt;nxz(rzE^BdCM#bU^5pMO(F& z8}~Z-(W0(cY+5a7m;7X65(W!WMNfZ{60*R_(t}*DA_m^ zp+D85#x)ax1xbgf*3u7#ZFm};Z&ds$0cnL}qi4RN{U(KMt#?rwOAt|>9x8#L+78Xa zxfGG_<|MxLhy^B+@IY~>%C=Knck$3@R2UhnfV+}Sd9bZgpy0CIrbQ6-cRc3ewJx%{ zs&BZ1H4@i`1A{tVEteH0YAL6d3{($R2A88BV{okv1qF1d0xwM5_e)2l+Y$j9cy34( znH-nYu%mJ66PRy&f0Le=jc^qx6kb447uqkiXWz43wPAX*ys{RN&LYNc{^epyk0;FT zv1<=DIENG#yGs*c4c#!gB140SFy;t1o$v;YJ3GUmpsFDdc*~z6&p+5gJZOoU1oFx; zs}YVU;dEgnLiz~J zt$!lIf{ad+&rV@Ov_vdS*C_=QjovhWB7uYWoACl#8Q`eR!sF6{w6Q?B?CXhU^ zI&R?jG?^(JWGQy?h$h}nxTQ@T)Wk`l zt21Dyx?;dPGnG%oX&I_3?r{{ba_G+ft+w-=alg9DToLi z24_sTwm~^`F~t#LjnE}tn9#Aq+u>w;bw%p2C}u{$MxLG`;`;0YTkyG`4WzG_c=&~# zs_4)^>C@8oRtstugsM*ZVTCd*9GvO+HQRVtdU#l7@Noy&69NoAxvN9z0e8-UlFH1~ zkfjKoVI;5xE2fA%cWVy_KDbLOZW%zmClLMqe8VUgaShZJZF!u=bS1~YRBUnAG5}!) zS}$b0n!SXBJuB-&8!}9TWm96Ih#`wY=V$+U6}Rx|R&wWVx{aYRUQ5|UGV1smr1lf; z&#GPe^{Px6MoQ@f$)o1YgzMGN-`-c8jZWXzl-z1$g-fy|66K%gmzGUEJEe5YPp8UZg)v^KEReI` zuMivwWUoTL@tCD2M%9Q%CiY$nC%+de8U}m(6^G!xZb1N!gQU(ihnlk+HmTiwC*RS( z*MGz5vuRsFPKKIFIYlT%6zPV`(9mFl+XE8u?Nb829k}61A_H|@#FHKwv=ja0h#^8% z*Ip1~PU0mG#xn&13kJ`rqqKw>le(`ZWaHx?J`*s^j~hbqLO(t0n8o_klO>wf_a zK6!^VFvy1)BVY}nTv4c_0(=ugm_9WYKAVQnSplY7#YW;e?E=@zCr53KDeLTKm3gCb zhI$1E3cd2a%Z&0GHemqdx!ZX5sm|pzO!s0$;peGCUnJ_r2BzUdC`@TX7@Byq{61={ji>kXQz*9;zcS05l~tfu`&YdD=+(eVUVnJ5@89ll zhW>1kx0)hrf!IaC-xY4f(;-3&Rj|TvCd4ZWE|9dM{QmFP;RZGnFD(&nhT)IvP`@%` zK?zSqtP?m!pGI#eJVX`>n`;PIIJ|1VwsQ^8iLISihC4+(G!y3}VNBJ!QdFAwk>h8D z5aRs=_4QK+^tVqoxCi#52onOBRVMgW(K6d?A>se{TyZf8lpX|@kCiQ;qE6}S-akMhhV z>OS1<=!b+s7lhUI|JJfxLpSqTXM)AYW_n`@&+OLecUYct-W%~`5mq)sDIRXDvAOlO zC}XuRUYo`rP=^l?0xjcm>ABKX8(JG7(eV#Uv;iTO8MSGuSawus>KyRJ)uXk8491c2 zaqCpKzHP$$j2R}b&EJW{JDx2u@o7xXQd+w!mMM`XLqbV}Y?HHEN7(P0GnZa!2O*p+ z^LfuIyDQ0X9U0(kHVi&}5w%ED@SMr=j}cU4>Xdv*Un}>l^-Sn#s^(U1jod0Ov_-kq zZS5JqWXzC|y#N8TMS1V&8O@K{s-9)+=&KYaILlJglgJlw5dn^y0|*X|c1Ox~P?oEJ z0)Ah3$B%R5IK5}cb6fOb3RcG{bblLi@R)fQUpa2zoOd3#)Wdq)T}A}(vh-?vD{qL_ zWxPkt*?4WfBhl-7$+c*D-ycOU7*lCh2E3jZ2Aq$X$f{=Q67jchD`!F4>k7v)%z1I? zkFG=XK?%|s>fMuUlu9n$dEK6deq>^zRXu>329z?p5Ru6kD81Ka%>IDzr(>on!upQY z^g9QxCY#Vm+#hvb4ZKeAEpRn6ADIC^h)-x*Ok|FvqCZ`J{7Kcs*fjaR%qD1UEFp2= zQ`9lXv0Ws!j7;2OlZqzPF8No9hw!E)R*tB(kXH&WsZ%)>A~@HFkNlHYV`J#u@ib`3 z+a&4B3<}dS%xk(opl$_@o)1-~UG@lK@={RCditJ|q^oRMXUB<`0b~5B@KpU@I@TZJ z2?c1cL9=T?$D$iJ+;VcBw^lda04;VaY*m8JRMHCiw1i_7Bu z)?y1gjiZmC{cQZ5hj!zsqdOMsj(*ksbEpY&G@>moCvS<#OO$B_hE{2Flrkgs@0Xsv z6!g?T7o#lERkMOxpnVgLmXY{GS1r*_@81$p~i0vcK$#4m}kEZ-)y}N=-U=)FnH~|+dbHB z3mJP0?*D@u6#Tcqui$#{1^`F$NjS1|03Nskg9mD3`snpVcZtC}+ig$n&Q^P(7rH3T zXswWO$0#?7OQjRtaS$p%T;xadK6d#O4_K!3tkO$3+f}nQD6bev!MEz+4}_{X;0Zd9!)z zvS{J;`PzYJiM{9^fH5otWHxhqDz2=@G~tRE)7M*lT1AH)vS@dImuR6Q0J^@T{)g5l zhDkzG9ho4d_;?P8jxOx zkKi{$waY$k963;h+62s`Lm9?tyT0SDp7zL21*{GA4m#xo;2}YhaM<78fVc^XXXo_> zeiNs0WUCN95%6-axIPfjC~>+4u-ps!hX!}NMs@LeqWm$7vaU7LgRARn0-9^g-Aw5x z8%{(0wszq~?zAUkM^69mA+eEg7DZ;}a^^<;4>2_QH)4Ov!d%NX>#TAGl<#?bF?&?o zO|l@(Z~ryfv(1P5Nn_0fCMo??z$3tZRJc%0$y*{G5x2%+h!Z>b@elS(vKn_R(A3S1 zx#Kf#rlf{`BRNFEVEL(&yd?7Rmxe#0mRh)-`{kKR?l`psgvQvjM$qKBp6nEEtAVDR zi@Kb%`uKKnVv7WcZ@Az#ocAZ{ISkQXME$-Kg$mbI?eR#$ABAXjJIpR`8PXf8!?OaQ zx5oA~9NgrOIwGMCOS6ws+>>j|kRPXTE=&}{754-x)(e$}k=I4nAd-p0A3Mfl;3ln4 zQ-w-8l266Xo*TuMUJuZY5jgi>UzxAd0I2Y>*XIjOLuS_SzZaj{h4rW)T8}}ea@2qp# zA?P%g8&lSL@SW0j zZRO=<2Ak}U7*D)FPwKtV{F3^YN|Z)E?x_)YB)7lWP*rS3;^|N25&@_U#;@lb`b~TG zT@D1d&@CXo6J6yqzuFJ?kA_aw-r^-?sCSElEnxiexAvGpu8Zd+O@-G8v{1F)9L?6R zwyq)|l9RpQVn)Ah^dj_;vyhhh;eJ1z@dr~fqwS#uk%N0-!CvCbyS*#O4HMhnmwS3G zMnAAiHL^PKc;o(bR)k>Ztpm*h=?CkmMpS>x0%H55C;mA{H(6l}u>;0DcdqvbjOEv-fSO^+i&gP>=76)v7fT?PNLz@hq-urYfq zx^zG^rpc|84hJY+E}+F(GwHH~k5(GW=y0>r6?|(oW9?iWS;V@9oZlGXep7B%Y9t3Fz#a>ko}bhE1RoheFBLn3+veMmC8*k%TX3 zWZNSpvchTV>}kswGlgmaCBrjMN|Bwe72Z2W3#SW+vn^H;0z zA5>Bs90g3>a-W@RI3$p#5zi@H;?@mZJC30XHfkxkz5@;-%Ufw3lo!&UFpo!BVG+{DsAX$GYQo}h?g zyao1(*4{=)ec({NW!U2VROC=FZY?`J*RYc}_s(sBj==&cdZ{V?69tV8Z9m&Ns>)Ys zrJV-D(oG2N%yri^oo+R~D#^r%EsC95*(SCauQxNQdKx*SmM}AZcS!kQu8xdAKi0hI zlu7!y(^hDpQrJ6C6uKYnba7tof8@IE}K!aB2)A@fJ22u1yDzf; z^McL@vgGGsHn-QPBe^Q~60$$4x2kPKlrtlLZj+WSy0)47PT}eIDD&YQicxMVbL0Aw zNseNyWS~CquTyO-vf^2cNYrX$eC&{3Q27e~67GwYs(qEBlD0why?@n|Dhd^2rTrI= z7=mq%^{7Cj^{7u#v^C5xN93iLD@881+Iig)>-Cbhav`{C{wc(xG#dN5rk_PQMkcD} zj$xfM6^dHTt)d+h_3>Yh=mEMEDExo1WX*}$r?tFLxnQOU(A^(@0?~W1_Bv@ZW;(Bb zdP^t!@>l*r3Gn>v^Rf6yRaI@d)1gi9Ierm)NaB`x{dLgE=~o;c)x798-%)#QHt*Y` zT087A@4NTr!{2Bzo|ww7XAX2TZDoMUyd(z-O(io7OHd(MYls>B{s~1VJ9!1WCM)k% z9}0qoR?(g4KHN4;=&s?SXK&^foV*{yMSw~>JE31$+iDF8 z#^xg1s=Rg#knV8L$W@0Ca9`#u%Aky>3Bxo`cr2{BfYqh zd<`eNBI?{T41N(AH{Z{3|C5|jSD;&6HtvEb!HwOwco+cmsNi^ps9$z9*LC_KW7c8jrV4`^XPIk>s{0w*5TP za1^jM#hG!UuVePnf7BB&)amr+d@&v@$`7jW5^lf+Ur~T6SPVbs!qdH9i;B&qf3655Z$;BSFJ52j|7GNQ_^i*YsCqK>&q z!P~#?U=id58c!lk8539b)FrkU{6+^(oSHKUhWWTa(acIovSJBD*Tz;@0hw#(oux|V zGS0wU$6+9xiCZ(5md50#H>5VUPj+iqT3<2H1=W=zkvJv6JSX`5tc+Ycn4tDlv*=%Y YYdMRvy^mxVBEd<2>0OxI7mjD0& literal 0 HcmV?d00001 From 9ce1cbc96e6585244afa603cbb056c99dd7be3fd Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Mon, 29 Jun 2026 00:41:32 +0000 Subject: [PATCH 06/54] =?UTF-8?q?feat(medical):=20per-specialist=20proxies?= =?UTF-8?q?=20=E2=86=92=20separate=20per-agent=20trajectories=20+=20native?= =?UTF-8?q?-vs-hosted=20runner?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route each Medical-Assistant node (supervisor/answer/guardrail) through its own LiteLLM proxy so each agent yields a separate llm_trajectory + usage, instead of one shared proxy emitting a single combined log. Adds run_native_vs_hosted.py to run the bench direct-vs-hosted and diff observability. --- examples/medical/medical_assistant.py | 30 +++-- examples/medical/run_native_vs_hosted.py | 138 +++++++++++++++++++++++ 2 files changed, 160 insertions(+), 8 deletions(-) create mode 100644 examples/medical/run_native_vs_hosted.py diff --git a/examples/medical/medical_assistant.py b/examples/medical/medical_assistant.py index 7f4c5901..95e35d1e 100644 --- a/examples/medical/medical_assistant.py +++ b/examples/medical/medical_assistant.py @@ -42,11 +42,25 @@ class S(TypedDict, total=False): trace: list -def _llm() -> ChatOpenAI: - base = os.environ.get("BENCHFLOW_PROVIDER_BASE_URL") or "https://api.deepseek.com/v1" - key = (os.environ.get("BENCHFLOW_PROVIDER_API_KEY") - or os.environ.get("DEEPSEEK_API_KEY") or "EMPTY") - model = os.environ.get("BENCHFLOW_PROVIDER_MODEL") or "deepseek/deepseek-v4-pro" +# agent_name -> (base_url, api_key, model). A runner sets one entry per +# specialist to route that agent through its OWN proxy, so each gets a separate +# llm_trajectory + usage. Empty => every node shares one endpoint (shared/native). +_AGENT_PROVIDERS: dict[str, tuple[str, str, str]] = {} + + +def set_agent_provider(agent: str, base: str, key: str, model: str) -> None: + _AGENT_PROVIDERS[agent] = (base, key, model) + + +def _llm(agent: str = "default") -> ChatOpenAI: + cfg = _AGENT_PROVIDERS.get(agent) + if cfg is not None: + base, key, model = cfg + else: + base = os.environ.get("BENCHFLOW_PROVIDER_BASE_URL") or "https://api.deepseek.com/v1" + key = (os.environ.get("BENCHFLOW_PROVIDER_API_KEY") + or os.environ.get("DEEPSEEK_API_KEY") or "EMPTY") + model = os.environ.get("BENCHFLOW_PROVIDER_MODEL") or "deepseek/deepseek-v4-pro" return ChatOpenAI(model=model, base_url=base, api_key=key, temperature=0.2, max_tokens=2048) @@ -56,7 +70,7 @@ def _log(state: S, node: str, note: str) -> None: def supervisor(state: S) -> S: - out = _llm().invoke([{"role": "user", "content": + out = _llm("supervisor").invoke([{"role": "user", "content": f"You route a medical question to a specialist. Question: {state['query']!r}. " "Reply with ONE word: 'kb' if a drug/condition knowledge base likely covers " "it, else 'web'."}]).content.lower() @@ -80,7 +94,7 @@ def web_search(state: S) -> S: def answer(state: S) -> S: - out = _llm().invoke([{"role": "user", "content": + out = _llm("answer").invoke([{"role": "user", "content": f"Question: {state['query']}\nEvidence: {state.get('evidence') or '(none)'}\n" "Answer concisely for a clinician. Then on a NEW line output " "'CONFIDENCE: <0.0-1.0>' for how well the evidence supports your answer."}] @@ -93,7 +107,7 @@ def answer(state: S) -> S: def guardrail(state: S) -> S: - out = _llm().invoke([{"role": "user", "content": + out = _llm("guardrail").invoke([{"role": "user", "content": f"Output guardrail. The assistant answered: {state['answer']!r}. Is it safe and " "appropriately cautious as medical information (not a diagnosis)? Reply yes/no."}] ).content.lower() diff --git a/examples/medical/run_native_vs_hosted.py b/examples/medical/run_native_vs_hosted.py new file mode 100644 index 00000000..81d75158 --- /dev/null +++ b/examples/medical/run_native_vs_hosted.py @@ -0,0 +1,138 @@ +"""Run the Medical-Assistant bench TWO ways and diff them. + + (a) NATIVE — every node calls deepseek DIRECTLY. No proxy, no BenchFlow + tracking: you get an answer, but no usage/cost and no llm_trajectory. + (b) HOSTED — every specialist (supervisor / answer / guardrail) gets its OWN + LiteLLM proxy, so each agent produces a SEPARATE llm_trajectory.jsonl + + its own usage/cost. This is the fix for "the multi-agent run only emits a + single trajectory": one shared proxy = one mixed log; one proxy per agent + = one log per agent. + + set -a; . ./sb-run.env; set +a + uv run python examples/medical/run_native_vs_hosted.py "side effects of metformin?" + # threshold 1.01 (default below) forces the confidence-gated web handoff so + # the 'answer' specialist fires twice — visible as 2 exchanges in its log. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys +from pathlib import Path + +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import medical_assistant as ma + +LLM_AGENTS = ("supervisor", "answer", "guardrail") # the nodes that call an LLM + + +def _summary(result: dict) -> dict: + return { + "path": " → ".join(t["node"] for t in result.get("trace", [])), + "route": result.get("route"), + "confidence": result.get("confidence"), + "safe": result.get("safe"), + "answer": (result.get("answer") or "").replace("\n", " ")[:200], + } + + +async def native_run(query: str) -> dict: + """Direct provider — no proxy. _AGENT_PROVIDERS empty → _llm() reads env.""" + ma._AGENT_PROVIDERS.clear() + os.environ["BENCHFLOW_PROVIDER_BASE_URL"] = "https://api.deepseek.com" + os.environ["BENCHFLOW_PROVIDER_API_KEY"] = os.environ["DEEPSEEK_API_KEY"] + os.environ["BENCHFLOW_PROVIDER_MODEL"] = "deepseek-v4-pro" # no prefix for direct + result = await asyncio.to_thread(ma.run, query) + return _summary(result) + + +async def hosted_run(query: str) -> tuple[dict, dict]: + """One proxy per specialist → a separate trajectory + usage per agent.""" + ma._AGENT_PROVIDERS.clear() + runtimes: dict = {} + print(f"starting {len(LLM_AGENTS)} per-agent LiteLLM proxies…", flush=True) + for ag in LLM_AGENTS: + env, rt = await ensure_litellm_runtime( + agent="deepagents", + agent_env={"DEEPSEEK_API_KEY": os.environ["DEEPSEEK_API_KEY"]}, + model="deepseek/deepseek-v4-pro", runtime=None, environment="local", + session_id=f"medical-{ag}", + ) + runtimes[ag] = rt + ma.set_agent_provider(ag, env["BENCHFLOW_PROVIDER_BASE_URL"], + env["BENCHFLOW_PROVIDER_API_KEY"], + env["BENCHFLOW_PROVIDER_MODEL"]) + print(f" {ag:<10}: {env['BENCHFLOW_PROVIDER_BASE_URL']}", flush=True) + + run_dir = Path("out/medical-hosted") + per_agent: dict = {} + try: + result = await asyncio.to_thread(ma.run, query) + await asyncio.sleep(1.5) # let each proxy's async callback flush + finally: + for ag in LLM_AGENTS: + await stop_provider_runtime(runtimes[ag]) # stop FIRST → parses callback log + for ag in LLM_AGENTS: # then read usage + trajectory (populated on stop) + usage = extract_usage(runtimes[ag]) + traj = getattr(getattr(runtimes[ag], "server", None), "trajectory", None) + n = len(traj.exchanges) if traj is not None else 0 + path = None + if traj is not None and traj.exchanges: + d = run_dir / ag / "trajectory" + d.mkdir(parents=True, exist_ok=True) + path = d / "llm_trajectory.jsonl" + path.write_text(traj.to_jsonl(redact_keys=True)) + per_agent[ag] = {"calls": n, "usage": usage, "trajectory": str(path) if path else None} + return _summary(result), per_agent + + +async def _main() -> None: + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required (the real upstream key)") + query = sys.argv[1] if len(sys.argv) > 1 else \ + "What are the main side effects of metformin?" + os.environ.setdefault("MEDICAL_CONFIDENCE_THRESHOLD", "1.01") # force web handoff + + print("=" * 72) + print("NATIVE (direct deepseek — no BenchFlow)") + print("=" * 72) + native = await native_run(query) + + print("\n" + "=" * 72) + print("HOSTED (one BenchFlow proxy PER specialist)") + print("=" * 72) + hosted, per_agent = await hosted_run(query) + + print("\n" + "=" * 72) + print(f"COMPARISON · query: {query!r}") + print("=" * 72) + print(f"{'':<14}{'NATIVE':<40}{'HOSTED':<40}") + for k in ("path", "route", "confidence", "safe"): + print(f"{k:<14}{str(native[k]):<40}{str(hosted[k]):<40}") + print(f"{'answer':<14}{native['answer'][:36]:<40}{hosted['answer'][:36]:<40}") + tot = sum((per_agent[a]['usage'] or {}).get('total_tokens', 0) or 0 for a in LLM_AGENTS) + cost = sum((per_agent[a]['usage'] or {}).get('cost_usd', 0) or 0 for a in LLM_AGENTS) + print(f"{'tracking':<14}{'none (no proxy → no usage/traj)':<40}" + f"{f'{tot} tok · ${cost:.5f} · per-agent':<40}") + + print("\n--- HOSTED: SEPARATE trajectory per agent ---") + for ag in LLM_AGENTS: + pa = per_agent[ag] + u = pa["usage"] or {} + print(f" {ag:<10} {pa['calls']} call(s) · " + f"tok={u.get('total_tokens')} cost=${u.get('cost_usd')}") + print(f" {pa['trajectory']}") + print("\nNATIVE produced NO trajectory and NO usage — that observability only " + "exists when the LLM is routed through BenchFlow.") + + +if __name__ == "__main__": + asyncio.run(_main()) From a02d2f55238b09c2a89276e50ed12f1378addec7 Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Mon, 29 Jun 2026 01:35:44 +0000 Subject: [PATCH 07/54] feat(medical): host Multi-Agent-Medical-Assistant as a BenchFlow benchmark MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Host the souvikmajumder26/Multi-Agent-Medical-Assistant pattern (LangGraph supervisor→specialists: router → KB/web specialists → confidence-gated handoff → guardrail) as a first-class BenchFlow benchmark that runs in a Docker sandbox with deepseek-v4-pro routed through the LiteLLM proxy. - medical-assistant ACP-shim agent (src/benchflow/agents/medical_acp_shim.py), registered in agents/registry.py; each graph node streamed as its own ACP tool_call step so the multi-agent structure is visible in the trajectory. - benchmarks/medical-assistant/: benchmark.yaml, 3 Harbor drug-safety tasks with deterministic keyword verifiers, job config, run script, README. - Verified: bench run in Docker → 3/3, errors=0; per-rollout llm_trajectory (proxy raw-LLM) + acp_trajectory (per-node steps) captured. The upstream's Azure-embeddings/Qdrant/CV/Tavily stack is out of scope on the deepseek-only path; the router→specialists→handoff→guardrail control flow is hosted faithfully with an in-process KB standing in for RAG. --- benchmarks/medical-assistant/README.md | 62 +++ benchmarks/medical-assistant/benchmark.yaml | 57 +++ .../medical-assistant-deepseek.yaml | 12 + .../run_medical_assistant.py | 55 +++ .../environment/Dockerfile | 12 + .../instruction.md | 3 + .../aspirin-secondary-prevention/task.toml | 18 + .../tests/ground_truth.json | 15 + .../tests/test.sh | 23 ++ .../environment/Dockerfile | 12 + .../ibuprofen-renal-caution/instruction.md | 3 + .../tasks/ibuprofen-renal-caution/task.toml | 18 + .../tests/ground_truth.json | 7 + .../ibuprofen-renal-caution/tests/test.sh | 23 ++ .../environment/Dockerfile | 12 + .../metformin-side-effects/instruction.md | 3 + .../tasks/metformin-side-effects/task.toml | 18 + .../tests/ground_truth.json | 16 + .../metformin-side-effects/tests/test.sh | 23 ++ src/benchflow/agents/medical_acp_shim.py | 373 ++++++++++++++++++ src/benchflow/agents/registry.py | 43 ++ 21 files changed, 808 insertions(+) create mode 100644 benchmarks/medical-assistant/README.md create mode 100644 benchmarks/medical-assistant/benchmark.yaml create mode 100644 benchmarks/medical-assistant/medical-assistant-deepseek.yaml create mode 100644 benchmarks/medical-assistant/run_medical_assistant.py create mode 100644 benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/environment/Dockerfile create mode 100644 benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/instruction.md create mode 100644 benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/task.toml create mode 100644 benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/ground_truth.json create mode 100755 benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/test.sh create mode 100644 benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/environment/Dockerfile create mode 100644 benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/instruction.md create mode 100644 benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/task.toml create mode 100644 benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/ground_truth.json create mode 100755 benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/test.sh create mode 100644 benchmarks/medical-assistant/tasks/metformin-side-effects/environment/Dockerfile create mode 100644 benchmarks/medical-assistant/tasks/metformin-side-effects/instruction.md create mode 100644 benchmarks/medical-assistant/tasks/metformin-side-effects/task.toml create mode 100644 benchmarks/medical-assistant/tasks/metformin-side-effects/tests/ground_truth.json create mode 100755 benchmarks/medical-assistant/tasks/metformin-side-effects/tests/test.sh create mode 100644 src/benchflow/agents/medical_acp_shim.py diff --git a/benchmarks/medical-assistant/README.md b/benchmarks/medical-assistant/README.md new file mode 100644 index 00000000..cf61fb9a --- /dev/null +++ b/benchmarks/medical-assistant/README.md @@ -0,0 +1,62 @@ +# Multi-Agent Medical Assistant — BenchFlow Adapter + +BenchFlow adapter that **hosts the [Multi-Agent-Medical-Assistant](https://github.com/souvikmajumder26/Multi-Agent-Medical-Assistant)** +agent pattern as a first-class BenchFlow benchmark: it runs in a Docker sandbox +with its LLM routed through BenchFlow's LiteLLM proxy, so per-rollout usage/cost +and the raw-LLM trajectory are captured, and the agent never sees the raw +provider key. + +## What is hosted + +The upstream is a LangGraph **supervisor → specialists** agent. This adapter +reproduces its control flow as a registered BenchFlow agent (`medical-assistant`, +an ACP shim at `src/benchflow/agents/medical_acp_shim.py`): + +``` +supervisor (router) + ├── retrieve_kb (RAG-style knowledge-base specialist) + └── web_search (fallback specialist) + → answer (emits a CONFIDENCE; low confidence → confidence-gated + handoff back to web_search) + → guardrail (output safety check) +``` + +Each graph node is streamed back as its own ACP `tool_call` step, so the +multi-agent structure (which specialist ran, in what order) is visible in the +captured trajectory. + +### Scope vs. the full upstream app + +The upstream app's heavy stack — **Azure OpenAI embeddings, Qdrant RAG, torch CV +imaging weights, Tavily web search** — is **out of scope** on the deepseek-only +proxy path. The RAG corpus is replaced by a small in-process knowledge base so +the full *router → specialists → confidence handoff → guardrail* control flow runs +end-to-end. The CV/imaging specialists and live web search are not included. See +`benchmark.yaml` (`hosting.not_included`). + +## Tasks & verification + +Three clinical drug-safety questions (`tasks/`), each verified **deterministically**: +the agent writes its answer to `/app/answer.md`, and `tests/test.sh` scores it +against the task's `ground_truth.json` keyword groups (OR within a group, AND +across groups). Reward = matched groups / total groups → `/logs/verifier/reward.txt`. + +| Task | Question | +|------|----------| +| `metformin-side-effects` | Main side effects of metformin | +| `aspirin-secondary-prevention` | When low-dose aspirin is used, and its main risk | +| `ibuprofen-renal-caution` | Cautions for ibuprofen use | + +## Run + +```bash +set -a; . ~/sb-run.env; set +a # provides DEEPSEEK_API_KEY for the proxy +python benchmarks/medical-assistant/run_medical_assistant.py + +# or via the CLI: +bench eval run --config benchmarks/medical-assistant/medical-assistant-deepseek.yaml +``` + +Every rollout runs in its own Docker sandbox (`environment: docker`). The agent is +installed into the sandbox (uv venv + `langgraph` + `langchain-openai`) and its +`deepseek-v4-pro` calls go through the LiteLLM proxy. diff --git a/benchmarks/medical-assistant/benchmark.yaml b/benchmarks/medical-assistant/benchmark.yaml new file mode 100644 index 00000000..14073847 --- /dev/null +++ b/benchmarks/medical-assistant/benchmark.yaml @@ -0,0 +1,57 @@ +# benchmark.yaml — standard benchmark descriptor for BenchFlow. +# +# Every benchmark in benchmarks// ships this file. It declares what the +# benchmark is, where it comes from, how tasks are verified, and how it runs. + +name: medical-assistant +description: "Multi-Agent Medical Assistant — a LangGraph supervisor→specialists agent (router → KB/web specialists → confidence-gated handoff → output guardrail) hosted by BenchFlow and answering clinical drug-safety questions" +url: https://github.com/souvikmajumder26/Multi-Agent-Medical-Assistant +repo: https://github.com/souvikmajumder26/Multi-Agent-Medical-Assistant +author: "souvikmajumder26 (pattern); BenchFlow (adapter)" +converted_by: BenchFlow + +# ── Tasks ──────────────────────────────────────────────────────────── +tasks: + count: 3 + categories: + - medical-qa + tags: [medical-assistant, langgraph, multi-agent, qa, drug-safety] + splits: + main: 3 + +# ── What is hosted ─────────────────────────────────────────────────── +# The agent is the Multi-Agent-Medical-Assistant PATTERN: a real langgraph +# StateGraph with a supervisor/router, a KB (RAG-style) specialist, a web-search +# specialist, a confidence-gated handoff, and an output guardrail. It runs as a +# BenchFlow ACP-shim agent (agent: medical-assistant) whose LLM is deepseek-v4-pro +# routed through the LiteLLM proxy. +hosting: + agent: medical-assistant # registered in src/benchflow/agents/registry.py + shim: src/benchflow/agents/medical_acp_shim.py + graph_nodes: [supervisor, retrieve_kb, web_search, answer, guardrail] + # The upstream's heavy stack (Azure embeddings, Qdrant, torch CV weights, + # Tavily) is OUT OF SCOPE on the deepseek-only proxy path — the RAG corpus is + # replaced by a small in-process knowledge base so the full router→specialists→ + # handoff→guardrail control flow runs end-to-end. The CV/imaging specialists and + # live web search are not included. + faithful_to_upstream: control-flow # router + specialists + confidence handoff + guardrail + not_included: [azure-embeddings, qdrant-rag, cv-imaging-weights, tavily-web-search] + +# ── Verification ───────────────────────────────────────────────────── +verification: + method: deterministic_script + reward: proportional # matched_keyword_groups / total_groups + aggregation: per-task + detail: > + Each task ships ground_truth.json with keyword groups (OR within a group, + AND across groups). The agent writes its answer to /app/answer.md; tests/test.sh + scores it and writes reward to /logs/verifier/reward.txt. No LLM judge. + anti_cheat: false + +# ── Environment ────────────────────────────────────────────────────── +environment: + base_image: "python:3.12-slim" + platform: linux/amd64 + allow_internet: true # agent install (PyPI/uv) + proxy reachability + cpus: 1 + memory_mb: 2048 diff --git a/benchmarks/medical-assistant/medical-assistant-deepseek.yaml b/benchmarks/medical-assistant/medical-assistant-deepseek.yaml new file mode 100644 index 00000000..41119915 --- /dev/null +++ b/benchmarks/medical-assistant/medical-assistant-deepseek.yaml @@ -0,0 +1,12 @@ +# Job config — run the medical-assistant benchmark with deepseek-v4-pro, +# every node's LLM routed through BenchFlow's LiteLLM proxy, in a Docker sandbox. +# +# bench eval run --config benchmarks/medical-assistant/medical-assistant-deepseek.yaml +# or: +# python benchmarks/medical-assistant/run_medical_assistant.py + +tasks_dir: benchmarks/medical-assistant/tasks +agent: medical-assistant +model: deepseek/deepseek-v4-pro +environment: docker # each rollout runs in its own Docker sandbox +concurrency: 2 diff --git a/benchmarks/medical-assistant/run_medical_assistant.py b/benchmarks/medical-assistant/run_medical_assistant.py new file mode 100644 index 00000000..0552c569 --- /dev/null +++ b/benchmarks/medical-assistant/run_medical_assistant.py @@ -0,0 +1,55 @@ +"""Run the medical-assistant benchmark via BenchFlow. + + set -a; . ~/sb-run.env; set +a # provides DEEPSEEK_API_KEY for the proxy + python benchmarks/medical-assistant/run_medical_assistant.py + + # or, equivalently, via the CLI: + bench eval run --config benchmarks/medical-assistant/medical-assistant-deepseek.yaml + +Every rollout runs in its own Docker sandbox (environment: docker). The +medical-assistant agent (a LangGraph supervisor->specialists graph) is installed +into the sandbox and its deepseek-v4-pro calls are routed through BenchFlow's +LiteLLM proxy, so usage/cost + the raw-LLM trajectory are captured per rollout. +""" + +import argparse +import asyncio +import logging +import sys +from pathlib import Path + +_SCRIPT_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SCRIPT_DIR.parents[1] +_SRC_ROOT = _REPO_ROOT / "src" +if str(_SRC_ROOT) not in sys.path: + sys.path.insert(0, str(_SRC_ROOT)) + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the medical-assistant benchmark.") + parser.add_argument( + "config", + nargs="?", + default=str(_SCRIPT_DIR / "medical-assistant-deepseek.yaml"), + help="BenchFlow evaluation YAML config.", + ) + return parser.parse_args() + + +async def main() -> None: + from benchflow.evaluation import Evaluation + + args = _parse_args() + job = Evaluation.from_yaml(args.config) + job._tasks_dir = _SCRIPT_DIR / "tasks" # absolute path → cwd-independent + # Dedicated jobs_dir so we never resume into another agent's results. + job._jobs_dir = _REPO_ROOT / "out" / "medical-bench" / "jobs" + job._job_name = "medical-assistant-deepseek" + result = await job.run() + print(f"\nScore: {result.passed}/{result.total} ({result.score:.1%})") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/environment/Dockerfile b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/environment/Dockerfile new file mode 100644 index 00000000..98916b10 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/environment/Dockerfile @@ -0,0 +1,12 @@ +# medical-assistant Q&A task — minimal sandbox. +# +# The medical-assistant agent (a LangGraph supervisor -> specialists graph) is +# installed at runtime by BenchFlow's registry install_cmd (uv venv + langgraph + +# langchain-openai), so this image only needs python3 (for the verifier), an +# `agent` user (the default sandbox_user), and writable /app + /logs dirs. The +# agent writes its answer to /app/answer.md, which the verifier scores. +FROM python:3.12-slim +RUN useradd -m -u 1000 agent || true +RUN mkdir -p /app /logs/verifier /logs/agent /logs/artifacts \ + && chown -R agent:agent /app /logs +WORKDIR /app diff --git a/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/instruction.md b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/instruction.md new file mode 100644 index 00000000..584c5849 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/instruction.md @@ -0,0 +1,3 @@ +You are a clinical decision-support assistant. Answer the following question accurately and concisely for a clinician, noting the key risks. + +Question: When is low-dose aspirin used, and what is its main risk? diff --git a/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/task.toml b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/task.toml new file mode 100644 index 00000000..9c9006d6 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/task.toml @@ -0,0 +1,18 @@ +version = "1.0" + +[metadata] +author_name = "benchflow" +difficulty = "easy" +category = "medical-qa" +tags = ["medical-assistant", "langgraph", "multi-agent", "qa"] + +[agent] +timeout_sec = 600 + +[verifier] +timeout_sec = 120 + +[environment] +cpus = 1 +memory_mb = 2048 +allow_internet = true diff --git a/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/ground_truth.json b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/ground_truth.json new file mode 100644 index 00000000..4435344e --- /dev/null +++ b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/ground_truth.json @@ -0,0 +1,15 @@ +{ + "question": "When is low-dose aspirin used, and what is its main risk?", + "keyword_groups": [ + [ + "cardiovascular", + "secondary prevention", + "antiplatelet", + "heart", + "stroke" + ], + [ + "bleed" + ] + ] +} diff --git a/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/test.sh b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/test.sh new file mode 100755 index 00000000..24a577b5 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/aspirin-secondary-prevention/tests/test.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Deterministic verifier: score /app/answer.md against this task's keyword groups +# (OR within a group, AND across groups). reward = matched_groups / total_groups. +set -uo pipefail +LOGS_DIR="${LOGS_DIR:-/logs/verifier}" +mkdir -p "$LOGS_DIR" +ANSWER="${ANSWER_FILE:-/app/answer.md}" +GT="$(dirname "$0")/ground_truth.json" +python3 - "$ANSWER" "$GT" "$LOGS_DIR/reward.txt" <<'PYEOF' +import json, os, sys +answer_path, gt_path, reward_path = sys.argv[1], sys.argv[2], sys.argv[3] +text = "" +if os.path.exists(answer_path): + text = open(answer_path, encoding="utf-8", errors="ignore").read().lower() +groups = json.load(open(gt_path))["keyword_groups"] +hit = sum(1 for grp in groups if any(kw.lower() in text for kw in grp)) +reward = hit / len(groups) if groups else 0.0 +open(reward_path, "w").write(f"{reward:.4f}\n") +print(f"matched {hit}/{len(groups)} keyword groups -> reward {reward:.4f}") +if not text: + print("WARN: /app/answer.md missing or empty") +PYEOF +cat "$LOGS_DIR/reward.txt" diff --git a/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/environment/Dockerfile b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/environment/Dockerfile new file mode 100644 index 00000000..98916b10 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/environment/Dockerfile @@ -0,0 +1,12 @@ +# medical-assistant Q&A task — minimal sandbox. +# +# The medical-assistant agent (a LangGraph supervisor -> specialists graph) is +# installed at runtime by BenchFlow's registry install_cmd (uv venv + langgraph + +# langchain-openai), so this image only needs python3 (for the verifier), an +# `agent` user (the default sandbox_user), and writable /app + /logs dirs. The +# agent writes its answer to /app/answer.md, which the verifier scores. +FROM python:3.12-slim +RUN useradd -m -u 1000 agent || true +RUN mkdir -p /app /logs/verifier /logs/agent /logs/artifacts \ + && chown -R agent:agent /app /logs +WORKDIR /app diff --git a/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/instruction.md b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/instruction.md new file mode 100644 index 00000000..92d8f3eb --- /dev/null +++ b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/instruction.md @@ -0,0 +1,3 @@ +You are a clinical decision-support assistant. Answer the following question accurately and concisely for a clinician, noting the key risks. + +Question: What cautions apply to ibuprofen use? diff --git a/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/task.toml b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/task.toml new file mode 100644 index 00000000..9c9006d6 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/task.toml @@ -0,0 +1,18 @@ +version = "1.0" + +[metadata] +author_name = "benchflow" +difficulty = "easy" +category = "medical-qa" +tags = ["medical-assistant", "langgraph", "multi-agent", "qa"] + +[agent] +timeout_sec = 600 + +[verifier] +timeout_sec = 120 + +[environment] +cpus = 1 +memory_mb = 2048 +allow_internet = true diff --git a/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/ground_truth.json b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/ground_truth.json new file mode 100644 index 00000000..eb8b1f80 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/ground_truth.json @@ -0,0 +1,7 @@ +{ + "question": "What cautions apply to ibuprofen use?", + "keyword_groups": [ + ["renal", "kidney", "nephro"], + ["anticoagulant", "bleed"] + ] +} diff --git a/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/test.sh b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/test.sh new file mode 100755 index 00000000..24a577b5 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/ibuprofen-renal-caution/tests/test.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Deterministic verifier: score /app/answer.md against this task's keyword groups +# (OR within a group, AND across groups). reward = matched_groups / total_groups. +set -uo pipefail +LOGS_DIR="${LOGS_DIR:-/logs/verifier}" +mkdir -p "$LOGS_DIR" +ANSWER="${ANSWER_FILE:-/app/answer.md}" +GT="$(dirname "$0")/ground_truth.json" +python3 - "$ANSWER" "$GT" "$LOGS_DIR/reward.txt" <<'PYEOF' +import json, os, sys +answer_path, gt_path, reward_path = sys.argv[1], sys.argv[2], sys.argv[3] +text = "" +if os.path.exists(answer_path): + text = open(answer_path, encoding="utf-8", errors="ignore").read().lower() +groups = json.load(open(gt_path))["keyword_groups"] +hit = sum(1 for grp in groups if any(kw.lower() in text for kw in grp)) +reward = hit / len(groups) if groups else 0.0 +open(reward_path, "w").write(f"{reward:.4f}\n") +print(f"matched {hit}/{len(groups)} keyword groups -> reward {reward:.4f}") +if not text: + print("WARN: /app/answer.md missing or empty") +PYEOF +cat "$LOGS_DIR/reward.txt" diff --git a/benchmarks/medical-assistant/tasks/metformin-side-effects/environment/Dockerfile b/benchmarks/medical-assistant/tasks/metformin-side-effects/environment/Dockerfile new file mode 100644 index 00000000..98916b10 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/metformin-side-effects/environment/Dockerfile @@ -0,0 +1,12 @@ +# medical-assistant Q&A task — minimal sandbox. +# +# The medical-assistant agent (a LangGraph supervisor -> specialists graph) is +# installed at runtime by BenchFlow's registry install_cmd (uv venv + langgraph + +# langchain-openai), so this image only needs python3 (for the verifier), an +# `agent` user (the default sandbox_user), and writable /app + /logs dirs. The +# agent writes its answer to /app/answer.md, which the verifier scores. +FROM python:3.12-slim +RUN useradd -m -u 1000 agent || true +RUN mkdir -p /app /logs/verifier /logs/agent /logs/artifacts \ + && chown -R agent:agent /app /logs +WORKDIR /app diff --git a/benchmarks/medical-assistant/tasks/metformin-side-effects/instruction.md b/benchmarks/medical-assistant/tasks/metformin-side-effects/instruction.md new file mode 100644 index 00000000..e5a8ec4f --- /dev/null +++ b/benchmarks/medical-assistant/tasks/metformin-side-effects/instruction.md @@ -0,0 +1,3 @@ +You are a clinical decision-support assistant. Answer the following question accurately and concisely for a clinician, noting the key risks. + +Question: What are the main side effects of metformin? diff --git a/benchmarks/medical-assistant/tasks/metformin-side-effects/task.toml b/benchmarks/medical-assistant/tasks/metformin-side-effects/task.toml new file mode 100644 index 00000000..9c9006d6 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/metformin-side-effects/task.toml @@ -0,0 +1,18 @@ +version = "1.0" + +[metadata] +author_name = "benchflow" +difficulty = "easy" +category = "medical-qa" +tags = ["medical-assistant", "langgraph", "multi-agent", "qa"] + +[agent] +timeout_sec = 600 + +[verifier] +timeout_sec = 120 + +[environment] +cpus = 1 +memory_mb = 2048 +allow_internet = true diff --git a/benchmarks/medical-assistant/tasks/metformin-side-effects/tests/ground_truth.json b/benchmarks/medical-assistant/tasks/metformin-side-effects/tests/ground_truth.json new file mode 100644 index 00000000..dcf0ca34 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/metformin-side-effects/tests/ground_truth.json @@ -0,0 +1,16 @@ +{ + "question": "What are the main side effects of metformin?", + "keyword_groups": [ + [ + "gastrointestinal", + "gi upset", + "gi ", + "nausea", + "diarrhea", + "diarrhoea" + ], + [ + "lactic acidosis" + ] + ] +} diff --git a/benchmarks/medical-assistant/tasks/metformin-side-effects/tests/test.sh b/benchmarks/medical-assistant/tasks/metformin-side-effects/tests/test.sh new file mode 100755 index 00000000..24a577b5 --- /dev/null +++ b/benchmarks/medical-assistant/tasks/metformin-side-effects/tests/test.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Deterministic verifier: score /app/answer.md against this task's keyword groups +# (OR within a group, AND across groups). reward = matched_groups / total_groups. +set -uo pipefail +LOGS_DIR="${LOGS_DIR:-/logs/verifier}" +mkdir -p "$LOGS_DIR" +ANSWER="${ANSWER_FILE:-/app/answer.md}" +GT="$(dirname "$0")/ground_truth.json" +python3 - "$ANSWER" "$GT" "$LOGS_DIR/reward.txt" <<'PYEOF' +import json, os, sys +answer_path, gt_path, reward_path = sys.argv[1], sys.argv[2], sys.argv[3] +text = "" +if os.path.exists(answer_path): + text = open(answer_path, encoding="utf-8", errors="ignore").read().lower() +groups = json.load(open(gt_path))["keyword_groups"] +hit = sum(1 for grp in groups if any(kw.lower() in text for kw in grp)) +reward = hit / len(groups) if groups else 0.0 +open(reward_path, "w").write(f"{reward:.4f}\n") +print(f"matched {hit}/{len(groups)} keyword groups -> reward {reward:.4f}") +if not text: + print("WARN: /app/answer.md missing or empty") +PYEOF +cat "$LOGS_DIR/reward.txt" diff --git a/src/benchflow/agents/medical_acp_shim.py b/src/benchflow/agents/medical_acp_shim.py new file mode 100644 index 00000000..7ed742b9 --- /dev/null +++ b/src/benchflow/agents/medical_acp_shim.py @@ -0,0 +1,373 @@ +#!/usr/bin/env python3 +"""ACP shim for the Multi-Agent Medical Assistant — runs a LangGraph +supervisor→specialists graph as an ACP server. + +Adapts the **Multi-Agent-Medical-Assistant** pattern +(https://github.com/souvikmajumder26/Multi-Agent-Medical-Assistant): a LangGraph +``StateGraph`` with a supervisor/router that dispatches to a knowledge-base +(RAG-style) or web-search specialist, an answer node that emits a confidence, a +**confidence-gated handoff** to the web specialist when confidence is low, and an +output **guardrail** node. The upstream's heavy stack (Azure embeddings, Qdrant, +torch CV weights, Tavily) is replaced with a small in-process knowledge base so +the whole multi-agent graph runs on the deepseek-only proxy path; the router → +specialists → confidence handoff → guardrail control flow is preserved. + +Architecture: + benchflow ACP client ←stdio→ this shim ←in-process→ LangGraph medical graph + ←ChatOpenAI→ BenchFlow LiteLLM proxy + +Each graph node is streamed back as its OWN ACP ``tool_call`` step, so the +multi-agent structure (which specialist ran, in what order) is visible in the +captured trajectory — even though every node's raw LLM call rides the single +per-rollout proxy log. The final answer is written to ``/answer.md`` for the +task verifier. + +Self-contained: like the other shims (deepagents, openclaw, harvey-lab), this +file is uploaded and run inside the sandbox with NO benchflow installed, so it +must NOT ``import benchflow``. It depends only on the stdlib plus langgraph / +langchain-openai installed into the shim's venv. +""" + +import json +import os +import re +import sys +import time +import uuid +from collections.abc import Mapping + +_DIAG_TRUNCATE = 2000 +_TOOL_INPUT_TRUNCATE = 500 +_TOOL_RESULT_TRUNCATE = 1000 +_DEFAULT_MODEL = "deepseek-v4-pro" +_DEFAULT_DEEPSEEK_BASE_URL = "https://api.deepseek.com/v1" + +# A tiny stand-in knowledge base — the RAG specialist's corpus (replaces the +# upstream Qdrant collection so the graph runs without embeddings). +_KB = { + "metformin": "Metformin is first-line for type 2 diabetes; common side effects " + "are GI upset (nausea, diarrhea) and, rarely, lactic acidosis.", + "aspirin": "Low-dose aspirin is used for secondary cardiovascular prevention; " + "bleeding risk rises with dose and with anticoagulants.", + "ibuprofen": "Ibuprofen is an NSAID for pain/inflammation; avoid in renal " + "impairment and use caution with anticoagulants.", +} + + +# ── ACP stdio I/O ───────────────────────────────────────────────────────────── + + +def send(msg): + sys.stdout.write(json.dumps(msg) + "\n") + sys.stdout.flush() + + +def recv(): + while True: + line = sys.stdin.readline() + if not line: + raise EOFError("stdin closed") + line = line.strip() + if not line: + continue + return json.loads(line) + + +# ── Provider / model resolution (deepseek over the OpenAI-compatible proxy) ──── + + +def resolve_base_url(env: Mapping[str, str] | None = None) -> str: + env = os.environ if env is None else env + for key in ("BENCHFLOW_PROVIDER_BASE_URL", "DEEPSEEK_BASE_URL"): + val = env.get(key) + if val and val.strip(): + return val.strip() + return _DEFAULT_DEEPSEEK_BASE_URL + + +def resolve_api_key(env: Mapping[str, str] | None = None) -> str: + env = os.environ if env is None else env + for key in ("BENCHFLOW_PROVIDER_API_KEY", "DEEPSEEK_API_KEY"): + val = env.get(key) + if val and val.strip(): + return val.strip() + return "" + + +def resolve_model_id(requested: str, env: Mapping[str, str] | None = None) -> str: + env = os.environ if env is None else env + model = (requested or "").strip() + if not model: + model = (env.get("BENCHFLOW_PROVIDER_MODEL") or "").strip() + if not model: + model = _DEFAULT_MODEL + if "/" in model: + model = model.split("/", 1)[1] + return model + + +def build_chat_model(model_id: str, env: Mapping[str, str] | None = None): + """deepseek-v4-pro is OpenAI-compatible → ChatOpenAI at the proxy base_url. + + Imported lazily so the module parses (and its pure helpers stay testable) + without langchain_openai installed in the host dev venv — it is only present + inside the shim's sandbox venv. + """ + from langchain_openai import ChatOpenAI + + env = os.environ if env is None else env + kwargs = { + "model": model_id, + "base_url": resolve_base_url(env), + "api_key": resolve_api_key(env) or "EMPTY", + "temperature": _float_env(env, "BENCHFLOW_MODEL_TEMPERATURE", 0.2), + "max_tokens": _int_env(env, "BENCHFLOW_MODEL_MAX_TOKENS", 2048), + } + return ChatOpenAI(**kwargs) + + +def _float_env(env, key, default): + raw = env.get(key) + if raw is None or not str(raw).strip(): + return default + try: + return float(raw) + except ValueError: + return default + + +def _int_env(env, key, default): + raw = env.get(key) + if raw is None or not str(raw).strip(): + return default + try: + return int(float(raw)) + except ValueError: + return default + + +# ── The LangGraph medical assistant (router → specialists → handoff → guardrail) + + +def build_graph(chat_model): + """Construct the supervisor→specialists StateGraph. Lazy import so the module + parses without langgraph in the host venv.""" + from typing import TypedDict + + from langgraph.graph import END, START, StateGraph + + class S(TypedDict, total=False): + query: str + route: str + evidence: str + answer: str + confidence: float + safe: bool + web_tried: bool + + threshold = _float_env(os.environ, "MEDICAL_CONFIDENCE_THRESHOLD", 0.6) + + def supervisor(state: S) -> S: + out = chat_model.invoke([{"role": "user", "content": + f"You route a medical question to a specialist. Question: {state['query']!r}. " + "Reply with ONE word: 'kb' if a drug/condition knowledge base likely covers " + "it, else 'web'."}]).content.lower() + return {"route": "web" if "web" in out else "kb"} + + def retrieve_kb(state: S) -> S: + q = state["query"].lower() + hits = [v for k, v in _KB.items() if k in q] + return {"evidence": " ".join(hits)} + + def web_search(state: S) -> S: + ev = f"[web] current guidance relevant to: {state['query']}" + return {"evidence": (state.get("evidence", "") + " " + ev).strip(), + "web_tried": True} + + def answer(state: S) -> S: + out = chat_model.invoke([{"role": "user", "content": + f"Question: {state['query']}\nEvidence: {state.get('evidence') or '(none)'}\n" + "Answer concisely for a clinician. Then on a NEW line output " + "'CONFIDENCE: <0.0-1.0>' for how well the evidence supports your answer."}] + ).content + m = re.search(r"CONFIDENCE:\s*([01](?:\.\d+)?)", out) + conf = float(m.group(1)) if m else 0.5 + ans = re.split(r"CONFIDENCE:", out)[0].strip() + return {"answer": ans, "confidence": conf} + + def guardrail(state: S) -> S: + out = chat_model.invoke([{"role": "user", "content": + f"Output guardrail. The assistant answered: {state['answer']!r}. Is it safe and " + "appropriately cautious as medical information (not a diagnosis)? Reply yes/no."}] + ).content.lower() + return {"safe": "yes" in out} + + def after_supervisor(state: S) -> str: + return "web_search" if state["route"] == "web" else "retrieve_kb" + + def after_answer(state: S) -> str: + if state.get("confidence", 0.0) < threshold and not state.get("web_tried"): + return "web_search" # confidence-gated handoff to the fallback specialist + return "guardrail" + + g = StateGraph(S) + for name, fn in (("supervisor", supervisor), ("retrieve_kb", retrieve_kb), + ("web_search", web_search), ("answer", answer), + ("guardrail", guardrail)): + g.add_node(name, fn) + g.add_edge(START, "supervisor") + g.add_conditional_edges("supervisor", after_supervisor, + {"retrieve_kb": "retrieve_kb", "web_search": "web_search"}) + g.add_edge("retrieve_kb", "answer") + g.add_edge("web_search", "answer") + g.add_conditional_edges("answer", after_answer, + {"web_search": "web_search", "guardrail": "guardrail"}) + g.add_edge("guardrail", END) + return g.compile() + + +def run_graph(model_id: str, query: str, cwd: str, session_id: str, + env: Mapping[str, str] | None = None) -> dict: + """Stream the medical graph, emitting one ACP tool_call per node, and write + the final answer to ``/answer.md`` for the verifier.""" + start = time.time() + chat_model = build_chat_model(model_id, env) + graph = build_graph(chat_model) + + final: dict = {"query": query} + order: list[str] = [] + for chunk in graph.stream({"query": query}, stream_mode="updates"): + for node_name, delta in (chunk or {}).items(): + order.append(node_name) + call_id = f"{node_name}_{len(order)}" + _emit_tool_call(session_id, call_id, node_name, delta or {}) + _emit_tool_result(session_id, call_id, json.dumps(delta or {})) + if isinstance(delta, dict): + final.update(delta) + + answer_text = final.get("answer") or "(no answer produced)" + _emit_text(session_id, answer_text) + out_path = os.path.join(cwd, "answer.md") + try: + with open(out_path, "w") as fh: + fh.write(answer_text + "\n") + except OSError as exc: + _emit_text(session_id, f"[medical-assistant: could not write {out_path}: {exc}]") + + return { + "path": " → ".join(order), + "route": final.get("route"), + "confidence": final.get("confidence"), + "safe": final.get("safe"), + "wall_clock_seconds": round(time.time() - start, 2), + } + + +# ── ACP notification helpers ────────────────────────────────────────────────── + + +def _emit_text(session_id: str, text: str): + send({"jsonrpc": "2.0", "method": "session/update", "params": { + "sessionId": session_id, "update": { + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": text[:_DIAG_TRUNCATE]}}}}) + + +def _emit_tool_call(session_id: str, call_id: str, node: str, delta): + try: + input_text = json.dumps(delta) + except (TypeError, ValueError): + input_text = str(delta) + send({"jsonrpc": "2.0", "method": "session/update", "params": { + "sessionId": session_id, "update": { + "sessionUpdate": "tool_call", "toolCallId": call_id, + "title": node, "kind": _node_kind(node), "status": "in_progress", + "input": input_text[:_TOOL_INPUT_TRUNCATE]}}}) + + +def _emit_tool_result(session_id: str, call_id: str, result: str): + send({"jsonrpc": "2.0", "method": "session/update", "params": { + "sessionId": session_id, "update": { + "sessionUpdate": "tool_call_update", "toolCallId": call_id, + "status": "completed", "content": [{"type": "content", "content": { + "type": "text", "text": result[:_TOOL_RESULT_TRUNCATE]}}]}}}) + + +def _node_kind(node: str) -> str: + n = (node or "").lower() + if "supervisor" in n: + return "think" + if "web" in n: + return "search" + if "retrieve" in n or "kb" in n: + return "read" + if "guardrail" in n: + return "other" + return "other" + + +# ── Main ACP loop ───────────────────────────────────────────────────────────── + + +def main(): + session_id = "medical-assistant-shim" + cwd = "/app" + model = "" + + while True: + try: + msg = recv() + except EOFError: + break + + method = msg.get("method", "") + req_id = msg.get("id") + params = msg.get("params", {}) + + if method == "initialize": + send({"jsonrpc": "2.0", "id": req_id, "result": { + "protocolVersion": 1, + "agentCapabilities": {"loadSession": False, + "promptCapabilities": {"image": False, "audio": False}}, + "agentInfo": {"name": "medical-assistant", "version": "1.0"}}}) + + elif method == "session/new": + cwd = params.get("cwd", "/app") + session_id = f"medical-{uuid.uuid4().hex[:8]}" + send({"jsonrpc": "2.0", "id": req_id, "result": {"sessionId": session_id}}) + + elif method == "session/set_model": + model = params.get("modelId", "") + send({"jsonrpc": "2.0", "id": req_id, "result": {}}) + + elif method == "session/prompt": + text = "" + for part in params.get("prompt", []): + if isinstance(part, dict) and part.get("type") == "text": + text += part.get("text", "") + model_id = resolve_model_id(model) + try: + metrics = run_graph(model_id, text.strip(), cwd, session_id) + _emit_text(session_id, f"[medical-assistant: path {metrics['path']} · " + f"confidence={metrics['confidence']} · safe={metrics['safe']} · " + f"{metrics['wall_clock_seconds']}s]") + except Exception as e: + _emit_text(session_id, f"[medical-assistant error: {type(e).__name__}: {e}]") + send({"jsonrpc": "2.0", "id": req_id, "result": {"stopReason": "end_turn"}}) + + elif method == "session/cancel": + send({"jsonrpc": "2.0", "id": req_id, "result": {}}) + + elif method == "session/request_permission": + options = params.get("options", []) + option_id = options[0].get("optionId", "default") if options else "default" + send({"jsonrpc": "2.0", "id": req_id, "result": { + "outcome": {"outcome": "selected", "optionId": option_id}}}) + + else: + if req_id is not None: + send({"jsonrpc": "2.0", "id": req_id, "result": {}}) + + +if __name__ == "__main__": + main() diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index 2825406e..f5c62a2e 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -214,6 +214,10 @@ def _js_agent_launch(binary: str, args: str = "") -> str: # Path to the deepagents ACP shim (runs LangChain's create_deep_agent as an ACP agent) _DEEPAGENTS_SHIM = (Path(__file__).parent / "deepagents_acp_shim.py").read_text() +# Path to the medical-assistant ACP shim (runs a LangGraph supervisor→specialists +# medical graph — the Multi-Agent-Medical-Assistant pattern — as an ACP agent) +_MEDICAL_SHIM = (Path(__file__).parent / "medical_acp_shim.py").read_text() + def _json_settings_merge(path: str, mutator: str) -> str: """Idempotent JSON-settings merge as a one-line bash snippet.""" @@ -766,6 +770,45 @@ class AgentConfig: # carries a registered provider prefix), with DEEPSEEK_* as a fallback # (auto_inherit_env propagates those). ), + "medical-assistant": AgentConfig( + name="medical-assistant", + description="Multi-Agent Medical Assistant — a LangGraph supervisor→" + "specialists graph (router → KB/web specialists → confidence-gated handoff " + "→ output guardrail), the Multi-Agent-Medical-Assistant pattern, run via " + "ACP shim driving deepseek-v4-pro through the OpenAI-compatible provider", + install_cmd=( + "export DEBIAN_FRONTEND=noninteractive && " + # Pinned interpreter via uv (same pattern as deepagents/openhands): task + # base images ship Python as old as 3.6/3.8, but langgraph needs >=3.11. + "( command -v curl >/dev/null 2>&1 || " + f" {_apt_install('curl', 'ca-certificates')} ) && " + "( command -v uv >/dev/null 2>&1 || " + " ( curl -LsSf https://astral.sh/uv/install.sh | sh >/dev/null 2>&1 ) ) && " + 'export PATH="$HOME/.local/bin:$PATH" && ' + "uv venv --python 3.12 /opt/benchflow/medical-venv && " + # langgraph drives the StateGraph; langchain-openai is the OpenAI- + # compatible chat model for deepseek-v4-pro. + "uv pip install -q --python /opt/benchflow/medical-venv/bin/python " + "langgraph langchain-openai && " + "chmod -R a+rX /opt/benchflow/medical-venv && " + "chmod o+x /root /root/.local /root/.local/share " + "/root/.local/share/uv /root/.local/share/uv/python 2>/dev/null; " + + _install_python_script( + f"{_BENCHFLOW_BIN_PREFIX}/medical-acp-shim", _MEDICAL_SHIM + ) + # Verify langgraph imports through the pinned venv (mirrors deepagents' + # import check) so a failed install fails loudly here, not at launch. + + " && /opt/benchflow/medical-venv/bin/python -c " + "'import langgraph, langchain_openai' >/dev/null 2>&1" + ), + launch_cmd=( + f"/opt/benchflow/medical-venv/bin/python {_BENCHFLOW_BIN_PREFIX}/medical-acp-shim" + ), + protocol="acp", + requires_env=[], # inferred from --model at runtime (DEEPSEEK_API_KEY, etc.) + # api_protocol / env_mapping intentionally empty — the shim builds an + # OpenAI-compatible ChatOpenAI from BENCHFLOW_PROVIDER_* (DEEPSEEK_* fallback). + ), "openhands": AgentConfig( name="openhands", description="OpenHands agent via ACP (multi-model, Python-based)", From b0febd2fe9c3c0ccf04aa1959cdea0645dcd44c7 Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Tue, 30 Jun 2026 00:49:19 +0000 Subject: [PATCH 08/54] feat(trajectories): multi-agent bf.* attribution + agent tree Track a multi-agent workflow's raw LLM calls through ONE pooled LiteLLM proxy as a structured trace tree, instead of one undifferentiated llm_trajectory. - providers/litellm_logging.py: the callback records per-call bf.* metadata (agent_id, agent_name, span_kind, parent_agent_id, run_id, session_id, and any other bf.* key) under request.body['bf'] instead of dropping all but model_group. One shared proxy log is now splittable per agent. - trajectories/agents.py: build_agent_tree() reconstructs an unmixed AgentTree from the tagged exchanges (group by agent, link by parent pointer), the same parent-pointer model LangSmith/OTel use. - examples/medical: the LangGraph slice tags each node via extra_body; verified end-to-end as one unmixed tree on local, docker, and daytona (run_agent_tree.py, run_in_sandbox.py). The single-trajectory problem is solved the industry way. - examples/arena/run_per_seat_proxy.py + arena/policy.py: per-seat proxy variant for black-box agents (separate trajectory per seat). - docs: reference/multi-agent-trajectory.md (the contract + industry comparison) and task-standard.md G7 (the normative bf.* metadata contract). Tests: tests/test_litellm_logging.py + tests/trajectories/test_agent_tree.py. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01JDzpphBuqpjMH23dVJEoYX --- docs/reference/multi-agent-trajectory.md | 256 +++++++++++++++++++++ docs/task-standard.md | 38 +++ examples/arena/run_per_seat_proxy.py | 97 ++++++++ examples/medical/medical_assistant.py | 42 +++- examples/medical/run_agent_tree.py | 150 ++++++++++++ examples/medical/run_in_sandbox.py | 178 ++++++++++++++ examples/medical/sandbox_env/Dockerfile | 7 + src/benchflow/arena/policy.py | 9 +- src/benchflow/providers/litellm_logging.py | 13 ++ src/benchflow/trajectories/__init__.py | 4 + src/benchflow/trajectories/agents.py | 109 +++++++++ tests/test_litellm_logging.py | 80 +++++++ tests/trajectories/test_agent_tree.py | 71 ++++++ 13 files changed, 1052 insertions(+), 2 deletions(-) create mode 100644 docs/reference/multi-agent-trajectory.md create mode 100644 examples/arena/run_per_seat_proxy.py create mode 100644 examples/medical/run_agent_tree.py create mode 100644 examples/medical/run_in_sandbox.py create mode 100644 examples/medical/sandbox_env/Dockerfile create mode 100644 src/benchflow/trajectories/agents.py create mode 100644 tests/trajectories/test_agent_tree.py diff --git a/docs/reference/multi-agent-trajectory.md b/docs/reference/multi-agent-trajectory.md new file mode 100644 index 00000000..a6594a57 --- /dev/null +++ b/docs/reference/multi-agent-trajectory.md @@ -0,0 +1,256 @@ +# Multi-agent trajectory tracking through the LiteLLM proxy + +Status: **design / target** (runtime-deferred, see `task-standard.md` G7). This +document records the industry comparison that motivates the design and the +contract BenchFlow should adopt so that a multi-agent workflow hosted through the +provider proxy produces **one structured trace tree** — agent identity *and* +agent-to-agent relationships — instead of one undifferentiated +`llm_trajectory.jsonl`. + +## Problem + +BenchFlow routes an agent's raw LLM calls through a loopback LiteLLM proxy. The +callback (`src/benchflow/providers/litellm_logging.py`) records `model + messages` +per call and keeps only `model_group` from request `metadata` +(`litellm_logging.py:130,150`) — every other tag is dropped. So when a +multi-agent workflow (a LangGraph supervisor→specialists graph, an Omnigent +session with sub-agents, concurrent arena seats) shares one proxy, the result is a +**flat, undifferentiated** log: you can see *that* N calls happened, not *which +agent* made each, nor how the agents relate. + +The earlier medical-bench workaround — **one proxy per agent → one file per agent** +(`out/medical-hosted//trajectory/llm_trajectory.jsonl`) — preserves +identity *by filename* but **destroys the relationship structure** (which agent +spawned/handed-off-to which) that every observability tool treats as first-class. +It also doesn't scale (one subprocess proxy per agent) and can't represent a +dynamic spawn tree. It was a demonstration, not the design. + +## What the industry does (surveyed, verified) + +A deep multi-source survey (OpenTelemetry GenAI semconv, LangSmith/LangGraph, +Langfuse, OpenLLMetry/Traceloop, Langtrace, and framework SDKs) returns one +**unanimous** structural answer: + +| System | Data model | Per-agent identity carrier | Relationship mechanism | Session/seat grouping | Capture | +|---|---|---|---|---|---| +| **OTel GenAI semconv** | one trace = tree of **typed** spans | `gen_ai.agent.id` / `gen_ai.agent.name` + span name `invoke_agent {name}` + `gen_ai.operation.name` | parent/child span **nesting** (`execute_tool` nests under `invoke_agent`; agents nest under `invoke_workflow`) | `gen_ai.conversation.id` (Conditionally Required; **never synthesize** a fallback) | in-process instrumentation; cross-process via `traceparent` | +| **LangSmith / LangGraph** | run tree | `run_type` + `name` (no `run_type=agent`; an agent is a `chain` run) | `parent_run_id` + `dotted_order` + `child_run_ids` | `thread_id`; `trace_id` = root run id | LangChain callbacks (`run_id`/`parent_run_id`) | +| **Langfuse** | trace + nested observations | typed observations (generation/span/event) + `name` | `parent_observation_id` | `session_id` | SDK / decorators / OTel | +| **OpenLLMetry / Traceloop** | OTel spans | `traceloop.span.kind ∈ {workflow,task,agent,tool}` + `entity.name` + `workflow.name` | OTel context nesting via `run_id`/`parent_run_id` | — | SDK monkey-patch + LangChain `BaseCallbackHandler`; injects `extra_headers` to **propagate** context | +| **Langtrace** | OTel spans ("adhere to OTEL") | OTel `gen_ai.*` attributes | OTel parent nesting | — | OTel instrumentation | +| **OpenAI Agents SDK** | trace + spans | agent span; **handoffs** are edges | root span via `Runner` | `group_id` | SDK tracing | +| **LiteLLM proxy** (BenchFlow today) | flat callback log | `metadata` body field + headers — but callback keeps only `model_group` | none natively (must add a parent pointer in metadata) | `metadata` | proxy callback (`StandardLoggingPayload`) | + +Three load-bearing facts, each confirmed 3-0 across independent verifiers: + +1. **A multi-agent run is ONE trace = a tree of typed spans joined by explicit + parent pointers — never a flat event log.** Universal across OTel, LangSmith, + Langfuse, Langtrace, OpenLLMetry. +2. **Per-agent identity is carried *on each call* (span attributes + span name), + so a single shared stream is differentiated per-call, not per-stream.** Agent / + tool / LLM-call / orchestration are first-class, **distinct span types** + (`gen_ai.operation.name`: `chat`, `invoke_agent`, `execute_tool`, + `invoke_workflow`, …; LangSmith `run_type`: chain/llm/tool/retriever/…). +3. **Relationships are captured at call time via parent/child nesting, not inferred + post-hoc** — `execute_tool` under `invoke_agent`; LangChain `run_id` → + `parent_run_id`. Session/seat grouping uses a single real conversation id; the + OTel spec **forbids synthetic fallbacks** (no UUID / trace-id / content hash). + +**Critical caveat for BenchFlow:** none of these tools solve the +"shared collector loses the tag" problem *with an HTTP proxy*. They instrument +**in-process** (framework callbacks, SDK monkey-patching) where the active +agent/parent context is already known, and propagate it across process boundaries +via OTel context (`traceparent`). BenchFlow's proxy sits *outside* the agent +process, so the agent context must be **explicitly attached to each request** — the +proxy cannot recover it otherwise. (Also: OTel GenAI agent conventions are +*experimental* / SHOULD-level; only `gen_ai.operation.name` and +`gen_ai.provider.name` are strictly Required. `gen_ai.agent.*` officially live on +agent-lifecycle spans, not every raw chat span — attaching them per raw call is a +deliberate, reasonable BenchFlow extension.) + +## The design: one pooled proxy + per-call metadata → one trace tree + +Adopt the dominant industry shape, adapted for an out-of-process proxy: + +**1. Each LLM request carries agent context in `metadata`** (the LiteLLM request +body field, which the proxy forwards to the logging callback — see *Verification* +below). A minimal, OTel-aligned schema: + +```jsonc +"metadata": { + "bf.agent_id": "answer", // stable id of the calling agent/node (~ gen_ai.agent.id) + "bf.agent_name": "answer", // human label (~ gen_ai.agent.name) + "bf.span_kind": "chat", // chat | invoke_agent | execute_tool | invoke_workflow (~ gen_ai.operation.name) + "bf.parent_agent_id": "supervisor", // parent pointer → reconstructs the tree (~ parent_run_id) + "bf.session_id": "medical-run-1", // real conversation/seat id, NEVER synthesized (~ gen_ai.conversation.id) + "bf.run_id": "answer#2" // this call's id, so children can point at it +} +``` + +**2. The callback records a span row** instead of today's bare model+messages line: +`litellm_logging.py:_base_record` stops discarding `metadata` and persists +`bf.agent_id` / `bf.agent_name` / `bf.span_kind` / `bf.parent_agent_id` / +`bf.session_id` / `bf.run_id` alongside the existing `model_group`. Every proxied +call becomes one typed span with a parent pointer. + +**3. The trajectory becomes a tree.** `trajectory_from_litellm_callback_log` +reconstructs parent/child structure from `bf.parent_agent_id` / `bf.run_id` +exactly as LangSmith does from `parent_run_id` / `dotted_order`. One +`llm_trajectory.jsonl` then holds the whole multi-agent run, splittable per agent +**and** navigable as a tree — no per-agent files, no lost relationships. + +Why this over "one proxy per agent": separate files preserve identity by filename +only and throw away the parent/child + handoff structure that is the *point* of a +multi-agent trace. The pooled-proxy + per-call-tag model is what every surveyed +tool does. + +### The uniform adapter + +BenchFlow cannot adopt a single framework as "the" multi-agent host (Omnigent, +the closest candidate, explicitly does **not** host LangGraph/CrewAI/AutoGen — see +below). The uniform layer is instead a **thin BenchFlow-side contract**: each +framework's native per-agent + parent context is mapped onto the one `metadata` +schema by a small per-framework shim, before the call leaves the agent process. + +| Framework | Native per-agent + parent context the shim maps from | +|---|---| +| LangChain / LangGraph | `BaseCallbackHandler` `run_id` / `parent_run_id` / node name → `bf.*` (already exposed) | +| OpenAI Agents SDK | `Runner` / root span + agent name + handoff edges | +| Omnigent | its internal conversation tree (`parent_conversation_id`/`root_conversation_id`/`agent_id`) → `bf.*` | +| custom (e.g. our medical slice) | the node passes its own name as `bf.agent_id` when it builds the `ChatOpenAI` call | +| AutoGen / CrewAI / Swarm | **under-evidenced** — per-agent + handoff exposure to an interceptor not yet confirmed; needs a per-framework spike before claiming support | + +The shim's only job is the mapping. Identity and relationships already exist +inside every framework; BenchFlow just needs them attached to the request. + +### Unified `bf.*` vocabulary (consolidated with the adapter proposal) + +The adapter proposal (PR #847) independently specified a richer per-call +attribution set. To avoid forking two schemas, that vocabulary is folded into the +one `bf.*` namespace. The callback captures **any** `bf.*` key generically (it +strips the `bf.` prefix from every metadata key), so the extended dimensions flow +through with **no code change** (verified in `tests/test_litellm_logging.py`). + +| `bf.*` key | status | meaning | OTel / #847 analogue | +|---|---|---|---| +| `agent_id` / `agent_name` | implemented | the calling agent/node | `gen_ai.agent.id` / `.name` | +| `parent_agent_id` | implemented | parent pointer (tree edge) | `parent_run_id` / `framework_parent_id` | +| `run_id` | implemented | this call's id | `llm.call_id` | +| `session_id` | implemented | conversation/seat/rollout id (never synthesized) | `gen_ai.conversation.id` / `rollout_id` | +| `span_kind` | implemented | `chat`/`invoke_agent`/`execute_tool`/`invoke_workflow` — the `relation` hook | `gen_ai.operation.name` | +| `role` | extended (#847) | declared role (planner/implementer/reviewer) | `role` | +| `scene` / `turn_index` | extended (#847) | scene + turn within a scene | `scene` / `turn` | +| `team_id` | extended (#847) | team / sub-graph grouping | `team_id` | +| `framework` / `framework_node_id` | extended (#847) | framework name + native node id | `framework` / `framework_node_id` | +| `trace_id` | extended (#847) | cross-process / cross-framework trace correlation | `trace_id` | + +### Uniform adapter protocol (from #847) + +A hosted framework is wrapped by a thin adapter — a *normalizer*, not a +reimplementation — with four operations: + +- `detect(task_dir, spec)` — can this adapter host the workflow? +- `prepare(ctx)` — install deps, write the LiteLLM env, compile the launch command. +- `run(ctx)` — run the external workflow **inside the BenchFlow sandbox**. +- `collect(ctx)` — gather native logs, normalize to BenchFlow events + graph edges, + return an `AdapterTraceBundle` (framework-native raw logs under + `trajectory/raw//` + coverage diagnostics: attribution quality, missing + metadata, unsupported relations, redaction state). + +Declared under a `benchflow.multi_agent` block (target / `benchflow:`-namespaced): +`capture_raw_llm: required` **fails closed** when zero LLM calls are captured; +`relationship_graph: required` persists `trajectory/agent_graph.json` (agent / team / +sub-graph nodes + edges carrying a `relation` ∈ {delegates, supervises, reviews, +handoff, parallel_child, fan_in, …}). When a framework **cannot** inject per-call +metadata, fall back to **one LiteLLM virtual key / route alias per role** as the +attribution channel. Report multi-agent lift only against a **matched single-agent +baseline** (same task set, tools, answer contract, usage + logging). + +The implemented `build_agent_tree()` is the minimal in-memory realization of this +graph (one parent pointer); `agent_graph.json` is its richer persisted form, with +`bf.span_kind` as the existing per-edge `relation` hook. + +## Does Omnigent satisfy this? (assessed separately, verified) + +Short answer: **not as wired today.** Omnigent the *framework* is a strong +multi-agent foundation, but the integration does not carry that across the +BenchFlow boundary. Per-requirement: + +- **Multi-agent support** — framework yes (recursive `AgentSpec.sub_agents`, + `sys_session_send` spawns children, parent-linked conversation tree); the BF + adapter drives a **single one-shot harness** (`omnigent run -p`) and only `pi` + is wired end-to-end. *Partial.* +- **Per-agent attribution through the proxy** — **fails.** Omnigent's HTTP adapter + sends only `{model, messages, tools?, stream?, **extra}` with `Content-Type` + + `Authorization` headers — **no agent id on the wire** — and BenchFlow's callback + would drop it anyway. `agent_id` lives only in Omnigent's DB. +- **Relationships** — Omnigent models a real parent-linked tree + shared OTel trace + internally (`db_models.py`), but it is **not populated by the one-shot CLI path** + (each `omnigent run -p` is a fresh, unlinked conversation, no `traceparent`) and + is never exported to BenchFlow. *Partial / not wired.* +- **Proxy routing of ALL traffic** — only `pi` (OpenAI-wire) is proven proxied; + native CLI sub-harnesses (Claude Code/Codex) use a separate + `HARNESS_*_GATEWAY_BASE_URL` / `claude_gateway_shim` channel the adapter never + sets, and sub-agent spawns pass no gateway override → **bypass risk**. +- **Uniform host for many frameworks** — hosts coding-agent CLIs + `claude_sdk` + + `agents_sdk`, but LangGraph/CrewAI/AutoGen/LangChain/DeepAgents are explicitly + "Not natively supported." *Not a universal host.* + +Two paths to make Omnigent conformant (complementary): **(A)** inject `agent_id` +into proxy metadata + have the callback record it (the contract above) — best for +per-call cost/attribution; **(B)** run Omnigent in server/session mode and export +its native conversation tree, joined to proxy usage by `agent_id` — best for the +relationship graph. `agent_id` from (A) is the clean join key for (B). + +## Implementation phases + +- **P0 ✅ DONE — callback records attribution.** `litellm_logging.py:_base_record` + now extracts every `bf.*` key from request `metadata` into `request.body['bf']` + (prefix stripped) instead of keeping only `model_group`. The trajectory importer + carries it through unchanged. Unit-tested in `tests/test_litellm_logging.py`. +- **P1 ✅ DONE — medical bench as proof, one proxy.** `examples/medical`'s `_llm` + tags each node's call with `bf.*` via `extra_body` (the verified client + mechanism); `trajectories/build_agent_tree()` reconstructs the unmixed agent + tree. Verified end-to-end on **three backends**, all producing the identical + tree `supervisor(1) → answer(2) → guardrail(1)`, `unmixed_ok: true`: + `run_agent_tree.py` (local, docker-proxy) and `run_in_sandbox.py` (agent running + **inside** a docker container and a remote daytona sandbox, through the proxy). + Unit-tested in `tests/trajectories/test_agent_tree.py`. +- **P2 — tree-shaped canonical trajectory.** Add a multi-agent trajectory kind that + reconstructs + emits the tree (parent pointers / dotted-order); fix + `n_tool_calls` to count real tool spans. +- **P3 — OTLP export option.** Optionally emit OTLP spans so existing backends + (Langfuse, Arize Phoenix/OpenInference, OpenLLMetry collectors) can render the + tree, instead of (or alongside) the bespoke JSONL. +- **Omnigent track.** Wire `HARNESS_*_GATEWAY_*` at the proxy to close the bypass; + drive server/session mode; export the conversation tree. + +## Verification (the research's #1 open question) + +The whole design rests on: *does a request-body `metadata` field survive to the +LiteLLM proxy logging callback?* This was the survey's top open question. It is +**confirmed empirically** against BenchFlow's own loopback proxy: a single chat +completion sent with `metadata: {bf.agent_id, bf.agent_name, bf.span_kind, +bf.parent_agent_id, bf.session_id, bf.run_id}` had **all six fields arrive intact** +at the callback (verdict PASS). The callback already *reads* +`litellm_params.metadata` (`litellm_logging.py:130`) — LiteLLM merges the request +body's `metadata` into it — so the only missing piece is *recording* the agent +fields instead of keeping just `model_group` (P0). + +## Open questions + +- AutoGen / CrewAI / OpenAI Swarm: how each exposes per-agent identity + explicit + handoff **edges** to an interceptor (needs a per-framework spike). +- Bespoke JSONL vs native OTLP export (P3) — what is lost/gained by staying custom. +- Concurrent seats vs supervisor→specialist: sibling sub-trees under one trace, OTel + span links, or distinct `bf.session_id` per seat? + +## Sources + +OpenTelemetry GenAI spans + agent spans +(`opentelemetry.io/docs/specs/semconv/gen-ai/`), LangSmith run-data-format +(`docs.langchain.com/langsmith/run-data-format`), Langfuse data model +(`langfuse.com/docs/observability/data-model`), OpenLLMetry/Traceloop semantic +conventions (`traceloop.com/docs/openllmetry`), Langtrace +(`github.com/Scale3-Labs/langtrace`). 27 sources fetched; 22 claims confirmed, 3 +refuted, across 110 research agents. diff --git a/docs/task-standard.md b/docs/task-standard.md index 152a98b4..0a05204c 100644 --- a/docs/task-standard.md +++ b/docs/task-standard.md @@ -685,6 +685,43 @@ Richer team semantics such as role membership enforcement, summaries, handoff artifacts, parallel teams, branch routing, and full trajectory sharing are parsed as draft surface but must fail closed until a runtime owns them. +### Multi-agent trajectory tracking (target, G7) + +A multi-agent interaction (`multi-agent-sequential`, `arena-concurrent`, or a +hosted multi-agent framework) must produce **one trace tree**, not one +undifferentiated `llm_trajectory.jsonl`. Every surveyed observability standard +(OpenTelemetry GenAI, LangSmith, Langfuse, OpenLLMetry) models a run as a tree of +**typed spans joined by parent pointers**, with per-agent identity carried *on each +call* — never as a flat event log, and never as one-file-per-agent (which keeps +identity by filename but discards the relationships). See +[`reference/multi-agent-trajectory.md`](reference/multi-agent-trajectory.md). + +Because BenchFlow's provider proxy sits **outside** the agent process, agent +context cannot be recovered post-hoc: each proxied raw-LLM call MUST carry it in +the request `metadata` (verified to survive intact to the LiteLLM callback). The +contract: + +| `metadata` field | meaning | OTel analogue | +|---|---|---| +| `bf.agent_id` / `bf.agent_name` | the calling agent/node | `gen_ai.agent.id` / `gen_ai.agent.name` | +| `bf.span_kind` | `chat` \| `invoke_agent` \| `execute_tool` \| `invoke_workflow` | `gen_ai.operation.name` | +| `bf.parent_agent_id` / `bf.run_id` | parent pointer + this call's id (reconstruct the tree) | `parent_run_id` / `dotted_order` | +| `bf.session_id` | real conversation/seat id — **never synthesized** | `gen_ai.conversation.id` | + +The LiteLLM callback records these as a span row; the trajectory importer rebuilds +the tree from `bf.parent_agent_id` / `bf.run_id`. The callback captures **any** +`bf.*` key generically, so richer dimensions (`bf.role`, `bf.scene`, +`bf.turn_index`, `bf.team_id`, `bf.framework`, `bf.framework_node_id`, +`bf.trace_id`) need no further code. A **uniform adapter** (`detect` / `prepare` / +`run` / `collect` under a `benchflow.multi_agent` block, with `capture_raw_llm: +required` failing closed and `agent_graph.json` for relations) maps each framework's +native per-agent + parent context (LangGraph `run_id`/`parent_run_id`; OpenAI Agents +SDK root span; Omnigent's conversation tree) onto this one schema — BenchFlow does +not adopt any single framework as *the* host. Runtime is deferred (G7); the contract +is declared so multi-agent tasks parse without loss. See +[`reference/multi-agent-trajectory.md`](reference/multi-agent-trajectory.md) for the +full vocabulary + adapter protocol. + Simulated users and nudges should be explicit about runtime type: ```yaml @@ -847,6 +884,7 @@ Current implementation status: | `reward.json` precedence | yes | partial | prefer JSON when present and reject both-present mismatches | | metrics aggregate policy | yes | partial | `mean`, `weighted_mean`, and `weighted_sum`; richer engines remain target work | | `arena-concurrent` interaction (G4) | no | no | add interaction-mode schema now; A2A bridge for the agent-under-test + concurrently-running assessor at M2 | +| multi-agent trajectory tracking (G7) | no | no | per-call `bf.*` agent-attribution metadata on proxied raw-LLM calls (verified to reach the callback) → one structured trace tree; callback records agent/parent/conversation, importer rebuilds the tree; uniform per-framework metadata shim; M1/M2 | | hybrid reward envelope (G1) | partial | no | declared cross-surface product/sum of factors; M1 | | GAIN aggregation (G2) | no | no | dynamic live baseline + ceiling; M1 | | leaderboard-submission (G5) | partial | no | hosted / hidden external scorer with durable result record; M1 | diff --git a/examples/arena/run_per_seat_proxy.py b/examples/arena/run_per_seat_proxy.py new file mode 100644 index 00000000..c479e753 --- /dev/null +++ b/examples/arena/run_per_seat_proxy.py @@ -0,0 +1,97 @@ +"""Separate trajectories per agent — ONE LiteLLM proxy PER seat. + +The shared-proxy run writes a single `llm_trajectory.jsonl` with every seat's +calls mixed in (and the callback records only model+messages, so a per-call agent +tag does NOT survive). To track each concurrent agent separately, give each its +OWN `ensure_litellm_runtime` → its own callback log → its own +`/trajectory/llm_trajectory.jsonl` + its own usage/cost. This is exactly how +a BenchFlow multi-role rollout isolates roles. + + set -a; . ./sb-run.env; set +a + uv run python examples/arena/run_per_seat_proxy.py 3 +""" + +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +import httpx + +from benchflow.arena import ProxyChatPolicy, SeatTrajectory, run_arena +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from floor_deepseek import HighCardFloor, pick_number, render_for + + +async def _main() -> None: + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required") + n = int(sys.argv[1]) if len(sys.argv) > 1 else 3 + seats = [f"seat-{i}" for i in range(n)] + run_dir = Path("out/arena-per-seat") + + # 1) ONE proxy per seat — separate callback log → separate trajectory + usage + runtimes: dict = {} + envs: dict = {} + print(f"starting {n} per-seat LiteLLM proxies…", flush=True) + for s in seats: + env, rt = await ensure_litellm_runtime( + agent="deepagents", + agent_env={"DEEPSEEK_API_KEY": os.environ["DEEPSEEK_API_KEY"]}, + model="deepseek/deepseek-v4-pro", runtime=None, environment="local", + session_id=f"arena-{s}", + ) + runtimes[s], envs[s] = rt, env + print(f" {s}: {env['BENCHFLOW_PROVIDER_BASE_URL']}", flush=True) + + floor = HighCardFloor(n) + tr = SeatTrajectory(run_dir) + usages: dict = {} + try: + async with httpx.AsyncClient() as http: + policies = { + s: ProxyChatPolicy( + s, http, render=render_for(s), pick=pick_number, + base=envs[s]["BENCHFLOW_PROVIDER_BASE_URL"], # this seat's proxy + api_key=envs[s]["BENCHFLOW_PROVIDER_API_KEY"], + model=envs[s]["BENCHFLOW_PROVIDER_MODEL"], + temperature=0.9, max_tokens=256, recorder=tr, + ) + for s in seats + } + res = await run_arena(seats, floor, lambda s: policies[s], + deadline_s=180.0, poll_s=0.05) + await asyncio.sleep(1.5) # let each proxy's async callback flush + finally: + for s in seats: + await stop_provider_runtime(runtimes[s]) # stop FIRST → parses callback log + for s in seats: # then read each proxy's usage + trajectory (populated on stop) + usages[s] = extract_usage(runtimes[s]) + traj = getattr(getattr(runtimes[s], "server", None), "trajectory", None) + if traj is not None and traj.exchanges: + d = run_dir / s / "trajectory" + d.mkdir(parents=True, exist_ok=True) + (d / "llm_trajectory.jsonl").write_text(traj.to_jsonl(redact_keys=True)) + + st = floor.standings() + print("\nstandings :", st, "· conserved", sum(st.values()), + "· status", {s: r["status"] for s, r in res.items()}) + print("=== SEPARATE per-agent trajectories + usage ===") + for s in seats: + p = run_dir / s / "trajectory" / "llm_trajectory.jsonl" + ex = len(p.read_text().splitlines()) if p.exists() else 0 + u = usages.get(s, {}) + print(f" {s}: {p} ({ex} exchange) " + f"tokens={u.get('total_tokens')} cost=${u.get('cost_usd')}") + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/examples/medical/medical_assistant.py b/examples/medical/medical_assistant.py index 95e35d1e..f77782c9 100644 --- a/examples/medical/medical_assistant.py +++ b/examples/medical/medical_assistant.py @@ -52,6 +52,41 @@ def set_agent_provider(agent: str, base: str, key: str, model: str) -> None: _AGENT_PROVIDERS[agent] = (base, key, model) +# Agent-tree edges for this graph: who dispatched/answered to whom. The router is +# the root; the answer specialist is its child; the guardrail reviews the answer. +# This is the parent map sent as bf.parent_agent_id so ONE shared proxy log +# reconstructs into an unmixed agent tree (no per-agent proxies needed). +_AGENT_PARENT = {"supervisor": None, "answer": "supervisor", "guardrail": "answer"} +_run_counter: dict[str, int] = {} + + +def reset_run(session_id: str | None = None) -> str: + """Start a fresh agent-tree run: clear per-agent call counters and pin the + session id stamped on every ``bf.*`` tag this run.""" + _run_counter.clear() + sid = session_id or os.environ.get("BF_SESSION_ID") or "medical-run" + os.environ["BF_SESSION_ID"] = sid + return sid + + +def _bf_metadata(agent: str) -> dict: + """The per-call ``bf.*`` attribution for one agent's LLM call (see + docs/reference/multi-agent-trajectory.md).""" + n = _run_counter.get(agent, 0) + 1 + _run_counter[agent] = n + bf = { + "bf.agent_id": agent, + "bf.agent_name": agent, + "bf.span_kind": "chat", + "bf.run_id": f"{agent}#{n}", + "bf.session_id": os.environ.get("BF_SESSION_ID") or "medical-run", + } + parent = _AGENT_PARENT.get(agent) + if parent: + bf["bf.parent_agent_id"] = parent + return bf + + def _llm(agent: str = "default") -> ChatOpenAI: cfg = _AGENT_PROVIDERS.get(agent) if cfg is not None: @@ -61,8 +96,12 @@ def _llm(agent: str = "default") -> ChatOpenAI: key = (os.environ.get("BENCHFLOW_PROVIDER_API_KEY") or os.environ.get("DEEPSEEK_API_KEY") or "EMPTY") model = os.environ.get("BENCHFLOW_PROVIDER_MODEL") or "deepseek/deepseek-v4-pro" + # extra_body lands `metadata` as a top-level request-body field (verified to + # reach the proxy callback; model_kwargs does NOT survive). One proxy, tagged + # per agent → an unmixed agent tree. return ChatOpenAI(model=model, base_url=base, api_key=key, - temperature=0.2, max_tokens=2048) + temperature=0.2, max_tokens=2048, + extra_body={"metadata": _bf_metadata(agent)}) def _log(state: S, node: str, note: str) -> None: @@ -144,4 +183,5 @@ def build_graph(): def run(query: str) -> S: + reset_run() # fresh per-agent bf.run_id counters for this run return build_graph().invoke({"query": query, "trace": []}) diff --git a/examples/medical/run_agent_tree.py b/examples/medical/run_agent_tree.py new file mode 100644 index 00000000..a32d39fe --- /dev/null +++ b/examples/medical/run_agent_tree.py @@ -0,0 +1,150 @@ +"""MVP: a multi-agent workflow through ONE BenchFlow proxy → an unmixed agent tree. + +Runs the LangGraph medical assistant (supervisor → answer → guardrail) through a +SINGLE LiteLLM proxy. Each node tags its call with bf.* metadata (agent id, parent +pointer, run id, session id); the proxy callback records it; we reconstruct the +AGENT TREE from the one shared trajectory and assert the agents are NOT mixed. + +Same runner, three backends — pass the proxy environment so behaviour can be +compared identically across local / docker / daytona: + + set -a; . ~/sb-run.env; set +a + uv run python examples/medical/run_agent_tree.py --env local + uv run python examples/medical/run_agent_tree.py --env docker + uv run python examples/medical/run_agent_tree.py --env daytona +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +from pathlib import Path + +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) +from benchflow.trajectories import build_agent_tree + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import medical_assistant as ma + +EXPECTED_PARENTS = {"supervisor": None, "answer": "supervisor", "guardrail": "answer"} + + +def _render(tree, usage_by_agent) -> list[str]: + lines: list[str] = [] + + def walk(node, depth: int) -> None: + u = usage_by_agent.get(node.agent_id, {}) + lines.append( + f"{' ' * depth}└─ {node.agent_id} " + f"({len(node.exchanges)} call(s)" + + (f", {u.get('tokens')} tok" if u.get("tokens") else "") + + ")" + ) + for child in node.children: + walk(child, depth + 1) + + for root in tree.roots: + walk(root, 0) + for orphan in tree.orphans: + lines.append(f"[orphan] {orphan.agent_id} (parent " + f"{orphan.parent_agent_id!r} absent)") + walk(orphan, 0) + return lines + + +def _assert_unmixed(tree) -> list[str]: + """Return a list of problems; empty list == the tree is correct + unmixed.""" + problems: list[str] = [] + seen = {n.agent_id for n in tree.nodes()} + if seen != {"supervisor", "answer", "guardrail"}: + problems.append(f"agents = {sorted(seen)}, expected supervisor/answer/guardrail") + if [r.agent_id for r in tree.roots] != ["supervisor"]: + problems.append(f"roots = {[r.agent_id for r in tree.roots]}, expected [supervisor]") + for node in tree.nodes(): + # NOT MIXED: every exchange under this node was made BY this agent + for ex in node.exchanges: + aid = (ex.request.body.get("bf") or {}).get("agent_id") + if aid != node.agent_id: + problems.append(f"MIXED: {node.agent_id} node holds a call from {aid!r}") + # parent link matches the declared graph edge + if EXPECTED_PARENTS.get(node.agent_id) != node.parent_agent_id: + problems.append( + f"{node.agent_id} parent = {node.parent_agent_id!r}, " + f"expected {EXPECTED_PARENTS.get(node.agent_id)!r}" + ) + return problems + + +async def _main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--env", default="local", choices=["local", "docker", "daytona"]) + p.add_argument("--query", default="What are the main side effects of metformin?") + args = p.parse_args() + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required") + os.environ.setdefault("MEDICAL_CONFIDENCE_THRESHOLD", "1.01") # force the web handoff + + print(f"== starting ONE LiteLLM proxy (environment={args.env}) ==", flush=True) + agent_env, runtime = await ensure_litellm_runtime( + agent="deepagents", + agent_env={"DEEPSEEK_API_KEY": os.environ["DEEPSEEK_API_KEY"]}, + model="deepseek/deepseek-v4-pro", runtime=None, environment=args.env, + session_id=f"medical-tree-{args.env}", + ) + ma._AGENT_PROVIDERS.clear() # ONE proxy, not one-per-agent + os.environ.update(agent_env) # the graph's ChatOpenAI now hits the proxy + ma.reset_run(session_id=f"medical-{args.env}") + print(" proxy base:", agent_env.get("BENCHFLOW_PROVIDER_BASE_URL"), flush=True) + + result: dict = {} + try: + result = await asyncio.to_thread(ma.run, args.query) + await asyncio.sleep(1.5) # let the proxy callback flush + finally: + await stop_provider_runtime(runtime) + usage = extract_usage(runtime) + traj = getattr(getattr(runtime, "server", None), "trajectory", None) + + tree = build_agent_tree(traj) if traj is not None else build_agent_tree([]) + path = " → ".join(t["node"] for t in result.get("trace", [])) + print("\nagent path :", path) + print("proxy calls:", len(traj.exchanges) if traj else 0, + "| usage:", json.dumps(usage)) + print("\nAGENT TREE (one shared proxy, reconstructed from bf.*):") + for line in _render(tree, {}): + print(" " + line) + + problems = _assert_unmixed(tree) + out_dir = Path(f"out/medical-tree-{args.env}") + (out_dir / "trajectory").mkdir(parents=True, exist_ok=True) + if traj is not None and traj.exchanges: + (out_dir / "trajectory" / "llm_trajectory.jsonl").write_text( + traj.to_jsonl(redact_keys=True)) + (out_dir / "agent_tree.json").write_text(json.dumps({ + "environment": args.env, + "roots": [r.agent_id for r in tree.roots], + "agents": {n.agent_id: {"parent": n.parent_agent_id, "calls": len(n.exchanges)} + for n in tree.nodes()}, + "unmixed_ok": not problems, + }, indent=2)) + print(f"\npersisted: {out_dir}/trajectory/llm_trajectory.jsonl + agent_tree.json") + + if problems: + print("\n❌ TREE CHECK FAILED:") + for pr in problems: + print(" -", pr) + return 1 + print(f"\n✅ [{args.env}] one proxy · {len(tree.roots)} root · agents NOT mixed · " + "tree = supervisor → answer → guardrail") + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(_main())) diff --git a/examples/medical/run_in_sandbox.py b/examples/medical/run_in_sandbox.py new file mode 100644 index 00000000..ad78c20b --- /dev/null +++ b/examples/medical/run_in_sandbox.py @@ -0,0 +1,178 @@ +"""Run the multi-agent medical workflow INSIDE a BenchFlow sandbox, end to end. + +The whole LangGraph workflow (supervisor → answer → guardrail) executes isolated +in a Docker or Daytona sandbox, routing every agent's LLM call through the +BenchFlow LiteLLM proxy. Each node tags its call with bf.* metadata; the proxy +records it; after the run we download the trajectory and reconstruct the AGENT +TREE, asserting the agents are NOT mixed — identical behaviour on both backends. + +Topology (a runtime detail; the tracking outcome is identical): + - docker : agent runs in a container; proxy on the host docker bridge, reached + from the container. + - daytona: agent runs in a remote sandbox; proxy runs sandbox-local (127.0.0.1). + + set -a; . ~/sb-run.env; set +a # DEEPSEEK_API_KEY + set -a; . ~/.daytona.env; set +a # DAYTONA_API_KEY (daytona only) + uv run python examples/medical/run_in_sandbox.py --env docker + uv run python examples/medical/run_in_sandbox.py --env daytona +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +import tempfile +from pathlib import Path + +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) +from benchflow.task.config import SandboxConfig +from benchflow.trajectories import build_agent_tree + +HERE = Path(__file__).resolve().parent +ENV_DIR = HERE / "sandbox_env" +sys.path.insert(0, str(HERE)) +from run_agent_tree import EXPECTED_PARENTS, _render # noqa: E402 - after sys.path + +# Runs inside the sandbox: import the uploaded graph, run it through the proxy +# (BENCHFLOW_PROVIDER_* come in via the exec env), print the path as one line. +_DRIVER = r''' +import json, os, sys +sys.path.insert(0, "/app") +import medical_assistant as ma +os.environ.setdefault("MEDICAL_CONFIDENCE_THRESHOLD", "1.01") +ma._AGENT_PROVIDERS.clear() # ONE shared proxy, tagged per agent +ma.reset_run(session_id=os.environ.get("BF_SESSION_ID", "medical-sandbox")) +result = ma.run(os.environ.get("MED_QUERY", "What are the main side effects of metformin?")) +print("TRACE_JSON:" + json.dumps({ + "path": [t["node"] for t in result.get("trace", [])], + "route": result.get("route"), "confidence": result.get("confidence"), + "safe": result.get("safe"), +})) +''' + + +def _make_sandbox(env: str): + cfg = SandboxConfig(allow_internet=True) + common = dict(environment_dir=ENV_DIR, environment_name="medicalmvp", + session_id=f"med-{env}", rollout_paths=None, task_env_config=cfg) + if env == "docker": + from benchflow.sandbox.docker import DockerSandbox + return DockerSandbox(**common) + from benchflow.sandbox.daytona import DaytonaSandbox + return DaytonaSandbox(**common) + + +def _assert_unmixed(tree) -> list[str]: + problems: list[str] = [] + seen = {n.agent_id for n in tree.nodes()} + if seen != {"supervisor", "answer", "guardrail"}: + problems.append(f"agents = {sorted(seen)}, expected supervisor/answer/guardrail") + if [r.agent_id for r in tree.roots] != ["supervisor"]: + problems.append(f"roots = {[r.agent_id for r in tree.roots]}, expected [supervisor]") + for node in tree.nodes(): + for ex in node.exchanges: + aid = (ex.request.body.get("bf") or {}).get("agent_id") + if aid != node.agent_id: + problems.append(f"MIXED: {node.agent_id} node holds a call from {aid!r}") + if EXPECTED_PARENTS.get(node.agent_id) != node.parent_agent_id: + problems.append(f"{node.agent_id} parent={node.parent_agent_id!r}, " + f"expected {EXPECTED_PARENTS.get(node.agent_id)!r}") + return problems + + +async def _main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--env", default="docker", choices=["docker", "daytona"]) + p.add_argument("--query", default="What are the main side effects of metformin?") + args = p.parse_args() + if not os.environ.get("DEEPSEEK_API_KEY"): + raise SystemExit("DEEPSEEK_API_KEY required") + + sandbox = _make_sandbox(args.env) + print(f"== provisioning {args.env} sandbox (image w/ langgraph deps) ==", flush=True) + await sandbox.start(force_build=False) + runtime = None + trace = None + try: + # Proxy: sandbox-local for daytona, host-bridge for docker (reached from + # the agent container). Either way the agent only sees BENCHFLOW_PROVIDER_*. + sandbox_arg = sandbox if args.env == "daytona" else None + proxy_environment = "daytona" if args.env == "daytona" else "docker" + agent_env, runtime = await ensure_litellm_runtime( + agent="deepagents", + agent_env={"DEEPSEEK_API_KEY": os.environ["DEEPSEEK_API_KEY"]}, + model="deepseek/deepseek-v4-pro", runtime=None, + environment=proxy_environment, + session_id=f"medical-sandbox-{args.env}", sandbox=sandbox_arg, + ) + print(" proxy base (agent-visible):", + agent_env.get("BENCHFLOW_PROVIDER_BASE_URL"), flush=True) + + # Upload the graph + a tiny driver into the sandbox. + await sandbox.upload_file(str(HERE / "medical_assistant.py"), + "/app/medical_assistant.py") + tmp = Path(tempfile.mkdtemp()) + (tmp / "driver.py").write_text(_DRIVER) + await sandbox.upload_file(str(tmp / "driver.py"), "/app/driver.py") + + run_env = dict(agent_env) + run_env.update({ + "MEDICAL_CONFIDENCE_THRESHOLD": "1.01", + "BF_SESSION_ID": f"medical-{args.env}", + "MED_QUERY": args.query, + }) + print("== running the medical graph INSIDE the sandbox ==", flush=True) + res = await sandbox.exec("python3 /app/driver.py", env=run_env, timeout_sec=300) + for line in (res.stdout or "").splitlines(): + if line.startswith("TRACE_JSON:"): + trace = json.loads(line[len("TRACE_JSON:"):]) + print(" driver rc =", res.return_code, + "| in-sandbox path:", " → ".join(trace["path"]) if trace else "(none)") + if res.return_code != 0: + print(" STDERR:", (res.stderr or "")[:900]) + await asyncio.sleep(1.5) # let the proxy callback flush + finally: + if runtime is not None: + await stop_provider_runtime(runtime) # downloads callback log -> trajectory + + usage = extract_usage(runtime) if runtime is not None else {} + traj = getattr(getattr(runtime, "server", None), "trajectory", None) + tree = build_agent_tree(traj) if traj is not None else build_agent_tree([]) + print("\nproxy calls:", len(traj.exchanges) if traj else 0, "| usage:", json.dumps(usage)) + print(f"AGENT TREE (one shared proxy, in {args.env} sandbox):") + for line in _render(tree, {}): + print(" " + line) + + problems = _assert_unmixed(tree) + out_dir = Path(f"out/medical-sandbox-{args.env}") + (out_dir / "trajectory").mkdir(parents=True, exist_ok=True) + if traj is not None and traj.exchanges: + (out_dir / "trajectory" / "llm_trajectory.jsonl").write_text( + traj.to_jsonl(redact_keys=True)) + (out_dir / "agent_tree.json").write_text(json.dumps({ + "environment": args.env, + "ran_in_sandbox": True, + "roots": [r.agent_id for r in tree.roots], + "agents": {n.agent_id: {"parent": n.parent_agent_id, "calls": len(n.exchanges)} + for n in tree.nodes()}, + "unmixed_ok": not problems, + }, indent=2)) + + await sandbox.stop(delete=True) + if problems or not trace: + print("\n❌ FAILED:", "; ".join(problems) or "no in-sandbox trace") + return 1 + print(f"\n✅ [{args.env}] agent ran IN sandbox · one proxy · agents NOT mixed · " + "tree = supervisor → answer → guardrail") + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(_main())) diff --git a/examples/medical/sandbox_env/Dockerfile b/examples/medical/sandbox_env/Dockerfile new file mode 100644 index 00000000..8ed2f3cd --- /dev/null +++ b/examples/medical/sandbox_env/Dockerfile @@ -0,0 +1,7 @@ +# Minimal image for running the LangGraph medical multi-agent workflow inside a +# BenchFlow sandbox (docker or daytona). Deps only — the graph code is uploaded +# at run time so iterating on it needs no rebuild. +FROM python:3.12-slim +RUN pip install --no-cache-dir "langgraph>=0.2" "langchain-openai>=0.2" && \ + python -c "import langgraph, langchain_openai; print('deps ok')" +WORKDIR /app diff --git a/src/benchflow/arena/policy.py b/src/benchflow/arena/policy.py index 911077e0..64c4eae5 100644 --- a/src/benchflow/arena/policy.py +++ b/src/benchflow/arena/policy.py @@ -54,13 +54,20 @@ def __init__( render: Callable[[Observation], str], pick: Callable[[str, list[dict[str, Any]]], dict[str, Any]], model: str | None = None, + base: str | None = None, + api_key: str | None = None, temperature: float = 0.7, max_tokens: int = 256, recorder: SeatTrajectory | None = None, ) -> None: self.seat, self.http = seat, http self.render, self.pick = render, pick - self.base, self.key, self.model = provider_config(model) + # Explicit base/api_key/model win — pass this seat's OWN proxy here to get + # a separate trajectory + usage per agent; else fall back to the shared env. + pc_base, pc_key, pc_model = provider_config(model) + self.base = base or pc_base + self.key = api_key or pc_key + self.model = model or pc_model self.temperature, self.max_tokens = temperature, max_tokens self.recorder = recorder diff --git a/src/benchflow/providers/litellm_logging.py b/src/benchflow/providers/litellm_logging.py index 776bedf5..f7e93ea7 100644 --- a/src/benchflow/providers/litellm_logging.py +++ b/src/benchflow/providers/litellm_logging.py @@ -144,6 +144,19 @@ def _base_record(self, kwargs: dict[str, Any], start_time: Any, end_time: Any) - if value is not None: request_body[key] = value request_body = {k: v for k, v in request_body.items() if v is not None} + # Multi-agent attribution: a hosted multi-agent workflow tags each call + # with bf.* fields in request metadata (agent id/name, span kind, parent + # pointer, run id, conversation id). Record them under request.body["bf"] + # so one shared proxy log reconstructs into an unmixed agent tree. The + # "bf." prefix isolates ours from litellm-internal metadata keys. + if isinstance(metadata, dict): + bf = { + key[3:]: value + for key, value in metadata.items() + if isinstance(key, str) and key.startswith("bf.") + } + if bf: + request_body["bf"] = bf return { "request_model": kwargs.get("model"), "provider_model": litellm_params.get("model") or kwargs.get("model"), diff --git a/src/benchflow/trajectories/__init__.py b/src/benchflow/trajectories/__init__.py index 3565aead..655d5430 100644 --- a/src/benchflow/trajectories/__init__.py +++ b/src/benchflow/trajectories/__init__.py @@ -13,6 +13,7 @@ rather than HTTP / OTLP traffic. """ +from .agents import AgentNode, AgentTree, build_agent_tree from .tree import RolloutNode, RolloutTree, Step, branch_points, trajectory from .types import LLMExchange, LLMRequest, LLMResponse, Trajectory @@ -22,6 +23,9 @@ "LLMRequest", "LLMResponse", "Trajectory", + "AgentNode", + "AgentTree", + "build_agent_tree", # tree-native Rollout data model "RolloutNode", "RolloutTree", diff --git a/src/benchflow/trajectories/agents.py b/src/benchflow/trajectories/agents.py new file mode 100644 index 00000000..ceaac47f --- /dev/null +++ b/src/benchflow/trajectories/agents.py @@ -0,0 +1,109 @@ +"""Reconstruct a multi-agent run as an agent tree from one shared proxy log. + +A hosted multi-agent workflow routes every agent's raw LLM call through one +BenchFlow proxy, tagging each call with ``bf.*`` fields in request metadata +(recorded by the callback under ``request.body['bf']`` — see +``providers/litellm_logging.py``). This module turns that flat, ordered exchange +list back into the **agent tree**: one node per agent holding only its own calls +(unmixed), linked to its parent via ``bf.parent_agent_id`` — the same +parent-pointer reconstruction LangSmith does from ``parent_run_id``. See +``docs/reference/multi-agent-trajectory.md``. + +Pure data + pure functions: no I/O, no async. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from dataclasses import dataclass, field + +from .types import LLMExchange, Trajectory + +__all__ = ["AgentNode", "AgentTree", "build_agent_tree"] + + +def _bf(exchange: LLMExchange) -> dict: + body = exchange.request.body + bf = body.get("bf") if isinstance(body, dict) else None + return bf if isinstance(bf, dict) else {} + + +@dataclass(eq=False) +class AgentNode: + """One agent in the run, holding only its own LLM exchanges (unmixed).""" + + agent_id: str + agent_name: str = "" + parent_agent_id: str | None = None + exchanges: list[LLMExchange] = field(default_factory=list) + children: list[AgentNode] = field(default_factory=list) + + +@dataclass(eq=False) +class AgentTree: + """The agent forest for one run. + + ``roots`` are agents with no (in-run) parent — a supervisor, or each + independent seat in a concurrent arena. ``orphans`` are agents whose + ``parent_agent_id`` names an agent not present in this run (kept as roots of + their own sub-tree rather than silently dropped). + """ + + roots: list[AgentNode] = field(default_factory=list) + orphans: list[AgentNode] = field(default_factory=list) + + def nodes(self) -> Iterator[AgentNode]: + """Yield every agent node, pre-order from each root (then orphans).""" + stack: list[AgentNode] = list(reversed(self.roots + self.orphans)) + while stack: + node = stack.pop() + yield node + stack.extend(reversed(node.children)) + + def find(self, agent_id: str) -> AgentNode | None: + for node in self.nodes(): + if node.agent_id == agent_id: + return node + return None + + +def build_agent_tree(source: Trajectory | Iterable[LLMExchange]) -> AgentTree: + """Group exchanges by ``bf.agent_id`` and link them by ``bf.parent_agent_id``. + + Exchanges keep their original order within each agent node. An agent is a + root when it has no ``parent_agent_id`` (or points at itself); it is an orphan + when its parent is named but absent from the run. + """ + exchanges = ( + source.exchanges if isinstance(source, Trajectory) else list(source) + ) + nodes: dict[str, AgentNode] = {} + order: list[str] = [] # first-seen agent order, for stable output + for exchange in exchanges: + bf = _bf(exchange) + agent_id = bf.get("agent_id") or "unknown" + node = nodes.get(agent_id) + if node is None: + node = AgentNode( + agent_id=agent_id, + agent_name=bf.get("agent_name") or agent_id, + parent_agent_id=bf.get("parent_agent_id") or None, + ) + nodes[agent_id] = node + order.append(agent_id) + elif not node.agent_name and bf.get("agent_name"): + node.agent_name = bf["agent_name"] + node.exchanges.append(exchange) + + roots: list[AgentNode] = [] + orphans: list[AgentNode] = [] + for agent_id in order: + node = nodes[agent_id] + parent_id = node.parent_agent_id + if not parent_id or parent_id == agent_id: + roots.append(node) + elif parent_id in nodes: + nodes[parent_id].children.append(node) + else: + orphans.append(node) + return AgentTree(roots=roots, orphans=orphans) diff --git a/tests/test_litellm_logging.py b/tests/test_litellm_logging.py index 94c679b4..84c4e3e8 100644 --- a/tests/test_litellm_logging.py +++ b/tests/test_litellm_logging.py @@ -99,6 +99,86 @@ def test_callback_log_preserves_bedrock_reasoning_effort_in_request_body(): assert trajectory.exchanges[0].request.body["reasoning_effort"] == "max" +def _instantiate_callback_logger(): + """Exec the proxy-side callback module source and return a logger instance.""" + namespace: dict = {} + exec(callback_module_source(), namespace) + return namespace["BenchFlowLiteLLMLogger"]() + + +def test_callback_records_bf_agent_attribution_from_request_metadata(): + """A multi-agent run tags each LLM call with ``bf.*`` fields in request + ``metadata``; the proxy callback must record them (under + ``request.body['bf']``) so one shared proxy log can be split into an unmixed + agent tree. Today the callback keeps only ``model_group`` and drops the rest.""" + from datetime import datetime + + logger = _instantiate_callback_logger() + record = logger._base_record( + kwargs={ + "model": "benchflow-deepseek-v4-pro", + "messages": [{"role": "user", "content": "side effects?"}], + "metadata": { + "model_group": "benchflow-deepseek-v4-pro", # litellm-internal, untouched + "bf.agent_id": "answer", + "bf.agent_name": "answer", + "bf.span_kind": "chat", + "bf.parent_agent_id": "supervisor", + "bf.run_id": "answer#1", + "bf.session_id": "medical-run-1", + }, + }, + start_time=datetime(2026, 6, 29, 10, 0, 0), + end_time=datetime(2026, 6, 29, 10, 0, 1), + ) + + assert record["request"]["body"]["bf"] == { + "agent_id": "answer", + "agent_name": "answer", + "span_kind": "chat", + "parent_agent_id": "supervisor", + "run_id": "answer#1", + "session_id": "medical-run-1", + } + + +def test_callback_records_extended_bf_vocabulary_without_code_change(): + """The callback captures ANY ``bf.*`` key generically, so the richer + attribution dimensions from the adapter proposal (#847) — role, scene, + turn_index, team_id, framework, framework_node_id, trace_id — flow through + with no code change. One shared ``bf.*`` vocabulary across PRs, not two.""" + from datetime import datetime + + logger = _instantiate_callback_logger() + record = logger._base_record( + kwargs={ + "model": "m", + "messages": [], + "metadata": { + "bf.agent_id": "planner", + "bf.role": "planner", + "bf.scene": "review", + "bf.turn_index": 2, + "bf.team_id": "core", + "bf.framework": "langgraph", + "bf.framework_node_id": "plan_node", + "bf.trace_id": "t-abc", + }, + }, + start_time=datetime(2026, 6, 29, 10, 0, 0), + end_time=datetime(2026, 6, 29, 10, 0, 1), + ) + + bf = record["request"]["body"]["bf"] + assert bf["role"] == "planner" + assert bf["scene"] == "review" + assert bf["turn_index"] == 2 + assert bf["team_id"] == "core" + assert bf["framework"] == "langgraph" + assert bf["framework_node_id"] == "plan_node" + assert bf["trace_id"] == "t-abc" + + def test_litellm_failure_records_become_error_exchanges(): record = { "event": "failure", diff --git a/tests/trajectories/test_agent_tree.py b/tests/trajectories/test_agent_tree.py new file mode 100644 index 00000000..db74f36b --- /dev/null +++ b/tests/trajectories/test_agent_tree.py @@ -0,0 +1,71 @@ +"""Reconstructing an unmixed agent tree from bf.*-tagged proxy exchanges.""" + +from __future__ import annotations + +from benchflow.trajectories import LLMExchange, LLMRequest, LLMResponse, Trajectory +from benchflow.trajectories.agents import build_agent_tree + + +def _ex(agent_id: str, parent: str | None, run_id: str, *, span_kind: str = "chat") -> LLMExchange: + bf = { + "agent_id": agent_id, + "agent_name": agent_id, + "span_kind": span_kind, + "run_id": run_id, + "session_id": "run-1", + } + if parent is not None: + bf["parent_agent_id"] = parent + return LLMExchange( + request=LLMRequest(body={"model": "m", "messages": [], "bf": bf}), + response=LLMResponse( + body={"choices": [{"message": {"role": "assistant", "content": "x"}}]} + ), + ) + + +def test_build_agent_tree_groups_unmixed_and_links_parents(): + # medical-shaped: supervisor(root) -> answer(x2, web handoff) -> guardrail + exchanges = [ + _ex("supervisor", None, "supervisor#1"), + _ex("answer", "supervisor", "answer#1"), + _ex("answer", "supervisor", "answer#2"), + _ex("guardrail", "answer", "guardrail#1"), + ] + tree = build_agent_tree(Trajectory(session_id="run-1", exchanges=exchanges)) + + assert {n.agent_id for n in tree.nodes()} == {"supervisor", "answer", "guardrail"} + assert [r.agent_id for r in tree.roots] == ["supervisor"] + + supervisor = tree.find("supervisor") + answer = tree.find("answer") + guardrail = tree.find("guardrail") + assert supervisor and answer and guardrail + + # tree shape: supervisor -> answer -> guardrail + assert [c.agent_id for c in supervisor.children] == ["answer"] + assert [c.agent_id for c in answer.children] == ["guardrail"] + assert guardrail.children == [] + + # NOT MIXED: each agent node holds only its own calls + assert len(supervisor.exchanges) == 1 + assert len(answer.exchanges) == 2 + assert len(guardrail.exchanges) == 1 + for node in tree.nodes(): + assert all( + ex.request.body["bf"]["agent_id"] == node.agent_id for ex in node.exchanges + ) + + +def test_build_agent_tree_handles_concurrent_seats_as_sibling_roots(): + # arena-shaped: independent seats, no parent -> distinct sibling root sub-trees, + # never merged into one mixed node. + exchanges = [ + _ex("seat-0", None, "seat-0#1"), + _ex("seat-1", None, "seat-1#1"), + _ex("seat-0", None, "seat-0#2"), + ] + tree = build_agent_tree(exchanges) + assert sorted(r.agent_id for r in tree.roots) == ["seat-0", "seat-1"] + assert len(tree.find("seat-0").exchanges) == 2 + assert len(tree.find("seat-1").exchanges) == 1 From 5c004de0bc13dc4c04b1254b894d23c98180c3cf Mon Sep 17 00:00:00 2001 From: symphony-bot Date: Tue, 30 Jun 2026 00:49:56 +0000 Subject: [PATCH 09/54] feat(casino): real concurrent multi-agent ACP floor + live town viewer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run N real autonomous ACP agents playing ONE shared casinobench World concurrently — the live realization of the deferred arena-concurrent mode. - run_floor.py: starts casinobench's shared multi-seat World on the host, then runs a roster of seats with asyncio.gather; each seat is a connect_acp agent in its own DockerSandbox reaching the World over the docker bridge. Subscription agents (codex / claude-code) upload their host auth per seat and call their provider directly (acp_trajectory only); deepseek/proxy seats also get a per-seat raw llm_trajectory. Trajectories stream live (survive the wall-clock timeout); agents choose games freely. - town_snapshot.py: a live Stanford-Town-style viewer — casinobench's render_html canvas floor in live mode (agents walking to game stations) with a click-to-open per-agent run dossier injected; serves same-origin JSON so a Cloudflare tunnel can publish it, and falls back to the persisted run after the World stops. - agent_env/: the casino-agent-seat image — Node + codex-acp + claude-agent-acp (via benchflow's install commands) + the casino seven-tool CLI. The shim and the casino CLI package are assembled into the build context (gitignored; see README). Verified live with 10 concurrent agents (codex/gpt-5.5 + claude/sonnet-4-6 + claude/haiku-4-5) across all 18 casinobench games. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01JDzpphBuqpjMH23dVJEoYX --- examples/casino/README.md | 45 ++++ examples/casino/agent_env/.gitignore | 5 + examples/casino/agent_env/Dockerfile | 50 ++++ .../casino/agent_env/install-node-agents.sh | 7 + examples/casino/run_floor.py | 246 ++++++++++++++++++ examples/casino/town_snapshot.py | 207 +++++++++++++++ 6 files changed, 560 insertions(+) create mode 100644 examples/casino/README.md create mode 100644 examples/casino/agent_env/.gitignore create mode 100644 examples/casino/agent_env/Dockerfile create mode 100644 examples/casino/agent_env/install-node-agents.sh create mode 100644 examples/casino/run_floor.py create mode 100644 examples/casino/town_snapshot.py diff --git a/examples/casino/README.md b/examples/casino/README.md new file mode 100644 index 00000000..131f693f --- /dev/null +++ b/examples/casino/README.md @@ -0,0 +1,45 @@ +# Real multi-agent casino floor (arena-concurrent, on BenchFlow) + +N **real autonomous ACP agents** play ONE shared [casinobench](https://github.com/benchflow-ai/casinobench) +World concurrently — the live realization of the deferred `arena-concurrent` mode. +Each seat is a benchflow ACP agent in its OWN sandbox, all competing on one +leaderboard; each agent's raw + ACP trajectory is captured per seat. + +- `run_floor.py` — starts casinobench's shared World on the host, then runs a + roster of seats concurrently (`asyncio.gather`), each a `connect_acp` agent in a + `DockerSandbox` reaching the World over the docker bridge. Subscription agents + (codex / claude-code) get their auth uploaded per seat and produce + `acp_trajectory.jsonl`; deepseek/proxy seats also get a per-seat raw + `llm_trajectory.jsonl`. +- `town_snapshot.py` — serves a live Stanford-Town-style floor viewer: + casinobench's `render_html` canvas board (agents walking to game stations) in + live mode, with a click-to-open per-agent **run dossier** injected. Polls the + World, falls back to the persisted run when it ends, and feeds same-origin JSON + so a Cloudflare tunnel can publish it. +- `agent_env/` — the seat image (`casino-agent-seat`): Node + `codex-acp` + + `claude-agent-acp` (via benchflow's install commands) + the `casino` seven-tool + CLI. The deepagents shim and the casino CLI package are **assembled** into the + build context (gitignored — see `agent_env/.gitignore`). + +## Setup +This example depends on a local casinobench checkout (a separate repo). Assemble +the seat-image build context, then build it: + +```bash +CB=~/casinobench # your casinobench checkout +cp src/benchflow/agents/deepagents_acp_shim.py examples/casino/agent_env/deepagents-acp-shim +cp -r "$CB/packages/environments/casino" examples/casino/agent_env/casino-pkg +docker build -t casino-agent-seat:latest examples/casino/agent_env + +set -a; . ~/sb-run.env; set +a # DEEPSEEK_API_KEY (proxy seats) +# codex/claude seats use the host's ~/.codex/auth.json + ~/.claude/.credentials.json subscriptions +uv run python examples/casino/run_floor.py --world-port 9100 +# in another shell, publish the live viewer: +cd "$CB" && uv run python /examples/casino/town_snapshot.py \ + http://127.0.0.1:9100 /out/casino-floor/all-games ./serve & +cloudflared tunnel --url http://localhost:8899 # serving ./serve +``` + +The roster (agents × models) and the seat prompt are at the top of `run_floor.py`. +Only the models a subscription actually exposes work (e.g. codex→`gpt-5.5`, +claude→`claude-sonnet-4-6`/`claude-haiku-4-5`); others are rejected by the plan. diff --git a/examples/casino/agent_env/.gitignore b/examples/casino/agent_env/.gitignore new file mode 100644 index 00000000..468e5b14 --- /dev/null +++ b/examples/casino/agent_env/.gitignore @@ -0,0 +1,5 @@ +# Assembled into the build context at setup time, not checked in (see ../README.md): +# deepagents-acp-shim — copied from src/benchflow/agents/deepagents_acp_shim.py +# casino-pkg/ — copied from /packages/environments/casino +deepagents-acp-shim +casino-pkg/ diff --git a/examples/casino/agent_env/Dockerfile b/examples/casino/agent_env/Dockerfile new file mode 100644 index 00000000..784a9a20 --- /dev/null +++ b/examples/casino/agent_env/Dockerfile @@ -0,0 +1,50 @@ +# Agent-seat image for the multi-agent casino floor. +# +# Bakes in everything ONE deepagents ACP seat needs so `connect_acp` can run it +# with no per-seat install: the deepagents venv + ACP shim (the autonomous agent), +# and the `casino` seven-tool CLI (the play surface — the agent shells out to +# `casino observe` / `casino act` against $CASINO_URL). The graph/world itself is +# the shared casinobench World service, reached over HTTP. +FROM python:3.12-slim + +RUN apt-get update && apt-get install -y --no-install-recommends \ + curl ca-certificates git && rm -rf /var/lib/apt/lists/* + +# uv (pinned interpreter for deepagents, exactly as benchflow's install_cmd does) +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.local/bin:${PATH}" + +# deepagents harness venv (the autonomous ReAct agent) + the OpenAI-compatible +# chat model dep, verified to import. +RUN uv venv --python 3.12 /opt/benchflow/deepagents-venv && \ + uv pip install -q --python /opt/benchflow/deepagents-venv/bin/python \ + deepagents langchain-openai && \ + /opt/benchflow/deepagents-venv/bin/python -c "import deepagents, langchain_openai; print('deepagents ok')" + +# the ACP shim benchflow launches on stdio +COPY deepagents-acp-shim /opt/benchflow/bin/deepagents-acp-shim +RUN chmod a+rx /opt/benchflow/bin/deepagents-acp-shim && \ + chmod -R a+rX /opt/benchflow/deepagents-venv + +# the casino seven-tool CLI on PATH (the agent's play surface). The CLI only +# makes HTTP calls to $CASINO_URL (click + httpx) — the casinobench engine is a +# server-side dep of the package, not the CLI, so install --no-deps + the two +# runtime imports the CLI actually uses. +COPY casino-pkg /opt/casino-pkg +RUN pip install --no-cache-dir "click>=8.0" "httpx>=0.27.0" && \ + pip install --no-cache-dir --no-deps /opt/casino-pkg && \ + casino --help >/dev/null && echo "casino cli ok" + +# the node ACP agents (codex-acp + claude-agent-acp) baked in via benchflow's own +# install commands — Node 22.20.0 + the npm packages + wrappers at +# /opt/benchflow/bin/{codex-acp,claude-agent-acp}. Subscription auth is uploaded +# per-seat at run time (not baked). connect_acp's force_build=False then needs no +# per-seat install. +RUN apt-get update && apt-get install -y --no-install-recommends tar xz-utils && \ + rm -rf /var/lib/apt/lists/* +COPY install-node-agents.sh /opt/install-node-agents.sh +RUN sh /opt/install-node-agents.sh && \ + test -x /opt/benchflow/bin/codex-acp && test -x /opt/benchflow/bin/claude-agent-acp && \ + echo "node acp agents ok" + +WORKDIR /app diff --git a/examples/casino/agent_env/install-node-agents.sh b/examples/casino/agent_env/install-node-agents.sh new file mode 100644 index 00000000..25b69eeb --- /dev/null +++ b/examples/casino/agent_env/install-node-agents.sh @@ -0,0 +1,7 @@ +#!/bin/sh +set -e +export PATH="/opt/benchflow/node/bin:/opt/benchflow/bin:$PATH" +echo "=== installing codex-acp ===" +export DEBIAN_FRONTEND=noninteractive; BF_NODE_DIR=/opt/benchflow/node; BF_NODE_VERSION=22.20.0; if [ ! -x "$BF_NODE_DIR/bin/node" ]; then if ! command -v curl >/dev/null 2>&1 || ! command -v tar >/dev/null 2>&1 || ! command -v xz >/dev/null 2>&1; then if command -v apt-get >/dev/null 2>&1; then apt-get update -qq && apt-get install -y -qq curl ca-certificates tar xz-utils; elif command -v dnf >/dev/null 2>&1; then dnf -y install curl ca-certificates tar xz; elif command -v apk >/dev/null 2>&1; then apk add --no-cache curl ca-certificates tar xz; else echo 'BenchFlow JS agent bootstrap requires curl, tar, and xz' >&2; exit 127; fi; fi; arch="$(uname -m)"; case "$arch" in x86_64|amd64) node_arch=x64 ;; aarch64|arm64) node_arch=arm64 ;; *) echo "Unsupported architecture for Node.js: $arch" >&2; exit 1 ;; esac; tmp="$(mktemp -d)"; mkdir -p /opt/benchflow; curl -fsSLo "$tmp/node.tar.xz" "https://nodejs.org/dist/v${BF_NODE_VERSION}/node-v${BF_NODE_VERSION}-linux-${node_arch}.tar.xz"; rm -rf "$BF_NODE_DIR"; mkdir -p "$BF_NODE_DIR"; tar -xJf "$tmp/node.tar.xz" -C "$BF_NODE_DIR" --strip-components=1 --no-same-owner; rm -rf "$tmp"; fi; export PATH="/opt/benchflow/node/bin:$PATH"; "$BF_NODE_DIR/bin/node" --version; "$BF_NODE_DIR/bin/npm" --version && mkdir -p /opt/benchflow/js-agents /opt/benchflow/bin && export PATH="/opt/benchflow/bin:/opt/benchflow/js-agents/bin:/opt/benchflow/node/bin:$PATH" && ( /opt/benchflow/node/bin/npm install -g --prefix /opt/benchflow/js-agents @agentclientprotocol/codex-acp@0.0.45 ) && printf '%s\n' '#!/bin/sh' 'exec /opt/benchflow/node/bin/node /opt/benchflow/js-agents/bin/codex-acp "$@"' > /opt/benchflow/bin/codex-acp && chmod +x /opt/benchflow/bin/codex-acp && chmod -R a+rX /opt/benchflow && [ -x /opt/benchflow/js-agents/bin/codex-acp ] && [ -x /opt/benchflow/bin/codex-acp ] +echo "=== installing claude-agent-acp ===" +export DEBIAN_FRONTEND=noninteractive; BF_NODE_DIR=/opt/benchflow/node; BF_NODE_VERSION=22.20.0; if [ ! -x "$BF_NODE_DIR/bin/node" ]; then if ! command -v curl >/dev/null 2>&1 || ! command -v tar >/dev/null 2>&1 || ! command -v xz >/dev/null 2>&1; then if command -v apt-get >/dev/null 2>&1; then apt-get update -qq && apt-get install -y -qq curl ca-certificates tar xz-utils; elif command -v dnf >/dev/null 2>&1; then dnf -y install curl ca-certificates tar xz; elif command -v apk >/dev/null 2>&1; then apk add --no-cache curl ca-certificates tar xz; else echo 'BenchFlow JS agent bootstrap requires curl, tar, and xz' >&2; exit 127; fi; fi; arch="$(uname -m)"; case "$arch" in x86_64|amd64) node_arch=x64 ;; aarch64|arm64) node_arch=arm64 ;; *) echo "Unsupported architecture for Node.js: $arch" >&2; exit 1 ;; esac; tmp="$(mktemp -d)"; mkdir -p /opt/benchflow; curl -fsSLo "$tmp/node.tar.xz" "https://nodejs.org/dist/v${BF_NODE_VERSION}/node-v${BF_NODE_VERSION}-linux-${node_arch}.tar.xz"; rm -rf "$BF_NODE_DIR"; mkdir -p "$BF_NODE_DIR"; tar -xJf "$tmp/node.tar.xz" -C "$BF_NODE_DIR" --strip-components=1 --no-same-owner; rm -rf "$tmp"; fi; export PATH="/opt/benchflow/node/bin:$PATH"; "$BF_NODE_DIR/bin/node" --version; "$BF_NODE_DIR/bin/npm" --version && mkdir -p /opt/benchflow/js-agents /opt/benchflow/bin && export PATH="/opt/benchflow/bin:/opt/benchflow/js-agents/bin:/opt/benchflow/node/bin:$PATH" && ( /opt/benchflow/node/bin/npm install -g --prefix /opt/benchflow/js-agents @agentclientprotocol/claude-agent-acp@0.40.0 ) && printf '%s\n' '#!/bin/sh' 'exec /opt/benchflow/node/bin/node /opt/benchflow/js-agents/bin/claude-agent-acp "$@"' > /opt/benchflow/bin/claude-agent-acp && chmod +x /opt/benchflow/bin/claude-agent-acp && chmod -R a+rX /opt/benchflow && [ -x /opt/benchflow/js-agents/bin/claude-agent-acp ] && [ -x /opt/benchflow/bin/claude-agent-acp ] diff --git a/examples/casino/run_floor.py b/examples/casino/run_floor.py new file mode 100644 index 00000000..d9d74407 --- /dev/null +++ b/examples/casino/run_floor.py @@ -0,0 +1,246 @@ +"""Real multi-agent casino floor on BenchFlow — heterogeneous, all games. + +N autonomous ACP agents play ONE shared casinobench World concurrently, each in +its OWN DockerSandbox (the casino-agent-seat image: node ACP agents + the `casino` +seven-tool CLI). The default roster is 4 subscription seats: + - 2x codex-acp on gpt-5.5 (ChatGPT subscription) + - 2x claude-agent-acp on sonnet-4-6 (Claude subscription) + +Each subscription agent calls its provider directly (oauth) → it produces an +`acp_trajectory.jsonl` (its `casino` tool-calls/thinking); there is NO raw +`llm_trajectory` for subscription seats (that needs an API key fronted by the +proxy — only proxy-routed deepseek seats get one). The World runs on the host; +agents reach it over the docker bridge gateway (the same path the proxy uses). + + set -a; . ~/sb-run.env; set +a + uv run python examples/casino/run_floor.py +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import socket +import subprocess +from pathlib import Path + +import httpx + +from benchflow.acp.runtime import AgentPromptTimeoutError, connect_acp, execute_prompts +from benchflow.agents.credentials import upload_subscription_auth +from benchflow.agents.registry import AGENTS +from benchflow.providers import ( + ensure_litellm_runtime, + extract_usage, + stop_provider_runtime, +) +from benchflow.providers.litellm_runtime import _docker_host_address +from benchflow.sandbox.docker import DockerSandbox +from benchflow.task.config import SandboxConfig +from benchflow.trajectories._capture import TrajectoryWriter, make_trajectory_sink + +HERE = Path(__file__).resolve().parent +AGENT_ENV_DIR = HERE / "agent_env" +CASINOBENCH = Path(os.environ.get("CASINOBENCH_DIR", "/home/liu.10379/casinobench")) +BRIDGE = _docker_host_address() # the host gateway agent containers can reach + +# roster: (agent, model, label, count) → seats