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:
@@ -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,6 +68,7 @@ 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()
|
||||||
|
with track_token_usage() as usage_tracker:
|
||||||
metric_scores = asyncio.run(
|
metric_scores = asyncio.run(
|
||||||
self.metric_pipeline.score_samples(
|
self.metric_pipeline.score_samples(
|
||||||
samples,
|
samples,
|
||||||
@@ -75,6 +77,7 @@ class Evaluator:
|
|||||||
)
|
)
|
||||||
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(
|
||||||
|
|||||||
@@ -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, {})
|
||||||
Reference in New Issue
Block a user