diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 61ed058bd1..6dd43ee8ac 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -4,8 +4,8 @@ import json import inspect from types import TracebackType -from typing import TYPE_CHECKING, Any, Generic, TypeVar, Iterator, Optional, AsyncIterator, cast -from typing_extensions import Self, Protocol, TypeGuard, override, get_origin, runtime_checkable +from typing import TYPE_CHECKING, Any, Generic, TypeVar, Callable, Iterator, Optional, AsyncIterator, cast +from typing_extensions import Self, Protocol, TypeGuard, get_args, override, get_origin, runtime_checkable import httpx2 @@ -56,6 +56,7 @@ def __stream__(self) -> Iterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data + normalize_stream_event = _make_stream_event_normalizer(cast_to) iterator = self._iter_events() try: @@ -84,6 +85,9 @@ def __stream__(self) -> Iterator[_T]: yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response) else: data = sse.json() + if normalize_stream_event is not None: + data = normalize_stream_event(data) + if is_mapping(data) and data.get("error"): message = None error = data.get("error") @@ -166,6 +170,7 @@ async def __stream__(self) -> AsyncIterator[_T]: cast_to = cast(Any, self._cast_to) response = self.response process_data = self._client._process_response_data + normalize_stream_event = _make_stream_event_normalizer(cast_to) iterator = self._iter_events() try: @@ -194,6 +199,9 @@ async def __stream__(self) -> AsyncIterator[_T]: yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response) else: data = sse.json() + if normalize_stream_event is not None: + data = normalize_stream_event(data) + if is_mapping(data) and data.get("error"): message = None error = data.get("error") @@ -402,6 +410,31 @@ def is_stream_class_type(typ: type) -> TypeGuard[type[Stream[object]] | type[Asy return inspect.isclass(origin) and issubclass(origin, (Stream, AsyncStream)) +def _make_stream_event_normalizer(cast_to: object) -> Callable[[object], object] | None: + if not _is_response_stream_event_type(cast_to): + return None + + from .lib.streaming.responses._stream_event_normalizer import ResponseStreamEventNormalizer + + return ResponseStreamEventNormalizer().normalize + + +def _is_response_stream_event_type(cast_to: object) -> bool: + annotated_args = get_args(cast_to) + if not annotated_args: + return False + + event_types = get_args(annotated_args[0]) + return any( + getattr(event_type, "__module__", None) + in { + "openai.types.responses.response_function_call_arguments_done_event", + "openai.types.beta.beta_response_function_call_arguments_done_event", + } + for event_type in event_types + ) + + def extract_stream_chunk_type( stream_cls: type, *, diff --git a/src/openai/lib/streaming/responses/_stream_event_normalizer.py b/src/openai/lib/streaming/responses/_stream_event_normalizer.py new file mode 100644 index 0000000000..f8fb5ec713 --- /dev/null +++ b/src/openai/lib/streaming/responses/_stream_event_normalizer.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from ...._utils import is_mapping + + +class ResponseStreamEventNormalizer: + def __init__(self) -> None: + self._function_call_names_by_item_id: dict[str, str] = {} + + def normalize(self, data: object) -> object: + if not is_mapping(data): + return data + + event_type = data.get("type") + if event_type == "response.output_item.added": + self._remember_function_call_name(data) + elif event_type == "response.function_call_arguments.done" and "name" not in data: + return self._with_function_call_name(data) + + return data + + def _remember_function_call_name(self, data: object) -> None: + if not is_mapping(data): + return + + item = data.get("item") + if not is_mapping(item) or item.get("type") != "function_call": + return + + item_id = item.get("id") + name = item.get("name") + if isinstance(item_id, str) and isinstance(name, str): + self._function_call_names_by_item_id[item_id] = name + + def _with_function_call_name(self, data: object) -> object: + if not is_mapping(data): + return data + + item_id = data.get("item_id") + if not isinstance(item_id, str): + return data + + name = self._function_call_names_by_item_id.get(item_id) + if name is None: + return data + + return {**data, "name": name} diff --git a/tests/lib/responses/test_responses.py b/tests/lib/responses/test_responses.py index 43879942b8..24e7b48898 100644 --- a/tests/lib/responses/test_responses.py +++ b/tests/lib/responses/test_responses.py @@ -1,7 +1,9 @@ from __future__ import annotations +from typing import cast from typing_extensions import TypeVar +import httpx2 import pytest from inline_snapshot import snapshot @@ -10,7 +12,7 @@ from openai._types import omit from openai._utils import assert_signatures_in_sync from openai._models import construct_type_unchecked -from openai.types.responses import Response +from openai.types.responses import Response, ResponseFunctionCallArgumentsDoneEvent from openai.lib._parsing._responses import parse_response from ...conftest import base_url @@ -92,3 +94,38 @@ def test_parse_method_definition_in_sync(sync: bool, client: OpenAI, async_clien checking_client.responses.parse, exclude_params={"tools"}, ) + + +@pytest.mark.respx2(base_url=base_url) +@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"]) +async def test_streamed_function_call_arguments_done_uses_output_item_name( + sync: bool, + client: OpenAI, + async_client: AsyncOpenAI, + respx2_mock: MockRouter, +) -> None: + respx2_mock.post("/responses").mock( + return_value=httpx2.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"type":"response.output_item.added","item":{"id":"fc_test","type":"function_call",' + b'"call_id":"call_test","name":"get_weather","arguments":""},"output_index":0,' + b'"sequence_number":1}\n\n' + b'data: {"type":"response.function_call_arguments.done","arguments":"{}",' + b'"item_id":"fc_test","output_index":0,"sequence_number":2}\n\n' + b"data: [DONE]\n\n" + ), + ) + ) + + if sync: + stream = client.responses.create(model="gpt-4o", input="call a tool", stream=True) + events = list(stream) + else: + stream = await async_client.responses.create(model="gpt-4o", input="call a tool", stream=True) + events = [event async for event in stream] + + done_event = cast(ResponseFunctionCallArgumentsDoneEvent, events[1]) + assert done_event.type == "response.function_call_arguments.done" + assert done_event.name == "get_weather" diff --git a/tests/test_streaming.py b/tests/test_streaming.py index ae6c0590f7..d157b32f51 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -1,5 +1,6 @@ from __future__ import annotations +import sys from typing import Iterator, AsyncIterator import httpx2 @@ -7,6 +8,7 @@ from openai import OpenAI, AsyncOpenAI from openai._streaming import Stream, AsyncStream, ServerSentEvent +from openai.types.chat.chat_completion_chunk import ChatCompletionChunk @pytest.mark.asyncio @@ -216,6 +218,33 @@ def body() -> Iterator[bytes]: assert sse.json() == {"content": "известни"} +def test_unrelated_stream_does_not_import_responses_normalizer( + monkeypatch: pytest.MonkeyPatch, + client: OpenAI, +) -> None: + for module_name in list(sys.modules): + if module_name.startswith("openai.lib.streaming.responses"): + monkeypatch.delitem(sys.modules, module_name) + + stream = Stream( + cast_to=ChatCompletionChunk, + client=client, + response=httpx2.Response( + 200, + content=( + b'data: {"id":"chatcmpl_test","object":"chat.completion.chunk","created":1,' + b'"model":"gpt-4o","choices":[]}\n\n' + b"data: [DONE]\n\n" + ), + ), + ) + + chunks = list(stream) + + assert chunks[0].id == "chatcmpl_test" + assert "openai.lib.streaming.responses._stream_event_normalizer" not in sys.modules + + async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]: for chunk in iter: yield chunk