diff --git a/dspy/utils/usage_tracker.py b/dspy/utils/usage_tracker.py index fa3c8f7447..2643645a70 100644 --- a/dspy/utils/usage_tracker.py +++ b/dspy/utils/usage_tracker.py @@ -37,12 +37,11 @@ def _merge_usage_entries(self, usage_entry1, usage_entry2) -> dict[str, dict[str result = dict(usage_entry2) for k, v in usage_entry1.items(): - if k in result: - if isinstance(v, dict): - result[k] = self._merge_usage_entries(result[k], v) - else: - result[k] = result[k] or 0 - result[k] += v if v else 0 + current_v = result.get(k) + if isinstance(v, dict): + result[k] = self._merge_usage_entries(current_v, v) + else: + result[k] = (current_v or 0) + (v or 0) return result def add_usage(self, lm: str, usage_entry: dict): diff --git a/tests/utils/test_usage_tracker.py b/tests/utils/test_usage_tracker.py index 2ce9211479..86b2d9ba56 100644 --- a/tests/utils/test_usage_tracker.py +++ b/tests/utils/test_usage_tracker.py @@ -157,3 +157,16 @@ def test_track_usage_context_manager(): assert "openai/gpt-4o-mini" in total_usage assert len(total_usage.keys()) == 1 assert isinstance(total_usage["openai/gpt-4o-mini"], dict) + + +def test_merge_usage_entries_with_new_keys(): + """Ensure merging usage entries preserves unseen keys.""" + tracker = UsageTracker() + + tracker.add_usage("model-x", {"prompt_tokens": 5}) + tracker.add_usage("model-x", {"completion_tokens": 2}) + + total_usage = tracker.get_total_tokens() + + assert total_usage["model-x"]["prompt_tokens"] == 5 + assert total_usage["model-x"]["completion_tokens"] == 2