siemens_ragas/tests/test_evaluator_token_usage.py

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