"""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, {})