From 3d359ba954b91bb524a9c1d97c4a9ef65aa4a435 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 5 Sep 2026 17:11:08 -0700 Subject: [PATCH] refactor(native): centralize typed route eligibility --- .../crates/ai-gateway/src/io/responses_ws.rs | 5 +- .../crates/core/src/chat_completions/mod.rs | 45 ++++++++- .../core/src/chat_completions/prepare.rs | 19 +++- .../crates/core/src/chat_completions/tests.rs | 73 +++++++------- .../src/chat_completions/transformation.rs | 2 - litellm-rust/crates/core/src/eligibility.rs | 99 +++++++++++++++++++ litellm-rust/crates/core/src/lib.rs | 1 + .../anthropic/chat_completions/tests.rs | 7 +- .../chat_completions/transformation.rs | 10 -- .../bedrock/chat_completions/tests.rs | 7 +- .../chat_completions/transformation.rs | 6 -- litellm-rust/crates/python-bridge/src/lib.rs | 25 ++--- .../src/routes/audio_transcription.rs | 34 ++----- .../src/routes/chat_completions.rs | 44 ++------- .../python-bridge/src/routes/definition.rs | 47 +++++---- .../python-bridge/src/routes/messages.rs | 30 ++---- .../crates/python-bridge/src/routes/ocr.rs | 33 ++----- litellm/rust_bridge/chat_completions.py | 85 ++++++++++------ litellm/rust_bridge/messages.py | 19 ++-- litellm/rust_bridge/protocols.py | 9 +- litellm/rust_bridge/runtime.py | 14 ++- litellm/rust_bridge/transcription.py | 21 ++-- .../test_rust_bridge_messages.py | 8 +- tests/test_litellm/ocr/test_rust_bridge.py | 9 +- .../responses/test_rust_bridge_websocket.py | 10 +- .../test_audio_transcription_rust_bridge.py | 11 +++ 26 files changed, 372 insertions(+), 301 deletions(-) create mode 100644 litellm-rust/crates/core/src/eligibility.rs diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 4c915063a13..ee907a17869 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -41,10 +41,7 @@ impl ResponsesWebSocketConnection { _context: &LiteLlmRequestContext, ) -> Result { if !litellm_core::responses::websocket::native_websocket_supported( - options - .custom_llm_provider - .as_deref() - .unwrap_or("openai"), + options.custom_llm_provider.as_deref().unwrap_or("openai"), ) { return Err(Error::Unsupported("unsupported native WebSocket provider")); } diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 7987542eb57..8d51f2826b6 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -7,6 +7,7 @@ //! calls the provider, and returns a typed OpenAI-shaped response. use crate::Error; +use crate::eligibility::native_route_decline; use crate::request_context::LiteLlmRequestContext; use crate::request_options::RequestOptions; mod client; @@ -22,7 +23,8 @@ use serde_json::{Map, Value}; use handler::execute_chat_completions_provider_call; use prepare::{parse_messages, resolve_provider_config, resolve_request}; -use types::{ChatCompletionsRequest, ChatCompletionsResponse}; +use transformation::{ChatCompletionsProviderConfig, Unsupported}; +use types::{ChatCompletionsRequest, ChatCompletionsResponse, ChatMessage}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn chat_completions( @@ -45,9 +47,10 @@ pub fn chat_completions_decline_reason( custom_llm_provider: Option<&str>, messages: Value, optional_params: &Map, + options: &RequestOptions, context: &LiteLlmRequestContext, ) -> Option<&'static str> { - let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else { + let Ok((_, provider, config)) = resolve_provider_config(model, custom_llm_provider) else { return Some("provider is not on the rust chat completions path"); }; let Ok(messages) = parse_messages(messages) else { @@ -56,9 +59,41 @@ pub fn chat_completions_decline_reason( if messages.is_empty() { return Some("empty message list"); } - config - .unsupported_reason(&messages, optional_params, context) - .map(|reason| reason.0) + unsupported_reason( + provider, + config, + &messages, + optional_params, + options, + context, + ) + .map(|reason| reason.0) +} + +fn unsupported_reason( + provider: &str, + config: &dyn ChatCompletionsProviderConfig, + messages: &[ChatMessage], + optional_params: &Map, + options: &RequestOptions, + context: &LiteLlmRequestContext, +) -> Option { + native_route_decline(true, &context.capabilities) + .map(|reason| Unsupported(reason.reason())) + .or_else(|| match provider { + "anthropic" => options + .anthropic + .as_ref() + .and_then(|anthropic| anthropic.user_id.as_ref()) + .map(|_| Unsupported("LiteLLM user metadata")), + "bedrock" => options + .bedrock + .as_ref() + .is_some_and(|bedrock| !bedrock.request_metadata_fields.is_empty()) + .then_some(Unsupported("LiteLLM request metadata forwarding")), + _ => None, + }) + .or_else(|| config.unsupported_reason(messages, optional_params)) } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index f6ef44c2c0c..2820c6ddb60 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -16,7 +16,7 @@ use super::types::{ pub(super) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, -) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> { +) -> Result<(String, &'a str, &'static dyn ChatCompletionsProviderConfig), Error> { let provider_info = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { custom_llm_provider.map(|provider| CustomLlmProvider { @@ -31,7 +31,11 @@ pub(super) fn resolve_provider_config<'a>( })?; let config = chat_completions_provider_config(provider_info.custom_llm_provider) .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; - Ok((provider_info.model.to_string(), config)) + Ok(( + provider_info.model.to_string(), + provider_info.custom_llm_provider, + config, + )) } pub(super) fn parse_messages(messages: Value) -> Result, Error> { @@ -44,7 +48,7 @@ pub(super) fn resolve_request( options: RequestOptions, context: &LiteLlmRequestContext, ) -> Result { - let (model, config) = + let (model, provider, config) = resolve_provider_config(request.model, options.custom_llm_provider.as_deref()) .map_err(|_| Error::Declined("provider is not on the rust chat completions path"))?; let messages = @@ -52,7 +56,14 @@ pub(super) fn resolve_request( if messages.is_empty() { return Err(Error::Declined("empty message list")); } - if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params, context) { + if let Some(reason) = super::unsupported_reason( + provider, + config, + &messages, + &request.optional_params, + &options, + context, + ) { return Err(Error::Declined(reason.0)); } Ok(ResolvedChatCompletionsRequest { diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index ec822c1f37b..00a5457743e 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -1,5 +1,5 @@ use crate::request_context::LiteLlmRequestContext; -use crate::request_options::{BedrockOptions, RequestOptions}; +use crate::request_options::{AnthropicOptions, BedrockOptions, RequestOptions}; use serde_json::{Map, Value, json}; use crate::error::Error; @@ -518,6 +518,7 @@ fn decline_reason( provider, messages, ¶ms, + &RequestOptions::default(), &LiteLlmRequestContext::default(), ) } @@ -888,37 +889,52 @@ mod round_trip { fn preflight_and_execution_share_provider_metadata_eligibility() { let messages = json!([{"role": "user", "content": "hi"}]); let cases = [ - ("anthropic", json!({"user_id": "u-123"}), vec![], true), - ("anthropic", json!({"user_id": null}), vec![], false), - ("anthropic", json!({"trace_id": "t-1"}), vec![], false), ( "anthropic", - json!({}), - vec!["user_api_key_team_id".into()], - false, + RequestOptions { + custom_llm_provider: Some("anthropic".into()), + anthropic: Some(AnthropicOptions { + user_id: Some("u-123".into()), + }), + ..Default::default() + }, + true, ), #[cfg(feature = "bedrock-auth")] - ("bedrock", json!({"user_id": "u-123"}), vec![], false), - #[cfg(feature = "bedrock-auth")] ( "bedrock", - json!({}), - vec!["user_api_key_team_id".into()], + RequestOptions { + custom_llm_provider: Some("bedrock".into()), + bedrock: Some(BedrockOptions { + request_metadata_fields: vec!["user_api_key_team_id".into()], + ..Default::default() + }), + ..Default::default() + }, true, ), + ( + "anthropic", + RequestOptions { + custom_llm_provider: Some("anthropic".into()), + bedrock: Some(BedrockOptions { + request_metadata_fields: vec!["user_api_key_team_id".into()], + ..Default::default() + }), + ..Default::default() + }, + false, + ), ]; - for (provider, metadata, request_metadata_fields, expected_decline) in cases { - let context = LiteLlmRequestContext { - metadata: metadata.as_object().cloned(), - request_metadata_fields, - ..Default::default() - }; + for (provider, options, expected_decline) in cases { + let context = LiteLlmRequestContext::default(); let params = Map::new(); let preflight = super::chat_completions_decline_reason( "claude-sonnet-4-5", Some(provider), messages.clone(), ¶ms, + &options, &context, ); let execution = resolve_request( @@ -926,11 +942,8 @@ fn preflight_and_execution_share_provider_metadata_eligibility() { model: "claude-sonnet-4-5", messages: messages.clone(), optional_params: params, - options: RequestOptions { - custom_llm_provider: Some(provider.into()), - ..Default::default() - }, }, + options, &context, ); assert_eq!( @@ -941,21 +954,3 @@ fn preflight_and_execution_share_provider_metadata_eligibility() { assert_eq!(execution.is_err(), expected_decline, "{provider} execution"); } } - -#[test] -fn anthropic_preflight_preserves_litellm_metadata_scope() { - let context = LiteLlmRequestContext { - litellm_metadata: json!({"user_id": "u-123"}).as_object().cloned(), - ..Default::default() - }; - assert_eq!( - super::chat_completions_decline_reason( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - &Map::new(), - &context, - ), - None, - ); -} diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index b71adb8a8cd..d7b9704c46c 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -1,5 +1,4 @@ use crate::Error; -use crate::request_context::LiteLlmRequestContext; use serde_json::{Map, Value}; use super::types::{ @@ -76,7 +75,6 @@ pub trait ChatCompletionsProviderConfig: Sync { &self, messages: &[ChatMessage], optional_params: &Map, - _context: &LiteLlmRequestContext, ) -> Option { unsupported_param( self.supported_openai_params(), diff --git a/litellm-rust/crates/core/src/eligibility.rs b/litellm-rust/crates/core/src/eligibility.rs new file mode 100644 index 00000000000..febd25d7e57 --- /dev/null +++ b/litellm-rust/crates/core/src/eligibility.rs @@ -0,0 +1,99 @@ +use crate::request_context::RequestCapabilities; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NativeRouteDecline { + UnsupportedProvider, + Streaming, + AgenticHook, + CustomClient, + NativeResponseFormat, +} + +impl NativeRouteDecline { + pub const fn reason(self) -> &'static str { + match self { + Self::UnsupportedProvider => "unsupported native provider", + Self::Streaming => "native streaming is unavailable", + Self::AgenticHook => "native agentic hooks are unavailable", + Self::CustomClient => "native custom clients are unavailable", + Self::NativeResponseFormat => "native OCR response format is unavailable", + } + } +} + +pub fn native_route_decline( + provider_supported: bool, + capabilities: &RequestCapabilities, +) -> Option { + if !provider_supported { + return Some(NativeRouteDecline::UnsupportedProvider); + } + if capabilities.stream { + return Some(NativeRouteDecline::Streaming); + } + if capabilities.has_agentic_hook { + return Some(NativeRouteDecline::AgenticHook); + } + if capabilities.has_custom_client { + return Some(NativeRouteDecline::CustomClient); + } + (capabilities.request_format.as_deref() == Some("native")) + .then_some(NativeRouteDecline::NativeResponseFormat) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn admission_preserves_precedence_and_accepts_supported_unary_calls() { + let all_unsupported = RequestCapabilities { + stream: true, + has_agentic_hook: true, + has_custom_client: true, + request_format: Some("native".into()), + }; + assert_eq!( + native_route_decline(false, &all_unsupported), + Some(NativeRouteDecline::UnsupportedProvider) + ); + assert_eq!( + native_route_decline(true, &all_unsupported), + Some(NativeRouteDecline::Streaming) + ); + assert_eq!( + native_route_decline(true, &RequestCapabilities::default()), + None + ); + } + + #[test] + fn admission_rejects_each_unsupported_capability() { + let cases = [ + ( + RequestCapabilities { + has_agentic_hook: true, + ..Default::default() + }, + NativeRouteDecline::AgenticHook, + ), + ( + RequestCapabilities { + has_custom_client: true, + ..Default::default() + }, + NativeRouteDecline::CustomClient, + ), + ( + RequestCapabilities { + request_format: Some("native".into()), + ..Default::default() + }, + NativeRouteDecline::NativeResponseFormat, + ), + ]; + for (capabilities, expected) in cases { + assert_eq!(native_route_decline(true, &capabilities), Some(expected)); + } + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 952dd94ba19..daf6ea89b79 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -3,6 +3,7 @@ pub mod caching; pub mod call_lifecycle; pub mod chat_completions; pub mod constants; +pub mod eligibility; pub mod error; pub mod http_utils; pub mod messages; diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs index 8142374182c..b22de6c47de 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::Error; -use crate::request_context::LiteLlmRequestContext; use serde_json::json; fn messages(value: Value) -> Vec { @@ -27,11 +26,7 @@ fn transform_response(body: Value) -> Result { } fn reason(msgs: Value, opts: Value) -> Option { - ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason( - &messages(msgs), - ¶ms(opts), - &LiteLlmRequestContext::default(), - ) + ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } #[test] diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index f4a1d043457..a7d5a8ad0cf 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -1,4 +1,3 @@ -use crate::request_context::LiteLlmRequestContext; use serde_json::{Map, Value, json}; use crate::chat_completions::conversation::{Conversation, build_conversation}; @@ -127,18 +126,9 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { &self, messages: &[ChatMessage], optional_params: &Map, - context: &LiteLlmRequestContext, ) -> Option { unsupported_param(self.supported_openai_params(), &[], optional_params) .or_else(|| messages.iter().find_map(unsupported_message)) - .or_else(|| { - context - .metadata - .as_ref() - .and_then(|metadata| metadata.get("user_id")) - .is_some_and(|value| !value.is_null()) - .then_some(Unsupported("LiteLLM user metadata")) - }) // Anthropic rejects a request whose first turn is not a user turn. // Python only repairs that under `litellm.modify_params`, which the // core cannot observe, so decline instead of guessing. diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs index 7746881de49..c86f061b9ca 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -1,6 +1,5 @@ use super::*; use crate::Error; -use crate::request_context::LiteLlmRequestContext; use serde_json::json; fn messages(value: Value) -> Vec { @@ -33,11 +32,7 @@ fn transform_response(body: Value) -> Result { } fn reason(msgs: Value, opts: Value) -> Option { - BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason( - &messages(msgs), - ¶ms(opts), - &LiteLlmRequestContext::default(), - ) + BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts)) } #[test] diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index 39e552f7b0a..7be3d108d44 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -1,4 +1,3 @@ -use crate::request_context::LiteLlmRequestContext; use serde_json::{Map, Value, json}; use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation}; @@ -177,7 +176,6 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { &self, messages: &[ChatMessage], optional_params: &Map, - context: &LiteLlmRequestContext, ) -> Option { unsupported_param( self.supported_openai_params(), @@ -185,10 +183,6 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { optional_params, ) .or_else(|| messages.iter().find_map(unsupported_message)) - .or_else(|| { - (!context.request_metadata_fields.is_empty()) - .then_some(Unsupported("LiteLLM request metadata forwarding")) - }) // Python's Converse translation drops blank text blocks instead of // substituting the placeholder the shared conversation builder // applies, so decline blank text rather than diverge. diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index a9f838d5bc6..5bca907a3cb 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -35,18 +35,14 @@ impl ResponsesWebSocketConnection { options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> { - if let Some(reason) = responses_websocket_decline( - "responses websocket", + let provider_supported = litellm_core::responses::websocket::native_websocket_supported( options.provider("openai"), - false, - false, - false, - None, - ) { + ); + let context: litellm_core::request_context::LiteLlmRequestContext = context.into(); + if let Some(reason) = routes::definition::request_decline(provider_supported, &context) { return Err(crate::errors::RustBridgeDeclined::new_err(reason)); } let options: litellm_core::request_options::RequestOptions = options.into(); - let context: litellm_core::request_context::LiteLlmRequestContext = context.into(); let request = ResponsesWebSocketRequest { url: request.url }; pyo3_async_runtimes::tokio::future_into_py(py, async move { let inner = RustResponsesWebSocketConnection::connect(request, &options, &context) @@ -79,21 +75,16 @@ impl ResponsesWebSocketConnection { } #[pyfunction] -#[pyo3(signature = (_model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))] +#[pyo3(signature = (_model, custom_llm_provider, *, context))] fn responses_websocket_decline( _model: &str, custom_llm_provider: &str, - stream: bool, - has_agentic_hook: bool, - has_custom_client: bool, - request_format: Option<&str>, + context: NativeRequestContext, ) -> Option { + let context: litellm_core::request_context::LiteLlmRequestContext = context.into(); routes::definition::request_decline( litellm_core::responses::websocket::native_websocket_supported(custom_llm_provider), - stream, - has_agentic_hook, - has_custom_client, - request_format, + &context, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 76e72e8bf50..068e7f66421 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -22,24 +22,13 @@ fn prepare_transcription( options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + Send + 'static> { - if let Some(reason) = transcription_decline( - &input.model, - input.options.provider("bedrock"), - input - .optional_params - .get("stream") - .and_then(Value::as_bool) - .unwrap_or(false), - false, - false, - input - .optional_params - .get("response_format") - .and_then(Value::as_str), - ) { + let provider_supported = litellm_core::audio_transcription::transcription_provider_supported( + options.provider("bedrock"), + ); + let context: LiteLlmRequestContext = context.into(); + if let Some(reason) = super::definition::request_decline(provider_supported, &context) { return Err(crate::errors::RustBridgeDeclined::new_err(reason)); } - let context: LiteLlmRequestContext = context.into(); let audio = input.audio; Ok(async move { run_route( @@ -56,21 +45,16 @@ fn prepare_transcription( } #[pyfunction] -#[pyo3(signature = (_model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))] +#[pyo3(signature = (_model, custom_llm_provider, *, context))] fn transcription_decline( _model: &str, custom_llm_provider: &str, - stream: bool, - has_agentic_hook: bool, - has_custom_client: bool, - request_format: Option<&str>, + context: NativeRequestContext, ) -> Option { + let context: LiteLlmRequestContext = context.into(); super::definition::request_decline( litellm_core::audio_transcription::transcription_provider_supported(custom_llm_provider), - stream, - has_agentic_hook, - has_custom_client, - request_format, + &context, ) } 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 9b48ef2cfd7..f1a13b2ee46 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -5,6 +5,7 @@ use litellm_core::chat_completions::chat_completions as run_route; use litellm_core::chat_completions::chat_completions_decline_reason; use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; use litellm_core::request_context::LiteLlmRequestContext; +use litellm_core::request_options::RequestOptions; use pyo3::prelude::*; use serde_json::{Map, Value}; use std::future::Future; @@ -39,52 +40,22 @@ fn prepare_chat_completions( }) } -fn preflight_context(context: &Bound<'_, PyAny>) -> PyResult { - let metadata = context.getattr("metadata")?; - let user_id = if metadata.is_none() { - None - } else { - match metadata.get_item("user_id") { - Ok(value) => Some(if value.is_none() { - Value::Null - } else { - Value::Bool(true) - }), - Err(error) if error.is_instance_of::(context.py()) => { - None - } - Err(error) => return Err(error), - } - }; - Ok(LiteLlmRequestContext { - metadata: user_id.map(|value| Map::from_iter([("user_id".into(), value)])), - request_metadata_fields: context.getattr("request_metadata_fields")?.extract()?, - ..Default::default() - }) -} - #[pyfunction] -#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None, *, context, stream=false, has_custom_client=false, has_agentic_hook=false))] +#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None, *, options, context))] #[allow( clippy::too_many_arguments, - reason = "PyO3 preserves chat preflight arguments alongside request features" + reason = "PyO3 preserves chat preflight inputs alongside separated options and context" )] fn chat_completions_decline( model: String, #[pyo3(from_py_with = litellm_python_interop::from_py)] messages: Value, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, custom_llm_provider: Option, - context: &Bound<'_, PyAny>, - stream: bool, - has_custom_client: bool, - has_agentic_hook: bool, + options: NativeRequestOptions, + context: NativeRequestContext, ) -> PyResult> { - if let Some(reason) = - super::definition::request_decline(true, stream, has_agentic_hook, has_custom_client, None) - { - return Ok(Some(reason)); - } - let context = preflight_context(context)?; + let context: LiteLlmRequestContext = context.into(); + let options: RequestOptions = options.into(); let optional_params = match optional_params { None | Some(Value::Null) => Map::new(), Some(Value::Object(params)) => params, @@ -99,6 +70,7 @@ fn chat_completions_decline( custom_llm_provider.as_deref(), messages, &optional_params, + &options, &context, ) .map(str::to_string)) diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index f8770a7754d..899969b5787 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -100,25 +100,10 @@ pub(super) fn add_function( pub(crate) fn request_decline( provider_supported: bool, - stream: bool, - has_agentic_hook: bool, - has_custom_client: bool, - request_format: Option<&str>, + context: &litellm_core::request_context::LiteLlmRequestContext, ) -> Option { - let reason = if !provider_supported { - Some("unsupported native provider") - } else if stream { - Some("native streaming is unavailable") - } else if has_agentic_hook { - Some("native agentic hooks are unavailable") - } else if has_custom_client { - Some("native custom clients are unavailable") - } else if request_format == Some("native") { - Some("native OCR response format is unavailable") - } else { - None - }; - reason.map(str::to_string) + litellm_core::eligibility::native_route_decline(provider_supported, &context.capabilities) + .map(|reason| reason.reason().to_string()) } #[cfg(test)] @@ -310,16 +295,20 @@ for route, provider in ( ('responses_websocket', 'openai'), ): decline = getattr(routes, route + '_decline') - assert decline('model', provider) is None, route + assert decline('model', provider, context=context) is None, route for flag in ('stream', 'has_agentic_hook', 'has_custom_client'): - assert decline('model', provider, **{flag: True}) is not None, (route, flag) - reason = decline('model', 'unsupported-native-provider') + flagged_context = replace( + context, + capabilities=replace(context.capabilities, **{flag: True}), + ) + assert decline('model', provider, context=flagged_context) is not None, (route, flag) + reason = decline('model', 'unsupported-native-provider', context=context) assert reason is not None, route request = Request( messages=[], body={}, audio={}, document={}, optional_params={}, url='invalid-url-must-not-be-used', - options=Options(custom_llm_provider='unsupported-native-provider'), ) + unsupported_options = Options(custom_llm_provider='unsupported-native-provider') functions = ( (routes.ResponsesWebSocketConnection.connect,) if route == 'responses_websocket' @@ -327,14 +316,22 @@ for route, provider in ( ) for execute in functions: try: - execute(request, context=context) + execute(request, options=unsupported_options, context=context) except Exception as error: assert type(error).__name__ == 'RustBridgeDeclined', (route, error) assert str(error) == reason, (route, reason, error) else: raise AssertionError('unsupported request reached provider execution') -assert routes.ocr_decline('model', 'mistral', request_format='native') is not None -assert routes.ocr_decline('model', 'mistral', request_format='litellm') is None +native_context = replace( + context, + capabilities=replace(context.capabilities, request_format='native'), +) +litellm_context = replace( + context, + capabilities=replace(context.capabilities, request_format='litellm'), +) +assert routes.ocr_decline('model', 'mistral', context=native_context) is not None +assert routes.ocr_decline('model', 'mistral', context=litellm_context) is None ", Some(&locals), Some(&locals), diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index f323d005f30..2ac4f222a86 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -20,21 +20,12 @@ fn prepare_messages( options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + Send + 'static> { - if let Some(reason) = messages_decline( - &input.model, - input.options.provider("anthropic"), - input - .body - .get("stream") - .and_then(Value::as_bool) - .unwrap_or(false), - false, - false, - input.body.get("response_format").and_then(Value::as_str), - ) { + let provider_supported = + litellm_core::messages::messages_provider_supported(options.provider("anthropic")); + let context: LiteLlmRequestContext = context.into(); + if let Some(reason) = super::definition::request_decline(provider_supported, &context) { return Err(crate::errors::RustBridgeDeclined::new_err(reason)); } - let context: LiteLlmRequestContext = context.into(); let body = required_value("body", input.body, Value::is_object, "dict")?; Ok(async move { run_route( @@ -50,21 +41,16 @@ fn prepare_messages( } #[pyfunction] -#[pyo3(signature = (_model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))] +#[pyo3(signature = (_model, custom_llm_provider, *, context))] fn messages_decline( _model: &str, custom_llm_provider: &str, - stream: bool, - has_agentic_hook: bool, - has_custom_client: bool, - request_format: Option<&str>, + context: NativeRequestContext, ) -> Option { + let context: LiteLlmRequestContext = context.into(); super::definition::request_decline( litellm_core::messages::messages_provider_supported(custom_llm_provider), - stream, - has_agentic_hook, - has_custom_client, - request_format, + &context, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 436f7a60c91..139f0ba48b4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -23,24 +23,14 @@ fn prepare_ocr( options: NativeRequestOptions, context: NativeRequestContext, ) -> PyResult> + Send + 'static> { - if let Some(reason) = ocr_decline( + let provider_supported = litellm_ai_gateway::io::ocr::ocr_provider_supported( &input.model, - input.options.provider("mistral"), - input - .optional_params - .get("stream") - .and_then(Value::as_bool) - .unwrap_or(false), - false, - false, - input - .optional_params - .get("req_format") - .and_then(Value::as_str), - ) { + options.provider("mistral"), + ); + let context: LiteLlmRequestContext = context.into(); + if let Some(reason) = super::definition::request_decline(provider_supported, &context) { return Err(crate::errors::RustBridgeDeclined::new_err(reason)); } - let context: LiteLlmRequestContext = context.into(); let document = input.document; Ok(async move { run_route( @@ -61,21 +51,16 @@ fn prepare_ocr( } #[pyfunction] -#[pyo3(signature = (model, custom_llm_provider, *, stream=false, has_agentic_hook=false, has_custom_client=false, request_format=None))] +#[pyo3(signature = (model, custom_llm_provider, *, context))] fn ocr_decline( model: &str, custom_llm_provider: &str, - stream: bool, - has_agentic_hook: bool, - has_custom_client: bool, - request_format: Option<&str>, + context: NativeRequestContext, ) -> Option { + let context: LiteLlmRequestContext = context.into(); super::definition::request_decline( litellm_ai_gateway::io::ocr::ocr_provider_supported(model, custom_llm_provider), - stream, - has_agentic_hook, - has_custom_client, - request_format, + &context, ) } diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions.py index b3c749aff52..68a8100c0f7 100644 --- a/litellm/rust_bridge/chat_completions.py +++ b/litellm/rust_bridge/chat_completions.py @@ -15,12 +15,12 @@ from __future__ import annotations import json from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from contextlib import nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -40,9 +40,12 @@ from litellm.rust_bridge.request import ( NativeBedrockOptions, NativeChatCompletionsRequest, NativePreCallDetails, + NativeRequestCapabilities, NativeRequestContext, NativeRequestOptions, PreparedNativeCall, + anthropic_options, + bedrock_options, call_native, ) from litellm.rust_bridge.runtime import ( @@ -63,8 +66,6 @@ from litellm.types.utils import ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) - RUST_RESPONSE_HEADER: Final = "x-litellm-rust" @@ -146,15 +147,35 @@ def set_rust_chat_completions( _CHAT_PREFLIGHT.override(decline) -def _preflight_context(litellm_params: Mapping[str, object] | None) -> NativeRequestContext: - metadata: Final = litellm_params.get("metadata") if litellm_params is not None else None - try: - entries: Final = _LITELLM_METADATA_ADAPTER.validate_python(metadata) - except ValidationError: - return NativeRequestContext(request_metadata_fields=get_bedrock_request_metadata_fields()) +def _provider_eligibility_options( + provider: str | None, + litellm_params: Mapping[str, object] | None, + optional_params: Mapping[str, object], +) -> NativeRequestOptions: + bedrock: Final = ( + replace( + bedrock_options(optional_params), + request_metadata_fields=get_bedrock_request_metadata_fields(), + ) + if provider == "bedrock" + else None + ) + anthropic: Final = anthropic_options(litellm_params) if provider == "anthropic" else None + return NativeRequestOptions(custom_llm_provider=provider, bedrock=bedrock, anthropic=anthropic) + + +def _eligibility_context( + *, + stream: bool, + has_custom_client: bool = False, + has_agentic_hook: bool = False, +) -> NativeRequestContext: return NativeRequestContext( - metadata=entries, - request_metadata_fields=get_bedrock_request_metadata_fields(), + capabilities=NativeRequestCapabilities( + stream=stream, + has_custom_client=has_custom_client, + has_agentic_hook=has_agentic_hook, + ) ) @@ -180,8 +201,8 @@ def rust_chat_completions_accepts( messages=messages, optional_params=optional_params, custom_llm_provider=custom_llm_provider, - context=_preflight_context(litellm_params), - stream=bool(stream), + options=_provider_eligibility_options(custom_llm_provider, litellm_params, optional_params), + context=_eligibility_context(stream=bool(stream)), ), ) @@ -360,10 +381,12 @@ class _ChatOperation: messages=ctx.messages, optional_params=ctx.optional_params, custom_llm_provider=ctx.custom_llm_provider, - context=_preflight_context(ctx.litellm_params), - stream=bool(ctx.stream), - has_custom_client=ctx.client is not None or ctx.shared_session is not None, - has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging), + options=_provider_eligibility_options(ctx.custom_llm_provider, ctx.litellm_params, ctx.optional_params), + context=_eligibility_context( + stream=bool(ctx.stream), + has_custom_client=ctx.client is not None or ctx.shared_session is not None, + has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging), + ), ), ) @@ -413,23 +436,25 @@ class _ChatOperation: } ctx.logging.pre_call(input=ctx.messages, api_key=key, additional_args=log_details) self.pre_call_logged = True + provider_options: Final = _provider_eligibility_options(ctx.custom_llm_provider, ctx.litellm_params, params) return PreparedNativeCall( NativeChatCompletionsRequest( model=ctx.model, messages=ctx.messages, - optional_params=provider_request_params(params), - options=NativeRequestOptions( - api_key=key, - api_base=base, - custom_llm_provider=ctx.custom_llm_provider, - extra_headers=headers, - timeout_seconds=timeout_to_seconds( - float(ctx.timeout) if isinstance(ctx.timeout, str) else ctx.timeout - ), - provider_connection=provider_connection_params(params), - ), + optional_params=params, + ), + options=replace( + provider_options, + api_key=key, + api_base=base, + extra_headers=headers, + timeout_seconds=timeout_to_seconds(float(ctx.timeout) if isinstance(ctx.timeout, str) else ctx.timeout), + ), + context=_eligibility_context( + stream=bool(ctx.stream), + has_custom_client=ctx.client is not None or ctx.shared_session is not None, + has_agentic_hook=BaseLLMHTTPHandler.has_agentic_completion_hook(ctx.logging), ), - context=_preflight_context(ctx.litellm_params), ) def fallback(self) -> _CompletionDispatchResult: diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 8cc0b354912..fc33f9bf6a5 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -265,16 +265,17 @@ class _MessagesOperation: NativeMessagesRequest( model=self.model, body=request_body, - options=NativeRequestOptions( - api_key=self.api_key, - api_base=base, - custom_llm_provider=self.provider, - extra_headers=headers, - timeout_seconds=timeout_to_seconds( - BaseLLMHTTPHandler.resolve_anthropic_messages_timeout(self.params, False, self.provider) - ), + ), + options=NativeRequestOptions( + api_key=self.api_key, + api_base=base, + custom_llm_provider=self.provider, + extra_headers=headers, + timeout_seconds=timeout_to_seconds( + BaseLLMHTTPHandler.resolve_anthropic_messages_timeout(self.params, False, self.provider) ), - ) + ), + context=NativeRequestContext(), ) def fallback(self) -> MessagesResult: diff --git a/litellm/rust_bridge/protocols.py b/litellm/rust_bridge/protocols.py index a7c335fd555..07489679671 100644 --- a/litellm/rust_bridge/protocols.py +++ b/litellm/rust_bridge/protocols.py @@ -32,10 +32,8 @@ class RustChatCompletionsDecline(Protocol): optional_params: Mapping[str, object] | None, custom_llm_provider: str | None, *, + options: NativeRequestOptions, context: NativeRequestContext, - stream: bool, - has_custom_client: bool = False, - has_agentic_hook: bool = False, ) -> str | None: ... @@ -64,10 +62,7 @@ class RustRouteDecline(Protocol): model: str, custom_llm_provider: str, *, - stream: bool = False, - has_agentic_hook: bool = False, - has_custom_client: bool = False, - request_format: str | None = None, + context: NativeRequestContext, ) -> str | None: ... diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index 86d5faacbc3..d687fb6ef0a 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -13,6 +13,7 @@ from litellm.rust_bridge.bindings import ( native_exception_types, ) from litellm.rust_bridge.protocols import NativeModule, RustRouteDecline +from litellm.rust_bridge.request import NativeRequestCapabilities, NativeRequestContext BindingT = TypeVar("BindingT") SelectedT = TypeVar("SelectedT") @@ -518,13 +519,18 @@ def assess_route( has_custom_client: bool = False, request_format: str | None = None, ) -> PythonFallback | None: - return binding.assess( - check=lambda decline: decline( - model, - provider, + context: Final = NativeRequestContext( + capabilities=NativeRequestCapabilities( stream=stream, has_agentic_hook=has_agentic_hook, has_custom_client=has_custom_client, request_format=request_format, + ) + ) + return binding.assess( + check=lambda decline: decline( + model, + provider, + context=context, ), ) diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py index 6624da230db..3a10f60624d 100644 --- a/litellm/rust_bridge/transcription.py +++ b/litellm/rust_bridge/transcription.py @@ -219,16 +219,17 @@ class _TranscriptionOperation: NativeTranscriptionRequest( model=self.model, audio=audio, - optional_params=provider_request_params(self.optional_params), - options=NativeRequestOptions( - api_key=key, - api_base=base, - custom_llm_provider=self.provider, - extra_headers=self.headers, - timeout_seconds=timeout_to_seconds(self.timeout), - provider_connection=provider_connection_params(self.optional_params), - ), - ) + optional_params=self.optional_params, + ), + options=NativeRequestOptions( + api_key=key, + api_base=base, + custom_llm_provider=self.provider, + extra_headers=self.headers, + timeout_seconds=timeout_to_seconds(self.timeout), + bedrock=bedrock_options(self.optional_params), + ), + context=NativeRequestContext(), ) def fallback(self) -> TranscriptionResult: 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 039978c241f..dc5d29d78bc 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -108,10 +108,10 @@ def _reset_rust_flag(): configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL rust_messages.set_rust_messages( - decline=lambda model, custom_llm_provider, **features: ( + decline=lambda model, custom_llm_provider, *, context: ( "unsupported feature" - if any(features.get(key) for key in ("stream", "has_agentic_hook", "has_custom_client")) - or features.get("request_format") == "native" + if any(getattr(context.capabilities, key) for key in ("stream", "has_agentic_hook", "has_custom_client")) + or context.capabilities.request_format == "native" else None ) ) @@ -285,7 +285,7 @@ def test_public_messages_invalid_response_does_not_fallback(monkeypatch, respons python = PythonMessages() monkeypatch.setattr(module, "base_llm_http_handler", python) litellm.rust(True) - rust_messages.set_rust_messages(messages=lambda request, *, context, callback_adapter=None: response) + rust_messages.set_rust_messages(messages=lambda request, *, options, context: response) with pytest.raises(ValidationError): litellm.anthropic.messages.create( model="anthropic/test-model", diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index b6480983c15..9e0de8e6157 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -213,7 +213,14 @@ def _reset_rust_flag(): rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None) configuration.reset_rust_configuration() rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - rust_bridge.set_rust_ocr(decline=lambda model, custom_llm_provider, **features: "unsupported feature" if any(features.get(key) for key in ("stream", "has_agentic_hook", "has_custom_client")) or features.get("request_format") == "native" else None) + rust_bridge.set_rust_ocr( + decline=lambda model, custom_llm_provider, *, context: ( + "unsupported feature" + if any(getattr(context.capabilities, key) for key in ("stream", "has_agentic_hook", "has_custom_client")) + or context.capabilities.request_format == "native" + else None + ) + ) yield rust_bridge.set_rust_ocr(ocr=None, aocr=None, decline=None) configuration.reset_rust_configuration() diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index c278ecbe309..93d56f012c8 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -43,10 +43,10 @@ def reset_responses_websocket(): responses_websocket.set_rust_responses_websocket(connection=None, decline=None) configuration.reset_rust_configuration() responses_websocket.set_rust_responses_websocket( - decline=lambda model, custom_llm_provider, **features: ( + decline=lambda model, custom_llm_provider, *, context: ( "unsupported feature" - if any(features.get(key) for key in ("stream", "has_agentic_hook", "has_custom_client")) - or features.get("request_format") == "native" + if any(getattr(context.capabilities, key) for key in ("stream", "has_agentic_hook", "has_custom_client")) + or context.capabilities.request_format == "native" else None ) ) @@ -131,9 +131,9 @@ async def test_connection_dispatch_cleans_up_without_reconnecting(native, sessio class Native: @classmethod - async def connect(cls, request, *, context, callback_adapter=None): + async def connect(cls, request, *, options, context): connections.append("native") - assert request.options.custom_llm_provider == "azure" + assert options.custom_llm_provider == "azure" return native_socket @asynccontextmanager diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index 6b58859155a..c7864b1319a 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -10,6 +10,17 @@ from litellm.rust_bridge.runtime import Handled rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") +@pytest.fixture(autouse=True) +def reset_rust_transcription() -> None: + rust_bridge.configure_rust_transcription( + transcription=None, + atranscription=None, + decline=lambda model, custom_llm_provider, *, context: None, + ) + yield + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None, decline=None) + + class SyncBridge: def __init__(self) -> None: self.calls: list[dict[str, object]] = []