From 24c907aeb35f6e1e7092b3e55a1565068e797d4e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 25 Jun 2026 11:27:49 -0700 Subject: [PATCH] Add Rust realtime bridge entrypoint --- .github/workflows/test-unit-misc.yml | 1 + litellm-rust/Cargo.lock | 2 + litellm-rust/Cargo.toml | 2 +- litellm-rust/crates/providers/src/realtime.rs | 136 ++++++- litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/lib.rs | 105 ++++- litellm/llms/openai/realtime/handler.py | 15 +- litellm/ocr/rust_bridge.py | 24 +- litellm/realtime_api/main.py | 57 +++ litellm/realtime_api/rust_bridge.py | 175 +++++++++ tests/test_litellm/realtime_api/__init__.py | 1 + .../realtime_api/test_rust_bridge.py | 359 ++++++++++++++++++ uv.lock | 6 +- 13 files changed, 865 insertions(+), 20 deletions(-) create mode 100644 litellm/realtime_api/rust_bridge.py create mode 100644 tests/test_litellm/realtime_api/__init__.py create mode 100644 tests/test_litellm/realtime_api/test_rust_bridge.py diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 2226d519331..8531aa0b757 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -34,6 +34,7 @@ jobs: tests/test_litellm/interactions tests/test_litellm/ocr tests/test_litellm/passthrough + tests/test_litellm/realtime_api tests/test_litellm/sandbox tests/test_litellm/vector_stores tests/test_litellm/test_*.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5f9c3a0b270..554d69f497b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -595,12 +595,14 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "futures-util", "litellm-core", "litellm-providers", "pyo3", "pyo3-async-runtimes", "serde_json", "tokio", + "tokio-tungstenite", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 63c7d6e8b9d..bb7d7079124 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -20,6 +20,6 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "rus serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" thiserror = "2.0" -tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "sync"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } diff --git a/litellm-rust/crates/providers/src/realtime.rs b/litellm-rust/crates/providers/src/realtime.rs index f2946db4137..be4cf58253d 100644 --- a/litellm-rust/crates/providers/src/realtime.rs +++ b/litellm-rust/crates/providers/src/realtime.rs @@ -8,16 +8,22 @@ use std::time::Duration; +use std::sync::Arc; + +use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{SinkExt, StreamExt}; use litellm_core::error::CoreError; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::CoreResult; use tokio_tungstenite::connect_async; +use tokio_tungstenite::connect_async_with_config; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue}; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; use crate::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; @@ -29,6 +35,22 @@ const DEFAULT_TIMEOUT_SECS: u64 = 60; const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; +pub type UpstreamWs = WebSocketStream>; +type UpstreamTx = SplitSink; +type UpstreamRx = SplitStream; + +pub enum RealtimePayload { + Text(String), + Binary(Vec), +} + +fn websocket_config(max_message_size: Option) -> Option { + max_message_size.map(|max_message_size| WebSocketConfig { + max_message_size: Some(max_message_size), + ..WebSocketConfig::default() + }) +} + /// Resolve the OpenAI API key from the explicit param or the environment. /// /// Blank/whitespace values are treated as absent (guard at resolution time). @@ -50,6 +72,107 @@ fn is_terminal_event(event: &RealtimeEvent) -> bool { event.event_type == "response.done" || event.event_type == "error" } +/// Dial an arbitrary realtime WebSocket URL with caller-supplied headers. +/// +/// Python owns URL/header construction for the proxy path so query params, +/// OpenAI-Beta passthrough, and logging stay identical to the existing route. +pub async fn dial_url( + url: &str, + headers: &[(String, String)], + max_message_size: Option, +) -> CoreResult { + let mut request = url + .into_client_request() + .map_err(|err| CoreError::Network(err.to_string()))?; + for (name, value) in headers { + let header_name = HeaderName::from_bytes(name.as_bytes()) + .map_err(|err| CoreError::Auth(format!("invalid header name '{name}': {err}")))?; + let header_value = HeaderValue::from_str(value) + .map_err(|err| CoreError::Auth(format!("invalid header value for '{name}': {err}")))?; + request.headers_mut().insert(header_name, header_value); + } + let config = websocket_config(max_message_size); + let (upstream, _response) = connect_async_with_config(request, config, false) + .await + .map_err(|err| match err { + tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Network(format!( + "upstream websocket HTTP status {}", + response.status().as_u16() + )), + other => CoreError::Network(other.to_string()), + })?; + Ok(upstream) +} + +/// Long-lived handle to an upstream realtime WebSocket. +/// +/// The split halves are protected independently so Python can await sends and +/// receives concurrently from the existing bidirectional forwarding tasks. +pub struct UpstreamHandle { + tx: Arc>, + rx: Arc>, +} + +impl UpstreamHandle { + pub fn new(upstream: UpstreamWs) -> Self { + let (tx, rx) = upstream.split(); + Self { + tx: Arc::new(tokio::sync::Mutex::new(tx)), + rx: Arc::new(tokio::sync::Mutex::new(rx)), + } + } + + pub async fn send_text(&self, text: String) -> CoreResult<()> { + self.send_payload(RealtimePayload::Text(text)).await + } + + pub async fn send_payload(&self, payload: RealtimePayload) -> CoreResult<()> { + let mut tx = self.tx.lock().await; + let message = match payload { + RealtimePayload::Text(text) => Message::Text(text), + RealtimePayload::Binary(bytes) => Message::Binary(bytes), + }; + tx.send(message) + .await + .map_err(|err| CoreError::Network(err.to_string())) + } + + pub async fn recv_text(&self) -> CoreResult> { + loop { + match self.recv_payload().await? { + Some(RealtimePayload::Text(text)) => return Ok(Some(text)), + Some(RealtimePayload::Binary(_)) => continue, + None => return Ok(None), + } + } + } + + pub async fn recv_payload(&self) -> CoreResult> { + let mut rx = self.rx.lock().await; + loop { + let Some(message) = rx.next().await else { + return Ok(None); + }; + match message.map_err(|err| CoreError::Network(err.to_string()))? { + Message::Text(text) => return Ok(Some(RealtimePayload::Text(text))), + Message::Binary(bytes) => return Ok(Some(RealtimePayload::Binary(bytes))), + Message::Close(_) => return Ok(None), + Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => continue, + } + } + } + + pub async fn close(&self) -> CoreResult<()> { + let mut tx = self.tx.lock().await; + match tx.close().await { + Ok(()) => Ok(()), + Err(tokio_tungstenite::tungstenite::Error::ConnectionClosed) + | Err(tokio_tungstenite::tungstenite::Error::AlreadyClosed) => Ok(()), + Err(err) => Err(CoreError::Network(err.to_string())), + } + } +} + /// Invoke the OpenAI realtime API end to end over a WebSocket. /// /// Sends each `input_events` entry after passing it through @@ -155,6 +278,17 @@ mod tests { ))); } + #[test] + fn websocket_config_applies_max_message_size() { + assert!(websocket_config(None).is_none()); + assert_eq!( + websocket_config(Some(1024)) + .expect("config") + .max_message_size, + Some(1024) + ); + } + /// Live end-to-end check against OpenAI. Ignored by default (CI never runs /// it); run explicitly with `OPENAI_API_KEY` set: /// `cargo test -p litellm-providers realtime_invokes_openai -- --ignored --nocapture` diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index d0a7cd56cc6..606ff50d7bc 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -16,3 +16,5 @@ pyo3 = { workspace = true, features = ["extension-module"] } pyo3-async-runtimes.workspace = true serde_json.workspace = true tokio.workspace = true +tokio-tungstenite.workspace = true +futures-util.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index aa1fb2e0127..c44ae0a37d5 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,9 +1,10 @@ +use std::sync::Arc; use std::time::Duration; use litellm_core::error::CoreError; -use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::exceptions::{PyRuntimeError, PyStopAsyncIteration, PyValueError}; use pyo3::prelude::*; -use pyo3::types::{PyAny, PyDict}; +use pyo3::types::{PyAny, PyBytes, PyDict}; use serde_json::{Map, Value}; mod gil; @@ -181,10 +182,110 @@ fn gil_stats(py: Python<'_>) -> PyResult> { Ok(stats.into_any().unbind()) } +fn core_error_to_runtime(err: CoreError) -> PyErr { + PyRuntimeError::new_err(err.to_string()) +} + +fn realtime_payload_to_py( + py: Python<'_>, + payload: litellm_providers::realtime::RealtimePayload, +) -> PyResult> { + match payload { + litellm_providers::realtime::RealtimePayload::Text(text) => { + Ok(text.into_pyobject(py)?.into_any().unbind()) + } + litellm_providers::realtime::RealtimePayload::Binary(bytes) => { + Ok(PyBytes::new(py, &bytes).into_any().unbind()) + } + } +} + +#[pyclass] +struct RustRealtimeUpstream { + inner: Arc, +} + +#[pymethods] +impl RustRealtimeUpstream { + fn send<'py>( + &self, + py: Python<'py>, + message: Bound<'py, PyAny>, + ) -> PyResult> { + let payload = if let Ok(text) = message.extract::() { + litellm_providers::realtime::RealtimePayload::Text(text) + } else if let Ok(bytes) = message.extract::>() { + litellm_providers::realtime::RealtimePayload::Binary(bytes) + } else { + return Err(PyValueError::new_err( + "realtime message must be str or bytes", + )); + }; + let inner = self.inner.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + inner + .send_payload(payload) + .await + .map_err(core_error_to_runtime) + }) + } + + fn recv<'py>(&self, py: Python<'py>) -> PyResult> { + let inner = self.inner.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + match inner.recv_payload().await { + Ok(Some(payload)) => Python::with_gil(|py| realtime_payload_to_py(py, payload)), + Ok(None) => Err(PyStopAsyncIteration::new_err("upstream closed")), + Err(err) => Err(core_error_to_runtime(err)), + } + }) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + let inner = self.inner.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + inner.close().await.map_err(core_error_to_runtime) + }) + } + + fn __aiter__(slf: Py) -> Py { + slf + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + self.recv(py) + } +} + +#[pyfunction] +#[pyo3(signature = (url, headers, max_size=None))] +fn realtime_connect<'py>( + py: Python<'py>, + url: String, + headers: Bound<'py, PyDict>, + max_size: Option, +) -> PyResult> { + let headers_vec: Vec<(String, String)> = headers + .iter() + .map(|(key, value)| Ok((key.extract::()?, value.extract::()?))) + .collect::>()?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let upstream = litellm_providers::realtime::dial_url(&url, &headers_vec, max_size) + .await + .map_err(core_error_to_runtime)?; + Ok(RustRealtimeUpstream { + inner: Arc::new(litellm_providers::realtime::UpstreamHandle::new(upstream)), + }) + }) +} + #[pymodule] fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; module.add_function(wrap_pyfunction!(aocr, module)?)?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + module.add_function(wrap_pyfunction!(realtime_connect, module)?)?; + module.add_class::()?; Ok(()) } diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 6751004f1b1..e7c3477be32 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -4,7 +4,7 @@ This file contains the calling OpenAI's `/v1/realtime` endpoint. This requires websockets, and is currently only supported on LiteLLM Proxy. """ -from typing import Any, Optional, cast +from typing import Any, AsyncContextManager, Callable, Optional, cast from litellm._logging import _redact_string, verbose_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -18,6 +18,8 @@ from ....litellm_core_utils.realtime_streaming import ( from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..openai import OpenAIChatCompletion +BackendConnect = Callable[..., AsyncContextManager[Any]] + class OpenAIRealtime(OpenAIChatCompletion): """ @@ -107,6 +109,7 @@ class OpenAIRealtime(OpenAIChatCompletion): query_params: Optional[RealtimeQueryParams] = None, user_api_key_dict: Optional[Any] = None, litellm_metadata: Optional[dict] = None, + backend_connect: Optional[BackendConnect] = None, **kwargs: Any, ): import websockets @@ -144,14 +147,18 @@ class OpenAIRealtime(OpenAIChatCompletion): additional_args={ "api_base": url, "headers": headers, - "complete_input_dict": {"query_params": query_params}, + "complete_input_dict": { + "query_params": query_params, + "rust_bridge": backend_connect is not None, + }, }, ) - async with websockets.connect( # type: ignore + connector: BackendConnect = backend_connect or websockets.connect # type: ignore[assignment] + async with connector( url, additional_headers=headers, # type: ignore max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, - ssl=ssl_config, + ssl=None if backend_connect is not None else ssl_config, ) as backend_ws: realtime_streaming = RealTimeStreaming( websocket, diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 1e3312c1473..4a723f36d49 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -1,18 +1,19 @@ """ Optional Rust-backed OCR path. -Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint -then routes supported Mistral calls through the compiled ``litellm_python_bridge`` -extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. - -No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` -can import it statically without forming an import cycle. +Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` +entrypoint then routes supported Mistral calls through the compiled +``litellm_python_bridge`` extension, which performs the whole OCR call +(URL, headers, HTTP, parse) in Rust. The same public helper also toggles the +optional Rust realtime bridge without importing top-level ``litellm`` here. """ from __future__ import annotations from typing import Awaitable, Final, Protocol, cast +from litellm.realtime_api.rust_bridge import RustRealtimeConnect, set_rust_realtime + class RustOcr(Protocol): """Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint.""" @@ -64,11 +65,12 @@ def use_litellm_rust( *, ocr: RustOcr | None | _Unset = _UNSET, aocr: RustAocr | None | _Unset = _UNSET, + realtime: RustRealtimeConnect | None | _Unset = _UNSET, ) -> None: - """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension. + """Route supported calls through the Rust ``litellm_python_bridge`` extension. - ``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled - extension is loaded on demand and any previously injected bridge is + ``ocr``, ``aocr``, and ``realtime`` inject bridge callables; when omitted the + compiled extension is loaded on demand and any previous injection is preserved. Pass ``None`` explicitly to clear a prior injection. """ global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl @@ -77,6 +79,10 @@ def use_litellm_rust( _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr + if isinstance(realtime, _Unset): + set_rust_realtime(enabled) + else: + set_rust_realtime(enabled, connect=realtime) def rust_ocr_enabled() -> bool: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index f6f0a92def0..edbcc4e218b 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -1,6 +1,7 @@ """Abstraction function for OpenAI's realtime API""" import os +from ssl import SSLContext from typing import Any, Dict, Optional, cast import litellm @@ -31,6 +32,13 @@ from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig from ..llms.vertex_ai.vertex_llm_base import VertexBase from ..llms.xai.realtime.handler import XAIRealtime from ..utils import client as wrapper_client +from .rust_bridge import ( + RustBackendConnectFactory, + load_rust_realtime, + rust_backend_connect_factory, + rust_realtime_enabled, + rust_supports_ssl_config, +) azure_realtime = AzureOpenAIRealtime() openai_realtime = OpenAIRealtime() @@ -40,6 +48,52 @@ vertex_llm_base = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() +def _is_default_ssl_verify(value: object) -> bool: + if value is None or value is True: + return True + if isinstance(value, str): + return value.lower() == "true" + return False + + +def _custom_realtime_tls_configured() -> bool: + ssl_verify = os.getenv("SSL_VERIFY", litellm.ssl_verify) + return ( + not _is_default_ssl_verify(ssl_verify) + or os.getenv("SSL_CERT_FILE") is not None + or os.getenv("SSL_CERTIFICATE") is not None + or litellm.ssl_certificate is not None + or os.getenv("SSL_SECURITY_LEVEL", litellm.ssl_security_level) is not None + or os.getenv("SSL_ECDH_CURVE", litellm.ssl_ecdh_curve) is not None + ) + + +def _maybe_rust_backend_connect() -> Optional[RustBackendConnectFactory]: + """Return a Rust-backed ``websockets.connect``-compatible factory.""" + if not rust_realtime_enabled(): + return None + connect = load_rust_realtime() + if connect is None: + from litellm._logging import verbose_logger + + verbose_logger.debug( + "Rust realtime bridge unavailable; falling back to Python path" + ) + return None + ssl_config = get_shared_realtime_ssl_context() + if _custom_realtime_tls_configured() or not ( + rust_supports_ssl_config(ssl_config) or isinstance(ssl_config, SSLContext) + ): + from litellm._logging import verbose_logger + + verbose_logger.debug( + "Rust realtime path cannot honor custom TLS configuration; " + "falling back to Python path" + ) + return None + return rust_backend_connect_factory(connect) + + def _build_litellm_metadata(kwargs: dict) -> dict: """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" metadata: dict = {**(kwargs.get("litellm_metadata") or {})} @@ -415,6 +469,8 @@ async def _arealtime( or get_secret_str("OPENAI_API_KEY") ) + backend_connect = _maybe_rust_backend_connect() + await openai_realtime.async_realtime( model=model, websocket=websocket, @@ -426,6 +482,7 @@ async def _arealtime( query_params=query_params, user_api_key_dict=kwargs.get("user_api_key_dict"), litellm_metadata=_build_litellm_metadata(kwargs), + backend_connect=backend_connect, ) elif _custom_llm_provider == "bedrock": # Extract AWS parameters from kwargs diff --git a/litellm/realtime_api/rust_bridge.py b/litellm/realtime_api/rust_bridge.py new file mode 100644 index 00000000000..9f7e7f65662 --- /dev/null +++ b/litellm/realtime_api/rust_bridge.py @@ -0,0 +1,175 @@ +""" +Optional Rust-backed realtime WebSocket upstream. + +Enable with ``litellm.use_litellm_rust(realtime=...)`` or +``litellm.use_litellm_rust()``. The proxy realtime route keeps URL construction, +OpenAI-Beta forwarding, pre-call logging, guardrails, and spend logging in +Python, while the upstream WebSocket transport is supplied by the compiled +``litellm_python_bridge`` extension. +""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from ssl import SSLContext +from types import TracebackType +from typing import ( + AsyncContextManager, + AsyncIterator, + Awaitable, + Callable, + Final, + Mapping, + Optional, + Protocol, + cast, + runtime_checkable, +) + + +@runtime_checkable +class RustRealtimeUpstream(Protocol): + def send(self, message: str | bytes) -> Awaitable[None]: ... + def recv(self) -> Awaitable[str | bytes]: ... + def close(self) -> Awaitable[None]: ... + + +class RustRealtimeConnect(Protocol): + def __call__( + self, url: str, headers: dict[str, str], max_size: Optional[int] = None + ) -> Awaitable[RustRealtimeUpstream]: ... + + +class _Unset: + """Sentinel so ``realtime=None`` clears while omission preserves.""" + + +_UNSET: Final[_Unset] = _Unset() + +_rust_realtime_enabled = False +_rust_realtime_impl: RustRealtimeConnect | None = None + + +def set_rust_realtime( + enabled: bool, *, connect: RustRealtimeConnect | None | _Unset = _UNSET +) -> None: + global _rust_realtime_enabled, _rust_realtime_impl + _rust_realtime_enabled = enabled + if not isinstance(connect, _Unset): + _rust_realtime_impl = connect + + +def rust_realtime_enabled() -> bool: + return _rust_realtime_enabled + + +def load_rust_realtime() -> RustRealtimeConnect | None: + if _rust_realtime_impl is not None: + return _rust_realtime_impl + try: + import litellm_python_bridge + except ImportError: + return None + return cast( + RustRealtimeConnect, getattr(litellm_python_bridge, "realtime_connect", None) + ) + + +class RustBackendWebsocket: + """Subset of ``websockets.ClientConnection`` used by ``RealTimeStreaming``.""" + + def __init__(self, upstream: RustRealtimeUpstream) -> None: + self._upstream = upstream + self._closed = False + + async def send(self, message: str | bytes) -> None: + await self._upstream.send(message) + + async def recv(self, *_args: object, **_kwargs: object) -> str | bytes: + try: + return await self._upstream.recv() + except StopAsyncIteration: + import websockets + + raise websockets.exceptions.ConnectionClosedOK(None, None) from None + + def __aiter__(self) -> RustBackendWebsocket: + return self + + async def __anext__(self) -> str | bytes: + return await self._upstream.recv() + + async def close(self, *_args: object, **_kwargs: object) -> None: + if self._closed: + return + self._closed = True + await self._upstream.close() + + async def __aenter__(self) -> RustBackendWebsocket: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + _ = (exc_type, exc, traceback) + await self.close() + + +RustBackendConnectFactory = Callable[..., AsyncContextManager[RustBackendWebsocket]] + + +def rust_backend_connect_factory( + connect: RustRealtimeConnect, +) -> RustBackendConnectFactory: + @asynccontextmanager + async def _connect( + url: str, + *, + additional_headers: Mapping[str, str], + max_size: Optional[int] = None, + ssl: bool | SSLContext | str | None = None, + ) -> AsyncIterator[RustBackendWebsocket]: + if not rust_supports_ssl_config(ssl): + raise NotImplementedError( + "Rust realtime path does not support custom TLS configuration" + ) + try: + upstream = await connect(url, dict(additional_headers), max_size) + except RuntimeError as exc: + status_code = _rust_realtime_http_status_code(str(exc)) + if status_code is not None: + import websockets.exceptions as websocket_exceptions + from websockets.datastructures import Headers + + raise websocket_exceptions.InvalidStatusCode( + status_code, Headers() + ) from exc + raise + backend = RustBackendWebsocket(upstream) + try: + yield backend + finally: + await backend.close() + + return _connect + + +def rust_supports_ssl_config(ssl_config: bool | SSLContext | str | None) -> bool: + return ssl_config is None or ssl_config is True + + +def _rust_realtime_http_status_code(message: str) -> int | None: + prefix = "upstream websocket HTTP status " + start = message.find(prefix) + if start == -1: + return None + start += len(prefix) + end = start + while end < len(message) and message[end].isdigit(): + end += 1 + if end == start: + return None + return int(message[start:end]) diff --git a/tests/test_litellm/realtime_api/__init__.py b/tests/test_litellm/realtime_api/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm/realtime_api/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm/realtime_api/test_rust_bridge.py b/tests/test_litellm/realtime_api/test_rust_bridge.py new file mode 100644 index 00000000000..bed11e318aa --- /dev/null +++ b/tests/test_litellm/realtime_api/test_rust_bridge.py @@ -0,0 +1,359 @@ +"""Tests for the optional Rust-backed realtime path.""" + +from __future__ import annotations + +import importlib +import asyncio +import ssl +import sys +import types +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm + +realtime_main = importlib.import_module("litellm.realtime_api.main") +rust_bridge = importlib.import_module("litellm.realtime_api.rust_bridge") + + +class RecordingUpstream: + def __init__(self, *, recv_frames: list[str | bytes] | None = None) -> None: + self.sent: list[str | bytes] = [] + self.closed = False + self._recv_frames = list(recv_frames or []) + + def send(self, message: str | bytes): + async def _send() -> None: + self.sent.append(message) + + return _send() + + def recv(self): + async def _recv() -> str | bytes: + if not self._recv_frames: + raise StopAsyncIteration("upstream closed") + return self._recv_frames.pop(0) + + return _recv() + + def close(self): + async def _close() -> None: + self.closed = True + + return _close() + + +class RecordingConnect: + def __init__(self, *, recv_frames: list[str | bytes] | None = None) -> None: + self.calls: list[dict[str, Any]] = [] + self.upstream = RecordingUpstream(recv_frames=recv_frames) + + def __call__(self, url: str, headers: dict[str, str], max_size: int | None = None): + self.calls.append({"url": url, "headers": headers, "max_size": max_size}) + + async def _connect() -> RecordingUpstream: + return self.upstream + + return _connect() + + +@pytest.fixture(autouse=True) +def _reset_rust_flag(): + ssl_verify = litellm.ssl_verify + ssl_certificate = litellm.ssl_certificate + ssl_security_level = litellm.ssl_security_level + ssl_ecdh_curve = litellm.ssl_ecdh_curve + rust_bridge.set_rust_realtime(False, connect=None) + yield + litellm.ssl_verify = ssl_verify + litellm.ssl_certificate = ssl_certificate + litellm.ssl_security_level = ssl_security_level + litellm.ssl_ecdh_curve = ssl_ecdh_curve + rust_bridge.set_rust_realtime(False, connect=None) + + +def test_use_litellm_rust_toggles_realtime_flag(): + assert rust_bridge.rust_realtime_enabled() is False + litellm.use_litellm_rust() + assert rust_bridge.rust_realtime_enabled() is True + litellm.use_litellm_rust(False) + assert rust_bridge.rust_realtime_enabled() is False + + +def test_load_rust_realtime_returns_injected_impl(): + bridge = RecordingConnect() + litellm.use_litellm_rust(True, realtime=bridge) + assert rust_bridge.load_rust_realtime() is bridge + + +def test_toggle_without_realtime_arg_preserves_injected_impl(): + bridge = RecordingConnect() + litellm.use_litellm_rust(True, realtime=bridge) + + litellm.use_litellm_rust(False) + assert rust_bridge.load_rust_realtime() is bridge + litellm.use_litellm_rust(True) + assert rust_bridge.load_rust_realtime() is bridge + + +def test_explicit_realtime_none_clears_injected_impl(monkeypatch): + monkeypatch.delitem(sys.modules, "litellm_python_bridge", raising=False) + bridge = RecordingConnect() + litellm.use_litellm_rust(True, realtime=bridge) + + litellm.use_litellm_rust(True, realtime=None) + + assert rust_bridge.load_rust_realtime() is None + + +def test_load_rust_realtime_none_when_extension_absent(monkeypatch): + monkeypatch.delitem(sys.modules, "litellm_python_bridge", raising=False) + litellm.use_litellm_rust(True) + assert rust_bridge.load_rust_realtime() is None + + +def test_load_rust_realtime_uses_compiled_extension(monkeypatch): + fake_module = types.ModuleType("litellm_python_bridge") + fake_module.realtime_connect = lambda url, headers, max_size=None: None # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) + + litellm.use_litellm_rust(True) + assert rust_bridge.load_rust_realtime() is fake_module.realtime_connect + + +@pytest.mark.asyncio +async def test_rust_backend_connect_factory_matches_websockets_shape(): + bridge = RecordingConnect() + factory = rust_bridge.rust_backend_connect_factory(bridge) + + async with factory( + "wss://api.openai.com/v1/realtime?model=gpt-realtime", + additional_headers={"Authorization": "Bearer sk-test"}, + max_size=1234, + ssl=None, + ) as backend_ws: + assert isinstance(backend_ws, rust_bridge.RustBackendWebsocket) + + assert bridge.calls == [ + { + "url": "wss://api.openai.com/v1/realtime?model=gpt-realtime", + "headers": {"Authorization": "Bearer sk-test"}, + "max_size": 1234, + } + ] + assert bridge.upstream.closed is True + + +@pytest.mark.asyncio +async def test_rust_backend_connect_factory_maps_http_status(): + import websockets.exceptions as websocket_exceptions + + def connect(url: str, headers: dict[str, str], max_size: int | None = None): + async def _connect() -> RecordingUpstream: + _ = (url, headers, max_size) + raise RuntimeError("upstream websocket HTTP status 401") + + return _connect() + + factory = rust_bridge.rust_backend_connect_factory(connect) + + with pytest.raises(websocket_exceptions.InvalidStatusCode) as exc: + async with factory("wss://example.test", additional_headers={}, max_size=1): + pass + assert exc.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_rust_backend_connect_factory_rejects_custom_ssl_context(): + bridge = RecordingConnect() + factory = rust_bridge.rust_backend_connect_factory(bridge) + + with pytest.raises(NotImplementedError): + async with factory( + "wss://example.test", + additional_headers={}, + ssl=ssl.create_default_context(), + ): + pass + + +@pytest.mark.asyncio +async def test_rust_backend_websocket_recv_raises_connection_closed_ok(): + import websockets + + backend = rust_bridge.RustBackendWebsocket(RecordingUpstream()) + + with pytest.raises(websockets.exceptions.ConnectionClosed): + await backend.recv() + + +@pytest.mark.asyncio +async def test_rust_backend_websocket_preserves_binary_frames(): + upstream = RecordingUpstream(recv_frames=[b"abc"]) + backend = rust_bridge.RustBackendWebsocket(upstream) + + await backend.send(b"client-binary") + + assert await backend.recv() == b"abc" + assert upstream.sent == [b"client-binary"] + + +def test_maybe_rust_backend_connect_returns_none_when_disabled(): + assert realtime_main._maybe_rust_backend_connect() is None + + +def test_maybe_rust_backend_connect_returns_factory_when_enabled(): + bridge = RecordingConnect() + litellm.use_litellm_rust(True, realtime=bridge) + factory = realtime_main._maybe_rust_backend_connect() + assert factory is not None and callable(factory) + + +def test_maybe_rust_backend_connect_falls_back_when_bridge_missing(monkeypatch): + monkeypatch.delitem(sys.modules, "litellm_python_bridge", raising=False) + litellm.use_litellm_rust(True) + assert realtime_main._maybe_rust_backend_connect() is None + + +def test_maybe_rust_backend_connect_falls_back_when_verify_off(monkeypatch): + bridge = RecordingConnect() + litellm.use_litellm_rust(True, realtime=bridge) + monkeypatch.setattr( + realtime_main, + "get_shared_realtime_ssl_context", + lambda: False, + ) + assert realtime_main._maybe_rust_backend_connect() is None + + +def test_maybe_rust_backend_connect_falls_back_for_custom_tls(monkeypatch): + bridge = RecordingConnect() + litellm.use_litellm_rust(True, realtime=bridge) + monkeypatch.setattr(litellm, "ssl_verify", ssl.create_default_context()) + monkeypatch.setattr( + realtime_main, + "get_shared_realtime_ssl_context", + ssl.create_default_context, + ) + assert realtime_main._maybe_rust_backend_connect() is None + + monkeypatch.setattr(litellm, "ssl_verify", "/tmp/custom-ca.pem") + monkeypatch.setattr( + realtime_main, + "get_shared_realtime_ssl_context", + lambda: "/tmp/custom-ca.pem", + ) + assert realtime_main._maybe_rust_backend_connect() is None + + +@pytest.mark.asyncio +async def test_openai_handler_passes_backend_connect_to_dial(monkeypatch): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.llms.openai.realtime.handler import OpenAIRealtime + + handler = OpenAIRealtime() + captured: dict[str, Any] = {} + + class FakeBackend: + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return None + + def fake_connect(url: str, *, additional_headers, max_size, ssl): + captured.update( + { + "url": url, + "additional_headers": dict(additional_headers), + "max_size": max_size, + "ssl": ssl, + } + ) + return FakeBackend() + + class FakeClientWebsocket: + scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + + async def close(self, *args, **kwargs): + return None + + async def fake_forward(self): + captured["streaming_started"] = True + + monkeypatch.setattr( + "litellm.litellm_core_utils.realtime_streaming.RealTimeStreaming.bidirectional_forward", + fake_forward, + ) + + logging_obj = LiteLLMLogging( + model="gpt-realtime", + messages=[], + stream=False, + call_type="realtime", + start_time=datetime.now(), + litellm_call_id="test-call", + function_id="fn", + ) + pre_call_payload: dict[str, Any] = {} + original_pre_call = logging_obj.pre_call + + def capture_pre_call(*args: Any, **kwargs: Any): + pre_call_payload.update(kwargs) + return original_pre_call(*args, **kwargs) + + monkeypatch.setattr(logging_obj, "pre_call", capture_pre_call) + + await handler.async_realtime( + model="gpt-realtime", + websocket=FakeClientWebsocket(), + logging_obj=logging_obj, + api_base="https://api.openai.com/", + api_key="sk-test", + backend_connect=fake_connect, + ) + + assert captured["url"].startswith( + "wss://api.openai.com/v1/realtime?model=gpt-realtime" + ) + assert captured["additional_headers"]["Authorization"] == "Bearer sk-test" + assert captured["additional_headers"]["OpenAI-Beta"] == "realtime=v1" + assert captured["ssl"] is None + assert captured["streaming_started"] is True + complete_input = pre_call_payload["additional_args"]["complete_input_dict"] + assert complete_input["rust_bridge"] is True + + +@pytest.mark.asyncio +async def test_realtime_logging_flushes_through_rust_backend_adapter(): + from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws = rust_bridge.RustBackendWebsocket( + RecordingUpstream( + recv_frames=[ + '{"type":"session.created","session":{"id":"sess_1"}}', + '{"type":"response.done","response":{"output":[]}}', + ] + ) + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + await streaming.backend_to_client_send_messages() + await asyncio.sleep(0) + + assert client_ws.send_text.await_count == 2 + logging_obj.async_success_handler.assert_awaited_once() + logged_events = logging_obj.async_success_handler.call_args.args[0] + assert [event["type"] for event in logged_events] == [ + "session.created", + "response.done", + ] diff --git a/uv.lock b/uv.lock index cac1696bf34..0d10d0d7b70 100644 --- a/uv.lock +++ b/uv.lock @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.0" +version = "4.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] [[package]]