diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 726d2f484da..72f5e70eea5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1960,6 +1960,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-types", "litellm-auth", + "litellm-http", "moka", "reqwest 0.12.28", "serde_json", @@ -2142,6 +2143,7 @@ dependencies = [ "reqwest 0.12.28", "rstest", "rustls 0.23.42", + "serde", "serde_json", "thiserror 2.0.19", "tokio", diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index d998b647960..1f27c7bc990 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-auth.workspace = true +litellm-http.workspace = true moka = { workspace = true, features = ["sync"] } serde_json.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index 3b6b73bc6a9..bbcb0f016c8 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -20,8 +20,7 @@ use super::constants::{ AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME, AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, - BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, - SIGV4_COMPUTED_HEADER_NAMES, + DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES, }; const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); @@ -451,11 +450,12 @@ pub fn is_sigv4_computed_header(name: &str) -> bool { SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str()) } -pub fn sign_bedrock_post( +pub fn sign_post( url: &str, body: &[u8], headers: &BTreeMap, region: &str, + service: &str, credentials: &Credentials, signing_time: SystemTime, ) -> Result, Error> { @@ -463,7 +463,7 @@ pub fn sign_bedrock_post( let params = v4::SigningParams::builder() .identity(&identity) .region(region) - .name(BEDROCK_SERVICE) + .name(service) .time(signing_time) .settings(SigningSettings::default()) .build() @@ -534,22 +534,28 @@ fn is_bedrock_region(value: &str) -> bool { .all(|char| char.is_ascii_alphanumeric() || char == '-') } +/// The region a caller configured: `aws_region_name`, then the model's own +/// region, then the environment. Each service decides what a missing one means. +pub fn resolve_aws_region( + model_region: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + optional_params + .get("aws_region_name") + .and_then(Value::as_str) + .or(model_region) + .map(str::to_string) + .or_else(|| env_lookup(AWS_REGION_NAME)) + .or_else(|| env_lookup(AWS_REGION)) +} + 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)) + resolve_aws_region(model_region, optional_params, env_lookup) .unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string()) } @@ -609,11 +615,36 @@ pub fn host_supplied_credentials(optional_params: &Map) -> Option #[cfg(test)] mod tests { use super::*; + use crate::constants::BEDROCK_SERVICE; fn no_env(_: &str) -> Option { None } + #[test] + fn a_region_comes_from_the_call_then_the_model_then_the_environment() { + let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]); + let region_name = |key: &str| (key == AWS_REGION_NAME).then(|| "ap-south-1".to_string()); + let region = |key: &str| (key == AWS_REGION).then(|| "sa-east-1".to_string()); + + let resolved = [ + resolve_aws_region(Some("us-east-2"), ¶ms, ®ion_name), + resolve_aws_region(Some("us-east-2"), &Map::new(), ®ion_name), + resolve_aws_region(None, &Map::new(), ®ion_name), + resolve_aws_region(None, &Map::new(), ®ion), + resolve_aws_region(None, &Map::new(), &no_env), + ]; + + assert_eq!( + resolved.map(|region| region.unwrap_or_else(|| "none".into())), + ["eu-west-1", "us-east-2", "ap-south-1", "sa-east-1", "none"] + ); + assert_eq!( + resolve_bedrock_region(None, &Map::new(), &no_env), + DEFAULT_BEDROCK_REGION + ); + } + fn parity_inputs() -> (String, Vec, BTreeMap) { ( "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" @@ -811,11 +842,12 @@ mod tests { None, "test", ); - let signed = sign_bedrock_post( + let signed = sign_post( &url, &body, &signable, "us-east-1", + BEDROCK_SERVICE, &credentials, SystemTime::UNIX_EPOCH, ) @@ -843,11 +875,12 @@ mod tests { None, "test", ); - let signed = sign_bedrock_post( + let signed = sign_post( &url, &body, &headers, "us-east-1", + BEDROCK_SERVICE, &credentials, UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), ) @@ -878,11 +911,12 @@ mod tests { None, "test", ); - let signed = sign_bedrock_post( + let signed = sign_post( &url, &body, &headers, "us-east-1", + BEDROCK_SERVICE, &credentials, UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), ) @@ -915,11 +949,12 @@ mod tests { let url = format!( "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" ); - let signed_headers = sign_bedrock_post( + let signed_headers = sign_post( &url, &body, &headers, region, + BEDROCK_SERVICE, &credentials, SystemTime::now(), )?; diff --git a/litellm-rust/crates/auth-aws/src/lib.rs b/litellm-rust/crates/auth-aws/src/lib.rs index 264592ccb2e..0fe0b390110 100644 --- a/litellm-rust/crates/auth-aws/src/lib.rs +++ b/litellm-rust/crates/auth-aws/src/lib.rs @@ -1,6 +1,9 @@ mod aws; pub mod constants; mod error; +mod signer; pub use aws::*; +pub use aws_credential_types::Credentials; pub use error::Error; +pub use signer::SigV4Signer; diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs new file mode 100644 index 00000000000..46d6fb5c391 --- /dev/null +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -0,0 +1,183 @@ +use std::{collections::BTreeMap, time::SystemTime}; + +use aws_credential_types::Credentials; +use litellm_http::outbound::{RequestSigner, UnsignedRequest}; +use serde_json::{Map, Value}; + +use crate::{ + Error, aws_auth_config, aws_signature_headers, host_supplied_credentials, + is_sigv4_computed_header, resolve_credentials, sign_post, +}; + +/// SigV4 over the serialized body. Credentials are resolved up front, since +/// they do not depend on the body; the signature waits for the final bytes. +#[derive(Clone, Debug)] +pub struct SigV4Signer { + region: String, + service: &'static str, + credentials: Credentials, + clock: fn() -> SystemTime, +} + +impl SigV4Signer { + pub fn new(region: String, service: &'static str, credentials: Credentials) -> Self { + Self { + region, + service, + credentials, + clock: SystemTime::now, + } + } + + pub fn with_clock(self, clock: fn() -> SystemTime) -> Self { + Self { clock, ..self } + } + + /// A host with its own resolution chain hands credentials down in + /// `optional_params`; only derive them here when it supplied none. + pub async fn resolve( + region: String, + service: &'static str, + optional_params: &Map, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + let credentials = match host_supplied_credentials(optional_params) { + Some(credentials) => credentials, + None => { + resolve_credentials(aws_auth_config(optional_params, env_lookup), env_lookup) + .await? + } + }; + Ok(Self::new(region, service, credentials)) + } +} + +impl RequestSigner for SigV4Signer { + fn sign( + &self, + request: UnsignedRequest<'_>, + ) -> Result, litellm_http::Error> { + // Sending a caller's copy next to the computed one is rejected by AWS. + if let Some((name, _)) = request + .headers + .iter() + .find(|(name, _)| is_sigv4_computed_header(name)) + { + return Err(litellm_http::Error::ComputedHeader(name.clone())); + } + let headers: BTreeMap = request.headers.iter().cloned().collect(); + sign_post( + request.url, + request.body, + &aws_signature_headers(&headers), + &self.region, + self.service, + &self.credentials, + (self.clock)(), + ) + .map(|signature| signature.into_iter().collect()) + .map_err(|error| litellm_http::Error::Signature(error.to_string())) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, UNIX_EPOCH}; + + use litellm_http::outbound::OutboundRequest; + use serde_json::json; + + use super::*; + + fn fixed_clock() -> SystemTime { + UNIX_EPOCH + Duration::from_secs(1_700_000_000) + } + + fn signer(service: &'static str) -> SigV4Signer { + SigV4Signer::new( + "us-east-1".into(), + service, + Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + ) + .with_clock(fixed_clock) + } + + fn authorization(body: &Value, service: &'static str) -> String { + OutboundRequest::signed_json( + "https://textract.us-east-1.amazonaws.com/".into(), + vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], + body, + None, + &signer(service), + ) + .unwrap() + .header("Authorization") + .unwrap() + .to_string() + } + + #[test] + fn the_signature_verifies_against_the_bytes_that_are_sent() { + let sent = OutboundRequest::signed_json( + "https://textract.us-east-1.amazonaws.com/".into(), + vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], + &json!({"Document": {"Bytes": "aGk="}}), + None, + &signer("textract"), + ) + .unwrap(); + let unsigned: BTreeMap = sent + .headers() + .iter() + .filter(|(name, _)| !is_sigv4_computed_header(name)) + .cloned() + .collect(); + let recomputed = sign_post( + sent.url(), + sent.body(), + &aws_signature_headers(&unsigned), + "us-east-1", + "textract", + &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + fixed_clock(), + ) + .unwrap(); + + assert_eq!( + sent.header("Authorization"), + Some(recomputed["Authorization"].as_str()) + ); + } + + #[test] + fn the_signature_depends_on_the_body_and_the_service() { + let original = authorization(&json!({"text": "card 4111"}), "textract"); + + assert_ne!( + original, + authorization(&json!({"text": "card [REDACTED]"}), "textract") + ); + assert_ne!( + original, + authorization(&json!({"text": "card 4111"}), "bedrock") + ); + assert!(original.contains("/us-east-1/textract/aws4_request")); + } + + #[test] + fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() { + let error = OutboundRequest::signed_json( + "https://textract.us-east-1.amazonaws.com/".into(), + vec![("authorization".into(), "Bearer caller".into())], + &json!({}), + None, + &signer("textract"), + ) + .unwrap_err(); + + assert_eq!( + error, + litellm_http::Error::ComputedHeader("authorization".into()) + ); + } +} diff --git a/litellm-rust/crates/auth/src/http.rs b/litellm-rust/crates/auth/src/http.rs index 7d20991d838..dd87d00e70f 100644 --- a/litellm-rust/crates/auth/src/http.rs +++ b/litellm-rust/crates/auth/src/http.rs @@ -40,13 +40,22 @@ pub fn apply_credential( ) } -/// How the upstream call is authenticated. API-key strategies are resolved in -/// `prepare`; SigV4 needs the serialized body, so the handler signs it. +/// How the upstream call is authenticated. API-key strategies become headers +/// in `prepare`; SigV4 covers the serialized body, so it is applied where the +/// outbound request is built. #[derive(Clone, Debug, PartialEq, Eq)] pub enum RequestAuth { - Header { name: &'static str, value: String }, - Bearer { token: String }, - AwsSigV4 { region: String }, + Header { + name: &'static str, + value: String, + }, + Bearer { + token: String, + }, + AwsSigV4 { + region: String, + service: &'static str, + }, } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/audio_transcription/error.rs b/litellm-rust/crates/core/src/audio_transcription/error.rs index 122cbab358f..81b57af2c6c 100644 --- a/litellm-rust/crates/core/src/audio_transcription/error.rs +++ b/litellm-rust/crates/core/src/audio_transcription/error.rs @@ -24,6 +24,8 @@ pub enum Error { #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] + Http(#[from] litellm_http::Error), + #[error(transparent)] Aws(#[from] litellm_auth_aws::Error), } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 503cc922966..a1862f341a5 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,4 +1,4 @@ -use litellm_http::request::{http_request, truncate_error_body}; +use litellm_http::request::truncate_error_body; use serde_json::Value; use super::{Error, client::http_client}; @@ -7,17 +7,18 @@ use crate::audio_transcription::types::ProviderAudioTranscriptionRequest; pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { - let body = serde_json::to_vec(&request.body) - .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; - let headers = signed_headers(&request, &body).await?; - let mut request_builder = http_client().post(&request.url).body(body); - for (key, value) in headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - let response = http_request(request_builder).await.map_err(|error| { + let response = crate::outbound::outbound_request::( + &request.auth, + request.url.clone(), + request.upstream_headers.clone(), + &request.body, + request.timeout, + &request.optional_params, + ) + .await? + .send(http_client()) + .await + .map_err(|error| { Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; let status = response.status(); @@ -37,33 +38,3 @@ pub async fn execute_audio_transcription_provider_call( .transform_audio_transcription_response(&request.model, response_json)? .into_json()) } - -async fn signed_headers( - request: &ProviderAudioTranscriptionRequest, - body: &[u8], -) -> Result, Error> { - use std::{collections::BTreeMap, time::SystemTime}; - - use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post}; - use litellm_llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth; - - 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()) -} diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 829617d26bd..807993c38b7 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,9 +1,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_http::request::{has_header, string_headers}; use litellm_llms::{ - base_llm::audio_transcription::transformation::{ - AudioTranscriptionAuth, BaseAudioTranscriptionConfig, - }, + base_llm::audio_transcription::transformation::{BaseAudioTranscriptionConfig, RequestAuth}, bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG, }; @@ -43,11 +41,14 @@ pub fn prepare_audio_transcription_provider_call( let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers("audio transcription", request.extra_headers)?; let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?; - if matches!(auth, AudioTranscriptionAuth::Bearer) - && !has_header(&headers, "authorization") - && let Some(api_key) = request.api_key - { - headers.push(("Authorization".to_string(), format!("Bearer {api_key}"))); + match &auth { + RequestAuth::Bearer { token } if !has_header(&headers, "authorization") => { + headers.push(("Authorization".to_string(), format!("Bearer {token}"))); + } + RequestAuth::Header { name, value } if !has_header(&headers, name) => { + headers.push(((*name).to_string(), value.clone())); + } + RequestAuth::Bearer { .. } | RequestAuth::Header { .. } | RequestAuth::AwsSigV4 { .. } => {} } if !has_header(&headers, "content-type") { headers.push(("Content-Type".to_string(), "application/json".to_string())); diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index ca09dd945be..0d87483c9bf 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_llms::base_llm::audio_transcription::transformation::{ - AudioTranscriptionAuth, BaseAudioTranscriptionConfig, + BaseAudioTranscriptionConfig, RequestAuth, }; use serde_json::{Map, Value}; @@ -24,7 +24,7 @@ pub struct ProviderAudioTranscriptionRequest { pub url: String, pub body: Value, pub upstream_headers: Vec<(String, String)>, - pub auth: AudioTranscriptionAuth, + pub auth: RequestAuth, pub optional_params: Map, pub timeout: Option, } diff --git a/litellm-rust/crates/core/src/chat_completions/error.rs b/litellm-rust/crates/core/src/chat_completions/error.rs index 122cbab358f..81b57af2c6c 100644 --- a/litellm-rust/crates/core/src/chat_completions/error.rs +++ b/litellm-rust/crates/core/src/chat_completions/error.rs @@ -24,6 +24,8 @@ pub enum Error { #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] + Http(#[from] litellm_http::Error), + #[error(transparent)] Aws(#[from] litellm_auth_aws::Error), } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index b73d4838760..de926c715d5 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,5 +1,5 @@ -use litellm_http::request::{http_request, truncate_error_body}; -use litellm_llms::base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData}; +use litellm_http::{outbound::OutboundRequest, request::truncate_error_body}; +use litellm_llms::base_llm::chat::transformation::ProviderChatResponseData; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; @@ -12,22 +12,9 @@ pub(super) async fn execute_chat_completions_provider_call( request: ResolvedChatCompletionsRequest<'_>, ) -> Result { let request = prepare_provider_request(request)?; - let body = serde_json::to_vec(&request.body).map_err(|err| { - Error::InvalidRequest(format!( - "failed to serialize chat completions request: {err}" - )) - })?; - let headers = signed_headers(&request, &body).await?; + let outbound = outbound_request(&request).await?; - let mut request_builder = http_client().post(&request.url).body(body); - for (key, value) in &headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { + let response = outbound.send(http_client()).await.map_err(|err| { // Failing to establish the connection means the request never went out, // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. @@ -77,57 +64,24 @@ pub(super) fn as_response_error(err: Error) -> Error { } } -pub(super) async fn signed_headers( +pub(super) async fn outbound_request( request: &ProviderChatCompletionsRequest, - body: &[u8], -) -> Result, Error> { - use std::{collections::BTreeMap, time::SystemTime}; - - use litellm_auth_aws::{ - 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? +) -> Result { + crate::outbound::outbound_request( + &request.auth, + request.url.clone(), + request.upstream_headers.clone(), + &request.body, + request.timeout, + &request.optional_params, + ) + .await + .map_err(|error| match error { + // Python drops the caller's copy and prefers a forwarded Authorization + // over the signature, so leave the request to it. + Error::Http(litellm_http::Error::ComputedHeader(_)) => { + Error::Unsupported("request forwards a header AWS SigV4 computes") } - }; - 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()) + other => other, + }) } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index d0aa1e88011..c8e6365121e 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,6 +1,6 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_http::request::has_header; -use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; +use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth}; use litellm_types::llms::openai::ChatMessage; use serde_json::Value; @@ -67,7 +67,7 @@ fn validate_environment( request: &ResolvedChatCompletionsRequest<'_>, model: &str, config: &dyn BaseConfig, -) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> { +) -> Result<(Vec<(String, String)>, RequestAuth), Error> { let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers(request.extra_headers.clone())?; let auth = config.auth( @@ -77,7 +77,7 @@ fn validate_environment( &env_lookup, )?; match &auth { - ChatCompletionsAuth::Header { name, value } => { + RequestAuth::Header { name, value } => { // The deployment's credential replaces whatever the caller forwarded // under the same name, mirroring Python's // `{**headers, **anthropic_headers}`: letting a request header win @@ -92,7 +92,7 @@ fn validate_environment( headers.push(((*name).to_string(), value.clone())); } } - ChatCompletionsAuth::Bearer { token } => { + RequestAuth::Bearer { token } => { // Bedrock's `get_request_headers` assigns `headers["Authorization"]` // unconditionally once a bearer token resolves, so the deployment's // identity outranks whatever the caller forwarded. Keeping the @@ -105,7 +105,7 @@ fn validate_environment( headers.push(("authorization".to_string(), format!("Bearer {token}"))); } // SigV4 signs the serialized body, so the handler adds its headers. - ChatCompletionsAuth::AwsSigV4 { .. } => {} + RequestAuth::AwsSigV4 { .. } => {} } for (name, value) in config.default_headers() { diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index dcaa3397add..dd5938cf168 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -1,4 +1,4 @@ -use litellm_llms::base_llm::chat::transformation::ChatCompletionsAuth; +use litellm_llms::base_llm::chat::transformation::RequestAuth; use serde_json::{Map, Value, json}; use super::{ @@ -90,7 +90,7 @@ fn adds_the_auth_and_default_headers() { ); assert!(matches!( prepared.auth, - ChatCompletionsAuth::Header { + RequestAuth::Header { name: "x-api-key", .. } @@ -289,8 +289,9 @@ fn prepares_a_bedrock_call_without_resolving_credentials() { ); assert_eq!( prepared.auth, - ChatCompletionsAuth::AwsSigV4 { - region: "us-east-1".to_string() + RequestAuth::AwsSigV4 { + region: "us-east-1".to_string(), + service: "bedrock", } ); // SigV4 signs the serialized body, so prepare must not have added an @@ -326,15 +327,14 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { json!("abc-123"), )])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#) + let signed = super::handler::outbound_request(&prepared) .await .expect("signs"); let authorization = signed - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.clone()) - .expect("carries an authorization header"); + .header("authorization") + .expect("carries an authorization header") + .to_string(); assert!( authorization.starts_with("AWS4-HMAC-SHA256"), "expected a SigV4 signature, got {authorization}" @@ -346,6 +346,7 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { // It still goes on the wire, it is just not part of the signature. assert!( signed + .headers() .iter() .any(|(name, value)| name == "x-request-id" && value == "abc-123"), "forwarded header was dropped instead of reattached" @@ -376,7 +377,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() { call.api_key = None; call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#) + let error = super::handler::outbound_request(&prepared) .await .expect_err("{forwarded} should decline instead of being signed"); assert!( @@ -466,7 +467,7 @@ fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { .expect("prepares"); assert_eq!( prepared.auth, - ChatCompletionsAuth::Bearer { + RequestAuth::Bearer { token: "sk-test".to_string() } ); diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 882611d5862..3b74cf5dace 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; +use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth}; use litellm_types::llms::openai::ChatMessage; use serde_json::{Map, Value}; @@ -38,7 +38,7 @@ pub struct ProviderChatCompletionsRequest { pub url: String, pub body: Value, pub upstream_headers: Vec<(String, String)>, - pub auth: ChatCompletionsAuth, + pub auth: RequestAuth, pub optional_params: Map, pub timeout: Option, } diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index e3e2fb48721..afe5ea595aa 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -4,6 +4,7 @@ pub mod constants; pub mod error; pub mod messages; pub mod ocr; +mod outbound; pub mod responses; pub use error::Error; diff --git a/litellm-rust/crates/core/src/ocr/arguments.rs b/litellm-rust/crates/core/src/ocr/arguments.rs index 43f1c6d6d43..05aa345f01d 100644 --- a/litellm-rust/crates/core/src/ocr/arguments.rs +++ b/litellm-rust/crates/core/src/ocr/arguments.rs @@ -15,6 +15,18 @@ const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ "azure_federated_token_file", "enable_azure_ad_token_refresh", ]; +const AWS_AUTH_OPTION_FIELDS: &[&str] = &[ + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_region_name", + "aws_session_name", + "aws_profile_name", + "aws_role_name", + "aws_web_identity_token", + "aws_sts_endpoint", + "aws_external_id", +]; const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[ "vertex_credentials", "vertex_ai_credentials", @@ -35,6 +47,7 @@ pub fn consumed_optional_param_names( let (model, config) = resolve_provider_config(model, custom_llm_provider)?; let provider_fields = config.get_supported_ocr_params(&model); let auth_fields: &[&str] = match config { + OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS, OcrConfigKind::AzureAi | OcrConfigKind::AzureDocumentIntelligence | OcrConfigKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS, @@ -57,6 +70,9 @@ pub(crate) fn is_secret_param(name: &str) -> bool { | "azure_federated_token_file" | "vertex_credentials" | "vertex_ai_credentials" + | "aws_secret_access_key" + | "aws_session_token" + | "aws_web_identity_token" ) } diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index c977f721a70..f298f106a5f 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -8,6 +8,10 @@ pub mod route; pub mod types; pub mod wire; +#[cfg(test)] +#[path = "../../tests/aws_textract_ocr.rs"] +mod aws_textract_tests; + #[cfg(test)] #[path = "../../tests/azure_ai_ocr.rs"] mod azure_ai_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index f1e1dcaaa6d..715aedc69df 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -19,7 +19,10 @@ pub(crate) fn prepare_request( Some("MISTRAL_AZURE_API_BASE"), ), OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")), - OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None), + OcrProvider::AwsTextract + | OcrProvider::Cohere + | OcrProvider::Reducto + | OcrProvider::VertexAi => (None, None), }; let secret = |name: &str| client.secrets().truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index ee9ba76928d..34e2a77b6d1 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -1,5 +1,9 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; use litellm_llms::{ + aws_textract::ocr::{ + analyze_transformation::TextractAnalyzeDocumentConfig, + transformation::TextractDetectTextConfig, + }, azure_ai::ocr::{ cohere_parse_transformation::AzureAICohereParseConfig, document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig, @@ -25,6 +29,14 @@ use strum::{EnumString, IntoStaticStr}; macro_rules! with_config { ($kind:expr, $config:ident => $body:expr) => { match $kind { + OcrConfigKind::AwsTextract => { + let $config = TextractDetectTextConfig; + $body + } + OcrConfigKind::AwsTextractAnalyze => { + let $config = TextractAnalyzeDocumentConfig; + $body + } OcrConfigKind::Cohere => { let $config = CohereParseConfig; $body @@ -67,6 +79,8 @@ macro_rules! with_config { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum OcrConfigKind { + AwsTextract, + AwsTextractAnalyze, Cohere, Mistral, AzureAi, @@ -81,6 +95,7 @@ pub(crate) enum OcrConfigKind { impl OcrConfigKind { pub(crate) const fn provider(self) -> OcrProvider { match self { + Self::AwsTextract | Self::AwsTextractAnalyze => OcrProvider::AwsTextract, Self::Cohere => OcrProvider::Cohere, Self::Mistral => OcrProvider::Mistral, Self::AzureAi | Self::AzureCohere | Self::AzureDocumentIntelligence => { @@ -141,6 +156,7 @@ pub fn get_health_check_document( #[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, PartialEq, Eq)] #[strum(serialize_all = "snake_case")] pub(crate) enum OcrProvider { + AwsTextract, Cohere, Mistral, AzureAi, @@ -162,6 +178,10 @@ pub(crate) fn resolve_provider_config( .parse::() .map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?; let config = match ocr_provider { + OcrProvider::AwsTextract if provider.model.eq_ignore_ascii_case("analyze-document") => { + OcrConfigKind::AwsTextractAnalyze + } + OcrProvider::AwsTextract => OcrConfigKind::AwsTextract, OcrProvider::Cohere => OcrConfigKind::Cohere, OcrProvider::Mistral => OcrConfigKind::Mistral, OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { @@ -419,6 +439,9 @@ mod tests { } #[rstest] + #[case("aws_textract/detect-document-text", OcrConfigKind::AwsTextract)] + #[case("aws_textract/analyze-document", OcrConfigKind::AwsTextractAnalyze)] + #[case("aws_textract/Analyze-Document", OcrConfigKind::AwsTextractAnalyze)] #[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)] #[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)] #[case("azure_ai/Cohere-parse-v5", OcrConfigKind::AzureCohere)] diff --git a/litellm-rust/crates/core/src/outbound.rs b/litellm-rust/crates/core/src/outbound.rs new file mode 100644 index 00000000000..7fc90084e6f --- /dev/null +++ b/litellm-rust/crates/core/src/outbound.rs @@ -0,0 +1,30 @@ +use std::time::Duration; + +use litellm_auth::RequestAuth; +use litellm_auth_aws::SigV4Signer; +use litellm_http::outbound::OutboundRequest; +use serde_json::{Map, Value}; + +/// Header credentials are already in `headers`; SigV4 is applied here, over the +/// bytes that are sent. +pub(crate) async fn outbound_request( + auth: &RequestAuth, + url: String, + headers: Vec<(String, String)>, + body: &Value, + timeout: Option, + optional_params: &Map, +) -> Result +where + E: From + From, +{ + let RequestAuth::AwsSigV4 { region, service } = auth else { + return Ok(OutboundRequest::json(url, headers, body, timeout)?); + }; + let env_lookup = |key: &str| std::env::var(key).ok(); + let signer = + SigV4Signer::resolve(region.clone(), service, optional_params, &env_lookup).await?; + Ok(OutboundRequest::signed_json( + url, headers, body, timeout, &signer, + )?) +} diff --git a/litellm-rust/crates/core/tests/aws_textract_ocr.rs b/litellm-rust/crates/core/tests/aws_textract_ocr.rs new file mode 100644 index 00000000000..c536317ad5c --- /dev/null +++ b/litellm-rust/crates/core/tests/aws_textract_ocr.rs @@ -0,0 +1,193 @@ +use std::{collections::BTreeMap, time::SystemTime}; + +use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use litellm_llms::base_llm::ocr::error::Error; +use serde_json::{Value, json}; +use time::{PrimitiveDateTime, format_description}; + +use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, header, mock_server, perform_ocr_with, request_body, + wire_request_with_document, + }, + types::LiteLLMOcrRequest, +}; + +const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; +const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; + +fn textract_request(base: &str) -> LiteLLMOcrRequest { + textract_request_for("aws_textract/detect-document-text", base) +} + +fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) +} + +fn textract_response() -> MockResponse { + MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) +} + +/// Recomputes SigV4 over the bytes the server received, at the time the client claimed. +fn expected_authorization(url: &str, raw_request: &str) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| { + ( + name.to_string(), + header(raw_request, name).unwrap().to_string(), + ) + }) + .collect(); + let body = raw_request.split_once("\r\n\r\n").unwrap().1; + sign_post( + url, + body.as_bytes(), + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() +} + +#[tokio::test] +async fn the_request_is_signed_for_textract_and_lines_become_the_page() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + + let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + header(&raw, "content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); +} + +#[tokio::test] +async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); +} + +#[tokio::test] +async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { + let (base, _, server) = mock_server(vec![MockResponse { + status: 400, + headers: vec![], + body: json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + }]) + .await; + + let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap_err(); + server.await.unwrap(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); +} + +#[tokio::test] +async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + let request = textract_request_for("aws_textract/analyze-document", &base); + + let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!( + request_body(&raw)["FeatureTypes"], + json!(["LAYOUT", "TABLES"]) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); +} diff --git a/litellm-rust/crates/core/tests/cohere_ocr.rs b/litellm-rust/crates/core/tests/cohere_ocr.rs index fc1203f0980..12824f58b1d 100644 --- a/litellm-rust/crates/core/tests/cohere_ocr.rs +++ b/litellm-rust/crates/core/tests/cohere_ocr.rs @@ -38,7 +38,7 @@ mod transformation { ) .await .unwrap(); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!( body, json!({ @@ -75,7 +75,7 @@ mod transformation { ) .await .unwrap(); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!(body["output_format"], "markdown"); assert!(body.get("req_format").is_none()); } diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 399b7cac39a..035f3fe944d 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -160,17 +160,16 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) .await .unwrap(); - assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr"); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); assert_eq!( - vertex_http.url().as_str(), + vertex_http.url(), "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" ); for http in [&direct_http, &vertex_http] { - assert_eq!(http.method(), reqwest::Method::POST); - assert_eq!(http.headers()["authorization"], "Bearer test-key"); - assert_eq!(http.headers()["content-type"], "application/json"); - assert_eq!(http.timeout(), Some(&Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!( body, json!({ @@ -250,9 +249,9 @@ mod transformation { .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) .await .unwrap(); - assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr"); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); assert_eq!( - vertex_http.url().as_str(), + vertex_http.url(), "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" ); let http = if use_vertex { @@ -260,11 +259,10 @@ mod transformation { } else { &direct_http }; - assert_eq!(http.method(), reqwest::Method::POST); - assert_eq!(http.headers()["authorization"], "Bearer test-key"); - assert_eq!(http.headers()["content-type"], "application/json"); - assert_eq!(http.timeout(), Some(&Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); assert_eq!( body, json!({ diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index d4457f5685c..cad5aa87e49 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -14,6 +14,7 @@ litellm-core-utils.workspace = true hyper-util.workspace = true reqwest.workspace = true rustls.workspace = true +serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 697d0cf59c8..e06f7c00cf5 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -8,6 +8,12 @@ pub enum Error { InvalidPem { path: PathBuf, message: String }, #[error("could not build the HTTP client: {0}")] Client(String), + #[error("request body could not be serialized: {0}")] + RequestBody(String), + #[error("request forwards a header the signer computes: {0}")] + ComputedHeader(String), + #[error("request signing failed: {0}")] + Signature(String), } impl From for Error { diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index c6d9959348d..6f62a00175c 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -1,6 +1,7 @@ mod config; mod error; pub mod media; +pub mod outbound; mod pool; mod proxy; pub mod request; diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs new file mode 100644 index 00000000000..d100bdf624b --- /dev/null +++ b/litellm-rust/crates/http/src/outbound.rs @@ -0,0 +1,210 @@ +//! The request a route hands to the transport. The body is serialized once, +//! when the request is built, and a [`RequestSigner`] sees those exact bytes. +//! +//! Host hooks may rewrite the wire request (redaction, guardrails) and a +//! signature such as AWS SigV4 covers the body, so a route builds this after +//! its hooks ran and cannot change or re-serialize it afterwards. + +use std::time::Duration; + +use serde::Serialize; + +use crate::{ + Error, + request::{HeaderPolicy, has_header, with_headers}, +}; + +#[derive(Clone, Copy, Debug)] +pub struct UnsignedRequest<'a> { + pub url: &'a str, + pub headers: &'a [(String, String)], + pub body: &'a [u8], +} + +/// Returns the headers to add to the request; it never sees a mutable request. +pub trait RequestSigner: Send + Sync { + fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error>; +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OutboundRequest { + url: String, + headers: Vec<(String, String)>, + body: Vec, + timeout: Option, +} + +impl OutboundRequest { + pub fn json( + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, + timeout: Option, + ) -> Result { + Self::build(url, headers, body, timeout, None) + } + + pub fn signed_json( + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, + timeout: Option, + signer: &dyn RequestSigner, + ) -> Result { + Self::build(url, headers, body, timeout, Some(signer)) + } + + fn build( + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, + timeout: Option, + signer: Option<&dyn RequestSigner>, + ) -> Result { + let body = + serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; + let content_type = (!has_header(&headers, "content-type")) + .then(|| ("content-type".to_string(), "application/json".to_string())); + let unsigned: Vec<(String, String)> = headers.into_iter().chain(content_type).collect(); + let signature = signer + .map(|signer| { + signer.sign(UnsignedRequest { + url: &url, + headers: &unsigned, + body: &body, + }) + }) + .transpose()? + .unwrap_or_default(); + Ok(Self { + url, + headers: unsigned.into_iter().chain(signature).collect(), + body, + timeout, + }) + } + + pub fn url(&self) -> &str { + &self.url + } + + pub fn headers(&self) -> &[(String, String)] { + &self.headers + } + + pub fn header(&self, name: &str) -> Option<&str> { + self.headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) + } + + pub fn body(&self) -> &[u8] { + &self.body + } + + pub fn timeout(&self) -> Option { + self.timeout + } + + pub async fn send(self, client: &reqwest::Client) -> Result { + let builder = with_headers( + client.post(&self.url).body(self.body), + &self.headers, + HeaderPolicy::All, + ); + match self.timeout { + Some(timeout) => builder.timeout(timeout), + None => builder, + } + .send() + .await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use serde_json::json; + + use super::*; + + #[derive(Default)] + struct Recording(Mutex>); + + impl RequestSigner for Recording { + fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { + *self.0.lock().unwrap() = request.body.to_vec(); + Ok(vec![("authorization".into(), "signed".into())]) + } + } + + #[test] + fn the_signer_sees_exactly_the_bytes_that_are_sent() { + let signer = Recording::default(); + let request = OutboundRequest::signed_json( + "https://provider.test/".into(), + vec![("x-caller".into(), "kept".into())], + &json!({"b": 1, "a": [true, null]}), + None, + &signer, + ) + .unwrap(); + + assert_eq!(request.body(), signer.0.lock().unwrap().as_slice()); + assert_eq!(request.header("authorization"), Some("signed")); + assert_eq!(request.header("x-caller"), Some("kept")); + } + + #[test] + fn the_content_type_is_part_of_what_the_signer_sees() { + struct RequiresContentType; + impl RequestSigner for RequiresContentType { + fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { + has_header(request.headers, "content-type") + .then(Vec::new) + .ok_or_else(|| Error::Signature("content-type was not signed".into())) + } + } + + let defaulted = OutboundRequest::signed_json( + "u".into(), + Vec::new(), + &json!({}), + None, + &RequiresContentType, + ) + .unwrap(); + assert_eq!(defaulted.header("content-type"), Some("application/json")); + + let provider = OutboundRequest::signed_json( + "u".into(), + vec![("Content-Type".into(), "application/x-amz-json-1.1".into())], + &json!({}), + None, + &RequiresContentType, + ) + .unwrap(); + assert_eq!( + provider.header("content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!(provider.headers().len(), 1); + } + + #[test] + fn a_signer_failure_produces_no_request() { + struct Refuses; + impl RequestSigner for Refuses { + fn sign(&self, _request: UnsignedRequest<'_>) -> Result, Error> { + Err(Error::ComputedHeader("authorization".into())) + } + } + + assert_eq!( + OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses), + Err(Error::ComputedHeader("authorization".into())) + ); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs b/litellm-rust/crates/llms/src/anthropic/chat/tests.rs index 40e89c52c3c..3777347d240 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/tests.rs @@ -428,7 +428,7 @@ fn resolves_the_messages_url_and_x_api_key_auth() { config .auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None) .expect("auth resolves"), - ChatCompletionsAuth::Header { + RequestAuth::Header { name: "x-api-key", value: "sk-x".to_string() } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 21fa4e9f82e..6fc4f00b981 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -16,7 +16,7 @@ use crate::{ }, }, base_llm::chat::transformation::{ - BaseConfig, ChatCompletionsAuth, Error, ProviderChatRequestData, ProviderChatResponseData, + BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, Unsupported, unsupported_message, unsupported_param, }, }; @@ -137,8 +137,8 @@ impl BaseConfig for AnthropicConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(ChatCompletionsAuth::Header { + ) -> Result { + Ok(RequestAuth::Header { name: "x-api-key", value: resolve_anthropic_api_key(api_key, env_lookup)?, }) diff --git a/litellm-rust/crates/llms/src/aws_textract/mod.rs b/litellm-rust/crates/llms/src/aws_textract/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md b/litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md new file mode 100644 index 00000000000..4913f89924a --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md @@ -0,0 +1,12 @@ +- https://docs.aws.amazon.com/textract/latest/APIReference/Welcome.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_Operations.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_DetectDocumentText.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_AnalyzeDocument.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_StartDocumentTextDetection.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_Document.md +- https://docs.aws.amazon.com/textract/latest/APIReference/API_Block.md +- https://docs.aws.amazon.com/textract/latest/dg/what-is.md +- https://docs.aws.amazon.com/textract/latest/dg/sync.md +- https://docs.aws.amazon.com/textract/latest/dg/async.md +- https://docs.aws.amazon.com/textract/latest/dg/how-it-works-document-layout.md +- https://docs.aws.amazon.com/textract/latest/dg/limits.md diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs new file mode 100644 index 00000000000..65c7688ea0c --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs @@ -0,0 +1,426 @@ +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use litellm_core_utils::call_arguments::{CallArguments, parse_options}; +use serde::{Deserialize, Serialize}; + +use super::common_utils::{ + Block, DocumentMetadata, HEALTH_CHECK_IMAGE_DATA_URI, TextractDocument, TextractEnvironment, + document_bytes, endpoint, environment, error_class, inline_document, lines_by_page, +}; +use crate::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrRequestContext, + OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, + }, +}; + +const ANALYZE_DOCUMENT_TARGET: &str = "Textract.AnalyzeDocument"; +const DEFAULT_FEATURE_TYPES: [&str; 2] = ["LAYOUT", "TABLES"]; + +#[derive(Default, Deserialize)] +pub struct AnalyzeDocumentOptions { + pub feature_types: Option>, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct AnalyzeDocumentRequest { + #[serde(rename = "Document")] + pub document: TextractDocument, + #[serde(rename = "FeatureTypes")] + pub feature_types: Vec, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub struct AnalyzeDocumentResponse { + #[serde(default)] + blocks: Vec, + document_metadata: Option, +} + +/// Synchronous `AnalyzeDocument`: layout and tables rendered as markdown. +#[derive(Clone, Copy, Debug, Default)] +pub struct TextractAnalyzeDocumentConfig; + +impl BaseOcrConfig for TextractAnalyzeDocumentConfig { + type OcrParams = AnalyzeDocumentOptions; + type ProviderRequest = AnalyzeDocumentRequest; + type Environment = TextractEnvironment; + + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &["feature_types"] + } + + fn get_health_check_document(&self) -> OcrDocument { + OcrDocument::ImageUrl { + image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(), + extra_fields: Default::default(), + } + } + + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + _model: &str, + ) -> Result { + Ok(parse_options(non_default_params)?) + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + environment(request, ANALYZE_DOCUMENT_TARGET).await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _optional_params: &AnalyzeDocumentOptions, + environment: &TextractEnvironment, + ) -> Result { + Ok(endpoint(request, environment)) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + optional_params: &AnalyzeDocumentOptions, + _headers: &[(String, String)], + ) -> Result { + Ok(AnalyzeDocumentRequest { + document: document_bytes(&document)?, + feature_types: optional_params.feature_types.clone().unwrap_or_else(|| { + DEFAULT_FEATURE_TYPES + .iter() + .map(|feature| feature.to_string()) + .collect() + }), + }) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &AnalyzeDocumentOptions, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let document = inline_document(document, context).await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result { + decode_and_normalize_response(model, raw_response, request_format, normalize_response) + } + + fn get_error_class( + &self, + error_message: String, + status_code: u16, + headers: Vec<(String, String)>, + ) -> Error { + error_class(error_message, status_code, headers) + } +} + +fn normalize_response( + model: &str, + response: AnalyzeDocumentResponse, +) -> Result { + let blocks = &response.blocks; + let has_layout = blocks.iter().any(is_layout); + let page_markdown: Vec<(i64, String)> = if has_layout { + let by_id: HashMap<&str, &Block> = blocks + .iter() + .map(|block| (block.id.as_str(), block)) + .collect(); + let pages: BTreeSet = blocks.iter().map(Block::page).collect(); + pages + .into_iter() + .map(|page| (page, layout_markdown(blocks, page, &by_id))) + .filter(|(_, markdown)| !markdown.is_empty()) + .collect() + } else { + lines_by_page(blocks) + }; + let pages: Vec = page_markdown + .into_iter() + .map(|(page, markdown)| OcrPage { + index: page - 1, + markdown, + ..Default::default() + }) + .collect(); + let pages_processed = response + .document_metadata + .and_then(|metadata| metadata.pages) + .or_else(|| i64::try_from(pages.len()).ok()); + Ok(LiteLLMOcrResponse { + usage_info: Some(OcrUsageInfo { + pages_processed, + ..Default::default() + }), + ..LiteLLMOcrResponse::new(model, pages) + }) +} + +fn is_layout(block: &Block) -> bool { + block.block_type.starts_with("LAYOUT_") +} + +/// Layout blocks arrive in reading order. A list's items are repeated as +/// top-level `LAYOUT_TEXT` blocks, and a `LAYOUT_TABLE` only links to the +/// table's lines, so the nth layout table on a page takes the nth `TABLE`. +fn layout_markdown(blocks: &[Block], page: i64, by_id: &HashMap<&str, &Block>) -> String { + let on_page = || blocks.iter().filter(move |block| block.page() == page); + let list_items: BTreeSet<&str> = on_page() + .filter(|block| block.block_type == "LAYOUT_LIST") + .flat_map(Block::children) + .collect(); + let tables: Vec<&Block> = on_page() + .filter(|block| block.block_type == "TABLE") + .collect(); + let table_ordinal: HashMap<&str, usize> = on_page() + .filter(|block| block.block_type == "LAYOUT_TABLE") + .enumerate() + .map(|(ordinal, block)| (block.id.as_str(), ordinal)) + .collect(); + let sections: Vec = on_page() + .filter(|block| is_layout(block) && !list_items.contains(block.id.as_str())) + .map(|block| match block.block_type.as_str() { + "LAYOUT_TITLE" => format!("# {}", text_of(block, by_id, " ")), + "LAYOUT_SECTION_HEADER" => format!("## {}", text_of(block, by_id, " ")), + "LAYOUT_LIST" => block + .children() + .filter_map(|id| by_id.get(id)) + .map(|item| format!("- {}", strip_bullet(&text_of(item, by_id, " ")))) + .collect::>() + .join("\n"), + "LAYOUT_TABLE" => table_ordinal + .get(block.id.as_str()) + .and_then(|ordinal| tables.get(*ordinal)) + .map(|table| table_markdown(table, by_id)) + .unwrap_or_else(|| text_of(block, by_id, "\n")), + _ => text_of(block, by_id, " "), + }) + .filter(|section| !section.trim().is_empty()) + .collect(); + sections.join("\n\n") +} + +fn text_of(block: &Block, by_id: &HashMap<&str, &Block>, separator: &str) -> String { + match &block.text { + Some(text) => text.clone(), + None => block + .children() + .filter_map(|id| by_id.get(id)) + .map(|child| text_of(child, by_id, separator)) + .filter(|text| !text.is_empty()) + .collect::>() + .join(separator), + } +} + +fn strip_bullet(item: &str) -> &str { + item.trim_start_matches(['-', '*', '\u{2022}', '\u{00b7}']) + .trim_start() +} + +fn table_markdown(table: &Block, by_id: &HashMap<&str, &Block>) -> String { + let cells: BTreeMap<(usize, usize), String> = table + .children() + .filter_map(|id| by_id.get(id)) + .filter(|cell| cell.block_type == "CELL") + .filter_map(|cell| { + Some(( + (cell.row_index?, cell.column_index?), + text_of(cell, by_id, " ").replace('|', "\\|"), + )) + }) + .collect(); + let columns = cells.keys().map(|(_, column)| *column).max().unwrap_or(0); + let rows: BTreeSet = cells.keys().map(|(row, _)| *row).collect(); + let render = |row: usize| { + let values: Vec<&str> = (1..=columns) + .map(|column| cells.get(&(row, column)).map_or("", String::as_str)) + .collect(); + format!("| {} |", values.join(" | ")) + }; + let divider = format!("|{}", " --- |".repeat(columns)); + rows.iter() + .enumerate() + .flat_map(|(position, row)| { + std::iter::once(render(*row)).chain((position == 0).then(|| divider.clone())) + }) + .collect::>() + .join("\n") +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + + use super::*; + + fn markdown(blocks: Value) -> Vec<(i64, String)> { + TextractAnalyzeDocumentConfig + .transform_ocr_response( + "analyze-document", + &serde_json::to_vec(&json!({"DocumentMetadata": {"Pages": 1}, "Blocks": blocks})) + .unwrap(), + OcrResponseFormat::Litellm, + ) + .unwrap() + .pages + .into_iter() + .map(|page| (page.index, page.markdown)) + .collect() + } + + fn child(ids: &[&str]) -> Value { + json!([{"Type": "CHILD", "Ids": ids}]) + } + + fn line(id: &str, text: &str) -> Value { + json!({"Id": id, "BlockType": "LINE", "Text": text}) + } + + fn word(id: &str, text: &str) -> Value { + json!({"Id": id, "BlockType": "WORD", "Text": text}) + } + + fn cell(id: &str, row: usize, column: usize, words: &[&str]) -> Value { + json!({"Id": id, "BlockType": "CELL", "RowIndex": row, "ColumnIndex": column, + "Relationships": child(words)}) + } + + #[test] + fn layout_becomes_headings_paragraphs_and_a_list_without_repeating_its_items() { + let pages = markdown(json!([ + line("l1", "Quarterly Report"), + line("l2", "This report lists"), + line("l3", "the invoices."), + line("l4", "Line items"), + line("l5", "- Pay within 30 days"), + line("l6", "\u{2022} Quote the number"), + {"Id": "t", "BlockType": "LAYOUT_TITLE", "Relationships": child(&["l1"])}, + {"Id": "p", "BlockType": "LAYOUT_TEXT", "Relationships": child(&["l2", "l3"])}, + {"Id": "h", "BlockType": "LAYOUT_SECTION_HEADER", "Relationships": child(&["l4"])}, + {"Id": "ul", "BlockType": "LAYOUT_LIST", "Relationships": child(&["i1", "i2"])}, + {"Id": "i1", "BlockType": "LAYOUT_TEXT", "Relationships": child(&["l5"])}, + {"Id": "i2", "BlockType": "LAYOUT_TEXT", "Relationships": child(&["l6"])} + ])); + + assert_eq!( + pages, + vec![( + 0, + "# Quarterly Report\n\nThis report lists the invoices.\n\n## Line items\n\n- Pay within 30 days\n- Quote the number".to_string() + )] + ); + } + + #[test] + fn a_layout_table_is_rendered_from_the_table_cells_in_row_and_column_order() { + let pages = markdown(json!([ + line("l1", "Invoice"), line("l2", "Total"), line("l3", "12345"), line("l4", "a|b"), + word("w1", "Invoice"), word("w2", "Total"), word("w3", "12345"), word("w4", "a|b"), + {"Id": "tb", "BlockType": "TABLE", "Relationships": [ + {"Type": "CHILD", "Ids": ["c4", "c1", "c3", "c2"]}, + {"Type": "TABLE_TITLE", "Ids": ["title"]} + ]}, + cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]), + cell("c3", 2, 1, &["w3"]), cell("c4", 2, 2, &["w4"]), + {"Id": "lt", "BlockType": "LAYOUT_TABLE", "Relationships": child(&["l1", "l2", "l3", "l4"])} + ])); + + assert_eq!( + pages, + vec![( + 0, + "| Invoice | Total |\n| --- | --- |\n| 12345 | a\\|b |".to_string() + )] + ); + } + + #[test] + fn a_layout_table_without_table_blocks_keeps_its_lines() { + let pages = markdown(json!([ + line("l1", "Invoice Total"), + line("l2", "12345 67.89"), + {"Id": "lt", "BlockType": "LAYOUT_TABLE", "Relationships": child(&["l1", "l2"])} + ])); + + assert_eq!(pages, vec![(0, "Invoice Total\n12345 67.89".to_string())]); + } + + #[test] + fn a_response_without_layout_blocks_falls_back_to_lines() { + let pages = markdown(json!([ + line("l1", "first"), + word("w1", "first"), + line("l2", "second") + ])); + + assert_eq!(pages, vec![(0, "first\nsecond".to_string())]); + } + + #[test] + fn each_page_gets_its_own_markdown_and_its_own_tables() { + let pages = markdown(json!([ + {"Id": "a", "BlockType": "LINE", "Text": "one", "Page": 1}, + {"Id": "b", "BlockType": "LINE", "Text": "two", "Page": 2}, + {"Id": "w", "BlockType": "WORD", "Text": "cell", "Page": 2}, + {"Id": "t1", "BlockType": "LAYOUT_TEXT", "Page": 1, "Relationships": child(&["a"])}, + {"Id": "tb", "BlockType": "TABLE", "Page": 2, "Relationships": child(&["c"])}, + {"Id": "c", "BlockType": "CELL", "Page": 2, "RowIndex": 1, "ColumnIndex": 1, + "Relationships": child(&["w"])}, + {"Id": "lt", "BlockType": "LAYOUT_TABLE", "Page": 2, "Relationships": child(&["b"])} + ])); + + assert_eq!( + pages, + vec![(0, "one".to_string()), (1, "| cell |\n| --- |".to_string())] + ); + } + + #[test] + fn feature_types_default_to_layout_and_tables_and_can_be_overridden() { + let document = || OcrDocument::ImageUrl { + image_url: "data:image/png;base64,aGk=".into(), + extra_fields: Default::default(), + }; + let request = |options: Value| { + let arguments: CallArguments = serde_json::from_value(options).unwrap(); + let params = TextractAnalyzeDocumentConfig + .map_ocr_params(&arguments, "analyze-document") + .unwrap(); + serde_json::to_value( + TextractAnalyzeDocumentConfig + .transform_ocr_request("analyze-document", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap() + }; + + assert_eq!( + request(json!({})), + json!({"Document": {"Bytes": "aGk="}, "FeatureTypes": ["LAYOUT", "TABLES"]}) + ); + assert_eq!( + request(json!({"feature_types": ["FORMS"]}))["FeatureTypes"], + json!(["FORMS"]) + ); + } +} diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs new file mode 100644 index 00000000000..bc1b5eb66d9 --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs @@ -0,0 +1,247 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_auth_aws::{SigV4Signer, resolve_aws_region}; +use litellm_http::outbound::RequestSigner; +use serde::{Deserialize, Serialize}; + +use crate::base_llm::ocr::{ + document::{InlineDocument, inline_remote_document}, + error::Error, + transformation::{ + OCR_INLINE_MAX_BYTES, OcrDocument, OcrEnvironment, OcrRequestContext, PreparedOcrRequest, + }, +}; + +const TEXTRACT_SERVICE: &str = "textract"; +const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1"; +const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException"; + +pub(super) const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; + +#[derive(Debug, Deserialize, Serialize)] +pub struct TextractDocument { + #[serde(rename = "Bytes")] + pub bytes: String, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub(super) struct Block { + #[serde(default)] + pub id: String, + pub block_type: String, + pub text: Option, + pub page: Option, + pub row_index: Option, + pub column_index: Option, + #[serde(default)] + pub relationships: Vec, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub(super) struct Relationship { + pub r#type: String, + #[serde(default)] + pub ids: Vec, +} + +impl Block { + /// The synchronous API omits `Page` because it only ever reads one. + pub fn page(&self) -> i64 { + self.page.unwrap_or(1) + } + + pub fn children(&self) -> impl Iterator { + self.relationships + .iter() + .filter(|relationship| relationship.r#type == "CHILD") + .flat_map(|relationship| relationship.ids.iter().map(String::as_str)) + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub(super) struct DocumentMetadata { + pub pages: Option, +} + +pub struct TextractEnvironment { + headers: Vec<(String, String)>, + region: String, + signer: SigV4Signer, +} + +impl OcrEnvironment for TextractEnvironment { + fn headers(&self) -> &[(String, String)] { + &self.headers + } + + fn signer(&self) -> Option<&dyn RequestSigner> { + Some(&self.signer) + } +} + +pub(super) async fn environment( + request: &PreparedOcrRequest, + target: &'static str, +) -> Result { + let env_lookup = |name: &str| request.connection.secret(name); + let region = + resolve_aws_region(None, &request.optional_params, &env_lookup).ok_or_else(|| { + Error::InvalidRequest( + "Missing AWS region - pass aws_region_name or set AWS_REGION_NAME or AWS_REGION" + .into(), + ) + })?; + let signer = SigV4Signer::resolve( + region.clone(), + TEXTRACT_SERVICE, + &request.optional_params, + &env_lookup, + ) + .await + .map_err(litellm_auth::Error::from)?; + Ok(TextractEnvironment { + headers: request + .connection + .extra_headers + .iter() + .cloned() + .chain([ + ("X-Amz-Target".into(), target.into()), + ("Content-Type".into(), AWS_JSON_CONTENT_TYPE.into()), + ]) + .collect(), + region, + signer, + }) +} + +pub(super) fn endpoint(request: &PreparedOcrRequest, environment: &TextractEnvironment) -> String { + request + .connection + .api_base + .clone() + .unwrap_or_else(|| format!("https://textract.{}.amazonaws.com/", environment.region)) +} + +pub(super) fn document_bytes(document: &OcrDocument) -> Result { + let inline = InlineDocument::parse(document.source())?.ok_or(Error::InvalidDataUri)?; + Ok(TextractDocument { + bytes: STANDARD.encode(inline.decode(OCR_INLINE_MAX_BYTES)?), + }) +} + +pub(super) async fn inline_document( + document: OcrDocument, + context: OcrRequestContext<'_>, +) -> Result { + inline_remote_document( + context.client.document_fetcher(), + document, + context.connection, + ) + .await +} + +#[derive(Deserialize)] +struct AwsError { + #[serde(rename = "__type", default)] + kind: String, + #[serde(rename = "Message", alias = "message", default)] + message: String, +} + +/// Textract answers a multi-page PDF or TIFF with a bare "unsupported document +/// format", which reads like a corrupt file. Say what the limit is. +pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, String)>) -> Error { + let unsupported = serde_json::from_str::(&body) + .ok() + .filter(|error| error.kind.ends_with(UNSUPPORTED_DOCUMENT)); + Error::Provider { + status, + body: match unsupported { + Some(error) => format!( + "{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; multi-page documents are not supported", + error.message + ), + None => body, + }, + headers, + } +} + +pub(super) fn lines_by_page(blocks: &[Block]) -> Vec<(i64, String)> { + let pages: std::collections::BTreeSet = blocks.iter().map(Block::page).collect(); + pages + .into_iter() + .map(|page| { + let lines: Vec<&str> = blocks + .iter() + .filter(|block| block.block_type == "LINE" && block.page() == page) + .filter_map(|block| block.text.as_deref()) + .collect(); + (page, lines.join("\n")) + }) + .filter(|(_, markdown)| !markdown.is_empty()) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_multi_page_rejection_names_the_single_page_limit_and_keeps_the_status() { + let error = error_class( + r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), + 400, + vec![("x-amzn-requestid".into(), "abc".into())], + ); + + let Error::Provider { + status, + body, + headers, + } = error + else { + panic!("expected a provider error"); + }; + assert_eq!(status, 400); + assert!(body.contains("Request has unsupported document format")); + assert!(body.contains("single-page PDF or TIFF")); + assert_eq!(headers, vec![("x-amzn-requestid".into(), "abc".into())]); + } + + #[test] + fn a_namespaced_exception_type_is_recognized() { + let Error::Provider { body, .. } = error_class( + r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","message":"bad"}"# + .into(), + 400, + Vec::new(), + ) else { + panic!("expected a provider error"); + }; + assert!(body.contains("multi-page documents are not supported")); + } + + #[test] + fn other_provider_errors_pass_through_untouched() { + for body in [ + r#"{"__type":"AccessDeniedException","Message":"no"}"#, + "bad gateway", + ] { + let Error::Provider { + body: reported, + status, + .. + } = error_class(body.into(), 403, Vec::new()) + else { + panic!("expected a provider error"); + }; + assert_eq!(reported, body); + assert_eq!(status, 403); + } + } +} diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs new file mode 100644 index 00000000000..ef07c1f24e1 --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs @@ -0,0 +1,3 @@ +pub mod analyze_transformation; +pub mod common_utils; +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs new file mode 100644 index 00000000000..148b7bed439 --- /dev/null +++ b/litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs @@ -0,0 +1,247 @@ +use litellm_core_utils::call_arguments::CallArguments; +use serde::{Deserialize, Serialize}; + +use super::common_utils::{ + Block, DocumentMetadata, HEALTH_CHECK_IMAGE_DATA_URI, TextractDocument, TextractEnvironment, + document_bytes, endpoint, environment, error_class, inline_document, lines_by_page, +}; +use crate::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrRequestContext, + OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, + }, +}; + +const DETECT_DOCUMENT_TEXT_TARGET: &str = "Textract.DetectDocumentText"; + +#[derive(Debug, Deserialize, Serialize)] +pub struct DetectDocumentTextRequest { + #[serde(rename = "Document")] + pub document: TextractDocument, +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +pub struct DetectDocumentTextResponse { + #[serde(default)] + blocks: Vec, + document_metadata: Option, +} + +/// Synchronous `DetectDocumentText`: plain lines from one image or single-page document. +#[derive(Clone, Copy, Debug, Default)] +pub struct TextractDetectTextConfig; + +impl BaseOcrConfig for TextractDetectTextConfig { + type OcrParams = (); + type ProviderRequest = DetectDocumentTextRequest; + type Environment = TextractEnvironment; + + fn get_health_check_document(&self) -> OcrDocument { + OcrDocument::ImageUrl { + image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(), + extra_fields: Default::default(), + } + } + + fn map_ocr_params( + &self, + _non_default_params: &CallArguments, + _model: &str, + ) -> Result<(), Error> { + Ok(()) + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + environment(request, DETECT_DOCUMENT_TEXT_TARGET).await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _optional_params: &(), + environment: &TextractEnvironment, + ) -> Result { + Ok(endpoint(request, environment)) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + _optional_params: &(), + _headers: &[(String, String)], + ) -> Result { + Ok(DetectDocumentTextRequest { + document: document_bytes(&document)?, + }) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &(), + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let document = inline_document(document, context).await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result { + decode_and_normalize_response(model, raw_response, request_format, normalize_response) + } + + fn get_error_class( + &self, + error_message: String, + status_code: u16, + headers: Vec<(String, String)>, + ) -> Error { + error_class(error_message, status_code, headers) + } +} + +fn normalize_response( + model: &str, + response: DetectDocumentTextResponse, +) -> Result { + let pages: Vec = lines_by_page(&response.blocks) + .into_iter() + .map(|(page, markdown)| OcrPage { + index: page - 1, + markdown, + ..Default::default() + }) + .collect(); + let pages_processed = response + .document_metadata + .and_then(|metadata| metadata.pages) + .or_else(|| i64::try_from(pages.len()).ok()); + Ok(LiteLLMOcrResponse { + usage_info: Some(OcrUsageInfo { + pages_processed, + ..Default::default() + }), + ..LiteLLMOcrResponse::new(model, pages) + }) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn normalize(response: serde_json::Value) -> LiteLLMOcrResponse { + TextractDetectTextConfig + .transform_ocr_response( + "detect-document-text", + &serde_json::to_vec(&response).unwrap(), + OcrResponseFormat::Litellm, + ) + .unwrap() + } + + #[test] + fn lines_become_one_markdown_page_and_words_are_not_repeated() { + let response = normalize(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"BlockType": "PAGE"}, + {"BlockType": "LINE", "Text": "Invoice 12345"}, + {"BlockType": "WORD", "Text": "Invoice"}, + {"BlockType": "WORD", "Text": "12345"}, + {"BlockType": "LINE", "Text": "total 67.89"} + ] + })); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "Invoice 12345\ntotal 67.89"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); + } + + #[test] + fn lines_are_grouped_by_their_page_in_page_order() { + let response = normalize(json!({ + "DocumentMetadata": {"Pages": 2}, + "Blocks": [ + {"BlockType": "LINE", "Text": "second", "Page": 2}, + {"BlockType": "LINE", "Text": "first", "Page": 1}, + {"BlockType": "LINE", "Text": "also second", "Page": 2} + ] + })); + + let pages: Vec<(i64, &str)> = response + .pages + .iter() + .map(|page| (page.index, page.markdown.as_str())) + .collect(); + assert_eq!(pages, vec![(0, "first"), (1, "second\nalso second")]); + } + + #[test] + fn a_multi_page_rejection_is_explained_to_the_caller() { + let error = TextractDetectTextConfig.get_error_class( + r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(), + 400, + Vec::new(), + ); + + assert!( + error + .to_string() + .contains("multi-page documents are not supported") + ); + } + + #[test] + fn the_request_carries_the_document_bytes_without_the_data_uri_envelope() { + let request = TextractDetectTextConfig + .transform_ocr_request( + "detect-document-text", + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,aGVsbG8=".into(), + extra_fields: Default::default(), + }, + &(), + &[], + ) + .unwrap(); + + assert_eq!( + serde_json::to_value(request).unwrap(), + json!({"Document": {"Bytes": "aGVsbG8="}}) + ); + } + + #[test] + fn a_remote_url_is_refused_by_the_sync_transform() { + let error = TextractDetectTextConfig + .transform_ocr_request( + "detect-document-text", + OcrDocument::DocumentUrl { + document_url: "https://example.com/a.pdf".into(), + extra_fields: Default::default(), + }, + &(), + &[], + ) + .unwrap_err(); + + assert!(matches!(error, Error::InvalidDataUri)); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs b/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs index dd4588732be..1257bbf0d6a 100644 --- a/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs @@ -21,14 +21,7 @@ impl AudioTranscriptionResponseData { } } -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum AudioTranscriptionAuth { - Bearer, - AwsSigV4 { - region: String, - service: &'static str, - }, -} +pub use litellm_auth::RequestAuth; pub trait BaseAudioTranscriptionConfig: Sync { fn get_supported_openai_params(&self) -> &'static [&'static str]; @@ -70,5 +63,5 @@ pub trait BaseAudioTranscriptionConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; + ) -> Result; } diff --git a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs index ac0450c25f0..c7d1a27c71e 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs @@ -41,14 +41,7 @@ pub const STREAM_PARAM: &str = "stream"; /// presence does not make a request untranslatable. const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; -/// How the upstream call is authenticated. API-key strategies are resolved in -/// `prepare`; SigV4 needs the serialized body, so the handler signs it. -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum ChatCompletionsAuth { - Header { name: &'static str, value: String }, - Bearer { token: String }, - AwsSigV4 { region: String }, -} +pub use litellm_auth::RequestAuth; /// Why a request cannot be served by the Rust path. /// @@ -91,7 +84,7 @@ pub trait BaseConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; + ) -> Result; fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[("content-type", "application/json")] diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 9fce387beb5..08f217351ba 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -100,6 +100,8 @@ pub enum Error { Params(#[from] litellm_core_utils::params::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), + #[error(transparent)] + Http(#[from] litellm_http::Error), } impl From for Error { @@ -155,6 +157,7 @@ impl Error { | Self::InvalidProvider(_) | Self::Params(_) | Self::Headers(_) + | Self::Http(_) ) } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 91fb6461770..245261d9f92 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -5,7 +5,7 @@ use litellm_host::event::WireRequest; use litellm_http::{ ClientVariant, HttpClientConfig, HttpClientPool, media::{MediaFetcher, UrlPolicy}, - request::{HeaderPolicy, execute_http_request, with_headers}, + outbound::{OutboundRequest, RequestSigner}, transport, }; use serde::{Serialize, de::DeserializeOwned}; @@ -117,8 +117,9 @@ pub async fn ocr( ) -> Result { let http = config.prepare_request(request, client, hooks).await?; let url = http.url().to_string(); - let headers = request_headers(&http)?; - let response = execute_http_request(client.provider_http(), http) + let headers = http.headers().to_vec(); + let response = http + .send(client.provider_http()) .await .map_err(transport_error)?; if !response.status().is_success() { @@ -153,21 +154,6 @@ pub async fn ocr( .await } -fn request_headers(request: &reqwest::Request) -> Result, Error> { - request - .headers() - .iter() - .map(|(name, value)| { - value - .to_str() - .map(|value| (name.to_string(), value.to_string())) - .map_err(|_| Error::RequestField { - path: "headers".into(), - }) - }) - .collect() -} - pub async fn read_json_response( response: reqwest::Response, native: bool, @@ -222,13 +208,13 @@ pub fn transport_error(error: reqwest::Error) -> Error { pub async fn transform_request_body( config: &C, - client: &OcrClient, request: &PreparedOcrRequest, url: &str, headers: &[(String, String)], body: B, + signer: Option<&dyn RequestSigner>, hooks: &dyn CallHooks, -) -> Result { +) -> Result { let composed = litellm_core_utils::call_arguments::compose_body( &request.optional_params, &body, @@ -244,7 +230,17 @@ pub async fn transform_request_body( }); } config.validate_request_body(&changed.body)?; - build_http_request(client, request, url, &changed.headers, &changed.body) + let timeout = Some(request.connection.timeout); + Ok(match signer { + Some(signer) => OutboundRequest::signed_json( + url.into(), + changed.headers, + &changed.body, + timeout, + signer, + ), + None => OutboundRequest::json(url.into(), changed.headers, &changed.body, timeout), + }?) } fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest { @@ -255,22 +251,18 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq } } -pub fn build_http_request( - client: &OcrClient, +pub fn build_http_request( request: &PreparedOcrRequest, - url: &str, - headers: &[(String, String)], - body: &B, -) -> Result { - let builder = client - .provider_http() - .post(url) - .json(body) - .timeout(request.connection.timeout); - with_headers(builder, headers, HeaderPolicy::All) - .build() - .map_err(transport::Error::from) - .map_err(Error::from) + url: String, + headers: Vec<(String, String)>, + body: &impl Serialize, +) -> Result { + Ok(OutboundRequest::json( + url, + headers, + body, + Some(request.connection.timeout), + )?) } pub async fn guardrail_document( diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 3960282b580..e02a4b7f266 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -6,6 +6,7 @@ use litellm_core_utils::{ serde_compat::{FiniteF64, LaxI64}, settings::ProcessEnvironment, }; +use litellm_http::outbound::{OutboundRequest, RequestSigner}; use serde::{ Deserialize, Serialize, de::{DeserializeOwned, IntoDeserializer}, @@ -394,6 +395,10 @@ const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQ /// (headers at minimum; Vertex also carries the project id). pub trait OcrEnvironment: Send + Sync { fn headers(&self) -> &[(String, String)]; + + fn signer(&self) -> Option<&dyn RequestSigner> { + None + } } impl OcrEnvironment for Vec<(String, String)> { @@ -536,7 +541,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> impl Future> + Send { + ) -> impl Future> + Send { async move { let params = self.map_ocr_params(&request.optional_params, &request.model)?; let environment = self.validate_environment(request, client).await?; @@ -554,7 +559,16 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { }, ) .await?; - transform_request_body(self, client, request, &url, headers, body, hooks).await + transform_request_body( + self, + request, + &url, + headers, + body, + environment.signer(), + hooks, + ) + .await } } } diff --git a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs index 39734d844da..cfabcb12341 100644 --- a/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs @@ -8,8 +8,8 @@ use serde_json::{Map, Value, json}; use crate::base_llm::{ audio_transcription::transformation::{ - AudioTranscriptionAuth, AudioTranscriptionRequestData, AudioTranscriptionResponseData, - BaseAudioTranscriptionConfig, + AudioTranscriptionRequestData, AudioTranscriptionResponseData, + BaseAudioTranscriptionConfig, RequestAuth, }, chat::transformation::Error, }; @@ -136,9 +136,9 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { + ) -> Result { let (_, model_region) = bedrock_model_id_and_region(model); - Ok(AudioTranscriptionAuth::AwsSigV4 { + Ok(RequestAuth::AwsSigV4 { region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), service: BEDROCK_SERVICE, }) diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 23c6c5c61bd..09c456f1d0a 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,6 +1,6 @@ use litellm_auth_aws::{ bedrock_model_id_and_region, - constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}, + constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, resolve_bedrock_region, }; use litellm_core_utils::{ @@ -17,8 +17,8 @@ use litellm_types::{ use serde_json::{Map, Value, json}; use crate::base_llm::chat::transformation::{ - BaseConfig, ChatCompletionsAuth, Error, ProviderChatRequestData, ProviderChatResponseData, - Unsupported, unsupported_message, unsupported_param, + BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, Unsupported, + unsupported_message, unsupported_param, }; /// Converse parameter names, post `map_openai_params`, that the Rust path can @@ -186,7 +186,7 @@ impl BaseConfig for AmazonConverseConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { + ) -> Result { // Python reads `api_key` as the Bedrock bearer token and consults the // env only when the caller passed none, so a caller-supplied empty key // falls through to SigV4 without reaching for the environment. An @@ -199,11 +199,12 @@ impl BaseConfig for AmazonConverseConfig { } .filter(|token| !token.is_empty()); if let Some(token) = bearer { - return Ok(ChatCompletionsAuth::Bearer { token }); + return Ok(RequestAuth::Bearer { token }); } let (_, model_region) = bedrock_model_id_and_region(model); - Ok(ChatCompletionsAuth::AwsSigV4 { + Ok(RequestAuth::AwsSigV4 { region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), + service: BEDROCK_SERVICE, }) } diff --git a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs b/litellm-rust/crates/llms/src/bedrock/chat/tests.rs index cca5cbda41a..d7ecde47c6b 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/tests.rs @@ -281,8 +281,9 @@ fn signs_with_sigv4_in_the_resolved_region() { &|_| None ) .expect("auth resolves"), - ChatCompletionsAuth::AwsSigV4 { - region: "eu-central-1".to_string() + RequestAuth::AwsSigV4 { + region: "eu-central-1".to_string(), + service: "bedrock", } ); } @@ -306,11 +307,12 @@ fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() { ) .expect("auth resolves") }; - let bearer = |token: &str| ChatCompletionsAuth::Bearer { + let bearer = |token: &str| RequestAuth::Bearer { token: token.to_string(), }; - let sigv4 = ChatCompletionsAuth::AwsSigV4 { + let sigv4 = RequestAuth::AwsSigV4 { region: "eu-central-1".to_string(), + service: "bedrock", }; // A caller-supplied key is the bearer token, and outranks the env. diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 8d1bb366ed4..701eaff4374 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -1,4 +1,5 @@ pub mod anthropic; +pub mod aws_textract; pub mod azure_ai; pub mod base_llm; pub mod bedrock; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 307ba697316..5272be97c24 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -5,6 +5,7 @@ use litellm_core_utils::{ params::OpaqueParams, url_utils::ApiUrl, }; +use litellm_http::outbound::OutboundRequest; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value, json}; @@ -166,7 +167,7 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -251,7 +252,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -264,7 +265,7 @@ async fn prepare_upload_request, -) -> Result { +) -> Result { let params = config.map_ocr_params(&request.optional_params, &request.model)?; let headers = config.validate_environment(request, client).await?; let url = config.get_complete_url(request, ¶ms, &headers)?; @@ -286,7 +287,7 @@ async fn prepare_upload_request Result { diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 19d28f76b6f..6c5a65173e3 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -58,6 +58,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { audio_transcription::Error::InvalidProvider(_) | audio_transcription::Error::InvalidRequest(_) | audio_transcription::Error::Headers(_) + | audio_transcription::Error::Http(_) | audio_transcription::Error::InvalidType { .. } | audio_transcription::Error::MissingField(_) | audio_transcription::Error::Aws(_) => true, @@ -68,6 +69,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { chat_completions::Error::InvalidProvider(_) | chat_completions::Error::InvalidRequest(_) | chat_completions::Error::Headers(_) + | chat_completions::Error::Http(_) | chat_completions::Error::InvalidType { .. } | chat_completions::Error::MissingField(_) | chat_completions::Error::Aws(_) => true, @@ -105,6 +107,7 @@ pub(crate) fn chat_completions_error_to_pyerr(error: chat_completions::Error) -> | Error::InvalidType { .. } | Error::MissingField(_) | Error::Headers(_) + | Error::Http(_) | Error::Transport(TransportError::Connect(_)) => { RustBridgeDeclined::new_err(error.to_string()) } diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 80c93273d1e..55b19458b7a 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -53,7 +53,9 @@ _PYTHON_AOCR: Final = cast( # cast-ok: forward the original call shape through def _context(request: LiteLLMOcrRequest) -> Context: - return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model) + prefix, separator, _ = request.model.partition("/") + provider: Final = request.custom_llm_provider or (prefix if separator else None) + return Context(Route.OCR, provider=provider, model=request.model) _DISPATCH: Final = PublicDispatch( diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index e6673ec99aa..d7e100dd630 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -209,6 +209,94 @@ ], "default_model_placeholder": "claude-3-opus" }, + { + "provider": "AWS_Textract", + "provider_display_name": "Amazon Textract", + "litellm_provider": "aws_textract", + "credential_fields": [ + { + "key": "aws_access_key_id", + "label": "AWS Access Key ID", + "placeholder": null, + "tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_secret_access_key", + "label": "AWS Secret Access Key", + "placeholder": null, + "tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_session_token", + "label": "AWS Session Token", + "placeholder": null, + "tooltip": "Temporary credentials session token. You can provide the raw token or the environment variable (e.g. `os.environ/MY_SESSION_TOKEN`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "aws_region_name", + "label": "AWS Region Name", + "placeholder": "us-east-1", + "tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_session_name", + "label": "AWS Session Name", + "placeholder": "my-session", + "tooltip": "Name for the AWS session. You can provide the raw value or the environment variable (e.g. `os.environ/MY_SESSION_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_profile_name", + "label": "AWS Profile Name", + "placeholder": "default", + "tooltip": "AWS profile name to use for authentication. You can provide the raw value or the environment variable (e.g. `os.environ/MY_PROFILE_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_role_name", + "label": "AWS Role Name", + "placeholder": "MyRole", + "tooltip": "AWS IAM role name to assume. You can provide the raw value or the environment variable (e.g. `os.environ/MY_ROLE_NAME`).", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "aws_web_identity_token", + "label": "AWS Web Identity Token", + "placeholder": null, + "tooltip": "Web identity token for OIDC authentication. You can provide the raw token or the environment variable (e.g. `os.environ/MY_WEB_IDENTITY_TOKEN`).", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "detect-document-text" + }, { "provider": "BedrockMantle", "provider_display_name": "Amazon Bedrock Mantle", diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index d843a874fe3..8794ff2db95 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -58,6 +58,7 @@ class Rule: Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( + Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), Rule(Route.OCR, Rollout.RUST_OPT_OUT), Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 378f8fec9d5..d416e2af33a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4020,6 +4020,7 @@ class LlmProviders(str, Enum): BYTEZ = "bytez" REPLICATE = "replicate" REDUCTO = "reducto" + AWS_TEXTRACT = "aws_textract" RUNWAYML = "runwayml" AWS_POLLY = "aws_polly" TRANSCRIBE = "transcribe" diff --git a/tests/test_litellm/ocr/test_dispatch.py b/tests/test_litellm/ocr/test_dispatch.py index 14d3368f869..e54d4070ba8 100644 --- a/tests/test_litellm/ocr/test_dispatch.py +++ b/tests/test_litellm/ocr/test_dispatch.py @@ -387,3 +387,39 @@ async def test_public_aocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPat NATIVE_AOCR.reset() assert result is expected assert [request.model for request in captured] == ["mistral/mistral-ocr-latest"] + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + ( + ("aws_textract/detect-document-text", None, "native"), + ("detect-document-text", "aws_textract", "native"), + ("mistral/mistral-ocr-latest", None, "python"), + ("mistral/mistral-ocr-latest", "aws_textract", "native"), + ("aws_textract", None, "python"), + ), +) +def test_provider_scoped_rule_sees_the_provider_named_by_the_model_prefix( + model: str, custom_llm_provider: str | None, expected: str +) -> None: + rules: Final[Rules] = ( + Rule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), + Rule(Route.OCR, Rollout.PYTHON_ONLY), + ) + document: Final[Mapping[str, object]] = {"type": "image_url", "image_url": "data:image/png;base64,YQ=="} + kwargs: Final[Mapping[str, object]] = ( + {} if custom_llm_provider is None else {"custom_llm_provider": custom_llm_provider} + ) + python_response: Final = response("python") + native_response: Final = response("native") + + result: Final = _DISPATCH.run( + (model, document), + kwargs, + python=lambda *_args, **_kwargs: python_response, + binding=ocr_binding(lambda *_args, **_kwargs: native_response), + native=lambda _hook, _request, _args, _kwargs: native_response, + rules=rules, + ) + + assert cast(OCRResponse, result).model == expected # noqa: TID251 # sync dispatch returns the response itself diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index e9fdbf859f4..147e863baf5 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -85,3 +85,15 @@ def test_first_matching_rule_respects_every_constraint(context: Context, expecte ) assert catalog.decision(context, rules) is expected + + +@pytest.mark.parametrize("process", (None, False, True)) +@pytest.mark.parametrize("environment", (None, "0", "1")) +def test_textract_ocr_has_no_python_path_to_opt_out_to( + monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None +) -> None: + configuration.rust(process) + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + + assert catalog.decision(Context(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED