diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 934b2269a69..27eb1d3e656 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,10 +1,10 @@ -use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; use crate::error::Error; use crate::http_utils::http_request; use super::client::http_client; use super::common_utils::truncate_error_body; use super::prepare::prepare_provider_request; +use super::transformation::AnthropicMessagesProviderConfig; use super::types::{AnthropicMessagesResponse, MessagesRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] @@ -42,15 +42,25 @@ pub(super) async fn execute_messages_provider_call( request.config.transform_response(&request.model, response) } +/// The streaming-capability gate. A provider whose config does not opt into +/// streaming declines ([`Error::Unsupported`]) before any request goes out, so +/// hosts treat it as "fall back", not "fail". +pub(super) fn ensure_streaming_supported( + config: &dyn AnthropicMessagesProviderConfig, +) -> Result<(), Error> { + if config.supports_streaming() { + return Ok(()); + } + Err(Error::Unsupported( + "streaming messages is not supported for this provider", + )) +} + pub(super) async fn execute_messages_provider_stream( request: MessagesRequest<'_>, ) -> Result { 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(), - )); - } + ensure_streaming_supported(request.config)?; // 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() { diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index ec83d03f535..6e3d80aaf72 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -46,7 +46,6 @@ pub(super) fn prepare_provider_request( let url = config.complete_url(request.api_base, &model, &env_lookup)?; Ok(ProviderMessagesRequest { - provider: provider.to_string(), model, config, url, diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index d23b56c2868..1e13d3b2d5b 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -9,8 +9,12 @@ use crate::error::Error; use super::common_utils::{ has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, }; +use super::handler::ensure_streaming_supported; use super::messages; +use super::transformation::AnthropicMessagesProviderConfig; use super::types::MessagesRequest; +use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; +use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -60,6 +64,37 @@ fn provider_config_resolves_anthropic_and_azure_ai() { assert!(messages_provider_config("openai").is_none()); } +#[test] +fn streaming_gate_accepts_opted_in_providers_and_declines_the_rest() { + struct NonStreamingConfig; + + impl AnthropicMessagesProviderConfig for NonStreamingConfig { + fn complete_url( + &self, + _api_base: Option<&str>, + _model: &str, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + unreachable!("the streaming gate must not touch the config's URLs") + } + + fn resolve_api_key( + &self, + _api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + unreachable!("the streaming gate must not resolve credentials") + } + } + + assert!(ensure_streaming_supported(&ANTHROPIC_MESSAGES_CONFIG).is_ok()); + assert!(ensure_streaming_supported(&AZURE_ANTHROPIC_MESSAGES_CONFIG).is_ok()); + assert!(matches!( + ensure_streaming_supported(&NonStreamingConfig).expect_err("default config declines"), + Error::Unsupported(_) + )); +} + #[test] fn truncate_error_body_caps_long_payloads() { let body = "x".repeat(400); @@ -510,6 +545,97 @@ async fn messages_stream_frames_yields_complete_frames_and_forces_stream_flag() assert_eq!(sent_body["stream"], Value::Bool(true)); } +#[tokio::test] +async fn messages_stream_frames_streams_azure_ai_through_the_anthropic_endpoint() { + 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; + let body = concat!( + "event: message_start\ndata: {\"type\":\"message_start\"}\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"); + socket + .write_all(body.as_bytes()) + .await + .expect("writes body"); + request + }); + + let mut stream = messages_stream_frames(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "hi"}] + }), + api_key: Some("sk-azure"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("azure_ai"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("azure_ai 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: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + ] + ); + + let request = server.await.expect("server task completes"); + let (head, body) = request.split_once("\r\n\r\n").expect("has body"); + assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}"); + let head_lower = head.to_ascii_lowercase(); + assert!(head_lower.contains("x-api-key: sk-azure"), "{head}"); + 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_declines_unsupported_provider_before_the_call() { + use super::messages_stream_frames; + + let err = match messages_stream_frames(MessagesRequest { + model: "gpt-4o", + body: json!({"model": "gpt-4o", "max_tokens": 8, "messages": []}), + api_key: Some("sk"), + // Nothing listens here: a decline must return before any connection. + api_base: Some("http://127.0.0.1:1"), + custom_llm_provider: Some("openai"), + extra_headers: None, + timeout: Some(Duration::from_millis(50)), + }) + .await + { + Err(err) => err, + Ok(_) => panic!("unsupported provider must decline before the stream"), + }; + + assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); +} + #[tokio::test] async fn messages_stream_frames_maps_upstream_error_status() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index a5904c085a0..aea1e2cc491 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -45,6 +45,13 @@ pub trait AnthropicMessagesProviderConfig: Sync { ] } + /// Whether this provider serves the SSE streaming variant of the route. + /// Opt-in: a config that does not override this declines streaming so the + /// host falls back to its own implementation instead of failing. + fn supports_streaming(&self) -> bool { + false + } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn transform_request( &self, diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index b9f807c29fd..8fcaf1f3481 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -16,7 +16,6 @@ pub struct MessagesRequest<'a> { } pub(super) struct ProviderMessagesRequest { - pub(super) provider: String, pub(super) model: String, pub(super) config: &'static dyn AnthropicMessagesProviderConfig, pub(super) url: String, diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index f31b961e78a..f58782ee33a 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -68,6 +68,10 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { fn auth_strategy(&self) -> MessagesAuthStrategy { MessagesAuthStrategy::Header("x-api-key") } + + fn supports_streaming(&self) -> bool { + true + } } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index b8ca10461fb..a159056dc13 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -164,6 +164,10 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { self.anthropic.auth_strategy() } + fn supports_streaming(&self) -> bool { + self.anthropic.supports_streaming() + } + fn accepts_bearer_auth(&self) -> bool { true } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 914e2e1e033..13e000150a1 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -33,7 +33,7 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { /// Everything raised before the request goes out is safe for the host to retry /// on its own path; anything after it is not, because the provider has already /// done the work and billed for it. -pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { +pub(crate) fn fallback_error_to_pyerr(err: Error) -> PyErr { match err { Error::Unsupported(_) | Error::Auth(_) diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 68b7762cb10..297a2f7dff9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -8,7 +8,7 @@ use litellm_core::chat_completions::{ use pyo3::prelude::*; use serde_json::Value; -use crate::errors::chat_completions_error_to_pyerr; +use crate::errors::fallback_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; fn prepare_chat_completions( @@ -86,6 +86,6 @@ bridge_route! { timeout_seconds: Option, }, prepare = prepare_chat_completions, - errors = chat_completions_error_to_pyerr, + errors = fallback_error_to_pyerr, extra = [chat_completions_decline], } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs b/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs index fbe2873d3f8..7b882cd7193 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages_stream.rs @@ -17,7 +17,7 @@ use pyo3::types::{PyAny, PyBytes}; use serde_json::Value; use tokio::sync::Mutex; -use crate::errors::core_error_to_pyerr; +use crate::errors::{core_error_to_pyerr, fallback_error_to_pyerr}; use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; /// One complete SSE frame, converted to `bytes` on the attached thread that @@ -115,6 +115,10 @@ fn amessages_stream<'py>( extra_headers, timeout, } = options; + // Only the request start maps through the fallback contract: declines + // (unsupported provider, bad request) happen before the provider is + // called, so the host may retry on its own path. Failures surfaced by + // `__anext__` are mid-stream and keep `core_error_to_pyerr`. let frames = messages_stream_frames(MessagesRequest { model: &model, body, @@ -125,7 +129,7 @@ fn amessages_stream<'py>( timeout, }) .await - .map_err(core_error_to_pyerr)?; + .map_err(fallback_error_to_pyerr)?; Ok(MessagesStream { frames: Arc::new(Mutex::new(Some(frames))), }) @@ -187,12 +191,12 @@ mod tests { } /// An SSE upstream that sends `body` in chunks separated by short delays - /// so frames span network chunk boundaries. + /// so frames span network chunk boundaries. Returns the captured request. async fn streaming_upstream( listener: TcpListener, status_line: &'static str, body: &'static str, - ) { + ) -> String { let (mut socket, _) = listener.accept().await.expect("accepts request"); let request = read_http_request(&mut socket).await; assert!( @@ -217,6 +221,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(10)).await; socket.write_all(chunk).await.expect("writes body chunk"); } + request } fn register_stream_module(py: Python<'_>) -> Bound<'_, PyModule> { @@ -233,6 +238,18 @@ mod tests { locals .set_item("url", url) .expect("URL should enter Python locals"); + locals + .set_item( + "Declined", + py.get_type::(), + ) + .expect("declined exception should enter Python locals"); + locals + .set_item( + "UpstreamError", + py.get_type::(), + ) + .expect("upstream exception 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"); @@ -324,8 +341,9 @@ async def exercise(): api_base=url, custom_llm_provider="anthropic", ) - except RuntimeError as error: - assert "429" in str(error), str(error) + except UpstreamError as error: + assert not isinstance(error, Declined), "a provider 429 is not a decline" + assert error.args[0] == 429, error.args else: raise AssertionError("upstream error should surface from the awaitable") @@ -340,6 +358,100 @@ asyncio.run(exercise()) .expect("server task should not panic"); } + #[test] + fn stream_route_streams_azure_ai_through_the_anthropic_endpoint() { + 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-azure", + api_base=url, + custom_llm_provider="azure_ai", + ) + 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()) +"#, + ); + }); + + let request = runtime + .block_on(async { tokio::time::timeout(Duration::from_secs(5), server).await }) + .expect("server should finish") + .expect("server task should not panic"); + let head = request + .split_once("\r\n\r\n") + .expect("request should have a body") + .0 + .to_ascii_lowercase(); + assert!( + head.starts_with("post /anthropic/v1/messages "), + "azure_ai must stream through the anthropic endpoint: {head}" + ); + assert!(head.contains("x-api-key: sk-azure"), "{head}"); + } + + #[test] + fn stream_route_declines_unsupported_provider_before_any_request() { + Python::initialize(); + Python::attach(|py| { + let module = register_stream_module(py); + run_async_code( + py, + &module, + "http://127.0.0.1:1", + r#" +import asyncio + +async def exercise(): + try: + await runtime.amessages_stream( + model="gpt-4o", + body={"model": "gpt-4o", "max_tokens": 8, "messages": []}, + api_key="sk", + api_base=url, + custom_llm_provider="openai", + ) + except Declined as error: + assert not isinstance(error, UpstreamError), str(error) + else: + raise AssertionError("unsupported provider must decline from the awaitable") + +asyncio.run(exercise()) +"#, + ); + }); + } + #[test] fn stream_route_ends_iteration_after_a_mid_stream_failure() { Python::initialize(); 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 705c755be81..aa3aff333a0 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -479,6 +479,29 @@ async def test_stream_gate_returns_sse_frames_and_keeps_stream_flag(): assert call["timeout_seconds"] == 30.0 +@pytest.mark.asyncio +async def test_stream_gate_routes_azure_ai_to_the_native_stream(): + bridge = RecordingAsyncMessagesStream() + rust_messages_stream.set_rust_amessages_stream(amessages_stream=bridge) + + response = await _stream_gate( + custom_llm_provider="azure_ai", + litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=True), + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + headers={"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"}, + ) + + 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["custom_llm_provider"] == "azure_ai" + assert call["api_key"] == "sk-azure" + assert call["body"]["stream"] is True + + @pytest.mark.asyncio async def test_stream_gate_falls_back_to_python_when_bridge_raises(): bridge = RaisingAsyncMessagesStream()