From 8b896e4e7f6e349f9654fd64a40bda386697b675 Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 2 Jul 2026 14:31:43 +0800 Subject: [PATCH] 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 | 82 ++++++++++++++++++++++++++++++ tests/test_token_tracker.py | 83 +++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 rag_eval/metrics/token_tracker.py create mode 100644 tests/test_token_tracker.py diff --git a/rag_eval/metrics/token_tracker.py b/rag_eval/metrics/token_tracker.py new file mode 100644 index 0000000..0c112a7 --- /dev/null +++ b/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() diff --git a/tests/test_token_tracker.py b/tests/test_token_tracker.py new file mode 100644 index 0000000..1a0b318 --- /dev/null +++ b/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} + }