Repository: stanfordnlp/dspy (python)

Issue:
[Bug] Track Usage not working w/ FastAPI+Bedrock+dspy.Streamify
### What happened?

When using `dspy.settings.configure(track_usage=True)` using a AWS Bedrock model in a FastAPI application and using streaming of the answer, I get the following error:

```python
Exception in thread Thread-1 (producer):
  + Exception Group Traceback (most recent call last):
  |   File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/threading.py", line 1075, in _bootstrap_inner
  |     self.run()
  |   File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/threading.py", line 1012, in run
  |     self._target(*self._args, **self._kwargs)
  |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/streaming/streamify.py", line 243, in producer
  |     context.run(asyncio.run, runner())
  |   File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/asyncio/runners.py", line 194, in run
  |     return runner.run(main)
  |            ^^^^^^^^^^^^^^^^
  |   File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/asyncio/runners.py", line 118, in run
  |     return self._loop.run_until_complete(task)
  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  |   File "uvloop/loop.pyx", line 1518, in uvloop.loop.Loop.run_until_complete
  |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/streaming/streamify.py", line 237, in runner
  |     async for item in async_generator:
  |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/streaming/streamify.py", line 177, in async_streamer
  |     async with create_task_group() as tg, send_stream, receive_stream:
  |                ^^^^^^^^^^^^^^^^^^^
  |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/anyio/_backends/_asyncio.py", line 772, in __aexit__
  |     raise BaseExceptionGroup(
  | ExceptionGroup: unhandled errors in a TaskGroup (1 sub-exception)
  +-+---------------- 1 ----------------
    | Traceback (most recent call last):
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/streaming/streamify.py", line 171, in generator
    |     prediction = await program(*args, **kwargs)
    |                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/utils/asyncify.py", line 63, in async_program
    |     return await call_async(*args, **kwargs)
    |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/asyncer/_main.py", line 382, in wrapper
    |     return await run_sync(
    |            ^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/asyncer/_compat.py", line 24, in run_sync
    |     return await anyio.to_thread.run_sync(
    |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/anyio/to_thread.py", line 56, in run_sync
    |     return await get_async_backend().run_sync_in_worker_thread(
    |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/anyio/_backends/_asyncio.py", line 2470, in run_sync_in_worker_thread
    |     return await future
    |            ^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/anyio/_backends/_asyncio.py", line 967, in run
    |     result = context.run(func, *args)
    |              ^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/utils/asyncify.py", line 57, in wrapped_program
    |     return program(*a, **kw)
    |            ^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/utils/callback.py", line 343, in sync_wrapper
    |     raise exception
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/utils/callback.py", line 339, in sync_wrapper
    |     results = fn(instance, *args, **kwargs)
    |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/primitives/program.py", line 29, in __call__
    |     output.set_lm_usage(usage_tracker.get_total_tokens())
    |                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/utils/usage_tracker.py", line 59, in get_total_tokens
    |     total_usage = self._merge_usage_entries(total_usage, usage_entry)
    |                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/Users/cmeyer/code/dspy_poc/.venv/lib/python3.12/site-packages/dspy/utils/usage_tracker.py", line 45, in _merge_usage_entries
    |     result[k] += v if v else 0
    | TypeError: unsupported operand type(s) for +=: 'dict' and 'int'
    +------------------------------------
```


### Steps to reproduce



