83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
|
"""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()
|