From 0b8817afbbf0abd7dda5bf968e34cced300350db Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:09:41 -0700 Subject: [PATCH] perf(bedrock): audio transcription via rust core (py->rust bridge) (#33990) * feat(bedrock): add audio transcription via Converse with py->rust bridge Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(ci): exclude rust transcription rollout flag from docs check Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(bedrock): await rust/python fallback in async transcription dispatch Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(bedrock): cover audio transcription rust dispatch Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(bedrock): route audio transcription through rust Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(bedrock): move rust transcription dispatch out of main Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(bedrock): include rust transcription coverage shard Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm-rust/CLAUDE.md | 7 + .../PROVIDER_CODING_STANDARDS.md | 10 +- litellm-rust/crates/ai-gateway/Cargo.toml | 2 +- .../src/audio_transcription/common_utils.rs | 48 +++ .../src/audio_transcription/handler.rs | 89 +++++ .../src/audio_transcription/hooks.rs | 300 +++++++++++++++++ .../ai-gateway/src/audio_transcription/mod.rs | 25 ++ .../src/audio_transcription/prepare.rs | 55 ++++ .../src/audio_transcription/tests.rs | 53 +++ .../src/audio_transcription/types.rs | 58 ++++ .../crates/ai-gateway/src/{ocr => }/client.rs | 6 +- .../ai-gateway/src/io/audio_transcription.rs | 1 + litellm-rust/crates/ai-gateway/src/io/mod.rs | 1 + litellm-rust/crates/ai-gateway/src/lib.rs | 2 + .../crates/ai-gateway/src/ocr/common_utils.rs | 2 +- .../crates/ai-gateway/src/ocr/handler.rs | 2 +- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 1 - .../core/src/audio_transcription/mod.rs | 2 + .../src/audio_transcription/transformation.rs | 57 ++++ .../core/src/audio_transcription/types.rs | 20 ++ litellm-rust/crates/core/src/lib.rs | 1 + .../providers/bedrock/audio_transcription.rs | 310 ++++++++++++++++++ .../core/src/providers/bedrock/aws_base.rs | 6 +- .../core/src/providers/bedrock/constants.rs | 4 + .../crates/core/src/providers/bedrock/mod.rs | 2 + litellm-rust/crates/python-bridge/CLAUDE.md | 8 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/src/lib.rs | 92 ++++++ .../bedrock/audio_transcription/__init__.py | 84 +++++ litellm/main.py | 55 +++- litellm/rust_bridge/ocr.py | 11 + litellm/rust_bridge/transcription.py | 148 +++++++++ .../test_audio_transcription_rust_bridge.py | 151 +++++++++ 33 files changed, 1586 insertions(+), 29 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs create mode 100644 litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs rename litellm-rust/crates/ai-gateway/src/{ocr => }/client.rs (60%) create mode 100644 litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/mod.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/transformation.rs create mode 100644 litellm-rust/crates/core/src/audio_transcription/types.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs create mode 100644 litellm/llms/bedrock/audio_transcription/__init__.py create mode 100644 litellm/rust_bridge/transcription.py create mode 100644 tests/test_litellm/test_audio_transcription_rust_bridge.py diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index 519b1d205ef..0659e63df39 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -62,6 +62,13 @@ Not allowed in `core`: Python owns rollout state and fallback while Rust is being introduced. Rust paths must be off by default until parity tests prove equivalence with Python. +A new provider/route may instead be implemented rust-only with no Python +reference; then the Python interface is a thin dispatch that calls Rust with no +fallback, and you state the rust-only choice explicitly in the PR. Either way +the Python side stays minimal (it only marshals inputs and calls the Rust +interface), never add a per-route feature flag, and never push provider +dispatch into `litellm/main.py`; put it in a thin dispatch class under +`litellm/llms///`. ## Production Bar diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index c7980b11147..ed44dc4c729 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -39,11 +39,17 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST 19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity. 20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping. -21. Rust paths stay off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. +21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR. + +## Python bridge (SDK side) + +22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust. +23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms///` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method. +24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_`. ## Checks before push -22. Run, and keep green: +25. Run, and keep green: ```bash cd litellm-rust cargo fmt --check diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index c15af4cc478..541beabe170 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -14,7 +14,7 @@ path = "src/main.rs" required-features = ["server"] [dependencies] -litellm-core.workspace = true +litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs new file mode 100644 index 00000000000..270d5c2d97a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs @@ -0,0 +1,48 @@ +use std::collections::BTreeMap; + +use litellm_core::CoreResult; +use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; +use litellm_core::error::CoreError; +use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; +use serde_json::{Map, Value}; + +pub(super) fn audio_transcription_provider_config( + provider: &str, +) -> Option<&'static dyn AudioTranscriptionProviderConfig> { + match provider { + "bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG), + _ => None, + } +} + +pub(super) fn string_headers( + headers: Option>, +) -> CoreResult> { + headers + .unwrap_or_default() + .into_iter() + .map(|(key, value)| { + value + .as_str() + .map(|value| (key.clone(), value.to_string())) + .ok_or_else(|| { + CoreError::InvalidRequest(format!( + "audio transcription extra_headers.{key} must be a string" + )) + }) + }) + .collect() +} + +pub(super) fn has_header(headers: &BTreeMap, name: &str) -> bool { + headers.keys().any(|key| key.eq_ignore_ascii_case(name)) +} + +pub(super) fn truncate_error_body(body: &str) -> String { + let truncated: String = body.chars().take(256).collect(); + if truncated.chars().count() == body.chars().count() { + truncated + } else { + format!("{truncated}... (truncated)") + } +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs new file mode 100644 index 00000000000..33c13550f58 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs @@ -0,0 +1,89 @@ +use std::time::SystemTime; + +use litellm_core::CoreResult; +use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; +use litellm_core::error::CoreError; +use litellm_core::providers::bedrock::audio_transcription::aws_auth_config; +use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; +use serde_json::Value; + +use super::common_utils::truncate_error_body; +use super::types::ProviderAudioTranscriptionRequest; +use crate::client::http_client; + +pub(crate) async fn execute_audio_transcription_provider_call( + request: ProviderAudioTranscriptionRequest, +) -> CoreResult { + let body = serde_json::to_vec(&request.body).map_err(|error| { + CoreError::InvalidRequest(format!("invalid audio request body: {error}")) + })?; + let mut request_builder = http_client().post(&request.url).body(body.clone()); + for (key, value) in &request.upstream_headers { + request_builder = request_builder.header(key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); + } + let response = request_builder + .send() + .await + .map_err(|error| CoreError::Network(error.to_string()))?; + let status = response.status(); + let text = response + .text() + .await + .map_err(|error| CoreError::Network(error.to_string()))?; + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + let response_json: Value = serde_json::from_str(&text).map_err(|error| { + CoreError::InvalidResponse(format!("invalid audio response JSON: {error}")) + })?; + Ok(request + .config + .transform_transcription_response(&request.model, response_json)? + .into_json()) +} + +pub(crate) async fn sign_request( + request: &ProviderAudioTranscriptionRequest, + optional_params: &serde_json::Map, +) -> CoreResult { + let env_lookup = environment_lookup; + let auth = request + .config + .auth_strategy(&request.model, optional_params, &env_lookup)?; + let body = serde_json::to_vec(&request.body).map_err(|error| { + CoreError::InvalidRequest(format!("invalid audio request body: {error}")) + })?; + let mut headers = super::common_utils::string_headers(None)?; + headers.insert("Content-Type".to_string(), "application/json".to_string()); + headers.extend(request.upstream_headers.iter().cloned()); + match auth { + AudioTranscriptionAuth::Bearer => {} + AudioTranscriptionAuth::AwsSigV4 { region, .. } => { + let credentials = + resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup) + .await?; + headers.extend(sign_bedrock_post( + &request.url, + &body, + &headers, + ®ion, + &credentials, + SystemTime::now(), + )?); + } + } + Ok(ProviderAudioTranscriptionRequest { + upstream_headers: headers.into_iter().collect(), + ..request.clone() + }) +} + +pub(super) fn environment_lookup(key: &str) -> Option { + std::env::var(key).ok() +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs new file mode 100644 index 00000000000..8b6896f3846 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -0,0 +1,300 @@ +use std::future::Future; +use std::pin::Pin; + +use litellm_core::CoreResult; +use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; +use litellm_core::error::CoreError; +use serde_json::{Map, Value, json}; + +use super::common_utils::{audio_transcription_provider_config, has_header, string_headers}; +use super::handler::sign_request; +use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +use crate::integrations::custom_guardrail::{ + CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, +}; +use crate::integrations::custom_logger::{ + CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails, +}; +use crate::integrations::types::{ + RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload, +}; + +pub(crate) struct AudioTranscriptionLifecycleHooks { + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, +} + +type AudioFuture<'a, T> = Pin> + Send + 'a>>; +type AudioLogFuture<'a> = Pin + Send + 'a>>; + +impl AudioTranscriptionLifecycleHooks { + pub(crate) fn new( + logger_runner: CustomLoggerRunner, + guardrail_runner: CustomGuardrailRunner, + request_metadata: RequestMetadata, + ) -> Self { + Self { + logger_runner, + guardrail_runner, + request_metadata, + } + } + + async fn run_pre_call_guardrails( + &self, + request: PreparedAudioTranscriptionRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + let (guardrail_request, _) = self + .guardrail_runner + .run_pre_call( + &guardrail_context(&self.request_metadata), + GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": request.custom_llm_provider, + "audio": request.audio, + "optional_params": request.optional_params, + })), + ) + .await + .map_err(guardrail_error_to_core_error)?; + let Value::Object(mut data) = guardrail_request.data else { + return Err(CoreError::InvalidRequest( + "audio transcription pre_call guardrail must return an object".to_string(), + )); + }; + let audio = data.remove("audio").ok_or_else(|| { + CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string()) + })?; + let optional_params = match data.remove("optional_params") { + Some(Value::Object(value)) => value, + Some(_) => { + return Err(CoreError::InvalidRequest( + "audio transcription optional_params must be an object".to_string(), + )); + } + None => Map::new(), + }; + Ok(PreparedAudioTranscriptionRequest { + audio, + optional_params, + ..request + }) + } + + async fn prepare_provider_request( + &self, + request: PreparedAudioTranscriptionRequest, + ) -> CoreResult { + let config = audio_transcription_provider_config(&request.custom_llm_provider) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + let env_lookup = super::handler::environment_lookup; + let headers = string_headers(request.extra_headers)?; + let url = config.complete_url( + request.api_base.as_deref(), + &request.model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_transcription_params(&request.optional_params); + let body = config.transform_transcription_request( + &request.model, + request.audio, + filtered_params, + )?; + let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?; + let mut upstream_headers = headers.into_iter().collect::>(); + if matches!(auth, AudioTranscriptionAuth::Bearer) + && !has_header( + &upstream_headers + .iter() + .cloned() + .collect::>(), + "authorization", + ) + && let Some(api_key) = request.api_key.as_deref() + { + upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}"))); + } + let provider_request = ProviderAudioTranscriptionRequest { + model: request.model, + config, + url, + body: body.body, + upstream_headers, + timeout: request.timeout, + }; + let provider_request = self.run_during_call_guardrails(provider_request).await?; + sign_request(&provider_request, &request.optional_params).await + } + + async fn run_during_call_guardrails( + &self, + request: ProviderAudioTranscriptionRequest, + ) -> CoreResult { + if self.guardrail_runner.is_empty() { + return Ok(request); + } + let (guardrail_request, _) = self + .guardrail_runner + .run_during_call( + &guardrail_context(&self.request_metadata), + GuardrailRequest::new(json!({ + "model": request.model, + "custom_llm_provider": "bedrock", + "url": request.url, + "body": request.body, + })), + ) + .await + .map_err(guardrail_error_to_core_error)?; + let Value::Object(mut data) = guardrail_request.data else { + return Err(CoreError::InvalidRequest( + "audio transcription during_call guardrail must return an object".to_string(), + )); + }; + let body = data.remove("body").ok_or_else(|| { + CoreError::InvalidRequest("audio transcription guardrail removed body".to_string()) + })?; + Ok(ProviderAudioTranscriptionRequest { body, ..request }) + } + + fn logging_payload( + &self, + context: &CallLifecycleContext, + timing: &CallLifecycleTiming, + ) -> StandardLoggingPayload { + StandardLoggingPayload { + id: context.litellm_call_id.clone(), + litellm_call_id: context.litellm_call_id.clone(), + call_type: context.call_type.clone(), + model: context.model.clone(), + custom_llm_provider: context.custom_llm_provider.clone(), + response_cost: 0.0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + start_time: timing.start_time, + end_time: timing.end_time, + stream: false, + metadata: StandardLoggingMetadata { + user_api_key_hash: self.request_metadata.user_api_key_hash.clone(), + user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(), + user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(), + ..Default::default() + }, + messages: None, + } + } +} + +impl CallLifecycleHooks + for AudioTranscriptionLifecycleHooks +{ + type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>; + type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>; + type SuccessFuture<'a> = AudioLogFuture<'a>; + type FailureFuture<'a> = AudioLogFuture<'a>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedAudioTranscriptionRequest, + ) -> Self::PreCallFuture<'a> { + Box::pin(async move { self.run_pre_call_guardrails(request).await }) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedAudioTranscriptionRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(async move { self.prepare_provider_request(request).await }) + } + + fn async_log_success_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + response: &'a Value, + timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + self.logger_runner + .async_log_success_event( + &ModelCallDetails::from_standard_logging_payload( + self.logging_payload(context, timing), + ), + &CallbackValue::new("audio_transcription", response.clone()), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } + + fn async_log_failure_event<'a>( + &'a self, + context: &'a CallLifecycleContext, + error: &'a CoreError, + timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + Box::pin(async move { + if self.logger_runner.is_empty() { + return; + } + let logging_error = LoggingError { + message: error.to_string(), + kind: core_error_kind(error).to_string(), + }; + self.logger_runner + .async_log_failure_event( + &ModelCallDetails::from_standard_logging_payload( + self.logging_payload(context, timing), + ) + .with_failure_error(logging_error.clone()), + Some(&CallbackValue::new( + "error", + json!({"message": logging_error.message, "kind": logging_error.kind}), + )), + CallbackTiming::new(timing.start_time, timing.end_time), + ) + .await; + }) + } +} + +fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { + GuardrailContext { + call_type: CallType::Other("audio_transcription".to_string()), + selected_guardrails: Vec::new(), + metadata: std::collections::HashMap::new(), + user_api_key_hash: metadata.user_api_key_hash.clone(), + user_api_key_user_id: metadata.user_api_key_user_id.clone(), + user_api_key_team_id: metadata.user_api_key_team_id.clone(), + trace_parent: None, + } +} + +fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { + CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +} + +fn core_error_kind(error: &CoreError) -> &'static str { + match error { + CoreError::Auth(_) => "AuthError", + CoreError::InvalidProvider(_) => "InvalidProvider", + CoreError::InvalidRequest(_) => "InvalidRequest", + CoreError::InvalidType { .. } => "InvalidType", + CoreError::MissingField(_) => "MissingField", + CoreError::Http { .. } => "HttpError", + CoreError::InvalidResponse(_) => "InvalidResponse", + CoreError::Network(_) => "NetworkError", + CoreError::Routing(_) => "RoutingError", + } +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs new file mode 100644 index 00000000000..5d33d912c40 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -0,0 +1,25 @@ +use litellm_core::CoreResult; +use litellm_core::call_lifecycle::CallLifecycle; +use serde_json::Value; + +mod common_utils; +mod handler; +mod hooks; +mod prepare; +mod types; + +pub use types::AudioTranscriptionRequest; + +use handler::execute_audio_transcription_provider_call; +use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call}; + +pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult { + let PreparedAudioTranscriptionCall { request, hooks } = + prepare_audio_transcription_call(request); + CallLifecycle::default() + .run_request(request, &hooks, execute_audio_transcription_provider_call) + .await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs new file mode 100644 index 00000000000..a475d58635f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs @@ -0,0 +1,55 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + +use super::hooks::AudioTranscriptionLifecycleHooks; +use super::types::{AudioTranscriptionRequest, PreparedAudioTranscriptionRequest}; +use crate::integrations::custom_guardrail::CustomGuardrailRunner; +use crate::integrations::custom_logger::CustomLoggerRunner; + +pub(crate) struct PreparedAudioTranscriptionCall { + pub(crate) request: PreparedAudioTranscriptionRequest, + pub(crate) hooks: AudioTranscriptionLifecycleHooks, +} + +pub(crate) fn prepare_audio_transcription_call( + request: AudioTranscriptionRequest<'_>, +) -> PreparedAudioTranscriptionCall { + let call_id = request + .litellm_call_id + .map(str::to_string) + .unwrap_or_else(new_audio_transcription_call_id); + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .unwrap_or(CustomLlmProvider { + model: request.model, + custom_llm_provider: "bedrock", + }); + PreparedAudioTranscriptionCall { + request: PreparedAudioTranscriptionRequest { + model: provider_info.model.to_string(), + custom_llm_provider: provider_info.custom_llm_provider.to_string(), + litellm_call_id: call_id, + audio: request.audio, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + }, + hooks: AudioTranscriptionLifecycleHooks::new( + CustomLoggerRunner::new(request.callbacks), + CustomGuardrailRunner::new(request.guardrails), + request.request_metadata, + ), + } +} + +fn new_audio_transcription_call_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_nanos()); + format!("audio-transcription-{timestamp}-{sequence}") +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs new file mode 100644 index 00000000000..5df04708b7d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs @@ -0,0 +1,53 @@ +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::thread; + +use serde_json::{Map, json}; + +use super::{AudioTranscriptionRequest, audio_transcription}; + +#[tokio::test] +async fn bedrock_request_is_signed_and_contains_audio() { + let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); + let address = listener.local_addr().expect("address"); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("connection"); + let mut request = Vec::new(); + let mut buffer = [0_u8; 16_384]; + let count = stream.read(&mut buffer).expect("request"); + request.extend_from_slice(&buffer[..count]); + let request = String::from_utf8_lossy(&request); + assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse")); + assert!(request.contains("authorization: AWS4-HMAC-SHA256")); + assert!(request.contains("x-amz-date:")); + assert!(request.contains("\"bytes\":\"AQI=\"")); + assert!(request.contains("Transcribe the audio. Respond with only the transcript.")); + let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}"; + stream.write_all(response).expect("response"); + }); + + let optional_params = Map::from_iter([ + ("aws_access_key_id".to_string(), json!("access-key")), + ("aws_secret_access_key".to_string(), json!("secret-key")), + ("aws_region_name".to_string(), json!("us-east-1")), + ]); + let api_base = format!("http://{address}"); + let response = audio_transcription(AudioTranscriptionRequest { + model: "mistral.voxtral-mini-3b-2507", + audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}), + api_key: None, + api_base: Some(&api_base), + custom_llm_provider: Some("bedrock"), + extra_headers: None, + optional_params, + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await + .expect("transcription"); + assert_eq!(response, json!({"text": "hello"})); + server.join().expect("server"); +} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs new file mode 100644 index 00000000000..9697aa98b0a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs @@ -0,0 +1,58 @@ +use std::sync::Arc; +use std::time::Duration; + +use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use serde_json::{Map, Value}; + +use crate::integrations::custom_guardrail::CustomGuardrail; +use crate::integrations::custom_logger::CustomLogger; +use crate::integrations::types::RequestMetadata; + +pub struct AudioTranscriptionRequest<'a> { + pub model: &'a str, + pub audio: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, + pub callbacks: Vec>, + pub guardrails: Vec>, + pub request_metadata: RequestMetadata, + pub litellm_call_id: Option<&'a str>, +} + +pub(crate) struct PreparedAudioTranscriptionRequest { + pub(crate) model: String, + pub(crate) custom_llm_provider: String, + pub(crate) litellm_call_id: String, + pub(crate) audio: Value, + pub(crate) api_key: Option, + pub(crate) api_base: Option, + pub(crate) extra_headers: Option>, + pub(crate) optional_params: Map, + pub(crate) timeout: Option, +} + +impl CallLifecycleRequest for PreparedAudioTranscriptionRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "audio_transcription", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} + +#[derive(Clone)] +pub(crate) struct ProviderAudioTranscriptionRequest { + pub(crate) model: String, + pub(crate) config: &'static dyn AudioTranscriptionProviderConfig, + pub(crate) url: String, + pub(crate) body: Value, + pub(crate) upstream_headers: Vec<(String, String)>, + pub(crate) timeout: Option, +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/client.rs b/litellm-rust/crates/ai-gateway/src/client.rs similarity index 60% rename from litellm-rust/crates/ai-gateway/src/ocr/client.rs rename to litellm-rust/crates/ai-gateway/src/client.rs index 79cc7816227..ff2606f0229 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/client.rs +++ b/litellm-rust/crates/ai-gateway/src/client.rs @@ -1,13 +1,13 @@ use std::sync::OnceLock; use std::time::Duration; -const OCR_TIMEOUT_SECS: u64 = 600; +const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600; -pub(super) fn http_client() -> &'static reqwest::Client { +pub(crate) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() - .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) + .timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS)) .build() .expect("failed to build reqwest client") }) diff --git a/litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs b/litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs new file mode 100644 index 00000000000..80d9e401a5f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs @@ -0,0 +1 @@ +pub use crate::audio_transcription::{AudioTranscriptionRequest, audio_transcription}; diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index 9cbfa568121..6129a808965 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,3 +1,4 @@ +pub mod audio_transcription; pub mod messages; pub mod ocr; pub mod realtime; diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 25aac3c495b..c44d661c29e 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -11,6 +11,8 @@ //! binary turns on. The `python-config` feature additionally pulls in [`python`] //! for the load-time config reader. +pub mod audio_transcription; +mod client; pub mod io; pub mod messages; pub mod ocr; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index 7d164a80137..9bc2818b6e7 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{ VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, }; -use super::client::http_client; +use crate::client::http_client; const ERROR_BODY_MAX_CHARS: usize = 256; const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 381d22e9cea..1de34eb400e 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -3,9 +3,9 @@ use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; -use super::client::http_client; use super::common_utils::{poll_document_intelligence, truncate_error_body}; use super::types::ProviderOcrRequest; +use crate::client::http_client; pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { let mut request_builder = http_client().post(&request.url).json(&request.body); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index ad346bc0c64..c4c13e2300c 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -2,7 +2,6 @@ use litellm_core::CoreResult; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; -mod client; mod common_utils; mod handler; mod hooks; diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs new file mode 100644 index 00000000000..eab34c13843 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -0,0 +1,57 @@ +use serde_json::{Map, Value}; + +use crate::CoreResult; + +use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AudioTranscriptionAuth { + Bearer, + AwsSigV4 { + region: String, + service: &'static str, + }, +} + +pub trait AudioTranscriptionProviderConfig: Sync { + fn supported_transcription_params(&self) -> &'static [&'static str]; + + fn map_transcription_params(&self, params: &Map) -> Map { + params + .iter() + .filter(|(key, _)| { + self.supported_transcription_params() + .contains(&key.as_str()) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() + } + + fn transform_transcription_request( + &self, + model: &str, + audio: Value, + optional_params: Map, + ) -> CoreResult; + + fn transform_transcription_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult; + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn auth_strategy( + &self, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs new file mode 100644 index 00000000000..3a9e1ecd88c --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -0,0 +1,20 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AudioTranscriptionRequestData { + pub body: Value, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct AudioTranscriptionResponseData { + pub text: String, +} + +impl AudioTranscriptionResponseData { + pub fn into_json(self) -> Value { + serde_json::json!({ + "text": self.text, + }) + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 3989fb441bc..51ea19750ea 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod audio_transcription; pub mod caching; pub mod call_lifecycle; pub mod constants; diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs new file mode 100644 index 00000000000..86eb589e2c0 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -0,0 +1,310 @@ +use serde_json::{Map, Value, json}; + +use crate::audio_transcription::transformation::{ + AudioTranscriptionAuth, AudioTranscriptionProviderConfig, +}; +use crate::audio_transcription::types::{ + AudioTranscriptionRequestData, AudioTranscriptionResponseData, +}; +use crate::error::{CoreError, CoreResult, json_type_name}; + +use super::aws_base::AwsAuthConfig; +use super::constants::{ + AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE, + DEFAULT_BEDROCK_REGION, +}; + +const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"]; + +pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig = + BedrockAudioTranscriptionConfig; + +pub struct BedrockAudioTranscriptionConfig; + +pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { + let mut stripped = model; + for prefix in ["bedrock/converse/", "bedrock/", "converse/"] { + if let Some(value) = stripped.strip_prefix(prefix) { + stripped = value; + break; + } + } + let mut region = None; + if let Some((candidate, remainder)) = stripped.split_once('/') + && is_bedrock_region(candidate) + { + region = Some(candidate.to_string()); + stripped = remainder; + } + for prefix in ["nova-2/", "nova/"] { + if let Some(value) = stripped.strip_prefix(prefix) { + stripped = value; + break; + } + } + if region.is_none() { + region = stripped + .strip_prefix("arn:") + .and_then(|value| value.split(':').nth(3)) + .filter(|value| !value.is_empty()) + .map(str::to_string); + } + (stripped.to_string(), region) +} + +fn is_bedrock_region(value: &str) -> bool { + value.len() > 3 + && value.contains('-') + && value + .chars() + .all(|char| char.is_ascii_alphanumeric() || char == '-') +} + +pub fn resolve_bedrock_region( + model_region: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + if let Some(region) = optional_params + .get("aws_region_name") + .and_then(Value::as_str) + { + return region.to_string(); + } + if let Some(region) = model_region { + return region.to_string(); + } + env_lookup(AWS_REGION_NAME) + .or_else(|| env_lookup(AWS_REGION)) + .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) +} + +fn audio_fields(audio: Value) -> CoreResult<(String, String)> { + let object = audio.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&audio), + })?; + let data = object + .get("data") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or(CoreError::MissingField("audio.data"))?; + let format = object + .get("format") + .and_then(Value::as_str) + .filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg")) + .ok_or_else(|| { + CoreError::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string()) + })?; + Ok((data.to_string(), format.to_string())) +} + +fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a str> { + params + .get(key) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) +} + +impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { + fn supported_transcription_params(&self) -> &'static [&'static str] { + SUPPORTED_PARAMS + } + + fn transform_transcription_request( + &self, + _model: &str, + audio: Value, + optional_params: Map, + ) -> CoreResult { + let (data, format) = audio_fields(audio)?; + let mut instruction = "Transcribe the audio. Respond with only the transcript.".to_string(); + if let Some(language) = optional_string(&optional_params, "language") { + instruction.push_str(&format!(" The audio language is {language}.")); + } + if let Some(prompt) = optional_string(&optional_params, "prompt") { + instruction.push_str(&format!(" Additional context: {prompt}")); + } + let mut inference_config = Map::from_iter([("maxTokens".to_string(), json!(4096))]); + if let Some(temperature) = optional_params.get("temperature") { + inference_config.insert("temperature".to_string(), temperature.clone()); + } + Ok(AudioTranscriptionRequestData { + body: json!({ + "messages": [{ + "role": "user", + "content": [ + {"audio": {"format": format, "source": {"bytes": data}}}, + {"text": instruction} + ] + }], + "system": [{"text": "You are a transcription assistant."}], + "inferenceConfig": inference_config, + }), + }) + } + + fn transform_transcription_response( + &self, + _model: &str, + response_json: Value, + ) -> CoreResult { + let content = response_json + .get("output") + .and_then(|value| value.get("message")) + .and_then(|value| value.get("content")) + .and_then(Value::as_array) + .ok_or_else(|| { + CoreError::InvalidResponse("Bedrock response has no output content".to_string()) + })?; + let mut text = String::new(); + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + Ok(AudioTranscriptionResponseData { text }) + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let (model_id, model_region) = bedrock_model_id_and_region(model); + let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); + let endpoint = optional_params + .get("aws_bedrock_runtime_endpoint") + .and_then(Value::as_str) + .or(api_base) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion)); + Ok(format!( + "{}/model/{model_id}/converse", + endpoint.trim_end_matches('/') + )) + } + + fn auth_strategy( + &self, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let (_, model_region) = bedrock_model_id_and_region(model); + Ok(AudioTranscriptionAuth::AwsSigV4 { + region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), + service: BEDROCK_SERVICE, + }) + } +} + +pub fn aws_auth_config( + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> AwsAuthConfig { + let value = |key: &str| { + optional_params + .get(key) + .and_then(Value::as_str) + .map(str::to_string) + }; + let env = |key: &str| env_lookup(key); + AwsAuthConfig { + access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")), + secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")), + session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")), + region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)), + session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")), + profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")), + role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")), + web_identity_token: value("aws_web_identity_token") + .or_else(|| env("AWS_WEB_IDENTITY_TOKEN")), + sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")), + external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + #[test] + fn request_matches_python_shape() { + let params = Map::from_iter([ + ("language".to_string(), json!("en")), + ("prompt".to_string(), json!("Speaker names")), + ("temperature".to_string(), json!(0)), + ("timestamp_granularities".to_string(), json!(["word"])), + ]); + let params = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.map_transcription_params(¶ms); + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG + .transform_transcription_request( + "mistral.voxtral-mini-3b-2507", + json!({"data": "AQI=", "format": "wav", "filename": "sample.wav"}), + params, + ) + .expect("request"); + assert_eq!( + result.body, + json!({ + "messages": [{ + "role": "user", + "content": [ + {"audio": {"format": "wav", "source": {"bytes": "AQI="}}}, + {"text": "Transcribe the audio. Respond with only the transcript. The audio language is en. Additional context: Speaker names"} + ] + }], + "system": [{"text": "You are a transcription assistant."}], + "inferenceConfig": {"maxTokens": 4096, "temperature": 0} + }) + ); + } + + #[test] + fn response_concatenates_content_blocks() { + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG + .transform_transcription_response( + "model", + json!({"output": {"message": {"content": [{"text": "hello "}, {"text": "world"}]}}}), + ) + .expect("response"); + assert_eq!(result.text, "hello world"); + assert_eq!(result.into_json(), json!({"text": "hello world"})); + } + + #[test] + fn invalid_audio_is_rejected() { + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_transcription_request( + "model", + json!({"data": "AQI="}), + Map::new(), + ); + assert!(result.is_err()); + } + + #[test] + fn region_and_url_precedence_match_python() { + let params = Map::from_iter([("aws_region_name".to_string(), json!("eu-west-1"))]); + let url = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG + .complete_url( + None, + "bedrock/us-east-1/mistral.voxtral-mini-3b-2507", + ¶ms, + &no_env, + ) + .expect("url"); + assert_eq!( + url, + "https://bedrock-runtime.eu-west-1.amazonaws.com/model/mistral.voxtral-mini-3b-2507/converse" + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index c5995732e41..dc036a3cf21 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -52,7 +52,7 @@ pub struct AwsAuthConfig { } impl AwsAuthConfig { - fn with_environment(self, env_lookup: &dyn Fn(&str) -> Option) -> Self { + fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option + Sync)) -> Self { Self { access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), secret_access_key: self @@ -144,7 +144,7 @@ fn same_role_arns(target: &str, caller: &str) -> bool { pub fn classify_auth( config: AwsAuthConfig, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> AwsAuthFlow { let config = config.with_environment(env_lookup); if let (Some(token), Some(role), Some(session_name)) = ( @@ -194,7 +194,7 @@ pub fn classify_auth( pub async fn resolve_credentials( config: AwsAuthConfig, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> CoreResult { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs index a08ae9de146..785295207e7 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/constants.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/constants.rs @@ -2,6 +2,7 @@ pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID"; pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; +pub const AWS_REGION: &str = "AWS_REGION"; pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; @@ -12,3 +13,6 @@ pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; pub const BEDROCK_SERVICE: &str = "bedrock"; pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session"; +pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2"; +pub const BEDROCK_RUNTIME_ENDPOINT_TEMPLATE: &str = + "https://bedrock-runtime.{region}.amazonaws.com"; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs index 8027260a78b..b09675ad7dd 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/mod.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/mod.rs @@ -2,5 +2,7 @@ //! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled //! separately. +#[cfg(feature = "bedrock-auth")] +pub mod audio_transcription; pub mod aws_base; mod constants; diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index efa1a554c9c..e5d021ec25b 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -17,7 +17,13 @@ Python-compatible dictionaries. - Provider dispatch belongs in Rust route modules such as `litellm_providers::ocr`, not in this PyO3 crate. - Python owns rollout state and fallback. Rust should return errors; Python - decides whether to raise or fall back. + decides whether to raise or fall back. For a rust-only provider/route (no + Python reference), the Python side is a thin dispatch that calls Rust and + raises when the bridge is unavailable, with no fallback. +- Keep the Python interface minimal (well under 100 lines per route): it only + marshals inputs and calls Rust. Do not add per-route feature flags, and do + not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch + class under `litellm/llms///`. ## Data Handling diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 83e163c38f1..20a9ba789ce 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -10,7 +10,7 @@ name = "_native" crate-type = ["cdylib"] [dependencies] -litellm-core.workspace = true +litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } pyo3 = { workspace = true, features = ["extension-module"] } pyo3-async-runtimes.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 07429667644..ee9bdd0b81f 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,6 +1,9 @@ use std::collections::HashMap; use std::time::Duration; +use litellm_ai_gateway::io::audio_transcription::{ + AudioTranscriptionRequest, audio_transcription as run_audio_transcription, +}; use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages}; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; @@ -244,6 +247,93 @@ fn aocr( }) } +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn transcription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let audio = py_to_json(py, audio.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + let result = gil::release_gil(py, || { + pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( + AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }, + )) + }); + match result { + Ok(value) => json_to_py(py, value), + Err(err) => Err(core_error_to_pyerr(err)), + } +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn atranscription( + py: Python<'_>, + model: String, + audio: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let audio = py_to_json(py, audio.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = run_audio_transcription(AudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: custom_llm_provider.as_deref(), + extra_headers, + optional_params, + timeout, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + }) + .await + .map_err(core_error_to_pyerr)?; + Python::attach(|py| json_to_py(py, value)) + }) +} + type MarshaledMessagesInputs = (Value, Option>, Option); fn marshal_messages_inputs( @@ -341,6 +431,8 @@ fn gil_stats(py: Python<'_>) -> PyResult> { fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; module.add_function(wrap_pyfunction!(aocr, module)?)?; + module.add_function(wrap_pyfunction!(transcription, module)?)?; + module.add_function(wrap_pyfunction!(atranscription, module)?)?; module.add_function(wrap_pyfunction!(messages, module)?)?; module.add_function(wrap_pyfunction!(amessages, module)?)?; module.add_class::()?; diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py new file mode 100644 index 00000000000..f2e58df3015 --- /dev/null +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -0,0 +1,84 @@ +import base64 +from typing import Union + +import httpx + +from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.rust_bridge import transcription as rust_transcription_bridge +from litellm.types.utils import FileTypes, TranscriptionResponse + + +class BedrockAudioTranscriptionRustDispatch: + @staticmethod + def _audio_payload(audio_file: FileTypes) -> dict[str, object]: + processed_audio = process_audio_file(audio_file) + formats = { + "audio/flac": "flac", + "audio/mpeg": "mp3", + "audio/mp3": "mp3", + "audio/ogg": "ogg", + "audio/wav": "wav", + "audio/x-wav": "wav", + } + audio_format = formats.get(processed_audio.content_type) or ( + processed_audio.filename.rsplit(".", 1)[-1].lower() if "." in processed_audio.filename else "" + ) + if audio_format not in {"wav", "mp3", "flac", "ogg"}: + raise ValueError(f"Unsupported Bedrock audio format for file {processed_audio.filename!r}") + return { + "data": base64.b64encode(processed_audio.file_content).decode("ascii"), + "format": audio_format, + "filename": processed_audio.filename, + } + + def audio_transcriptions( + self, + *, + model: str, + audio_file: FileTypes, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, + ) -> TranscriptionResponse: + rust_response = rust_transcription_bridge.transcription( + model=model, + audio=self._audio_payload(audio_file), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) + if rust_response is None: + raise RuntimeError("Rust audio transcription bridge is unavailable") + return TranscriptionResponse(**rust_response) + + async def async_audio_transcriptions( + self, + *, + model: str, + audio_file: FileTypes, + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, + ) -> TranscriptionResponse: + rust_response = await rust_transcription_bridge.atranscription( + model=model, + audio=self._audio_payload(audio_file), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) + if rust_response is None: + raise RuntimeError("Rust audio transcription bridge is unavailable") + return TranscriptionResponse(**rust_response) diff --git a/litellm/main.py b/litellm/main.py index 3584297b35f..fb05a375111 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -27,7 +27,6 @@ from typing import ( TYPE_CHECKING, Any, AsyncIterator, - Callable, Coroutine, Dict, Iterable, @@ -81,22 +80,19 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.chat_completion_agentic_loop import ( - maybe_run_chat_completion_agentic_loop, -) from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, ) -from litellm.litellm_core_utils.completion_timeout import CompletionTimeout -from litellm.litellm_core_utils.request_timeout_resolver import ( - get_configured_request_timeout, +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, ) +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, OPTIONAL_KWARGS_KEYS, ) -from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -112,6 +108,9 @@ from litellm.litellm_core_utils.mock_functions import ( from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) +from litellm.litellm_core_utils.request_timeout_resolver import ( + get_configured_request_timeout, +) from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -213,7 +212,6 @@ from .llms.bedrock.embed.embedding import BedrockEmbedding from .llms.bedrock.image_edit.handler import BedrockImageEdit from .llms.bedrock.image_generation.image_handler import BedrockImageGeneration from .llms.bytez.chat.transformation import BytezChatConfig -from .llms.gdc.chat.transformation import GDCGeminiConfig from .llms.clarifai.chat.transformation import ClarifaiConfig from .llms.codestral.completion.handler import CodestralTextCompletion from .llms.cohere.embed import handler as cohere_embed @@ -222,24 +220,25 @@ from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from .llms.custom_llm import CustomLLM, custom_chat_llm_router from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha, palm +from .llms.gdc.chat.transformation import GDCGeminiConfig from .llms.gemini.common_utils import get_api_key_from_env from .llms.groq.chat.handler import GroqChatCompletion from .llms.heroku.chat.transformation import HerokuChatConfig from .llms.huggingface.embedding.handler import HuggingFaceEmbedding from .llms.lemonade.chat.transformation import LemonadeChatConfig from .llms.nlp_cloud.chat.handler import completion as nlp_cloud_chat_completion -from .llms.oci.chat.transformation import OCIChatConfig -from .llms.ollama.completion import handler as ollama -from .llms.oobabooga.chat import oobabooga -from .llms.openai.completion.handler import OpenAITextCompletion -from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler -from .llms.openai.openai import OpenAIChatCompletion from .llms.nvidia_riva.audio_transcription.handler import ( NvidiaRivaAudioTranscription, ) from .llms.nvidia_riva.audio_transcription.transformation import ( NvidiaRivaAudioTranscriptionConfig, ) +from .llms.oci.chat.transformation import OCIChatConfig +from .llms.ollama.completion import handler as ollama +from .llms.oobabooga.chat import oobabooga +from .llms.openai.completion.handler import OpenAITextCompletion +from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler +from .llms.openai.openai import OpenAIChatCompletion from .llms.openai.transcriptions.handler import OpenAIAudioTranscription from .llms.openai_like.chat.handler import OpenAILikeChatHandler from .llms.openai_like.embedding.handler import OpenAILikeEmbeddingHandler @@ -7722,6 +7721,32 @@ def transcription( headers=extra_headers, provider_config=provider_config, # type: ignore[arg-type] ) + elif custom_llm_provider == "bedrock": + from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch + + dispatch = BedrockAudioTranscriptionRustDispatch() + if atranscription: + response = dispatch.async_audio_transcriptions( + model=model, + audio_file=file, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) + else: + response = dispatch.audio_transcriptions( + model=model, + audio_file=file, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout=timeout, + ) elif provider_config is not None: response = base_llm_http_handler.audio_transcriptions( model=model, diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index e9139a634f1..91aac6c1232 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -72,17 +72,28 @@ def use_litellm_rust( messages: RustMessages | None | _Unset = _UNSET, amessages: RustAmessages | None | _Unset = _UNSET, responses_websocket: Any | None | _Unset = _UNSET, + transcription: Any | None | _Unset = _UNSET, + atranscription: Any | None | _Unset = _UNSET, ) -> None: global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl configuring_ocr = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset) configuring_messages = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset) configuring_responses_websocket = not isinstance(responses_websocket, _Unset) + configuring_transcription = not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset) if configuring_ocr or (not configuring_messages and not configuring_responses_websocket): _rust_ocr_enabled = enabled if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): _rust_aocr_impl = aocr + if configuring_transcription: + from litellm.rust_bridge.transcription import configure_rust_transcription + + configure_rust_transcription( + enabled=enabled, + transcription=transcription, + atranscription=atranscription, + ) if not configuring_messages and not configuring_responses_websocket: return if configuring_messages: diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription.py new file mode 100644 index 00000000000..44bb42e4104 --- /dev/null +++ b/litellm/rust_bridge/transcription.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Awaitable, Final, Protocol, Union, cast + +import httpx + +from litellm.rust_bridge.timeouts import timeout_to_seconds + + +class RustTranscription(Protocol): + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + raise NotImplementedError + + +class RustAtranscription(Protocol): + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: + raise NotImplementedError + + +class _Unset: + pass + + +_UNSET: Final[_Unset] = _Unset() + + +@dataclass +class _RustTranscriptionState: + transcription: RustTranscription | None = None + atranscription: RustAtranscription | None = None + + +_STATE = _RustTranscriptionState() + + +def configure_rust_transcription( + enabled: bool = True, + *, + transcription: RustTranscription | None | _Unset = _UNSET, + atranscription: RustAtranscription | None | _Unset = _UNSET, +) -> None: + if not isinstance(transcription, _Unset): + _STATE.transcription = transcription + if not isinstance(atranscription, _Unset): + _STATE.atranscription = atranscription + + +def load_rust_transcription() -> RustTranscription | None: + if _STATE.transcription is not None: + return _STATE.transcription + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + return ( + None + if native_bridge is None + else cast( # cast-ok: native extension protocol is runtime-defined + RustTranscription, getattr(native_bridge, "transcription", None) + ) + ) + + +def load_rust_atranscription() -> RustAtranscription | None: + if _STATE.atranscription is not None: + return _STATE.atranscription + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + return ( + None + if native_bridge is None + else cast( # cast-ok: native extension protocol is runtime-defined + RustAtranscription, getattr(native_bridge, "atranscription", None) + ) + ) + + +def transcription( + *, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_transcription = load_rust_transcription() + if rust_transcription is None: + return None + return rust_transcription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_to_seconds(timeout), + ) + + +async def atranscription( + *, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout: Union[float, httpx.Timeout] | None, +) -> dict[str, object] | None: + rust_atranscription = load_rust_atranscription() + if rust_atranscription is None: + return None + return await rust_atranscription( + model=model, + audio=audio, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_to_seconds(timeout), + ) diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py new file mode 100644 index 00000000000..bbeb6c38f78 --- /dev/null +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -0,0 +1,151 @@ +import importlib + +import pytest + +import litellm +from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch + +rust_bridge = importlib.import_module("litellm.rust_bridge.transcription") + + +class SyncBridge: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append({"model": model, "audio": audio, "optional_params": optional_params}) + return {"text": "hello"} + + +class AsyncBridge: + async def __call__( + self, + model: str, + audio: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> dict[str, object]: + return {"text": "async"} + + +def test_enabled_sync_bridge_receives_audio() -> None: + bridge = SyncBridge() + rust_bridge.configure_rust_transcription(True, transcription=bridge) + result = rust_bridge.transcription( + model="mistral.voxtral-mini-3b-2507", + audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={"temperature": 0}, + timeout=5.0, + ) + assert result == {"text": "hello"} + assert bridge.calls[0]["audio"] == {"data": "AQI=", "format": "wav", "filename": "audio.wav"} + + +@pytest.mark.asyncio +async def test_enabled_async_bridge() -> None: + rust_bridge.configure_rust_transcription(True, atranscription=AsyncBridge()) + result = await rust_bridge.atranscription( + model="mistral.voxtral-mini-3b-2507", + audio={"data": "AQI=", "format": "wav", "filename": "audio.wav"}, + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={}, + timeout=None, + ) + assert result == {"text": "async"} + + +def test_loader_returns_none_without_native_extension(monkeypatch: pytest.MonkeyPatch) -> None: + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) + monkeypatch.setattr("litellm.rust_bridge.get_native_bridge", lambda: None) + assert rust_bridge.load_rust_transcription() is None + assert rust_bridge.load_rust_atranscription() is None + + +def test_dispatch_sync_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(rust_bridge, "transcription", lambda **_: None) + + with pytest.raises(RuntimeError, match="bridge is unavailable"): + BedrockAudioTranscriptionRustDispatch().audio_transcriptions( + model="bedrock/mistral.voxtral-mini-3b-2507", + audio_file=("audio.wav", b"audio", "audio/wav"), + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={}, + timeout=5, + ) + + +@pytest.mark.asyncio +async def test_dispatch_async_path_requires_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + async def unavailable(**_: object) -> None: + return None + + monkeypatch.setattr(rust_bridge, "atranscription", unavailable) + + with pytest.raises(RuntimeError, match="bridge is unavailable"): + await BedrockAudioTranscriptionRustDispatch().async_audio_transcriptions( + model="bedrock/mistral.voxtral-mini-3b-2507", + audio_file=("audio.wav", b"audio", "audio/wav"), + api_key=None, + api_base=None, + custom_llm_provider="bedrock", + extra_headers=None, + optional_params={}, + timeout=5, + ) + + +def test_bedrock_transcription_uses_rust_only_path() -> None: + rust_bridge.configure_rust_transcription( + transcription=lambda **_: {"text": "rust"}, + atranscription=None, + ) + try: + response = litellm.transcription( + model="bedrock/mistral.voxtral-mini-3b-2507", + file=("audio.wav", b"audio", "audio/wav"), + ) + finally: + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) + + assert response.text == "rust" + + +@pytest.mark.asyncio +async def test_bedrock_atranscription_uses_rust_only_path() -> None: + async def rust_response(**_: object) -> dict[str, object]: + return {"text": "rust"} + + rust_bridge.configure_rust_transcription(transcription=None, atranscription=rust_response) + try: + response = await litellm.atranscription( + model="bedrock/mistral.voxtral-mini-3b-2507", + file=("audio.wav", b"audio", "audio/wav"), + ) + finally: + rust_bridge.configure_rust_transcription(transcription=None, atranscription=None) + + assert response.text == "rust"