benchmark.py 37 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023
  1. import asyncio
  2. import json
  3. import time
  4. import unicodedata
  5. from collections.abc import AsyncIterator, Callable, Mapping
  6. from enum import StrEnum
  7. from types import MappingProxyType
  8. from typing import Annotated, Any, Literal, TypeAlias
  9. from urllib.parse import urlsplit
  10. from uuid import uuid4
  11. import httpx
  12. from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
  13. from agent_lab.application.contracts import (
  14. AgentParams,
  15. DebugRunRequest,
  16. EventAgentParams,
  17. )
  18. from agent_lab.application.runtime import DebugRuntime
  19. from agent_lab.application.queues import RuntimeQueues
  20. from agent_lab.application.tools import build_default_tool_registry
  21. from agent_lab.domain.events import ToolCallEvent
  22. from agent_lab.domain.messages import ChatMessage, StreamItem, TokenUsage
  23. from agent_lab.infrastructure.chat_client import OpenAICompatibleChatClient
  24. from agent_lab.infrastructure.sqlite_store import SQLiteSessionStore
  25. class BenchmarkMode(StrEnum):
  26. DUAL_AGENT = "dual_agent"
  27. CHAT_AGENT_TOOLS = "chat_agent_tools"
  28. class BenchmarkCaseId(StrEnum):
  29. ORDINARY_CHAT = "ordinary_chat"
  30. DEVICE_VOLUME_SILENT = "device_volume_silent"
  31. CALENDAR_SCHEDULE_TEMPLATE = "calendar_schedule_template"
  32. WEB_SEARCH_TWO_ANSWERS = "web_search_two_answers"
  33. SESSION_TERMINATE = "session_terminate"
  34. PARALLEL_VOLUME_SCHEDULE = "parallel_volume_schedule"
  35. JsonScalar: TypeAlias = str | int | float | bool | None
  36. BENCHMARK_MODES: tuple[BenchmarkMode, ...] = tuple(BenchmarkMode)
  37. BENCHMARK_CASE_IDS: tuple[BenchmarkCaseId, ...] = tuple(BenchmarkCaseId)
  38. BENCHMARK_SYSTEM_PROMPT = "Exercise the requested business scenario."
  39. BENCHMARK_MAX_TOKENS = 128
  40. BENCHMARK_MAX_EVENT_LOOPS = 1
  41. BENCHMARK_MAX_PARALLEL_EVENTS = 2
  42. BENCHMARK_BATCH_TIMEOUT_SECONDS = 1.0
  43. class _StrictBenchmarkModel(BaseModel):
  44. model_config = ConfigDict(
  45. extra="forbid",
  46. strict=True,
  47. hide_input_in_errors=True,
  48. )
  49. class _FrozenStrictBenchmarkModel(BaseModel):
  50. model_config = ConfigDict(
  51. extra="forbid",
  52. strict=True,
  53. hide_input_in_errors=True,
  54. frozen=True,
  55. )
  56. def _validate_base_url(value: str) -> str:
  57. if not value or value != value.strip() or any(char.isspace() for char in value):
  58. raise ValueError("base_url must be a non-empty HTTP(S) URL")
  59. if any(unicodedata.category(char).startswith("C") for char in value):
  60. raise ValueError("base_url must not contain Unicode category C characters")
  61. try:
  62. parsed = urlsplit(value)
  63. hostname = parsed.hostname
  64. except ValueError as exc:
  65. raise ValueError("base_url must be a valid HTTP(S) URL") from exc
  66. if parsed.scheme not in {"http", "https"} or not parsed.netloc or not hostname:
  67. raise ValueError("base_url must use http or https and include a host")
  68. if parsed.username is not None or parsed.password is not None or "@" in parsed.netloc:
  69. raise ValueError("base_url must not contain userinfo")
  70. if "?" in value:
  71. raise ValueError("base_url must not contain a query")
  72. if "#" in value:
  73. raise ValueError("base_url must not contain a fragment")
  74. try:
  75. transport_url = httpx.URL(value)
  76. _ = (transport_url.host, transport_url.port, parsed.port)
  77. except (httpx.InvalidURL, ValueError) as exc:
  78. raise ValueError("base_url must be a valid HTTP(S) URL") from exc
  79. return value
  80. def _validate_model(value: str) -> str:
  81. stripped = value.strip()
  82. if not stripped:
  83. raise ValueError("model must not be empty")
  84. return stripped
  85. class BenchmarkTarget(_StrictBenchmarkModel):
  86. base_url: str
  87. model: str
  88. _base_url_validator = field_validator("base_url")(_validate_base_url)
  89. _model_validator = field_validator("model")(_validate_model)
  90. class BenchmarkConfig(BenchmarkTarget):
  91. schema_version: Literal[1]
  92. runs_per_case: int = Field(default=1, ge=1)
  93. modes: list[Annotated[BenchmarkMode, Field(strict=False)]] = Field(
  94. default_factory=lambda: list(BENCHMARK_MODES),
  95. min_length=1,
  96. )
  97. cases: list[Annotated[BenchmarkCaseId, Field(strict=False)]] = Field(
  98. default_factory=lambda: list(BENCHMARK_CASE_IDS),
  99. min_length=1,
  100. )
  101. @field_validator("schema_version", mode="before")
  102. @classmethod
  103. def require_integer_schema_version(cls, value: object) -> object:
  104. if type(value) is not int or value != 1:
  105. raise ValueError("schema_version must be the integer 1")
  106. return value
  107. @field_validator("modes", "cases")
  108. @classmethod
  109. def reject_duplicates(
  110. cls,
  111. values: list[BenchmarkMode | BenchmarkCaseId],
  112. ) -> list[BenchmarkMode | BenchmarkCaseId]:
  113. if len(values) != len(set(values)):
  114. raise ValueError("values must not contain duplicates")
  115. return values
  116. class BenchmarkModelCallTiming(_StrictBenchmarkModel):
  117. call_index: int = Field(ge=1)
  118. call_kind: Literal["chat_completion", "argument_fallback"]
  119. first_item_kind: str | None
  120. provider_ttft_ms: int | None = Field(ge=0)
  121. visible_ttft_ms: int | None = Field(ge=0)
  122. elapsed_ms: int | None = Field(ge=0)
  123. usage: TokenUsage | None
  124. class BenchmarkRunResult(_StrictBenchmarkModel):
  125. case_id: BenchmarkCaseId
  126. mode: BenchmarkMode
  127. iteration: int = Field(ge=1)
  128. status: Literal["passed", "failed"]
  129. initial_provider_ttft_ms: int | None = Field(ge=0)
  130. visible_ttft_ms: int | None = Field(ge=0)
  131. turn_wall_time_ms: int | None = Field(ge=0)
  132. prompt_tokens: int | None = Field(ge=0)
  133. completion_tokens: int | None = Field(ge=0)
  134. total_tokens: int | None = Field(ge=0)
  135. cached_tokens: int | None = Field(ge=0)
  136. model_call_count: int | None = Field(ge=0)
  137. fallback_count: int | None = Field(ge=0)
  138. tool_count: int | None = Field(ge=0)
  139. event_names: list[str]
  140. event_sources: list[str]
  141. tool_statuses: list[str]
  142. tool_latencies_ms: list[int | None]
  143. semantic_failures: list[str]
  144. error: str | None
  145. model_calls: list[BenchmarkModelCallTiming]
  146. class TimingChatClient:
  147. def __init__(
  148. self,
  149. inner: object,
  150. *,
  151. clock: Callable[[], float] = time.perf_counter,
  152. ) -> None:
  153. self.inner = inner
  154. self.clock = clock
  155. self.timings: list[BenchmarkModelCallTiming] = []
  156. self._next_call_index = 0
  157. async def stream_chat(
  158. self,
  159. messages: list[ChatMessage],
  160. tools: list[dict[str, Any]],
  161. params: AgentParams,
  162. tool_choice: dict[str, Any] | None = None,
  163. ) -> AsyncIterator[StreamItem]:
  164. self._next_call_index += 1
  165. call_index = self._next_call_index
  166. started_at = self.clock()
  167. first_item_at: float | None = None
  168. first_item_kind: str | None = None
  169. first_message_delta_at: float | None = None
  170. usage: TokenUsage | None = None
  171. try:
  172. stream_chat = getattr(self.inner, "stream_chat")
  173. async for item in stream_chat(
  174. messages=messages,
  175. tools=tools,
  176. params=params,
  177. tool_choice=tool_choice,
  178. ):
  179. item_at = self.clock()
  180. if first_item_at is None:
  181. first_item_at = item_at
  182. first_item_kind = item.kind
  183. if (
  184. item.kind == "message_delta"
  185. and (item.content or "").strip()
  186. and first_message_delta_at is None
  187. ):
  188. first_message_delta_at = item_at
  189. if item.kind == "usage" and item.usage is not None:
  190. usage = item.usage
  191. yield item
  192. finally:
  193. timing = BenchmarkModelCallTiming(
  194. call_index=call_index,
  195. call_kind=(
  196. "argument_fallback"
  197. if tool_choice is not None
  198. else "chat_completion"
  199. ),
  200. first_item_kind=first_item_kind,
  201. provider_ttft_ms=self._elapsed_ms(started_at, first_item_at),
  202. visible_ttft_ms=self._elapsed_ms(
  203. started_at,
  204. first_message_delta_at,
  205. ),
  206. elapsed_ms=self._elapsed_ms(started_at, self.clock()),
  207. usage=usage,
  208. )
  209. self.timings.append(timing)
  210. self.timings.sort(key=lambda item: item.call_index)
  211. async def aclose(self) -> None:
  212. close = getattr(self.inner, "aclose", None)
  213. if close is not None:
  214. await close()
  215. @staticmethod
  216. def _elapsed_ms(started_at: float, ended_at: float | None) -> int | None:
  217. if ended_at is None:
  218. return None
  219. return max(0, round((ended_at - started_at) * 1000))
  220. class BenchmarkExpectation(_FrozenStrictBenchmarkModel):
  221. visible_messages: tuple[str, ...]
  222. event_names: tuple[str, ...]
  223. answer_count: int = Field(ge=1)
  224. expected_model_call_count: int = Field(ge=1)
  225. expected_fallback_count: int = Field(ge=0)
  226. terminal: bool
  227. first_reply_before_tool: bool
  228. @model_validator(mode="after")
  229. def match_answer_count(self) -> "BenchmarkExpectation":
  230. if self.answer_count != len(self.visible_messages):
  231. raise ValueError("answer_count must match visible_messages")
  232. return self
  233. class BenchmarkMockEvent(_FrozenStrictBenchmarkModel):
  234. id: str
  235. name: str
  236. arguments: tuple[tuple[str, JsonScalar], ...]
  237. @model_validator(mode="after")
  238. def validate_event(self) -> "BenchmarkMockEvent":
  239. if not self.id.strip():
  240. raise ValueError("mock event id must not be empty")
  241. if not self.name.strip():
  242. raise ValueError("mock event name must not be empty")
  243. keys = tuple(key for key, _ in self.arguments)
  244. if any(not key.strip() for key in keys):
  245. raise ValueError("mock event argument names must not be empty")
  246. if len(keys) != len(set(keys)):
  247. raise ValueError("mock event arguments must not contain duplicate names")
  248. return self
  249. class BenchmarkCase(_FrozenStrictBenchmarkModel):
  250. case_id: BenchmarkCaseId
  251. user_message: str
  252. enabled_tools: tuple[str, ...]
  253. expectation: BenchmarkExpectation
  254. mock_visible_messages: tuple[str, ...]
  255. mock_events: tuple[BenchmarkMockEvent, ...]
  256. @model_validator(mode="after")
  257. def validate_case_semantics(self) -> "BenchmarkCase":
  258. if not self.user_message.strip():
  259. raise ValueError("user_message must not be empty")
  260. if len(self.enabled_tools) != len(set(self.enabled_tools)):
  261. raise ValueError("enabled_tools must not contain duplicates")
  262. mock_event_names = tuple(event.name for event in self.mock_events)
  263. if mock_event_names != self.expectation.event_names:
  264. raise ValueError("mock event order must match expected event order")
  265. if self.expectation.first_reply_before_tool:
  266. if not self.mock_events or not self.mock_visible_messages:
  267. raise ValueError(
  268. "first_reply_before_tool requires mock text and at least one event"
  269. )
  270. return self
  271. def _mock_event(
  272. event_id: str,
  273. name: str,
  274. arguments: tuple[tuple[str, JsonScalar], ...] = (),
  275. ) -> BenchmarkMockEvent:
  276. return BenchmarkMockEvent(id=event_id, name=name, arguments=arguments)
  277. _SCHEDULE_ARGUMENTS: tuple[tuple[str, JsonScalar], ...] = (
  278. ("title", "Design review"),
  279. ("start_at", "2026-07-14T09:30:00+08:00"),
  280. ("timezone", "Asia/Shanghai"),
  281. )
  282. _SCHEDULE_CONFIRMATION = (
  283. "Scheduled Design review for 2026-07-14T09:30:00+08:00 "
  284. "(Asia/Shanghai)."
  285. )
  286. BENCHMARK_CASE_CATALOG: Mapping[BenchmarkCaseId, BenchmarkCase] = MappingProxyType(
  287. {
  288. BenchmarkCaseId.ORDINARY_CHAT: BenchmarkCase(
  289. case_id=BenchmarkCaseId.ORDINARY_CHAT,
  290. user_message="Hello.",
  291. enabled_tools=(
  292. "session.terminate",
  293. "device.volume.adjust",
  294. "calendar.schedule.create",
  295. "knowledge.web.search",
  296. ),
  297. expectation=BenchmarkExpectation(
  298. visible_messages=("Ordinary answer.",),
  299. event_names=(),
  300. answer_count=1,
  301. expected_model_call_count=1,
  302. expected_fallback_count=0,
  303. terminal=False,
  304. first_reply_before_tool=False,
  305. ),
  306. mock_visible_messages=("Ordinary answer.",),
  307. mock_events=(),
  308. ),
  309. BenchmarkCaseId.DEVICE_VOLUME_SILENT: BenchmarkCase(
  310. case_id=BenchmarkCaseId.DEVICE_VOLUME_SILENT,
  311. user_message="set volume to 30",
  312. enabled_tools=("device.volume.adjust",),
  313. expectation=BenchmarkExpectation(
  314. visible_messages=("I will set the volume to 30.",),
  315. event_names=("device.volume.adjust",),
  316. answer_count=1,
  317. expected_model_call_count=1,
  318. expected_fallback_count=0,
  319. terminal=False,
  320. first_reply_before_tool=True,
  321. ),
  322. mock_visible_messages=("I will set the volume to 30.",),
  323. mock_events=(
  324. _mock_event(
  325. "volume-1",
  326. "device.volume.adjust",
  327. (("mode", "absolute"), ("value", 30)),
  328. ),
  329. ),
  330. ),
  331. BenchmarkCaseId.CALENDAR_SCHEDULE_TEMPLATE: BenchmarkCase(
  332. case_id=BenchmarkCaseId.CALENDAR_SCHEDULE_TEMPLATE,
  333. user_message=(
  334. 'schedule "Design review" at 2026-07-14T09:30:00+08:00 '
  335. "timezone Asia/Shanghai"
  336. ),
  337. enabled_tools=("calendar.schedule.create",),
  338. expectation=BenchmarkExpectation(
  339. visible_messages=("I will schedule that.", _SCHEDULE_CONFIRMATION),
  340. event_names=("calendar.schedule.create",),
  341. answer_count=2,
  342. expected_model_call_count=1,
  343. expected_fallback_count=0,
  344. terminal=False,
  345. first_reply_before_tool=True,
  346. ),
  347. mock_visible_messages=("I will schedule that.",),
  348. mock_events=(
  349. _mock_event(
  350. "schedule-1",
  351. "calendar.schedule.create",
  352. _SCHEDULE_ARGUMENTS,
  353. ),
  354. ),
  355. ),
  356. BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS: BenchmarkCase(
  357. case_id=BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS,
  358. user_message="event batch safety",
  359. enabled_tools=("knowledge.web.search",),
  360. expectation=BenchmarkExpectation(
  361. visible_messages=("I will check.", "Grounded search update."),
  362. event_names=("knowledge.web.search",),
  363. answer_count=2,
  364. expected_model_call_count=2,
  365. expected_fallback_count=0,
  366. terminal=False,
  367. first_reply_before_tool=True,
  368. ),
  369. mock_visible_messages=("I will check.", "Grounded search update."),
  370. mock_events=(
  371. _mock_event(
  372. "search-1",
  373. "knowledge.web.search",
  374. (("query", "event batch safety"),),
  375. ),
  376. ),
  377. ),
  378. BenchmarkCaseId.SESSION_TERMINATE: BenchmarkCase(
  379. case_id=BenchmarkCaseId.SESSION_TERMINATE,
  380. user_message="end this session",
  381. enabled_tools=("session.terminate",),
  382. expectation=BenchmarkExpectation(
  383. visible_messages=("Goodbye.",),
  384. event_names=("session.terminate",),
  385. answer_count=1,
  386. expected_model_call_count=1,
  387. expected_fallback_count=0,
  388. terminal=True,
  389. first_reply_before_tool=True,
  390. ),
  391. mock_visible_messages=("Goodbye.",),
  392. mock_events=(_mock_event("terminate-1", "session.terminate"),),
  393. ),
  394. BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE: BenchmarkCase(
  395. case_id=BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE,
  396. user_message=(
  397. 'set volume to 30 and schedule "Design review" at '
  398. "2026-07-14T09:30:00+08:00 timezone Asia/Shanghai"
  399. ),
  400. enabled_tools=(
  401. "device.volume.adjust",
  402. "calendar.schedule.create",
  403. ),
  404. expectation=BenchmarkExpectation(
  405. visible_messages=(
  406. "I will set the volume to 30 and schedule that.",
  407. _SCHEDULE_CONFIRMATION,
  408. ),
  409. event_names=(
  410. "device.volume.adjust",
  411. "calendar.schedule.create",
  412. ),
  413. answer_count=2,
  414. expected_model_call_count=1,
  415. expected_fallback_count=0,
  416. terminal=False,
  417. first_reply_before_tool=True,
  418. ),
  419. mock_visible_messages=(
  420. "I will set the volume to 30 and schedule that.",
  421. ),
  422. mock_events=(
  423. _mock_event(
  424. "parallel-volume",
  425. "device.volume.adjust",
  426. (("mode", "absolute"), ("value", 30)),
  427. ),
  428. _mock_event(
  429. "parallel-schedule",
  430. "calendar.schedule.create",
  431. _SCHEDULE_ARGUMENTS,
  432. ),
  433. ),
  434. ),
  435. }
  436. )
  437. def _normalize_benchmark_mode(mode: BenchmarkMode | str) -> BenchmarkMode:
  438. try:
  439. return BenchmarkMode(mode)
  440. except ValueError as exc:
  441. raise ValueError(f"unsupported benchmark mode: {mode}") from exc
  442. def build_benchmark_request(
  443. case: BenchmarkCase,
  444. mode: BenchmarkMode | str,
  445. model: str,
  446. ) -> DebugRunRequest:
  447. normalized_mode = _normalize_benchmark_mode(mode)
  448. validated_model = _validate_model(model)
  449. return DebugRunRequest(
  450. user_message=case.user_message,
  451. system_prompts=[BENCHMARK_SYSTEM_PROMPT],
  452. pre_messages=[],
  453. chat_agent=AgentParams(
  454. model=validated_model,
  455. temperature=0.0,
  456. max_tokens=BENCHMARK_MAX_TOKENS,
  457. ),
  458. event_agent=EventAgentParams(
  459. model=validated_model,
  460. temperature=0.0,
  461. max_tokens=BENCHMARK_MAX_TOKENS,
  462. enabled_tools=list(case.enabled_tools),
  463. max_event_loops=BENCHMARK_MAX_EVENT_LOOPS,
  464. max_parallel_events=BENCHMARK_MAX_PARALLEL_EVENTS,
  465. batch_timeout_seconds=BENCHMARK_BATCH_TIMEOUT_SECONDS,
  466. ),
  467. tool_invocation_mode=normalized_mode.value,
  468. )
  469. def build_mock_rounds(
  470. case: BenchmarkCase,
  471. mode: BenchmarkMode | str,
  472. ) -> list[list[StreamItem]]:
  473. normalized_mode = _normalize_benchmark_mode(mode)
  474. first_round: list[StreamItem] = []
  475. if case.mock_visible_messages:
  476. first_round.append(StreamItem.message_delta(case.mock_visible_messages[0]))
  477. for mock_event in case.mock_events:
  478. arguments = dict(mock_event.arguments)
  479. event = ToolCallEvent(
  480. id=mock_event.id,
  481. name=mock_event.name,
  482. arguments=arguments,
  483. raw_arguments=json.dumps(
  484. arguments,
  485. sort_keys=True,
  486. separators=(",", ":"),
  487. ),
  488. )
  489. event_item = (
  490. StreamItem.text_event(event)
  491. if normalized_mode is BenchmarkMode.DUAL_AGENT
  492. else StreamItem.provider_tool_call(event)
  493. )
  494. first_round.append(event_item)
  495. rounds = [first_round]
  496. rounds.extend(
  497. [StreamItem.message_delta(message)]
  498. for message in case.mock_visible_messages[1:]
  499. )
  500. return rounds
  501. class MockBenchmarkChatClient:
  502. def __init__(
  503. self,
  504. case: BenchmarkCase,
  505. mode: BenchmarkMode | str,
  506. ) -> None:
  507. self.case = case
  508. self.mode = _normalize_benchmark_mode(mode)
  509. self.rounds = build_mock_rounds(case, self.mode)
  510. self.calls = 0
  511. async def stream_chat(
  512. self,
  513. messages: list[ChatMessage],
  514. tools: list[dict[str, Any]],
  515. params: AgentParams,
  516. tool_choice: dict[str, Any] | None = None,
  517. ) -> AsyncIterator[StreamItem]:
  518. del messages, tools, params, tool_choice
  519. call_index = self.calls
  520. if call_index >= len(self.rounds):
  521. raise RuntimeError(
  522. f"mock benchmark stream exhausted: {self.case.case_id}"
  523. )
  524. self.calls += 1
  525. yield StreamItem.raw_response_chunk(
  526. {
  527. "mock": True,
  528. "case_id": self.case.case_id.value,
  529. "mode": self.mode.value,
  530. "call_index": self.calls,
  531. }
  532. )
  533. for item in self.rounds[call_index]:
  534. yield item
  535. yield StreamItem.usage_item(
  536. TokenUsage(
  537. prompt_tokens=10,
  538. completion_tokens=5,
  539. total_tokens=15,
  540. cached_tokens=2,
  541. )
  542. )
  543. async def aclose(self) -> None:
  544. return None
  545. BenchmarkClientFactory: TypeAlias = Callable[[BenchmarkConfig, str], object]
  546. class _BenchmarkOutputQueue(asyncio.Queue[dict[str, Any]]):
  547. def __init__(self, clock: Callable[[], float]) -> None:
  548. super().__init__()
  549. self.clock = clock
  550. self.first_message_delta_at: float | None = None
  551. async def put(self, item: dict[str, Any]) -> None:
  552. content = item.get("content")
  553. if (
  554. item.get("type") == "message_delta"
  555. and isinstance(content, str)
  556. and content.strip()
  557. and self.first_message_delta_at is None
  558. ):
  559. self.first_message_delta_at = self.clock()
  560. await super().put(item)
  561. class BenchmarkRunner:
  562. def __init__(
  563. self,
  564. config: BenchmarkConfig,
  565. api_key: str | None,
  566. mock: bool = False,
  567. client_factory: BenchmarkClientFactory | None = None,
  568. clock: Callable[[], float] = time.perf_counter,
  569. ) -> None:
  570. if not mock and not (api_key or "").strip():
  571. raise ValueError("api_key is required for live benchmark runs")
  572. self.config = config
  573. self.api_key = (api_key or "").strip()
  574. self.mock = mock
  575. self.client_factory = client_factory or self._default_client_factory
  576. self.clock = clock
  577. async def run(self) -> list[BenchmarkRunResult]:
  578. results: list[BenchmarkRunResult] = []
  579. for case_id in self.config.cases:
  580. case = BENCHMARK_CASE_CATALOG[case_id]
  581. for mode in self.config.modes:
  582. for iteration in range(1, self.config.runs_per_case + 1):
  583. results.append(await self._run_one(case, mode, iteration))
  584. return results
  585. async def _run_one(
  586. self,
  587. case: BenchmarkCase,
  588. mode: BenchmarkMode,
  589. iteration: int,
  590. ) -> BenchmarkRunResult:
  591. store = SQLiteSessionStore(":memory:")
  592. runtime: DebugRuntime | None = None
  593. timing_client: TimingChatClient | None = None
  594. result: BenchmarkRunResult | None = None
  595. close_error: Exception | None = None
  596. try:
  597. inner_client = (
  598. MockBenchmarkChatClient(case, mode)
  599. if self.mock
  600. else self.client_factory(self.config, self.api_key)
  601. )
  602. timing_client = TimingChatClient(inner_client, clock=self.clock)
  603. session_id = f"benchmark-{uuid4().hex}"
  604. request = build_benchmark_request(
  605. case,
  606. mode,
  607. self.config.model,
  608. ).model_copy(update={"session_id": session_id})
  609. output_queue = _BenchmarkOutputQueue(self.clock)
  610. runtime = DebugRuntime(
  611. timing_client,
  612. queues=RuntimeQueues(output=output_queue),
  613. registry=build_default_tool_registry(),
  614. session_store=store,
  615. clock=self.clock,
  616. )
  617. turn_started_at = self.clock()
  618. outputs, visible_ttft_ms, runtime_error = await self._consume_turn(
  619. runtime,
  620. request,
  621. turn_started_at,
  622. output_queue,
  623. )
  624. audits = store.list_audit_logs(session_id)
  625. usage = store.usage_summary(session_id)
  626. result = self._build_result(
  627. case=case,
  628. mode=mode,
  629. iteration=iteration,
  630. outputs=outputs,
  631. audits=audits,
  632. usage=usage,
  633. timings=timing_client.timings,
  634. visible_ttft_ms=visible_ttft_ms,
  635. runtime_error=runtime_error,
  636. )
  637. except Exception as exc:
  638. result = self._failed_result(
  639. case.case_id,
  640. mode,
  641. iteration,
  642. error=str(exc),
  643. timings=timing_client.timings if timing_client is not None else [],
  644. )
  645. finally:
  646. if runtime is not None:
  647. try:
  648. await runtime.aclose()
  649. except Exception as exc:
  650. close_error = exc
  651. if timing_client is not None:
  652. try:
  653. await timing_client.aclose()
  654. except Exception as exc:
  655. close_error = close_error or exc
  656. assert result is not None
  657. if close_error is not None:
  658. error = str(close_error)
  659. if result.error:
  660. error = f"{result.error}; close failed: {error}"
  661. result = result.model_copy(update={"status": "failed", "error": error})
  662. return result
  663. async def _consume_turn(
  664. self,
  665. runtime: DebugRuntime,
  666. request: DebugRunRequest,
  667. turn_started_at: float,
  668. output_queue: _BenchmarkOutputQueue,
  669. ) -> tuple[list[dict[str, Any]], int | None, str | None]:
  670. queues = runtime.start_session(request)
  671. outputs: list[dict[str, Any]] = []
  672. visible_ttft_ms: int | None = None
  673. runtime_error: str | None = None
  674. while True:
  675. message = await queues.output.get()
  676. outputs.append(message)
  677. message_type = message.get("type")
  678. if message_type == "error":
  679. runtime_error = str(message.get("message") or "runtime error")
  680. break
  681. if message_type == "turn_completed":
  682. break
  683. if output_queue.first_message_delta_at is not None:
  684. visible_ttft_ms = max(
  685. 0,
  686. round(
  687. (output_queue.first_message_delta_at - turn_started_at) * 1000
  688. ),
  689. )
  690. return outputs, visible_ttft_ms, runtime_error
  691. def _build_result(
  692. self,
  693. *,
  694. case: BenchmarkCase,
  695. mode: BenchmarkMode,
  696. iteration: int,
  697. outputs: list[dict[str, Any]],
  698. audits: list[dict[str, Any]],
  699. usage: dict[str, Any],
  700. timings: list[BenchmarkModelCallTiming],
  701. visible_ttft_ms: int | None,
  702. runtime_error: str | None,
  703. ) -> BenchmarkRunResult:
  704. session_usage = usage["session"]
  705. calls = usage["calls"]
  706. event_audits = [
  707. audit for audit in audits if audit["event"] == "chat_event_detected"
  708. ]
  709. batch_audits = [
  710. audit for audit in audits if audit["event"] == "event_batch_results"
  711. ]
  712. batch_results = [
  713. item
  714. for audit in batch_audits
  715. for item in audit["details"].get("results", [])
  716. ]
  717. event_names = [audit["details"]["event_name"] for audit in event_audits]
  718. event_sources = [audit["details"]["event_source"] for audit in event_audits]
  719. tool_statuses = [str(item["status"]) for item in batch_results]
  720. tool_calls = [call for call in calls if call["call_kind"] == "tool_execution"]
  721. model_call_count = sum(
  722. call["call_kind"] == "chat_completion" for call in calls
  723. )
  724. fallback_count = int(session_usage["fallback_count"])
  725. tool_count = int(session_usage["tool_count"])
  726. semantic_failures = self._semantic_failures(
  727. case=case,
  728. mode=mode,
  729. outputs=outputs,
  730. audits=audits,
  731. event_names=event_names,
  732. event_sources=event_sources,
  733. tool_statuses=tool_statuses,
  734. model_call_count=model_call_count,
  735. fallback_count=fallback_count,
  736. tool_count=tool_count,
  737. timings=timings,
  738. check_mock_text=self.mock,
  739. )
  740. initial_call = next(
  741. (timing for timing in timings if timing.call_kind == "chat_completion"),
  742. None,
  743. )
  744. status = "failed" if runtime_error or semantic_failures else "passed"
  745. return BenchmarkRunResult(
  746. case_id=case.case_id,
  747. mode=mode,
  748. iteration=iteration,
  749. status=status,
  750. initial_provider_ttft_ms=(
  751. initial_call.provider_ttft_ms if initial_call is not None else None
  752. ),
  753. visible_ttft_ms=visible_ttft_ms,
  754. turn_wall_time_ms=session_usage["turn_wall_time_ms"],
  755. prompt_tokens=int(session_usage["prompt_tokens"]),
  756. completion_tokens=int(session_usage["completion_tokens"]),
  757. total_tokens=int(session_usage["total_tokens"]),
  758. cached_tokens=int(session_usage["cached_tokens"]),
  759. model_call_count=model_call_count,
  760. fallback_count=fallback_count,
  761. tool_count=tool_count,
  762. event_names=event_names,
  763. event_sources=event_sources,
  764. tool_statuses=tool_statuses,
  765. tool_latencies_ms=[call["tool_latency_ms"] for call in tool_calls],
  766. semantic_failures=semantic_failures,
  767. error=runtime_error,
  768. model_calls=list(timings),
  769. )
  770. def _semantic_failures(
  771. self,
  772. *,
  773. case: BenchmarkCase,
  774. mode: BenchmarkMode,
  775. outputs: list[dict[str, Any]],
  776. audits: list[dict[str, Any]],
  777. event_names: list[str],
  778. event_sources: list[str],
  779. tool_statuses: list[str],
  780. model_call_count: int,
  781. fallback_count: int,
  782. tool_count: int,
  783. timings: list[BenchmarkModelCallTiming],
  784. check_mock_text: bool,
  785. ) -> list[str]:
  786. failures: list[str] = []
  787. answers, first_answer_index, first_tool_index = self._visible_answers(outputs)
  788. expected = case.expectation
  789. if len(answers) != expected.answer_count:
  790. failures.append(
  791. f"answer_count expected {expected.answer_count}, got {len(answers)}"
  792. )
  793. if any(not answer.strip() for answer in answers):
  794. failures.append("visible answers must be non-empty")
  795. if check_mock_text and tuple(answers) != expected.visible_messages:
  796. failures.append(
  797. f"mock visible messages expected {expected.visible_messages!r}, "
  798. f"got {tuple(answers)!r}"
  799. )
  800. if expected.first_reply_before_tool and not (
  801. first_answer_index is not None
  802. and first_tool_index is not None
  803. and first_answer_index < first_tool_index
  804. ):
  805. failures.append("first visible answer must precede tool results")
  806. if tuple(event_names) != expected.event_names:
  807. failures.append(
  808. f"event names expected {expected.event_names!r}, got {tuple(event_names)!r}"
  809. )
  810. expected_source = (
  811. "text_event"
  812. if mode is BenchmarkMode.DUAL_AGENT
  813. else "provider_resolved"
  814. )
  815. expected_sources = (expected_source,) * len(expected.event_names)
  816. if tuple(event_sources) != expected_sources:
  817. failures.append(
  818. f"event sources expected {expected_sources!r}, "
  819. f"got {tuple(event_sources)!r}"
  820. )
  821. expected_statuses = ("success",) * len(expected.event_names)
  822. if tuple(tool_statuses) != expected_statuses:
  823. failures.append(
  824. f"tool statuses expected {expected_statuses!r}, "
  825. f"got {tuple(tool_statuses)!r}"
  826. )
  827. if tool_count != len(expected.event_names):
  828. failures.append(
  829. f"tool_count expected {len(expected.event_names)}, got {tool_count}"
  830. )
  831. timing_chat_count = sum(
  832. timing.call_kind == "chat_completion" for timing in timings
  833. )
  834. timing_fallback_count = sum(
  835. timing.call_kind == "argument_fallback" for timing in timings
  836. )
  837. if model_call_count != expected.expected_model_call_count:
  838. failures.append(
  839. "model_call_count expected "
  840. f"{expected.expected_model_call_count}, ledger got {model_call_count}"
  841. )
  842. if timing_chat_count != expected.expected_model_call_count:
  843. failures.append(
  844. "model_call_count expected "
  845. f"{expected.expected_model_call_count}, timing got {timing_chat_count}"
  846. )
  847. if fallback_count != expected.expected_fallback_count:
  848. failures.append(
  849. "fallback_count expected "
  850. f"{expected.expected_fallback_count}, ledger got {fallback_count}"
  851. )
  852. if timing_fallback_count != expected.expected_fallback_count:
  853. failures.append(
  854. "fallback_count expected "
  855. f"{expected.expected_fallback_count}, timing got "
  856. f"{timing_fallback_count}"
  857. )
  858. if model_call_count != timing_chat_count:
  859. failures.append(
  860. f"model_call_count ledger={model_call_count}, timing={timing_chat_count}"
  861. )
  862. if fallback_count != timing_fallback_count:
  863. failures.append(
  864. f"fallback_count ledger={fallback_count}, timing={timing_fallback_count}"
  865. )
  866. terminal = any(audit["event"] == "terminal_completed" for audit in audits)
  867. if terminal is not expected.terminal:
  868. failures.append(
  869. f"terminal expected {expected.terminal}, got {terminal}"
  870. )
  871. return failures
  872. @staticmethod
  873. def _visible_answers(
  874. outputs: list[dict[str, Any]],
  875. ) -> tuple[list[str], int | None, int | None]:
  876. answers: list[str] = []
  877. current: list[str] = []
  878. first_answer_index: int | None = None
  879. first_tool_index: int | None = None
  880. for index, message in enumerate(outputs):
  881. message_type = message.get("type")
  882. if message_type == "message_delta":
  883. content = str(message.get("content") or "")
  884. if first_answer_index is None and content.strip():
  885. first_answer_index = index
  886. current.append(content)
  887. continue
  888. if message_type == "tool_result":
  889. if first_tool_index is None:
  890. first_tool_index = index
  891. answer = "".join(current).strip()
  892. if answer:
  893. answers.append(answer)
  894. current = []
  895. answer = "".join(current).strip()
  896. if answer:
  897. answers.append(answer)
  898. return answers, first_answer_index, first_tool_index
  899. @staticmethod
  900. def _failed_result(
  901. case_id: BenchmarkCaseId,
  902. mode: BenchmarkMode,
  903. iteration: int,
  904. *,
  905. error: str,
  906. timings: list[BenchmarkModelCallTiming],
  907. ) -> BenchmarkRunResult:
  908. initial_call = next(
  909. (timing for timing in timings if timing.call_kind == "chat_completion"),
  910. None,
  911. )
  912. return BenchmarkRunResult(
  913. case_id=case_id,
  914. mode=mode,
  915. iteration=iteration,
  916. status="failed",
  917. initial_provider_ttft_ms=(
  918. initial_call.provider_ttft_ms if initial_call is not None else None
  919. ),
  920. visible_ttft_ms=None,
  921. turn_wall_time_ms=None,
  922. prompt_tokens=None,
  923. completion_tokens=None,
  924. total_tokens=None,
  925. cached_tokens=None,
  926. model_call_count=None,
  927. fallback_count=None,
  928. tool_count=None,
  929. event_names=[],
  930. event_sources=[],
  931. tool_statuses=[],
  932. tool_latencies_ms=[],
  933. semantic_failures=[],
  934. error=error,
  935. model_calls=list(timings),
  936. )
  937. @staticmethod
  938. def _default_client_factory(
  939. config: BenchmarkConfig,
  940. api_key: str,
  941. ) -> OpenAICompatibleChatClient:
  942. return OpenAICompatibleChatClient(
  943. api_key=api_key,
  944. base_url=config.base_url,
  945. default_model=config.model,
  946. request_timeout_seconds=60.0,
  947. )