From e364e30b3c2f7fa70f9c28fb6b3948945f3083a4 Mon Sep 17 00:00:00 2001 From: Earl St Sauver Date: Sun, 18 May 2025 15:47:51 +0200 Subject: [PATCH 1/2] Bugfix: Ensure keys in both entities are merged This PR fixes asymmetric merging in _merge_usage_entries so keys existing only in the first entry are now preserved. --- dspy/utils/usage_tracker.py | 12 ++++++------ tests/utils/test_usage_tracker.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/dspy/utils/usage_tracker.py b/dspy/utils/usage_tracker.py index fa3c8f7447..0b845ad020 100644 --- a/dspy/utils/usage_tracker.py +++ b/dspy/utils/usage_tracker.py @@ -37,12 +37,12 @@ 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 + result[k] += v if v else 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 From 79f5cefd65f10f4293dccc41933dc039649f9bd2 Mon Sep 17 00:00:00 2001 From: Earl St Sauver Date: Tue, 27 May 2025 07:57:42 +0200 Subject: [PATCH 2/2] Update dspy/utils/usage_tracker.py Accepted copilot suggestion Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- dspy/utils/usage_tracker.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dspy/utils/usage_tracker.py b/dspy/utils/usage_tracker.py index 0b845ad020..2643645a70 100644 --- a/dspy/utils/usage_tracker.py +++ b/dspy/utils/usage_tracker.py @@ -41,8 +41,7 @@ def _merge_usage_entries(self, usage_entry1, usage_entry2) -> dict[str, dict[str if isinstance(v, dict): result[k] = self._merge_usage_entries(current_v, v) else: - result[k] = current_v or 0 - result[k] += v if v else 0 + result[k] = (current_v or 0) + (v or 0) return result def add_usage(self, lm: str, usage_entry: dict):