From 6a2bbf82394d5f9eedb136037429a1a938b1ccf9 Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 2 Jul 2026 15:03:25 +0800 Subject: [PATCH] feat(token-tracking): capture token usage in /api/score/async job runs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../plans/2026-07-02-token-usage-tracking.md | 8 +- .../test_score_job_manager_token_usage.py | 73 +++++++++++++++++++ webapp/services/score_job_manager.py | 31 ++++---- 3 files changed, 94 insertions(+), 18 deletions(-) create mode 100644 tests/webapp/test_score_job_manager_token_usage.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 index 448e9d8..8259531 100644 --- a/docs/superpowers/plans/2026-07-02-token-usage-tracking.md +++ b/docs/superpowers/plans/2026-07-02-token-usage-tracking.md @@ -1190,7 +1190,7 @@ def test_run_writes_token_usage_to_metadata(tmp_path, monkeypatch): return {m: 0.9 for m in kwargs["metrics"]} monkeypatch.setattr( - "webapp.services.score_job_manager.inline_scorer.score", _fake_score + "webapp.services.inline_scorer.inline_scorer.score", _fake_score ) request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) @@ -1216,7 +1216,7 @@ def test_run_writes_empty_token_usage_when_nothing_recorded(tmp_path, monkeypatc return {m: 0.9 for m in kwargs["metrics"]} monkeypatch.setattr( - "webapp.services.score_job_manager.inline_scorer.score", _fake_score + "webapp.services.inline_scorer.inline_scorer.score", _fake_score ) request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) @@ -1399,7 +1399,7 @@ def test_session_accumulates_token_usage_across_calls(tmp_path, monkeypatch): return {m: 0.9 for m in kwargs["metrics"]} monkeypatch.setattr( - "webapp.services.session_score_manager.inline_scorer.score", _fake_score + "webapp.services.inline_scorer.inline_scorer.score", _fake_score ) request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) @@ -1432,7 +1432,7 @@ def test_session_first_call_writes_token_usage_from_scratch(tmp_path, monkeypatc return {m: 0.9 for m in kwargs["metrics"]} monkeypatch.setattr( - "webapp.services.session_score_manager.inline_scorer.score", _fake_score + "webapp.services.inline_scorer.inline_scorer.score", _fake_score ) request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"]) diff --git a/tests/webapp/test_score_job_manager_token_usage.py b/tests/webapp/test_score_job_manager_token_usage.py new file mode 100644 index 0000000..5b138dd --- /dev/null +++ b/tests/webapp/test_score_job_manager_token_usage.py @@ -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"] == {} diff --git a/webapp/services/score_job_manager.py b/webapp/services/score_job_manager.py index 6c9de9e..b16d541 100644 --- a/webapp/services/score_job_manager.py +++ b/webapp/services/score_job_manager.py @@ -108,6 +108,7 @@ class ScoreJobManager: # Lazy imports to keep web server bootable if ragas is not installed. 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 @@ -130,20 +131,21 @@ class ScoreJobManager: started_at = utc_now_iso() 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 = {} + 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 = {} latency_ms = int((time.monotonic() - t0) * 1000) finished_at = utc_now_iso() @@ -201,6 +203,7 @@ class ScoreJobManager: valid_samples=[sample], invalid_samples=[], score_rows=[score_row], + token_usage=usage_tracker.as_dict(), ) write_run_artifacts(result)