diff --git a/dspy/utils/usage_tracker.py b/dspy/utils/usage_tracker.py index 1d3d13c638..e724e1de0e 100644 --- a/dspy/utils/usage_tracker.py +++ b/dspy/utils/usage_tracker.py @@ -45,7 +45,7 @@ def _merge_usage_entries( current_v = result.get(k) if isinstance(v, dict) or isinstance(current_v, dict): result[k] = self._merge_usage_entries(current_v, v) - else: + elif current_v is not None or v is not None: result[k] = (current_v or 0) + (v or 0) return result diff --git a/tests/utils/test_usage_tracker.py b/tests/utils/test_usage_tracker.py index 28e35e40ab..f3ae612738 100644 --- a/tests/utils/test_usage_tracker.py +++ b/tests/utils/test_usage_tracker.py @@ -190,6 +190,16 @@ def test_merge_usage_entries_with_none_values(): "prompt_tokens": 800, "completion_tokens": 100, "total_tokens": 900, + "prompt_tokens_details": None, + "completion_tokens_details": None, + }, + }, + { + "model": "gpt-4o-mini", + "usage": { + "prompt_tokens": 100, + "completion_tokens": 100, + "total_tokens": 200, "prompt_tokens_details": {"cached_tokens": 50, "audio_tokens": 50}, "completion_tokens_details": None, }, @@ -216,9 +226,9 @@ def test_merge_usage_entries_with_none_values(): total_usage = tracker.get_total_tokens() - assert total_usage["gpt-4o-mini"]["prompt_tokens"] == 2717 - assert total_usage["gpt-4o-mini"]["completion_tokens"] == 246 - assert total_usage["gpt-4o-mini"]["total_tokens"] == 2963 + assert total_usage["gpt-4o-mini"]["prompt_tokens"] == 2817 + assert total_usage["gpt-4o-mini"]["completion_tokens"] == 346 + assert total_usage["gpt-4o-mini"]["total_tokens"] == 3163 assert total_usage["gpt-4o-mini"]["prompt_tokens_details"]["cached_tokens"] == 50 assert total_usage["gpt-4o-mini"]["prompt_tokens_details"]["audio_tokens"] == 50 assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["reasoning_tokens"] == 1