From 21d7aae62f9c997a62e2187aad14a9bcd8e7eb4c Mon Sep 17 00:00:00 2001 From: Aastha Joshi Date: Tue, 11 Aug 2026 00:39:17 -0700 Subject: [PATCH] refactor(rust): make transcription state transitions immutable --- .../src/audio_transcription/hooks.rs | 16 ++-- .../core/src/audio_transcription/handler.rs | 95 ++++++++++--------- .../core/src/audio_transcription/prepare.rs | 16 +++- .../core/src/audio_transcription/types.rs | 9 +- 4 files changed, 71 insertions(+), 65 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 36da54555e5..2198ac7a7d6 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -69,17 +69,17 @@ impl AudioTranscriptionLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; - let Value::Object(mut data) = guardrail_request.data else { + let Value::Object(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(|| { + let audio = data.get("audio").cloned().ok_or_else(|| { CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string()) })?; - let optional_params = match data.remove("optional_params") { + let optional_params = match data.get("optional_params").cloned() { Some(Value::Object(value)) => value, Some(_) => { return Err(CoreError::InvalidRequest( @@ -100,7 +100,7 @@ impl AudioTranscriptionLifecycleHooks { &self, request: PreparedAudioTranscriptionRequest, ) -> CoreResult { - let mut provider_request = + let provider_request = prepare_audio_transcription_request(CoreAudioTranscriptionRequest { model: &request.model, audio: request.audio, @@ -135,19 +135,17 @@ impl AudioTranscriptionLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; - let Value::Object(mut data) = guardrail_request.data else { + let Value::Object(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(|| { + let body = data.get("body").cloned().ok_or_else(|| { CoreError::InvalidRequest("audio transcription guardrail removed body".to_string()) })?; - provider_request.replace_body(body); - - Ok(provider_request) + Ok(provider_request.with_body(body)) } fn logging_payload( diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 23ed5c5586c..59636966767 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,12 +1,13 @@ use crate::CoreResult; use crate::error::CoreError; +#[cfg(feature = "bedrock-auth")] +use std::collections::BTreeMap; #[cfg(feature = "bedrock-auth")] use std::time::SystemTime; #[cfg(feature = "bedrock-auth")] use crate::providers::bedrock::audio_transcription::aws_auth_config; - #[cfg(feature = "bedrock-auth")] use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; @@ -17,23 +18,23 @@ use super::transformation::AudioTranscriptionAuth; use super::types::{AudioTranscriptionResponseData, PreparedAudioTranscriptionRequest}; pub(super) async fn execute_audio_transcription_provider_call( - mut request: PreparedAudioTranscriptionRequest, + request: PreparedAudioTranscriptionRequest, ) -> CoreResult { - sign_request(&mut request).await?; - 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); + let headers = authenticated_headers(&request, &body).await?; - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } + let request_builder = headers.into_iter().fold( + http_client().post(&request.url).body(body), + |builder, (key, value)| builder.header(key, value), + ); - if let Some(timeout) = request.timeout { - request_builder = request_builder.timeout(timeout); - } + let request_builder = match request.timeout { + Some(timeout) => request_builder.timeout(timeout), + None => request_builder, + }; let response = request_builder .send() @@ -65,7 +66,10 @@ pub(super) async fn execute_audio_transcription_provider_call( .transform_transcription_response(&request.model, response_json) } -async fn sign_request(request: &mut PreparedAudioTranscriptionRequest) -> CoreResult<()> { +async fn authenticated_headers( + request: &PreparedAudioTranscriptionRequest, + body: &[u8], +) -> CoreResult> { let env_lookup = environment_lookup; let auth = @@ -74,45 +78,42 @@ async fn sign_request(request: &mut PreparedAudioTranscriptionRequest) -> CoreRe .auth_strategy(&request.model, &request.optional_params, &env_lookup)?; match auth { - AudioTranscriptionAuth::Bearer => { - ensure_content_type(request); - Ok(()) - } - + AudioTranscriptionAuth::Bearer => Ok(with_content_type(&request.upstream_headers)), AudioTranscriptionAuth::AwsSigV4 { region, .. } => { - sign_aws_sigv4_request(request, ®ion).await + aws_sigv4_headers(request, body, ®ion).await } } } -fn ensure_content_type(request: &mut PreparedAudioTranscriptionRequest) { - if !request - .upstream_headers +fn with_content_type(headers: &[(String, String)]) -> Vec<(String, String)> { + if headers .iter() .any(|(key, _)| key.eq_ignore_ascii_case("content-type")) { - request - .upstream_headers - .push(("Content-Type".to_string(), "application/json".to_string())); + headers.to_vec() + } else { + headers + .iter() + .cloned() + .chain(std::iter::once(( + "Content-Type".to_string(), + "application/json".to_string(), + ))) + .collect() } } #[cfg(feature = "bedrock-auth")] -async fn sign_aws_sigv4_request( - request: &mut PreparedAudioTranscriptionRequest, +async fn aws_sigv4_headers( + request: &PreparedAudioTranscriptionRequest, + body: &[u8], region: &str, -) -> CoreResult<()> { +) -> CoreResult> { let env_lookup = environment_lookup; - let body = serde_json::to_vec(&request.body).map_err(|error| { - CoreError::InvalidRequest(format!("invalid audio request body: {error}")) - })?; - - let mut headers = std::collections::BTreeMap::new(); - - headers.insert("Content-Type".to_string(), "application/json".to_string()); - - headers.extend(request.upstream_headers.iter().cloned()); + let headers = std::iter::once(("Content-Type".to_string(), "application/json".to_string())) + .chain(request.upstream_headers.iter().cloned()) + .collect::>(); let credentials = resolve_credentials( aws_auth_config(&request.optional_params, &env_lookup), @@ -120,25 +121,29 @@ async fn sign_aws_sigv4_request( ) .await?; - headers.extend(sign_bedrock_post( + let signed_headers = sign_bedrock_post( &request.url, - &body, + body, &headers, region, &credentials, SystemTime::now(), - )?); + )?; - request.upstream_headers = headers.into_iter().collect(); - - Ok(()) + Ok(headers + .into_iter() + .chain(signed_headers.into_iter()) + .collect::>() + .into_iter() + .collect()) } #[cfg(not(feature = "bedrock-auth"))] -async fn sign_aws_sigv4_request( - _request: &mut PreparedAudioTranscriptionRequest, +async fn aws_sigv4_headers( + _request: &PreparedAudioTranscriptionRequest, + _body: &[u8], _region: &str, -) -> CoreResult<()> { +) -> CoreResult> { Err(CoreError::InvalidProvider( "AWS SigV4 authentication requires the `bedrock-auth` feature".to_string(), )) diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 842057e8b8f..831fc014498 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -22,7 +22,7 @@ pub(super) fn prepare_audio_transcription_call( let env_lookup = environment_lookup; - let mut headers = string_headers(request.extra_headers)?; + let headers = string_headers(request.extra_headers)?; let url = config.complete_url( request.api_base, @@ -38,12 +38,20 @@ pub(super) fn prepare_audio_transcription_call( let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?; - if matches!(auth, AudioTranscriptionAuth::Bearer) + let headers = if matches!(auth, AudioTranscriptionAuth::Bearer) && !has_header(&headers, "authorization") && let Some(api_key) = request.api_key { - headers.insert("Authorization".to_string(), format!("Bearer {api_key}")); - } + headers + .into_iter() + .chain(std::iter::once(( + "Authorization".to_string(), + format!("Bearer {api_key}"), + ))) + .collect() + } else { + headers + }; Ok(PreparedAudioTranscriptionRequest { model, diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index d0a888f1dd5..2a5a25a4bac 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -54,13 +54,8 @@ impl PreparedAudioTranscriptionRequest { &self.body } - /// Replace the provider-transformed request body. - /// - /// This exists for host lifecycle integrations such as during-call - /// guardrails. Authentication is intentionally performed after this step - /// so signatures always cover the final body sent upstream. - pub fn replace_body(&mut self, body: Value) { - self.body = body; + pub fn with_body(self, body: Value) -> Self { + Self { body, ..self } } }