84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
|
"""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}
|
||
|
|
}
|