feat: updated tests

Author: WJPBProjects

tests/models/test_litellm_chatcompletions_stream.py

@@ -8,7 +8,11 @@
     ChoiceDeltaToolCall,
     ChoiceDeltaToolCallFunction,
 )
-from openai.types.completion_usage import CompletionUsage
+from openai.types.completion_usage import (
+    CompletionTokensDetails,
+    CompletionUsage,
+    PromptTokensDetails,
+)
 from openai.types.responses import (
     Response,
     ResponseFunctionToolCall,
@@ -46,7 +50,13 @@ async def test_stream_response_yields_events_for_text_content(monkeypatch) -> No
         model="fake",
         object="chat.completion.chunk",
         choices=[Choice(index=0, delta=ChoiceDelta(content="llo"))],
-        usage=CompletionUsage(completion_tokens=5, prompt_tokens=7, total_tokens=12),
+        usage=CompletionUsage(
+            completion_tokens=5,
+            prompt_tokens=7,
+            total_tokens=12,
+            completion_tokens_details=CompletionTokensDetails(reasoning_tokens=2),
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=5),
+        ),
     )
 
     async def fake_stream() -> AsyncIterator[ChatCompletionChunk]:

tests/models/test_litellm_extra_body.py

@@ -22,12 +22,14 @@ async def fake_acompletion(model, messages=None, **kwargs):
         captured.update(kwargs)
         msg = Message(role="assistant", content="ok")
         choice = Choices(index=0, message=msg)
-        return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
+        return ModelResponse(
+            choices=[choice],
+            usage=Usage(0, 0, 0),
+        )
 
     monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
     settings = ModelSettings(
-        temperature=0.1,
-        extra_body={"cached_content": "some_cache", "foo": 123}
+        temperature=0.1, extra_body={"cached_content": "some_cache", "foo": 123}
     )
     model = LitellmModel(model="test-model")
 

tests/test_extra_headers.py

@@ -1,6 +1,7 @@
 import pytest
 from openai.types.chat.chat_completion import ChatCompletion, Choice
 from openai.types.chat.chat_completion_message import ChatCompletionMessage
+from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails
 
 from agents import ModelSettings, ModelTracing, OpenAIChatCompletionsModel, OpenAIResponsesModel
 
@@ -17,21 +18,29 @@ class DummyResponses:
         async def create(self, **kwargs):
             nonlocal called_kwargs
             called_kwargs = kwargs
+
             class DummyResponse:
                 id = "dummy"
                 output = []
                 usage = type(
-                    "Usage", (), {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+                    "Usage",
+                    (),
+                    {
+                        "input_tokens": 0,
+                        "output_tokens": 0,
+                        "total_tokens": 0,
+                        "input_tokens_details": InputTokensDetails(cached_tokens=0),
+                        "output_tokens_details": OutputTokensDetails(reasoning_tokens=0),
+                    },
                 )()
+
             return DummyResponse()
 
     class DummyClient:
         def __init__(self):
             self.responses = DummyResponses()
 
-
-
-    model = OpenAIResponsesModel(model="gpt-4", openai_client=DummyClient()) # type: ignore
+    model = OpenAIResponsesModel(model="gpt-4", openai_client=DummyClient())  # type: ignore
     extra_headers = {"X-Test-Header": "test-value"}
     await model.get_response(
         system_instructions=None,
@@ -47,7 +56,6 @@ def __init__(self):
     assert called_kwargs["extra_headers"]["X-Test-Header"] == "test-value"
 
 
-
 @pytest.mark.allow_call_model_methods
 @pytest.mark.asyncio
 async def test_extra_headers_passed_to_openai_client():
@@ -76,7 +84,7 @@ def __init__(self):
             self.chat = type("_Chat", (), {"completions": DummyCompletions()})()
             self.base_url = "https://api.openai.com"
 
-    model = OpenAIChatCompletionsModel(model="gpt-4", openai_client=DummyClient()) # type: ignore
+    model = OpenAIChatCompletionsModel(model="gpt-4", openai_client=DummyClient())  # type: ignore
     extra_headers = {"X-Test-Header": "test-value"}
     await model.get_response(
         system_instructions=None,

tests/test_responses_tracing.py

@@ -2,6 +2,7 @@
 from inline_snapshot import snapshot
 from openai import AsyncOpenAI
 from openai.types.responses import ResponseCompletedEvent
+from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails
 
 from agents import ModelSettings, ModelTracing, OpenAIResponsesModel, trace
 from agents.tracing.span_data import ResponseSpanData
@@ -16,10 +17,25 @@ def is_disabled(self):
 
 
 class DummyUsage:
-    def __init__(self, input_tokens=1, output_tokens=1, total_tokens=2):
+    def __init__(
+        self,
+        input_tokens: int = 1,
+        input_tokens_details: InputTokensDetails | None = None,
+        output_tokens: int = 1,
+        output_tokens_details: OutputTokensDetails | None = None,
+        total_tokens: int = 2,
+    ):
         self.input_tokens = input_tokens
         self.output_tokens = output_tokens
         self.total_tokens = total_tokens
+        self.input_tokens_details = (
+            input_tokens_details if input_tokens_details else InputTokensDetails(cached_tokens=0)
+        )
+        self.output_tokens_details = (
+            output_tokens_details
+            if output_tokens_details
+            else OutputTokensDetails(reasoning_tokens=0)
+        )
 
 
 class DummyResponse: