130 lines
4.5 KiB
Python
130 lines
4.5 KiB
Python
"""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, {})
|