Here is my reproductible error:
```python
from dataclasses import asdict

import dspy
import ujson
from dspy.streaming import StatusMessageProvider, StreamListener, StreamResponse, StatusMessage
from fastapi.responses import StreamingResponse
from fastapi import FastAPI, HTTPException
import dspy


class MyStatusMessageProvider(StatusMessageProvider):
    def tool_start_status_message(self, instance, inputs):
        return f"`Pythia is using tool {instance.name}...`"

    def tool_end_status_message(self, outputs):
        return ""

    def lm_start_status_message(self, instance, inputs):
        return "`Pythia is thinking...`"

def evaluate_math(expression: str):
    return dspy.PythonInterpreter({}).execute(expression)

def search_wikipedia(query: str):
    results = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')(query, k=3)
    return [x['text'] for x in results]

def streaming_response(streamer):
    """
    Custom streaming response function that handles StatusMessage objects.
    This implements the same functionality as dspy_streaming_response but adds handling for StatusMessage objects.
    """
    from dspy.primitives.prediction import Prediction
    import litellm

    for value in streamer:
        if isinstance(value, StatusMessage):
            data = {"status": value.message}
            yield f"data: {ujson.dumps(data)}\n\n"
        elif isinstance(value, Prediction):
            data = {"prediction": {k: v for k, v in value.items(include_dspy=False)}}
            yield f"data: {ujson.dumps(data)}\n\n"
        elif isinstance(value, litellm.ModelResponseStream):
            data = {"chunk": value.json()}
            yield f"data: {ujson.dumps(data)}\n\n"
        elif isinstance(value, StreamResponse):
            data = {"chunk": {k: v for k, v in asdict(value).items()}}
            yield f"data: {ujson.dumps(data)}\n\n"
        elif isinstance(value, str) and value.startswith("data:"):
            yield value
        else:
            data = {"unknown": str(value)}
            yield f"data: {ujson.dumps(data)}\n\n"
    yield "data: [DONE]\n\n"

app = FastAPI(
    title="DSPy Program API",
    description="A simple API serving a DSPy Chain of Thought program",
    version="1.0.0"
)
# Configure DSPy with LiteLLM
dspy.configure(
    lm=dspy.LM(
        provider="litellm",
        model="bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0",
        aws_region_name="us-east-1",
        cache=False,
    )
)
dspy.settings.configure(track_usage=True)
react = dspy.ReAct("question -> answer: str", tools=[evaluate_math, search_wikipedia])

dspy_streamer = dspy.streamify(
    react,
    status_message_provider=MyStatusMessageProvider(),
    stream_listeners=[
        StreamListener(signature_field_name="reasoning"),
        StreamListener(signature_field_name="answer"),
    ],
    async_streaming=False,
)

@app.post("/predict")
def predict():
    stream = dspy_streamer(question="What is 9362158 divided by the year of birth of David Gregory of Kinnairdy castle?")
    return StreamingResponse(streaming_response(stream), media_type="text/event-stream")
```

I might do some more testing with other provider or without streaming

Result of more testing cases:
❌ FastAPI + Bedrock + dspy.streamify
✅ FastAPI + OpenAI + dspy.streamify
✅ FastAPI + Bedrock (without streaming)
✅ Bedrock + dspy.streamify (simple script without API)


### DSPy version

2.6.24

Pull request:
Fix usage tracking when None exists as value
fix #8265 

In some providers (only Bedrock for now), the value of certain keys in the usage dict can be `None`, this PR handles this case.

Patch:
diff --git a/dspy/adapters/chat_adapter.py b/dspy/adapters/chat_adapter.py
index 5d5f2400..97bb173c 100644
--- a/dspy/adapters/chat_adapter.py
+++ b/dspy/adapters/chat_adapter.py
@@ -141,12 +141,14 @@ class ChatAdapter(Adapter):
         outputs: dict[str, Any],
         missing_field_message=None,
     ) -> str:
