Fix TypeError when tracking usage with Anthropic models returning Pydantic objects (#8978)

* Initial plan

* Fix TypeError when merging Anthropic CacheCreation objects in usage tracker

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>

* Enhance _flatten_usage_entry to convert Pydantic models on first add

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>

* Fix potential TypeError when both usage entries are None

Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>

* simplify

* small fix

* lint

* robust version handling

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: TomeHirata <33407409+TomeHirata@users.noreply.github.com>
Co-authored-by: chenmoneygithub <chen.qian@databricks.com>

Author: 3 people

dspy/utils/usage_tracker.py

@@ -4,6 +4,8 @@
 from contextlib import contextmanager
 from typing import Any, Generator
 
+from pydantic import BaseModel
+
 from dspy.dsp.utils.settings import settings
 
 
@@ -21,15 +23,18 @@ def __init__(self):
         self.usage_data = defaultdict(list)
 
     def _flatten_usage_entry(self, usage_entry: dict[str, Any]) -> dict[str, Any]:
-        result = dict(usage_entry)
-
-        if completion_tokens_details := result.get("completion_tokens_details"):
-            result["completion_tokens_details"] = dict(completion_tokens_details)
-        if prompt_tokens_details := result.get("prompt_tokens_details"):
-            result["prompt_tokens_details"] = dict(prompt_tokens_details)
+        result = {}
+        for key, value in usage_entry.items():
+            if isinstance(value, BaseModel):
+                # Convert Pydantic models to dicts, like `PromptTokensDetailsWrapper` from litellm.
+                result[key] = value.model_dump()
+            else:
+                result[key] = value
         return result
 
-    def _merge_usage_entries(self, usage_entry1: dict[str, Any] | None, usage_entry2: dict[str, Any] | None) -> dict[str, Any]:
+    def _merge_usage_entries(
+        self, usage_entry1: dict[str, Any] | None, usage_entry2: dict[str, Any] | None
+    ) -> dict[str, Any]:
         if usage_entry1 is None or len(usage_entry1) == 0:
             return dict(usage_entry2)
         if usage_entry2 is None or len(usage_entry2) == 0:

tests/utils/test_usage_tracker.py

@@ -1,3 +1,5 @@
+from pydantic import BaseModel
+
 import dspy
 from dspy.utils.usage_tracker import UsageTracker, track_usage
 
@@ -221,3 +223,105 @@ def test_merge_usage_entries_with_none_values():
     assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["audio_tokens"] == 1
     assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["accepted_prediction_tokens"] == 1
     assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["rejected_prediction_tokens"] == 1
+
+
+def test_merge_usage_entries_with_pydantic_models():
+    """Test merging usage entries with Pydantic model objects, like `PromptTokensDetailsWrapper` from litellm."""
+    tracker = UsageTracker()
+
+    # Here we define a simplified version of the Pydantic models from litellm to avoid the dependency change on litellm.
+    class CacheCreationTokenDetails(BaseModel):
+        ephemeral_5m_input_tokens: int
+        ephemeral_1h_input_tokens: int
+
+    class PromptTokensDetailsWrapper(BaseModel):
+        audio_tokens: int | None
+        cached_tokens: int
+        text_tokens: int | None
+        image_tokens: int | None
+        cache_creation_tokens: int
+        cache_creation_token_details: CacheCreationTokenDetails
+
+    # Add usage entries for different models
+    usage_entries = [
+        {
+            "model": "gpt-4o-mini",
+            "usage": {
+                "prompt_tokens": 1117,
+                "completion_tokens": 46,
+                "total_tokens": 1163,
+                "prompt_tokens_details": PromptTokensDetailsWrapper(
+                    audio_tokens=None,
+                    cached_tokens=3,
+                    text_tokens=None,
+                    image_tokens=None,
+                    cache_creation_tokens=0,
+                    cache_creation_token_details=CacheCreationTokenDetails(
+                        ephemeral_5m_input_tokens=5, ephemeral_1h_input_tokens=0
+                    ),
+                ),
+                "completion_tokens_details": {},
+            },
+        },
+        {
+            "model": "gpt-4o-mini",
+            "usage": {
+                "prompt_tokens": 800,
+                "completion_tokens": 100,
+                "total_tokens": 900,
+                "prompt_tokens_details": PromptTokensDetailsWrapper(
+                    audio_tokens=None,
+                    cached_tokens=3,
+                    text_tokens=None,
+                    image_tokens=None,
+                    cache_creation_tokens=0,
+                    cache_creation_token_details=CacheCreationTokenDetails(
+                        ephemeral_5m_input_tokens=5, ephemeral_1h_input_tokens=0
+                    ),
+                ),
+                "completion_tokens_details": None,
+            },
+        },
+        {
+            "model": "gpt-4o-mini",
+            "usage": {
+                "prompt_tokens": 800,
+                "completion_tokens": 100,
+                "total_tokens": 900,
+                "prompt_tokens_details": PromptTokensDetailsWrapper(
+                    audio_tokens=None,
+                    cached_tokens=3,
+                    text_tokens=None,
+                    image_tokens=None,
+                    cache_creation_tokens=0,
+                    cache_creation_token_details=CacheCreationTokenDetails(
+                        ephemeral_5m_input_tokens=5, ephemeral_1h_input_tokens=0
+                    ),
+                ),
+                "completion_tokens_details": {
+                    "reasoning_tokens": 1,
+                    "audio_tokens": 1,
+                    "accepted_prediction_tokens": 1,
+                    "rejected_prediction_tokens": 1,
+                },
+            },
+        },
+    ]
+
+    for entry in usage_entries:
+        tracker.add_usage(entry["model"], entry["usage"])
+
+    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_details"]["cached_tokens"] == 9
+    assert (
+        total_usage["gpt-4o-mini"]["prompt_tokens_details"]["cache_creation_token_details"]["ephemeral_5m_input_tokens"]
+        == 15
+    )
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["reasoning_tokens"] == 1
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["audio_tokens"] == 1
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["accepted_prediction_tokens"] == 1
+    assert total_usage["gpt-4o-mini"]["completion_tokens_details"]["rejected_prediction_tokens"] == 1