| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243 |
- import asyncio
- import json
- import time
- import unicodedata
- from collections.abc import AsyncIterator, Callable, Mapping
- from enum import StrEnum
- from types import MappingProxyType
- from typing import Annotated, Any, Literal, TypeAlias
- from urllib.parse import urlsplit
- from uuid import uuid4
- import httpx
- from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
- from agent_lab.application.contracts import (
- AgentParams,
- DebugRunRequest,
- EventAgentParams,
- )
- from agent_lab.application.events import (
- InMemoryCalendarScheduleAdapter,
- InMemoryDeviceVolumeAdapter,
- )
- from agent_lab.application.runtime import DebugRuntime
- from agent_lab.application.queues import RuntimeQueues
- from agent_lab.application.tools import build_default_tool_registry
- from agent_lab.domain.events import ToolCallEvent
- from agent_lab.domain.messages import ChatMessage, StreamItem, TokenUsage
- from agent_lab.infrastructure.chat_client import OpenAICompatibleChatClient
- from agent_lab.infrastructure.sqlite_store import SQLiteSessionStore
- class BenchmarkMode(StrEnum):
- DUAL_AGENT = "dual_agent"
- CHAT_AGENT_TOOLS = "chat_agent_tools"
- class BenchmarkCaseId(StrEnum):
- ORDINARY_CHAT = "ordinary_chat"
- DEVICE_VOLUME_SILENT = "device_volume_silent"
- CALENDAR_SCHEDULE_TEMPLATE = "calendar_schedule_template"
- WEB_SEARCH_TWO_ANSWERS = "web_search_two_answers"
- SESSION_TERMINATE = "session_terminate"
- PARALLEL_VOLUME_SCHEDULE = "parallel_volume_schedule"
- JsonScalar: TypeAlias = str | int | float | bool | None
- BENCHMARK_MODES: tuple[BenchmarkMode, ...] = tuple(BenchmarkMode)
- BENCHMARK_CASE_IDS: tuple[BenchmarkCaseId, ...] = tuple(BenchmarkCaseId)
- BENCHMARK_SYSTEM_PROMPT = "Exercise the requested business scenario."
- BENCHMARK_MAX_TOKENS = 128
- BENCHMARK_MAX_EVENT_LOOPS = 1
- BENCHMARK_MAX_PARALLEL_EVENTS = 2
- BENCHMARK_BATCH_TIMEOUT_SECONDS = 1.0
- BENCHMARK_PARALLEL_BARRIER_TIMEOUT_SECONDS = 0.05
- class _StrictBenchmarkModel(BaseModel):
- model_config = ConfigDict(
- extra="forbid",
- strict=True,
- hide_input_in_errors=True,
- )
- class _FrozenStrictBenchmarkModel(BaseModel):
- model_config = ConfigDict(
- extra="forbid",
- strict=True,
- hide_input_in_errors=True,
- frozen=True,
- )
- def _validate_base_url(value: str) -> str:
- if not value or value != value.strip() or any(char.isspace() for char in value):
- raise ValueError("base_url must be a non-empty HTTP(S) URL")
- if any(unicodedata.category(char).startswith("C") for char in value):
- raise ValueError("base_url must not contain Unicode category C characters")
- try:
- parsed = urlsplit(value)
- hostname = parsed.hostname
- except ValueError as exc:
- raise ValueError("base_url must be a valid HTTP(S) URL") from exc
- if parsed.scheme not in {"http", "https"} or not parsed.netloc or not hostname:
- raise ValueError("base_url must use http or https and include a host")
- if parsed.username is not None or parsed.password is not None or "@" in parsed.netloc:
- raise ValueError("base_url must not contain userinfo")
- if "?" in value:
- raise ValueError("base_url must not contain a query")
- if "#" in value:
- raise ValueError("base_url must not contain a fragment")
- try:
- transport_url = httpx.URL(value)
- _ = (transport_url.host, transport_url.port, parsed.port)
- except (httpx.InvalidURL, ValueError) as exc:
- raise ValueError("base_url must be a valid HTTP(S) URL") from exc
- return value
- def _validate_model(value: str) -> str:
- stripped = value.strip()
- if not stripped:
- raise ValueError("model must not be empty")
- return stripped
- class BenchmarkTarget(_StrictBenchmarkModel):
- base_url: str
- model: str
- _base_url_validator = field_validator("base_url")(_validate_base_url)
- _model_validator = field_validator("model")(_validate_model)
- class BenchmarkConfig(BenchmarkTarget):
- schema_version: Literal[1]
- runs_per_case: int = Field(default=1, ge=1)
- modes: list[Annotated[BenchmarkMode, Field(strict=False)]] = Field(
- default_factory=lambda: list(BENCHMARK_MODES),
- min_length=1,
- )
- cases: list[Annotated[BenchmarkCaseId, Field(strict=False)]] = Field(
- default_factory=lambda: list(BENCHMARK_CASE_IDS),
- min_length=1,
- )
- @field_validator("schema_version", mode="before")
- @classmethod
- def require_integer_schema_version(cls, value: object) -> object:
- if type(value) is not int or value != 1:
- raise ValueError("schema_version must be the integer 1")
- return value
- @field_validator("modes", "cases")
- @classmethod
- def reject_duplicates(
- cls,
- values: list[BenchmarkMode | BenchmarkCaseId],
- ) -> list[BenchmarkMode | BenchmarkCaseId]:
- if len(values) != len(set(values)):
- raise ValueError("values must not contain duplicates")
- return values
- class BenchmarkModelCallTiming(_StrictBenchmarkModel):
- call_index: int = Field(ge=1)
- call_kind: Literal["chat_completion", "argument_fallback"]
- first_item_kind: str | None
- provider_ttft_ms: int | None = Field(ge=0)
- visible_ttft_ms: int | None = Field(ge=0)
- elapsed_ms: int | None = Field(ge=0)
- usage: TokenUsage | None
- class BenchmarkRunResult(_StrictBenchmarkModel):
- case_id: BenchmarkCaseId
- mode: BenchmarkMode
- iteration: int = Field(ge=1)
- status: Literal["passed", "failed"]
- initial_provider_ttft_ms: int | None = Field(ge=0)
- visible_ttft_ms: int | None = Field(ge=0)
- turn_wall_time_ms: int | None = Field(ge=0)
- prompt_tokens: int | None = Field(ge=0)
- completion_tokens: int | None = Field(ge=0)
- total_tokens: int | None = Field(ge=0)
- cached_tokens: int | None = Field(ge=0)
- model_call_count: int | None = Field(ge=0)
- fallback_count: int | None = Field(ge=0)
- tool_count: int | None = Field(ge=0)
- event_names: list[str]
- batch_event_names: list[str]
- tool_event_names: list[str]
- event_sources: list[str]
- tool_statuses: list[str]
- tool_latencies_ms: list[int | None]
- semantic_failures: list[str]
- error: str | None
- model_calls: list[BenchmarkModelCallTiming]
- parallel_events_overlapped: bool | None = None
- class BenchmarkParallelEventProbe:
- _EXPECTED_EVENTS = frozenset(
- {"device.volume.adjust", "calendar.schedule.create"}
- )
- def __init__(
- self,
- *,
- barrier_timeout_seconds: float = BENCHMARK_PARALLEL_BARRIER_TIMEOUT_SECONDS,
- ) -> None:
- self.barrier_timeout_seconds = barrier_timeout_seconds
- self._volume = InMemoryDeviceVolumeAdapter()
- self._calendar = InMemoryCalendarScheduleAdapter()
- self._active: set[str] = set()
- self._release = asyncio.Event()
- self._lock = asyncio.Lock()
- self.overlapped = False
- async def adjust(
- self,
- event_id: str,
- *,
- mode: str,
- value: int | None = None,
- delta: int | None = None,
- ) -> dict[str, Any]:
- await self._await_overlap("device.volume.adjust")
- return self._volume.adjust(
- event_id,
- mode=mode,
- value=value,
- delta=delta,
- )
- async def create(
- self,
- event_id: str,
- *,
- title: str,
- start_at: str,
- timezone: str,
- recurrence: str | None = None,
- reminder_minutes: int | None = None,
- ) -> dict[str, Any]:
- await self._await_overlap("calendar.schedule.create")
- return self._calendar.create(
- event_id,
- title=title,
- start_at=start_at,
- timezone=timezone,
- recurrence=recurrence,
- reminder_minutes=reminder_minutes,
- )
- async def _await_overlap(self, event_name: str) -> None:
- async with self._lock:
- self._active.add(event_name)
- if self._EXPECTED_EVENTS.issubset(self._active):
- self.overlapped = True
- self._release.set()
- try:
- await asyncio.wait_for(
- self._release.wait(),
- timeout=self.barrier_timeout_seconds,
- )
- except TimeoutError:
- pass
- finally:
- async with self._lock:
- self._active.discard(event_name)
- class TimingChatClient:
- def __init__(
- self,
- inner: object,
- *,
- clock: Callable[[], float] = time.perf_counter,
- ) -> None:
- self.inner = inner
- self.clock = clock
- self.timings: list[BenchmarkModelCallTiming] = []
- self._next_call_index = 0
- async def stream_chat(
- self,
- messages: list[ChatMessage],
- tools: list[dict[str, Any]],
- params: AgentParams,
- tool_choice: dict[str, Any] | None = None,
- ) -> AsyncIterator[StreamItem]:
- self._next_call_index += 1
- call_index = self._next_call_index
- started_at = self.clock()
- first_item_at: float | None = None
- first_item_kind: str | None = None
- first_message_delta_at: float | None = None
- usage: TokenUsage | None = None
- provider_elapsed_seconds = 0.0
- iterator: AsyncIterator[StreamItem] | None = None
- provider_exception_active = False
- try:
- stream_chat = getattr(self.inner, "stream_chat")
- iterator = stream_chat(
- messages=messages,
- tools=tools,
- params=params,
- tool_choice=tool_choice,
- ).__aiter__()
- while True:
- wait_started_at = self.clock()
- try:
- item = await iterator.__anext__()
- except StopAsyncIteration:
- provider_elapsed_seconds += max(
- 0.0,
- self.clock() - wait_started_at,
- )
- break
- except BaseException:
- provider_exception_active = True
- provider_elapsed_seconds += max(
- 0.0,
- self.clock() - wait_started_at,
- )
- raise
- provider_elapsed_seconds += max(
- 0.0,
- self.clock() - wait_started_at,
- )
- item_at = self.clock()
- if first_item_at is None:
- first_item_at = item_at
- first_item_kind = item.kind
- if (
- item.kind == "message_delta"
- and (item.content or "").strip()
- and first_message_delta_at is None
- ):
- first_message_delta_at = item_at
- if item.kind == "usage" and item.usage is not None:
- usage = item.usage
- yield item
- finally:
- try:
- close = getattr(iterator, "aclose", None)
- if close is not None:
- try:
- await close()
- except BaseException:
- if not provider_exception_active:
- raise
- finally:
- timing = BenchmarkModelCallTiming(
- call_index=call_index,
- call_kind=(
- "argument_fallback"
- if tool_choice is not None
- else "chat_completion"
- ),
- first_item_kind=first_item_kind,
- provider_ttft_ms=self._elapsed_ms(started_at, first_item_at),
- visible_ttft_ms=self._elapsed_ms(
- started_at,
- first_message_delta_at,
- ),
- elapsed_ms=max(0, round(provider_elapsed_seconds * 1000)),
- usage=usage,
- )
- self.timings.append(timing)
- self.timings.sort(key=lambda item: item.call_index)
- async def aclose(self) -> None:
- close = getattr(self.inner, "aclose", None)
- if close is not None:
- await close()
- @staticmethod
- def _elapsed_ms(started_at: float, ended_at: float | None) -> int | None:
- if ended_at is None:
- return None
- return max(0, round((ended_at - started_at) * 1000))
- class BenchmarkExpectation(_FrozenStrictBenchmarkModel):
- visible_messages: tuple[str, ...]
- event_names: tuple[str, ...]
- answer_count: int = Field(ge=1)
- expected_model_call_count: int = Field(ge=1)
- expected_fallback_count: int = Field(ge=0)
- terminal: bool
- first_reply_before_tool: bool
- @model_validator(mode="after")
- def match_answer_count(self) -> "BenchmarkExpectation":
- if self.answer_count != len(self.visible_messages):
- raise ValueError("answer_count must match visible_messages")
- return self
- class BenchmarkMockEvent(_FrozenStrictBenchmarkModel):
- id: str
- name: str
- arguments: tuple[tuple[str, JsonScalar], ...]
- @model_validator(mode="after")
- def validate_event(self) -> "BenchmarkMockEvent":
- if not self.id.strip():
- raise ValueError("mock event id must not be empty")
- if not self.name.strip():
- raise ValueError("mock event name must not be empty")
- keys = tuple(key for key, _ in self.arguments)
- if any(not key.strip() for key in keys):
- raise ValueError("mock event argument names must not be empty")
- if len(keys) != len(set(keys)):
- raise ValueError("mock event arguments must not contain duplicate names")
- return self
- class BenchmarkCase(_FrozenStrictBenchmarkModel):
- case_id: BenchmarkCaseId
- user_message: str
- enabled_tools: tuple[str, ...]
- expectation: BenchmarkExpectation
- mock_visible_messages: tuple[str, ...]
- mock_events: tuple[BenchmarkMockEvent, ...]
- @model_validator(mode="after")
- def validate_case_semantics(self) -> "BenchmarkCase":
- if not self.user_message.strip():
- raise ValueError("user_message must not be empty")
- if len(self.enabled_tools) != len(set(self.enabled_tools)):
- raise ValueError("enabled_tools must not contain duplicates")
- mock_event_names = tuple(event.name for event in self.mock_events)
- if mock_event_names != self.expectation.event_names:
- raise ValueError("mock event order must match expected event order")
- if self.expectation.first_reply_before_tool:
- if not self.mock_events or not self.mock_visible_messages:
- raise ValueError(
- "first_reply_before_tool requires mock text and at least one event"
- )
- return self
- def _mock_event(
- event_id: str,
- name: str,
- arguments: tuple[tuple[str, JsonScalar], ...] = (),
- ) -> BenchmarkMockEvent:
- return BenchmarkMockEvent(id=event_id, name=name, arguments=arguments)
- _SCHEDULE_ARGUMENTS: tuple[tuple[str, JsonScalar], ...] = (
- ("title", "Design review"),
- ("start_at", "2026-07-14T09:30:00+08:00"),
- ("timezone", "Asia/Shanghai"),
- )
- _SCHEDULE_CONFIRMATION = (
- "Scheduled Design review for 2026-07-14T09:30:00+08:00 "
- "(Asia/Shanghai)."
- )
- BENCHMARK_CASE_CATALOG: Mapping[BenchmarkCaseId, BenchmarkCase] = MappingProxyType(
- {
- BenchmarkCaseId.ORDINARY_CHAT: BenchmarkCase(
- case_id=BenchmarkCaseId.ORDINARY_CHAT,
- user_message="Hello.",
- enabled_tools=(
- "session.terminate",
- "device.volume.adjust",
- "calendar.schedule.create",
- "knowledge.web.search",
- ),
- expectation=BenchmarkExpectation(
- visible_messages=("Ordinary answer.",),
- event_names=(),
- answer_count=1,
- expected_model_call_count=1,
- expected_fallback_count=0,
- terminal=False,
- first_reply_before_tool=False,
- ),
- mock_visible_messages=("Ordinary answer.",),
- mock_events=(),
- ),
- BenchmarkCaseId.DEVICE_VOLUME_SILENT: BenchmarkCase(
- case_id=BenchmarkCaseId.DEVICE_VOLUME_SILENT,
- user_message="set volume to 30",
- enabled_tools=("device.volume.adjust",),
- expectation=BenchmarkExpectation(
- visible_messages=("I will set the volume to 30.",),
- event_names=("device.volume.adjust",),
- answer_count=1,
- expected_model_call_count=1,
- expected_fallback_count=0,
- terminal=False,
- first_reply_before_tool=True,
- ),
- mock_visible_messages=("I will set the volume to 30.",),
- mock_events=(
- _mock_event(
- "volume-1",
- "device.volume.adjust",
- (("mode", "absolute"), ("value", 30)),
- ),
- ),
- ),
- BenchmarkCaseId.CALENDAR_SCHEDULE_TEMPLATE: BenchmarkCase(
- case_id=BenchmarkCaseId.CALENDAR_SCHEDULE_TEMPLATE,
- user_message=(
- 'schedule "Design review" at 2026-07-14T09:30:00+08:00 '
- "timezone Asia/Shanghai"
- ),
- enabled_tools=("calendar.schedule.create",),
- expectation=BenchmarkExpectation(
- visible_messages=("I will schedule that.", _SCHEDULE_CONFIRMATION),
- event_names=("calendar.schedule.create",),
- answer_count=2,
- expected_model_call_count=1,
- expected_fallback_count=0,
- terminal=False,
- first_reply_before_tool=True,
- ),
- mock_visible_messages=("I will schedule that.",),
- mock_events=(
- _mock_event(
- "schedule-1",
- "calendar.schedule.create",
- _SCHEDULE_ARGUMENTS,
- ),
- ),
- ),
- BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS: BenchmarkCase(
- case_id=BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS,
- user_message="event batch safety",
- enabled_tools=("knowledge.web.search",),
- expectation=BenchmarkExpectation(
- visible_messages=("I will check.", "Grounded search update."),
- event_names=("knowledge.web.search",),
- answer_count=2,
- expected_model_call_count=2,
- expected_fallback_count=0,
- terminal=False,
- first_reply_before_tool=True,
- ),
- mock_visible_messages=("I will check.", "Grounded search update."),
- mock_events=(
- _mock_event(
- "search-1",
- "knowledge.web.search",
- (("query", "event batch safety"),),
- ),
- ),
- ),
- BenchmarkCaseId.SESSION_TERMINATE: BenchmarkCase(
- case_id=BenchmarkCaseId.SESSION_TERMINATE,
- user_message="end this session",
- enabled_tools=("session.terminate",),
- expectation=BenchmarkExpectation(
- visible_messages=("Goodbye.",),
- event_names=("session.terminate",),
- answer_count=1,
- expected_model_call_count=1,
- expected_fallback_count=0,
- terminal=True,
- first_reply_before_tool=True,
- ),
- mock_visible_messages=("Goodbye.",),
- mock_events=(_mock_event("terminate-1", "session.terminate"),),
- ),
- BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE: BenchmarkCase(
- case_id=BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE,
- user_message=(
- 'set volume to 30 and schedule "Design review" at '
- "2026-07-14T09:30:00+08:00 timezone Asia/Shanghai"
- ),
- enabled_tools=(
- "device.volume.adjust",
- "calendar.schedule.create",
- ),
- expectation=BenchmarkExpectation(
- visible_messages=(
- "I will set the volume to 30 and schedule that.",
- _SCHEDULE_CONFIRMATION,
- ),
- event_names=(
- "device.volume.adjust",
- "calendar.schedule.create",
- ),
- answer_count=2,
- expected_model_call_count=1,
- expected_fallback_count=0,
- terminal=False,
- first_reply_before_tool=True,
- ),
- mock_visible_messages=(
- "I will set the volume to 30 and schedule that.",
- ),
- mock_events=(
- _mock_event(
- "parallel-volume",
- "device.volume.adjust",
- (("mode", "absolute"), ("value", 30)),
- ),
- _mock_event(
- "parallel-schedule",
- "calendar.schedule.create",
- _SCHEDULE_ARGUMENTS,
- ),
- ),
- ),
- }
- )
- def _normalize_benchmark_mode(mode: BenchmarkMode | str) -> BenchmarkMode:
- try:
- return BenchmarkMode(mode)
- except ValueError as exc:
- raise ValueError(f"unsupported benchmark mode: {mode}") from exc
- def build_benchmark_request(
- case: BenchmarkCase,
- mode: BenchmarkMode | str,
- model: str,
- ) -> DebugRunRequest:
- normalized_mode = _normalize_benchmark_mode(mode)
- validated_model = _validate_model(model)
- return DebugRunRequest(
- user_message=case.user_message,
- system_prompts=[BENCHMARK_SYSTEM_PROMPT],
- pre_messages=[],
- chat_agent=AgentParams(
- model=validated_model,
- temperature=0.0,
- max_tokens=BENCHMARK_MAX_TOKENS,
- ),
- event_agent=EventAgentParams(
- model=validated_model,
- temperature=0.0,
- max_tokens=BENCHMARK_MAX_TOKENS,
- enabled_tools=list(case.enabled_tools),
- max_event_loops=BENCHMARK_MAX_EVENT_LOOPS,
- max_parallel_events=BENCHMARK_MAX_PARALLEL_EVENTS,
- batch_timeout_seconds=BENCHMARK_BATCH_TIMEOUT_SECONDS,
- ),
- tool_invocation_mode=normalized_mode.value,
- )
- def build_mock_rounds(
- case: BenchmarkCase,
- mode: BenchmarkMode | str,
- ) -> list[list[StreamItem]]:
- normalized_mode = _normalize_benchmark_mode(mode)
- first_round: list[StreamItem] = []
- if case.mock_visible_messages:
- first_round.append(StreamItem.message_delta(case.mock_visible_messages[0]))
- for mock_event in case.mock_events:
- arguments = dict(mock_event.arguments)
- event = ToolCallEvent(
- id=mock_event.id,
- name=mock_event.name,
- arguments=arguments,
- raw_arguments=json.dumps(
- arguments,
- sort_keys=True,
- separators=(",", ":"),
- ),
- )
- event_item = (
- StreamItem.text_event(event)
- if normalized_mode is BenchmarkMode.DUAL_AGENT
- else StreamItem.provider_tool_call(event)
- )
- first_round.append(event_item)
- rounds = [first_round]
- rounds.extend(
- [StreamItem.message_delta(message)]
- for message in case.mock_visible_messages[1:]
- )
- return rounds
- class MockBenchmarkChatClient:
- def __init__(
- self,
- case: BenchmarkCase,
- mode: BenchmarkMode | str,
- ) -> None:
- self.case = case
- self.mode = _normalize_benchmark_mode(mode)
- self.rounds = build_mock_rounds(case, self.mode)
- self.calls = 0
- async def stream_chat(
- self,
- messages: list[ChatMessage],
- tools: list[dict[str, Any]],
- params: AgentParams,
- tool_choice: dict[str, Any] | None = None,
- ) -> AsyncIterator[StreamItem]:
- del messages, tools, params, tool_choice
- call_index = self.calls
- if call_index >= len(self.rounds):
- raise RuntimeError(
- f"mock benchmark stream exhausted: {self.case.case_id}"
- )
- self.calls += 1
- yield StreamItem.raw_response_chunk(
- {
- "mock": True,
- "case_id": self.case.case_id.value,
- "mode": self.mode.value,
- "call_index": self.calls,
- }
- )
- for item in self.rounds[call_index]:
- yield item
- yield StreamItem.usage_item(
- TokenUsage(
- prompt_tokens=10,
- completion_tokens=5,
- total_tokens=15,
- cached_tokens=2,
- )
- )
- async def aclose(self) -> None:
- return None
- BenchmarkClientFactory: TypeAlias = Callable[[BenchmarkConfig, str], object]
- BenchmarkParallelProbeFactory: TypeAlias = Callable[[], BenchmarkParallelEventProbe]
- class _BenchmarkOutputQueue(asyncio.Queue[dict[str, Any]]):
- def __init__(self, clock: Callable[[], float]) -> None:
- super().__init__()
- self.clock = clock
- self.first_message_delta_at: float | None = None
- async def put(self, item: dict[str, Any]) -> None:
- content = item.get("content")
- if (
- item.get("type") == "message_delta"
- and isinstance(content, str)
- and content.strip()
- and self.first_message_delta_at is None
- ):
- self.first_message_delta_at = self.clock()
- await super().put(item)
- class BenchmarkRunner:
- def __init__(
- self,
- config: BenchmarkConfig,
- api_key: str | None,
- mock: bool = False,
- client_factory: BenchmarkClientFactory | None = None,
- parallel_probe_factory: BenchmarkParallelProbeFactory | None = None,
- clock: Callable[[], float] = time.perf_counter,
- ) -> None:
- if not mock and not (api_key or "").strip():
- raise ValueError("api_key is required for live benchmark runs")
- self.config = config
- self.api_key = (api_key or "").strip()
- self.mock = mock
- self.client_factory = client_factory or self._default_client_factory
- self.parallel_probe_factory = (
- parallel_probe_factory or BenchmarkParallelEventProbe
- )
- self.clock = clock
- async def run(self) -> list[BenchmarkRunResult]:
- results: list[BenchmarkRunResult] = []
- for case_id in self.config.cases:
- case = BENCHMARK_CASE_CATALOG[case_id]
- for mode in self.config.modes:
- for iteration in range(1, self.config.runs_per_case + 1):
- results.append(await self._run_one(case, mode, iteration))
- return results
- async def _run_one(
- self,
- case: BenchmarkCase,
- mode: BenchmarkMode,
- iteration: int,
- ) -> BenchmarkRunResult:
- store = SQLiteSessionStore(":memory:")
- runtime: DebugRuntime | None = None
- timing_client: TimingChatClient | None = None
- parallel_probe: BenchmarkParallelEventProbe | None = None
- result: BenchmarkRunResult | None = None
- close_error: Exception | None = None
- try:
- inner_client = (
- MockBenchmarkChatClient(case, mode)
- if self.mock
- else self.client_factory(self.config, self.api_key)
- )
- timing_client = TimingChatClient(inner_client, clock=self.clock)
- session_id = f"benchmark-{uuid4().hex}"
- request = build_benchmark_request(
- case,
- mode,
- self.config.model,
- ).model_copy(update={"session_id": session_id})
- output_queue = _BenchmarkOutputQueue(self.clock)
- if case.case_id is BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE:
- parallel_probe = self.parallel_probe_factory()
- registry = build_default_tool_registry(
- device_volume_port=parallel_probe,
- calendar_schedule_port=parallel_probe,
- )
- else:
- registry = build_default_tool_registry()
- runtime = DebugRuntime(
- timing_client,
- queues=RuntimeQueues(output=output_queue),
- registry=registry,
- session_store=store,
- clock=self.clock,
- )
- turn_started_at = self.clock()
- outputs, visible_ttft_ms, runtime_error = await self._consume_turn(
- runtime,
- request,
- turn_started_at,
- output_queue,
- )
- audits = store.list_audit_logs(session_id)
- usage = store.usage_summary(session_id)
- result = self._build_result(
- case=case,
- mode=mode,
- iteration=iteration,
- outputs=outputs,
- audits=audits,
- usage=usage,
- timings=timing_client.timings,
- visible_ttft_ms=visible_ttft_ms,
- runtime_error=runtime_error,
- parallel_events_overlapped=(
- parallel_probe.overlapped
- if parallel_probe is not None
- else None
- ),
- )
- except Exception as exc:
- result = self._failed_result(
- case.case_id,
- mode,
- iteration,
- error=str(exc),
- timings=timing_client.timings if timing_client is not None else [],
- )
- finally:
- if runtime is not None:
- try:
- await runtime.aclose()
- except Exception as exc:
- close_error = exc
- if timing_client is not None:
- try:
- await timing_client.aclose()
- except Exception as exc:
- close_error = close_error or exc
- assert result is not None
- if close_error is not None:
- error = str(close_error)
- if result.error:
- error = f"{result.error}; close failed: {error}"
- result = result.model_copy(update={"status": "failed", "error": error})
- return result
- async def _consume_turn(
- self,
- runtime: DebugRuntime,
- request: DebugRunRequest,
- turn_started_at: float,
- output_queue: _BenchmarkOutputQueue,
- ) -> tuple[list[dict[str, Any]], int | None, str | None]:
- queues = runtime.start_session(request)
- outputs: list[dict[str, Any]] = []
- visible_ttft_ms: int | None = None
- runtime_error: str | None = None
- while True:
- message = await queues.output.get()
- outputs.append(message)
- message_type = message.get("type")
- if message_type == "error":
- runtime_error = str(message.get("message") or "runtime error")
- break
- if message_type == "turn_completed":
- break
- if output_queue.first_message_delta_at is not None:
- visible_ttft_ms = max(
- 0,
- round(
- (output_queue.first_message_delta_at - turn_started_at) * 1000
- ),
- )
- return outputs, visible_ttft_ms, runtime_error
- def _build_result(
- self,
- *,
- case: BenchmarkCase,
- mode: BenchmarkMode,
- iteration: int,
- outputs: list[dict[str, Any]],
- audits: list[dict[str, Any]],
- usage: dict[str, Any],
- timings: list[BenchmarkModelCallTiming],
- visible_ttft_ms: int | None,
- runtime_error: str | None,
- parallel_events_overlapped: bool | None = None,
- ) -> BenchmarkRunResult:
- session_usage = usage["session"]
- calls = usage["calls"]
- event_audits = [
- audit for audit in audits if audit["event"] == "chat_event_detected"
- ]
- batch_audits = [
- audit for audit in audits if audit["event"] == "event_batch_results"
- ]
- batch_results = [
- item
- for audit in batch_audits
- for item in audit["details"].get("results", [])
- ]
- event_names = [audit["details"]["event_name"] for audit in event_audits]
- batch_event_names = [
- str(item.get("event_name") or "") for item in batch_results
- ]
- event_sources = [audit["details"]["event_source"] for audit in event_audits]
- tool_statuses = [str(item["status"]) for item in batch_results]
- tool_calls = [call for call in calls if call["call_kind"] == "tool_execution"]
- tool_event_names = [
- str(call.get("event_name") or "") for call in tool_calls
- ]
- ledger_model_call_count = sum(
- call["call_kind"] == "chat_completion" for call in calls
- )
- ledger_fallback_count = int(session_usage["fallback_count"])
- timing_model_call_count = sum(
- timing.call_kind == "chat_completion" for timing in timings
- )
- timing_fallback_count = sum(
- timing.call_kind == "argument_fallback" for timing in timings
- )
- token_totals, missing_usage_count = self._token_totals(timings)
- tool_count = int(session_usage["tool_count"])
- semantic_failures = self._semantic_failures(
- case=case,
- mode=mode,
- outputs=outputs,
- audits=audits,
- event_names=event_names,
- batch_event_names=batch_event_names,
- tool_event_names=tool_event_names,
- event_sources=event_sources,
- tool_statuses=tool_statuses,
- model_call_count=ledger_model_call_count,
- fallback_count=ledger_fallback_count,
- tool_count=tool_count,
- timings=timings,
- check_mock_text=self.mock,
- missing_usage_count=missing_usage_count,
- parallel_events_overlapped=parallel_events_overlapped,
- )
- initial_call = next(
- (timing for timing in timings if timing.call_kind == "chat_completion"),
- None,
- )
- status = "failed" if runtime_error or semantic_failures else "passed"
- return BenchmarkRunResult(
- case_id=case.case_id,
- mode=mode,
- iteration=iteration,
- status=status,
- initial_provider_ttft_ms=(
- initial_call.provider_ttft_ms if initial_call is not None else None
- ),
- visible_ttft_ms=visible_ttft_ms,
- turn_wall_time_ms=session_usage["turn_wall_time_ms"],
- prompt_tokens=token_totals[0],
- completion_tokens=token_totals[1],
- total_tokens=token_totals[2],
- cached_tokens=token_totals[3],
- model_call_count=timing_model_call_count,
- fallback_count=timing_fallback_count,
- tool_count=tool_count,
- event_names=event_names,
- batch_event_names=batch_event_names,
- tool_event_names=tool_event_names,
- event_sources=event_sources,
- tool_statuses=tool_statuses,
- tool_latencies_ms=[call["tool_latency_ms"] for call in tool_calls],
- semantic_failures=semantic_failures,
- error=runtime_error,
- model_calls=list(timings),
- parallel_events_overlapped=parallel_events_overlapped,
- )
- def _semantic_failures(
- self,
- *,
- case: BenchmarkCase,
- mode: BenchmarkMode,
- outputs: list[dict[str, Any]],
- audits: list[dict[str, Any]],
- event_names: list[str],
- batch_event_names: list[str],
- tool_event_names: list[str],
- event_sources: list[str],
- tool_statuses: list[str],
- model_call_count: int,
- fallback_count: int,
- tool_count: int,
- timings: list[BenchmarkModelCallTiming],
- check_mock_text: bool,
- missing_usage_count: int = 0,
- parallel_events_overlapped: bool | None = None,
- ) -> list[str]:
- failures: list[str] = []
- answers, first_answer_index, first_tool_index = self._visible_answers(outputs)
- expected = case.expectation
- if len(answers) != expected.answer_count:
- failures.append(
- f"answer_count expected {expected.answer_count}, got {len(answers)}"
- )
- if any(not answer.strip() for answer in answers):
- failures.append("visible answers must be non-empty")
- if check_mock_text and tuple(answers) != expected.visible_messages:
- failures.append(
- f"mock visible messages expected {expected.visible_messages!r}, "
- f"got {tuple(answers)!r}"
- )
- if expected.first_reply_before_tool and not (
- first_answer_index is not None
- and first_tool_index is not None
- and first_answer_index < first_tool_index
- ):
- failures.append("first visible answer must precede tool results")
- if tuple(event_names) != expected.event_names:
- failures.append(
- f"event names expected {expected.event_names!r}, got {tuple(event_names)!r}"
- )
- if tuple(batch_event_names) != expected.event_names:
- failures.append(
- "batch event names expected "
- f"{expected.event_names!r}, got {tuple(batch_event_names)!r}"
- )
- if tuple(tool_event_names) != expected.event_names:
- failures.append(
- "tool event names expected "
- f"{expected.event_names!r}, got {tuple(tool_event_names)!r}"
- )
- expected_source = (
- "text_event"
- if mode is BenchmarkMode.DUAL_AGENT
- else "provider_resolved"
- )
- expected_sources = (expected_source,) * len(expected.event_names)
- if tuple(event_sources) != expected_sources:
- failures.append(
- f"event sources expected {expected_sources!r}, "
- f"got {tuple(event_sources)!r}"
- )
- expected_statuses = tuple(
- (event_name, "success") for event_name in expected.event_names
- )
- actual_statuses = tuple(
- (
- batch_event_names[index]
- if index < len(batch_event_names)
- else None,
- tool_statuses[index] if index < len(tool_statuses) else None,
- )
- for index in range(max(len(batch_event_names), len(tool_statuses)))
- )
- if actual_statuses != expected_statuses:
- failures.append(
- f"event batch statuses expected {expected_statuses!r}, "
- f"got {actual_statuses!r}"
- )
- if tool_count != len(expected.event_names):
- failures.append(
- f"tool_count expected {len(expected.event_names)}, got {tool_count}"
- )
- timing_chat_count = sum(
- timing.call_kind == "chat_completion" for timing in timings
- )
- timing_fallback_count = sum(
- timing.call_kind == "argument_fallback" for timing in timings
- )
- if model_call_count != expected.expected_model_call_count:
- failures.append(
- "model_call_count expected "
- f"{expected.expected_model_call_count}, ledger got {model_call_count}"
- )
- if timing_chat_count != expected.expected_model_call_count:
- failures.append(
- "model_call_count expected "
- f"{expected.expected_model_call_count}, timing got {timing_chat_count}"
- )
- if fallback_count != expected.expected_fallback_count:
- failures.append(
- "fallback_count expected "
- f"{expected.expected_fallback_count}, ledger got {fallback_count}"
- )
- if timing_fallback_count != expected.expected_fallback_count:
- failures.append(
- "fallback_count expected "
- f"{expected.expected_fallback_count}, timing got "
- f"{timing_fallback_count}"
- )
- if model_call_count != timing_chat_count:
- failures.append(
- f"model_call_count ledger={model_call_count}, timing={timing_chat_count}"
- )
- if fallback_count != timing_fallback_count:
- failures.append(
- f"fallback_count ledger={fallback_count}, timing={timing_fallback_count}"
- )
- if missing_usage_count:
- failures.append(
- f"model usage missing for {missing_usage_count} call(s)"
- )
- if (
- case.case_id is BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE
- and parallel_events_overlapped is not True
- ):
- failures.append("parallel event handlers did not overlap")
- terminal = any(audit["event"] == "terminal_completed" for audit in audits)
- if terminal is not expected.terminal:
- failures.append(
- f"terminal expected {expected.terminal}, got {terminal}"
- )
- return failures
- @staticmethod
- def _visible_answers(
- outputs: list[dict[str, Any]],
- ) -> tuple[list[str], int | None, int | None]:
- answers: list[str] = []
- current: list[str] = []
- first_answer_index: int | None = None
- first_tool_index: int | None = None
- for index, message in enumerate(outputs):
- message_type = message.get("type")
- if message_type == "message_delta":
- content = str(message.get("content") or "")
- if first_answer_index is None and content.strip():
- first_answer_index = index
- current.append(content)
- continue
- if message_type == "tool_result":
- if first_tool_index is None:
- first_tool_index = index
- answer = "".join(current).strip()
- if answer:
- answers.append(answer)
- current = []
- answer = "".join(current).strip()
- if answer:
- answers.append(answer)
- return answers, first_answer_index, first_tool_index
- @staticmethod
- def _failed_result(
- case_id: BenchmarkCaseId,
- mode: BenchmarkMode,
- iteration: int,
- *,
- error: str,
- timings: list[BenchmarkModelCallTiming],
- ) -> BenchmarkRunResult:
- initial_call = next(
- (timing for timing in timings if timing.call_kind == "chat_completion"),
- None,
- )
- model_call_count = sum(
- timing.call_kind == "chat_completion" for timing in timings
- )
- fallback_count = sum(
- timing.call_kind == "argument_fallback" for timing in timings
- )
- token_totals, missing_usage_count = BenchmarkRunner._token_totals(timings)
- return BenchmarkRunResult(
- case_id=case_id,
- mode=mode,
- iteration=iteration,
- status="failed",
- initial_provider_ttft_ms=(
- initial_call.provider_ttft_ms if initial_call is not None else None
- ),
- visible_ttft_ms=None,
- turn_wall_time_ms=None,
- prompt_tokens=token_totals[0],
- completion_tokens=token_totals[1],
- total_tokens=token_totals[2],
- cached_tokens=token_totals[3],
- model_call_count=model_call_count,
- fallback_count=fallback_count,
- tool_count=None,
- event_names=[],
- batch_event_names=[],
- tool_event_names=[],
- event_sources=[],
- tool_statuses=[],
- tool_latencies_ms=[],
- semantic_failures=(
- [f"model usage missing for {missing_usage_count} call(s)"]
- if missing_usage_count
- else []
- ),
- error=error,
- model_calls=list(timings),
- parallel_events_overlapped=None,
- )
- @staticmethod
- def _token_totals(
- timings: list[BenchmarkModelCallTiming],
- ) -> tuple[tuple[int | None, int | None, int | None, int | None], int]:
- if not timings:
- return (None, None, None, None), 0
- missing_usage_count = sum(timing.usage is None for timing in timings)
- if missing_usage_count:
- return (None, None, None, None), missing_usage_count
- usages = [timing.usage for timing in timings]
- assert all(usage is not None for usage in usages)
- complete_usages = [usage for usage in usages if usage is not None]
- return (
- sum(usage.prompt_tokens for usage in complete_usages),
- sum(usage.completion_tokens for usage in complete_usages),
- sum(usage.total_tokens for usage in complete_usages),
- sum(usage.cached_tokens for usage in complete_usages),
- ), 0
- @staticmethod
- def _default_client_factory(
- config: BenchmarkConfig,
- api_key: str,
- ) -> OpenAICompatibleChatClient:
- return OpenAICompatibleChatClient(
- api_key=api_key,
- base_url=config.base_url,
- default_model=config.model,
- request_timeout_seconds=60.0,
- )
|