-        return self.format_field_with_value(
+        assistant_message_content = self.format_field_with_value(
             {
                 FieldInfoWithName(name=k, info=v): outputs.get(k, missing_field_message)
                 for k, v in signature.output_fields.items()
             },
         )
+        assistant_message_content += "\n\n[[ ## completed ## ]]\n"
+        return assistant_message_content
 
     def parse(self, signature: Type[Signature], completion: str) -> dict[str, Any]:
         sections = [(None, [])]
diff --git a/dspy/adapters/types/audio.py b/dspy/adapters/types/audio.py
index 800017e7..b4637a06 100644
--- a/dspy/adapters/types/audio.py
+++ b/dspy/adapters/types/audio.py
@@ -44,10 +44,10 @@ class Audio(BaseType):
     @classmethod
     def validate_input(cls, values: Any) -> Any:
         """
-        Validate input for Audio, expecting 'data' and 'format' keys in dictionary.
+        Validate input for Audio, expecting 'data' and 'audio_format' keys in dictionary.
         """
         if isinstance(values, cls):
-            return {"data": values.data, "format": values.format}
+            return {"data": values.data, "audio_format": values.audio_format}
         return encode_audio(values)
 
     @classmethod
@@ -62,7 +62,7 @@ class Audio(BaseType):
             raise ValueError(f"Unsupported MIME type for audio: {mime_type}")
         audio_format = mime_type.split("/")[1]
         encoded_data = base64.b64encode(response.content).decode("utf-8")
-        return cls(data=encoded_data, format=audio_format)
+        return cls(data=encoded_data, audio_format=audio_format)
 
     @classmethod
     def from_file(cls, file_path: str) -> "Audio":
@@ -81,7 +81,7 @@ class Audio(BaseType):
 
         audio_format = mime_type.split("/")[1]
         encoded_data = base64.b64encode(file_data).decode("utf-8")
-        return cls(data=encoded_data, format=audio_format)
+        return cls(data=encoded_data, audio_format=audio_format)
 
     @classmethod
     def from_array(
@@ -102,44 +102,44 @@ class Audio(BaseType):
             subtype="PCM_16",
         )
         encoded_data = base64.b64encode(byte_buffer.getvalue()).decode("utf-8")
-        return cls(data=encoded_data, format=format)
+        return cls(data=encoded_data, audio_format=format)
 
     def __str__(self) -> str:
         return self.serialize_model()
 
     def __repr__(self) -> str:
         length = len(self.data)
-        return f"Audio(data=<AUDIO_BASE_64_ENCODED({length})>, format='{self.audio_format}')"
+        return f"Audio(data=<AUDIO_BASE_64_ENCODED({length})>, audio_format='{self.audio_format}')"
 
 def encode_audio(audio: Union[str, bytes, dict, "Audio", Any], sampling_rate: int = 16000, format: str = "wav") -> dict:
     """
-    Encode audio to a dict with 'data' and 'format'.
+    Encode audio to a dict with 'data' and 'audio_format'.
     
     Accepts: local file path, URL, data URI, dict, Audio instance, numpy array, or bytes (with known format).
     """
-    if isinstance(audio, dict) and "data" in audio and "format" in audio:
+    if isinstance(audio, dict) and "data" in audio and "audio_format" in audio:
         return audio
     elif isinstance(audio, Audio):
-        return {"data": audio.data, "format": audio.format}
+        return {"data": audio.data, "audio_format": audio.audio_format}
     elif isinstance(audio, str) and audio.startswith("data:audio/"):
         try:
             header, b64data = audio.split(",", 1)
             mime = header.split(";")[0].split(":")[1]
             audio_format = mime.split("/")[1]
-            return {"data": b64data, "format": audio_format}
+            return {"data": b64data, "audio_format": audio_format}
         except Exception as e:
             raise ValueError(f"Malformed audio data URI: {e}")
     elif isinstance(audio, str) and os.path.isfile(audio):
         a = Audio.from_file(audio)
-        return {"data": a.data, "format": a.format}
+        return {"data": a.data, "audio_format": a.audio_format}
     elif isinstance(audio, str) and audio.startswith("http"):
         a = Audio.from_url(audio)
-        return {"data": a.data, "format": a.format}
+        return {"data": a.data, "audio_format": a.audio_format}
     elif SF_AVAILABLE and hasattr(audio, "shape"):
         a = Audio.from_array(audio, sampling_rate=sampling_rate, format=format)
-        return {"data": a.data, "format": a.format}
+        return {"data": a.data, "audio_format": a.audio_format}
     elif isinstance(audio, bytes):
         encoded = base64.b64encode(audio).decode("utf-8")
-        return {"data": encoded, "format": format}
+        return {"data": encoded, "audio_format": format}
     else:
         raise ValueError(f"Unsupported type for encode_audio: {type(audio)}")
diff --git a/dspy/adapters/types/tool.py b/dspy/adapters/types/tool.py
index b17cabd6..71008d34 100644
--- a/dspy/adapters/types/tool.py
+++ b/dspy/adapters/types/tool.py
@@ -6,6 +6,7 @@ from jsonschema import ValidationError, validate
 from pydantic import BaseModel, TypeAdapter, create_model
 
 from dspy.adapters.types.base_type import BaseType
+from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import with_callbacks
 
 if TYPE_CHECKING:
@@ -160,12 +161,26 @@ class Tool(BaseType):
             },
         }
 
+    def _run_async_in_sync(self, coroutine):
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return asyncio.run(coroutine)
+
+        return loop.run_until_complete(coroutine)
+
     @with_callbacks
     def __call__(self, **kwargs):
         parsed_kwargs = self._validate_and_parse_args(**kwargs)
         result = self.func(**parsed_kwargs)
         if asyncio.iscoroutine(result):
-            raise ValueError("You are calling `__call__` on an async tool, please use `acall` instead.")
+            if settings.allow_tool_async_sync_conversion:
+                return self._run_async_in_sync(result)
+            else:
+                raise ValueError(
+                    "You are calling `__call__` on an async tool, please use `acall` instead or set "
+                    "`allow_async=True` to run the async tool in sync mode."
+                )
         return result
 
     @with_callbacks
