mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(python): adapt typed Rust streams
This commit is contained in:
Yujong Lee
GitHub
parent
22c50c0759
commit
6b49cb963f
6 changed files with 851 additions and 40 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
354
litellm/rust_bridge/streaming.py
Normal file
354
litellm/rust_bridge/streaming.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
327
tests/test_litellm/rust_bridge/test_streaming.py
Normal file
327
tests/test_litellm/rust_bridge/test_streaming.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue