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:
@@ -1190,7 +1190,7 @@ def test_run_writes_token_usage_to_metadata(tmp_path, monkeypatch):
|
|||||||
return {m: 0.9 for m in kwargs["metrics"]}
|
return {m: 0.9 for m in kwargs["metrics"]}
|
||||||
|
|
||||||
monkeypatch.setattr(
|
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"])
|
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"]}
|
return {m: 0.9 for m in kwargs["metrics"]}
|
||||||
|
|
||||||
monkeypatch.setattr(
|
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"])
|
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"]}
|
return {m: 0.9 for m in kwargs["metrics"]}
|
||||||
|
|
||||||
monkeypatch.setattr(
|
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"])
|
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"]}
|
return {m: 0.9 for m in kwargs["metrics"]}
|
||||||
|
|
||||||
monkeypatch.setattr(
|
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"])
|
request = ScoreRequest(question="q?", answer="a.", metrics=["answer_relevancy"])
|
||||||
|
|||||||
@@ -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"] == {}
|
||||||
@@ -108,6 +108,7 @@ class ScoreJobManager:
|
|||||||
|
|
||||||
# Lazy imports to keep web server bootable if ragas is not installed.
|
# Lazy imports to keep web server bootable if ragas is not installed.
|
||||||
from rag_eval.advisor import run_advisor
|
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.metrics.weights import compute_weighted_score
|
||||||
from rag_eval.reporting.writers import write_run_artifacts
|
from rag_eval.reporting.writers import write_run_artifacts
|
||||||
from rag_eval.settings import EvaluationSettings
|
from rag_eval.settings import EvaluationSettings
|
||||||
@@ -130,6 +131,7 @@ class ScoreJobManager:
|
|||||||
started_at = utc_now_iso()
|
started_at = utc_now_iso()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
with track_token_usage() as usage_tracker:
|
||||||
if effective:
|
if effective:
|
||||||
raw_scores = inline_scorer.score(
|
raw_scores = inline_scorer.score(
|
||||||
question=request.question,
|
question=request.question,
|
||||||
@@ -201,6 +203,7 @@ class ScoreJobManager:
|
|||||||
valid_samples=[sample],
|
valid_samples=[sample],
|
||||||
invalid_samples=[],
|
invalid_samples=[],
|
||||||
score_rows=[score_row],
|
score_rows=[score_row],
|
||||||
|
token_usage=usage_tracker.as_dict(),
|
||||||
)
|
)
|
||||||
|
|
||||||
write_run_artifacts(result)
|
write_run_artifacts(result)
|
||||||
|
|||||||
Reference in New Issue
Block a user