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

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-02 14:31:43 +08:00
parent 54cc58e157
commit 8b896e4e7f
2 changed files with 165 additions and 0 deletions
@@ -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}
+    }