From 38c85792b8ad74d8143b1b0cb52564cac0d17dcd Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 8 Sep 2026 15:02:39 -0700 Subject: [PATCH] wip --- litellm-rust/Cargo.lock | 22 ++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/core/AGENTS.md | 16 ++ litellm-rust/crates/core/Cargo.toml | 1 + .../core/src/audio_transcription/handler.rs | 43 +--- .../core/src/audio_transcription/lifecycle.rs | 8 +- .../core/src/audio_transcription/prepare.rs | 16 +- .../src/audio_transcription/transformation.rs | 15 ++ .../core/src/audio_transcription/types.rs | 18 +- .../core/src/chat_completions/common_utils.rs | 19 +- .../core/src/chat_completions/conversation.rs | 14 +- .../core/src/chat_completions/handler.rs | 66 +----- .../src/chat_completions/transformation.rs | 15 ++ .../integrations/custom_guardrail/types.rs | 14 +- .../src/integrations/custom_logger/mod.rs | 16 ++ .../src/integrations/custom_logger/types.rs | 45 +--- .../crates/core/src/lifecycle/terminal.rs | 2 +- .../crates/core/src/messages/common_utils.rs | 17 +- litellm-rust/crates/core/src/ocr/request.rs | 35 +-- .../crates/core/src/ocr/transformation.rs | 8 + .../core/src/providers/anthropic/auth.rs | 42 ++++ .../chat_completions/transformation.rs | 6 +- .../anthropic/messages/transformation.rs | 51 +---- .../core/src/providers/anthropic/mod.rs | 1 + .../providers/azure_ai/ocr/transformation.rs | 21 ++ .../providers/bedrock/audio_transcription.rs | 40 ++++ .../chat_completions/transformation.rs | 67 +++++- .../crates/core/src/providers/dispatch.rs | 216 ++++++++++++++++++ litellm-rust/crates/core/src/providers/mod.rs | 1 + .../crates/core/src/providers/openai/auth.rs | 11 + .../crates/core/src/providers/openai/mod.rs | 14 ++ .../openai/realtime/transformation.rs | 25 +- .../openai/responses/transformation.rs | 48 ++++ .../providers/vertex_ai/ocr/transformation.rs | 18 ++ .../crates/core/src/realtime/streaming.rs | 58 ++--- .../core/src/realtime/transformation.rs | 8 + .../crates/core/src/responses/types.rs | 26 +-- .../crates/core/src/responses/websocket.rs | 113 +++------ .../crates/core/tests/chat_completions.rs | 49 ++++ 39 files changed, 750 insertions(+), 456 deletions(-) create mode 100644 litellm-rust/crates/core/src/providers/anthropic/auth.rs create mode 100644 litellm-rust/crates/core/src/providers/dispatch.rs create mode 100644 litellm-rust/crates/core/src/providers/openai/auth.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f209c6f5e6e..1665951add1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1469,6 +1469,7 @@ dependencies = [ "rustls-native-certs", "serde", "serde_json", + "strum", "thiserror 2.0.19", "tokio", "tokio-tungstenite 0.24.0", @@ -2503,6 +2504,27 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "subtle" version = "2.6.1" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c610a5e4de4..75e70a20a56 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -38,6 +38,7 @@ rustls-native-certs = "0.8" serial_test = { version = "4.0.1", default-features = false } serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip", "preserve_order"] } +strum = { version = "0.28", features = ["derive"] } sha2 = "0.10" subtle = "2" thiserror = "2.0" diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index b16689cdabc..64d14e43b19 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -68,6 +68,22 @@ The exact spelling may be a method on `LiteLlm`. Public adapters may wrap that function but can never reimplement admission, callbacks, provider preparation or transport around it. +## Provider and route ownership + +`providers/dispatch.rs` selects a typed adapter for each supported provider and +route pair, including model-specific OCR variants. Routes own their contracts +and lifecycle sequencing; provider adapters own admission policy, URLs, +transformation and authorization. Keep provider selection out of handlers + +Provider modules share credential and protocol helpers across their route +adapters. Routes supply the exact settled bytes to the adapter's authorization +operation at the existing lifecycle phase. Shared helpers must never invoke +another public route or repeat its callbacks + +WebSocket execution carries its selected adapter through dialing and event +transformation. The existing Responses WebSocket entrypoint defaults to OpenAI; +realtime resolves its existing optional provider prefix through dispatch + ## Services, not a context Capabilities are supplied through focused trait implementations. Route-specific diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 0697b11a410..e78e29c4e7f 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,6 +15,7 @@ rustls.workspace = true rustls-native-certs.workspace = true serde.workspace = true serde_json.workspace = true +strum.workspace = true thiserror.workspace = true tracing.workspace = true tokio = { workspace = true, features = ["rt", "sync", "time"] } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index bf36a6b9fb7..1627b7c67d7 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -42,50 +42,9 @@ pub(super) async fn execute_audio_transcription_provider_call( .into_json()) } -#[cfg(feature = "bedrock-auth")] async fn signed_headers( request: &ProviderAudioTranscriptionRequest, body: &[u8], ) -> Result, Error> { - use std::collections::BTreeMap; - use std::time::SystemTime; - - use crate::audio_transcription::transformation::AudioTranscriptionAuth; - use crate::providers::bedrock::audio_transcription::aws_auth_config; - use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; - - let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else { - return Ok(request.upstream_headers.clone()); - }; - let env_lookup = |key: &str| std::env::var(key).ok(); - let credentials = resolve_credentials( - aws_auth_config(&request.optional_params, &env_lookup), - &env_lookup, - ) - .await?; - let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); - let signature = sign_bedrock_post( - &request.url, - body, - &unsigned, - region, - &credentials, - SystemTime::now(), - )?; - Ok(unsigned.into_iter().chain(signature).collect()) -} - -#[cfg(not(feature = "bedrock-auth"))] -async fn signed_headers( - request: &ProviderAudioTranscriptionRequest, - _body: &[u8], -) -> Result, Error> { - use crate::audio_transcription::transformation::AudioTranscriptionAuth; - - match request.auth { - AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported( - "AWS SigV4 requires the bedrock-auth feature", - )), - AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()), - } + request.config.authorize(request, body).await } diff --git a/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs b/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs index b073e3a6c4c..29f1fed879c 100644 --- a/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs +++ b/litellm-rust/crates/core/src/audio_transcription/lifecycle.rs @@ -16,7 +16,7 @@ use crate::lifecycle::{ ActionResult, CallLifecycle, CallLifecycleContext, Clock, ExecutedCall, RequestPolicy, TerminalDispatcher, TerminalRecord, }; -use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; +use crate::providers::dispatch::resolve_audio_route_provider; use super::handler::execute_audio_transcription_provider_call; use super::prepare::prepare_audio_transcription_provider_call; @@ -70,11 +70,7 @@ impl AudioRoute { services: &S, request: AudioRouteRequest<'_>, ) -> ExecutedCall { - let provider = get_custom_llm_provider(request.model, request.custom_llm_provider) - .unwrap_or(CustomLlmProvider { - model: request.model, - custom_llm_provider: "bedrock", - }); + let provider = resolve_audio_route_provider(request.model, request.custom_llm_provider); let context = CallLifecycleContext::new( "audio_transcription", provider.model, diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 25139da1d1c..4f6ab546bb2 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,22 +1,12 @@ +use crate::providers::dispatch::audio_transcription_provider_config as provider_config; + use crate::error::Error; use crate::http_utils::{has_header, string_headers}; -#[cfg(feature = "bedrock-auth")] -use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; -use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; +use super::transformation::AudioTranscriptionAuth; use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { - #[cfg(feature = "bedrock-auth")] - if provider == "bedrock" { - return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG); - } - let _ = provider; - None -} - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub(super) fn prepare_audio_transcription_provider_call( request: AudioTranscriptionRequest<'_>, diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index aa9846427dc..5581c4bf239 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -13,6 +13,21 @@ pub enum AudioTranscriptionAuth { } pub trait AudioTranscriptionProviderConfig: Sync { + fn authorize<'a>( + &'a self, + request: &'a super::types::ProviderAudioTranscriptionRequest, + _body: &'a [u8], + ) -> crate::providers::AuthorizationFuture<'a> { + Box::pin(async move { + match &request.auth { + AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported( + "AWS SigV4 requires the bedrock-auth feature", + )), + _ => Ok(request.upstream_headers.clone()), + } + }) + } + fn supported_transcription_params(&self) -> &'static [&'static str]; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index 5b8500275a3..7a7d6055d0e 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -33,16 +33,16 @@ pub struct AudioRouteRequest<'a> { #[derive(Clone)] pub struct ProviderAudioTranscriptionRequest { - pub(super) model: String, - pub(super) custom_llm_provider: String, - pub(super) config: &'static dyn AudioTranscriptionProviderConfig, - pub(super) url: String, - pub(super) body: Value, - pub(super) upstream_headers: Vec<(String, String)>, - pub(super) auth: AudioTranscriptionAuth, + pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) config: &'static dyn AudioTranscriptionProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) auth: AudioTranscriptionAuth, #[cfg(feature = "bedrock-auth")] - pub(super) optional_params: Map, - pub(super) timeout: Option, + pub(crate) optional_params: Map, + pub(crate) timeout: Option, } impl ProviderAudioTranscriptionRequest { diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index 69e5f175ad5..c97199e9922 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -1,26 +1,11 @@ +pub(crate) use crate::providers::dispatch::chat_completions_provider_config; + use crate::Error; use crate::http_utils::string_headers as shared_string_headers; -use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; use serde_json::{Map, Value}; -use super::transformation::ChatCompletionsProviderConfig; - const HEADER_CONTEXT: &str = "chat completions"; -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub(super) fn chat_completions_provider_config( - provider: &str, -) -> Option<&'static dyn ChatCompletionsProviderConfig> { - match provider { - "anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG), - #[cfg(feature = "bedrock-auth")] - "bedrock" => Some( - &crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, - ), - _ => None, - } -} - pub(super) fn string_headers( extra_headers: Option>, ) -> Result, Error> { diff --git a/litellm-rust/crates/core/src/chat_completions/conversation.rs b/litellm-rust/crates/core/src/chat_completions/conversation.rs index f7bdc60af37..cafbd3b9f84 100644 --- a/litellm-rust/crates/core/src/chat_completions/conversation.rs +++ b/litellm-rust/crates/core/src/chat_completions/conversation.rs @@ -11,24 +11,18 @@ //! accepts; anything richer is declined upstream by the capability gate. use crate::constants::EMPTY_TEXT_PLACEHOLDER; +use strum::AsRefStr; use super::types::{ChatMessage, ChatMessageContent}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(AsRefStr, Clone, Copy, Debug, PartialEq, Eq)] pub enum TurnRole { + #[strum(serialize = "user")] User, + #[strum(serialize = "assistant")] Assistant, } -impl TurnRole { - pub fn as_str(self) -> &'static str { - match self { - Self::User => "user", - Self::Assistant => "assistant", - } - } -} - #[derive(Clone, Debug, PartialEq, Eq)] pub struct Turn { pub role: TurnRole, diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 301cd5b81f9..b1070f4a12a 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -5,7 +5,6 @@ use crate::http_utils::{http_request, truncate_error_body}; use super::client::http_client; use super::request::build_provider_request; -use super::transformation::ChatCompletionsAuth; use super::types::{ ChatBodySnapshot, ChatCompletionsResponse, ChatEndpoint, ProviderChatCompletionsRequest, ProviderChatResponseData, ResolvedChatCompletionsRequest, SettledChatRequest, @@ -100,72 +99,9 @@ pub fn as_response_error(err: Error) -> Error { } } -#[cfg(feature = "bedrock-auth")] pub async fn signed_headers( request: &ProviderChatCompletionsRequest, body: &[u8], ) -> Result, Error> { - use std::collections::BTreeMap; - use std::time::SystemTime; - - use crate::providers::bedrock::aws_base::{ - aws_auth_config, aws_signature_headers, host_supplied_credentials, - is_sigv4_computed_header, resolve_credentials, sign_bedrock_post, - }; - - let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else { - return Ok(request.upstream_headers.clone()); - }; - // Reattaching a header the signer also emits would put both copies on the - // wire, and Bedrock rejects that pair. Python instead drops the caller's - // copy and prefers a forwarded Authorization over the signature, so leave - // the request to Python rather than serving it a different way here. - if request - .upstream_headers - .iter() - .any(|(name, _)| is_sigv4_computed_header(name)) - { - return Err(Error::Unsupported( - "request forwards a header AWS SigV4 computes", - )); - } - let env_lookup = |key: &str| std::env::var(key).ok(); - let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); - // A host with its own resolution chain hands the result down; only fall - // back to deriving credentials here when it supplied none. - let credentials = match host_supplied_credentials(&request.optional_params) { - Some(credentials) => credentials, - None => { - resolve_credentials( - aws_auth_config(&request.optional_params, &env_lookup), - &env_lookup, - ) - .await? - } - }; - let signature = sign_bedrock_post( - &request.url, - body, - &aws_signature_headers(&unsigned), - region, - &credentials, - SystemTime::now(), - )?; - // Every original header goes back on the wire alongside the computed ones, - // as Python reattaches them. The guard above already rejected the names - // that would collide, so no name appears twice. - Ok(unsigned.into_iter().chain(signature).collect()) -} - -#[cfg(not(feature = "bedrock-auth"))] -pub async fn signed_headers( - request: &ProviderChatCompletionsRequest, - _body: &[u8], -) -> Result, Error> { - match &request.auth { - ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported( - "AWS SigV4 requires the bedrock-auth feature", - )), - _ => Ok(request.upstream_headers.clone()), - } + request.config.authorize(request, body).await } diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index c050a358d46..5ad3726b06f 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -33,6 +33,21 @@ pub const STREAM_PARAM: &str = "stream"; const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; pub trait ChatCompletionsProviderConfig: Sync { + fn authorize<'a>( + &'a self, + request: &'a super::types::ProviderChatCompletionsRequest, + _body: &'a [u8], + ) -> crate::providers::AuthorizationFuture<'a> { + Box::pin(async move { + match &request.auth { + ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported( + "AWS SigV4 requires the bedrock-auth feature", + )), + _ => Ok(request.upstream_headers.clone()), + } + }) + } + fn request_body_behavior(&self) -> RequestBodyBehavior { RequestBodyBehavior::STRUCTURED_AT_SEND } diff --git a/litellm-rust/crates/core/src/integrations/custom_guardrail/types.rs b/litellm-rust/crates/core/src/integrations/custom_guardrail/types.rs index 825e56cc0d7..4bf7a98ae8c 100644 --- a/litellm-rust/crates/core/src/integrations/custom_guardrail/types.rs +++ b/litellm-rust/crates/core/src/integrations/custom_guardrail/types.rs @@ -3,27 +3,21 @@ use std::future::Future; use std::pin::Pin; use serde_json::Value; +use strum::AsRefStr; use crate::integrations::custom_logger::CallType; pub type GuardrailFuture<'a> = Pin> + Send + 'a>>; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(AsRefStr, Clone, Copy, Debug, PartialEq, Eq)] pub enum GuardrailEventHook { + #[strum(serialize = "pre_call")] PreCall, + #[strum(serialize = "during_call")] DuringCall, } -impl GuardrailEventHook { - pub fn as_str(&self) -> &'static str { - match self { - Self::PreCall => "pre_call", - Self::DuringCall => "during_call", - } - } -} - #[derive(Clone, Debug, PartialEq, Eq)] pub struct GuardrailError { pub message: String, diff --git a/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs b/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs index 9d38703dd57..de5ff62b3d3 100644 --- a/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs +++ b/litellm-rust/crates/core/src/integrations/custom_logger/mod.rs @@ -381,4 +381,20 @@ mod tests { assert_eq!(details.request_id, Some("req_ocr".to_string())); assert_eq!(details.litellm_call_id, Some("call_ocr".to_string())); } + + #[test] + fn call_type_strings_preserve_known_and_unknown_values() { + assert_eq!( + CallType::from("chat_completion").as_ref(), + "chat_completion" + ); + assert_eq!( + CallType::from("audio_transcription").as_ref(), + "audio_transcription" + ); + assert_eq!( + CallType::from("audio_transcription").to_string(), + "audio_transcription" + ); + } } diff --git a/litellm-rust/crates/core/src/integrations/custom_logger/types.rs b/litellm-rust/crates/core/src/integrations/custom_logger/types.rs index e68babc7821..b3f74a6394a 100644 --- a/litellm-rust/crates/core/src/integrations/custom_logger/types.rs +++ b/litellm-rust/crates/core/src/integrations/custom_logger/types.rs @@ -4,6 +4,7 @@ use std::pin::Pin; use serde::Serialize; use serde_json::Value; +use strum::{AsRefStr, Display, EnumString}; use crate::integrations::types::{StandardLoggingMetadata, StandardLoggingPayload}; @@ -15,48 +16,22 @@ pub struct CallbackDispatchReport { pub dropped: usize, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(AsRefStr, Clone, Debug, Display, EnumString, PartialEq, Eq)] pub enum CallType { + #[strum(serialize = "ocr")] Ocr, + #[strum(serialize = "realtime")] Realtime, + #[strum(serialize = "completion")] Completion, + #[strum(serialize = "acompletion")] Acompletion, + #[strum(serialize = "chat_completion")] ChatCompletion, + #[strum(default, transparent)] Other(String), } -impl CallType { - pub fn as_str(&self) -> &str { - match self { - Self::Ocr => "ocr", - Self::Realtime => "realtime", - Self::Completion => "completion", - Self::Acompletion => "acompletion", - Self::ChatCompletion => "chat_completion", - Self::Other(value) => value.as_str(), - } - } -} - -impl From<&str> for CallType { - fn from(value: &str) -> Self { - match value { - "ocr" => Self::Ocr, - "realtime" => Self::Realtime, - "completion" => Self::Completion, - "acompletion" => Self::Acompletion, - "chat_completion" => Self::ChatCompletion, - other => Self::Other(other.to_string()), - } - } -} - -impl std::fmt::Display for CallType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str(self.as_str()) - } -} - #[derive(Clone, Copy, Debug, PartialEq, Serialize)] pub struct CallbackTiming { pub start_time: f64, @@ -129,7 +104,7 @@ impl ModelCallDetails { Self { model: payload.model.clone(), custom_llm_provider: payload.custom_llm_provider.clone(), - call_type: CallType::from(payload.call_type.as_str()), + call_type: CallType::from(payload.call_type.as_ref()), metadata, extra_metadata: HashMap::new(), request_id, @@ -143,7 +118,7 @@ impl ModelCallDetails { pub fn with_standard_logging_payload(mut self, payload: StandardLoggingPayload) -> Self { self.model = payload.model.clone(); self.custom_llm_provider = payload.custom_llm_provider.clone(); - self.call_type = CallType::from(payload.call_type.as_str()); + self.call_type = CallType::from(payload.call_type.as_ref()); self.request_id = Some(payload.id.clone()); self.litellm_call_id = Some(payload.litellm_call_id.clone()); self.response_cost = Some(payload.response_cost); diff --git a/litellm-rust/crates/core/src/lifecycle/terminal.rs b/litellm-rust/crates/core/src/lifecycle/terminal.rs index 4b9b2128b29..976cfa1d202 100644 --- a/litellm-rust/crates/core/src/lifecycle/terminal.rs +++ b/litellm-rust/crates/core/src/lifecycle/terminal.rs @@ -111,7 +111,7 @@ impl From<&TerminalRecord> for ModelCallDetails { impl TerminalRecord { pub fn call_type(&self) -> CallType { - CallType::from(self.call_type.as_str()) + CallType::from(self.call_type.as_ref()) } } diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 6b688a3251e..55894d67391 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,26 +1,13 @@ +pub use crate::providers::dispatch::messages_provider_config; + use crate::Error; use crate::http_utils::string_headers as shared_string_headers; -use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; -use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; use serde_json::{Map, Value}; -use super::transformation::AnthropicMessagesProviderConfig; - pub use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body}; const HEADER_CONTEXT: &str = "messages"; -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub fn messages_provider_config( - provider: &str, -) -> Option<&'static dyn AnthropicMessagesProviderConfig> { - match provider { - "anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG), - "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG), - _ => None, - } -} - pub fn string_headers( extra_headers: Option>, ) -> Result, Error> { diff --git a/litellm-rust/crates/core/src/ocr/request.rs b/litellm-rust/crates/core/src/ocr/request.rs index 78623705165..56ac3c965bd 100644 --- a/litellm-rust/crates/core/src/ocr/request.rs +++ b/litellm-rust/crates/core/src/ocr/request.rs @@ -1,11 +1,10 @@ +pub(super) use crate::providers::dispatch::ocr_provider_config as provider_config; + use std::time::Duration; use serde_json::{Map, Value}; use crate::Error; -use crate::providers::azure_ai::ocr::transformation as azure_ai; -use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; -use crate::providers::vertex_ai::ocr::transformation as vertex_ai; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::transformation::{OcrProviderConfig, OcrResponseHandling}; @@ -45,7 +44,7 @@ fn check_admission_capabilities( request: &OcrAdmissionRequest, env_lookup: &dyn Fn(&str) -> Option, ) -> Result<(), Error> { - let (provider, config) = request_config(request)?; + let (_, config) = request_config(request)?; validate_capabilities(config)?; request .document @@ -59,20 +58,7 @@ fn check_admission_capabilities( .as_deref() .is_some_and(|key| !key.trim().is_empty()) || crate::http_utils::has_header(&request.extra_headers, "authorization"); - let configured = match provider.custom_llm_provider { - "azure_ai" => { - crate::http_utils::has_header(&request.extra_headers, "api-key") - || request - .azure_ad_token - .as_deref() - .is_some_and(|key| !key.trim().is_empty()) - || env_lookup("AZURE_AI_API_KEY").is_some_and(|key| !key.trim().is_empty()) - } - "vertex_ai" => ["VERTEX_AI_API_KEY", "VERTEXAI_API_KEY"] - .into_iter() - .any(|name| env_lookup(name).is_some_and(|key| !key.trim().is_empty())), - _ => false, - }; + let configured = config.has_configured_credentials(request, env_lookup); if !supplied && !configured { return Err(Error::Unsupported(operation)); } @@ -153,22 +139,11 @@ pub(super) fn validate_capabilities(config: &dyn OcrProviderConfig) -> Result<() } } -pub(super) fn provider_config( - provider: &str, - model: &str, -) -> Result<&'static dyn OcrProviderConfig, Error> { - match provider { - "mistral" => Ok(&MISTRAL_OCR_CONFIG), - "azure_ai" => azure_ai::config_for_model(model), - "vertex_ai" => vertex_ai::config_for_model(model), - _ => Err(Error::Unsupported("OCR provider")), - } -} - #[cfg(test)] mod tests { use super::*; use crate::ocr::types::OcrDocument; + use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; fn request() -> OcrAdmissionRequest { OcrAdmissionRequest { diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 6bb13771396..7d83ef4081b 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -29,6 +29,14 @@ pub trait OcrProviderConfig: Sync { OcrDocumentProjection::RetainedDocument } + fn has_configured_credentials( + &self, + _request: &super::types::OcrAdmissionRequest, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> bool { + false + } + fn credential_acquisition_operation(&self) -> Option<&'static str> { None } diff --git a/litellm-rust/crates/core/src/providers/anthropic/auth.rs b/litellm-rust/crates/core/src/providers/anthropic/auth.rs new file mode 100644 index 00000000000..a2382e020c9 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/anthropic/auth.rs @@ -0,0 +1,42 @@ +use crate::Error; + +const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; +const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; +const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; +const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; + +pub fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +pub fn resolve_anthropic_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + non_empty(api_key) + .map(str::to_string) + .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + Error::Auth( + "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ + environment variable" + .to_string(), + ) + }) +} + +pub fn complete_anthropic_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + let api_base = non_empty(api_base) + .map(str::to_string) + .or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty())) + .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()); + + let api_base = api_base.trim_end_matches('/'); + if api_base.ends_with(MESSAGES_PATH_SUFFIX) { + return api_base.to_string(); + } + format!("{api_base}{MESSAGES_PATH_SUFFIX}") +} 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 85f111f2094..bd5986ee314 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 @@ -11,9 +11,7 @@ use crate::chat_completions::types::{ }; use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX; use crate::error::Error; -use crate::providers::anthropic::messages::transformation::{ - complete_anthropic_url, resolve_anthropic_api_key, -}; +use crate::providers::anthropic::auth::{complete_anthropic_url, resolve_anthropic_api_key}; use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; @@ -49,7 +47,7 @@ fn anthropic_body(model: &str, conversation: &Conversation, params: Map>(), }) }) 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..91cff356d45 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -1,51 +1,13 @@ use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; - -const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; -const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; -const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; -const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; +pub use crate::providers::anthropic::auth::{ + complete_anthropic_url, non_empty, resolve_anthropic_api_key, +}; pub struct AnthropicMessagesConfig; pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig; -pub fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -pub fn resolve_anthropic_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - non_empty(api_key) - .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ - environment variable" - .to_string(), - ) - }) -} - -pub fn complete_anthropic_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - let api_base = non_empty(api_base) - .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty())) - .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()); - - let api_base = api_base.trim_end_matches('/'); - if api_base.ends_with(MESSAGES_PATH_SUFFIX) { - return api_base.to_string(); - } - format!("{api_base}{MESSAGES_PATH_SUFFIX}") -} - impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn complete_url( @@ -100,9 +62,8 @@ mod tests { #[test] fn url_falls_back_to_env_base() { - let with_env = |key: &str| { - (key == ANTHROPIC_API_BASE_ENV).then(|| "https://env.anthropic".to_string()) - }; + let with_env = + |key: &str| (key == "ANTHROPIC_API_BASE").then(|| "https://env.anthropic".to_string()); assert_eq!( complete_anthropic_url(Some(" "), &with_env), "https://env.anthropic/v1/messages" @@ -115,7 +76,7 @@ mod tests { resolve_anthropic_api_key(Some("sk-param"), &|_| None).unwrap(), "sk-param" ); - let with_env = |key: &str| (key == ANTHROPIC_API_KEY_ENV).then(|| "sk-env".to_string()); + let with_env = |key: &str| (key == "ANTHROPIC_API_KEY").then(|| "sk-env".to_string()); assert_eq!( resolve_anthropic_api_key(Some(" "), &with_env).unwrap(), "sk-env" diff --git a/litellm-rust/crates/core/src/providers/anthropic/mod.rs b/litellm-rust/crates/core/src/providers/anthropic/mod.rs index 0bb20991ff7..ccc0f1e3676 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/mod.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/mod.rs @@ -1,2 +1,3 @@ +pub mod auth; pub mod chat_completions; pub mod messages; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index 6fa932e4f0c..19597545ba9 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -516,6 +516,19 @@ fn transform_document_intelligence_response( } impl OcrProviderConfig for AzureAiOcrConfig { + fn has_configured_credentials( + &self, + request: &crate::ocr::types::OcrAdmissionRequest, + env_lookup: &dyn Fn(&str) -> Option, + ) -> bool { + crate::http_utils::has_header(&request.extra_headers, "api-key") + || request + .azure_ad_token + .as_deref() + .is_some_and(|key| !key.trim().is_empty()) + || env_lookup("AZURE_AI_API_KEY").is_some_and(|key| !key.trim().is_empty()) + } + fn document_projection(&self) -> OcrDocumentProjection { OcrDocumentProjection::ShallowCopyDocument } @@ -581,6 +594,14 @@ impl OcrProviderConfig for AzureAiOcrConfig { } impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { + fn has_configured_credentials( + &self, + request: &crate::ocr::types::OcrAdmissionRequest, + env_lookup: &dyn Fn(&str) -> Option, + ) -> bool { + AZURE_AI_OCR_CONFIG.has_configured_credentials(request, env_lookup) + } + fn document_projection(&self) -> OcrDocumentProjection { OcrDocumentProjection::Transformed } diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index 9bf1f73a74d..f388825cbb2 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -1,3 +1,4 @@ +use crate::audio_transcription::types::ProviderAudioTranscriptionRequest; use serde_json::{Map, Value, json}; use crate::audio_transcription::transformation::{ @@ -46,6 +47,14 @@ fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a } impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { + fn authorize<'a>( + &'a self, + request: &'a ProviderAudioTranscriptionRequest, + body: &'a [u8], + ) -> crate::providers::AuthorizationFuture<'a> { + Box::pin(signed_headers(request, body)) + } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_transcription_params(&self) -> &'static [&'static str] { SUPPORTED_PARAMS @@ -145,6 +154,37 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { } } +async fn signed_headers( + request: &ProviderAudioTranscriptionRequest, + body: &[u8], +) -> Result, Error> { + use std::collections::BTreeMap; + use std::time::SystemTime; + + use crate::audio_transcription::transformation::AudioTranscriptionAuth; + use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; + + let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else { + return Ok(request.upstream_headers.clone()); + }; + let env_lookup = |key: &str| std::env::var(key).ok(); + let credentials = resolve_credentials( + aws_auth_config(&request.optional_params, &env_lookup), + &env_lookup, + ) + .await?; + let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); + let signature = sign_bedrock_post( + &request.url, + body, + &unsigned, + region, + &credentials, + SystemTime::now(), + )?; + Ok(unsigned.into_iter().chain(signature).collect()) +} + #[cfg(test)] mod tests { use super::*; 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 9be2fec9bb8..af4614dece2 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,3 +1,4 @@ +use crate::chat_completions::types::ProviderChatCompletionsRequest; use serde_json::{Map, Value, json}; use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation}; @@ -61,7 +62,7 @@ fn converse_body(conversation: &Conversation, params: &Map) -> Va .iter() .map(|turn| { json!({ - "role": turn.role.as_str(), + "role": turn.role.as_ref(), "content": turn.texts.iter().map(|text| json!({"text": text})).collect::>(), }) }) @@ -105,6 +106,14 @@ fn has_blank_text(message: &ChatMessage) -> bool { } impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { + fn authorize<'a>( + &'a self, + request: &'a ProviderChatCompletionsRequest, + body: &'a [u8], + ) -> crate::providers::AuthorizationFuture<'a> { + Box::pin(signed_headers(request, body)) + } + fn request_body_behavior(&self) -> crate::lifecycle::RequestBodyBehavior { crate::lifecycle::RequestBodyBehavior::SERIALIZED_AT_BUILD } @@ -302,3 +311,59 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { }) } } + +async fn signed_headers( + request: &ProviderChatCompletionsRequest, + body: &[u8], +) -> Result, Error> { + use std::collections::BTreeMap; + use std::time::SystemTime; + + use crate::providers::bedrock::aws_base::{ + aws_auth_config, aws_signature_headers, host_supplied_credentials, + is_sigv4_computed_header, resolve_credentials, sign_bedrock_post, + }; + + let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else { + return Ok(request.upstream_headers.clone()); + }; + // Reattaching a header the signer also emits would put both copies on the + // wire, and Bedrock rejects that pair. Python instead drops the caller's + // copy and prefers a forwarded Authorization over the signature, so leave + // the request to Python rather than serving it a different way here. + if request + .upstream_headers + .iter() + .any(|(name, _)| is_sigv4_computed_header(name)) + { + return Err(Error::Unsupported( + "request forwards a header AWS SigV4 computes", + )); + } + let env_lookup = |key: &str| std::env::var(key).ok(); + let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); + // A host with its own resolution chain hands the result down; only fall + // back to deriving credentials here when it supplied none. + let credentials = match host_supplied_credentials(&request.optional_params) { + Some(credentials) => credentials, + None => { + resolve_credentials( + aws_auth_config(&request.optional_params, &env_lookup), + &env_lookup, + ) + .await? + } + }; + let signature = sign_bedrock_post( + &request.url, + body, + &aws_signature_headers(&unsigned), + region, + &credentials, + SystemTime::now(), + )?; + // Every original header goes back on the wire alongside the computed ones, + // as Python reattaches them. The guard above already rejected the names + // that would collide, so no name appears twice. + Ok(unsigned.into_iter().chain(signature).collect()) +} diff --git a/litellm-rust/crates/core/src/providers/dispatch.rs b/litellm-rust/crates/core/src/providers/dispatch.rs new file mode 100644 index 00000000000..f21181f9fbf --- /dev/null +++ b/litellm-rust/crates/core/src/providers/dispatch.rs @@ -0,0 +1,216 @@ +use super::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; +use super::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; +use super::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; +use super::azure_ai::ocr::transformation as azure_ai; +#[cfg(feature = "bedrock-auth")] +use super::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; +use super::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use super::vertex_ai::ocr::transformation as vertex_ai; +use crate::Error; +use crate::audio_transcription::transformation::AudioTranscriptionProviderConfig; +use crate::chat_completions::transformation::ChatCompletionsProviderConfig; +use crate::messages::transformation::AnthropicMessagesProviderConfig; +use crate::ocr::transformation::OcrProviderConfig; + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub fn chat_completions_provider_config( + provider: &str, +) -> Option<&'static dyn ChatCompletionsProviderConfig> { + match provider { + "anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG), + #[cfg(feature = "bedrock-auth")] + "bedrock" => Some( + &crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, + ), + _ => None, + } +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub fn messages_provider_config( + provider: &str, +) -> Option<&'static dyn AnthropicMessagesProviderConfig> { + match provider { + "anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG), + "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG), + _ => None, + } +} + +#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] +pub fn audio_transcription_provider_config( + provider: &str, +) -> Option<&'static dyn AudioTranscriptionProviderConfig> { + #[cfg(feature = "bedrock-auth")] + if provider == "bedrock" { + return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG); + } + let _ = provider; + None +} + +pub fn ocr_provider_config( + provider: &str, + model: &str, +) -> Result<&'static dyn OcrProviderConfig, Error> { + match provider { + "mistral" => Ok(&MISTRAL_OCR_CONFIG), + "azure_ai" => azure_ai::config_for_model(model), + "vertex_ai" => vertex_ai::config_for_model(model), + _ => Err(Error::Unsupported("OCR provider")), + } +} + +pub fn realtime_provider_config( + model: &str, +) -> Result< + ( + &str, + &'static (dyn crate::realtime::transformation::RealtimeProviderConfig + Sync), + ), + Error, +> { + let (provider, model) = model.split_once('/').unwrap_or(("openai", model)); + match provider { + "openai" => Ok(( + model, + &super::openai::realtime::transformation::OPENAI_REALTIME_CONFIG, + )), + _ => Err(Error::InvalidProvider(format!( + "realtime route does not support provider '{provider}'" + ))), + } +} + +pub fn responses_websocket_provider_config() +-> &'static dyn crate::responses::websocket::ResponsesWebSocketProviderConfig { + &super::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG +} + +pub(crate) fn resolve_audio_route_provider<'a>( + model: &'a str, + custom_llm_provider: Option<&'a str>, +) -> crate::routing_utils::provider::CustomLlmProvider<'a> { + use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider { + model, + custom_llm_provider: "bedrock", + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ocr::transformation::OcrResponseHandling; + + #[test] + fn unsupported_pairs_do_not_select_another_provider() { + let default_audio = resolve_audio_route_provider("model", None); + assert_eq!(default_audio.custom_llm_provider, "bedrock"); + assert_eq!(default_audio.model, "model"); + let explicit_audio = resolve_audio_route_provider("unknown/model", None); + assert_eq!(explicit_audio.custom_llm_provider, "unknown"); + assert!(audio_transcription_provider_config(explicit_audio.custom_llm_provider).is_none()); + for provider in ["unknown", "openai", "mistral", "reducto"] { + assert!(chat_completions_provider_config(provider).is_none()); + assert!(messages_provider_config(provider).is_none()); + assert!(audio_transcription_provider_config(provider).is_none()); + } + assert!(matches!( + ocr_provider_config("reducto", "model"), + Err(Error::Unsupported(_)) + )); + assert!(matches!( + realtime_provider_config("anthropic/model"), + Err(Error::InvalidProvider(_)) + )); + assert_eq!( + chat_completions_provider_config("bedrock").is_some(), + cfg!(feature = "bedrock-auth") + ); + assert_eq!( + audio_transcription_provider_config("bedrock").is_some(), + cfg!(feature = "bedrock-auth") + ); + } + + #[test] + fn ocr_selection_preserves_model_specific_protocols() { + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-layout") + .unwrap() + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert_eq!( + ocr_provider_config("azure_ai", "mistral") + .unwrap() + .response_handling(), + OcrResponseHandling::Json + ); + assert!( + ocr_provider_config("vertex_ai", "mistral") + .unwrap() + .requires_data_uri_document() + ); + assert!( + !ocr_provider_config("vertex_ai", "deepseek") + .unwrap() + .requires_data_uri_document() + ); + for provider in ["azure_ai", "vertex_ai"] { + assert!(matches!( + ocr_provider_config(provider, "cohere"), + Err(Error::Unsupported(_)) + )); + } + } + + #[test] + fn websocket_adapters_preserve_model_and_credential_policy() { + let (model, config) = realtime_provider_config("openai/model/variant").unwrap(); + assert_eq!(model, "model/variant"); + assert_eq!( + config.complete_url(Some("http://localhost"), model), + "ws://localhost/v1/realtime?model=model%2Fvariant" + ); + assert_eq!( + config + .resolve_api_key(Some(" explicit "), &|_| panic!("explicit key must win")) + .unwrap(), + "explicit" + ); + assert_eq!( + config + .resolve_api_key(Some(" "), &|key| (key == "OPENAI_API_KEY") + .then(|| " env ".into())) + .unwrap(), + " env " + ); + assert!(matches!( + config.resolve_api_key(None, &|_| None), + Err(Error::Auth(_)) + )); + + let config = responses_websocket_provider_config(); + let event = serde_json::from_value( + serde_json::json!({"type": "response.create", "model": "caller"}), + ) + .unwrap(); + let events = config + .transform_ws_request(&event, "deployment/model") + .unwrap() + .events; + assert_eq!(events[0].model(), Some("deployment/model")); + assert_eq!( + config + .resolve_api_key(Some(" explicit "), &|_| panic!("explicit key must win")) + .unwrap(), + "explicit" + ); + assert!(matches!( + config.resolve_api_key(None, &|_| Some(" ".into())), + Err(Error::Auth(_)) + )); + } +} diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index b71f3dbb5bb..56c88fbe53a 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -5,6 +5,7 @@ pub mod anthropic; pub mod azure_ai; #[cfg(feature = "bedrock-auth")] pub mod bedrock; +pub mod dispatch; pub mod mistral; pub mod openai; pub mod reducto; diff --git a/litellm-rust/crates/core/src/providers/openai/auth.rs b/litellm-rust/crates/core/src/providers/openai/auth.rs new file mode 100644 index 00000000000..9f33ab33d16 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/openai/auth.rs @@ -0,0 +1,11 @@ +use crate::Error; + +pub fn resolve_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + route: &str, +) -> Result { + api_key.map(str::trim).filter(|value| !value.is_empty()).map(str::to_string) + .or_else(|| env_lookup("OPENAI_API_KEY").filter(|value| !value.trim().is_empty())) + .ok_or_else(|| Error::Auth(format!("Missing OpenAI API Key - a {route} call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"))) +} diff --git a/litellm-rust/crates/core/src/providers/openai/mod.rs b/litellm-rust/crates/core/src/providers/openai/mod.rs index 62fcc50f2ac..371f09874ad 100644 --- a/litellm-rust/crates/core/src/providers/openai/mod.rs +++ b/litellm-rust/crates/core/src/providers/openai/mod.rs @@ -1,2 +1,16 @@ +pub mod auth; pub mod realtime; pub mod responses; + +fn percent_encode(value: &str) -> String { + let mut encoded = String::with_capacity(value.len()); + for byte in value.bytes() { + if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') { + encoded.push(byte as char); + } else { + encoded.push('%'); + encoded.push_str(&format!("{byte:02X}")); + } + } + encoded +} diff --git a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs index f1985f81b7d..ead76dc91a6 100644 --- a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -8,22 +8,7 @@ pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com"; /// Path appended to the resolved host base to reach the realtime endpoint. pub const OPENAI_REALTIME_PATH: &str = "/v1/realtime"; -/// Percent-encode a query value, escaping any char outside the RFC 3986 -/// unreserved set (`A-Za-z0-9-._~`). Keeps us dependency-free; common realtime -/// model slugs have no special chars, but this stays correct for the rest. -fn percent_encode(value: &str) -> String { - let mut encoded = String::with_capacity(value.len()); - for byte in value.bytes() { - let unreserved = byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~'); - if unreserved { - encoded.push(byte as char); - } else { - encoded.push('%'); - encoded.push_str(&format!("{byte:02X}")); - } - } - encoded -} +use crate::providers::openai::percent_encode; /// Build the realtime WebSocket URL, porting Python's `OpenAIRealtime._construct_url`. /// @@ -64,6 +49,14 @@ pub struct OpenAiRealtimeConfig; pub const OPENAI_REALTIME_CONFIG: OpenAiRealtimeConfig = OpenAiRealtimeConfig; impl RealtimeProviderConfig for OpenAiRealtimeConfig { + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + crate::providers::openai::auth::resolve_api_key(api_key, env_lookup, "realtime") + } + fn complete_url(&self, api_base: Option<&str>, model: &str) -> String { complete_url(api_base, model) } diff --git a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs index be86bb90311..61ac2f078da 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs @@ -1,4 +1,6 @@ use crate::Error; +use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; +use crate::providers::openai::percent_encode; use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; @@ -7,6 +9,14 @@ pub struct OpenAIResponsesWsConfig; pub const OPENAI_RESPONSES_WS_CONFIG: OpenAIResponsesWsConfig = OpenAIResponsesWsConfig; impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + crate::providers::openai::auth::resolve_api_key(api_key, env_lookup, "Responses WebSocket") + } + fn supports_native_websocket(&self) -> bool { true } @@ -30,6 +40,44 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { } } +pub fn complete_websocket_url(api_base: Option<&str>, model: &str, model_in_url: bool) -> String { + let base = api_base + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(OPENAI_RESPONSES_DEFAULT_API_BASE); + let (base, query) = base + .split_once('?') + .map_or((base, None), |(base, query)| (base, Some(query))); + let response_url = format!("{}{}", base.trim_end_matches('/'), OPENAI_RESPONSES_PATH); + let response_url = response_url + .strip_prefix("https://") + .map(|rest| format!("wss://{rest}")) + .or_else(|| { + response_url + .strip_prefix("http://") + .map(|rest| format!("ws://{rest}")) + }) + .unwrap_or(response_url); + let url = query.map_or_else( + || response_url.clone(), + |query| format!("{response_url}?{query}"), + ); + if !model_in_url + || query.is_some_and(|query| { + query + .split('&') + .any(|part| part.split('=').next() == Some("model")) + }) + { + return url; + } + format!( + "{url}{}model={}", + if query.is_some() { "&" } else { "?" }, + percent_encode(model) + ) +} + #[cfg(test)] mod tests { use super::*; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index 2b4941f459a..1fb965d77d6 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -220,6 +220,16 @@ fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> V } impl OcrProviderConfig for VertexAiOcrConfig { + fn has_configured_credentials( + &self, + _request: &crate::ocr::types::OcrAdmissionRequest, + env_lookup: &dyn Fn(&str) -> Option, + ) -> bool { + ["VERTEX_AI_API_KEY", "VERTEXAI_API_KEY"] + .into_iter() + .any(|name| env_lookup(name).is_some_and(|key| !key.trim().is_empty())) + } + fn document_projection(&self) -> OcrDocumentProjection { OcrDocumentProjection::ShallowCopyDocument } @@ -275,6 +285,14 @@ impl OcrProviderConfig for VertexAiOcrConfig { } impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { + fn has_configured_credentials( + &self, + request: &crate::ocr::types::OcrAdmissionRequest, + env_lookup: &dyn Fn(&str) -> Option, + ) -> bool { + VERTEX_AI_OCR_CONFIG.has_configured_credentials(request, env_lookup) + } + fn document_projection(&self) -> OcrDocumentProjection { OcrDocumentProjection::Transformed } diff --git a/litellm-rust/crates/core/src/realtime/streaming.rs b/litellm-rust/crates/core/src/realtime/streaming.rs index 89368d1c0df..974c702b4f0 100644 --- a/litellm-rust/crates/core/src/realtime/streaming.rs +++ b/litellm-rust/crates/core/src/realtime/streaming.rs @@ -23,20 +23,19 @@ use crate::lifecycle::{ CallLifecycleContext, Clock, CostInputs, ExecutedCall, RouteProjection, TerminalClassification, TerminalDispatcher, TerminalRecord, }; -use crate::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; +use crate::providers::dispatch::realtime_provider_config; use crate::realtime::transformation::RealtimeProviderConfig; use crate::realtime::types::RealtimeEvent; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const IDLE_TIMEOUT: Duration = Duration::from_secs(300); -const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; type Upstream = WebSocketStream>; static TLS_CONFIG: OnceLock> = OnceLock::new(); -#[derive(Clone, Eq)] +#[derive(Clone)] pub struct RealtimeConnectionSpec { + config: &'static (dyn RealtimeProviderConfig + Sync), model: String, api_key: String, api_base: Option, @@ -49,9 +48,11 @@ impl RealtimeConnectionSpec { api_base: Option<&str>, ) -> Result { let model = model.into(); + let (model, config) = realtime_provider_config(&model)?; Ok(Self { - model: openai_model(&model)?.to_string(), - api_key: resolve_api_key(api_key)?, + config, + model: model.to_string(), + api_key: config.resolve_api_key(api_key, &|key| std::env::var(key).ok())?, api_base: api_base.map(str::to_string), }) } @@ -61,6 +62,8 @@ impl RealtimeConnectionSpec { } } +impl Eq for RealtimeConnectionSpec {} + impl PartialEq for RealtimeConnectionSpec { fn eq(&self, other: &Self) -> bool { self.model == other.model @@ -233,14 +236,15 @@ where Out::Error: std::fmt::Display, { let WarmConnection { - connection: _, + connection, upstream, session_created, } = connection; + let config = connection.config; let (mut upstream_tx, mut upstream_rx) = upstream.split(); if !session_created.event_type.is_empty() { observation.observe(&session_created); - send_client_event(&mut client_out, &session_created, model).await?; + send_client_event(config, &mut client_out, &session_created, model).await?; } loop { tokio::select! { @@ -251,7 +255,7 @@ where "realtime client disconnected before provider completion", ); }; - for outbound in OPENAI_REALTIME_CONFIG.transform_realtime_request(&event, model)?.events { + for outbound in config.transform_realtime_request(&event, model)?.events { let payload = serde_json::to_string(&outbound) .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream_tx.send(Message::Text(payload)).await.map_err(ws_transport_error)?; @@ -270,7 +274,7 @@ where let event = serde_json::from_str::(&text) .map_err(|error| Error::InvalidResponse(error.to_string()))?; observation.observe(&event); - send_client_event(&mut client_out, &event, model).await?; + send_client_event(config, &mut client_out, &event, model).await?; if event.event_type == "error" || response_failed(&event) { return Err(RealtimeFailure::new( "ProviderError", @@ -346,6 +350,7 @@ fn provider_error_message(event: &RealtimeEvent) -> String { } async fn send_client_event( + config: &(dyn RealtimeProviderConfig + Sync), client_out: &mut Out, event: &RealtimeEvent, model: &str, @@ -354,10 +359,7 @@ where Out: Sink + Unpin, Out::Error: std::fmt::Display, { - for outbound in OPENAI_REALTIME_CONFIG - .transform_realtime_response(event, model)? - .events - { + for outbound in config.transform_realtime_response(event, model)?.events { client_out .send(outbound) .await @@ -389,7 +391,8 @@ async fn read_event(upstream: &mut Upstream) -> Result { } async fn dial_upstream(connection: &RealtimeConnectionSpec) -> Result { - let url = OPENAI_REALTIME_CONFIG + let url = connection + .config .complete_url(connection.api_base.as_deref(), connection.model.as_str()); let mut request = url.into_client_request().map_err(ws_transport_error)?; request.headers_mut().insert( @@ -433,31 +436,6 @@ fn tls_config() -> Result, Error> { Ok(Arc::clone(TLS_CONFIG.get_or_init(|| config))) } -fn resolve_api_key(api_key: Option<&str>) -> Result { - api_key - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .or_else(|| { - std::env::var(OPENAI_API_KEY_ENV) - .ok() - .filter(|value| !value.trim().is_empty()) - }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) -} - -fn openai_model(model: &str) -> Result<&str, Error> { - if let Some((provider, provider_model)) = model.split_once('/') { - if provider != "openai" { - return Err(Error::InvalidProvider(format!( - "realtime route does not support provider '{provider}'" - ))); - } - return Ok(provider_model); - } - Ok(model) -} - fn ws_handshake_error(error: WsError) -> Error { match error { WsError::Http(response) => Error::Http { diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs index b08084514ef..e6c441e2da8 100644 --- a/litellm-rust/crates/core/src/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -2,6 +2,14 @@ use crate::Error; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; pub trait RealtimeProviderConfig { + fn resolve_api_key( + &self, + _api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Err(Error::Unsupported("provider credential resolution")) + } + /// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`). /// Pure string construction only — no network, no env. fn complete_url(&self, api_base: Option<&str>, model: &str) -> String; diff --git a/litellm-rust/crates/core/src/responses/types.rs b/litellm-rust/crates/core/src/responses/types.rs index 4942309992e..9727d3954e9 100644 --- a/litellm-rust/crates/core/src/responses/types.rs +++ b/litellm-rust/crates/core/src/responses/types.rs @@ -1,37 +1,31 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_json::{Map, Value}; +use strum::AsRefStr; -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(AsRefStr, Clone, Debug, PartialEq, Eq)] pub enum ResponsesWsEventType { + #[strum(serialize = "response.create")] ResponseCreate, + #[strum(serialize = "response.created")] ResponseCreated, + #[strum(serialize = "response.completed")] ResponseCompleted, + #[strum(serialize = "response.failed")] ResponseFailed, + #[strum(serialize = "response.incomplete")] ResponseIncomplete, + #[strum(serialize = "error")] Error, + #[strum(default, transparent)] Other(String), } -impl ResponsesWsEventType { - pub fn as_str(&self) -> &str { - match self { - Self::ResponseCreate => "response.create", - Self::ResponseCreated => "response.created", - Self::ResponseCompleted => "response.completed", - Self::ResponseFailed => "response.failed", - Self::ResponseIncomplete => "response.incomplete", - Self::Error => "error", - Self::Other(value) => value, - } - } -} - impl Serialize for ResponsesWsEventType { fn serialize(&self, serializer: S) -> Result where S: Serializer, { - serializer.serialize_str(self.as_str()) + serializer.serialize_str(self.as_ref()) } } diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 337512bdbfe..35f842ac825 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -14,25 +14,35 @@ use tokio_tungstenite::{ }; use crate::Error; -use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::integrations::custom_logger::CallbackTiming; use crate::lifecycle::{ CostInputs, ExecutedCall, RouteProjection, TerminalClassification, TerminalDispatcher, TerminalRecord, }; -use crate::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; +use crate::providers::dispatch::responses_websocket_provider_config; use crate::responses::instrumentation::ResponsesWsInstrumentation; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); const IDLE_TIMEOUT: Duration = Duration::from_secs(300); -const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; type Upstream = WebSocketStream>; static TLS_CONFIG: OnceLock> = OnceLock::new(); +struct ProviderConnection { + upstream: Upstream, + config: &'static dyn ResponsesWebSocketProviderConfig, +} + pub trait ResponsesWebSocketProviderConfig: Sync { + fn resolve_api_key( + &self, + _api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Err(Error::Unsupported("provider credential resolution")) + } + fn supports_native_websocket(&self) -> bool { false } @@ -80,8 +90,9 @@ where Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { - let key = resolve_api_key(request.api_key.as_deref())?; - let upstream = dial_upstream(&request.model, &key, request.api_base.as_deref()).await?; + let config = responses_websocket_provider_config(); + let key = config.resolve_api_key(request.api_key.as_deref(), &|key| std::env::var(key).ok())?; + let upstream = dial_upstream(config, &request.model, &key, request.api_base.as_deref()).await?; let start_time = services.now(); let instrumentation = Arc::new(ResponsesWsInstrumentation::default()); let mut completion = @@ -234,7 +245,7 @@ impl ResponsesWsFailure { } async fn splice( - upstream: Upstream, + connection: ProviderConnection, model: &str, first_frame: Option, idle_timeout: Duration, @@ -247,9 +258,10 @@ where Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { + let ProviderConnection { upstream, config } = connection; let (mut upstream_tx, mut upstream_rx) = upstream.split(); if let Some(event) = first_frame { - send_provider_event(&mut upstream_tx, &event, model) + send_provider_event(config, &mut upstream_tx, &event, model) .await .map_err(ResponsesWsFailure::new)?; } @@ -263,7 +275,7 @@ where )); }; let event = event.map_err(ResponsesWsFailure::new)?; - send_provider_event(&mut upstream_tx, &event, model).await.map_err(ResponsesWsFailure::new)?; + send_provider_event(config, &mut upstream_tx, &event, model).await.map_err(ResponsesWsFailure::new)?; } message = upstream_rx.next() => { let Some(message) = message else { @@ -278,7 +290,7 @@ where .map_err(|error| ResponsesWsFailure::new(Error::InvalidResponse(error.to_string())))?; instrumentation.observe(&event); let terminal = instrumentation.terminal_classification(&event); - for outbound in OPENAI_RESPONSES_WS_CONFIG.transform_ws_response(&event, model) + for outbound in config.transform_ws_response(&event, model) .map_err(ResponsesWsFailure::new)?.events { client_out.send(outbound).await .map_err(|error| ResponsesWsFailure::session( @@ -306,14 +318,12 @@ where } async fn send_provider_event( + config: &dyn ResponsesWebSocketProviderConfig, upstream: &mut futures_util::stream::SplitSink, event: &ResponsesWsEvent, model: &str, ) -> Result<(), Error> { - for outbound in OPENAI_RESPONSES_WS_CONFIG - .transform_ws_request(event, model)? - .events - { + for outbound in config.transform_ws_request(event, model)?.events { let payload = serde_json::to_string(&outbound) .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream @@ -325,11 +335,12 @@ async fn send_provider_event( } async fn dial_upstream( + config: &'static dyn ResponsesWebSocketProviderConfig, model: &str, api_key: &str, api_base: Option<&str>, -) -> Result { - let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model); +) -> Result { + let url = config.complete_websocket_url(api_base, model); let mut request = url.into_client_request().map_err(ws_transport_error)?; request.headers_mut().insert( AUTHORIZATION, @@ -344,7 +355,9 @@ async fn dial_upstream( let result = tokio::time::timeout(CONNECT_TIMEOUT, connect) .await .map_err(|_| Error::Connect("Responses WebSocket connection timed out".to_string()))?; - result.map(|(socket, _)| socket).map_err(ws_handshake_error) + result + .map(|(upstream, _)| ProviderConnection { upstream, config }) + .map_err(ws_handshake_error) } fn tls_config() -> Result, Error> { @@ -391,69 +404,7 @@ fn ws_transport_error(error: WsError) -> Error { } } -fn resolve_api_key(api_key: Option<&str>) -> Result { - api_key - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .or_else(|| { - std::env::var(OPENAI_API_KEY_ENV) - .ok() - .filter(|value| !value.trim().is_empty()) - }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) -} - -pub fn complete_websocket_url(api_base: Option<&str>, model: &str, model_in_url: bool) -> String { - let base = api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or(OPENAI_RESPONSES_DEFAULT_API_BASE); - let (base, query) = base - .split_once('?') - .map_or((base, None), |(base, query)| (base, Some(query))); - let response_url = format!("{}{}", base.trim_end_matches('/'), OPENAI_RESPONSES_PATH); - let response_url = response_url - .strip_prefix("https://") - .map(|rest| format!("wss://{rest}")) - .or_else(|| { - response_url - .strip_prefix("http://") - .map(|rest| format!("ws://{rest}")) - }) - .unwrap_or(response_url); - let url = query.map_or_else( - || response_url.clone(), - |query| format!("{response_url}?{query}"), - ); - if !model_in_url - || query.is_some_and(|query| { - query - .split('&') - .any(|part| part.split('=').next() == Some("model")) - }) - { - return url; - } - format!( - "{url}{}model={}", - if query.is_some() { "&" } else { "?" }, - percent_encode(model) - ) -} - -fn percent_encode(value: &str) -> String { - value - .bytes() - .map(|byte| { - if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') { - format!("{}", byte as char) - } else { - format!("%{byte:02X}") - } - }) - .collect() -} +pub use crate::providers::openai::responses::transformation::complete_websocket_url; pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent { if !event.is_response_create() { @@ -662,7 +613,7 @@ mod tests { .unwrap(); assert!(matches!(result, ExecutedCall::Failure { .. })); assert_eq!( - output_rx.next().await.unwrap().event_type.as_str(), + output_rx.next().await.unwrap().event_type.as_ref(), expected_type ); assert_failure(services.as_ref(), kind, message); diff --git a/litellm-rust/crates/core/tests/chat_completions.rs b/litellm-rust/crates/core/tests/chat_completions.rs index 3471d2a8a27..e2b463a901c 100644 --- a/litellm-rust/crates/core/tests/chat_completions.rs +++ b/litellm-rust/crates/core/tests/chat_completions.rs @@ -832,3 +832,52 @@ mod round_trip { )); } } + +#[cfg(feature = "bedrock-auth")] +#[tokio::test] +async fn provider_authorization_signs_the_supplied_bytes_without_reserializing() { + use litellm_core::providers::bedrock::aws_base::{ + aws_signature_headers, host_supplied_credentials, sign_bedrock_post, + }; + use std::collections::BTreeMap; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16, "aws_access_key_id": "test-access", "aws_secret_access_key": "test-secret"}), + ); + call.api_key = None; + let built = build_chat_completions_request(call).unwrap(); + let bytes = b"{ \"settled\": true }\n"; + let before = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let signed: BTreeMap<_, _> = litellm_core::chat_completions::signed_headers(&built, bytes) + .await + .unwrap() + .into_iter() + .collect(); + let after = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let credentials = host_supplied_credentials(&built.optional_params).unwrap(); + let headers = aws_signature_headers(&built.upstream_headers.iter().cloned().collect()); + assert!((before..=after).any(|second| { + let expected = sign_bedrock_post( + &built.url, + bytes, + &headers, + "us-east-1", + &credentials, + UNIX_EPOCH + Duration::from_secs(second), + ) + .unwrap(); + expected + .iter() + .all(|(name, value)| signed.get(name) == Some(value)) + })); +}