refactor(websocket): add SDK-owned Responses policy helpers

src/openai/_websocket_connection.py

@@ -1,22 +1,191 @@
 from __future__ import annotations
 
+import json
 import time
 import random
 import logging
 import threading
 import contextlib
-from typing import TYPE_CHECKING, Callable, Awaitable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Mapping,
+    Callable,
+    Iterator,
+    Awaitable,
+    Generator,
+    AsyncIterator,
+    AsyncGenerator,
+    cast,
+)
 
 import anyio
 
 from ._types import Query, Headers
+from ._utils import is_given
+from ._models import FinalRequestOptions
+from ._exceptions import OpenAIError, WebSocketConnectionClosedError
 from ._send_queue import SendQueue, SendReceipt
+from ._base_client import _merge_mappings
 from .types.websocket_reconnection import ReconnectingEvent, ReconnectingOverrides, is_recoverable_close
 
 if TYPE_CHECKING:
+    import httpx2
     from websockets.sync.client import ClientConnection as WebSocketConnection
     from websockets.asyncio.client import ClientConnection as AsyncWebSocketConnection
 
+    from ._client import OpenAI, AsyncOpenAI
+
+
+def _websocket_request_options(
+    client: OpenAI | AsyncOpenAI, url: httpx2.URL, extra_query: Query, extra_headers: Headers
+) -> FinalRequestOptions:
+    try:
+        query_params = client._websocket_query_params
+    except AttributeError:
+        query_params = client._websocket_query_params_fallback
+    url = url.copy_with(params=query_params(url, extra_query))
+    url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme))
+    return FinalRequestOptions.construct(
+        method="get", url=str(url), headers=dict(extra_headers), security={"bearer_auth": True}
+    )
+
+
+def _websocket_connect_options(
+    client: OpenAI | AsyncOpenAI,
+    options: FinalRequestOptions,
+    connection_options: Mapping[str, Any],
+    default_headers: Headers,
+    logger: logging.Logger,
+) -> tuple[str, dict[str, str], dict[str, Any]]:
+    url = client._prepare_url(options.url).copy_merge_params(client.qs.stringify(cast(Any, options.params)))
+    url = url.copy_with(scheme={"http": "ws", "https": "wss"}.get(url.scheme, url.scheme))
+    headers = {
+        key.lower(): (key, value)
+        for header_set in (
+            client.auth_headers,
+            default_headers,
+            client.default_headers,
+            options.headers if is_given(options.headers) else {},
+        )
+        for key, value in header_set.items()
+    }
+    logger.debug("Connecting to WebSocket API")
+    if connection_options:
+        logger.debug("Custom WebSocket connection options provided")
+    websocket_options = dict(connection_options)
+    websocket_options.setdefault("max_size", None)
+    if websocket_options.get("max_size") is None:
+        websocket_options.setdefault("compression", None)
+    return str(url), _merge_mappings(dict(headers.values()), {}), websocket_options
+
+
+def _connect_websocket(
+    client: OpenAI,
+    prepare_url: Callable[[], httpx2.URL],
+    extra_query: Query,
+    extra_headers: Headers,
+    connection_options: Mapping[str, Any],
+    default_headers: Headers,
+    logger: logging.Logger,
+) -> WebSocketConnection:
+    try:
+        from websockets.sync.client import connect
+    except ImportError as exc:
+        raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc
+    options = client._prepare_options(_websocket_request_options(client, prepare_url(), extra_query, extra_headers))
+    url, headers, websocket_options = _websocket_connect_options(
+        client, options, connection_options, default_headers, logger
+    )
+    return connect(url, user_agent_header=None, additional_headers=headers, **websocket_options)
+
+
+async def _async_connect_websocket(
+    client: AsyncOpenAI,
+    prepare_url: Callable[[], httpx2.URL],
+    extra_query: Query,
+    extra_headers: Headers,
+    connection_options: Mapping[str, Any],
+    default_headers: Headers,
+    logger: logging.Logger,
+) -> AsyncWebSocketConnection:
+    try:
+        from .lib._websocket import _WebSocketConnect as connect
+    except ImportError as exc:
+        raise OpenAIError("You need to install `openai[realtime]` to use this method") from exc
+    options = await client._prepare_options(
+        _websocket_request_options(client, prepare_url(), extra_query, extra_headers)
+    )
+    url, headers, websocket_options = _websocket_connect_options(
+        client, options, connection_options, default_headers, logger
+    )
+    return await connect(url, user_agent_header=None, additional_headers=headers, **websocket_options)
+
+
+def _record_websocket_progress(
+    connection: WebSocketConnectionBase | AsyncWebSocketConnectionBase, event: object, *, raw: bool = False
+) -> None:
+    if raw:
+        if not connection._reconnect_attempt:
+            return
+        try:
+            event = json.loads(cast("str | bytes", event))
+        except (ValueError, RecursionError):
+            return
+    event_type = event.get("type") if isinstance(event, dict) else None if raw else getattr(event, "type", None)
+    if isinstance(event_type, str) and event_type and event_type != "error":
+        connection._reconnect_attempt = 0
+
+
+@contextlib.contextmanager
+def _websocket_entry_guard(
+    connection: WebSocketConnection, clear_connection: Callable[[], None]
+) -> Generator[None, None, None]:
+    try:
+        yield
+    except BaseException:
+        clear_connection()
+        try:
+            connection.close()
+        except BaseException:
+            # Preserve the exception that prevented successful entry.
+            pass
+        raise
+
+
+def _exit_websocket(connection: WebSocketConnectionBase | None, exc: BaseException | None) -> None:
+    if connection is not None:
+        try:
+            connection.close()
+        except Exception:
+            if exc is None:
+                raise
+
+
+async def _async_exit_websocket(connection: AsyncWebSocketConnectionBase | None, exc: BaseException | None) -> None:
+    if connection is not None:
+        try:
+            await connection.close()
+        except Exception:
+            if exc is None:
+                raise
+
+
+@contextlib.asynccontextmanager
+async def _async_websocket_entry_guard(
+    connection: AsyncWebSocketConnection, clear_connection: Callable[[], None]
+) -> AsyncGenerator[None, None]:
+    try:
+        yield
+    except BaseException:
+        clear_connection()
+        try:
+            await connection.close()
+        except BaseException:
+            # Preserve the exception that prevented successful entry.
+            pass
+        raise
+
 
 class WebSocketConnectionBase:
     def __init__(
@@ -52,6 +221,28 @@ def __init__(
         if interruptible_reconnect:
             self._close_event = threading.Event()
 
+    def _iter_events(self) -> Iterator[Any]:
+        from websockets.exceptions import ConnectionClosed, ConnectionClosedOK, ConnectionClosedError
+
+        while True:
+            try:
+                yield cast(Any, self).recv()
+            except ConnectionClosed as exc:
+                reconnect_exc = (
+                    ConnectionClosedError(exc.rcvd, exc.sent, exc.rcvd_then_sent)
+                    if isinstance(exc, ConnectionClosedOK)
+                    else exc
+                )
+                if not self._reconnect(reconnect_exc):
+                    unsent = self._send_queue.drain()
+                    if unsent:
+                        raise WebSocketConnectionClosedError(
+                            "WebSocket connection closed with unsent messages", unsent_messages=unsent
+                        ) from exc
+                    if isinstance(exc, ConnectionClosedOK):
+                        return
+                    raise
+
     def send_raw(self, data: bytes | str) -> None:
         self._send(data)
 
@@ -258,6 +449,28 @@ def __init__(
         self._send_queue = send_queue if send_queue is not None else SendQueue()
         self._logger = logger
 
+    async def _iter_events(self) -> AsyncIterator[Any]:
+        from websockets.exceptions import ConnectionClosed, ConnectionClosedOK, ConnectionClosedError
+
+        while True:
+            try:
+                yield await cast(Any, self).recv()
+            except ConnectionClosed as exc:
+                reconnect_exc = (
+                    ConnectionClosedError(exc.rcvd, exc.sent, exc.rcvd_then_sent)
+                    if isinstance(exc, ConnectionClosedOK)
+                    else exc
+                )
+                if not await self._reconnect(reconnect_exc):
+                    unsent = self._send_queue.drain()
+                    if unsent:
+                        raise WebSocketConnectionClosedError(
+                            "WebSocket connection closed with unsent messages", unsent_messages=unsent
+                        ) from exc
+                    if isinstance(exc, ConnectionClosedOK):
+                        return
+                    raise
+
     async def send_raw(self, data: bytes | str) -> None:
         await self._send(data)