Skip to content

Commit 99cb334

Browse files
committed
Revert uv.lock changes
1 parent 5e17bcf commit 99cb334

File tree

2 files changed

+15
-0
lines changed

2 files changed

+15
-0
lines changed

dspy/utils/usage_tracker.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ def _merge_usage_entries(self, usage_entry1, usage_entry2) -> dict[str, dict[str
4343
else:
4444
result[k] = result[k] or 0
4545
result[k] += v if v else 0
46+
else:
47+
result[k] = dict(v) if isinstance(v, dict) else v
4648
return result
4749

4850
def add_usage(self, lm: str, usage_entry: dict):

tests/utils/test_usage_tracker.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,3 +157,16 @@ def test_track_usage_context_manager():
157157
assert "openai/gpt-4o-mini" in total_usage
158158
assert len(total_usage.keys()) == 1
159159
assert isinstance(total_usage["openai/gpt-4o-mini"], dict)
160+
161+
162+
def test_merge_usage_entries_with_new_keys():
163+
"""Ensure merging usage entries preserves unseen keys."""
164+
tracker = UsageTracker()
165+
166+
tracker.add_usage("model-x", {"prompt_tokens": 5})
167+
tracker.add_usage("model-x", {"completion_tokens": 2})
168+
169+
total_usage = tracker.get_total_tokens()
170+
171+
assert total_usage["model-x"]["prompt_tokens"] == 5
172+
assert total_usage["model-x"]["completion_tokens"] == 2

0 commit comments

Comments
 (0)