diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..93e1c56f9a7 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1435,6 +1435,8 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "bytes", + "futures-util", "rand 0.8.7", "reqwest", "serde", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a13dd4c04b0..59d99040eff 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -31,7 +31,8 @@ serde_json = "1.0" sha2 = "0.10" subtle = "2" thiserror = "2.0" -tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net", "sync"] } +bytes = "1" 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"] } base64 = "0.22" diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 389dbd49505..c752487f561 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -6,6 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +bytes.workspace = true +futures-util.workspace = true rand.workspace = true reqwest.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 3633130528d..cddcf15e5b5 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -1,5 +1,10 @@ //! Header and upstream-body helpers shared by every route module. +use std::pin::Pin; +use std::task::{Context, Poll}; + +use bytes::Bytes; +use futures_util::stream::Stream; use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; @@ -12,6 +17,77 @@ pub async fn http_request( request.send().await } +/// A stream of complete server-sent-event frames. +/// +/// Wraps an upstream byte stream and reassembles it into frames delimited by a +/// blank line (`\n\n` or `\r\n\r\n`), so a host receives whole `event:`/`data:` +/// frames regardless of how the provider chunked the response. Frames that are +/// only whitespace (keep-alive newlines) are dropped. A trailing partial frame +/// at end-of-stream is emitted as-is rather than lost. +pub struct SseFrameStream { + inner: Pin> + Send>>, + buffer: Vec, + done: bool, +} + +impl SseFrameStream { + pub fn new(response: reqwest::Response) -> Self { + Self::from_byte_stream(Box::pin(response.bytes_stream())) + } + + fn from_byte_stream(inner: Pin> + Send>>) -> Self { + Self { + inner, + buffer: Vec::new(), + done: false, + } + } +} + +impl Stream for SseFrameStream { + type Item = Result, Error>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + if let Some(end) = find_frame_end(&self.buffer) { + let frame: Vec = self.buffer.drain(..end).collect(); + if !frame.iter().all(u8::is_ascii_whitespace) { + return Poll::Ready(Some(Ok(frame))); + } + continue; + } + if self.done { + if self.buffer.is_empty() { + return Poll::Ready(None); + } + let frame = std::mem::take(&mut self.buffer); + return Poll::Ready(Some(Ok(frame))); + } + match self.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => self.buffer.extend_from_slice(&chunk), + Poll::Ready(Some(Err(error))) => { + self.done = true; + return Poll::Ready(Some(Err(Error::Network(error.to_string())))); + } + Poll::Ready(None) => self.done = true, + Poll::Pending => return Poll::Pending, + } + } + } +} + +fn find_frame_end(buffer: &[u8]) -> Option { + for index in 0..buffer.len() { + if buffer[index..].starts_with(b"\r\n\r\n") { + return Some(index + 4); + } + if buffer[index..].starts_with(b"\n\n") { + return Some(index + 2); + } + } + None +} + /// Bound an upstream error body before it crosses a host boundary, so provider /// bodies stay data-minimized. pub fn truncate_error_body(body: &str) -> String { @@ -63,9 +139,29 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool { #[cfg(test)] mod tests { - use super::*; + use bytes::Bytes; + use futures_util::StreamExt; use serde_json::json; + use super::*; + + fn frame_stream(chunks: Vec<&[u8]>) -> SseFrameStream { + let chunks: Vec> = chunks + .into_iter() + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect(); + SseFrameStream::from_byte_stream(Box::pin(futures_util::stream::iter(chunks))) + } + + async fn collect_frames(chunks: Vec<&[u8]>) -> Vec> { + let mut stream = frame_stream(chunks); + let mut frames = Vec::new(); + while let Some(frame) = stream.next().await { + frames.push(frame.expect("frame should decode")); + } + frames + } + #[test] fn truncate_leaves_short_bodies_untouched() { assert_eq!(truncate_error_body("short"), "short"); @@ -116,4 +212,63 @@ mod tests { "Basic abc".to_string() )])); } + + #[tokio::test] + async fn sse_frames_split_on_blank_lines() { + let frames = collect_frames(vec![b"event: a\ndata: {}\n\nevent: b\ndata: {}\n\n"]).await; + assert_eq!( + frames, + vec![ + b"event: a\ndata: {}\n\n".to_vec(), + b"event: b\ndata: {}\n\n".to_vec() + ] + ); + } + + #[tokio::test] + async fn sse_frames_reassemble_across_chunk_boundaries() { + let frames = collect_frames(vec![ + b"event: message_star", + b"t\ndata: {\"a\":1}\n", + b"\nevent: message_stop\ndata: {}\n", + b"\n", + ]) + .await; + assert_eq!( + frames, + vec![ + b"event: message_start\ndata: {\"a\":1}\n\n".to_vec(), + b"event: message_stop\ndata: {}\n\n".to_vec() + ] + ); + } + + #[tokio::test] + async fn sse_frames_accept_crlf_delimiters() { + let frames = collect_frames(vec![b"event: a\r\ndata: {}\r\n\r\n"]).await; + assert_eq!(frames, vec![b"event: a\r\ndata: {}\r\n\r\n".to_vec()]); + } + + #[tokio::test] + async fn sse_frames_drop_whitespace_only_keep_alive_chunks() { + let frames = collect_frames(vec![b"\n\nevent: a\ndata: {}\n\n\n\n"]).await; + assert_eq!(frames, vec![b"event: a\ndata: {}\n\n".to_vec()]); + } + + #[tokio::test] + async fn sse_frames_emit_trailing_partial_frame_at_eof() { + let frames = collect_frames(vec![b"event: a\ndata: {}\n\nevent: b\ndata: {"]).await; + assert_eq!( + frames, + vec![ + b"event: a\ndata: {}\n\n".to_vec(), + b"event: b\ndata: {".to_vec() + ] + ); + } + + #[tokio::test] + async fn empty_input_yields_no_frames() { + assert!(collect_frames(vec![]).await.is_empty()); + } } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 61ff81bcdc8..934b2269a69 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -45,12 +45,17 @@ pub(super) async fn execute_messages_provider_call( pub(super) async fn execute_messages_provider_stream( request: MessagesRequest<'_>, ) -> Result { - let request = prepare_provider_request(request)?; + let mut request = prepare_provider_request(request)?; if request.provider != ANTHROPIC_MESSAGES_PROVIDER { return Err(Error::InvalidRequest( "streaming messages is not supported for this provider".to_string(), )); } + // The streaming entrypoints only make sense for `stream: true`; force it so + // a host cannot accidentally ask for SSE and receive a buffered JSON body. + if let Some(body) = request.body.as_object_mut() { + body.insert("stream".to_string(), serde_json::Value::Bool(true)); + } let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index cfa8bda1104..7640e1e3d1f 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -8,6 +8,7 @@ //! can splice the event stream to its own caller. use crate::Error; +use crate::http_utils::SseFrameStream; mod client; mod common_utils; mod handler; @@ -27,5 +28,15 @@ pub async fn messages_stream(request: MessagesRequest<'_>) -> Result) -> Result { + let response = execute_messages_provider_stream(request).await?; + Ok(SseFrameStream::new(response)) +} + #[cfg(test)] mod tests; diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index df9f7051011..d23b56c2868 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -439,3 +439,118 @@ async fn messages_rejects_unsupported_provider() { assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); } + +#[tokio::test] +async fn messages_stream_frames_yields_complete_frames_and_forces_stream_flag() { + use futures_util::StreamExt; + + use super::messages_stream_frames; + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + // Send the SSE body in multiple writes so frames span network chunks. + let body = concat!( + "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response head"); + let bytes = body.as_bytes(); + for chunk in [&bytes[..40], &bytes[40..90], &bytes[90..]] { + tokio::time::sleep(Duration::from_millis(10)).await; + socket.write_all(chunk).await.expect("writes body chunk"); + } + request + }); + + let mut stream = messages_stream_frames(MessagesRequest { + model: "claude-sonnet-4-5", + // `stream` deliberately omitted: the entrypoint must force it upstream. + body: json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "hi"}] + }), + api_key: Some("sk-ant"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("stream request succeeds"); + + let mut frames = Vec::new(); + while let Some(frame) = stream.next().await { + frames.push(String::from_utf8(frame.expect("frame decodes")).expect("frame is utf8")); + } + assert_eq!( + frames, + vec![ + "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ] + ); + + let request = server.await.expect("server task completes"); + let (_, body) = request.split_once("\r\n\r\n").expect("has body"); + let sent_body: Value = serde_json::from_str(body).expect("body is json"); + assert_eq!(sent_body["stream"], Value::Bool(true)); +} + +#[tokio::test] +async fn messages_stream_frames_maps_upstream_error_status() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + let body = "{\"error\":\"rate limited\"}"; + let response = format!( + "HTTP/1.1 429 Too Many Requests\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + + use super::messages_stream_frames; + + let err = match messages_stream_frames(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "hi"}], + "stream": true + }), + api_key: Some("sk-ant"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + }) + .await + { + Err(err) => err, + Ok(_) => panic!("provider error should propagate before the stream"), + }; + + assert!(matches!(err, Error::Http { status: 429, .. })); +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 637e5580170..08c89625d2c 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -26,7 +26,7 @@ pyo3.workspace = true pyo3-async-runtimes.workspace = true serde.workspace = true serde_json.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] criterion = "0.8.2" diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5f36a22370a..16a16cd6473 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -102,6 +102,8 @@ mod tests { "atranscription", "messages", "amessages", + "amessages_stream", + "MessagesStream", "chat_completions_decline", "chat_completions", "achat_completions", diff --git a/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs b/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs new file mode 100644 index 00000000000..fbe2873d3f8 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs @@ -0,0 +1,520 @@ +//! The streaming Anthropic Messages route: `amessages_stream`. +//! +//! Returns a Python async iterator (`MessagesStream`) that yields complete +//! SSE frames as `bytes`. There is deliberately no sync twin: the Python +//! consumer of this path is async-only, and no `bridge_route!` macro because +//! the route's value is a stream, not a terminal `Serialize`d response. + +use std::sync::Arc; + +use futures_util::StreamExt; +use litellm_core::http_utils::SseFrameStream; +use litellm_core::messages::messages_stream_frames; +use litellm_core::messages::types::MessagesRequest; +use pyo3::exceptions::{PyStopAsyncIteration, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyBytes}; +use serde_json::Value; +use tokio::sync::Mutex; + +use crate::errors::core_error_to_pyerr; +use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; + +/// One complete SSE frame, converted to `bytes` on the attached thread that +/// completes the awaitable (mirrors `Pythonized` in litellm-python-interop). +struct SseFrame(Vec); + +impl<'py> IntoPyObject<'py> for SseFrame { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(PyBytes::new(py, &self.0).into_any()) + } +} + +/// Async iterator over the upstream SSE frames. +/// +/// Pull-based: each `__anext__` awaits exactly one frame, so backpressure is +/// structural (nothing buffers beyond the current frame). Dropping the object +/// (GC, `break` out of `async for`) drops the upstream response body, which +/// aborts the provider request. Cancelling a task awaiting `__anext__` only +/// drops the frame fetch, so a later `__anext__` resumes without data loss. +#[pyclass] +struct MessagesStream { + frames: Arc>>, +} + +#[pymethods] +impl MessagesStream { + fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + let frames = Arc::clone(&self.frames); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let mut guard = frames.lock().await; + let frame = match guard.as_mut() { + Some(stream) => stream.next().await, + None => None, + }; + match frame { + Some(Ok(frame)) => Ok(SseFrame(frame)), + Some(Err(error)) => { + // The upstream stream is dead after a mid-stream failure. + *guard = None; + Err(core_error_to_pyerr(error)) + } + None => { + *guard = None; + Err(PyStopAsyncIteration::new_err(())) + } + } + }) + } +} + +#[pyfunction] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=false))] +#[allow(clippy::too_many_arguments)] +fn amessages_stream<'py>( + py: Python<'py>, + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option, + timeout_seconds: Option, + trace: bool, +) -> PyResult> { + // `trace_call` wraps a completed response; a stream has no terminal value + // to attach the trace to. Reject rather than silently drop the request. + if trace { + return Err(PyValueError::new_err( + "trace is not supported on the streaming route", + )); + } + let body = required_value("body", body, Value::is_object, "dict")?; + let options = RouteOptions::from_python(RouteOptionsInputs { + model, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout_seconds, + })?; + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let RouteOptions { + model, + api_key, + api_base, + custom_llm_provider, + extra_headers, + timeout, + } = options; + let frames = messages_stream_frames(MessagesRequest { + model: &model, + body, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + timeout, + }) + .await + .map_err(core_error_to_pyerr)?; + Ok(MessagesStream { + frames: Arc::new(Mutex::new(Some(frames))), + }) + }) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + crate::routes::definition::add_function( + module, + pyo3::wrap_pyfunction!(amessages_stream, module)?, + )?; + module.add_class::() +} + +#[cfg(test)] +mod tests { + use std::ffi::CString; + use std::time::Duration; + + use pyo3::types::{PyDict, PyModule}; + use tokio::io::AsyncWriteExt; + use tokio::net::TcpListener; + + use super::*; + + async fn read_http_request(socket: &mut tokio::net::TcpStream) -> String { + use tokio::io::AsyncReadExt; + + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let read = socket.read(&mut buffer).await.expect("reads request"); + if read == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..read]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let read = socket.read(&mut buffer).await.expect("reads body"); + if read == 0 { + break; + } + request.extend_from_slice(&buffer[..read]); + } + String::from_utf8(request).expect("request is utf8") + } + + /// An SSE upstream that sends `body` in chunks separated by short delays + /// so frames span network chunk boundaries. + async fn streaming_upstream( + listener: TcpListener, + status_line: &'static str, + body: &'static str, + ) { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + assert!( + request.contains("\"stream\":true"), + "upstream body must force stream:true: {request}" + ); + let response = format!( + "{status_line}\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response head"); + let bytes = body.as_bytes(); + let chunks: Vec<&[u8]> = match bytes.len() { + 0..=30 => vec![bytes], + 31..=70 => vec![&bytes[..30], &bytes[30..]], + _ => vec![&bytes[..30], &bytes[30..70], &bytes[70..]], + }; + for chunk in chunks { + tokio::time::sleep(Duration::from_millis(10)).await; + socket.write_all(chunk).await.expect("writes body chunk"); + } + } + + fn register_stream_module(py: Python<'_>) -> Bound<'_, PyModule> { + let module = PyModule::new(py, "messages_stream").expect("module should be created"); + register(&module).expect("stream route should register"); + module + } + + fn run_async_code(py: Python<'_>, module: &Bound<'_, PyModule>, url: &str, code: &str) { + let locals = PyDict::new(py); + locals + .set_item("runtime", module) + .expect("module should enter Python locals"); + locals + .set_item("url", url) + .expect("URL should enter Python locals"); + let source = CString::new(code).expect("Python source should not contain null bytes"); + py.run(&source, Some(&locals), Some(&locals)) + .expect("Python stream exercise should pass"); + } + + #[test] + fn stream_route_yields_complete_frames_in_order() { + Python::initialize(); + let runtime = pyo3_async_runtimes::tokio::get_runtime(); + let listener = runtime + .block_on(TcpListener::bind("127.0.0.1:0")) + .expect("listener should bind"); + let address = listener + .local_addr() + .expect("listener should have an address"); + let body = concat!( + "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", + "event: content_block_delta\ndata: {\"delta\":\"hi\"}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ); + let server = runtime.spawn(streaming_upstream(listener, "HTTP/1.1 200 OK", body)); + + Python::attach(|py| { + let module = register_stream_module(py); + run_async_code( + py, + &module, + &format!("http://{address}"), + r#" +import asyncio + +async def exercise(): + stream = await runtime.amessages_stream( + model="claude-sonnet-4-5", + body={"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}, + api_key="sk-ant", + api_base=url, + custom_llm_provider="anthropic", + ) + frames = [chunk async for chunk in stream] + assert frames == [ + b'event: message_start\ndata: {"type":"message_start"}\n\n', + b'event: content_block_delta\ndata: {"delta":"hi"}\n\n', + b'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ], frames + +asyncio.run(exercise()) +"#, + ); + }); + + runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await }) + .expect("server should finish") + .expect("server task should not panic"); + } + + #[test] + fn stream_route_surfaces_upstream_errors_before_the_stream() { + Python::initialize(); + let runtime = pyo3_async_runtimes::tokio::get_runtime(); + let listener = runtime + .block_on(TcpListener::bind("127.0.0.1:0")) + .expect("listener should bind"); + let address = listener + .local_addr() + .expect("listener should have an address"); + let server = runtime.spawn(streaming_upstream( + listener, + "HTTP/1.1 429 Too Many Requests", + "{\"error\":\"rate limited\"}", + )); + + Python::attach(|py| { + let module = register_stream_module(py); + run_async_code( + py, + &module, + &format!("http://{address}"), + r#" +import asyncio + +async def exercise(): + try: + await runtime.amessages_stream( + model="claude-sonnet-4-5", + body={"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}, + api_key="sk-ant", + api_base=url, + custom_llm_provider="anthropic", + ) + except RuntimeError as error: + assert "429" in str(error), str(error) + else: + raise AssertionError("upstream error should surface from the awaitable") + +asyncio.run(exercise()) +"#, + ); + }); + + runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await }) + .expect("server should finish") + .expect("server task should not panic"); + } + + #[test] + fn stream_route_ends_iteration_after_a_mid_stream_failure() { + Python::initialize(); + let runtime = pyo3_async_runtimes::tokio::get_runtime(); + let listener = runtime + .block_on(TcpListener::bind("127.0.0.1:0")) + .expect("listener should bind"); + let address = listener + .local_addr() + .expect("listener should have an address"); + // Advertise more bytes than are sent, then close: reqwest reports an + // incomplete body mid-stream after the buffered frames are delivered. + let server = runtime.spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _request = read_http_request(&mut socket).await; + let body = "event: message_start\ndata: {}\n\n"; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len() + 40, + body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + socket.shutdown().await.expect("closes socket"); + }); + + Python::attach(|py| { + let module = register_stream_module(py); + run_async_code( + py, + &module, + &format!("http://{address}"), + r#" +import asyncio + +async def exercise(): + stream = await runtime.amessages_stream( + model="claude-sonnet-4-5", + body={"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}, + api_key="sk-ant", + api_base=url, + custom_llm_provider="anthropic", + ) + frames = [] + failed = False + try: + async for chunk in stream: + frames.append(chunk) + except Exception: + failed = True + assert frames == [b'event: message_start\ndata: {}\n\n'], frames + assert failed, "incomplete body should raise mid-stream" + # After the failure the iterator is terminal. + try: + await stream.__anext__() + except StopAsyncIteration: + pass + else: + raise AssertionError("stream should be terminal after a failure") + +asyncio.run(exercise()) +"#, + ); + }); + + runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await }) + .expect("server should finish") + .expect("server task should not panic"); + } + + #[test] + fn stream_route_rejects_trace() { + Python::initialize(); + Python::attach(|py| { + let module = register_stream_module(py); + let error = amessages_stream( + py, + "claude-sonnet-4-5".to_string(), + serde_json::json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + None, + None, + None, + None, + None, + true, + ) + .expect_err("trace must be rejected"); + assert!(error.is_instance_of::(py)); + assert_eq!( + error.to_string(), + "ValueError: trace is not supported on the streaming route" + ); + let _module = module; + }); + } + + #[test] + fn stream_route_supports_cancel_and_resume_without_frame_loss() { + Python::initialize(); + let runtime = pyo3_async_runtimes::tokio::get_runtime(); + let listener = runtime + .block_on(TcpListener::bind("127.0.0.1:0")) + .expect("listener should bind"); + let address = listener + .local_addr() + .expect("listener should have an address"); + // First frame arrives after a delay; a cancelled `__anext__` must not + // consume it, and resuming must deliver it exactly once. + let first_frame = "event: message_start\ndata: {}\n\n"; + let second_frame = "event: message_stop\ndata: {}\n\n"; + let total = first_frame.len() + second_frame.len(); + let server = runtime.spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _request = read_http_request(&mut socket).await; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {total}\r\nconnection: close\r\n\r\n" + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response head"); + tokio::time::sleep(Duration::from_millis(300)).await; + socket + .write_all(first_frame.as_bytes()) + .await + .expect("writes first frame"); + tokio::time::sleep(Duration::from_millis(300)).await; + socket + .write_all(second_frame.as_bytes()) + .await + .expect("writes second frame"); + socket.shutdown().await.expect("closes socket"); + }); + + Python::attach(|py| { + let module = register_stream_module(py); + run_async_code( + py, + &module, + &format!("http://{address}"), + r#" +import asyncio + +async def exercise(): + stream = await runtime.amessages_stream( + model="claude-sonnet-4-5", + body={"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}, + api_key="sk-ant", + api_base=url, + custom_llm_provider="anthropic", + ) + next_chunk = asyncio.ensure_future(stream.__anext__()) + await asyncio.sleep(0.05) + next_chunk.cancel() + try: + await next_chunk + except asyncio.CancelledError: + pass + frames = [chunk async for chunk in stream] + assert frames == [ + b'event: message_start\ndata: {}\n\n', + b'event: message_stop\ndata: {}\n\n', + ], frames + +asyncio.run(asyncio.wait_for(exercise(), timeout=10)) +"#, + ); + }); + + runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await }) + .expect("server should finish") + .expect("server task should not panic"); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index bf611c26d44..00278428f6f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -6,11 +6,13 @@ mod definition; mod audio_transcription; mod chat_completions; mod messages; +mod messages_stream; mod ocr; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { ocr::register(module)?; audio_transcription::register(module)?; messages::register(module)?; + messages_stream::register(module)?; chat_completions::register(module) } diff --git a/litellm-rust/crates/python-bridge/src/routes/runtime.rs b/litellm-rust/crates/python-bridge/src/routes/runtime.rs deleted file mode 100644 index 87a0c3e0104..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/runtime.rs +++ /dev/null @@ -1,423 +0,0 @@ -use std::future::Future; -use std::panic::AssertUnwindSafe; -use std::time::Duration; - -use futures_util::FutureExt; -use litellm_core::error::Error; -use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::*; -use serde::Serialize; -use tokio::runtime::{Handle, Runtime}; -use tokio::time::{self, MissedTickBehavior}; - -pub(super) fn run_sync( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - run_sync_on( - py, - pyo3_async_runtimes::tokio::get_runtime(), - future, - map_error, - ) -} - -fn run_sync_on( - py: Python<'_>, - runtime: &Runtime, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + Send + 'static, -{ - 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 || runtime.block_on(wait_for_sync_result(future)))?; - let result = map_core_result(result, map_error)?; - Pythonized(result).into_pyobject(py).map(Bound::unbind) -} - -pub(super) fn run_async( - py: Python<'_>, - future: F, - map_error: fn(Error) -> PyErr, -) -> PyResult> -where - T: Serialize + Send + 'static, - F: Future> + 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(Pythonized(result)) - }) -} - -fn map_core_result(result: Result, map_error: fn(Error) -> PyErr) -> PyResult { - match result { - Ok(value) => Ok(value), - Err(error) => Err( - std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error))) - .map_err(panic_to_pyerr)?, - ), - } -} - -async fn catch_route_panic(future: F) -> PyResult> -where - F: Future>, -{ - AssertUnwindSafe(future) - .catch_unwind() - .await - .map_err(panic_to_pyerr) -} - -async fn wait_for_sync_result(future: F) -> PyResult> -where - F: Future>, -{ - let future = catch_route_panic(future); - tokio::pin!(future); - - let signal_interval = Duration::from_millis(50); - let mut signal_checks = - time::interval_at(time::Instant::now() + signal_interval, signal_interval); - signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay); - loop { - tokio::select! { - result = &mut future => return result, - _ = signal_checks.tick() => Python::attach(|py| py.check_signals())?, - } - } -} - -#[cfg(test)] -mod tests { - use std::ffi::CString; - use std::future::poll_fn; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, mpsc}; - use std::task::Poll; - use std::thread; - use std::time::Instant; - - use pyo3::panic::PanicException; - use pyo3::types::{PyDict, PyModule}; - use serde::Serializer; - use tokio::runtime::Builder; - - use super::*; - - fn runtime_error(error: Error) -> PyErr { - PyRuntimeError::new_err(error.to_string()) - } - - fn panicking_error_mapper(_error: Error) -> PyErr { - panic!("error mapper panicked") - } - - struct PanickingOutput; - - static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0); - - impl Serialize for PanickingOutput { - fn serialize(&self, _serializer: S) -> Result - where - S: Serializer, - { - panic!("serializer panicked") - } - } - - #[pyfunction] - fn async_serialization_panic(py: Python<'_>) -> PyResult> { - run_async(py, async { Ok(PanickingOutput) }, runtime_error) - } - - #[pyfunction] - fn async_runtime_probe(py: Python<'_>) -> PyResult> { - run_async( - py, - async { - ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst); - Ok(true) - }, - runtime_error, - ) - } - - #[pyfunction] - fn runtime_worker_count() -> usize { - pyo3_async_runtimes::tokio::get_runtime() - .metrics() - .num_workers() - } - - #[pyfunction] - fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool { - let completion_deadline = Instant::now() + Duration::from_secs(2); - while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions { - if Instant::now() >= completion_deadline { - return false; - } - thread::sleep(Duration::from_millis(1)); - } - - let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1); - pyo3_async_runtimes::tokio::get_runtime().spawn(async move { - let _ = heartbeat_tx.send(()); - }); - heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok() - } - - fn extract_bool(py: Python<'_>, result: PyResult>) -> bool { - result - .expect("route should complete") - .bind(py) - .extract() - .expect("result should convert") - } - - #[test] - fn sync_runner_polls_future_on_the_caller_thread() { - Python::initialize(); - Python::attach(|py| { - let caller_thread = std::thread::current().id(); - let result = run_sync( - py, - async move { Ok(std::thread::current().id() == caller_thread) }, - runtime_error, - ); - - assert!(extract_bool(py, result)); - }); - } - - #[test] - fn sync_runner_releases_gil_while_waiting() { - Python::initialize(); - Python::attach(|py| { - let result = run_sync( - py, - async { - let gil_acquired = tokio::time::timeout( - Duration::from_secs(2), - tokio::task::spawn_blocking(|| Python::attach(|_| true)), - ) - .await; - Ok(matches!(gil_acquired, Ok(Ok(true)))) - }, - runtime_error, - ); - - assert!(extract_bool(py, result)); - }); - } - - #[test] - fn sync_runner_rejects_calls_from_a_tokio_context() { - Python::initialize(); - let runtime = Builder::new_current_thread() - .enable_all() - .build() - .expect("runtime should build"); - - let error = runtime.block_on(async { - Python::attach(|py| { - run_sync::(py, async { Ok(true) }, runtime_error) - .expect_err("sync route should reject a nested Tokio runtime") - }) - }); - - assert_eq!( - error.to_string(), - "RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route" - ); - } - - #[test] - fn sync_runner_can_drive_a_current_thread_runtime() { - Python::initialize(); - let runtime = Builder::new_current_thread() - .enable_all() - .build() - .expect("runtime should build"); - Python::attach(|py| { - let result = run_sync_on( - py, - &runtime, - async { - tokio::task::yield_now().await; - Ok(true) - }, - runtime_error, - ); - assert!(extract_bool(py, result)); - }); - } - - #[test] - fn sync_runner_maps_a_panicked_future() { - Python::initialize(); - Python::attach(|py| { - let error = run_sync::( - py, - poll_fn(|_| -> Poll> { panic!("route future panicked") }), - runtime_error, - ) - .expect_err("panicked route should become a Python exception"); - - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "PanicException: route future panicked"); - }); - } - - #[test] - fn sync_runner_maps_a_panicked_error_mapper() { - Python::initialize(); - Python::attach(|py| { - let error = run_sync::( - py, - async { Err(Error::InvalidRequest("invalid".to_string())) }, - panicking_error_mapper, - ) - .expect_err("panicked mapper should become a Python exception"); - - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "PanicException: error mapper panicked"); - }); - } - - #[test] - fn sync_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { - let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) - .expect_err("serializer panic should become a Python exception"); - - assert!(error.is_instance_of::(py)); - assert_eq!(error.to_string(), "PanicException: serializer panicked"); - }); - } - - #[test] - fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { - Python::initialize(); - let barrier = Arc::new(tokio::sync::Barrier::new(2)); - let callers: Vec<_> = (0..2) - .map(|_| { - let barrier = Arc::clone(&barrier); - thread::spawn(move || { - Python::attach(|py| { - extract_bool( - py, - run_sync( - py, - async move { - Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait()) - .await - .is_ok()) - }, - runtime_error, - ), - ) - }) - }) - }) - .collect(); - let results: Vec<_> = callers - .into_iter() - .map(|caller| caller.join().expect("caller should not panic")) - .collect(); - - assert_eq!(results, vec![true, true]); - } - - #[test] - fn async_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { - let module = PyModule::new(py, "runtime").expect("module should be created"); - module - .add_function( - wrap_pyfunction!(async_serialization_panic, &module) - .expect("function should wrap"), - ) - .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(): - try: - await runtime.async_serialization_panic() - except BaseException as error: - assert type(error).__name__ == "PanicException" - assert str(error) == "serializer panicked" - else: - raise AssertionError("serializer panic was not raised") - -asyncio.run(exercise()) -"#, - ) - .expect("Python source should not contain null bytes"); - py.run(&code, Some(&locals), Some(&locals)) - .expect("serializer panic should reach the Python awaiter"); - }); - } - - #[test] - fn async_result_delivery_does_not_stall_tokio_workers() { - Python::initialize(); - ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); - Python::attach(|py| { - let module = PyModule::new(py, "runtime").expect("module should be created"); - for function in [ - wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"), - wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"), - wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"), - ] { - module - .add_function(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(): - worker_count = runtime.runtime_worker_count() - awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)] - assert runtime.runtime_is_responsive(worker_count) - assert await asyncio.gather(*awaitables) == [True] * worker_count - -asyncio.run(exercise()) -"#, - ) - .expect("Python source should not contain null bytes"); - py.run(&code, Some(&locals), Some(&locals)) - .expect("result delivery should leave Tokio workers responsive"); - }); - } -} diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 71a598a6fe7..2f847c46b4e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2224,35 +2224,55 @@ class BaseLLMHTTPHandler: }, ) - rust_messages_response: Final = await self._maybe_rust_anthropic_messages( - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - has_agentic_hook=self._has_agentic_completion_hook(logging_obj), - model=model, - api_key=api_key, - api_base=api_base, - headers=headers, - request_body=request_body, - timeout=self._resolve_anthropic_messages_timeout( + if stream: + # Streaming goes through the native SSE-frame route only; on any + # fallback (disabled, unavailable, bridge failure) the Python path + # below issues the one and only provider call. + rust_streaming_response: Final = await self._maybe_rust_anthropic_messages_stream( + custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, - stream=stream or False, - custom_llm_provider=custom_llm_provider, - ), - ) - if rust_messages_response is not None: - if stream: - return self._rust_anthropic_messages_fake_stream(rust_messages_response) - return await self._finalize_anthropic_messages_response( - initial_response=rust_messages_response, + has_agentic_hook=self._has_agentic_completion_hook(logging_obj), model=model, - messages=messages, - anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, api_key=api_key, - kwargs=kwargs, + api_base=api_base, + headers=headers, + request_body=request_body, + timeout=self._resolve_anthropic_messages_timeout( + litellm_params=litellm_params, + stream=True, + custom_llm_provider=custom_llm_provider, + ), ) + if rust_streaming_response is not None: + return rust_streaming_response + else: + rust_messages_response: Final = await self._maybe_rust_anthropic_messages( + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + has_agentic_hook=self._has_agentic_completion_hook(logging_obj), + model=model, + api_key=api_key, + api_base=api_base, + headers=headers, + request_body=request_body, + timeout=self._resolve_anthropic_messages_timeout( + litellm_params=litellm_params, + stream=False, + custom_llm_provider=custom_llm_provider, + ), + ) + if rust_messages_response is not None: + return await self._finalize_anthropic_messages_response( + initial_response=rust_messages_response, + model=model, + messages=messages, + anthropic_messages_provider_config=anthropic_messages_provider_config, + anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + api_key=api_key, + kwargs=kwargs, + ) response: Final = await self._async_post_anthropic_messages_with_http_error_retry( async_httpx_client=async_httpx_client, @@ -2433,21 +2453,65 @@ class BaseLLMHTTPHandler: return response_obj @staticmethod - def _rust_anthropic_messages_fake_stream( - rust_response: AnthropicMessagesResponse, - ) -> "AnthropicMessagesStreamingResponse": - from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( - FakeAnthropicMessagesStreamIterator, - ) + async def _maybe_rust_anthropic_messages_stream( + *, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + has_agentic_hook: bool, + model: str, + api_key: str | None, + api_base: str | None, + headers: dict, + request_body: dict, + timeout: float | httpx.Timeout | None, + ) -> "AnthropicMessagesStreamingResponse | None": + """Native streaming path: SSE frames flow through as they arrive. + + Returns ``None`` (host falls back to the Python streaming path) when + the bridge is disabled, unavailable, raises, or an agentic hook needs + the chunk-buffering wrapper the native route does not provide. + """ + if custom_llm_provider not in ("azure_ai", "anthropic"): + return None + from litellm.rust_bridge.configuration import rust_enabled + + raw_request_override: Final = litellm_params.get("rust") + request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None + if not rust_enabled(request_override=request_override): + return None + if has_agentic_hook: + return None + + from litellm.rust_bridge import messages_stream as rust_messages_stream_bridge + + # `stream` stays in the body here: the native route speaks SSE. + try: + rust_stream: Final = await rust_messages_stream_bridge.amessages_stream( + model=model, + body=request_body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=headers, + timeout=timeout, + ) + except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path + verbose_logger.debug( + "Rust Anthropic messages stream bridge raised %s; falling back to Python path", + type(rust_error).__name__, + ) + return None + if rust_stream is None: + return None + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( AnthropicMessagesStreamHiddenParams, AnthropicMessagesStreamingResponse, ) - completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response)) hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"}) return AnthropicMessagesStreamingResponse( - completion_stream=completion_stream, + completion_stream=rust_stream, hidden_params=hidden_params, ) diff --git a/litellm/rust_bridge/messages_stream.py b/litellm/rust_bridge/messages_stream.py new file mode 100644 index 00000000000..2cb2b24131c --- /dev/null +++ b/litellm/rust_bridge/messages_stream.py @@ -0,0 +1,125 @@ +"""Thin Python wrapper for the native Rust Anthropic Messages streaming bridge.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Awaitable +from dataclasses import dataclass +from types import TracebackType +from typing import Final, Protocol + +import httpx + +from litellm.rust_bridge.timeouts import timeout_to_seconds + + +class RustAmessagesStream(Protocol): + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> Awaitable[object]: + raise NotImplementedError + + +class RustMessagesStreamAdapter: + """Adapts the native SSE-frame async iterator to ``AsyncIterator[bytes]``. + + ``aclose`` drops the native stream so the Rust side releases the upstream + response body (aborting the provider request) instead of waiting for GC. + """ + + def __init__(self, native_stream: object) -> None: + self._native: object | None = native_stream + + def __aiter__(self) -> RustMessagesStreamAdapter: + return self + + async def __anext__(self) -> bytes: + native = self._native + if native is None: + raise StopAsyncIteration + return await native.__anext__() # type: ignore[attr-defined,no-any-return] + + async def aclose(self) -> None: + self._native = None + + async def __aenter__(self) -> RustMessagesStreamAdapter: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + await self.aclose() + + +class _Unset: + pass + + +_UNSET: Final[_Unset] = _Unset() + + +@dataclass(slots=True) +class _RustMessagesStreamState: + amessages_stream: RustAmessagesStream | None = None + + +_STATE: Final[_RustMessagesStreamState] = _RustMessagesStreamState() + + +def set_rust_amessages_stream( + *, + amessages_stream: RustAmessagesStream | None | _Unset = _UNSET, +) -> None: + if not isinstance(amessages_stream, _Unset): + _STATE.amessages_stream = amessages_stream + + +def load_rust_amessages_stream() -> RustAmessagesStream | None: + if _STATE.amessages_stream is not None: + return _STATE.amessages_stream + from litellm.rust_bridge import get_native_bridge + + native_bridge: Final = get_native_bridge() + if native_bridge is None: + return None + return getattr(native_bridge, "amessages_stream", None) # type: ignore[return-value] + + +async def amessages_stream( + *, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout: float | httpx.Timeout | None, +) -> AsyncIterator[bytes] | None: + """Start a native streaming call; returns an SSE-frame iterator, or ``None``. + + Unlike the non-streaming bridge, ``stream`` stays in the body: the native + route forces it upstream. Each yielded item is one complete SSE frame + (``event: ...\\ndata: {...}\\n\\n``) as ``bytes``. + """ + rust_amessages_stream: Final = load_rust_amessages_stream() + if rust_amessages_stream is None: + return None + native_stream: Final = await rust_amessages_stream( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + ) + return RustMessagesStreamAdapter(native_stream) diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 293f75b7592..705c755be81 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -1,7 +1,7 @@ """Tests for the optional Rust-backed Anthropic Messages path.""" import importlib -from typing import cast +from typing import AsyncIterator import httpx import pytest @@ -9,12 +9,10 @@ import pytest import litellm from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.rust_bridge import configuration -from litellm.types.llms.anthropic_messages.anthropic_response import ( - AnthropicMessagesResponse, -) from litellm.types.router import GenericLiteLLMParams rust_messages = importlib.import_module("litellm.rust_bridge.messages") +rust_messages_stream = importlib.import_module("litellm.rust_bridge.messages_stream") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") FAKE_MESSAGES_RESPONSE: dict[str, object] = { @@ -108,13 +106,86 @@ class RaisingAsyncMessages: raise RuntimeError("upstream request failed with status 400: bad request") +FAKE_SSE_FRAMES: tuple[bytes, ...] = ( + b'event: message_start\ndata: {"type":"message_start"}\n\n', + b'event: content_block_delta\ndata: {"delta":"hello world"}\n\n', + b'event: message_stop\ndata: {"type":"message_stop"}\n\n', +) + + +class FakeNativeStream: + """Stands in for the native `MessagesStream` async iterator.""" + + def __init__(self, frames: tuple[bytes, ...] = FAKE_SSE_FRAMES) -> None: + self._frames = frames + self._index = 0 + + def __aiter__(self) -> "FakeNativeStream": + return self + + async def __anext__(self) -> bytes: + if self._index >= len(self._frames): + raise StopAsyncIteration + frame = self._frames[self._index] + self._index += 1 + return frame + + +class RecordingAsyncMessagesStream: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> object: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return FakeNativeStream() + + +class ExplodingAsyncMessagesStream: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> object: + self.calls += 1 + raise AssertionError("stream bridge must not be called") + + +class RaisingAsyncMessagesStream: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> object: + self.calls += 1 + raise RuntimeError("upstream request failed with status 429: rate limited") + + @pytest.fixture(autouse=True) def _reset_rust_flag(): rust_messages.set_rust_messages(messages=None, amessages=None) + rust_messages_stream.set_rust_amessages_stream(amessages_stream=None) configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield rust_messages.set_rust_messages(messages=None, amessages=None) + rust_messages_stream.set_rust_amessages_stream(amessages_stream=None) configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -375,20 +446,149 @@ async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): assert bridge.calls[0]["body"] == REQUEST_BODY +def _stream_gate(**overrides): + kwargs = { + "custom_llm_provider": "anthropic", + "litellm_params": GenericLiteLLMParams(api_key="sk-ant", rust=True), + "has_agentic_hook": False, + "model": "claude-sonnet-4-5", + "api_key": "sk-ant", + "api_base": "https://api.anthropic.com", + "headers": {"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"}, + "request_body": {**REQUEST_BODY, "stream": True}, + "timeout": 30.0, + } + kwargs.update(overrides) + return BaseLLMHTTPHandler._maybe_rust_anthropic_messages_stream(**kwargs) + + @pytest.mark.asyncio -async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): - response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)) - stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response) +async def test_stream_gate_returns_sse_frames_and_keeps_stream_flag(): + bridge = RecordingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) - assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + response = await _stream_gate() - chunks = [chunk async for chunk in stream] - joined = b"".join(chunks) + assert response is not None + assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + chunks = [chunk async for chunk in response] + assert chunks == list(FAKE_SSE_FRAMES) + call = bridge.calls[0] + assert call["model"] == "claude-sonnet-4-5" + assert call["body"]["stream"] is True + assert call["timeout_seconds"] == 30.0 - assert b"event: message_start" in joined - assert b"event: content_block_delta" in joined - assert b"hello world" in joined - assert b"event: message_stop" in joined + +@pytest.mark.asyncio +async def test_stream_gate_falls_back_to_python_when_bridge_raises(): + bridge = RaisingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) + + response = await _stream_gate() + + assert response is None + assert bridge.calls == 1 + + +@pytest.mark.asyncio +async def test_stream_gate_skips_rust_when_flag_absent(): + bridge = ExplodingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) + + response = await _stream_gate( + litellm_params=GenericLiteLLMParams(api_key="sk-ant"), + ) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_stream_gate_skips_rust_for_agentic_hook(): + bridge = ExplodingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) + + response = await _stream_gate(has_agentic_hook=True) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_stream_gate_skips_rust_for_unsupported_provider(): + bridge = ExplodingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) + + response = await _stream_gate(custom_llm_provider="openai") + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_stream_gate_falls_back_when_bridge_unavailable(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + + response = await _stream_gate() + + assert response is None + + +@pytest.mark.asyncio +async def test_messages_stream_wrapper_forwards_args_and_adapts_frames(): + bridge = RecordingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) + + stream = await rust_messages_stream.amessages_stream( + model="claude-sonnet-4-5", + body={**REQUEST_BODY, "stream": True}, + api_key="sk-ant", + api_base="https://api.anthropic.com", + custom_llm_provider="anthropic", + extra_headers=None, + timeout=httpx.Timeout(600.0, read=42.0), + ) + + assert stream is not None + assert [chunk async for chunk in stream] == list(FAKE_SSE_FRAMES) + assert bridge.calls[0]["timeout_seconds"] == 42.0 + assert bridge.calls[0]["body"]["stream"] is True + + +@pytest.mark.asyncio +async def test_messages_stream_adapter_aclose_ends_iteration(): + adapter = rust_messages_stream.RustMessagesStreamAdapter(FakeNativeStream()) + assert await adapter.__anext__() == FAKE_SSE_FRAMES[0] + await adapter.aclose() + with pytest.raises(StopAsyncIteration): + await adapter.__anext__() + + +@pytest.mark.asyncio +async def test_messages_stream_wrapper_returns_none_when_bridge_absent(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + + stream = await rust_messages_stream.amessages_stream( + model="claude-sonnet-4-5", + body={**REQUEST_BODY, "stream": True}, + api_key="sk-ant", + api_base="https://api.anthropic.com", + custom_llm_provider="anthropic", + extra_headers=None, + timeout=30.0, + ) + + assert stream is None @pytest.mark.asyncio