diff --git a/dspy/dsp/utils/settings.py b/dspy/dsp/utils/settings.py
index fd7ce52f..fc0f7300 100644
--- a/dspy/dsp/utils/settings.py
+++ b/dspy/dsp/utils/settings.py
@@ -27,6 +27,8 @@ DEFAULT_CONFIG = dotdict(
     provide_traceback=False,  # Whether to include traceback information in error logs.
     num_threads=8,  # Number of threads to use for parallel processing.
     max_errors=10,  # Maximum errors before halting operations.
+    # If true, async tools can be called in sync mode by getting converted to sync.
+    allow_tool_async_sync_conversion=False,
 )
 
 # Global base configuration and owner tracking
diff --git a/dspy/predict/react.py b/dspy/predict/react.py
index cc2a8a02..9b19482b 100644
--- a/dspy/predict/react.py
+++ b/dspy/predict/react.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Callable, Literal
+from typing import TYPE_CHECKING, Any, Callable, Literal, Type
 
 from litellm import ContextWindowExceededError
 
@@ -10,11 +10,33 @@ from dspy.signatures.signature import ensure_signature
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from dspy.signatures.signature import Signature
+
 
 class ReAct(Module):
-    def __init__(self, signature, tools: list[Callable], max_iters=5):
+    def __init__(self, signature: Type["Signature"], tools: list[Callable], max_iters: int = 10):
         """
-        `tools` is either a list of functions, callable classes, or `dspy.Tool` instances.
+        ReAct stands for "Reasoning and Acting," a popular paradigm for building tool-using agents.
+        In this approach, the language model is iteratively provided with a list of tools and has
+        to reason about the current situation. The model decides whether to call a tool to gather more
+        information or to finish the task based on its reasoning process. The DSPy version of ReAct is
+        generalized to work over any signature, thanks to signature polymorphism.
+
+        Args:
+            signature: The signature of the module, which defines the input and output of the react module.
+            tools (list[Callable]): A list of functions, callable objects, or `dspy.Tool` instances.
+            max_iters (Optional[int]): The maximum number of iterations to run. Defaults to 10.
+
+        Example:
+
+        ```python
+        def get_weather(city: str) -> str:
+            return f"The weather in {city} is sunny."
+
+        react = dspy.ReAct(signature="question->answer", tools=[get_weather])
+        pred = react(question="What is the weather in Tokyo?")
+        ```
         """
         super().__init__()
         self.signature = signature = ensure_signature(signature)
diff --git a/dspy/utils/usage_tracker.py b/dspy/utils/usage_tracker.py
index 2643645a..5199f3ad 100644
--- a/dspy/utils/usage_tracker.py
+++ b/dspy/utils/usage_tracker.py
@@ -38,7 +38,7 @@ class UsageTracker:
         result = dict(usage_entry2)
         for k, v in usage_entry1.items():
             current_v = result.get(k)
-            if isinstance(v, dict):
+            if isinstance(v, dict) or isinstance(current_v, dict):
                 result[k] = self._merge_usage_entries(current_v, v)
             else:
                 result[k] = (current_v or 0) + (v or 0)


Test patch:
diff --git a/tests/adapters/test_chat_adapter.py b/tests/adapters/test_chat_adapter.py
index 3d8a67b7..969a8f17 100644
--- a/tests/adapters/test_chat_adapter.py
+++ b/tests/adapters/test_chat_adapter.py
@@ -244,6 +244,9 @@ def test_chat_adapter_formats_image_with_few_shot_examples():
     # 1 system message, 2 few shot examples (1 user and assistant message for each example), 1 user message
     assert len(messages) == 6
 
+    assert "[[ ## completed ## ]]\n" in messages[2]["content"]
+    assert "[[ ## completed ## ]]\n" in messages[4]["content"]
+
     assert {"type": "image_url", "image_url": {"url": "https://example.com/image1.jpg"}} in messages[1]["content"]
     assert {"type": "image_url", "image_url": {"url": "https://example.com/image2.jpg"}} in messages[3]["content"]
     assert {"type": "image_url", "image_url": {"url": "https://example.com/image3.jpg"}} in messages[5]["content"]
@@ -350,3 +353,26 @@ def test_chat_adapter_with_tool():
     # Tool arguments format should be included in the user message
     assert "{'city': {'type': 'string'}}" in messages[1]["content"]
     assert "{'country': {'type': 'string'}, 'year': {'type': 'integer'}}" in messages[1]["content"]
