From 613d167e81b8ac9827e4b68f6b58687cd4b10669 Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 2 Jul 2026 14:36:59 +0800 Subject: [PATCH] feat(token-tracking): add HTTP response hook and attach_usage_hook, wire into build_models Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../plans/2026-07-02-token-usage-tracking.md | 1995 +++++++++++++++++ rag_eval/metrics/factory.py | 47 + tests/test_token_usage_hook.py | 75 + 3 files changed, 2117 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-02-token-usage-tracking.md create mode 100644 tests/test_token_usage_hook.py diff --git a/docs/superpowers/plans/2026-07-02-token-usage-tracking.md b/docs/superpowers/plans/2026-07-02-token-usage-tracking.md new file mode 100644 index 0000000..448e9d8 --- /dev/null +++ b/docs/superpowers/plans/2026-07-02-token-usage-tracking.md @@ -0,0 +1,1995 @@ +# Judge LLM Token 用量追踪 Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 按 LLM 模型名累计记录 RAGAS 评分调用与优化顾问 LLM 分析调用的原始 token 用量(input/output/调用次数,不换算金额),持久化到 `metadata.json`/`summary.md`,并在 Web 报告详情页展示。 + +**Architecture:** 在 HTTP 层(httpx 响应钩子)拦截所有 `AsyncOpenAI` 请求的响应体,读取其中的 `usage` + `model` 字段。用 `contextvars.ContextVar` 保存"当前活跃的统计器",因为 `InlineScorer` 会跨请求缓存 `AsyncOpenAI` 客户端,不能把统计器固定绑死在客户端构造时刻。三条落盘路径(CLI scenario 运行、`/api/score/async`、`/api/score/session_async`)各自在评分调用外包一层 `track_token_usage()`,把汇总结果写进 `EvaluationResult.token_usage` → `metadata.json`。 + +**Tech Stack:** Python 3.12、httpx(openai SDK 的传递依赖,已验证 0.28.1 可用)、pytest + unittest(现有测试框架混用)、pandas(既有报告聚合逻辑)。 + +## Global Constraints + +- 只记录原始 token 数(input_tokens / output_tokens / calls),**不做金额换算**,不改 LLM Profile 结构。 +- 只统计**落盘成 run 目录**的路径:CLI `main.py --scenario`、`/api/score/async`、`/api/score/session_async`。**不**统计不落盘的同步 `/api/score`。 +- 统计范围覆盖:RAGAS 评分调用(judge model + embedding model)+ 优化顾问 LLM 分析调用(`advisor/llm_analyzer.py`)。**不**统计 `dataset_builder` 题库生成路径。 +- `session_async` 场景下 token 用量必须**跨调用累加**,不能每次覆盖为最后一次的值。 +- HTTP 钩子解析失败、网关不返回 `usage`/`model` 字段时必须**静默跳过**,绝不能让评分调用报错或变慢。 +- 已核实技术事实:`openai==1.102.0` 的 `AsyncOpenAI()._client` 是 `httpx.AsyncClient` 子类(`AsyncHttpxClientWrapper`),自带 `event_hooks = {"request": [], "response": []}` 可直接追加钩子。 +- 设计依据:`docs/superpowers/specs/2026-07-02-token-usage-tracking-design.md`(本计划的每个任务都对应该文档的某一节,不要偏离)。 +- 测试约定:本仓库用 `asyncio.run(...)` 包裹被测异步调用,而不是 `@pytest.mark.asyncio`(参考 `tests/test_advisor_llm_analyzer.py`)。运行测试用 `C:\software\Python312\python.exe -m pytest tests/ -v`。 + +--- + +### Task 1: `TokenUsageTracker` 核心与 contextvar + +**Files:** +- Create: `rag_eval/metrics/token_tracker.py` +- Test: `tests/test_token_tracker.py` + +**Interfaces:** +- Produces: + - `class TokenUsageTracker` — `.record(model: str, input_tokens: int, output_tokens: int) -> None`;`.as_dict() -> dict[str, dict[str, int]]`;`.merge_into(existing: dict[str, dict[str, int]]) -> dict[str, dict[str, int]]` + - `track_token_usage() -> AbstractContextManager[TokenUsageTracker]`(`@contextmanager`) + - `get_current_tracker() -> TokenUsageTracker | None` + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_token_tracker.py`: + +```python +"""Tests for the per-run token usage accumulator and its context-scoped activation.""" +from __future__ import annotations + +from rag_eval.metrics.token_tracker import ( + TokenUsageTracker, + get_current_tracker, + track_token_usage, +) + + +def test_record_accumulates_input_output_and_calls(): + tracker = TokenUsageTracker() + tracker.record("gpt-5", 100, 50) + tracker.record("gpt-5", 20, 10) + assert tracker.as_dict() == { + "gpt-5": {"input_tokens": 120, "output_tokens": 60, "calls": 2} + } + + +def test_record_groups_by_model_name(): + tracker = TokenUsageTracker() + tracker.record("gpt-5", 100, 50) + tracker.record("Qwen/Qwen3-Embedding-4B", 30, 0) + result = tracker.as_dict() + assert set(result.keys()) == {"gpt-5", "Qwen/Qwen3-Embedding-4B"} + assert result["Qwen/Qwen3-Embedding-4B"] == { + "input_tokens": 30, "output_tokens": 0, "calls": 1 + } + + +def test_record_defaults_blank_model_to_unknown(): + tracker = TokenUsageTracker() + tracker.record("", 10, 5) + assert "unknown" in tracker.as_dict() + + +def test_merge_into_sums_with_existing_totals(): + tracker = TokenUsageTracker() + tracker.record("gpt-5", 100, 50) + existing = {"gpt-5": {"input_tokens": 200, "output_tokens": 100, "calls": 3}} + merged = tracker.merge_into(existing) + assert merged == {"gpt-5": {"input_tokens": 300, "output_tokens": 150, "calls": 4}} + + +def test_merge_into_keeps_models_only_in_existing(): + tracker = TokenUsageTracker() + tracker.record("gpt-5", 10, 5) + existing = {"other-model": {"input_tokens": 1, "output_tokens": 1, "calls": 1}} + merged = tracker.merge_into(existing) + assert merged["other-model"] == {"input_tokens": 1, "output_tokens": 1, "calls": 1} + assert merged["gpt-5"] == {"input_tokens": 10, "output_tokens": 5, "calls": 1} + + +def test_merge_into_does_not_mutate_existing_dict(): + tracker = TokenUsageTracker() + tracker.record("gpt-5", 10, 5) + existing = {"gpt-5": {"input_tokens": 1, "output_tokens": 1, "calls": 1}} + tracker.merge_into(existing) + assert existing == {"gpt-5": {"input_tokens": 1, "output_tokens": 1, "calls": 1}} + + +def test_get_current_tracker_returns_none_outside_context(): + assert get_current_tracker() is None + + +def test_track_token_usage_activates_and_resets_context(): + assert get_current_tracker() is None + with track_token_usage() as tracker: + assert get_current_tracker() is tracker + tracker.record("gpt-5", 1, 1) + assert get_current_tracker() is None + + +def test_track_token_usage_nested_contexts_are_isolated(): + with track_token_usage() as outer: + outer.record("outer-model", 5, 5) + with track_token_usage() as inner: + inner.record("inner-model", 1, 1) + assert get_current_tracker() is inner + assert get_current_tracker() is outer + assert outer.as_dict() == { + "outer-model": {"input_tokens": 5, "output_tokens": 5, "calls": 1} + } +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_token_tracker.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'rag_eval.metrics.token_tracker'` + +- [ ] **Step 3: Write minimal implementation** + +Create `rag_eval/metrics/token_tracker.py`: + +```python +"""Per-run token usage accumulation, keyed by model name. + +RAGAS 0.4.3's `ragas.metrics.collections` + instructor code path does not +expose real token counts (`ragas/cost.py` only serves the legacy langchain +`evaluate()` path). This module provides a context-scoped accumulator that +the HTTP response hook in `rag_eval.metrics.factory` feeds into, so token +counts survive across the AsyncOpenAI client caching used by InlineScorer +(see docs/superpowers/specs/2026-07-02-token-usage-tracking-design.md). +""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field +from typing import Iterator + + +@dataclass +class TokenUsageTracker: + """Accumulates input/output token counts and call counts, grouped by model name.""" + + _totals: dict[str, dict[str, int]] = field(default_factory=dict) + + def record(self, model: str, input_tokens: int, output_tokens: int) -> None: + """Add one API call's usage to the running total for `model`.""" + key = model or "unknown" + bucket = self._totals.setdefault( + key, {"input_tokens": 0, "output_tokens": 0, "calls": 0} + ) + bucket["input_tokens"] += int(input_tokens) + bucket["output_tokens"] += int(output_tokens) + bucket["calls"] += 1 + + def as_dict(self) -> dict[str, dict[str, int]]: + """Return a plain-dict snapshot: {model: {input_tokens, output_tokens, calls}}.""" + return {model: dict(usage) for model, usage in self._totals.items()} + + def merge_into(self, existing: dict[str, dict[str, int]]) -> dict[str, dict[str, int]]: + """Return a new dict combining `existing` accumulated totals with this tracker's totals. + + Used by session-scoped scoring (one call at a time) to keep a running + total across multiple calls instead of overwriting with just the latest call. + Does not mutate `existing`. + """ + merged: dict[str, dict[str, int]] = { + model: dict(usage) for model, usage in existing.items() + } + for model, usage in self.as_dict().items(): + bucket = merged.setdefault( + model, {"input_tokens": 0, "output_tokens": 0, "calls": 0} + ) + bucket["input_tokens"] += usage["input_tokens"] + bucket["output_tokens"] += usage["output_tokens"] + bucket["calls"] += usage["calls"] + return merged + + +_current_tracker: ContextVar[TokenUsageTracker | None] = ContextVar( + "_current_tracker", default=None +) + + +@contextmanager +def track_token_usage() -> Iterator[TokenUsageTracker]: + """Activate a fresh TokenUsageTracker for the duration of the `with` block. + + Any AsyncOpenAI client with `attach_usage_hook()` applied that makes a + call while this context is active will have its usage recorded here. + Safe to nest; the innermost tracker is active within its own block. + """ + tracker = TokenUsageTracker() + token = _current_tracker.set(tracker) + try: + yield tracker + finally: + _current_tracker.reset(token) + + +def get_current_tracker() -> TokenUsageTracker | None: + """Return the currently active tracker, or None if no `track_token_usage()` block is active.""" + return _current_tracker.get() +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_token_tracker.py -v` +Expected: `10 passed` + +- [ ] **Step 5: Commit** + +```powershell +git add rag_eval/metrics/token_tracker.py tests/test_token_tracker.py +git commit -m "feat(token-tracking): add TokenUsageTracker with context-scoped activation" +``` + +--- + +### Task 2: HTTP 响应钩子 + `attach_usage_hook()`,接入 `build_models()` + +**Files:** +- Modify: `rag_eval/metrics/factory.py` +- Test: `tests/test_token_usage_hook.py` + +**Interfaces:** +- Consumes: `get_current_tracker()`, `TokenUsageTracker`, `track_token_usage()` from Task 1 (`rag_eval.metrics.token_tracker`) +- Produces: + - `attach_usage_hook(client: AsyncOpenAI) -> None` — idempotent; other tasks (3, and indirectly 6/7/8 via `build_models`) call this. + - `_usage_response_hook(response: httpx.Response) -> None` (module-private, but imported directly by tests for unit coverage) + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_token_usage_hook.py`: + +```python +"""Tests for the token-usage HTTP response hook and attach_usage_hook wiring.""" +from __future__ import annotations + +import asyncio +import json + +import httpx + +from rag_eval.metrics.factory import _usage_response_hook, attach_usage_hook +from rag_eval.metrics.token_tracker import track_token_usage + + +def _fake_response(payload: dict | None) -> httpx.Response: + """Build a real httpx.Response with a JSON (or broken) body for hook testing.""" + content = b"not json" if payload is None else json.dumps(payload).encode("utf-8") + return httpx.Response(200, content=content, request=httpx.Request("POST", "http://test/x")) + + +class TestUsageResponseHook: + def test_records_usage_when_tracker_active(self): + with track_token_usage() as tracker: + response = _fake_response({ + "model": "gpt-5", + "usage": {"prompt_tokens": 120, "completion_tokens": 45}, + }) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == { + "gpt-5": {"input_tokens": 120, "output_tokens": 45, "calls": 1} + } + + def test_noop_when_no_tracker_active(self): + response = _fake_response({"model": "gpt-5", "usage": {"prompt_tokens": 1, "completion_tokens": 1}}) + # Must not raise even though no tracker is active. + asyncio.run(_usage_response_hook(response)) + + def test_noop_when_response_has_no_usage_field(self): + with track_token_usage() as tracker: + response = _fake_response({"model": "gpt-5"}) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == {} + + def test_noop_on_non_json_response(self): + with track_token_usage() as tracker: + response = _fake_response(None) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == {} + + def test_embedding_response_without_completion_tokens_defaults_output_to_zero(self): + """Embeddings responses omit completion_tokens; output should default to 0.""" + with track_token_usage() as tracker: + response = _fake_response({ + "model": "Qwen/Qwen3-Embedding-4B", + "usage": {"prompt_tokens": 30, "total_tokens": 30}, + }) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == { + "Qwen/Qwen3-Embedding-4B": {"input_tokens": 30, "output_tokens": 0, "calls": 1} + } + + +class TestAttachUsageHook: + def test_attaches_hook_to_client_event_hooks(self): + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1") + attach_usage_hook(client) + assert _usage_response_hook in client._client.event_hooks["response"] + + def test_idempotent_when_called_twice_on_same_client(self): + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1") + attach_usage_hook(client) + attach_usage_hook(client) + assert client._client.event_hooks["response"].count(_usage_response_hook) == 1 +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_token_usage_hook.py -v` +Expected: FAIL with `ImportError: cannot import name '_usage_response_hook' from 'rag_eval.metrics.factory'` + +- [ ] **Step 3: Write minimal implementation** + +Edit `rag_eval/metrics/factory.py`. Find the existing single import line: + +```python +from openai import AsyncOpenAI +``` + +Replace with (adds the new `httpx` import right before it): + +```python +import httpx +from openai import AsyncOpenAI +``` + +Find: + +```python +from rag_eval.shared.models import Scenario +``` + +Replace with (adds the new `token_tracker` import right after it): + +```python +from rag_eval.shared.models import Scenario + +from .token_tracker import get_current_tracker +``` + +Add these two functions right after the module-level `logger = logging.getLogger(...)` line and before `_resolve_openai_client_kwargs`: + +```python +async def _usage_response_hook(response: httpx.Response) -> None: + """Record token usage from an OpenAI-compatible HTTP response, if a tracker is active. + + Applies to both chat-completions and embeddings responses since both + return top-level `model` and `usage` fields in OpenAI-compatible APIs. + Never raises — a broken/incompatible gateway response must not affect scoring. + """ + tracker = get_current_tracker() + if tracker is None: + return + try: + await response.aread() + data = response.json() + usage = data.get("usage") + if not usage: + # Gateway did not report usage at all — skip rather than record a + # misleading 0/0 call. + return + model = data.get("model") or "unknown" + tracker.record( + model, + int(usage.get("prompt_tokens", 0) or 0), + int(usage.get("completion_tokens", 0) or 0), + ) + except Exception: # noqa: BLE001 + logger.debug("[factory] usage hook failed to parse response", exc_info=True) + + +def attach_usage_hook(client: AsyncOpenAI) -> None: + """Attach the token-usage response hook to an AsyncOpenAI client (idempotent). + + Safe to call multiple times on the same client (e.g. when judge and + embedding models share one client) — the hook is only appended once. + """ + httpx_client = getattr(client, "_client", None) + if httpx_client is None or not hasattr(httpx_client, "event_hooks"): + return + hooks = httpx_client.event_hooks.setdefault("response", []) + if _usage_response_hook not in hooks: + hooks.append(_usage_response_hook) +``` + +Now wire it into `build_models()`. Find: + +```python + llm_client = AsyncOpenAI(**llm_kwargs) + # Only allocate a second client when the embedding model needs different settings. + emb_client = AsyncOpenAI(**emb_kwargs) if emb_kwargs != llm_kwargs else llm_client +``` + +Replace with: + +```python + llm_client = AsyncOpenAI(**llm_kwargs) + attach_usage_hook(llm_client) + # Only allocate a second client when the embedding model needs different settings. + emb_client = AsyncOpenAI(**emb_kwargs) if emb_kwargs != llm_kwargs else llm_client + attach_usage_hook(emb_client) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_token_usage_hook.py -v` +Expected: `7 passed` + +Also re-run the pre-existing `build_models` tests to confirm no regression (they patch `AsyncOpenAI` with a fake class that has no `_client` attribute, so `attach_usage_hook` must degrade gracefully via the `getattr(client, "_client", None)` guard): + +Run: `C:\software\Python312\python.exe -m pytest tests/test_build_models_separate_clients.py -v` +Expected: `3 passed` + +- [ ] **Step 5: Commit** + +```powershell +git add rag_eval/metrics/factory.py tests/test_token_usage_hook.py +git commit -m "feat(token-tracking): add HTTP response hook and attach_usage_hook, wire into build_models" +``` + +--- + +### Task 3: 优化顾问 LLM 分析调用接入用量钩子 + +**Files:** +- Modify: `rag_eval/advisor/llm_analyzer.py` +- Test: `tests/test_advisor_llm_analyzer.py` (extend existing file) + +**Interfaces:** +- Consumes: `attach_usage_hook(client: AsyncOpenAI) -> None` from Task 2 (`rag_eval.metrics.factory`) + +- [ ] **Step 1: Write the failing test** + +Add to `tests/test_advisor_llm_analyzer.py` (after `test_analyze_does_not_close_injected_client`, before `test_is_reasoning_model_detection`): + +```python +def test_analyze_attaches_usage_hook_to_self_created_client(monkeypatch) -> None: + """A self-created client gets the token-usage hook attached (not the injected-client path).""" + captured: dict = {} + fake = _FakeClient(captured) + hook_calls = [] + + import openai + import rag_eval.metrics.factory as factory_mod + + monkeypatch.setattr(openai, "AsyncOpenAI", lambda **kwargs: fake) + monkeypatch.setattr( + factory_mod, "resolve_openai_client_kwargs", lambda *a, **k: {"api_key": "x"} + ) + monkeypatch.setattr(factory_mod, "attach_usage_hook", lambda c: hook_calls.append(c)) + + asyncio.run(analyze([_diagnosis()], "scn", "gpt-4o", _Settings())) + + assert hook_calls == [fake] + + +def test_analyze_does_not_attach_hook_for_injected_client() -> None: + """An injected chat_client is assumed to already have the hook attached by its owner.""" + hook_calls = [] + import rag_eval.metrics.factory as factory_mod + import unittest.mock as mock + + with mock.patch.object(factory_mod, "attach_usage_hook", lambda c: hook_calls.append(c)): + fake = _FakeClient({}) + asyncio.run(analyze([_diagnosis()], "scn", "gpt-4o", _Settings(), chat_client=fake)) + + assert hook_calls == [] +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_advisor_llm_analyzer.py -v -k usage_hook` +Expected: FAIL — `test_analyze_attaches_usage_hook_to_self_created_client` fails with `assert [] == [fake]` (hook never called). + +- [ ] **Step 3: Write minimal implementation** + +In `rag_eval/advisor/llm_analyzer.py`, find the self-created client branch inside `analyze()`: + +```python + client = chat_client + owns_client = False + if client is None: + from openai import AsyncOpenAI + + from rag_eval.metrics.factory import resolve_openai_client_kwargs + + client = AsyncOpenAI(**resolve_openai_client_kwargs(judge_model, settings)) + owns_client = True +``` + +Replace with: + +```python + client = chat_client + owns_client = False + if client is None: + from openai import AsyncOpenAI + + from rag_eval.metrics.factory import attach_usage_hook, resolve_openai_client_kwargs + + client = AsyncOpenAI(**resolve_openai_client_kwargs(judge_model, settings)) + attach_usage_hook(client) + owns_client = True +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_advisor_llm_analyzer.py -v` +Expected: all tests pass (previous tests + 2 new ones) + +- [ ] **Step 5: Commit** + +```powershell +git add rag_eval/advisor/llm_analyzer.py tests/test_advisor_llm_analyzer.py +git commit -m "feat(token-tracking): attach usage hook to advisor's self-created LLM client" +``` + +--- + +### Task 4: `EvaluationResult.token_usage` 字段 + `metadata.json` 持久化 + +**Files:** +- Modify: `rag_eval/shared/models.py` +- Modify: `rag_eval/reporting/writers.py` +- Test: `tests/test_token_usage_persistence.py` + +**Interfaces:** +- Produces: `EvaluationResult.token_usage: dict[str, dict[str, int]]` (default `{}`), consumed by Tasks 5, 6, 7, 8, 9. + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_token_usage_persistence.py`: + +```python +"""Tests that EvaluationResult.token_usage is persisted into metadata.json.""" +from __future__ import annotations + +import json +from pathlib import Path + +from rag_eval.reporting.writers import write_run_artifacts +from rag_eval.shared.models import DatasetConfig, EvaluationResult, RuntimeConfig, Scenario + + +def _scenario(tmp_path: Path) -> Scenario: + return Scenario( + scenario_name="token-persist-test", + mode="offline", + dataset=DatasetConfig(path=tmp_path / "dataset.csv"), + judge_model="gpt-5", + embedding_model="embedding-model", + metrics=["faithfulness"], + output_dir=tmp_path / "outputs", + runtime=RuntimeConfig(batch_size=1), + ) + + +def test_evaluation_result_defaults_token_usage_to_empty_dict(tmp_path: Path) -> None: + result = EvaluationResult( + scenario=_scenario(tmp_path), + run_id="run-1", + started_at="t0", + finished_at="t1", + valid_samples=[], + invalid_samples=[], + score_rows=[], + ) + assert result.token_usage == {} + + +def test_write_run_artifacts_persists_token_usage(tmp_path: Path) -> None: + scenario = _scenario(tmp_path) + result = EvaluationResult( + scenario=scenario, + run_id="run-2", + started_at="t0", + finished_at="t1", + valid_samples=[], + invalid_samples=[], + score_rows=[], + token_usage={"gpt-5": {"input_tokens": 100, "output_tokens": 40, "calls": 2}}, + ) + + write_run_artifacts(result) + + metadata_path = scenario.output_dir / "run-2" / "metadata.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + assert metadata["token_usage"] == { + "gpt-5": {"input_tokens": 100, "output_tokens": 40, "calls": 2} + } + + +def test_write_run_artifacts_writes_empty_token_usage_when_unset(tmp_path: Path) -> None: + scenario = _scenario(tmp_path) + result = EvaluationResult( + scenario=scenario, + run_id="run-3", + started_at="t0", + finished_at="t1", + valid_samples=[], + invalid_samples=[], + score_rows=[], + ) + + write_run_artifacts(result) + + metadata_path = scenario.output_dir / "run-3" / "metadata.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + assert metadata["token_usage"] == {} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_token_usage_persistence.py -v` +Expected: FAIL — `test_evaluation_result_defaults_token_usage_to_empty_dict` fails with `TypeError: EvaluationResult.__init__() got an unexpected keyword argument 'token_usage'` (for the other two tests, `KeyError: 'token_usage'` on the metadata assertion). + +- [ ] **Step 3: Write minimal implementation** + +In `rag_eval/shared/models.py`, find the `EvaluationResult` dataclass: + +```python +@dataclass(slots=True) +class EvaluationResult: + """Aggregate result object returned after a scenario completes.""" + + scenario: Scenario + run_id: str + started_at: str + finished_at: str + valid_samples: list[NormalizedSample] + invalid_samples: list[InvalidSample] + score_rows: list[dict[str, Any]] +``` + +Replace with: + +```python +@dataclass(slots=True) +class EvaluationResult: + """Aggregate result object returned after a scenario completes.""" + + scenario: Scenario + run_id: str + started_at: str + finished_at: str + valid_samples: list[NormalizedSample] + invalid_samples: list[InvalidSample] + score_rows: list[dict[str, Any]] + # Token usage grouped by model name: {model: {input_tokens, output_tokens, calls}}. + # Populated by callers via rag_eval.metrics.token_tracker.track_token_usage(). + token_usage: dict[str, dict[str, int]] = field(default_factory=dict) +``` + +In `rag_eval/reporting/writers.py`, find: + +```python + metadata = { + "run_id": result.run_id, + "scenario_name": result.scenario.scenario_name, + "mode": result.scenario.mode, + "judge_model": result.scenario.judge_model, + "embedding_model": result.scenario.embedding_model, + "started_at": result.started_at, + "finished_at": result.finished_at, + "dataset": result.scenario.dataset.path.as_posix(), + "valid_samples": len(result.valid_samples), + "invalid_samples": len(result.invalid_samples), + } +``` + +Replace with: + +```python + metadata = { + "run_id": result.run_id, + "scenario_name": result.scenario.scenario_name, + "mode": result.scenario.mode, + "judge_model": result.scenario.judge_model, + "embedding_model": result.scenario.embedding_model, + "started_at": result.started_at, + "finished_at": result.finished_at, + "dataset": result.scenario.dataset.path.as_posix(), + "valid_samples": len(result.valid_samples), + "invalid_samples": len(result.invalid_samples), + "token_usage": result.token_usage, + } +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_token_usage_persistence.py -v` +Expected: `3 passed` + +Also re-run the broader offline/online suites to confirm the new dataclass field (with a default) doesn't break any existing `EvaluationResult(...)` construction call sites: + +Run: `C:\software\Python312\python.exe -m pytest tests/test_offline_eval.py tests/test_online_eval.py -v` +Expected: all pass (unchanged) + +- [ ] **Step 5: Commit** + +```powershell +git add rag_eval/shared/models.py rag_eval/reporting/writers.py tests/test_token_usage_persistence.py +git commit -m "feat(token-tracking): add EvaluationResult.token_usage and persist to metadata.json" +``` + +--- + +### Task 5: `summary.md` 新增 "Token 用量" 小节 + +**Files:** +- Modify: `rag_eval/reporting/summary.py` +- Test: `tests/test_reporting_summary_token_usage.py` + +**Interfaces:** +- Consumes: `EvaluationResult.token_usage` from Task 4 +- Produces: `_token_usage_section(token_usage: dict[str, dict[str, int]]) -> list[str]` (module-private helper, used by `build_summary_markdown`) + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_reporting_summary_token_usage.py`: + +```python +"""Tests for the '## Token 用量' section rendered by build_summary_markdown.""" +from __future__ import annotations + +from pathlib import Path + +from rag_eval.reporting.summary import build_summary_markdown +from rag_eval.shared.models import DatasetConfig, EvaluationResult, RuntimeConfig, Scenario + + +def _scenario(tmp_path: Path) -> Scenario: + return Scenario( + scenario_name="summary-token-test", + mode="offline", + dataset=DatasetConfig(path=tmp_path / "dataset.csv"), + judge_model="gpt-5", + embedding_model="embedding-model", + metrics=["faithfulness"], + output_dir=tmp_path / "outputs", + runtime=RuntimeConfig(batch_size=1), + ) + + +def test_summary_includes_token_usage_table(tmp_path: Path) -> None: + scenario = _scenario(tmp_path) + result = EvaluationResult( + scenario=scenario, + run_id="run-1", + started_at="t0", + finished_at="t1", + valid_samples=[], + invalid_samples=[], + score_rows=[{"sample_id": "s1", "faithfulness": 0.9, "error": ""}], + token_usage={ + "gpt-5": {"input_tokens": 12450, "output_tokens": 3200, "calls": 60}, + "Qwen/Qwen3-Embedding-4B": {"input_tokens": 45000, "output_tokens": 0, "calls": 30}, + }, + ) + + markdown = build_summary_markdown(result) + + assert "## Token 用量" in markdown + assert "gpt-5" in markdown + assert "12450" in markdown + assert "Qwen/Qwen3-Embedding-4B" in markdown + assert "45000" in markdown + + +def test_summary_shows_fallback_text_when_token_usage_empty(tmp_path: Path) -> None: + scenario = _scenario(tmp_path) + result = EvaluationResult( + scenario=scenario, + run_id="run-2", + started_at="t0", + finished_at="t1", + valid_samples=[], + invalid_samples=[], + score_rows=[{"sample_id": "s1", "faithfulness": 0.9, "error": ""}], + ) + + markdown = build_summary_markdown(result) + + assert "## Token 用量" in markdown + assert "未记录 token 用量" in markdown +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_reporting_summary_token_usage.py -v` +Expected: FAIL — `assert "## Token 用量" in markdown` fails (section not present yet). + +- [ ] **Step 3: Write minimal implementation** + +In `rag_eval/reporting/summary.py`, add this helper function after `_table_from_frame` and before `build_summary_markdown`: + +```python +def _token_usage_section(token_usage: dict[str, dict[str, int]]) -> list[str]: + """Render the '## Token 用量' section as a list of markdown lines.""" + lines = ["", "## Token 用量", ""] + if not token_usage: + lines.append("未记录 token 用量。") + return lines + lines.append("| 模型 | input_tokens | output_tokens | 调用次数 |") + lines.append("|---|---|---|---|") + for model in sorted(token_usage): + usage = token_usage[model] + lines.append( + f"| {model} | {usage.get('input_tokens', 0)} " + f"| {usage.get('output_tokens', 0)} | {usage.get('calls', 0)} |" + ) + return lines +``` + +Then find the two `return` statements in `build_summary_markdown`. First, the empty-scores early return: + +```python + if scores.empty: + lines.append("No valid samples were scored.") + return "\n".join(lines) + "\n" +``` + +Replace with: + +```python + if scores.empty: + lines.append("No valid samples were scored.") + lines.extend(_token_usage_section(result.token_usage)) + return "\n".join(lines) + "\n" +``` + +Second, the final return at the end of the function: + +```python + detail_columns = ["sample_id", *result.scenario.metrics, "weighted_score", "error"] + existing_columns = [c for c in detail_columns if c in scores.columns] + detail = scores[existing_columns] + lines.extend([ + "", + "## Per-sample Scores", + "", + "```text", + _table_from_frame(detail), + "```", + ]) + return "\n".join(lines) + "\n" +``` + +Replace with: + +```python + detail_columns = ["sample_id", *result.scenario.metrics, "weighted_score", "error"] + existing_columns = [c for c in detail_columns if c in scores.columns] + detail = scores[existing_columns] + lines.extend([ + "", + "## Per-sample Scores", + "", + "```text", + _table_from_frame(detail), + "```", + ]) + lines.extend(_token_usage_section(result.token_usage)) + return "\n".join(lines) + "\n" +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_reporting_summary_token_usage.py -v` +Expected: `2 passed` + +- [ ] **Step 5: Commit** + +```powershell +git add rag_eval/reporting/summary.py tests/test_reporting_summary_token_usage.py +git commit -m "feat(token-tracking): add Token usage section to summary.md" +``` + +--- + +### Task 6: CLI `main.py --scenario` 评测流程接入统计 + +**Files:** +- Modify: `rag_eval/execution/evaluator.py` +- Test: `tests/test_evaluator_token_usage.py` + +**Interfaces:** +- Consumes: `track_token_usage()` from Task 1; `EvaluationResult.token_usage` field from Task 4. + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_evaluator_token_usage.py`: + +```python +"""Tests verifying the CLI evaluation flow captures token usage from metric scoring.""" +from __future__ import annotations + +import shutil +import unittest +from pathlib import Path + +import pandas as pd + +from rag_eval.execution.evaluator import Evaluator +from rag_eval.metrics.pipeline import MetricPipeline +from rag_eval.metrics.token_tracker import get_current_tracker +from rag_eval.shared.models import DatasetConfig, RuntimeConfig, Scenario + + +class FakeMetricWithUsage: + """Fake RAGAS metric that records token usage like a real HTTP-hooked call would.""" + + def __init__(self, value: float, model: str, input_tokens: int, output_tokens: int): + self.value = value + self.model = model + self.input_tokens = input_tokens + self.output_tokens = output_tokens + + async def ascore(self, **kwargs): + tracker = get_current_tracker() + if tracker is not None: + tracker.record(self.model, self.input_tokens, self.output_tokens) + + class Result: + def __init__(self, value: float): + self.value = value + + return Result(self.value) + + +class PlainFakeMetric: + """Fake metric that never records usage (simulates a hook that captured nothing).""" + + async def ascore(self, **kwargs): + class Result: + value = 0.9 + + return Result() + + +class EvaluatorTokenUsageTests(unittest.TestCase): + def setUp(self) -> None: + root = Path("tests/.tmp").resolve() + root.mkdir(parents=True, exist_ok=True) + self.temp_dir = root / self._testMethodName + shutil.rmtree(self.temp_dir, ignore_errors=True) + self.temp_dir.mkdir(parents=True, exist_ok=True) + + def tearDown(self) -> None: + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def _write_offline_dataset(self, path: Path, rows: list[dict]) -> None: + pd.DataFrame(rows).to_csv(path, index=False) + + def test_evaluate_populates_token_usage_from_metric_calls(self) -> None: + dataset_path = self.temp_dir / "offline.csv" + self._write_offline_dataset(dataset_path, [ + { + "sample_id": "sample-1", + "question": "What is the policy scope?", + "answer": "It covers all employees.", + "contexts": '["Context A"]', + "ground_truth": "It covers all employees.", + }, + { + "sample_id": "sample-2", + "question": "What about contractors?", + "answer": "Contractors are excluded.", + "contexts": '["Context B"]', + "ground_truth": "Contractors are excluded.", + }, + ]) + + scenario = Scenario( + scenario_name="token-usage-test", + mode="offline", + dataset=DatasetConfig(path=dataset_path), + judge_model="gpt-5", + embedding_model="embedding-model", + metrics=["faithfulness"], + output_dir=self.temp_dir / "outputs", + runtime=RuntimeConfig(batch_size=1), + ) + pipeline = MetricPipeline( + metrics={"faithfulness": FakeMetricWithUsage(0.8, "gpt-5", 100, 40)} + ) + evaluator = Evaluator(scenario=scenario, metric_pipeline=pipeline) + + result = evaluator.evaluate() + + # Two samples each recorded one call → totals sum across both. + self.assertEqual( + result.token_usage, + {"gpt-5": {"input_tokens": 200, "output_tokens": 80, "calls": 2}}, + ) + + def test_evaluate_defaults_to_empty_token_usage_when_nothing_recorded(self) -> None: + dataset_path = self.temp_dir / "offline.csv" + self._write_offline_dataset(dataset_path, [ + { + "sample_id": "sample-1", + "question": "What is the policy scope?", + "answer": "It covers all employees.", + "contexts": '["Context A"]', + "ground_truth": "It covers all employees.", + }, + ]) + + scenario = Scenario( + scenario_name="token-usage-empty-test", + mode="offline", + dataset=DatasetConfig(path=dataset_path), + judge_model="gpt-5", + embedding_model="embedding-model", + metrics=["faithfulness"], + output_dir=self.temp_dir / "outputs", + runtime=RuntimeConfig(batch_size=1), + ) + pipeline = MetricPipeline(metrics={"faithfulness": PlainFakeMetric()}) + evaluator = Evaluator(scenario=scenario, metric_pipeline=pipeline) + + result = evaluator.evaluate() + self.assertEqual(result.token_usage, {}) +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_evaluator_token_usage.py -v` +Expected: FAIL — `AssertionError: {} != {'gpt-5': {...}}` (evaluator does not wrap scoring in `track_token_usage()` yet, so nothing gets captured and `result.token_usage` stays `{}` from the Task 4 default). + +- [ ] **Step 3: Write minimal implementation** + +In `rag_eval/execution/evaluator.py`, add the import near the other `rag_eval` imports: + +```python +from rag_eval.metrics.pipeline import MetricPipeline +from rag_eval.metrics.token_tracker import track_token_usage +from rag_eval.metrics.weights import compute_weighted_score, resolve_weight +``` + +Find: + +```python + logger.info("[eval] scoring %d samples with metric pipeline ...", len(samples)) + t0 = time.monotonic() + metric_scores = asyncio.run( + self.metric_pipeline.score_samples( + samples, + max_concurrency=self.scenario.runtime.metric_limit(), + ) + ) + elapsed = time.monotonic() - t0 + logger.info("[eval] metric scoring done elapsed=%.1fs", elapsed) +``` + +Replace with: + +```python + logger.info("[eval] scoring %d samples with metric pipeline ...", len(samples)) + t0 = time.monotonic() + with track_token_usage() as usage_tracker: + metric_scores = asyncio.run( + self.metric_pipeline.score_samples( + samples, + max_concurrency=self.scenario.runtime.metric_limit(), + ) + ) + elapsed = time.monotonic() - t0 + logger.info("[eval] metric scoring done elapsed=%.1fs", elapsed) + logger.info("[eval] token_usage=%s", usage_tracker.as_dict()) +``` + +Find the final `return EvaluationResult(...)`: + +```python + return EvaluationResult( + scenario=self.scenario, + run_id=run_id, + started_at=started_at, + finished_at=finished_at, + valid_samples=samples, + invalid_samples=invalid_samples, + score_rows=score_rows, + ) +``` + +Replace with: + +```python + return EvaluationResult( + scenario=self.scenario, + run_id=run_id, + started_at=started_at, + finished_at=finished_at, + valid_samples=samples, + invalid_samples=invalid_samples, + score_rows=score_rows, + token_usage=usage_tracker.as_dict(), + ) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_evaluator_token_usage.py -v` +Expected: `2 passed` + +Also re-run the broader offline/online suites to confirm no regression: + +Run: `C:\software\Python312\python.exe -m pytest tests/test_offline_eval.py tests/test_online_eval.py -v` +Expected: all pass (unchanged) + +- [ ] **Step 5: Commit** + +```powershell +git add rag_eval/execution/evaluator.py tests/test_evaluator_token_usage.py +git commit -m "feat(token-tracking): wrap CLI evaluator metric scoring in track_token_usage" +``` + +--- + +### Task 7: `/api/score/async`(`score_job_manager.py`)接入统计 + +**Files:** +- Modify: `webapp/services/score_job_manager.py` +- Test: `tests/webapp/test_score_job_manager_token_usage.py` + +**Interfaces:** +- Consumes: `track_token_usage()` from Task 1; `EvaluationResult.token_usage` from Task 4. + +- [ ] **Step 1: Write the failing test** + +Create `tests/webapp/test_score_job_manager_token_usage.py`: + +```python +"""Tests that /api/score/async persists token usage captured during scoring.""" +from __future__ import annotations + +import json +import time + +from webapp.models import ScoreRequest +from webapp.services.score_job_manager import ScoreJobManager + + +def _wait_for_status(mgr: ScoreJobManager, job_id: str, timeout: float = 2.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + status = mgr.get(job_id) + if status is not None and status.status in ("completed", "failed"): + return status + time.sleep(0.02) + raise TimeoutError(f"job {job_id} did not complete in time") + + +def test_run_writes_token_usage_to_metadata(tmp_path, monkeypatch): + """_run() wraps inline_scorer.score in track_token_usage and persists totals.""" + from rag_eval.metrics.token_tracker import get_current_tracker + + mgr = ScoreJobManager( + output_dir=tmp_path / "score-async", + index_dir=tmp_path / "score-jobs", + max_workers=1, + ) + + def _fake_score(**kwargs): + tracker = get_current_tracker() + if tracker is not None: + tracker.record("gpt-5", 120, 45) + return {m: 0.9 for m in kwargs["metrics"]} + + monkeypatch.setattr( + "webapp.services.score_job_manager.inline_scorer.score", _fake_score + ) + + request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) + status = mgr.submit(request) + final_status = _wait_for_status(mgr, status.job_id) + + assert final_status.status == "completed" + run_dir = tmp_path / "score-async" / final_status.run_id + metadata = json.loads((run_dir / "metadata.json").read_text(encoding="utf-8")) + assert metadata["token_usage"] == { + "gpt-5": {"input_tokens": 120, "output_tokens": 45, "calls": 1} + } + + +def test_run_writes_empty_token_usage_when_nothing_recorded(tmp_path, monkeypatch): + mgr = ScoreJobManager( + output_dir=tmp_path / "score-async", + index_dir=tmp_path / "score-jobs", + max_workers=1, + ) + + def _fake_score(**kwargs): + return {m: 0.9 for m in kwargs["metrics"]} + + monkeypatch.setattr( + "webapp.services.score_job_manager.inline_scorer.score", _fake_score + ) + + request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) + status = mgr.submit(request) + final_status = _wait_for_status(mgr, status.job_id) + + run_dir = tmp_path / "score-async" / final_status.run_id + metadata = json.loads((run_dir / "metadata.json").read_text(encoding="utf-8")) + assert metadata["token_usage"] == {} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/webapp/test_score_job_manager_token_usage.py -v` +Expected: FAIL — `KeyError: 'token_usage'` is not raised (Task 4 already ensures the key exists), but the first test fails with `assert {} == {'gpt-5': {...}}` since `_run()` does not yet wrap scoring in `track_token_usage()`. + +- [ ] **Step 3: Write minimal implementation** + +In `webapp/services/score_job_manager.py`, inside `_run()`, add the import to the existing lazy-import block: + +```python + from rag_eval.advisor import run_advisor + from rag_eval.metrics.weights import compute_weighted_score + from rag_eval.reporting.writers import write_run_artifacts + from rag_eval.settings import EvaluationSettings +``` + +becomes: + +```python + from rag_eval.advisor import run_advisor + from rag_eval.metrics.token_tracker import track_token_usage + from rag_eval.metrics.weights import compute_weighted_score + from rag_eval.reporting.writers import write_run_artifacts + from rag_eval.settings import EvaluationSettings +``` + +Find: + +```python + try: + if effective: + raw_scores = inline_scorer.score( + question=request.question, + answer=request.answer, + contexts=request.contexts_as_list(), + ground_truth=request.ground_truth, + metrics=effective, + judge_model=judge_model, + embedding_model=embedding_model, + settings=settings, + judge_language=judge_language, + ) + else: + raw_scores = {} +``` + +Replace with: + +```python + try: + with track_token_usage() as usage_tracker: + if effective: + raw_scores = inline_scorer.score( + question=request.question, + answer=request.answer, + contexts=request.contexts_as_list(), + ground_truth=request.ground_truth, + metrics=effective, + judge_model=judge_model, + embedding_model=embedding_model, + settings=settings, + judge_language=judge_language, + ) + else: + raw_scores = {} +``` + +Find the `EvaluationResult(` construction: + +```python + result = EvaluationResult( + scenario=scenario, + run_id=run_id, + started_at=started_at, + finished_at=finished_at, + valid_samples=[sample], + invalid_samples=[], + score_rows=[score_row], + ) +``` + +Replace with: + +```python + result = EvaluationResult( + scenario=scenario, + run_id=run_id, + started_at=started_at, + finished_at=finished_at, + valid_samples=[sample], + invalid_samples=[], + score_rows=[score_row], + token_usage=usage_tracker.as_dict(), + ) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/webapp/test_score_job_manager_token_usage.py -v` +Expected: `2 passed` + +Also re-run the existing async score jobs API tests to confirm no regression: + +Run: `C:\software\Python312\python.exe -m pytest tests/webapp/test_score_jobs_api.py -v` +Expected: all pass (unchanged) + +- [ ] **Step 5: Commit** + +```powershell +git add webapp/services/score_job_manager.py tests/webapp/test_score_job_manager_token_usage.py +git commit -m "feat(token-tracking): capture token usage in /api/score/async job runs" +``` + +--- + +### Task 8: `/api/score/session_async`(`session_score_manager.py`)接入统计并跨调用累加 + +**Files:** +- Modify: `webapp/services/session_score_manager.py` +- Test: `tests/webapp/test_session_score_manager_token_usage.py` + +**Interfaces:** +- Consumes: `track_token_usage()`, `TokenUsageTracker.merge_into(...)` from Task 1; `EvaluationResult.token_usage` from Task 4. +- Produces: `SessionScoreJobManager._read_metadata(run_dir: Path) -> dict[str, Any]` (private helper, mirrors the existing `_read_score_rows`). + +- [ ] **Step 1: Write the failing test** + +Create `tests/webapp/test_session_score_manager_token_usage.py`: + +```python +"""Tests that session-grouped async scoring accumulates token usage across calls.""" +from __future__ import annotations + +import json +import time + +from webapp.models import ScoreRequest +from webapp.services.session_score_manager import SessionScoreJobManager + + +def _wait_for_call_count(mgr: SessionScoreJobManager, session_id: str, expected: int, timeout: float = 2.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + session = mgr.get_session(session_id) + if session is not None and session.call_count >= expected: + all_done = all(j.status in ("completed", "failed") for j in session.jobs) + if all_done: + return session + time.sleep(0.02) + raise TimeoutError(f"session {session_id} did not reach {expected} completed calls in time") + + +def test_session_accumulates_token_usage_across_calls(tmp_path, monkeypatch): + from rag_eval.metrics.token_tracker import get_current_tracker + + mgr = SessionScoreJobManager( + output_dir=tmp_path / "score-session", + index_dir=tmp_path / "score-session-jobs", + max_workers=1, + ) + + call_usages = iter([(100, 40), (30, 10)]) + + def _fake_score(**kwargs): + tracker = get_current_tracker() + input_tok, output_tok = next(call_usages) + if tracker is not None: + tracker.record("gpt-5", input_tok, output_tok) + return {m: 0.9 for m in kwargs["metrics"]} + + monkeypatch.setattr( + "webapp.services.session_score_manager.inline_scorer.score", _fake_score + ) + + request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) + + _, run_id = mgr.submit("session-token-test", request) + _wait_for_call_count(mgr, "session-token-test", 1) + mgr.submit("session-token-test", request) + _wait_for_call_count(mgr, "session-token-test", 2) + + run_dir = tmp_path / "score-session" / run_id + metadata = json.loads((run_dir / "metadata.json").read_text(encoding="utf-8")) + assert metadata["token_usage"] == { + "gpt-5": {"input_tokens": 130, "output_tokens": 50, "calls": 2} + } + + +def test_session_first_call_writes_token_usage_from_scratch(tmp_path, monkeypatch): + from rag_eval.metrics.token_tracker import get_current_tracker + + mgr = SessionScoreJobManager( + output_dir=tmp_path / "score-session", + index_dir=tmp_path / "score-session-jobs", + max_workers=1, + ) + + def _fake_score(**kwargs): + tracker = get_current_tracker() + if tracker is not None: + tracker.record("gpt-5", 50, 20) + return {m: 0.9 for m in kwargs["metrics"]} + + monkeypatch.setattr( + "webapp.services.session_score_manager.inline_scorer.score", _fake_score + ) + + request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) + _, run_id = mgr.submit("session-first-call-test", request) + _wait_for_call_count(mgr, "session-first-call-test", 1) + + run_dir = tmp_path / "score-session" / run_id + metadata = json.loads((run_dir / "metadata.json").read_text(encoding="utf-8")) + assert metadata["token_usage"] == {"gpt-5": {"input_tokens": 50, "output_tokens": 20, "calls": 1}} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/webapp/test_session_score_manager_token_usage.py -v` +Expected: FAIL — both tests fail with `assert {} == {'gpt-5': {...}}` (the manager does not wrap scoring in `track_token_usage()` yet). + +- [ ] **Step 3: Write minimal implementation** + +In `webapp/services/session_score_manager.py`, add `_read_metadata` right after the existing `_read_score_rows` method: + +```python + def _read_score_rows(self, run_dir: Path) -> list[dict[str, Any]]: + """Read existing scores.csv rows, returning empty list if file doesn't exist.""" + scores_path = run_dir / "scores.csv" + if not scores_path.is_file(): + return [] + try: + frame = pd.read_csv(scores_path) + return frame.where(pd.notnull(frame), None).to_dict("records") + except (OSError, ValueError): + return [] + + def _read_metadata(self, run_dir: Path) -> dict[str, Any]: + """Read this session's existing metadata.json, returning {} if absent/unreadable.""" + metadata_path = run_dir / "metadata.json" + if not metadata_path.is_file(): + return {} + try: + return json.loads(metadata_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} +``` + +Add the import to the existing lazy-import block inside `_run()`: + +```python + from rag_eval.advisor import run_advisor + from rag_eval.metrics.weights import compute_weighted_score + from rag_eval.reporting.writers import write_run_artifacts +``` + +becomes: + +```python + from rag_eval.advisor import run_advisor + from rag_eval.metrics.token_tracker import track_token_usage + from rag_eval.metrics.weights import compute_weighted_score + from rag_eval.reporting.writers import write_run_artifacts +``` + +Find: + +```python + try: + # --- Scoring (can run concurrently for the same session) ---------- + if effective: + raw_scores = inline_scorer.score( + question=request.question, + answer=request.answer, + contexts=request.contexts_as_list(), + ground_truth=request.ground_truth, + metrics=effective, + judge_model=judge_model, + embedding_model=embedding_model, + settings=settings, + judge_language=judge_language, + ) + else: + raw_scores = {} +``` + +Replace with: + +```python + try: + # --- Scoring (can run concurrently for the same session) ---------- + with track_token_usage() as usage_tracker: + if effective: + raw_scores = inline_scorer.score( + question=request.question, + answer=request.answer, + contexts=request.contexts_as_list(), + ground_truth=request.ground_truth, + metrics=effective, + judge_model=judge_model, + embedding_model=embedding_model, + settings=settings, + judge_language=judge_language, + ) + else: + raw_scores = {} +``` + +Find, inside the `with session_lock:` block: + +```python + session_lock = self._get_session_lock(session_id) + with session_lock: + run_dir = self._output_dir / run_id + run_dir.mkdir(parents=True, exist_ok=True) + + # Read all existing rows, then append the new one + existing_rows = self._read_score_rows(run_dir) +``` + +Replace with: + +```python + session_lock = self._get_session_lock(session_id) + with session_lock: + run_dir = self._output_dir / run_id + run_dir.mkdir(parents=True, exist_ok=True) + + # Merge this call's token usage into the session's running total, so + # repeated calls accumulate instead of overwriting (mirrors the + # scores.csv append-only accumulation below). + existing_metadata = self._read_metadata(run_dir) + merged_token_usage = usage_tracker.merge_into( + existing_metadata.get("token_usage", {}) + ) + + # Read all existing rows, then append the new one + existing_rows = self._read_score_rows(run_dir) +``` + +Find the `EvaluationResult(` construction: + +```python + result = EvaluationResult( + scenario=scenario, + run_id=run_id, + started_at=started_at_val if isinstance(started_at_val, str) else finished_at, + finished_at=finished_at, + valid_samples=valid_samples, + invalid_samples=[], + score_rows=all_rows, + ) +``` + +Replace with: + +```python + result = EvaluationResult( + scenario=scenario, + run_id=run_id, + started_at=started_at_val if isinstance(started_at_val, str) else finished_at, + finished_at=finished_at, + valid_samples=valid_samples, + invalid_samples=[], + score_rows=all_rows, + token_usage=merged_token_usage, + ) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/webapp/test_session_score_manager_token_usage.py -v` +Expected: `2 passed` + +Also re-run the existing session score jobs API tests to confirm no regression: + +Run: `C:\software\Python312\python.exe -m pytest tests/webapp/test_session_score_jobs_api.py -v` +Expected: all pass (unchanged) + +- [ ] **Step 5: Commit** + +```powershell +git add webapp/services/session_score_manager.py tests/webapp/test_session_score_manager_token_usage.py +git commit -m "feat(token-tracking): accumulate token usage across session_async calls" +``` + +--- + +### Task 9: Web 报告层透传 `token_usage` + +**Files:** +- Modify: `webapp/models.py` +- Modify: `webapp/services/report_builder.py` +- Test: `tests/test_report_builder_token_usage.py` + +**Interfaces:** +- Consumes: `metadata.json`'s `token_usage` key (written by Tasks 4/7/8) +- Produces: `ReportData.token_usage: dict[str, dict[str, int]]`, consumed by Task 10's `report.js`. + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_report_builder_token_usage.py`: + +```python +"""Tests for token_usage passthrough in the webapp report builder.""" +from __future__ import annotations + +import json +from pathlib import Path + +from webapp.services.report_builder import build_report + + +def _write_minimal_run(run_dir: Path, token_usage: dict | None) -> None: + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "scores.csv").write_text( + "sample_id,faithfulness\ns1,0.9\n", encoding="utf-8" + ) + (run_dir / "summary.md").write_text("summary", encoding="utf-8") + metadata = {"run_id": run_dir.name} + if token_usage is not None: + metadata["token_usage"] = token_usage + (run_dir / "metadata.json").write_text(json.dumps(metadata), encoding="utf-8") + + +def test_build_report_passes_through_token_usage(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + _write_minimal_run( + run_dir, + token_usage={"gpt-5": {"input_tokens": 100, "output_tokens": 50, "calls": 2}}, + ) + + report = build_report(run_dir, ["faithfulness"]) + + assert report.token_usage == { + "gpt-5": {"input_tokens": 100, "output_tokens": 50, "calls": 2} + } + + +def test_build_report_defaults_token_usage_to_empty_dict(tmp_path: Path) -> None: + run_dir = tmp_path / "run" + _write_minimal_run(run_dir, token_usage=None) + + report = build_report(run_dir, ["faithfulness"]) + + assert report.token_usage == {} + + +def test_build_report_early_return_branch_still_surfaces_token_usage(tmp_path: Path) -> None: + """metrics=[] forces the early-return branch; token_usage must still surface.""" + run_dir = tmp_path / "run" + _write_minimal_run( + run_dir, + token_usage={"gpt-5": {"input_tokens": 5, "output_tokens": 5, "calls": 1}}, + ) + + report = build_report(run_dir, []) + + assert report.token_usage == {"gpt-5": {"input_tokens": 5, "output_tokens": 5, "calls": 1}} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_report_builder_token_usage.py -v` +Expected: FAIL — `AttributeError: 'ReportData' object has no attribute 'token_usage'` + +- [ ] **Step 3: Write minimal implementation** + +In `webapp/models.py`, add the field to `ReportData` (after `doc_weights`): + +```python +class ReportData(BaseModel): + """Aggregated report payload rendered by the report detail page.""" + + metrics: list[str] = Field(default_factory=list) + metric_means: dict[str, float | None] = Field(default_factory=dict) + distributions: dict[str, list[DistributionBin]] = Field(default_factory=dict) + groupings: dict[str, list[GroupStat]] = Field(default_factory=dict) + lowest_samples: list[SampleScore] = Field(default_factory=list) + summary_markdown: str = "" + advice_markdown: str = "" # optimization_advice.md content (empty if not generated) + weighted_score_mean: float | None = Field( + default=None, + description="加权综合得分均值(metric_weights × doc_weights 共同作用)。", + ) + metric_weights: dict[str, float] = Field( + default_factory=dict, + description="该次运行使用的指标权重配置(来自 scenario.snapshot.yaml)。", + ) + doc_weights: dict[str, float] = Field( + default_factory=dict, + description="该次运行使用的文档权重配置(来自 scenario.snapshot.yaml)。", + ) + token_usage: dict[str, dict[str, int]] = Field( + default_factory=dict, + description="按模型累计的 token 用量:{model: {input_tokens, output_tokens, calls}}。", + ) +``` + +In `webapp/services/report_builder.py`, find `build_report`: + +```python +def build_report(run_dir: Path, metrics: list[str]) -> ReportData: + """Build the full aggregated report payload for one run directory.""" + frame = run_reader.read_scores_frame(run_dir) + summary_markdown = run_reader.read_summary_markdown(run_dir) + advice_markdown = run_reader.read_advice_markdown(run_dir) + metric_weights, doc_weights = _read_weights_from_snapshot(run_dir) + + if frame.empty or not metrics: + return ReportData( + metrics=metrics, + metric_means={metric: None for metric in metrics}, + summary_markdown=summary_markdown, + advice_markdown=advice_markdown, + metric_weights=metric_weights, + doc_weights=doc_weights, + ) +``` + +Replace with: + +```python +def build_report(run_dir: Path, metrics: list[str]) -> ReportData: + """Build the full aggregated report payload for one run directory.""" + frame = run_reader.read_scores_frame(run_dir) + summary_markdown = run_reader.read_summary_markdown(run_dir) + advice_markdown = run_reader.read_advice_markdown(run_dir) + metric_weights, doc_weights = _read_weights_from_snapshot(run_dir) + # Read once up front so both the empty-frame and full branches can surface it. + metadata = run_reader._read_json(run_dir / "metadata.json") + token_usage = metadata.get("token_usage") or {} + + if frame.empty or not metrics: + return ReportData( + metrics=metrics, + metric_means={metric: None for metric in metrics}, + summary_markdown=summary_markdown, + advice_markdown=advice_markdown, + metric_weights=metric_weights, + doc_weights=doc_weights, + token_usage=token_usage, + ) +``` + +Further down in the same function, find: + +```python + # Cross-run history: scores of the same question in *other* runs (Approach A — + # on-demand global scan, excluding the run currently being viewed). + metadata = run_reader._read_json(run_dir / "metadata.json") + current_run_id = str(metadata.get("run_id") or run_dir.name) + history_index = question_history.build_question_history_index( + exclude_run_id=current_run_id + ) + + return ReportData( + metrics=metrics, + metric_means=rounded_means, + distributions=distributions, + groupings=_groupings(frame, metrics), + lowest_samples=_lowest_samples(frame, metrics, history_index), + summary_markdown=summary_markdown, + advice_markdown=advice_markdown, + weighted_score_mean=_round_or_none(overall_ws), + metric_weights=metric_weights, + doc_weights=doc_weights, + ) +``` + +Replace with (removes the now-duplicate `metadata` read, keeps `current_run_id` derivation): + +```python + # Cross-run history: scores of the same question in *other* runs (Approach A — + # on-demand global scan, excluding the run currently being viewed). + current_run_id = str(metadata.get("run_id") or run_dir.name) + history_index = question_history.build_question_history_index( + exclude_run_id=current_run_id + ) + + return ReportData( + metrics=metrics, + metric_means=rounded_means, + distributions=distributions, + groupings=_groupings(frame, metrics), + lowest_samples=_lowest_samples(frame, metrics, history_index), + summary_markdown=summary_markdown, + advice_markdown=advice_markdown, + weighted_score_mean=_round_or_none(overall_ws), + metric_weights=metric_weights, + doc_weights=doc_weights, + token_usage=token_usage, + ) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `C:\software\Python312\python.exe -m pytest tests/test_report_builder_token_usage.py -v` +Expected: `3 passed` + +Also re-run the existing report builder tests to confirm no regression: + +Run: `C:\software\Python312\python.exe -m pytest tests/test_webapp_report_builder.py -v` +Expected: all pass (unchanged) + +- [ ] **Step 5: Commit** + +```powershell +git add webapp/models.py webapp/services/report_builder.py tests/test_report_builder_token_usage.py +git commit -m "feat(token-tracking): surface token_usage in ReportData" +``` + +--- + +### Task 10: Web 报告详情页新增 "Token 用量" 面板 + +**Files:** +- Modify: `webapp/static/index.html` +- Modify: `webapp/static/js/report.js` + +**Interfaces:** +- Consumes: `ReportData.token_usage` (JSON field `token_usage` from Task 9's API response) + +This task has no automated test — the repository has no JavaScript test harness (vanilla JS, no `package.json`/Jest). Verify manually per Step 3 below. + +- [ ] **Step 1: Add the panel container to `index.html`** + +In `webapp/static/index.html`, find the optimization-advice section closing tag inside `#report-content`: + +```html + + + + +``` + +Replace with: + +```html + + + + + +
+ + +``` + +- [ ] **Step 2: Add `renderTokenUsage` to `report.js` and call it from `render()`** + +In `webapp/static/js/report.js`, find the `render()` method's body: + +```js + const detail = await API.runDetail(runId); + Report.currentDetail = detail; + Report.renderMeta(detail.summary); + Report.renderMetricCards(detail.summary, detail.report); + Report.renderDistribution(detail.report); + Report.renderGroupings(detail.report); + Report.renderLowest(detail.report); + Report.renderAdvice(detail.summary, detail.report); + content.style.opacity = "1"; +``` + +Replace with: + +```js + const detail = await API.runDetail(runId); + Report.currentDetail = detail; + Report.renderMeta(detail.summary); + Report.renderMetricCards(detail.summary, detail.report); + Report.renderDistribution(detail.report); + Report.renderGroupings(detail.report); + Report.renderLowest(detail.report); + Report.renderAdvice(detail.summary, detail.report); + Report.renderTokenUsage(detail.report); + content.style.opacity = "1"; +``` + +Add the new method right after `_drawGroupTable` (reuses the existing `group-table` CSS class for visual consistency): + +```js + // 渲染"Token 用量"面板:按模型分组的 input/output/调用次数表格。 + renderTokenUsage(report) { + const wrap = document.getElementById("token-usage-wrap"); + if (!wrap) return; + const usage = report.token_usage || {}; + const models = Object.keys(usage).sort(); + if (models.length === 0) { + wrap.innerHTML = '

