diff --git a/src/openai/lib/responses_websocket/README.md b/src/openai/lib/responses_websocket/README.md index a43d05c859..b429a42596 100644 --- a/src/openai/lib/responses_websocket/README.md +++ b/src/openai/lib/responses_websocket/README.md @@ -39,6 +39,28 @@ async def run(): return response ``` +The same helpers accept native `client.beta.responses.connect()` connections. +The constructor infers the connection family: Beta lanes accept Beta requests, +return Beta events, and collect a native `BetaResponse`. Stable connections and +bare session/lane annotations retain their existing stable types and defaults. +Use the native connection class as the type argument when explicitly annotating +a Beta session or lane. The async helpers follow the same rule. + +For caller-fed Beta events, select their response model when constructing the +accumulator: + +```python +from openai.types.beta.beta_response import BetaResponse +from openai.lib.responses_websocket import ResponsesWebSocketAccumulator + +beta_accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) +# Feed events from one Beta connection or managed lane, as in the example below. +``` + +The default `ResponsesWebSocketAccumulator()` keeps stable response types. +`reset()` retains the selected family. Beta finals preserve their native fields, +including opaque and encrypted data; the helper does not add parsed Beta models. + The session owns the connection's read side and lifetime. Its one reader uses the connection's raw receive, event parsing, and recovery methods. While the session is active, receive through its lanes instead of calling the connection's `recv`, diff --git a/src/openai/lib/responses_websocket/_accumulator.py b/src/openai/lib/responses_websocket/_accumulator.py index 5fcc674d78..bb7a3eed35 100644 --- a/src/openai/lib/responses_websocket/_accumulator.py +++ b/src/openai/lib/responses_websocket/_accumulator.py @@ -1,13 +1,16 @@ from __future__ import annotations from copy import deepcopy -from typing import Any, Literal, cast +from typing import Any, Generic, Literal, cast from dataclasses import field, dataclass +from typing_extensions import TypeVar, overload -from ._session import ResponsesWebSocketError, _field +from ._session import ResponsesWebSocketError, _field, _ServerEvent from ..._compat import PYDANTIC_V1, model_copy from ..._models import BaseModel from ...types.responses import Response +from ...types.beta.beta_response import BetaResponse +from ...types.beta.beta_responses_server_event import BetaResponsesServerEvent from ...types.responses.responses_server_event import ResponsesServerEvent _RichField = Literal["summary", "reasoning", "code", "command", "stdout", "stderr"] @@ -42,6 +45,8 @@ "response.shell_call_output_content.done": "shell_call_output", } +_ResponseT = TypeVar("_ResponseT", Response, BetaResponse, default=Response) + @dataclass(frozen=True) class ResponsesWebSocketOutput: @@ -89,7 +94,7 @@ class _Output: shell_output: dict[str, dict[str, object]] = field(default_factory=dict[str, dict[str, object]]) -class ResponsesWebSocketAccumulator: +class ResponsesWebSocketAccumulator(Generic[_ResponseT]): """Caller-fed text/tool inputs and detailed reasoning, code and shell progress. Feed typed events from a connection or lane's recv(). Use one accumulator per @@ -99,7 +104,16 @@ class ResponsesWebSocketAccumulator: server response, including failed/incomplete and missing/null/empty output. """ - def __init__(self) -> None: + @overload + def __init__( + self: ResponsesWebSocketAccumulator[Response], *, response_type: type[Response] = Response + ) -> None: ... + + @overload + def __init__(self: ResponsesWebSocketAccumulator[BetaResponse], *, response_type: type[BetaResponse]) -> None: ... + + def __init__(self, *, response_type: type[Response] | type[BetaResponse] = Response) -> None: + self._response_type = response_type self._stream_id: str | None = None self._response_id: str | None = None self._terminal_type: str | None = None @@ -108,7 +122,7 @@ def __init__(self) -> None: # non-negative indices without decimal string conversion limits. self._output: dict[str, _Output] = {} self._response: dict[str, object] | None = None - self._final: Response | None = None + self._final: _ResponseT | None = None self._error: Exception | None = None self._shell_command_counts: dict[str, int] = {} @@ -231,7 +245,7 @@ def detailed_snapshot(self) -> dict[str, Any]: "output": output, } - def get_final_response(self) -> Response: + def get_final_response(self) -> _ResponseT: """Return an independent copy of the received terminal response, never a partial success.""" error = self._error if error is not None: @@ -240,7 +254,13 @@ def get_final_response(self) -> Response: raise RuntimeError("No terminal response has been received") return model_copy(self._final, deep=True) - def add_event(self, event: ResponsesServerEvent) -> None: + @overload + def add_event(self: ResponsesWebSocketAccumulator[Response], event: ResponsesServerEvent) -> None: ... + + @overload + def add_event(self: ResponsesWebSocketAccumulator[BetaResponse], event: BetaResponsesServerEvent) -> None: ... + + def add_event(self, event: _ServerEvent) -> None: """Observe an event without consuming or changing it; unknown events are ignored. Final fields replace deltas. A supplied output list, even [], replaces @@ -315,7 +335,7 @@ def add_event(self, event: ResponsesServerEvent) -> None: terminal = kind in {"response.completed", "response.failed", "response.incomplete"} if kind in {"response.created", "response.queued", "response.in_progress"} or terminal: response = _field(event, "response") - if not isinstance(response, Response): + if not isinstance(response, self._response_type): error = ValueError("WebSocket event is missing a valid response") if terminal: self._error = error @@ -351,7 +371,7 @@ def add_event(self, event: ResponsesServerEvent) -> None: self._response_id = response_id or self._response_id self._bound, self._stream_id = True, stream_id if terminal: - self._final = model_copy(response, deep=True) + self._final = cast(_ResponseT, model_copy(response, deep=True)) self._terminal_type = kind return if kind in _RICH_TEXT_EVENTS or kind in _RICH_STATUS_EVENTS or kind in _RICH_PART_EVENTS: @@ -487,7 +507,7 @@ def add_event(self, event: ResponsesServerEvent) -> None: elif kind == "response.custom_tool_call_input.done": item.input = [value] - def _add_rich_event(self, event: ResponsesServerEvent, kind: str) -> None: + def _add_rich_event(self, event: _ServerEvent, kind: str) -> None: index = _index_field(event, "output_index") # Shell-command events identify their call by output_index only. An # unmodeled item_id extra must stay raw, not become a routing guard. diff --git a/src/openai/lib/responses_websocket/_session.py b/src/openai/lib/responses_websocket/_session.py index f7a080dece..9839f8dfdf 100644 --- a/src/openai/lib/responses_websocket/_session.py +++ b/src/openai/lib/responses_websocket/_session.py @@ -6,11 +6,11 @@ import asyncio import threading from types import TracebackType -from typing import Any, cast +from typing import Any, Generic, Mapping, cast from weakref import WeakSet from collections import deque from dataclasses import dataclass -from typing_extensions import Self +from typing_extensions import Self, TypeVar, overload from ..._types import omit from ..._utils import maybe_transform @@ -19,11 +19,32 @@ from ..._exceptions import WebSocketConnectionClosedError from ...types.responses import Response, ResponseStreamEvent from ..streaming.responses import ResponseStreamState +from ...types.beta.beta_response import BetaResponse from ...resources.responses.responses import ResponsesConnection, AsyncResponsesConnection +from ...resources.beta.responses.responses import ( + ResponsesConnection as BetaResponsesConnection, + AsyncResponsesConnection as AsyncBetaResponsesConnection, +) +from ...types.beta.beta_responses_client_event import BetaResponsesClientEvent +from ...types.beta.beta_responses_server_event import BetaResponsesServerEvent from ...types.responses.responses_client_event import ResponsesClientEvent from ...types.responses.responses_server_event import ResponsesServerEvent +from ...types.beta.beta_responses_client_event_param import BetaResponsesClientEventParam from ...types.responses.responses_client_event_param import ResponsesClientEventParam +_ServerEvent = ResponsesServerEvent | BetaResponsesServerEvent +_ClientEvent = ( + ResponsesClientEvent | ResponsesClientEventParam | BetaResponsesClientEvent | BetaResponsesClientEventParam +) +_Response = Response | BetaResponse +_SyncConnectionT = TypeVar( + "_SyncConnectionT", ResponsesConnection, BetaResponsesConnection, default=ResponsesConnection +) +_AsyncConnectionT = TypeVar( + "_AsyncConnectionT", AsyncResponsesConnection, AsyncBetaResponsesConnection, default=AsyncResponsesConnection +) +_EventT = TypeVar("_EventT", bound=_ServerEvent, default=ResponsesServerEvent) + @dataclass(frozen=True) class ResponsesWebSocketLimits: @@ -55,10 +76,10 @@ class ResponsesWebSocketBufferError(RuntimeError): """A session or response exceeded an application-selected buffer limit.""" -class ResponsesWebSocketError(RuntimeError): +class ResponsesWebSocketError(RuntimeError, Generic[_EventT]): """A request-scoped protocol error. The original event is available as ``event``.""" - def __init__(self, event: ResponsesServerEvent) -> None: + def __init__(self, event: _EventT) -> None: super().__init__("The Responses WebSocket request returned an error event") self.event = event @@ -83,15 +104,18 @@ def _field(event: object, name: str) -> Any: return cast("dict[str, Any]", event).get(name) if isinstance(event, dict) else getattr(event, name, None) -def _command(event: ResponsesClientEvent | ResponsesClientEventParam, stream_id: str | None) -> tuple[str, bool]: - data = event.to_dict(exclude_unset=True) if isinstance(event, BaseModel) else dict(event) +def _command(event: _ClientEvent, stream_id: str | None, *, beta: bool) -> tuple[str, bool]: + data = ( + event.to_dict(exclude_unset=True) if isinstance(event, BaseModel) else dict(cast("Mapping[str, object]", event)) + ) is_create = bool(data.get("type") == "response.create") if is_create: if "stream_id" in data and data["stream_id"] != stream_id: raise ValueError("The command stream_id does not match its lane") if stream_id is not None: data["stream_id"] = stream_id - return json.dumps(maybe_transform(data, ResponsesClientEventParam)), is_create + param_type = BetaResponsesClientEventParam if beta else ResponsesClientEventParam + return json.dumps(maybe_transform(data, param_type)), is_create class _Accumulator: @@ -99,10 +123,10 @@ def __init__(self, limit: int | None) -> None: self.limit = limit self.size = 0 self.state: ResponseStreamState[object] | None = None - self.final: Response | None = None + self.final: _Response | None = None self.error: Exception | None = None - def add(self, event: ResponsesServerEvent, size: int) -> None: + def add(self, event: _ServerEvent, size: int) -> None: kind = _field(event, "type") if kind == "response.created": self.size = 0 @@ -135,18 +159,26 @@ def add(self, event: ResponsesServerEvent, size: int) -> None: self.state = None self.final = None raise ResponsesWebSocketBufferError("Response accumulation exceeded max_response_bytes") + # Reuse the isolated event for the compatible-event bridge and the + # native Beta final. The shared engine still validates completion; its + # intermediate ParsedResponse is not the public Beta result. + isolated = event + if (self.state is not None and compatible) or ( + terminal and isinstance(_field(event, "response"), BetaResponse) + ): + isolated = model_copy(event, deep=True) if self.state is not None and compatible: - # Accumulation must not mutate the raw event visible to callers. - self.state.handle_event(cast(ResponseStreamEvent, model_copy(event, deep=True))) + self.state.handle_event(cast(ResponseStreamEvent, isolated)) if terminal: - self.final = cast(Response, _field(event, "response")) - if kind == "response.completed" and self.state is not None: - self.final = self.state._completed_response - elif self.state is not None and getattr(self.final, "output", None) is None: - self.final = model_copy(self.final) - self.final.output = [ - self.state._completed_output[index] for index in sorted(self.state._completed_output) - ] + self.final = cast(_Response, _field(isolated, "response")) + if not isinstance(self.final, BetaResponse): + if kind == "response.completed" and self.state is not None: + self.final = self.state._completed_response + elif self.state is not None and getattr(self.final, "output", None) is None: + self.final = model_copy(self.final) + self.final.output = [ + self.state._completed_output[index] for index in sorted(self.state._completed_output) + ] self.state = None self.size = 0 @@ -154,7 +186,7 @@ def add(self, event: ResponsesServerEvent, size: int) -> None: class _LaneState: def __init__(self, stream_id: str | None, limit: int | None) -> None: self.stream_id = stream_id - self.queue: deque[tuple[ResponsesServerEvent, int]] = deque() + self.queue: deque[tuple[_ServerEvent, int]] = deque() self.size = 0 self.closed = False self.failure: ResponsesWebSocketBufferError | None = None @@ -170,7 +202,10 @@ class _Router: def __init__( self, limits: ResponsesWebSocketLimits, - connection: ResponsesConnection | AsyncResponsesConnection, + connection: ResponsesConnection + | AsyncResponsesConnection + | BetaResponsesConnection + | AsyncBetaResponsesConnection, *, isolate_lane_overflow: bool = False, ) -> None: @@ -232,7 +267,7 @@ def reconnected(self, connection: object) -> None: # including when it was registered before reconnect. del self.lanes[stream_id] - def dispatch(self, event: ResponsesServerEvent, size: int) -> None: + def dispatch(self, event: _ServerEvent, size: int) -> None: with self.condition: if self.ended: return @@ -290,7 +325,7 @@ def finish(self, error: BaseException | None = None) -> None: lane.wake.set() self.condition.notify_all() - def pop(self, lane: _LaneState) -> ResponsesServerEvent | None: + def pop(self, lane: _LaneState) -> _ServerEvent | None: # The caller holds the condition. No await occurs after dequeueing. if lane.failure is not None: raise lane.failure @@ -333,23 +368,38 @@ def prepare(self, lane: _LaneState, is_create: bool) -> None: lane.in_flight = True -class ResponsesWebSocketLane: +class ResponsesWebSocketLane(Generic[_SyncConnectionT]): """A routed view of one session. ``close`` detaches it without closing the socket. Use a single consumer per lane. Cancel a synchronous wait with ``recv(timeout=...)``; a timeout never consumes an event. Different lanes may be consumed concurrently. """ - def __init__(self, session: ResponsesWebSocketSession, state: _LaneState) -> None: + def __init__(self, session: ResponsesWebSocketSession[_SyncConnectionT], state: _LaneState) -> None: self._session = session self._state = state - def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> None: + @overload + def send( + self: ResponsesWebSocketLane[ResponsesConnection], event: ResponsesClientEvent | ResponsesClientEventParam + ) -> None: ... + + @overload + def send( + self: ResponsesWebSocketLane[BetaResponsesConnection], + event: BetaResponsesClientEvent | BetaResponsesClientEventParam, + ) -> None: ... + + def send(self, event: _ClientEvent) -> None: # Pin the transport before checking recovery so this write cannot migrate. connection = self._session.connection._connection self._session._check_send(self._state) # Serializing supported iterables can call back into the connection. - data, is_create = _command(event, self._state.stream_id) + data, is_create = _command( + event, + self._state.stream_id, + beta=isinstance(self._session.connection, (BetaResponsesConnection, AsyncBetaResponsesConnection)), + ) with self._session.connection._connection_lock: if self._session.connection._is_reconnecting: raise RuntimeError("The WebSocket connection is reconnecting; restore state before sending") @@ -361,7 +411,20 @@ def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> None: # Send directly to this physical socket so lane requests never enter it. connection.send(data) - def recv(self, *, timeout: float | None = None) -> ResponsesServerEvent: + @overload + def recv( + self: ResponsesWebSocketLane[ResponsesConnection], *, timeout: float | None = None + ) -> ResponsesServerEvent: ... + + @overload + def recv( + self: ResponsesWebSocketLane[BetaResponsesConnection], *, timeout: float | None = None + ) -> BetaResponsesServerEvent: ... + + @overload + def recv(self, *, timeout: float | None = None) -> _ServerEvent: ... + + def recv(self, *, timeout: float | None = None) -> _ServerEvent: if self._session._thread is None: raise RuntimeError("Enter the WebSocket session before receiving events") deadline = None if timeout is None else time.monotonic() + timeout @@ -376,7 +439,17 @@ def recv(self, *, timeout: float | None = None) -> ResponsesServerEvent: raise TimeoutError("Timed out waiting for a Responses WebSocket event") router.condition.wait(remaining) - def get_final_response(self, *, timeout: float | None = None) -> Response: + @overload + def get_final_response( + self: ResponsesWebSocketLane[ResponsesConnection], *, timeout: float | None = None + ) -> Response: ... + + @overload + def get_final_response( + self: ResponsesWebSocketLane[BetaResponsesConnection], *, timeout: float | None = None + ) -> BetaResponse: ... + + def get_final_response(self, *, timeout: float | None = None) -> _Response: """Consume this response's remaining events; failed/incomplete are final results.""" deadline = None if timeout is None else time.monotonic() + timeout generation = self._state.consumed_generation @@ -404,7 +477,7 @@ def close(self) -> None: self._session._router.detach(self._state) -class ResponsesWebSocketSession: +class ResponsesWebSocketSession(Generic[_SyncConnectionT]): """Own one existing connection and route its events with one background reader. Enter the session before sending. While active, consume events only through @@ -425,14 +498,14 @@ class ResponsesWebSocketSession: """ def __init__( - self, connection: ResponsesConnection, *, limits: ResponsesWebSocketLimits, isolate_lane_overflow: bool = False + self, connection: _SyncConnectionT, *, limits: ResponsesWebSocketLimits, isolate_lane_overflow: bool = False ) -> None: self.connection = connection self._router = _Router(limits, connection, isolate_lane_overflow=isolate_lane_overflow) self._thread: threading.Thread | None = None self._closed = False self._owns_connection = False - self.default = ResponsesWebSocketLane(self, self._router.default) + self.default: ResponsesWebSocketLane[_SyncConnectionT] = ResponsesWebSocketLane(self, self._router.default) def __enter__(self) -> Self: if self._thread is not None or self._closed: @@ -464,7 +537,7 @@ def __exit__( if exc is None: raise - def lane(self, stream_id: str | None = None) -> ResponsesWebSocketLane: + def lane(self, stream_id: str | None = None) -> ResponsesWebSocketLane[_SyncConnectionT]: lane = ResponsesWebSocketLane(self, self._router.register(stream_id)) if stream_id is None: self.default = lane @@ -525,23 +598,48 @@ def close(self) -> None: self._release_connection() -class AsyncResponsesWebSocketLane: +class AsyncResponsesWebSocketLane(Generic[_AsyncConnectionT]): """An async routed lane. Canceling a recv/final-response wait never closes it.""" - def __init__(self, session: AsyncResponsesWebSocketSession, state: _LaneState) -> None: + def __init__(self, session: AsyncResponsesWebSocketSession[_AsyncConnectionT], state: _LaneState) -> None: self._session = session self._state = state - async def send(self, event: ResponsesClientEvent | ResponsesClientEventParam) -> None: + @overload + async def send( + self: AsyncResponsesWebSocketLane[AsyncResponsesConnection], + event: ResponsesClientEvent | ResponsesClientEventParam, + ) -> None: ... + + @overload + async def send( + self: AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection], + event: BetaResponsesClientEvent | BetaResponsesClientEventParam, + ) -> None: ... + + async def send(self, event: _ClientEvent) -> None: # Pin the transport before checking recovery so this write cannot migrate. connection = self._session.connection._connection self._session._check_send(self._state) - data, is_create = _command(event, self._state.stream_id) + data, is_create = _command( + event, + self._state.stream_id, + beta=isinstance(self._session.connection, (BetaResponsesConnection, AsyncBetaResponsesConnection)), + ) self._session._router.prepare(self._state, is_create) # Bypass the legacy recovery queue, including a reconnect starting now. await connection.send(data) - async def recv(self) -> ResponsesServerEvent: + @overload + async def recv(self: AsyncResponsesWebSocketLane[AsyncResponsesConnection]) -> ResponsesServerEvent: ... + + @overload + async def recv(self: AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection]) -> BetaResponsesServerEvent: ... + + @overload + async def recv(self) -> _ServerEvent: ... + + async def recv(self) -> _ServerEvent: if self._session._task is None: raise RuntimeError("Enter the WebSocket session before receiving events") router = self._session._router @@ -553,7 +651,13 @@ async def recv(self) -> ResponsesServerEvent: self._state.wake.clear() await self._state.wake.wait() - async def get_final_response(self) -> Response: + @overload + async def get_final_response(self: AsyncResponsesWebSocketLane[AsyncResponsesConnection]) -> Response: ... + + @overload + async def get_final_response(self: AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection]) -> BetaResponse: ... + + async def get_final_response(self) -> _Response: generation = self._state.consumed_generation while ( self._state.accumulator.final is not None @@ -577,12 +681,12 @@ def close(self) -> None: self._session._router.detach(self._state) -class AsyncResponsesWebSocketSession: +class AsyncResponsesWebSocketSession(Generic[_AsyncConnectionT]): """Async counterpart of ResponsesWebSocketSession, owning the connection lifetime.""" def __init__( self, - connection: AsyncResponsesConnection, + connection: _AsyncConnectionT, *, limits: ResponsesWebSocketLimits, isolate_lane_overflow: bool = False, @@ -592,7 +696,9 @@ def __init__( self._task: asyncio.Task[None] | None = None self._close_task: asyncio.Task[None] | None = None self._closed = False - self.default = AsyncResponsesWebSocketLane(self, self._router.default) + self.default: AsyncResponsesWebSocketLane[_AsyncConnectionT] = AsyncResponsesWebSocketLane( + self, self._router.default + ) async def __aenter__(self) -> Self: if self._task is not None or self._closed: @@ -610,7 +716,7 @@ async def __aexit__( if exc is None: raise - def lane(self, stream_id: str | None = None) -> AsyncResponsesWebSocketLane: + def lane(self, stream_id: str | None = None) -> AsyncResponsesWebSocketLane[_AsyncConnectionT]: lane = AsyncResponsesWebSocketLane(self, self._router.register(stream_id)) if stream_id is None: self.default = lane diff --git a/tests/lib/responses/test_websocket_beta_helpers.py b/tests/lib/responses/test_websocket_beta_helpers.py new file mode 100644 index 0000000000..f218afe1fa --- /dev/null +++ b/tests/lib/responses/test_websocket_beta_helpers.py @@ -0,0 +1,224 @@ +from __future__ import annotations + +import json +import warnings +from copy import deepcopy +from typing import Any, AsyncGenerator +from contextlib import asynccontextmanager + +import httpx2 +import pytest +from websockets.sync.server import ServerConnection + +from openai import OpenAI, AsyncOpenAI +from openai._compat import model_parse +from openai._models import BaseModel +from openai.types.responses import Response +from openai.lib.responses_websocket import ( + ResponsesWebSocketError, + ResponsesWebSocketSession, + ResponsesWebSocketAccumulator, + AsyncResponsesWebSocketSession, +) +from openai.types.beta.beta_response import BetaResponse +from openai.types.beta.beta_response_failed_event import BetaResponseFailedEvent +from openai.types.beta.beta_response_created_event import BetaResponseCreatedEvent +from openai.types.beta.beta_responses_client_event import ResponseCreate as BetaResponseCreate +from openai.types.beta.beta_responses_server_event import ( + BetaResponseWsError, + BetaResponseWsCreated, + BetaResponseWsCompleted, +) +from openai.types.beta.beta_response_completed_event import BetaResponseCompletedEvent +from openai.types.beta.beta_response_incomplete_event import BetaResponseIncompleteEvent +from openai.types.beta.beta_response_output_item_done_event import BetaResponseOutputItemDoneEvent + +from .test_websocket_session import LIMITS, Driver, bypass_proxy as bypass_proxy, script_server, response_event + + +@asynccontextmanager +async def beta_session(mode: str, url: str) -> AsyncGenerator[Driver]: + if mode == "sync": + with ( + OpenAI(api_key="fake-beta-key", base_url=url, http_client=httpx2.Client(trust_env=False)) as client, + client.beta.responses.connect() as connection, + ResponsesWebSocketSession(connection, limits=LIMITS) as session, + ): + yield Driver(session, mode) + else: + async with ( + AsyncOpenAI( + api_key="fake-beta-key", base_url=url, http_client=httpx2.AsyncClient(trust_env=False) + ) as client, + client.beta.responses.connect() as connection, + AsyncResponsesWebSocketSession(connection, limits=LIMITS) as session, + ): + yield Driver(session, mode) + + +def reasoning(marker: str) -> dict[str, Any]: + return { + "type": "reasoning", + "id": "reason", + "status": "completed", + "summary": [{"type": "summary_text", "text": marker}], + "encrypted_content": "opaque-" + marker, + "agent": {"agent_name": "synthetic-agent"}, + "opaque_item": {"nested": [marker, {"keep": True}]}, + } + + +def beta_frame(kind: str, output: list[dict[str, Any]], sequence: int) -> dict[str, Any]: + event = response_event(kind, "beta", output=output) + event["sequence_number"] = sequence + event["opaque_event"] = {"nested": ["keep", None]} + event["response"]["opaque_response"] = {"nested": [kind]} + return event + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("terminal", ["completed", "failed", "incomplete"]) +@pytest.mark.parametrize("scaffold", [False, True]) +@pytest.mark.parametrize("full_output", [False, True]) +async def test_beta_native_final_and_caller_fed_family( + mode: str, + terminal: str, + scaffold: bool, + full_output: bool, +) -> None: + frames = [beta_frame("created", [], 1)] if scaffold else [] + if scaffold: + frames.append( + { + "type": "response.output_item.done", + "sequence_number": 2, + "stream_id": "beta", + "output_index": 0, + "item": reasoning("item-done"), + } + ) + final = beta_frame(terminal, [reasoning("terminal")] if full_output else [], len(frames) + 1) + frames.append(final) + models = { + "completed": BetaResponseCompletedEvent, + "failed": BetaResponseFailedEvent, + "incomplete": BetaResponseIncompleteEvent, + } + model_parse(models[terminal], final) + if scaffold: + model_parse(BetaResponseCreatedEvent, frames[0]) + model_parse(BetaResponseOutputItemDoneEvent, frames[1]) + + def script(socket: ServerConnection) -> None: + assert json.loads(socket.recv(timeout=5)) == {"type": "response.create", "stream_id": "beta"} + for event in frames: + socket.send(json.dumps(event)) + + with script_server(script) as url: + async with beta_session(mode, url) as driver: + lane = driver.session.lane("beta") + command = BetaResponseCreate(type="response.create") if scaffold else {"type": "response.create"} + before_command = deepcopy(command.to_dict() if isinstance(command, BaseModel) else command) + await driver.call(lane, "send", command) + assert (command.to_dict() if isinstance(command, BaseModel) else command) == before_command + accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) + received = [] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + for expected in frames: + event = await driver.call(lane, "recv") + assert event.to_dict() == expected + received.append(event) + accumulator.add_event(event) + managed = await driver.call(lane, "get_final_response") + collected = accumulator.get_final_response() + assert isinstance(managed, BetaResponse) + assert managed is not received[-1].response + if full_output: + assert managed.output[0] is not received[-1].response.output[0] + assert isinstance(collected, BetaResponse) + assert managed.to_dict() == final["response"] + assert collected.to_dict() == final["response"] + assert not caught + assert [event.to_dict() for event in received] == frames + saved = accumulator.detailed_snapshot() + collected.output.clear() + assert accumulator.get_final_response().to_dict() == final["response"] + accumulator.reset() + assert saved["response"]["opaque_response"] == final["response"]["opaque_response"] + assert [event.to_dict() for event in received] == frames + assert managed.to_dict() == final["response"] + + +@pytest.mark.parametrize("beta", [False, True]) +def test_caller_fed_selected_family_reset_and_rejection(beta: bool) -> None: + # Runtime rejection from an untyped caller; separate native witnesses check + # the public positive and negative type contracts without casts or ignores. + accumulator: Any = ( + ResponsesWebSocketAccumulator(response_type=BetaResponse) if beta else ResponsesWebSocketAccumulator() + ) + chosen = BetaResponseCreatedEvent if beta else None + from openai.types.responses.responses_server_event import ResponseWsCreated + + model = chosen or ResponseWsCreated + wrong = ResponseWsCreated if beta else BetaResponseCreatedEvent + created = beta_frame("created", [], 1) + original = deepcopy(created) + accumulator.add_event(model_parse(model, created)) + snapshot = accumulator.detailed_snapshot() + with pytest.raises(ValueError, match="missing a valid response"): + accumulator.add_event(model_parse(wrong, created)) + assert accumulator.detailed_snapshot() == snapshot + accumulator.reset() + accumulator.add_event(model_parse(model, created)) + assert accumulator.detailed_snapshot() == snapshot + assert created == original + assert isinstance(model_parse(model, created).response, BetaResponse if beta else Response) + + +@pytest.mark.parametrize("mismatch", ["lane", "response"]) +def test_beta_caller_fed_identity_rejection_keeps_prior_snapshot(mismatch: str) -> None: + accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) + created = beta_frame("created", [], 1) + accumulator.add_event(model_parse(BetaResponseWsCreated, created)) + saved = accumulator.detailed_snapshot() + final = beta_frame("completed", [reasoning("authoritative")], 2) + other = deepcopy(final) + if mismatch == "lane": + other["stream_id"] = "other" + else: + other["response"]["id"] = "other-response" + with pytest.raises(ValueError, match="another"): + accumulator.add_event(model_parse(BetaResponseWsCompleted, other)) + assert accumulator.detailed_snapshot() == saved + with pytest.raises(RuntimeError, match="No terminal"): + accumulator.get_final_response() + accumulator.add_event(model_parse(BetaResponseWsCompleted, final)) + assert accumulator.get_final_response().to_dict() == final["response"] + assert saved["terminal_type"] is None + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +async def test_beta_protocol_error_preserves_original_native_event(mode: str) -> None: + error = { + "type": "error", + "stream_id": "beta", + "error": {"type": "server_error", "message": "synthetic detail", "code": "synthetic"}, + "agent": {"agent_name": "synthetic-agent"}, + "opaque": {"keep": True}, + } + model_parse(BetaResponseWsError, error) + + def script(socket: ServerConnection) -> None: + assert json.loads(socket.recv(timeout=5))["stream_id"] == "beta" + socket.send(json.dumps(error)) + + with script_server(script) as url: + async with beta_session(mode, url) as driver: + lane = driver.session.lane("beta") + await driver.call(lane, "send", BetaResponseCreate(type="response.create")) + with pytest.raises(ResponsesWebSocketError) as raised: + await driver.call(lane, "get_final_response") + assert isinstance(raised.value.event, BetaResponseWsError) + assert raised.value.event.to_dict() == error + assert str(raised.value) == "The Responses WebSocket request returned an error event" diff --git a/tests/lib/responses/websocket_beta_typing.py b/tests/lib/responses/websocket_beta_typing.py new file mode 100644 index 0000000000..795d2d5b89 --- /dev/null +++ b/tests/lib/responses/websocket_beta_typing.py @@ -0,0 +1,115 @@ +"""Public type-checker witnesses; these functions do not open live connections.""" + +from typing_extensions import assert_type + +from openai import OpenAI, AsyncOpenAI +from openai.types.responses import Response +from openai.lib.responses_websocket import ( + ResponsesWebSocketLane, + ResponsesWebSocketLimits, + ResponsesWebSocketSession, + AsyncResponsesWebSocketLane, + ResponsesWebSocketAccumulator, + AsyncResponsesWebSocketSession, +) +from openai.types.beta.beta_response import BetaResponse +from openai.resources.responses.responses import ResponsesConnection, AsyncResponsesConnection +from openai.resources.beta.responses.responses import ( + ResponsesConnection as BetaResponsesConnection, + AsyncResponsesConnection as AsyncBetaResponsesConnection, +) +from openai.types.beta.beta_responses_client_event import ResponseCreate as BetaResponseCreate +from openai.types.beta.beta_responses_server_event import BetaResponsesServerEvent +from openai.types.responses.responses_server_event import ResponsesServerEvent + + +def stable_sync(client: OpenAI, limits: ResponsesWebSocketLimits) -> None: + with client.responses.connect() as connection: + with ResponsesWebSocketSession(connection, limits=limits) as session: + lane = session.lane("stable") + assert_type(session, ResponsesWebSocketSession[ResponsesConnection]) + assert_type(lane, ResponsesWebSocketLane[ResponsesConnection]) + lane.send({"type": "response.create"}) + assert_type(lane.recv(), ResponsesServerEvent) + assert_type(lane.get_final_response(), Response) + accumulator = ResponsesWebSocketAccumulator() + accumulator.add_event(lane.recv()) + assert_type(accumulator.get_final_response(), Response) + + +def beta_sync(client: OpenAI, limits: ResponsesWebSocketLimits) -> None: + with client.beta.responses.connect() as connection: + with ResponsesWebSocketSession(connection, limits=limits) as session: + assert_type(session, ResponsesWebSocketSession[BetaResponsesConnection]) + assert_type(session.default, ResponsesWebSocketLane[BetaResponsesConnection]) + lane = session.lane("beta") + assert_type(lane, ResponsesWebSocketLane[BetaResponsesConnection]) + explicit_lane: ResponsesWebSocketLane[BetaResponsesConnection] = lane + assert_type(explicit_lane, ResponsesWebSocketLane[BetaResponsesConnection]) + lane.send(BetaResponseCreate(type="response.create")) + lane.send({"type": "response.create"}) + assert_type(lane.recv(), BetaResponsesServerEvent) + assert_type(lane.get_final_response(), BetaResponse) + accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) + accumulator.add_event(lane.recv()) + assert_type(accumulator.get_final_response(), BetaResponse) + explicit: ResponsesWebSocketAccumulator[BetaResponse] = ResponsesWebSocketAccumulator( + response_type=BetaResponse + ) + explicit.add_event(lane.recv()) + assert_type(explicit.get_final_response(), BetaResponse) + + +async def stable_async(client: AsyncOpenAI, limits: ResponsesWebSocketLimits) -> None: + async with client.responses.connect() as connection: + async with AsyncResponsesWebSocketSession(connection, limits=limits) as session: + assert_type(session, AsyncResponsesWebSocketSession[AsyncResponsesConnection]) + lane = session.lane("stable") + assert_type(lane, AsyncResponsesWebSocketLane[AsyncResponsesConnection]) + await lane.send({"type": "response.create"}) + assert_type(await lane.recv(), ResponsesServerEvent) + assert_type(await lane.get_final_response(), Response) + accumulator = ResponsesWebSocketAccumulator() + accumulator.add_event(await lane.recv()) + assert_type(accumulator.get_final_response(), Response) + + +async def beta_async(client: AsyncOpenAI, limits: ResponsesWebSocketLimits) -> None: + async with client.beta.responses.connect() as connection: + async with AsyncResponsesWebSocketSession(connection, limits=limits) as session: + assert_type(session, AsyncResponsesWebSocketSession[AsyncBetaResponsesConnection]) + assert_type(session.default, AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection]) + lane = session.lane("beta") + assert_type(lane, AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection]) + explicit_lane: AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection] = lane + assert_type(explicit_lane, AsyncResponsesWebSocketLane[AsyncBetaResponsesConnection]) + await lane.send(BetaResponseCreate(type="response.create")) + assert_type(await lane.recv(), BetaResponsesServerEvent) + assert_type(await lane.get_final_response(), BetaResponse) + accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) + accumulator.add_event(await lane.recv()) + assert_type(accumulator.get_final_response(), BetaResponse) + + +def caller_fed_stable(connection: ResponsesConnection) -> None: + accumulator = ResponsesWebSocketAccumulator() + accumulator.add_event(connection.recv()) + assert_type(accumulator.get_final_response(), Response) + + +def caller_fed_beta(connection: BetaResponsesConnection) -> None: + accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) + accumulator.add_event(connection.recv()) + assert_type(accumulator.get_final_response(), BetaResponse) + + +async def caller_fed_async_stable(connection: AsyncResponsesConnection) -> None: + accumulator = ResponsesWebSocketAccumulator() + accumulator.add_event(await connection.recv()) + assert_type(accumulator.get_final_response(), Response) + + +async def caller_fed_async_beta(connection: AsyncBetaResponsesConnection) -> None: + accumulator = ResponsesWebSocketAccumulator(response_type=BetaResponse) + accumulator.add_event(await connection.recv()) + assert_type(accumulator.get_final_response(), BetaResponse)