+
+
+def test_chat_adapter_formats_conversation_history():
+    class MySignature(dspy.Signature):
+        question: str = dspy.InputField()
+        history: dspy.History = dspy.InputField()
+        answer: str = dspy.OutputField()
+
+    history = dspy.History(
+        messages=[
+            {"question": "What is the capital of France?", "answer": "Paris"},
+            {"question": "What is the capital of Germany?", "answer": "Berlin"},
+        ]
+    )
+
+    adapter = dspy.ChatAdapter()
+    messages = adapter.format(MySignature, [], {"question": "What is the capital of France?", "history": history})
+
+    assert len(messages) == 6
+    assert messages[1]["content"] == "[[ ## question ## ]]\nWhat is the capital of France?"
+    assert messages[2]["content"] == "[[ ## answer ## ]]\nParis\n\n[[ ## completed ## ]]\n"
+    assert messages[3]["content"] == "[[ ## question ## ]]\nWhat is the capital of Germany?"
+    assert messages[4]["content"] == "[[ ## answer ## ]]\nBerlin\n\n[[ ## completed ## ]]\n"
diff --git a/tests/adapters/test_json_adapter.py b/tests/adapters/test_json_adapter.py
index cd561fe0..7db4a320 100644
--- a/tests/adapters/test_json_adapter.py
+++ b/tests/adapters/test_json_adapter.py
@@ -419,3 +419,26 @@ def test_json_adapter_with_tool():
             },
         },
     }
+
+
+def test_json_adapter_formats_conversation_history():
+    class MySignature(dspy.Signature):
+        question: str = dspy.InputField()
+        history: dspy.History = dspy.InputField()
+        answer: str = dspy.OutputField()
+
+    history = dspy.History(
+        messages=[
+            {"question": "What is the capital of France?", "answer": "Paris"},
+            {"question": "What is the capital of Germany?", "answer": "Berlin"},
+        ]
+    )
+
+    adapter = dspy.JSONAdapter()
+    messages = adapter.format(MySignature, [], {"question": "What is the capital of France?", "history": history})
+
+    assert len(messages) == 6
+    assert messages[1]["content"] == "[[ ## question ## ]]\nWhat is the capital of France?"
+    assert messages[2]["content"] == '{\n  "answer": "Paris"\n}'
+    assert messages[3]["content"] == "[[ ## question ## ]]\nWhat is the capital of Germany?"
+    assert messages[4]["content"] == '{\n  "answer": "Berlin"\n}'
diff --git a/tests/adapters/test_tool.py b/tests/adapters/test_tool.py
index de2f64d8..58938b87 100644
--- a/tests/adapters/test_tool.py
+++ b/tests/adapters/test_tool.py
@@ -5,6 +5,7 @@ import pytest
 from pydantic import BaseModel
 
 from dspy.adapters.types.tool import Tool
+import dspy
 
 
 # Test fixtures
@@ -363,3 +364,14 @@ async def test_async_concurrent_calls():
     # Check that it ran concurrently (should take ~0.1s, not ~0.5s)
     # We use 0.3s as threshold to account for some overhead
     assert end_time - start_time < 0.3
+
+
+def test_async_tool_call_in_sync_mode():
+    tool = Tool(async_dummy_function)
+    with dspy.context(allow_tool_async_sync_conversion=False):
+        with pytest.raises(ValueError):
+            result = tool(x=1, y="hello")
+
+    with dspy.context(allow_tool_async_sync_conversion=True):
+        result = tool(x=1, y="hello")
+        assert result == "hello 1"
diff --git a/tests/utils/test_usage_tracker.py b/tests/utils/test_usage_tracker.py
index 86b2d9ba..70244b32 100644
--- a/tests/utils/test_usage_tracker.py
+++ b/tests/utils/test_usage_tracker.py
@@ -170,3 +170,62 @@ def test_merge_usage_entries_with_new_keys():
 
     assert total_usage["model-x"]["prompt_tokens"] == 5
     assert total_usage["model-x"]["completion_tokens"] == 2
+
+
+def test_merge_usage_entries_with_none_values():
+    """Test tracking usage across multiple models."""
+    tracker = UsageTracker()
+
+    # 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": None,
+                "completion_tokens_details": {},
+            },
+        },
+        {
+            "model": "gpt-4o-mini",
+            "usage": {
+                "prompt_tokens": 800,
+                "completion_tokens": 100,
+                "total_tokens": 900,
+                "prompt_tokens_details": {"cached_tokens": 50, "audio_tokens": 50},
+                "completion_tokens_details": None,
+            },
+        },
+        {
+            "model": "gpt-4o-mini",
+            "usage": {
+                "prompt_tokens": 800,
+                "completion_tokens": 100,
+                "total_tokens": 900,
+                "prompt_tokens_details": None,
+                "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"] == 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
+    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