暂无 token 用量数据。

'; + return; + } + let head = "模型input_tokensoutput_tokens调用次数"; + let body = ""; + models.forEach((model) => { + const u = usage[model] || {}; + body += `${App.escape(model)}${u.input_tokens ?? 0}` + + `${u.output_tokens ?? 0}${u.calls ?? 0}`; + }); + wrap.innerHTML = `${head}${body}
`; + }, +``` + +- [ ] **Step 3: Manual verification** + +Run: `C:\software\Python312\python.exe -m pytest tests/ -v -k token_usage` (sanity check: all backend token-usage tests from Tasks 1–9 still pass together) +Expected: all pass + +Then start the web server and manually confirm the panel renders: + +```powershell +C:\software\Python312\python.exe -m uvicorn webapp.server:app --reload --port 8800 +``` + +Open `http://localhost:8800`, submit any scoring request (e.g. via `/api/score/async`), open「运行列表」→ pick that run → 报告详情页 should show a new "⑥ Token 用量" panel: either the fallback text "暂无 token 用量数据。" (if the configured LLM gateway doesn't echo a `usage` field) or a table with model/input/output/调用次数 rows. + +- [ ] **Step 4: Commit** + +```powershell +git add webapp/static/index.html webapp/static/js/report.js +git commit -m "feat(token-tracking): add Token usage panel to the report detail page" +``` + +--- + +## Final Verification + +- [ ] **Run the full test suite once all 10 tasks are complete** + +Run: `C:\software\Python312\python.exe -m pytest tests/ -v` +Expected: All tests pass, including the new `test_token_tracker.py`, `test_token_usage_hook.py`, `test_token_usage_persistence.py`, `test_reporting_summary_token_usage.py`, `test_evaluator_token_usage.py`, `tests/webapp/test_score_job_manager_token_usage.py`, `tests/webapp/test_session_score_manager_token_usage.py`, `test_report_builder_token_usage.py`, plus all pre-existing tests unchanged (any pre-existing failures unrelated to this feature are expected to remain as-is — do not attempt to fix them as part of this plan). + +- [ ] **Spec coverage check** + +Confirm every section of `docs/superpowers/specs/2026-07-02-token-usage-tracking-design.md` maps to a task: +- §3.1 (`TokenUsageTracker` + contextvar) → Task 1 +- §3.2 (HTTP hook + `attach_usage_hook`) → Task 2 +- §3.3 (挂载点: `build_models` + `llm_analyzer`) → Tasks 2, 3 +- §4.1 (CLI evaluator) → Task 6 +- §4.2 (`/api/score/async`) → Task 7 +- §4.3 (`/api/score/session_async` 累加) → Task 8 +- §4.4 (`write_run_artifacts` + `summary.md`) → Tasks 4, 5 +- §5 (报告层与 Web UI) → Tasks 9, 10 +- §6 (错误处理与边界情况) → covered inline in Tasks 2 (hook silently no-ops), 8 (session lock reuse), 9 (missing-key defaults) +- §7 (测试策略) → one test file per task, matching the spec's proposed file list +- §8 (非目标) → respected: no cost/$ conversion, no `/api/score` sync coverage, no dataset_builder coverage diff --git a/rag_eval/metrics/factory.py b/rag_eval/metrics/factory.py index 97edbf8..42d3a73 100644 --- a/rag_eval/metrics/factory.py +++ b/rag_eval/metrics/factory.py @@ -5,12 +5,15 @@ from __future__ import annotations import logging from typing import Any +import httpx from openai import AsyncOpenAI from rag_eval.compat import ensure_ragas_import_compat from rag_eval.settings import EvaluationSettings from rag_eval.shared.models import Scenario +from .token_tracker import get_current_tracker + ensure_ragas_import_compat() from ragas.embeddings.base import embedding_factory @@ -32,6 +35,48 @@ from .pipeline import MetricPipeline logger = logging.getLogger("rag_eval.metrics.factory") +async def _usage_response_hook(response: httpx.Response) -> None: + """Record token usage from an OpenAI-compatible HTTP response, if a tracker is active. + + Applies to both chat-completions and embeddings responses since both + return top-level `model` and `usage` fields in OpenAI-compatible APIs. + Never raises — a broken/incompatible gateway response must not affect scoring. + """ + tracker = get_current_tracker() + if tracker is None: + return + try: + await response.aread() + data = response.json() + usage = data.get("usage") + if not usage: + # Gateway did not report usage at all — skip rather than record a + # misleading 0/0 call. + return + model = data.get("model") or "unknown" + tracker.record( + model, + int(usage.get("prompt_tokens", 0) or 0), + int(usage.get("completion_tokens", 0) or 0), + ) + except Exception: # noqa: BLE001 + logger.debug("[factory] usage hook failed to parse response", exc_info=True) + + +def attach_usage_hook(client: AsyncOpenAI) -> None: + """Attach the token-usage response hook to an AsyncOpenAI client (idempotent). + + Safe to call multiple times on the same client (e.g. when judge and + embedding models share one client) — the hook is only appended once. + """ + httpx_client = getattr(client, "_client", None) + if httpx_client is None or not hasattr(httpx_client, "event_hooks"): + return + hooks = httpx_client.event_hooks.setdefault("response", []) + if _usage_response_hook not in hooks: + hooks.append(_usage_response_hook) + + def _resolve_openai_client_kwargs( model: str, settings: EvaluationSettings, @@ -126,8 +171,10 @@ def build_models( ) llm_client = AsyncOpenAI(**llm_kwargs) + attach_usage_hook(llm_client) # Only allocate a second client when the embedding model needs different settings. emb_client = AsyncOpenAI(**emb_kwargs) if emb_kwargs != llm_kwargs else llm_client + attach_usage_hook(emb_client) # RAGAS structured-output judge calls can be truncated by the upstream default # 1024 completion budget, especially for faithfulness and GPT-5 family models. diff --git a/tests/test_token_usage_hook.py b/tests/test_token_usage_hook.py new file mode 100644 index 0000000..2caaa43 --- /dev/null +++ b/tests/test_token_usage_hook.py @@ -0,0 +1,75 @@ +"""Tests for the token-usage HTTP response hook and attach_usage_hook wiring.""" +from __future__ import annotations + +import asyncio +import json + +import httpx + +from rag_eval.metrics.factory import _usage_response_hook, attach_usage_hook +from rag_eval.metrics.token_tracker import track_token_usage + + +def _fake_response(payload: dict | None) -> httpx.Response: + """Build a real httpx.Response with a JSON (or broken) body for hook testing.""" + content = b"not json" if payload is None else json.dumps(payload).encode("utf-8") + return httpx.Response(200, content=content, request=httpx.Request("POST", "http://test/x")) + + +class TestUsageResponseHook: + def test_records_usage_when_tracker_active(self): + with track_token_usage() as tracker: + response = _fake_response({ + "model": "gpt-5", + "usage": {"prompt_tokens": 120, "completion_tokens": 45}, + }) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == { + "gpt-5": {"input_tokens": 120, "output_tokens": 45, "calls": 1} + } + + def test_noop_when_no_tracker_active(self): + response = _fake_response({"model": "gpt-5", "usage": {"prompt_tokens": 1, "completion_tokens": 1}}) + # Must not raise even though no tracker is active. + asyncio.run(_usage_response_hook(response)) + + def test_noop_when_response_has_no_usage_field(self): + with track_token_usage() as tracker: + response = _fake_response({"model": "gpt-5"}) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == {} + + def test_noop_on_non_json_response(self): + with track_token_usage() as tracker: + response = _fake_response(None) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == {} + + def test_embedding_response_without_completion_tokens_defaults_output_to_zero(self): + """Embeddings responses omit completion_tokens; output should default to 0.""" + with track_token_usage() as tracker: + response = _fake_response({ + "model": "Qwen/Qwen3-Embedding-4B", + "usage": {"prompt_tokens": 30, "total_tokens": 30}, + }) + asyncio.run(_usage_response_hook(response)) + assert tracker.as_dict() == { + "Qwen/Qwen3-Embedding-4B": {"input_tokens": 30, "output_tokens": 0, "calls": 1} + } + + +class TestAttachUsageHook: + def test_attaches_hook_to_client_event_hooks(self): + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1") + attach_usage_hook(client) + assert _usage_response_hook in client._client.event_hooks["response"] + + def test_idempotent_when_called_twice_on_same_client(self): + from openai import AsyncOpenAI + + client = AsyncOpenAI(api_key="sk-test", base_url="http://localhost:1") + attach_usage_hook(client) + attach_usage_hook(client) + assert client._client.event_hooks["response"].count(_usage_response_hook) == 1