From 6b49cb963f80867fa82eeed0e435a32eef33563f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 07:22:57 -0700 Subject: [PATCH] feat(python): adapt typed Rust streams --- litellm/llms/custom_httpx/llm_http_handler.py | 10 +- litellm/rust_bridge/responses_websocket.py | 61 +-- litellm/rust_bridge/streaming.py | 354 ++++++++++++++++++ .../custom_httpx/test_llm_http_handler.py | 10 +- .../responses/test_rust_bridge_websocket.py | 129 ++++++- .../rust_bridge/test_streaming.py | 327 ++++++++++++++++ 6 files changed, 851 insertions(+), 40 deletions(-) create mode 100644 litellm/rust_bridge/streaming.py create mode 100644 tests/test_litellm/rust_bridge/test_streaming.py diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a00e0fa0ae1..2002a1023f3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -163,7 +163,11 @@ def _rust_responses_websocket_enabled( raw_request_override: Final = litellm_params.get("rust") request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return custom_llm_provider == "openai" and rust_enabled(request_override=request_override) + if custom_llm_provider != "openai" or not rust_enabled(request_override=request_override): + return False + from litellm.rust_bridge.streaming import supports_streaming + + return supports_streaming("responses", custom_llm_provider, "websocket") def _anthropic_messages_with_core_engine( @@ -6513,7 +6517,9 @@ class BaseLLMHTTPHandler: from litellm.rust_bridge import responses_websocket as rust_responses_websocket rust_execution: Final = await rust_responses_websocket.connect( - url=ws_url, + provider="openai", + api_key=api_key, + api_base=api_base, headers={str(key): str(value) for key, value in headers.items()}, timeout=timeout, ) diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 4c1c2330f5c..54e8ee4bb2a 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -2,11 +2,16 @@ from __future__ import annotations +import json +from collections.abc import Mapping +from types import MappingProxyType from typing import Final, Protocol import httpx +from pydantic import TypeAdapter from websockets.exceptions import ConnectionClosedOK +from litellm.rust_bridge import streaming from litellm.rust_bridge.bindings import UNSET, NativeBinding, Unset from litellm.rust_bridge.runtime import ( BridgeErrorContext, @@ -19,11 +24,13 @@ from litellm.rust_bridge.runtime import ( ) from litellm.rust_bridge.timeouts import timeout_to_seconds +_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + class RustResponsesWebSocket(Protocol): - async def send_text(self, text: str) -> None: ... + async def send_event(self, event: Mapping[str, object]) -> None: ... - async def recv_text(self) -> str | None: ... + async def recv_event(self) -> Mapping[str, object] | None: ... async def close(self) -> None: ... @@ -32,13 +39,15 @@ class RustResponsesWebSocketConnection(Protocol): @classmethod async def connect( cls, - url: str, - headers: dict[str, str], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, timeout_seconds: float | None, ) -> RustResponsesWebSocket: ... -_CONNECTION: Final = NativeBinding[type[RustResponsesWebSocketConnection]]("ResponsesWebSocketConnection") +_CONNECTION: Final = NativeBinding[type[RustResponsesWebSocketConnection]]("ResponsesWebSocketSession") def set_rust_responses_websocket( @@ -53,52 +62,62 @@ def load_rust_responses_websocket() -> type[RustResponsesWebSocketConnection] | class _ConnectionAdapter: - def __init__(self, connection: RustResponsesWebSocket): + def __init__(self, connection: RustResponsesWebSocket, context: BridgeErrorContext): self._connection: Final = connection + self._context: Final = context self.core_engine: Final = CoreEngine.RUST async def send(self, text: str) -> None: + event: Final = _EVENT_ADAPTER.validate_json(text) await acall( - lambda: self._connection.send_text(text), - BridgeErrorContext(route="responses websocket", provider="openai", model=""), + lambda: self._connection.send_event(event), + self._context, ) async def recv(self) -> str: - message: Final = await acall( - self._connection.recv_text, - BridgeErrorContext(route="responses websocket", provider="openai", model=""), + event: Final = await acall( + self._connection.recv_event, + self._context, ) - if message is None: + if event is None: raise ConnectionClosedOK(None, None) - return message + return json.dumps(dict(event), separators=(",", ":")) # mutable-ok: JSON needs a concrete dict async def close(self) -> None: await acall( self._connection.close, - BridgeErrorContext(route="responses websocket", provider="openai", model=""), + self._context, ) async def connect( *, - url: str, - headers: dict[str, str], + provider: str, + api_key: str | None, + api_base: str | None, + headers: Mapping[str, str], timeout: float | httpx.Timeout | None, ) -> ExecutionResult[_ConnectionAdapter | None]: + context: Final = BridgeErrorContext(route="responses websocket", provider=provider, model="") + if not streaming.supports_streaming("responses", provider, "websocket"): + return ExecutionResult(value=None, source=CoreEngine.PYTHON) connection_type: Final = load_rust_responses_websocket() + credentials: Final = None if api_key is None else MappingProxyType({"api_key": api_key}) native_call: Final = ( None if connection_type is None else lambda: connection_type.connect( - url=url, - headers=headers, - timeout_seconds=timeout_to_seconds(timeout), + provider, + credentials, + api_base, + headers, + timeout_to_seconds(timeout), ) ) return await ainvoke( native_call=native_call, fallback=async_none, - adapt=_ConnectionAdapter, + adapt=lambda connection: _ConnectionAdapter(connection, context), mode=FallbackMode.PYTHON, - context=BridgeErrorContext(route="responses websocket", provider="openai", model=""), + context=context, ) diff --git a/litellm/rust_bridge/streaming.py b/litellm/rust_bridge/streaming.py new file mode 100644 index 00000000000..1b7fc1bb5e8 --- /dev/null +++ b/litellm/rust_bridge/streaming.py @@ -0,0 +1,354 @@ +from __future__ import annotations + +import json +from collections.abc import AsyncIterator, Iterator, Mapping +from functools import lru_cache +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable + +import httpx + +from litellm.rust_bridge.bindings import UNSET, NativeBinding, Unset +from litellm.rust_bridge.runtime import ( + BridgeErrorContext, + RustHandled, + aattempt, + acall, + attempt, + call, +) +from litellm.rust_bridge.timeouts import timeout_to_seconds + +if TYPE_CHECKING: + from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponsesAPIStreamingResponse + +StreamApi: TypeAlias = Literal["chat_completions", "messages", "responses"] +StreamTransport: TypeAlias = Literal["http", "websocket"] +Event: TypeAlias = Mapping[str, object] + + +@runtime_checkable +class RustEventStream(Protocol): + @property + def metadata(self) -> Mapping[str, object]: ... + + def next_event(self) -> Event | None: ... + + async def anext_event(self) -> Event | None: ... + + def close(self) -> None: ... + + async def aclose(self) -> None: ... + + +class RustStreamOpen(Protocol): + def __call__( + self, + request: Mapping[str, object], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> object: ... + + +class RustAsyncStreamOpen(Protocol): + async def __call__( + self, + request: Mapping[str, object], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> object: ... + + +class RustStreamCapability(Protocol): + def __call__(self, api: str, provider: str, transport: str) -> bool: ... + + +_CAPABILITY: Final = NativeBinding[RustStreamCapability]("supports_streaming") +_CHAT: Final = NativeBinding[RustStreamOpen]("chat_completions_stream") +_ACHAT: Final = NativeBinding[RustAsyncStreamOpen]("achat_completions_stream") +_MESSAGES: Final = NativeBinding[RustStreamOpen]("messages_stream") +_AMESSAGES: Final = NativeBinding[RustAsyncStreamOpen]("amessages_stream") +_RESPONSES: Final = NativeBinding[RustStreamOpen]("responses_stream") +_ARESPONSES: Final = NativeBinding[RustAsyncStreamOpen]("aresponses_stream") + + +def set_rust_streaming( + *, + capability: RustStreamCapability | None | Unset = UNSET, + chat: RustStreamOpen | None | Unset = UNSET, + achat: RustAsyncStreamOpen | None | Unset = UNSET, + messages: RustStreamOpen | None | Unset = UNSET, + amessages: RustAsyncStreamOpen | None | Unset = UNSET, + responses: RustStreamOpen | None | Unset = UNSET, + aresponses: RustAsyncStreamOpen | None | Unset = UNSET, +) -> None: + """Override native stream bindings in tests; passing ``None`` restores discovery.""" + + _CAPABILITY.update(capability) + _CHAT.update(chat) + _ACHAT.update(achat) + _MESSAGES.update(messages) + _AMESSAGES.update(amessages) + _RESPONSES.update(responses) + _ARESPONSES.update(aresponses) + + +def supports_streaming(api: StreamApi, provider: str, transport: StreamTransport = "http") -> bool: + capability: Final = _CAPABILITY.load() + if capability is None: + return False + try: + return capability(api, provider, transport) is True + except Exception: # noqa: BLE001 # extension-defined errors cannot be named statically + return False + + +def _sync_opener(api: StreamApi) -> RustStreamOpen | None: + match api: + case "chat_completions": + return _CHAT.load() + case "messages": + return _MESSAGES.load() + case "responses": + return _RESPONSES.load() + + +def _async_opener(api: StreamApi) -> RustAsyncStreamOpen | None: + match api: + case "chat_completions": + return _ACHAT.load() + case "messages": + return _AMESSAGES.load() + case "responses": + return _ARESPONSES.load() + + +def _context(api: StreamApi, provider: str, request: Mapping[str, object]) -> BridgeErrorContext: + model: Final = request.get("model") + return BridgeErrorContext( + route=f"{api.replace('_', ' ')} stream", + provider=provider, + model=model if isinstance(model, str) else "", + ) + + +def _adapt_stream( + stream: object, + *, + context: BridgeErrorContext, +) -> TypedEventStreamAdapter: + if not isinstance(stream, RustEventStream): + raise TypeError("native stream opener returned an invalid stream") + return TypedEventStreamAdapter(stream, context) + + +class TypedEventStreamAdapter: + def __init__(self, stream: RustEventStream, context: BridgeErrorContext) -> None: + self._stream: Final = stream + self._context: Final = context + self.metadata: Final = stream.metadata + self._mode: Literal["sync", "async"] | None = None + + def _claim(self, mode: Literal["sync", "async"]) -> None: + if self._mode is None: + self._mode = mode + return + if self._mode != mode: + raise RuntimeError("native stream cannot mix synchronous and asynchronous consumption") + + def __iter__(self) -> Iterator[Event]: + self._claim("sync") + return self + + def __next__(self) -> Event: + self._claim("sync") + event: Final = call(self._stream.next_event, self._context) + if event is None: + raise StopIteration + return event + + def __aiter__(self) -> AsyncIterator[Event]: + self._claim("async") + return self + + async def __anext__(self) -> Event: + self._claim("async") + event: Final = await acall(self._stream.anext_event, self._context) + if event is None: + raise StopAsyncIteration + return event + + def close(self) -> None: + call(self._stream.close, self._context) + + async def aclose(self) -> None: + await acall(self._stream.aclose, self._context) + + +class MessagesSseStreamAdapter: + def __init__(self, events: TypedEventStreamAdapter) -> None: + self._events: Final = events + self.metadata: Final = events.metadata + + def __iter__(self) -> Iterator[bytes]: + return (_event_to_sse(event) for event in self._events) + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for event in self._events: + yield _event_to_sse(event) + + def close(self) -> None: + self._events.close() + + async def aclose(self) -> None: + await self._events.aclose() + + +class ResponsesSdkEventStreamAdapter: + def __init__(self, events: TypedEventStreamAdapter) -> None: + self._events: Final = events + self.metadata: Final = events.metadata + + def __iter__(self) -> Iterator[ResponsesAPIStreamingResponse]: + return (_responses_event_to_sdk(event) for event in self._events) + + async def __aiter__(self) -> AsyncIterator[ResponsesAPIStreamingResponse]: + async for event in self._events: + yield _responses_event_to_sdk(event) + + def close(self) -> None: + self._events.close() + + async def aclose(self) -> None: + await self._events.aclose() + + +def _responses_event_to_sdk(event: Event) -> ResponsesAPIStreamingResponse: + from litellm.types.llms.openai import GenericEvent + + event_type: Final = event.get("type") + model: Final = _responses_event_models().get(event_type) if isinstance(event_type, str) else None + return (model or GenericEvent).model_validate(event) + + +@lru_cache(maxsize=1) +def _responses_event_models() -> Mapping[str, type[BaseLiteLLMOpenAIResponseObject]]: + from litellm.types.llms import openai as openai_types + + return MappingProxyType( + { + "response.created": openai_types.ResponseCreatedEvent, + "response.in_progress": openai_types.ResponseInProgressEvent, + "response.completed": openai_types.ResponseCompletedEvent, + "response.failed": openai_types.ResponseFailedEvent, + "response.incomplete": openai_types.ResponseIncompleteEvent, + "response.reasoning_summary_part.added": openai_types.ResponsePartAddedEvent, + "response.reasoning_summary_text.delta": openai_types.ReasoningSummaryTextDeltaEvent, + "response.reasoning_summary_text.done": openai_types.ReasoningSummaryTextDoneEvent, + "response.reasoning_summary_part.done": openai_types.ReasoningSummaryPartDoneEvent, + "response.output_item.added": openai_types.OutputItemAddedEvent, + "response.output_item.done": openai_types.OutputItemDoneEvent, + "response.content_part.added": openai_types.ContentPartAddedEvent, + "response.content_part.done": openai_types.ContentPartDoneEvent, + "response.output_text.delta": openai_types.OutputTextDeltaEvent, + "response.output_text.annotation.added": openai_types.OutputTextAnnotationAddedEvent, + "response.output_text.done": openai_types.OutputTextDoneEvent, + "response.refusal.delta": openai_types.RefusalDeltaEvent, + "response.refusal.done": openai_types.RefusalDoneEvent, + "response.function_call_arguments.delta": openai_types.FunctionCallArgumentsDeltaEvent, + "response.function_call_arguments.done": openai_types.FunctionCallArgumentsDoneEvent, + "response.file_search_call.in_progress": openai_types.FileSearchCallInProgressEvent, + "response.file_search_call.searching": openai_types.FileSearchCallSearchingEvent, + "response.file_search_call.completed": openai_types.FileSearchCallCompletedEvent, + "response.web_search_call.in_progress": openai_types.WebSearchCallInProgressEvent, + "response.web_search_call.searching": openai_types.WebSearchCallSearchingEvent, + "response.web_search_call.completed": openai_types.WebSearchCallCompletedEvent, + "response.mcp_list_tools.in_progress": openai_types.MCPListToolsInProgressEvent, + "response.mcp_list_tools.completed": openai_types.MCPListToolsCompletedEvent, + "response.mcp_list_tools.failed": openai_types.MCPListToolsFailedEvent, + "response.mcp_call.in_progress": openai_types.MCPCallInProgressEvent, + "response.mcp_call_arguments.delta": openai_types.MCPCallArgumentsDeltaEvent, + "response.mcp_call_arguments.done": openai_types.MCPCallArgumentsDoneEvent, + "response.mcp_call.completed": openai_types.MCPCallCompletedEvent, + "response.mcp_call.failed": openai_types.MCPCallFailedEvent, + "image_generation.partial_image": openai_types.ImageGenerationPartialImageEvent, + "error": openai_types.ErrorEvent, + } + ) + + +def _event_to_sse(event: Event) -> bytes: + payload: Final = dict(event) # mutable-ok: JSON requires a concrete dict + return f"data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def open_stream( + *, + api: StreamApi, + provider: str, + request: Mapping[str, object], + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout: float | httpx.Timeout | None, +) -> TypedEventStreamAdapter | None: + if not supports_streaming(api, provider): + return None + opener: Final = _sync_opener(api) + context: Final = _context(api, provider, request) + result: Final = attempt( + native_call=( + None + if opener is None + else lambda: opener( + request, + provider, + credentials, + api_base, + extra_headers, + timeout_to_seconds(timeout), + ) + ), + adapt=lambda stream: _adapt_stream(stream, context=context), + context=context, + ) + return result.value if isinstance(result, RustHandled) else None + + +async def aopen_stream( + *, + api: StreamApi, + provider: str, + request: Mapping[str, object], + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout: float | httpx.Timeout | None, +) -> TypedEventStreamAdapter | None: + if not supports_streaming(api, provider): + return None + opener: Final = _async_opener(api) + context: Final = _context(api, provider, request) + result: Final = await aattempt( + native_call=( + None + if opener is None + else lambda: opener( + request, + provider, + credentials, + api_base, + extra_headers, + timeout_to_seconds(timeout), + ) + ), + adapt=lambda stream: _adapt_stream(stream, context=context), + context=context, + ) + return result.value if isinstance(result, RustHandled) else None diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 26f841c1146..e571cc4e263 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2680,10 +2680,16 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h (None, GenericLiteLLMParams(rust=True), False), ], ) -def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag( +def test_the_rust_responses_websocket_needs_provider_flag_and_typed_capability( custom_llm_provider, litellm_params, expected ): - assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected + from litellm.rust_bridge import streaming + + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) + try: + assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected + finally: + streaming.set_rust_streaming(capability=None) def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch): diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index 12f75bc32f0..744c5172689 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -1,29 +1,32 @@ from __future__ import annotations +from collections.abc import Mapping + import pytest +from litellm.exceptions import APIError from litellm.llms.custom_httpx.llm_http_handler import _rust_responses_websocket_enabled -from litellm.rust_bridge import bindings, configuration, responses_websocket +from litellm.rust_bridge import bindings, configuration, responses_websocket, streaming from litellm.types.router import GenericLiteLLMParams class _FakeNativeConnection: def __init__(self) -> None: - self.sent: list[str] = [] + self.sent: list[dict[str, object]] = [] self.closed = False - async def send_text(self, text: str) -> None: - self.sent.append(text) + async def send_event(self, event: Mapping[str, object]) -> None: + self.sent.append(dict(event)) - async def recv_text(self) -> str: - return "response.completed" + async def recv_event(self) -> Mapping[str, object]: + return {"type": "response.completed"} async def close(self) -> None: self.closed = True class _ClosedNativeConnection: - async def recv_text(self) -> None: + async def recv_event(self) -> None: return None @@ -31,19 +34,61 @@ class _FakeNativeBridge: @classmethod async def connect( cls, - *, - url: str, - headers: dict[str, str], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, timeout_seconds: float | None, ) -> _FakeNativeConnection: return _FakeNativeConnection() +class _Declined(Exception): + pass + + +class _Upstream(Exception): + pass + + +class _NativeErrors: + RustBridgeDeclined = _Declined + RustUpstreamError = _Upstream + + +class _DecliningNativeBridge: + @classmethod + async def connect( + cls, + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> _FakeNativeConnection: + raise _Declined("provider unsupported") + + +class _FailingNativeBridge: + @classmethod + async def connect( + cls, + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> _FakeNativeConnection: + raise _Upstream(503, "request may have executed") + + @pytest.fixture(autouse=True) def reset_responses_websocket(): + streaming.set_rust_streaming(capability=None) responses_websocket.set_rust_responses_websocket(connection=None) configuration.reset_rust_configuration() yield + streaming.set_rust_streaming(capability=None) responses_websocket.set_rust_responses_websocket(connection=None) configuration.reset_rust_configuration() @@ -51,16 +96,26 @@ def reset_responses_websocket(): def test_rust_websocket_bridge_is_disabled_without_flag() -> None: assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) assert not _rust_responses_websocket_enabled("anthropic", GenericLiteLLMParams(rust=True)) + assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) + + +def test_injected_typed_capability_enables_the_gate() -> None: + streaming.set_rust_streaming( + capability=lambda api, provider, transport: (api, provider, transport) == ("responses", "openai", "websocket") + ) + assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=True)) def test_explicit_false_overrides_process_enable() -> None: + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) configuration.use_litellm_rust(True) assert not _rust_responses_websocket_enabled("openai", GenericLiteLLMParams(rust=False)) def test_process_enable_applies_without_request_override() -> None: + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) configuration.use_litellm_rust(True) assert _rust_responses_websocket_enabled("openai", GenericLiteLLMParams()) @@ -68,7 +123,10 @@ def test_process_enable_applies_without_request_override() -> None: @pytest.mark.asyncio async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: - adapter = responses_websocket._ConnectionAdapter(_ClosedNativeConnection()) + adapter = responses_websocket._ConnectionAdapter( + _ClosedNativeConnection(), + responses_websocket.BridgeErrorContext(route="responses websocket", provider="openai", model=""), + ) with pytest.raises(responses_websocket.ConnectionClosedOK): await adapter.recv() @@ -77,9 +135,12 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: @pytest.mark.asyncio async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) result = await responses_websocket.connect( - url="wss://example.test/responses", + provider="openai", + api_key=None, + api_base="https://example.test", headers={}, timeout=None, ) @@ -91,10 +152,13 @@ async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) async def test_enabled_bridge_connects_and_adapts_socket( monkeypatch: pytest.MonkeyPatch, ) -> None: + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) responses_websocket.set_rust_responses_websocket(connection=_FakeNativeBridge) result = await responses_websocket.connect( - url="wss://example.test/responses", + provider="openai", + api_key="key", + api_base="https://example.test", headers={"Authorization": "Bearer key"}, timeout=1.0, ) @@ -103,6 +167,41 @@ async def test_enabled_bridge_connects_and_adapts_socket( connection = result.value assert connection is not None assert connection.core_engine is responses_websocket.CoreEngine.RUST - await connection.send("response.create") - assert await connection.recv() == "response.completed" + await connection.send('{"type":"response.create","model":"gpt-5"}') + assert await connection.recv() == '{"type":"response.completed"}' await connection.close() + + +@pytest.mark.asyncio +async def test_declined_connect_falls_back_before_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _NativeErrors()) + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) + responses_websocket.set_rust_responses_websocket(connection=_DecliningNativeBridge) + + result = await responses_websocket.connect( + provider="openai", + api_key="key", + api_base="https://example.test", + headers={}, + timeout=None, + ) + + assert result.value is None + assert result.source is responses_websocket.CoreEngine.PYTHON + + +@pytest.mark.asyncio +async def test_upstream_connect_failure_never_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _NativeErrors()) + streaming.set_rust_streaming(capability=lambda api, provider, transport: True) + responses_websocket.set_rust_responses_websocket(connection=_FailingNativeBridge) + + with pytest.raises(APIError, match="request may have executed") as caught: + await responses_websocket.connect( + provider="openai", + api_key="key", + api_base="https://example.test", + headers={}, + timeout=None, + ) + assert caught.value.headers["x-litellm-core"] == "rust" diff --git a/tests/test_litellm/rust_bridge/test_streaming.py b/tests/test_litellm/rust_bridge/test_streaming.py new file mode 100644 index 00000000000..585187fa5e9 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_streaming.py @@ -0,0 +1,327 @@ +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.exceptions import APIError +from litellm.rust_bridge import bindings, streaming +from litellm.rust_bridge.runtime import BridgeErrorContext + + +class _FakeEventStream: + def __init__(self, events: tuple[Mapping[str, object], ...]) -> None: + self.metadata: Final = { + "status_code": 200, + "provider": "anthropic", + "transport": "http", + "response_headers": [{"name": "x-test", "value": "ready"}], + } + self._events: Final = iter(events) + self.closed = False + + def next_event(self) -> Mapping[str, object] | None: + if self.closed: + return None + return next(self._events, None) + + async def anext_event(self) -> Mapping[str, object] | None: + return self.next_event() + + def close(self) -> None: + self.closed = True + + async def aclose(self) -> None: + self.close() + + +class _RecordingOpen: + def __init__(self, events: tuple[Mapping[str, object], ...]) -> None: + self._events: Final = events + self.calls = 0 + + def __call__( + self, + request: Mapping[str, object], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> _FakeEventStream: + self.calls += 1 + return _FakeEventStream(self._events) + + +class _RecordingAsyncOpen: + def __init__(self, events: tuple[Mapping[str, object], ...]) -> None: + self._events: Final = events + self.calls = 0 + + async def __call__( + self, + request: Mapping[str, object], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> _FakeEventStream: + self.calls += 1 + return _FakeEventStream(self._events) + + +class _Declined(Exception): + pass + + +class _Upstream(Exception): + pass + + +class _NativeErrors: + RustBridgeDeclined = _Declined + RustUpstreamError = _Upstream + + +class _FailingOpen: + def __init__(self, error: Exception) -> None: + self._error: Final = error + + def __call__( + self, + request: Mapping[str, object], + provider: str, + credentials: Mapping[str, str] | None, + api_base: str | None, + extra_headers: Mapping[str, str] | None, + timeout_seconds: float | None, + ) -> _FakeEventStream: + raise self._error + + +class _FailingEventStream(_FakeEventStream): + def next_event(self) -> Mapping[str, object] | None: + raise _Upstream(502, "stream interrupted") + + +@pytest.fixture(autouse=True) +def reset_bridge() -> Iterator[None]: + streaming.set_rust_streaming( + capability=None, + chat=None, + achat=None, + messages=None, + amessages=None, + responses=None, + aresponses=None, + ) + yield + streaming.set_rust_streaming( + capability=None, + chat=None, + achat=None, + messages=None, + amessages=None, + responses=None, + aresponses=None, + ) + + +def _chat_event(text: str) -> Mapping[str, object]: + return { + "text": text, + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + } + + +def _context() -> BridgeErrorContext: + return BridgeErrorContext(route="chat completions stream", provider="anthropic", model="claude") + + +def test_no_native_capability_keeps_every_provider_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + opener: Final = _RecordingOpen((_chat_event("unused"),)) + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + streaming.set_rust_streaming(chat=opener) + + result: Final = streaming.open_stream( + api="chat_completions", + provider="anthropic", + request={"model": "claude", "messages": []}, + credentials={"api_key": "test"}, + api_base=None, + extra_headers=None, + timeout=None, + ) + + assert result is None + assert opener.calls == 0 + + +def test_disabled_capability_never_calls_native() -> None: + opener: Final = _RecordingOpen((_chat_event("unused"),)) + streaming.set_rust_streaming(capability=lambda api, provider, transport: False, chat=opener) + + result: Final = streaming.open_stream( + api="chat_completions", + provider="anthropic", + request={"model": "claude", "messages": []}, + credentials=None, + api_base=None, + extra_headers=None, + timeout=None, + ) + + assert result is None + assert opener.calls == 0 + + +def test_declined_open_failure_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _NativeErrors()) + streaming.set_rust_streaming( + capability=lambda api, provider, transport: True, + chat=_FailingOpen(_Declined("unsupported request")), + ) + + result: Final = streaming.open_stream( + api="chat_completions", + provider="anthropic", + request={"model": "claude", "messages": []}, + credentials=None, + api_base=None, + extra_headers=None, + timeout=None, + ) + + assert result is None + + +def test_upstream_open_failure_never_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _NativeErrors()) + streaming.set_rust_streaming( + capability=lambda api, provider, transport: True, + chat=_FailingOpen(_Upstream(503, "connection closed after request")), + ) + + with pytest.raises(APIError, match="connection closed after request") as caught: + streaming.open_stream( + api="chat_completions", + provider="anthropic", + request={"model": "claude", "messages": []}, + credentials=None, + api_base=None, + extra_headers=None, + timeout=None, + ) + assert caught.value.headers["x-litellm-core"] == "rust" + + +def test_midstream_failure_preserves_rust_provenance(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bindings, "get_native_bridge", lambda: _NativeErrors()) + events: Final = streaming.TypedEventStreamAdapter(_FailingEventStream(()), _context()) + + with pytest.raises(APIError, match="stream interrupted") as caught: + next(events) + assert caught.value.headers["x-litellm-core"] == "rust" + + +def test_sync_typed_events_preserve_shape_metadata_and_close() -> None: + opener: Final = _RecordingOpen((_chat_event("one"), _chat_event("two"))) + streaming.set_rust_streaming(capability=lambda api, provider, transport: True, chat=opener) + result: Final = streaming.open_stream( + api="chat_completions", + provider="anthropic", + request={"model": "claude", "messages": []}, + credentials=None, + api_base=None, + extra_headers=None, + timeout=1.0, + ) + + assert result is not None + assert tuple(event["text"] for event in result) == ("one", "two") + assert result.metadata["provider"] == "anthropic" + result.close() + assert opener.calls == 1 + + +def test_chat_events_flow_through_custom_stream_wrapper() -> None: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream + + native: Final = _FakeEventStream((_chat_event("hello"),)) + events: Final = streaming.TypedEventStreamAdapter(native, _context()) + wrapper: Final = CustomStreamWrapper( + completion_stream=events, + model="claude", + logging_obj=MagicMock(), + custom_llm_provider="anthropic", + ) + + chunk: Final = next(wrapper) + assert isinstance(chunk, ModelResponseStream) + assert chunk.choices[0].delta.content == "hello" + + +@pytest.mark.asyncio +async def test_async_typed_events_and_cancellation() -> None: + opener: Final = _RecordingAsyncOpen((_chat_event("one"), _chat_event("two"))) + streaming.set_rust_streaming(capability=lambda api, provider, transport: True, achat=opener) + result: Final = await streaming.aopen_stream( + api="chat_completions", + provider="anthropic", + request={"model": "claude", "messages": []}, + credentials=None, + api_base=None, + extra_headers=None, + timeout=None, + ) + + assert result is not None + assert tuple(event["text"] for event in [event async for event in result]) == ("one", "two") + await result.aclose() + + +def test_messages_events_are_wrapped_in_existing_sse_bytes() -> None: + events: Final = streaming.TypedEventStreamAdapter(_FakeEventStream(({"type": "message_stop"},)), _context()) + messages: Final = streaming.MessagesSseStreamAdapter(events) + + assert tuple(messages) == (b'data: {"type":"message_stop"}\n\n',) + + +def test_responses_events_are_validated_into_existing_sdk_objects() -> None: + from litellm.types.llms.openai import OutputTextDeltaEvent + + native: Final = _FakeEventStream( + ( + { + "type": "response.output_text.delta", + "item_id": "item_1", + "output_index": 0, + "content_index": 0, + "delta": "hello", + }, + ) + ) + responses: Final = streaming.ResponsesSdkEventStreamAdapter(streaming.TypedEventStreamAdapter(native, _context())) + + event: Final = next(iter(responses)) + assert isinstance(event, OutputTextDeltaEvent) + assert event.delta == "hello" + + +@pytest.mark.asyncio +async def test_rejects_mixed_sync_and_async_consumption() -> None: + events: Final = streaming.TypedEventStreamAdapter(_FakeEventStream((_chat_event("one"),)), _context()) + assert tuple(events) == (_chat_event("one"),) + + with pytest.raises(RuntimeError, match="cannot mix"): + async for _ in events: + pass