From 22c50c07591ac9283857b480dceb5935a0ef00d4 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 07:13:58 -0700 Subject: [PATCH] feat(python-bridge): expose typed native streams --- .../crates/python-bridge/src/errors.rs | 15 + litellm-rust/crates/python-bridge/src/lib.rs | 10 + .../crates/python-bridge/src/routes/mod.rs | 5 +- .../python-bridge/src/routes/receiver.rs | 206 +++++++++ .../python-bridge/src/routes/runtime.rs | 100 +++++ .../python-bridge/src/routes/streaming.rs | 422 ++++++++++++++++++ .../rust_bridge/native_route_wheel_test.py | 90 ++++ 7 files changed, 847 insertions(+), 1 deletion(-) create mode 100644 litellm-rust/crates/python-bridge/src/routes/receiver.rs create mode 100644 litellm-rust/crates/python-bridge/src/routes/streaming.rs diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index d3e47937666..3e9bf9a3bbf 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -9,6 +9,10 @@ pyo3::create_exception!( "The route declined before calling the provider, so the host may retry on its own path." ); +pub(crate) fn declined(error: impl std::fmt::Display) -> PyErr { + RustBridgeDeclined::new_err(error.to_string()) +} + pyo3::create_exception!( _native, RustUpstreamError, @@ -34,6 +38,12 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { pub(crate) fn fallback_route_error_to_pyerr(err: Error) -> PyErr { match err { Error::Unsupported(_) => RustBridgeDeclined::new_err(err.to_string()), + other => executed_route_error_to_pyerr(other), + } +} + +pub(crate) fn executed_route_error_to_pyerr(err: Error) -> PyErr { + match err { Error::Http { status, body } => { RustUpstreamError::new_err((status, format!("{status}: {body}"))) } @@ -93,6 +103,11 @@ mod tests { .expect("upstream error should carry status and message"); assert_eq!(args, (expected.0, expected.1.to_string())); } + + let midstream = executed_route_error_to_pyerr(Error::Unsupported( + "a stream cannot fall back after opening", + )); + assert!(midstream.is_instance_of::(py)); }); } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 47df6b905a8..65ce75a7768 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -102,6 +102,16 @@ mod tests { "chat_completions_decline", "chat_completions", "achat_completions", + "ChatCompletionsEventStream", + "MessagesEventStream", + "ResponsesEventStream", + "ResponsesWebSocketSession", + "chat_completions_stream", + "achat_completions_stream", + "messages_stream", + "amessages_stream", + "responses_stream", + "aresponses_stream", "ResponsesWebSocketConnection", "gil_stats", ]; diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 3590ea3a769..1e64ae23248 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -8,10 +8,13 @@ mod audio_transcription; mod chat_completions; mod messages; mod ocr; +mod receiver; +mod streaming; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { ocr::register(module)?; audio_transcription::register(module)?; messages::register(module)?; - chat_completions::register(module) + chat_completions::register(module)?; + streaming::register(module) } diff --git a/litellm-rust/crates/python-bridge/src/routes/receiver.rs b/litellm-rust/crates/python-bridge/src/routes/receiver.rs new file mode 100644 index 00000000000..a97a256cfa2 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/receiver.rs @@ -0,0 +1,206 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use futures_util::{Stream, StreamExt, pin_mut}; +use litellm_core::Error; +use tokio::sync::{Mutex, mpsc}; +use tokio::task::JoinHandle; + +const BRIDGE_CHANNEL_CAPACITY: usize = 1; + +struct ReceiverState { + receiver: Mutex>>, + reading: AtomicBool, + closed: AtomicBool, + producer: std::sync::Mutex>>, +} + +impl Drop for ReceiverState { + fn drop(&mut self) { + if let Ok(producer) = self.producer.get_mut() + && let Some(producer) = producer.take() + { + producer.abort(); + } + } +} + +#[derive(Clone)] +pub(super) struct BridgeReceiver { + state: Arc>, +} + +struct ReadGuard<'a>(&'a AtomicBool); + +impl Drop for ReadGuard<'_> { + fn drop(&mut self) { + self.0.store(false, Ordering::Release); + } +} + +impl BridgeReceiver { + pub(super) fn from_stream(stream: S) -> Self + where + S: Stream> + Send + 'static, + { + let (sender, receiver) = mpsc::channel(BRIDGE_CHANNEL_CAPACITY); + let producer = pyo3_async_runtimes::tokio::get_runtime().spawn(async move { + pin_mut!(stream); + while let Some(item) = stream.next().await { + let terminal = item.is_err(); + if sender.send(item).await.is_err() || terminal { + return; + } + } + }); + Self { + state: Arc::new(ReceiverState { + receiver: Mutex::new(receiver), + reading: AtomicBool::new(false), + closed: AtomicBool::new(false), + producer: std::sync::Mutex::new(Some(producer)), + }), + } + } + + pub(super) async fn next(&self) -> Result, Error> { + if self.state.closed.load(Ordering::Acquire) { + return Ok(None); + } + if self.state.reading.swap(true, Ordering::AcqRel) { + return Err(Error::InvalidRequest( + "native stream does not support concurrent reads".to_string(), + )); + } + let _guard = ReadGuard(&self.state.reading); + match self.state.receiver.lock().await.recv().await { + Some(Ok(item)) => Ok(Some(item)), + Some(Err(error)) => Err(error), + None => Ok(None), + } + } + + pub(super) fn close(&self) { + if self.state.closed.swap(true, Ordering::AcqRel) { + return; + } + if let Ok(mut producer) = self.state.producer.lock() + && let Some(producer) = producer.take() + { + producer.abort(); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use futures_util::stream; + + use super::*; + + struct DropFlag(Arc); + + impl Drop for DropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn receiver_preserves_items_and_terminal_error() { + let receiver = BridgeReceiver::from_stream(stream::iter([ + Ok(vec![1_u8]), + Err(Error::Network("broken".to_string())), + Ok(vec![2_u8]), + ])); + + assert_eq!(receiver.next().await.expect("first item"), Some(vec![1])); + assert!(matches!( + receiver.next().await, + Err(Error::Network(message)) if message == "broken" + )); + assert_eq!(receiver.next().await.expect("closed after error"), None); + } + + #[tokio::test] + async fn close_unblocks_a_pending_read() { + let receiver = BridgeReceiver::>::from_stream(stream::pending()); + let pending = { + let receiver = receiver.clone(); + tokio::spawn(async move { receiver.next().await }) + }; + tokio::time::sleep(Duration::from_millis(10)).await; + receiver.close(); + + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), pending) + .await + .expect("read should unblock") + .expect("task should finish") + .expect("close is clean"), + None + ); + } + + #[tokio::test] + async fn capacity_one_stops_the_producer_from_draining_the_source() { + let polls = Arc::new(AtomicUsize::new(0)); + let source_polls = polls.clone(); + let source = stream::poll_fn(move |_| { + let item = source_polls.fetch_add(1, Ordering::SeqCst); + std::task::Poll::Ready(Some(Ok(item))) + }); + let receiver = BridgeReceiver::from_stream(source); + + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!(polls.load(Ordering::SeqCst), 2); + assert_eq!(receiver.next().await.expect("first item"), Some(0)); + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!(polls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn concurrent_reads_are_rejected() { + let receiver = BridgeReceiver::>::from_stream(stream::pending()); + let pending = { + let receiver = receiver.clone(); + tokio::spawn(async move { receiver.next().await }) + }; + tokio::time::sleep(Duration::from_millis(10)).await; + + assert!(matches!( + receiver.next().await, + Err(Error::InvalidRequest(message)) + if message == "native stream does not support concurrent reads" + )); + receiver.close(); + pending + .await + .expect("pending read task") + .expect("clean close"); + } + + #[tokio::test] + async fn dropping_the_last_receiver_cancels_the_producer() { + let dropped = Arc::new(AtomicBool::new(false)); + let producer_dropped = dropped.clone(); + let source = stream::once(async move { + let _flag = DropFlag(producer_dropped); + std::future::pending::, Error>>().await + }); + let receiver = BridgeReceiver::from_stream(source); + tokio::time::sleep(Duration::from_millis(10)).await; + + drop(receiver); + tokio::time::timeout(Duration::from_secs(1), async { + while !dropped.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("producer future should be dropped"); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs index 87a0c3e0104..75cce13bd36 100644 --- a/litellm-rust/crates/python-bridge/src/routes/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/routes/runtime.rs @@ -65,6 +65,71 @@ where }) } +struct AttachedConversion { + value: T, + convert: C, +} + +impl<'py, T, C> IntoPyObject<'py> for AttachedConversion +where + C: FnOnce(Python<'py>, T) -> PyResult>, +{ + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + std::panic::catch_unwind(AssertUnwindSafe(|| (self.convert)(py, self.value))) + .map_err(panic_to_pyerr)? + .map(|value| value.into_bound(py)) + } +} + +pub(super) fn run_sync_with( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, + convert: C, +) -> PyResult> +where + T: Send + 'static, + F: Future> + Send + 'static, + C: FnOnce(Python<'_>, T) -> PyResult>, +{ + if Handle::try_current().is_ok() { + return Err(PyRuntimeError::new_err( + "synchronous native routes cannot run from a Tokio context; use the async route", + )); + } + + let result = release_gil(py, move || { + pyo3_async_runtimes::tokio::get_runtime().block_on(wait_for_sync_result(future)) + })?; + let result = map_core_result(result, map_error)?; + std::panic::catch_unwind(AssertUnwindSafe(|| convert(py, result))).map_err(panic_to_pyerr)? +} + +pub(super) fn run_async_with( + py: Python<'_>, + future: F, + map_error: fn(Error) -> PyErr, + convert: C, +) -> PyResult> +where + T: Send + 'static, + F: Future> + Send + 'static, + C: for<'py> FnOnce(Python<'py>, T) -> PyResult> + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let result = catch_route_panic(future).await?; + let result = map_core_result(result, map_error)?; + Ok(AttachedConversion { + value: result, + convert, + }) + }) +} + fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { match result { Ok(value) => Ok(value), @@ -159,6 +224,13 @@ mod tests { ) } + #[pyfunction] + fn async_custom_mapping(py: Python<'_>) -> PyResult> { + run_async_with(py, async { Ok("mapped") }, runtime_error, |py, value| { + Ok(value.into_pyobject(py)?.unbind().into_any()) + }) + } + #[pyfunction] fn runtime_worker_count() -> usize { pyo3_async_runtimes::tokio::get_runtime() @@ -420,4 +492,32 @@ asyncio.run(exercise()) .expect("result delivery should leave Tokio workers responsive"); }); } + + #[test] + fn async_custom_mapping_runs_during_attached_result_delivery() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "runtime").expect("module should be created"); + module + .add_function(wrap_pyfunction!(async_custom_mapping, &module).expect("function")) + .expect("function should register"); + let locals = PyDict::new(py); + locals + .set_item("runtime", &module) + .expect("module should enter Python locals"); + let code = CString::new( + r#" +import asyncio + +async def exercise(): + assert await runtime.async_custom_mapping() == "mapped" + +asyncio.run(exercise()) +"#, + ) + .expect("Python source should not contain null bytes"); + py.run(&code, Some(&locals), Some(&locals)) + .expect("custom mapping should reach the Python awaiter"); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/streaming.rs b/litellm-rust/crates/python-bridge/src/routes/streaming.rs new file mode 100644 index 00000000000..7b4edc50d4d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/streaming.rs @@ -0,0 +1,422 @@ +use std::sync::Arc; + +use litellm_core::chat_completions::chat_completions_stream as run_chat_completions_stream; +use litellm_core::chat_completions::types::{ + ChatCompletionsStreamRequest, ChatCompletionsStreamRequestBody, ChatStreamEvent, +}; +use litellm_core::messages::messages_event_stream as run_messages_stream; +use litellm_core::messages::types::{ + AnthropicMessagesRequest, MessagesStreamEvent, MessagesStreamRequest, +}; +use litellm_core::responses::responses_stream as run_responses_stream; +use litellm_core::responses::responses_websocket as run_responses_websocket; +use litellm_core::responses::types::{ + ResponseCommand, ResponsesStreamEvent, ResponsesStreamRequest, ResponsesStreamRequestBody, + ResponsesWebSocketRequest, +}; +use litellm_core::responses::websocket::TypedResponsesWebSocketSession; +use litellm_core::streaming::{ + JsonObject, OpenedStream, ProviderCredentials, StreamMetadata, StreamProviderId, StreamTarget, + StreamTransportOptions, +}; +use litellm_python_interop::{from_py, to_py}; +use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration}; +use pyo3::prelude::*; +use pyo3::types::{PyModule, PyType}; +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; + +use crate::marshal::{marshal_headers, optional_timeout}; +use crate::routes::receiver::BridgeReceiver; +use crate::routes::runtime::{run_async, run_async_with, run_sync_with}; + +struct TypedEventReceiver { + metadata: StreamMetadata, + receiver: BridgeReceiver, +} + +impl TypedEventReceiver +where + E: Send + 'static, +{ + fn from_opened(opened: OpenedStream) -> Self { + Self { + metadata: opened.metadata, + receiver: BridgeReceiver::from_stream(opened.events), + } + } +} + +fn next_event( + py: Python<'_>, + receiver: BridgeReceiver, + stop_iteration: bool, +) -> PyResult> +where + E: Serialize + Send + 'static, +{ + run_sync_with( + py, + async move { receiver.next().await }, + crate::errors::executed_route_error_to_pyerr, + move |py, event| match event { + Some(event) => to_py(py, &event), + None if stop_iteration => Err(PyStopIteration::new_err(())), + None => Ok(py.None()), + }, + ) +} + +fn anext_event( + py: Python<'_>, + receiver: BridgeReceiver, + stop_iteration: bool, +) -> PyResult> +where + E: Serialize + Send + 'static, +{ + run_async_with( + py, + async move { receiver.next().await }, + crate::errors::executed_route_error_to_pyerr, + move |py, event| match event { + Some(event) => to_py(py, &event), + None if stop_iteration => Err(PyStopAsyncIteration::new_err(())), + None => Ok(py.None()), + }, + ) +} + +macro_rules! event_stream_class { + ($class:ident, $event:ty) => { + #[pyclass] + struct $class { + inner: TypedEventReceiver<$event>, + } + + impl From> for $class { + fn from(opened: OpenedStream<$event>) -> Self { + Self { + inner: TypedEventReceiver::from_opened(opened), + } + } + } + + #[pymethods] + impl $class { + #[getter] + fn metadata(&self, py: Python<'_>) -> PyResult> { + to_py(py, &self.inner.metadata) + } + + fn next_event(&self, py: Python<'_>) -> PyResult> { + next_event(py, self.inner.receiver.clone(), false) + } + + fn anext_event<'py>(&self, py: Python<'py>) -> PyResult> { + anext_event(py, self.inner.receiver.clone(), false) + } + + fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __next__(&self, py: Python<'_>) -> PyResult> { + next_event(py, self.inner.receiver.clone(), true) + } + + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + anext_event(py, self.inner.receiver.clone(), true) + } + + fn close(&self) { + self.inner.receiver.close(); + } + + fn aclose<'py>(&self, py: Python<'py>) -> PyResult> { + let receiver = self.inner.receiver.clone(); + run_async( + py, + async move { + receiver.close(); + Ok(()) + }, + crate::errors::executed_route_error_to_pyerr, + ) + } + } + }; +} + +event_stream_class!(ChatCompletionsEventStream, ChatStreamEvent); +event_stream_class!(MessagesEventStream, MessagesStreamEvent); +event_stream_class!(ResponsesEventStream, ResponsesStreamEvent); + +#[derive(Default, Deserialize)] +struct PythonProviderCredentials { + api_key: Option, + aws_access_key_id: Option, + aws_secret_access_key: Option, + aws_session_token: Option, +} + +struct PythonStreamTarget { + provider: String, + credentials: Option>, + api_base: Option, +} + +struct PythonStreamTransport { + extra_headers: Option>, + timeout_seconds: Option, +} + +fn parse_call( + py: Python<'_>, + request: Py, + target: PythonStreamTarget, + transport: PythonStreamTransport, +) -> PyResult<(B, StreamTarget, StreamTransportOptions)> +where + B: DeserializeOwned, +{ + (|| -> PyResult<(B, StreamTarget, StreamTransportOptions)> { + let body = from_py(request.bind(py))?; + let credentials = target + .credentials + .map(|value| from_py::(value.bind(py))) + .transpose()? + .unwrap_or_default(); + let provider = StreamProviderId::try_from(target.provider.as_str()) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + let target = StreamTarget::new( + provider, + ProviderCredentials::new( + credentials.api_key, + credentials.aws_access_key_id, + credentials.aws_secret_access_key, + credentials.aws_session_token, + ), + target.api_base, + ); + let extra_headers = transport + .extra_headers + .map(|value| from_py(value.bind(py))) + .transpose()?; + let forwarded_headers = marshal_headers(extra_headers)? + .into_iter() + .map(|(name, value)| litellm_core::streaming::Header { name, value }) + .collect(); + let transport = StreamTransportOptions::new( + forwarded_headers, + optional_timeout(transport.timeout_seconds), + ); + Ok((body, target, transport)) + })() + .map_err(crate::errors::declined) +} + +macro_rules! stream_openers { + ( + sync = $sync_name:ident, + asynchronous = $async_name:ident, + body = $body:ty, + request = $request:ident, + open = $open:path, + stream = $stream:ident + ) => { + #[pyfunction] + #[pyo3(signature = (request, provider, credentials=None, api_base=None, extra_headers=None, timeout_seconds=None))] + #[allow(clippy::too_many_arguments)] + fn $sync_name( + py: Python<'_>, + request: Py, + provider: String, + credentials: Option>, + api_base: Option, + extra_headers: Option>, + timeout_seconds: Option, + ) -> PyResult> { + let (body, target, transport) = parse_call::<$body>( + py, + request, + PythonStreamTarget { + provider, + credentials, + api_base, + }, + PythonStreamTransport { + extra_headers, + timeout_seconds, + }, + )?; + run_sync_with( + py, + async move { $open($request { body, target, transport }).await }, + crate::errors::fallback_route_error_to_pyerr, + |py, opened| Ok(Py::new(py, $stream::from(opened))?.into_any()), + ) + } + + #[pyfunction] + #[pyo3(signature = (request, provider, credentials=None, api_base=None, extra_headers=None, timeout_seconds=None))] + #[allow(clippy::too_many_arguments)] + fn $async_name( + py: Python<'_>, + request: Py, + provider: String, + credentials: Option>, + api_base: Option, + extra_headers: Option>, + timeout_seconds: Option, + ) -> PyResult> { + let (body, target, transport) = parse_call::<$body>( + py, + request, + PythonStreamTarget { + provider, + credentials, + api_base, + }, + PythonStreamTransport { + extra_headers, + timeout_seconds, + }, + )?; + run_async_with( + py, + async move { $open($request { body, target, transport }).await }, + crate::errors::fallback_route_error_to_pyerr, + |py, opened| Ok(Py::new(py, $stream::from(opened))?.into_any()), + ) + } + }; +} + +stream_openers! { + sync = chat_completions_stream, + asynchronous = achat_completions_stream, + body = ChatCompletionsStreamRequestBody, + request = ChatCompletionsStreamRequest, + open = run_chat_completions_stream, + stream = ChatCompletionsEventStream +} + +stream_openers! { + sync = messages_stream, + asynchronous = amessages_stream, + body = AnthropicMessagesRequest, + request = MessagesStreamRequest, + open = run_messages_stream, + stream = MessagesEventStream +} + +stream_openers! { + sync = responses_stream, + asynchronous = aresponses_stream, + body = ResponsesStreamRequestBody, + request = ResponsesStreamRequest, + open = run_responses_stream, + stream = ResponsesEventStream +} + +#[pyclass] +struct ResponsesWebSocketSession { + session: Arc, +} + +#[pymethods] +impl ResponsesWebSocketSession { + #[classmethod] + #[pyo3(signature = (provider, credentials=None, api_base=None, extra_headers=None, timeout_seconds=None))] + #[allow(clippy::too_many_arguments)] + fn connect<'py>( + _cls: &Bound<'py, PyType>, + py: Python<'py>, + provider: String, + credentials: Option>, + api_base: Option, + extra_headers: Option>, + timeout_seconds: Option, + ) -> PyResult> { + let empty_request = pyo3::types::PyDict::new(py).unbind().into_any(); + let (_, target, transport) = parse_call::( + py, + empty_request, + PythonStreamTarget { + provider, + credentials, + api_base, + }, + PythonStreamTransport { + extra_headers, + timeout_seconds, + }, + )?; + run_async_with( + py, + async move { + run_responses_websocket(ResponsesWebSocketRequest { target, transport }).await + }, + crate::errors::fallback_route_error_to_pyerr, + |py, session| { + Ok(Py::new( + py, + ResponsesWebSocketSession { + session: Arc::from(session), + }, + )? + .into_any()) + }, + ) + } + + fn send_event<'py>(&self, py: Python<'py>, command: Py) -> PyResult> { + let command: ResponseCommand = from_py(command.bind(py))?; + let session = self.session.clone(); + run_async( + py, + async move { session.send(command).await }, + crate::errors::executed_route_error_to_pyerr, + ) + } + + fn recv_event<'py>(&self, py: Python<'py>) -> PyResult> { + let session = self.session.clone(); + run_async_with( + py, + async move { session.recv().await }, + crate::errors::executed_route_error_to_pyerr, + |py, event| match event { + Some(event) => to_py(py, &event), + None => Ok(py.None()), + }, + ) + } + + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + let session = self.session.clone(); + run_async( + py, + async move { session.close().await }, + crate::errors::executed_route_error_to_pyerr, + ) + } +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_function(wrap_pyfunction!(chat_completions_stream, module)?)?; + module.add_function(wrap_pyfunction!(achat_completions_stream, module)?)?; + module.add_function(wrap_pyfunction!(messages_stream, module)?)?; + module.add_function(wrap_pyfunction!(amessages_stream, module)?)?; + module.add_function(wrap_pyfunction!(responses_stream, module)?)?; + module.add_function(wrap_pyfunction!(aresponses_stream, module)?) +} diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 2789def7784..8f51c33565e 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -11,6 +11,7 @@ import sys import tempfile import threading import zipfile +from collections.abc import Awaitable, Callable from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path @@ -253,11 +254,100 @@ async def exercise_async_concurrency(native: object, api_base: str) -> None: assert_success("messages", response) +def streaming_route_kwargs(route: str) -> dict[str, object]: + if route == "chat_completions": + return { + "request": { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + }, + "provider": "anthropic", + } + if route == "messages": + return { + "request": { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 16, + "stream": True, + }, + "provider": "anthropic", + } + if route == "responses": + return { + "request": {"model": "gpt-5", "input": "hello", "stream": True}, + "provider": "openai", + } + raise AssertionError(f"unknown streaming route: {route}") + + +def assert_streaming_decline(native: object, operation: Callable[[], object]) -> None: + try: + operation() + except native.RustBridgeDeclined: + return + raise AssertionError("disabled native streaming route did not decline") + + +async def assert_async_streaming_decline( + native: object, + operation: Callable[[], Awaitable[object]], +) -> None: + try: + await operation() + except native.RustBridgeDeclined: + return + raise AssertionError("disabled async native streaming route did not decline") + + +def exercise_disabled_streaming_surface(native: object) -> None: + expected: Final = ( + "ChatCompletionsEventStream", + "MessagesEventStream", + "ResponsesEventStream", + "ResponsesWebSocketSession", + "chat_completions_stream", + "achat_completions_stream", + "messages_stream", + "amessages_stream", + "responses_stream", + "aresponses_stream", + ) + for name in expected: + if not hasattr(native, name): + raise AssertionError(f"packaged native bridge is missing {name}") + + for route in ("chat_completions", "messages", "responses"): + function: Final = getattr(native, f"{route}_stream") + kwargs: Final = streaming_route_kwargs(route) + assert_streaming_decline( + native, + lambda function=function, kwargs=kwargs: function(**kwargs), + ) + + async def exercise_async_surface() -> None: + for route in ("chat_completions", "messages", "responses"): + function: Final = getattr(native, f"a{route}_stream") + kwargs: Final = streaming_route_kwargs(route) + await assert_async_streaming_decline( + native, + lambda function=function, kwargs=kwargs: function(**kwargs), + ) + await assert_async_streaming_decline( + native, + lambda: native.ResponsesWebSocketSession.connect(provider="openai"), + ) + + asyncio.run(exercise_async_surface()) + + def exercise_routes(native_path: Path, api_base: str) -> object: native: Final = load_native(native_path) exercise_sync(native, api_base) asyncio.run(exercise_async(native, api_base)) asyncio.run(exercise_async_concurrency(native, api_base)) + exercise_disabled_streaming_surface(native) return native