feat(token-tracking): capture token usage in /api/score/async job runs

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-02 15:03:25 +08:00
co-authored by Copilot
parent 5494840431
commit 6a2bbf8239
3 changed files with 94 additions and 18 deletions
@@ -0,0 +1,73 @@
"""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.inline_scorer.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.inline_scorer.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"] == {}