feat(token-tracking): add TokenUsageTracker with context-scoped activation

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

rag_eval/metrics/token_tracker.py

@@ -0,0 +1,82 @@
+"""Per-run token usage accumulation, keyed by model name.
+
+RAGAS 0.4.3's `ragas.metrics.collections` + instructor code path does not
+expose real token counts (`ragas/cost.py` only serves the legacy langchain
+`evaluate()` path). This module provides a context-scoped accumulator that
+the HTTP response hook in `rag_eval.metrics.factory` feeds into, so token
+counts survive across the AsyncOpenAI client caching used by InlineScorer
+(see docs/superpowers/specs/2026-07-02-token-usage-tracking-design.md).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import Iterator
+
+
+@dataclass
+class TokenUsageTracker:
+    """Accumulates input/output token counts and call counts, grouped by model name."""
+
+    _totals: dict[str, dict[str, int]] = field(default_factory=dict)
+
+    def record(self, model: str, input_tokens: int, output_tokens: int) -> None:
+        """Add one API call's usage to the running total for `model`."""
+        key = model or "unknown"
+        bucket = self._totals.setdefault(
+            key, {"input_tokens": 0, "output_tokens": 0, "calls": 0}
+        )
+        bucket["input_tokens"] += int(input_tokens)
+        bucket["output_tokens"] += int(output_tokens)
+        bucket["calls"] += 1
+
+    def as_dict(self) -> dict[str, dict[str, int]]:
+        """Return a plain-dict snapshot: {model: {input_tokens, output_tokens, calls}}."""
+        return {model: dict(usage) for model, usage in self._totals.items()}
+
+    def merge_into(self, existing: dict[str, dict[str, int]]) -> dict[str, dict[str, int]]:
+        """Return a new dict combining `existing` accumulated totals with this tracker's totals.
+
+        Used by session-scoped scoring (one call at a time) to keep a running
+        total across multiple calls instead of overwriting with just the latest call.
+        Does not mutate `existing`.
+        """
+        merged: dict[str, dict[str, int]] = {
+            model: dict(usage) for model, usage in existing.items()
+        }
+        for model, usage in self.as_dict().items():
+            bucket = merged.setdefault(
+                model, {"input_tokens": 0, "output_tokens": 0, "calls": 0}
+            )
+            bucket["input_tokens"] += usage["input_tokens"]
+            bucket["output_tokens"] += usage["output_tokens"]
+            bucket["calls"] += usage["calls"]
+        return merged
+
+
+_current_tracker: ContextVar[TokenUsageTracker | None] = ContextVar(
+    "_current_tracker", default=None
+)
+
+
+@contextmanager
+def track_token_usage() -> Iterator[TokenUsageTracker]:
+    """Activate a fresh TokenUsageTracker for the duration of the `with` block.
+
+    Any AsyncOpenAI client with `attach_usage_hook()` applied that makes a
+    call while this context is active will have its usage recorded here.
+    Safe to nest; the innermost tracker is active within its own block.
+    """
+    tracker = TokenUsageTracker()
+    token = _current_tracker.set(tracker)
+    try:
+        yield tracker
+    finally:
+        _current_tracker.reset(token)
+
+
+def get_current_tracker() -> TokenUsageTracker | None:
+    """Return the currently active tracker, or None if no `track_token_usage()` block is active."""
+    return _current_tracker.get()

tests/test_token_tracker.py

@@ -0,0 +1,83 @@
+"""Tests for the per-run token usage accumulator and its context-scoped activation."""
+from __future__ import annotations
+
+from rag_eval.metrics.token_tracker import (
+    TokenUsageTracker,
+    get_current_tracker,
+    track_token_usage,
+)
+
+
+def test_record_accumulates_input_output_and_calls():
+    tracker = TokenUsageTracker()
+    tracker.record("gpt-5", 100, 50)
+    tracker.record("gpt-5", 20, 10)
+    assert tracker.as_dict() == {
+        "gpt-5": {"input_tokens": 120, "output_tokens": 60, "calls": 2}
+    }
+
+
+def test_record_groups_by_model_name():
+    tracker = TokenUsageTracker()
+    tracker.record("gpt-5", 100, 50)
+    tracker.record("Qwen/Qwen3-Embedding-4B", 30, 0)
+    result = tracker.as_dict()
+    assert set(result.keys()) == {"gpt-5", "Qwen/Qwen3-Embedding-4B"}
+    assert result["Qwen/Qwen3-Embedding-4B"] == {
+        "input_tokens": 30, "output_tokens": 0, "calls": 1
+    }
+
+
+def test_record_defaults_blank_model_to_unknown():
+    tracker = TokenUsageTracker()
+    tracker.record("", 10, 5)
+    assert "unknown" in tracker.as_dict()
+
+
+def test_merge_into_sums_with_existing_totals():
+    tracker = TokenUsageTracker()
+    tracker.record("gpt-5", 100, 50)
+    existing = {"gpt-5": {"input_tokens": 200, "output_tokens": 100, "calls": 3}}
+    merged = tracker.merge_into(existing)
+    assert merged == {"gpt-5": {"input_tokens": 300, "output_tokens": 150, "calls": 4}}
+
+
+def test_merge_into_keeps_models_only_in_existing():
+    tracker = TokenUsageTracker()
+    tracker.record("gpt-5", 10, 5)
+    existing = {"other-model": {"input_tokens": 1, "output_tokens": 1, "calls": 1}}
+    merged = tracker.merge_into(existing)
+    assert merged["other-model"] == {"input_tokens": 1, "output_tokens": 1, "calls": 1}
+    assert merged["gpt-5"] == {"input_tokens": 10, "output_tokens": 5, "calls": 1}
+
+
+def test_merge_into_does_not_mutate_existing_dict():
+    tracker = TokenUsageTracker()
+    tracker.record("gpt-5", 10, 5)
+    existing = {"gpt-5": {"input_tokens": 1, "output_tokens": 1, "calls": 1}}
+    tracker.merge_into(existing)
+    assert existing == {"gpt-5": {"input_tokens": 1, "output_tokens": 1, "calls": 1}}
+
+
+def test_get_current_tracker_returns_none_outside_context():
+    assert get_current_tracker() is None
+
+
+def test_track_token_usage_activates_and_resets_context():
+    assert get_current_tracker() is None
+    with track_token_usage() as tracker:
+        assert get_current_tracker() is tracker
+        tracker.record("gpt-5", 1, 1)
+    assert get_current_tracker() is None
+
+
+def test_track_token_usage_nested_contexts_are_isolated():
+    with track_token_usage() as outer:
+        outer.record("outer-model", 5, 5)
+        with track_token_usage() as inner:
+            inner.record("inner-model", 1, 1)
+            assert get_current_tracker() is inner
+        assert get_current_tracker() is outer
+    assert outer.as_dict() == {
+        "outer-model": {"input_tokens": 5, "output_tokens": 5, "calls": 1}
+    }