Bugfix: Ensure keys in both entities are merged

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):

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