feat(token-tracking): add EvaluationResult.token_usage and persist to metadata.json

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 14:45:58 +08:00
parent 29e76d021c
commit b50b911de0
3 changed files with 79 additions and 0 deletions
@@ -0,0 +1,75 @@
+"""Tests that EvaluationResult.token_usage is persisted into metadata.json."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from rag_eval.reporting.writers import write_run_artifacts
+from rag_eval.shared.models import DatasetConfig, EvaluationResult, RuntimeConfig, Scenario
+
+
+def _scenario(tmp_path: Path) -> Scenario:
+    return Scenario(
+        scenario_name="token-persist-test",
+        mode="offline",
+        dataset=DatasetConfig(path=tmp_path / "dataset.csv"),
+        judge_model="gpt-5",
+        embedding_model="embedding-model",
+        metrics=["faithfulness"],
+        output_dir=tmp_path / "outputs",
+        runtime=RuntimeConfig(batch_size=1),
+    )
+
+
+def test_evaluation_result_defaults_token_usage_to_empty_dict(tmp_path: Path) -> None:
+    result = EvaluationResult(
+        scenario=_scenario(tmp_path),
+        run_id="run-1",
+        started_at="t0",
+        finished_at="t1",
+        valid_samples=[],
+        invalid_samples=[],
+        score_rows=[],
+    )
+    assert result.token_usage == {}
+
+
+def test_write_run_artifacts_persists_token_usage(tmp_path: Path) -> None:
+    scenario = _scenario(tmp_path)
+    result = EvaluationResult(
+        scenario=scenario,
+        run_id="run-2",
+        started_at="t0",
+        finished_at="t1",
+        valid_samples=[],
+        invalid_samples=[],
+        score_rows=[],
+        token_usage={"gpt-5": {"input_tokens": 100, "output_tokens": 40, "calls": 2}},
+    )
+
+    write_run_artifacts(result)
+
+    metadata_path = scenario.output_dir / "run-2" / "metadata.json"
+    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    assert metadata["token_usage"] == {
+        "gpt-5": {"input_tokens": 100, "output_tokens": 40, "calls": 2}
+    }
+
+
+def test_write_run_artifacts_writes_empty_token_usage_when_unset(tmp_path: Path) -> None:
+    scenario = _scenario(tmp_path)
+    result = EvaluationResult(
+        scenario=scenario,
+        run_id="run-3",
+        started_at="t0",
+        finished_at="t1",
+        valid_samples=[],
+        invalid_samples=[],
+        score_rows=[],
+    )
+
+    write_run_artifacts(result)
+
+    metadata_path = scenario.output_dir / "run-3" / "metadata.json"
+    metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+    assert metadata["token_usage"] == {}