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 响应结构(这是纯增量的可观测性能力)。
|