466b44df18
PR to enhance the `Usage` object and related logic, to support more granular token accounting, matching the details available in the [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) . Specifically, it: - Adds `input_tokens_details` and `output_tokens_details` fields to the `Usage` dataclass, storing detailed token breakdowns (e.g., `cached_tokens`, `reasoning_tokens`). - Flows this change through - Updates and extends tests to match - Adds a test for the Usage.add method ### Motivation - Aligns the SDK’s usage with the latest OpenAI responses API Usage object - Supports downstream use cases that require fine-grained token usage data (e.g., billing, analytics, optimization) requested by startups --------- Co-authored-by: Wulfie Bain <wulfie@openai.com>
100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
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
|
|
|
|
|
|
@pytest.mark.allow_call_model_methods
|
|
@pytest.mark.asyncio
|
|
async def test_extra_headers_passed_to_openai_responses_model():
|
|
"""
|
|
Ensure extra_headers in ModelSettings is passed to the OpenAIResponsesModel client.
|
|
"""
|
|
called_kwargs = {}
|
|
|
|
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,
|
|
"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
|
|
extra_headers = {"X-Test-Header": "test-value"}
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="hi",
|
|
model_settings=ModelSettings(extra_headers=extra_headers),
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
assert "extra_headers" in called_kwargs
|
|
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():
|
|
"""
|
|
Ensure extra_headers in ModelSettings is passed to the OpenAI client.
|
|
"""
|
|
called_kwargs = {}
|
|
|
|
class DummyCompletions:
|
|
async def create(self, **kwargs):
|
|
nonlocal called_kwargs
|
|
called_kwargs = kwargs
|
|
msg = ChatCompletionMessage(role="assistant", content="Hello")
|
|
choice = Choice(index=0, finish_reason="stop", message=msg)
|
|
return ChatCompletion(
|
|
id="resp-id",
|
|
created=0,
|
|
model="fake",
|
|
object="chat.completion",
|
|
choices=[choice],
|
|
usage=None,
|
|
)
|
|
|
|
class DummyClient:
|
|
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
|
|
extra_headers = {"X-Test-Header": "test-value"}
|
|
await model.get_response(
|
|
system_instructions=None,
|
|
input="hi",
|
|
model_settings=ModelSettings(extra_headers=extra_headers),
|
|
tools=[],
|
|
output_schema=None,
|
|
handoffs=[],
|
|
tracing=ModelTracing.DISABLED,
|
|
previous_response_id=None,
|
|
)
|
|
assert "extra_headers" in called_kwargs
|
|
assert called_kwargs["extra_headers"]["X-Test-Header"] == "test-value"
|