217 lines
13 KiB
Markdown
217 lines
13 KiB
Markdown
# Judge LLM Token 用量追踪设计
|
||
|
||
**日期**: 2026-07-02
|
||
**状态**: 已批准，待实现
|
||
**范围**: 按 LLM 模型累计记录评测过程中判官（judge）与向量化（embedding）模型的原始 token 用量（不换算金额），持久化到 run 产物并在报告详情页展示（架构设计 §10.5）。
|
||
|
||
---
|
||
|
||
## 1. 背景与目标
|
||
|
||
架构设计 §10.5 明确要求：
|
||
|
||
> 此类 LLM 评判指标对每条样本产生多次 LLM 调用，规模化运行的 token 开销需预先评估——一次数百条的回归，叠加每条多次评判调用，量级不小。
|
||
|
||
当前全链路（`MetricScore` / `EvaluationResult` / `metadata.json` / `summary.md` / Web 报告）**没有任何 token 用量记录**，规模化跑评测时无法评估实际消耗。本设计新增能力：
|
||
|
||
- 按**模型名**累计 input/output token 数与调用次数（不做金额换算，YAGNI）。
|
||
- 覆盖 RAGAS 评分调用（judge model + embedding model）与优化顾问 LLM 分析调用（`advisor/llm_analyzer.py`）。
|
||
- 只统计**落盘成 run 目录**的路径：CLI `main.py --scenario`、`/api/score/async`、`/api/score/session_async`；不落盘的同步 `/api/score` 不纳入范围。
|
||
- 持久化到 `metadata.json`，摘要写入 `summary.md`，报告详情页新增"Token 用量"面板。
|
||
|
||
**成功标准**：
|
||
- 跑完一次评测（任一入口），`metadata.json` 里能看到按模型分组的 `{input_tokens, output_tokens, calls}`。
|
||
- session_async 多次调用同一 session，token 用量随调用**累加**，不是只保留最后一次。
|
||
- 网关不返回 `usage` 字段时静默跳过，不影响评分主流程、不抛异常。
|
||
- 现有评分行为、既有测试逐字节不变（这是纯增量的可观测性功能）。
|
||
|
||
---
|
||
|
||
## 2. 现状确认（已核实的技术事实）
|
||
|
||
1. RAGAS 0.4.3 的 `ragas/cost.py`（`TokenUsageParser` / `CostCallbackHandler`）只服务旧版 langchain `evaluate()` 回调链路，**新版 `ragas.metrics.collections` + instructor 路径完全不走这条链路**，无法复用。
|
||
2. `ragas.llms.base.llm_factory()` 内部用 `instructor.from_openai(client, mode=instructor.Mode.JSON)` 包装我们传入的 `AsyncOpenAI`，`InstructorLLM.generate/agenerate` 只在返回前触发 RAGAS 自己的匿名遥测事件（`track(LLMUsageEvent(...))`），**不是真实 token 计数**，也不能拿到。
|
||
3. `embedding_factory(provider="openai", ...)` 走的是普通 `AsyncOpenAI.embeddings.create()`，不经过 instructor，跟 chat completions 是两条不同的调用路径 —— 说明"在 instructor 层挂钩子"这条路无法同时覆盖 judge 和 embedding 两种调用，必须在更底层（HTTP 层）统一拦截。
|
||
4. `webapp/services/inline_scorer.py` 的 `InlineScorer` 用 `(judge_model, embedding_model)` 做 key **跨请求缓存** `(llm, embeddings)` 客户端对象；`/api/score`、`/api/score/async`（`score_job_manager.py`）、`/api/score/session_async`（`session_score_manager.py`）三个接口都调用同一个 `inline_scorer` 单例。**因此统计器不能绑定在客户端构造时刻，必须按每次调用动态切换**。
|
||
5. 三条落盘路径最终都调用同一个 `rag_eval/reporting/writers.py::write_run_artifacts(EvaluationResult)` 生成 `metadata.json`，这是唯一需要改动的持久化落点（DRY）。
|
||
|
||
---
|
||
|
||
## 3. 架构与数据流
|
||
|
||
### 3.1 核心组件：`rag_eval/metrics/token_tracker.py`（新文件）
|
||
|
||
```python
|
||
@dataclass
|
||
class TokenUsageTracker:
|
||
"""按模型名累计一次 run/一次请求范围内的 token 用量。"""
|
||
_totals: dict[str, dict[str, int]] = field(default_factory=dict)
|
||
|
||
def record(self, model: str, input_tokens: int, output_tokens: int) -> None:
|
||
"""累加一次 API 调用的用量（未知模型名用 "unknown"）。"""
|
||
|
||
def as_dict(self) -> dict[str, dict[str, int]]:
|
||
"""返回 {model: {"input_tokens": int, "output_tokens": int, "calls": int}}。"""
|
||
|
||
def merge_into(self, existing: dict[str, dict[str, int]]) -> dict[str, dict[str, int]]:
|
||
"""把当前统计合并进已有汇总（用于 session 累加），返回合并后的新 dict。"""
|
||
```
|
||
|
||
一个模块级 `contextvars.ContextVar[TokenUsageTracker | None]`（默认 `None`）保存"当前活跃统计器"。配套：
|
||
|
||
```python
|
||
@contextmanager
|
||
def track_token_usage() -> Iterator[TokenUsageTracker]:
|
||
"""创建并激活一个新 tracker，退出时还原上下文，返回 tracker 供调用方读取汇总。"""
|
||
|
||
def get_current_tracker() -> TokenUsageTracker | None:
|
||
"""供 HTTP 钩子读取；未激活时返回 None（即不统计，用于 /api/score 等排除路径）。"""
|
||
```
|
||
|
||
`contextvars.ContextVar` 天然按 async task / 线程隔离：`session_score_manager` 的线程池并发场景、`MetricPipeline.score_samples()` 内部 `asyncio.gather` 并发场景都能正确工作 —— 同一次 run 内的并发样本共享同一个 tracker（这正是我们想要的：整个 run 的用量汇总到一起），不同 run/请求之间互不干扰。
|
||
|
||
### 3.2 HTTP 层拦截：`attach_usage_hook(client: AsyncOpenAI) -> None`
|
||
|
||
在 `rag_eval/metrics/factory.py` 新增辅助函数，对**已构造好**的 `AsyncOpenAI` 客户端做构造后挂载（不改变现有 `AsyncOpenAI(**kwargs)` 调用点的参数，只在其后追加一行）：
|
||
|
||
```python
|
||
async def _usage_response_hook(response: httpx.Response) -> None:
|
||
tracker = get_current_tracker()
|
||
if tracker is None:
|
||
return
|
||
try:
|
||
await response.aread()
|
||
data = response.json()
|
||
usage = data.get("usage") or {}
|
||
model = data.get("model") or "unknown"
|
||
tracker.record(
|
||
model,
|
||
int(usage.get("prompt_tokens", 0)),
|
||
int(usage.get("completion_tokens", 0)),
|
||
)
|
||
except Exception:
|
||
pass # 可观测性钩子绝不能影响评分主流程
|
||
|
||
|
||
def attach_usage_hook(client: AsyncOpenAI) -> None:
|
||
"""给已构造的 AsyncOpenAI 客户端追加 token 用量响应钩子（幂等，可重复调用）。"""
|
||
httpx_client = client._client # openai SDK 内部持有的 httpx.AsyncClient
|
||
hooks = httpx_client.event_hooks.setdefault("response", [])
|
||
if _usage_response_hook not in hooks:
|
||
hooks.append(_usage_response_hook)
|
||
```
|
||
|
||
调用方式：在 `AsyncOpenAI(**kwargs)` 构造完成后紧跟一行 `attach_usage_hook(llm_client)`。已在当前锁定版本 `openai==1.102.0` 上验证：`AsyncOpenAI()._client` 是 `AsyncHttpxClientWrapper`（`httpx.AsyncClient` 子类），且自带 `event_hooks = {"request": [], "response": []}` 可直接追加，无需额外适配。
|
||
|
||
- 用响应体里网关自己回填的 `model` 字段分组（chat completions 和 embeddings 响应体都带 `model` + `usage`），而不是我们请求时传的模型名 —— 这样即使 judge/embedding 共享同一个 client（`build_models()` 中 base_url/api_key 相同时的既有优化），也能按请求本身的 `model` 字段正确拆开统计。
|
||
- 网关不返回 `usage`/`model`、返回非 JSON、或任何解析异常，一律静默跳过（`except Exception: pass`），钩子失败不能让评分调用报错。
|
||
|
||
### 3.3 挂载点
|
||
|
||
- `rag_eval/metrics/factory.py::build_models()`：构造 `llm_client`、`emb_client` 时都调用 `attach_usage_hook`（对共享 client 场景只需挂一次）。
|
||
- `rag_eval/advisor/llm_analyzer.py::analyze()`：构造自己的 `AsyncOpenAI` 时同样调用 `attach_usage_hook`（从 `rag_eval.metrics.factory` 导入复用，不重复实现)。
|
||
|
||
---
|
||
|
||
## 4. 三条入口路径的接入方式
|
||
|
||
### 4.1 CLI `main.py --scenario`（`rag_eval/execution/evaluator.py`）
|
||
|
||
评测主流程包一层：
|
||
|
||
```python
|
||
with track_token_usage() as tracker:
|
||
# ...现有的 build_metric_pipeline + score_samples 流程...
|
||
result.token_usage = tracker.as_dict()
|
||
```
|
||
|
||
`EvaluationResult` 新增字段：
|
||
|
||
```python
|
||
token_usage: dict[str, dict[str, int]] = field(default_factory=dict)
|
||
```
|
||
|
||
### 4.2 `/api/score/async`（`webapp/services/score_job_manager.py`）
|
||
|
||
`_run()` 里调用 `inline_scorer.score(...)` 的地方外面包 `track_token_usage()`，把汇总结果塞进它构造的 `EvaluationResult.token_usage` 再调 `write_run_artifacts()`。单样本、单次落盘，语义等同 CLI。
|
||
|
||
### 4.3 `/api/score/session_async`（`webapp/services/session_score_manager.py`）—— 累加语义
|
||
|
||
session 场景每次调用都会重写整个 `metadata.json`（现有逻辑：重新汇总全部累计样本行）。Token 用量需要**跨调用累加**而不是只保留最后一次：
|
||
|
||
1. 调用前：`existing = _read_json(metadata_path).get("token_usage", {})`。
|
||
2. 本次调用用 `track_token_usage()` 统计。
|
||
3. 写回前：`merged = tracker.merge_into(existing)`，赋给 `EvaluationResult.token_usage = merged`。
|
||
|
||
这与 `scores.csv` 逐次追加行的既有累加模式保持一致。
|
||
|
||
### 4.4 `write_run_artifacts()`（`rag_eval/reporting/writers.py`）
|
||
|
||
`metadata` dict 新增一个键：
|
||
|
||
```python
|
||
"token_usage": result.token_usage, # {} 表示未统计或统计为空
|
||
```
|
||
|
||
`summary.md`（`rag_eval/reporting/summary.py::build_summary_markdown`）追加一段：
|
||
|
||
```
|
||
## Token 用量
|
||
|
||
| 模型 | input_tokens | output_tokens | 调用次数 |
|
||
|---|---|---|---|
|
||
| gpt-5 | 12450 | 3200 | 60 |
|
||
| Qwen/Qwen3-Embedding-4B | 45000 | 0 | 30 |
|
||
```
|
||
|
||
`token_usage` 为空时输出 `未记录 token 用量。`（保持既有"空数据降级文案"风格，参考 `_table_from_frame` 对空 DataFrame 的处理）。
|
||
|
||
---
|
||
|
||
## 5. 报告层与 Web UI
|
||
|
||
- `webapp/models.py::ReportData` 新增字段：
|
||
```python
|
||
token_usage: dict[str, dict[str, int]] = Field(default_factory=dict)
|
||
```
|
||
- `webapp/services/report_builder.py::build_report()`：从 `run_reader._read_json(run_dir / "metadata.json")` 读出的 `metadata` 里取 `token_usage` 塞进 `ReportData`（此函数已经在读 `metadata.json` 取 `run_id`，顺手多取一个字段即可，无需新增 IO）。
|
||
- `webapp/static/js/report.js`：新增 `renderTokenUsage(report)`，仿照 `renderGroupings` 的表格渲染方式，在报告详情页新增一个"Token 用量"面板：按模型一行，列为 input/output/调用次数；`token_usage` 为空时显示"暂无 token 用量数据"。对应在 `index.html` 加一个面板容器 `token-usage-wrap`，在 `Report.render()` 里追加调用。
|
||
|
||
---
|
||
|
||
## 6. 错误处理与边界情况
|
||
|
||
| 情况 | 处理 |
|
||
|---|---|
|
||
| 网关响应无 `usage` 字段 | 钩子静默跳过，不记录、不报错 |
|
||
| 网关响应非 JSON / 读取异常 | `except Exception: pass`，钩子绝不向上抛异常 |
|
||
| 未激活 tracker（如未来有调用方忘记包 `track_token_usage()`，或 `/api/score` 这类明确排除的路径） | `get_current_tracker()` 返回 `None`，钩子直接返回，等同于关闭统计 |
|
||
| `judge_model == embedding_model` 或共享同一个 `AsyncOpenAI` 客户端（`build_models()` 现有优化） | 用响应体 `model` 字段区分，两个模型各自正确累加，无需为共享客户端特殊处理 |
|
||
| 旧的 `metadata.json`（无 `token_usage` 键，历史 run） | `report_builder` 用 `metadata.get("token_usage", {})`，前端空数据降级展示，不报错 |
|
||
| session_async 并发多个不同 session | 各自线程内 `contextvars` 独立，互不干扰；同一 session 的 `metadata.json` 读-改-写已有per-session锁（`threading.Lock`）保护，新增的 `token_usage` 合并逻辑复用这把锁 |
|
||
|
||
---
|
||
|
||
## 7. 测试策略
|
||
|
||
新增/修改测试文件：
|
||
|
||
- `tests/test_token_tracker.py`：`TokenUsageTracker.record/as_dict/merge_into` 纯函数单测；`track_token_usage()` 上下文管理器激活/还原行为；未激活时 `get_current_tracker()` 返回 `None`。
|
||
- `tests/test_token_usage_hook.py`：用假的 `httpx.Response`（或轻量 mock transport）验证 `_usage_response_hook` 在有/无 `usage` 字段、非 JSON 响应时的行为；验证 `attach_usage_hook` 正确挂载到 `AsyncOpenAI` 客户端。
|
||
- `tests/test_evaluator.py` 或现有 `test_offline_eval.py` / `test_online_eval.py`：扩展断言 `EvaluationResult.token_usage` 字段存在且结构正确（mock LLM 客户端返回带 `usage` 的响应）。
|
||
- `tests/webapp/test_score_jobs_api.py`：断言 `/api/score/async` 产出的 `metadata.json` 含 `token_usage`。
|
||
- `tests/webapp/test_session_score_jobs_api.py`：断言连续两次调用同一 session，`token_usage` 按模型**累加**而非覆盖。
|
||
- `tests/test_webapp_report_builder.py`：断言 `ReportData.token_usage` 从 `metadata.json` 正确透传；`metadata.json` 缺失该键时默认空 dict。
|
||
|
||
全部使用 mock LLM 客户端 / 假 httpx 响应，不依赖真实网络调用（遵循仓库现有测试约定）。
|
||
|
||
---
|
||
|
||
## 8. 非目标（本轮明确不做）
|
||
|
||
- 不做金额（$/¥）换算，不改 LLM Profile 增加单价字段（用户已确认后续再说）。
|
||
- 不统计不落盘的同步 `/api/score` 接口。
|
||
- 不统计 `dataset_builder` 题库生成路径的 token 用量（范围已与用户确认排除）。
|
||
- 不做事前成本预估/dry-run 采样估算（用户已确认只要事后按模型累计的原始计数）。
|
||
- 不改变现有评分逻辑、指标计算或既有 API 响应结构（这是纯增量的可观测性能力）。
|