diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index d54b48d36..df54b6b19 100644 --- a/py/src/braintrust/integrations/openai/test_openai.py +++ b/py/src/braintrust/integrations/openai/test_openai.py @@ -6,6 +6,7 @@ import tempfile import time import zlib +from types import SimpleNamespace import openai import pytest @@ -14,6 +15,7 @@ from braintrust.integrations.openai.tracing import ( RAW_RESPONSE_HEADER, ChatCompletionWrapper, + ResponseWrapper, _materialize_logged_file_input, _process_attachments_in_chat_output, ) @@ -2994,3 +2996,21 @@ def test_chat_completion_non_stream_audio_is_materialized_as_attachment(self, me filename="generated_audio.wav", ) assert not memory_logger.pop() + + +@pytest.mark.parametrize("event_type", ["response.completed", "response.incomplete", "response.failed"]) +def test_responses_stream_terminal_event_provides_output_and_metrics(event_type): + output_item = SimpleNamespace(id="msg_1", type="message", role="assistant", status="incomplete", content=[]) + usage = SimpleNamespace(input_tokens=10, output_tokens=4, total_tokens=14) + all_results = [ + SimpleNamespace(type="response.created", response=SimpleNamespace(id="resp_1", output=[], usage=None)), + SimpleNamespace( + type=event_type, + response=SimpleNamespace(id="resp_1", output=[output_item], usage=usage), + ), + ] + + result = ResponseWrapper(None, None)._postprocess_streaming_results(all_results) + + assert result["output"] == [output_item] + assert result["metrics"] == {"prompt_tokens": 10, "completion_tokens": 4, "tokens": 14} diff --git a/py/src/braintrust/integrations/openai/tracing.py b/py/src/braintrust/integrations/openai/tracing.py index 5b2ed9692..233f0e2a3 100644 --- a/py/src/braintrust/integrations/openai/tracing.py +++ b/py/src/braintrust/integrations/openai/tracing.py @@ -1470,6 +1470,12 @@ async def gen(): return traced_stream +# A Responses API stream ends with one of these events. Streams that hit +# `max_output_tokens` or a content filter end with `response.incomplete`, and +# all three carry the final `response` object with output and usage. +_TERMINAL_RESPONSE_EVENT_TYPES = frozenset({"response.completed", "response.incomplete", "response.failed"}) + + class ResponseWrapper: def __init__( self, @@ -1652,7 +1658,7 @@ def _postprocess_streaming_results(cls, all_results: list[Any]) -> dict[str, Any if ( not usage and hasattr(result, "type") - and result.type == "response.completed" + and result.type in _TERMINAL_RESPONSE_EVENT_TYPES and hasattr(result, "response") ): response_dict = _try_to_dict(result.response) @@ -1682,7 +1688,7 @@ def _postprocess_streaming_results(cls, all_results: list[Any]) -> dict[str, Any output.append(item_data) continue - if result.type == "response.completed": + if result.type in _TERMINAL_RESPONSE_EVENT_TYPES: if hasattr(result, "response") and hasattr(result.response, "output"): if response_dict is None: response_dict = _try_to_dict(result.response)