diff --git a/rag_eval/execution/evaluator.py b/rag_eval/execution/evaluator.py index 5044730..6983e0e 100644 --- a/rag_eval/execution/evaluator.py +++ b/rag_eval/execution/evaluator.py @@ -12,6 +12,7 @@ from rag_eval.datasets.loader import load_dataset_records from rag_eval.datasets.normalizers import normalize_records from rag_eval.execution.concurrency import gather_with_limit 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.shared.models import EvaluationResult, InvalidSample, NormalizedSample, Scenario 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)) t0 = time.monotonic() - metric_scores = asyncio.run( - self.metric_pipeline.score_samples( - samples, - max_concurrency=self.scenario.runtime.metric_limit(), + with track_token_usage() as usage_tracker: + metric_scores = asyncio.run( + self.metric_pipeline.score_samples( + samples, + max_concurrency=self.scenario.runtime.metric_limit(), + ) ) - ) elapsed = time.monotonic() - t0 logger.info("[eval] metric scoring done elapsed=%.1fs", elapsed) + logger.info("[eval] token_usage=%s", usage_tracker.as_dict()) finished_at = utc_now_iso() score_rows = [self._merge_score(sample, score) for sample, score in zip(samples, metric_scores)] @@ -99,6 +102,7 @@ class Evaluator: valid_samples=samples, invalid_samples=invalid_samples, score_rows=score_rows, + token_usage=usage_tracker.as_dict(), ) async def _enrich_online_samples( diff --git a/tests/test_evaluator_token_usage.py b/tests/test_evaluator_token_usage.py new file mode 100644 index 0000000..a0b8591 --- /dev/null +++ b/tests/test_evaluator_token_usage.py @@ -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, {})