feat(token-tracking): wrap CLI evaluator metric scoring in track_token_usage

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-02 14:57:19 +08:00
co-authored by Copilot
parent cd044c8416
commit 5494840431
2 changed files with 138 additions and 5 deletions
+9 -5
View File
@@ -12,6 +12,7 @@ from rag_eval.datasets.loader import load_dataset_records
from rag_eval.datasets.normalizers import normalize_records from rag_eval.datasets.normalizers import normalize_records
from rag_eval.execution.concurrency import gather_with_limit from rag_eval.execution.concurrency import gather_with_limit
from rag_eval.metrics.pipeline import MetricPipeline from rag_eval.metrics.pipeline import MetricPipeline
from rag_eval.metrics.token_tracker import track_token_usage
from rag_eval.metrics.weights import compute_weighted_score, resolve_weight from rag_eval.metrics.weights import compute_weighted_score, resolve_weight
from rag_eval.shared.models import EvaluationResult, InvalidSample, NormalizedSample, Scenario from rag_eval.shared.models import EvaluationResult, InvalidSample, NormalizedSample, Scenario
from rag_eval.shared.utils import utc_now_iso from rag_eval.shared.utils import utc_now_iso
@@ -67,14 +68,16 @@ class Evaluator:
logger.info("[eval] scoring %d samples with metric pipeline ...", len(samples)) logger.info("[eval] scoring %d samples with metric pipeline ...", len(samples))
t0 = time.monotonic() t0 = time.monotonic()
metric_scores = asyncio.run( with track_token_usage() as usage_tracker:
self.metric_pipeline.score_samples( metric_scores = asyncio.run(
samples, self.metric_pipeline.score_samples(
max_concurrency=self.scenario.runtime.metric_limit(), samples,
max_concurrency=self.scenario.runtime.metric_limit(),
)
) )
)
elapsed = time.monotonic() - t0 elapsed = time.monotonic() - t0
logger.info("[eval] metric scoring done elapsed=%.1fs", elapsed) logger.info("[eval] metric scoring done elapsed=%.1fs", elapsed)
logger.info("[eval] token_usage=%s", usage_tracker.as_dict())
finished_at = utc_now_iso() finished_at = utc_now_iso()
score_rows = [self._merge_score(sample, score) for sample, score in zip(samples, metric_scores)] score_rows = [self._merge_score(sample, score) for sample, score in zip(samples, metric_scores)]
@@ -99,6 +102,7 @@ class Evaluator:
valid_samples=samples, valid_samples=samples,
invalid_samples=invalid_samples, invalid_samples=invalid_samples,
score_rows=score_rows, score_rows=score_rows,
token_usage=usage_tracker.as_dict(),
) )
async def _enrich_online_samples( async def _enrich_online_samples(
+129
View File
@@ -0,0 +1,129 @@
"""Tests verifying the CLI evaluation flow captures token usage from metric scoring."""
from __future__ import annotations
import shutil
import unittest
from pathlib import Path
import pandas as pd
from rag_eval.execution.evaluator import Evaluator
from rag_eval.metrics.pipeline import MetricPipeline
from rag_eval.metrics.token_tracker import get_current_tracker
from rag_eval.shared.models import DatasetConfig, RuntimeConfig, Scenario
class FakeMetricWithUsage:
"""Fake RAGAS metric that records token usage like a real HTTP-hooked call would."""
def __init__(self, value: float, model: str, input_tokens: int, output_tokens: int):
self.value = value
self.model = model
self.input_tokens = input_tokens
self.output_tokens = output_tokens
async def ascore(self, **kwargs):
tracker = get_current_tracker()
if tracker is not None:
tracker.record(self.model, self.input_tokens, self.output_tokens)
class Result:
def __init__(self, value: float):
self.value = value
return Result(self.value)
class PlainFakeMetric:
"""Fake metric that never records usage (simulates a hook that captured nothing)."""
async def ascore(self, **kwargs):
class Result:
value = 0.9
return Result()
class EvaluatorTokenUsageTests(unittest.TestCase):
def setUp(self) -> None:
root = Path("tests/.tmp").resolve()
root.mkdir(parents=True, exist_ok=True)
self.temp_dir = root / self._testMethodName
shutil.rmtree(self.temp_dir, ignore_errors=True)
self.temp_dir.mkdir(parents=True, exist_ok=True)
def tearDown(self) -> None:
shutil.rmtree(self.temp_dir, ignore_errors=True)
def _write_offline_dataset(self, path: Path, rows: list[dict]) -> None:
pd.DataFrame(rows).to_csv(path, index=False)
def test_evaluate_populates_token_usage_from_metric_calls(self) -> None:
dataset_path = self.temp_dir / "offline.csv"
self._write_offline_dataset(dataset_path, [
{
"sample_id": "sample-1",
"question": "What is the policy scope?",
"answer": "It covers all employees.",
"contexts": '["Context A"]',
"ground_truth": "It covers all employees.",
},
{
"sample_id": "sample-2",
"question": "What about contractors?",
"answer": "Contractors are excluded.",
"contexts": '["Context B"]',
"ground_truth": "Contractors are excluded.",
},
])
scenario = Scenario(
scenario_name="token-usage-test",
mode="offline",
dataset=DatasetConfig(path=dataset_path),
judge_model="gpt-5",
embedding_model="embedding-model",
metrics=["faithfulness"],
output_dir=self.temp_dir / "outputs",
runtime=RuntimeConfig(batch_size=1),
)
pipeline = MetricPipeline(
metrics={"faithfulness": FakeMetricWithUsage(0.8, "gpt-5", 100, 40)}
)
evaluator = Evaluator(scenario=scenario, metric_pipeline=pipeline)
result = evaluator.evaluate()
# Two samples each recorded one call → totals sum across both.
self.assertEqual(
result.token_usage,
{"gpt-5": {"input_tokens": 200, "output_tokens": 80, "calls": 2}},
)
def test_evaluate_defaults_to_empty_token_usage_when_nothing_recorded(self) -> None:
dataset_path = self.temp_dir / "offline.csv"
self._write_offline_dataset(dataset_path, [
{
"sample_id": "sample-1",
"question": "What is the policy scope?",
"answer": "It covers all employees.",
"contexts": '["Context A"]',
"ground_truth": "It covers all employees.",
},
])
scenario = Scenario(
scenario_name="token-usage-empty-test",
mode="offline",
dataset=DatasetConfig(path=dataset_path),
judge_model="gpt-5",
embedding_model="embedding-model",
metrics=["faithfulness"],
output_dir=self.temp_dir / "outputs",
runtime=RuntimeConfig(batch_size=1),
)
pipeline = MetricPipeline(metrics={"faithfulness": PlainFakeMetric()})
evaluator = Evaluator(scenario=scenario, metric_pipeline=pipeline)
result = evaluator.evaluate()
self.assertEqual(result.token_usage, {})