mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
feat(rust-bridge): stream Anthropic Messages SSE frames natively
Streaming requests through the Rust bridge previously buffered the full provider response and faked the SSE stream on the Python side, so the first chunk only arrived after the last. Add a native streaming path: - core: SseFrameStream reassembles upstream bytes into complete SSE frames (blank-line delimited, keep-alive frames dropped, trailing partial frame preserved); messages_stream_frames exposes it and forces stream:true upstream so a host cannot misuse the entrypoint. - bridge: amessages_stream returns a MessagesStream async iterator yielding one complete SSE frame as bytes per __anext__ (recv_text pattern: per-call future_into_py, pull-based backpressure). Cancelling __anext__ drops only the frame fetch; dropping the iterator aborts the upstream request; mid-stream failures map to exceptions then terminate iteration. trace=true is rejected (no terminal response). - python: rust_bridge.messages_stream wrapper + aclose-aware adapter; streaming in llm_http_handler now routes through the native stream gate when eligible (same rust gates + agentic-hook fallback) and the fake-stream path is removed. - delete the orphaned duplicate routes/runtime.rs (never declared).
This commit is contained in:
parent
1de960bce7
commit
490d7dd983
15 changed files with 1255 additions and 474 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -1435,6 +1435,8 @@ dependencies = [
|
|||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-types",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"rand 0.8.7",
|
||||
"reqwest",
|
||||
"serde",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Box<dyn Stream<Item = reqwest::Result<Bytes>> + Send>>,
|
||||
buffer: Vec<u8>,
|
||||
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<Box<dyn Stream<Item = reqwest::Result<Bytes>> + Send>>) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
buffer: Vec::new(),
|
||||
done: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for SseFrameStream {
|
||||
type Item = Result<Vec<u8>, Error>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
loop {
|
||||
if let Some(end) = find_frame_end(&self.buffer) {
|
||||
let frame: Vec<u8> = 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<usize> {
|
||||
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<reqwest::Result<Bytes>> = 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<Vec<u8>> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -45,12 +45,17 @@ pub(super) async fn execute_messages_provider_call(
|
|||
pub(super) async fn execute_messages_provider_stream(
|
||||
request: MessagesRequest<'_>,
|
||||
) -> Result<reqwest::Response, Error> {
|
||||
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 {
|
||||
|
|
|
|||
|
|
@ -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<reqwest::Re
|
|||
execute_messages_provider_stream(request).await
|
||||
}
|
||||
|
||||
/// The streaming variant that yields complete server-sent-event frames.
|
||||
///
|
||||
/// Unlike [`messages_stream`], which hands back the raw upstream response for a
|
||||
/// host to splice, this entrypoint keeps frame decoding in core and yields one
|
||||
/// complete `event:`/`data:` frame per item, reassembled across network chunks.
|
||||
pub async fn messages_stream_frames(request: MessagesRequest<'_>) -> Result<SseFrameStream, Error> {
|
||||
let response = execute_messages_provider_stream(request).await?;
|
||||
Ok(SseFrameStream::new(response))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
|
|||
|
|
@ -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, .. }));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -102,6 +102,8 @@ mod tests {
|
|||
"atranscription",
|
||||
"messages",
|
||||
"amessages",
|
||||
"amessages_stream",
|
||||
"MessagesStream",
|
||||
"chat_completions_decline",
|
||||
"chat_completions",
|
||||
"achat_completions",
|
||||
|
|
|
|||
520
litellm-rust/crates/python-bridge/src/routes/messages_stream.rs
Normal file
520
litellm-rust/crates/python-bridge/src/routes/messages_stream.rs
Normal file
|
|
@ -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<u8>);
|
||||
|
||||
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<Self::Output> {
|
||||
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<Mutex<Option<SseFrameStream>>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl MessagesStream {
|
||||
fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
slf
|
||||
}
|
||||
|
||||
fn __anext__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
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<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
#[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option<Value>,
|
||||
timeout_seconds: Option<f64>,
|
||||
trace: bool,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
// `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::<MessagesStream>()
|
||||
}
|
||||
|
||||
#[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::<usize>().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::<PyValueError>(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");
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<T, F>(
|
||||
py: Python<'_>,
|
||||
future: F,
|
||||
map_error: fn(Error) -> PyErr,
|
||||
) -> PyResult<Py<PyAny>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: Future<Output = Result<T, Error>> + Send + 'static,
|
||||
{
|
||||
run_sync_on(
|
||||
py,
|
||||
pyo3_async_runtimes::tokio::get_runtime(),
|
||||
future,
|
||||
map_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn run_sync_on<T, F>(
|
||||
py: Python<'_>,
|
||||
runtime: &Runtime,
|
||||
future: F,
|
||||
map_error: fn(Error) -> PyErr,
|
||||
) -> PyResult<Py<PyAny>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: Future<Output = Result<T, Error>> + 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<T, F>(
|
||||
py: Python<'_>,
|
||||
future: F,
|
||||
map_error: fn(Error) -> PyErr,
|
||||
) -> PyResult<Bound<'_, PyAny>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: Future<Output = Result<T, Error>> + 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<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
|
||||
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<T, F>(future: F) -> PyResult<Result<T, Error>>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>>,
|
||||
{
|
||||
AssertUnwindSafe(future)
|
||||
.catch_unwind()
|
||||
.await
|
||||
.map_err(panic_to_pyerr)
|
||||
}
|
||||
|
||||
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>>,
|
||||
{
|
||||
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<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
panic!("serializer panicked")
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
run_async(py, async { Ok(PanickingOutput) }, runtime_error)
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<Py<PyAny>>) -> 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::<bool, _>(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::<bool, _>(
|
||||
py,
|
||||
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
|
||||
runtime_error,
|
||||
)
|
||||
.expect_err("panicked route should become a Python exception");
|
||||
|
||||
assert!(error.is_instance_of::<PanicException>(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::<bool, _>(
|
||||
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::<PanicException>(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::<PanicException>(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");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
125
litellm/rust_bridge/messages_stream.py
Normal file
125
litellm/rust_bridge/messages_stream.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue