scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

feat(rust): add Amazon Textract to litellm.ocr and sign provider requests after host hooks

Browse source 
Add an aws_textract OCR provider on the Rust route, with no Python path. The
detect-document-text model returns plain lines and analyze-document renders
layout and tables as markdown. Both use Textract's synchronous API, so a
multi-page PDF or TIFF is rejected with an error that names the single-page
limit. A call with no region fails instead of falling back to Bedrock's default

SigV4 covers the request body, and host hooks can rewrite that body before it
is sent. litellm-http now has OutboundRequest, which serializes the body once,
shows those bytes to a RequestSigner and is the only thing a route can send.
Chat, audio transcription and OCR build it after their hooks ran, so a callback
that redacts the body still produces a valid Bedrock or Textract signature

ChatCompletionsAuth and AudioTranscriptionAuth are replaced by
litellm_auth::RequestAuth, and one helper in core turns it into a signed or
unsigned request. Audio transcription now signs only the AWS header set and
rejects a forwarded header that SigV4 computes, the same as chat

The OCR catalog routes aws_textract as Rust required, and the dispatch context
reads the provider from the model prefix so a provider scoped rule can match
This commit is contained in:
Yujong Lee 2026-09-19 09:11:37 -07:00
parent 362be56bb0
commit c404bed9f0
53 changed files with 1987 additions and 258 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/Cargo.lock generated
View file

@ -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",

1
litellm-rust/crates/auth-aws/Cargo.toml
View file

@ -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

73
litellm-rust/crates/auth-aws/src/aws.rs
View file

@ -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<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, 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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> 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<String, Value>) -> Option
#[cfg(test)]
mod tests {
use super::*;
use crate::constants::BEDROCK_SERVICE;
fn no_env(_: &str) -> Option<String> {
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"), &params, &region_name),
resolve_aws_region(Some("us-east-2"), &Map::new(), &region_name),
resolve_aws_region(None, &Map::new(), &region_name),
resolve_aws_region(None, &Map::new(), &region),
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<u8>, BTreeMap<String, String>) {
(
"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(),
)?;

3
litellm-rust/crates/auth-aws/src/lib.rs
View file

@ -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;

183
litellm-rust/crates/auth-aws/src/signer.rs Normal file
View file

@ -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<String, Value>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Self, Error> {
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<Vec<(String, String)>, 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<String, String> = 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<String, String> = 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())
);
}
}

19
litellm-rust/crates/auth/src/http.rs
View file

@ -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)]

2
litellm-rust/crates/core/src/audio_transcription/error.rs
View file

@ -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),
}

55
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -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<Value, Error> {
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::<Error>(
&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<Vec<(String, String)>, 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<String, String> = 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())
}

17
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -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()));

4
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -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<String, Value>,
pub timeout: Option<Duration>,
}

2
litellm-rust/crates/core/src/chat_completions/error.rs
View file

@ -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),
}

90
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -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<ChatCompletionsResponse, Error> {
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<Vec<(String, String)>, 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<String, String> = 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<OutboundRequest, Error> {
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,
})
}

10
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -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() {

23
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -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()
}
);

4
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -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<String, Value>,
pub timeout: Option<Duration>,
}

1
litellm-rust/crates/core/src/lib.rs
View file

@ -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;

16
litellm-rust/crates/core/src/ocr/arguments.rs
View file

@ -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"
)
}

4
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -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;

5
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -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(|| {

23
litellm-rust/crates/core/src/ocr/provider_config.rs
View file

@ -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::<OcrProvider>()
.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)]

30
litellm-rust/crates/core/src/outbound.rs Normal file
View file

@ -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<E>(
auth: &RequestAuth,
url: String,
headers: Vec<(String, String)>,
body: &Value,
timeout: Option<Duration>,
optional_params: &Map<String, Value>,
) -> Result<OutboundRequest, E>
where
E: From<litellm_http::Error> + From<litellm_auth_aws::Error>,
{
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,
)?)
}

193
litellm-rust/crates/core/tests/aws_textract_ocr.rs Normal file
View file

@ -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<String, String> = ["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");
}

4
litellm-rust/crates/core/tests/cohere_ocr.rs
View file

@ -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());
}

26
litellm-rust/crates/core/tests/vertex_ai_ocr.rs
View file

@ -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!({

1
litellm-rust/crates/http/Cargo.toml
View file

@ -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

6
litellm-rust/crates/http/src/error.rs
View file

@ -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<reqwest::Error> for Error {

1
litellm-rust/crates/http/src/lib.rs
View file

@ -1,6 +1,7 @@
mod config;
mod error;
pub mod media;
pub mod outbound;
mod pool;
mod proxy;
pub mod request;

210
litellm-rust/crates/http/src/outbound.rs Normal file
View file

@ -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<Vec<(String, String)>, Error>;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OutboundRequest {
url: String,
headers: Vec<(String, String)>,
body: Vec<u8>,
timeout: Option<Duration>,
}
impl OutboundRequest {
pub fn json(
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
timeout: Option<Duration>,
) -> Result<Self, Error> {
Self::build(url, headers, body, timeout, None)
}
pub fn signed_json(
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
timeout: Option<Duration>,
signer: &dyn RequestSigner,
) -> Result<Self, Error> {
Self::build(url, headers, body, timeout, Some(signer))
}
fn build(
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
timeout: Option<Duration>,
signer: Option<&dyn RequestSigner>,
) -> Result<Self, Error> {
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<Duration> {
self.timeout
}
pub async fn send(self, client: &reqwest::Client) -> Result<reqwest::Response, reqwest::Error> {
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<Vec<u8>>);
impl RequestSigner for Recording {
fn sign(&self, request: UnsignedRequest<'_>) -> Result<Vec<(String, String)>, 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<Vec<(String, String)>, 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<Vec<(String, String)>, Error> {
Err(Error::ComputedHeader("authorization".into()))
}
}
assert_eq!(
OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses),
Err(Error::ComputedHeader("authorization".into()))
);
}
}

2
litellm-rust/crates/llms/src/anthropic/chat/tests.rs
View file

@ -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()
}

6
litellm-rust/crates/llms/src/anthropic/chat/transformation.rs
View file

@ -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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
Ok(ChatCompletionsAuth::Header {
) -> Result<RequestAuth, Error> {
Ok(RequestAuth::Header {
name: "x-api-key",
value: resolve_anthropic_api_key(api_key, env_lookup)?,
})

1
litellm-rust/crates/llms/src/aws_textract/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod ocr;

12
litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md Normal file
View file

@ -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

426
litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs Normal file
View file

@ -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<Vec<String>>,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct AnalyzeDocumentRequest {
#[serde(rename = "Document")]
pub document: TextractDocument,
#[serde(rename = "FeatureTypes")]
pub feature_types: Vec<String>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct AnalyzeDocumentResponse {
#[serde(default)]
blocks: Vec<Block>,
document_metadata: Option<DocumentMetadata>,
}
/// 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<AnalyzeDocumentOptions, Error> {
Ok(parse_options(non_default_params)?)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<TextractEnvironment, Error> {
environment(request, ANALYZE_DOCUMENT_TARGET).await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &AnalyzeDocumentOptions,
environment: &TextractEnvironment,
) -> Result<String, Error> {
Ok(endpoint(request, environment))
}
fn transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
optional_params: &AnalyzeDocumentOptions,
_headers: &[(String, String)],
) -> Result<AnalyzeDocumentRequest, Error> {
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<AnalyzeDocumentRequest, Error> {
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<LiteLLMOcrResponse, Error> {
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<LiteLLMOcrResponse, Error> {
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<i64> = 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<OcrPage> = 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<String> = 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::<Vec<_>>()
.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::<Vec<_>>()
.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<usize> = 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::<Vec<_>>()
.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(), &params, &[])
.unwrap(),
)
.unwrap()
};
assert_eq!(
request(json!({})),
json!({"Document": {"Bytes": "aGk="}, "FeatureTypes": ["LAYOUT", "TABLES"]})
);
assert_eq!(
request(json!({"feature_types": ["FORMS"]}))["FeatureTypes"],
json!(["FORMS"])
);
}
}

247
litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs Normal file
View file

@ -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<String>,
pub page: Option<i64>,
pub row_index: Option<usize>,
pub column_index: Option<usize>,
#[serde(default)]
pub relationships: Vec<Relationship>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(super) struct Relationship {
pub r#type: String,
#[serde(default)]
pub ids: Vec<String>,
}
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<Item = &str> {
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<i64>,
}
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<TextractEnvironment, Error> {
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<TextractDocument, Error> {
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<OcrDocument, Error> {
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::<AwsError>(&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<i64> = 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"}"#,
"<html>bad gateway</html>",
] {
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);
}
}
}

3
litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs Normal file
View file

@ -0,0 +1,3 @@
pub mod analyze_transformation;
pub mod common_utils;
pub mod transformation;

247
litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs Normal file
View file

@ -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<Block>,
document_metadata: Option<DocumentMetadata>,
}
/// 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<TextractEnvironment, Error> {
environment(request, DETECT_DOCUMENT_TEXT_TARGET).await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &(),
environment: &TextractEnvironment,
) -> Result<String, Error> {
Ok(endpoint(request, environment))
}
fn transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
_optional_params: &(),
_headers: &[(String, String)],
) -> Result<DetectDocumentTextRequest, Error> {
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<DetectDocumentTextRequest, Error> {
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<LiteLLMOcrResponse, Error> {
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<LiteLLMOcrResponse, Error> {
let pages: Vec<OcrPage> = 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));
}
}

11
litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs
View file

@ -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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AudioTranscriptionAuth, Error>;
) -> Result<RequestAuth, Error>;
}

11
litellm-rust/crates/llms/src/base_llm/chat/transformation.rs
View file

@ -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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error>;
) -> Result<RequestAuth, Error>;
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[("content-type", "application/json")]

3
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -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<litellm_host::machine::MachineFault> for Error {
@ -155,6 +157,7 @@ impl Error {
| Self::InvalidProvider(_)
| Self::Params(_)
| Self::Headers(_)
| Self::Http(_)
)
}

64
litellm-rust/crates/llms/src/base_llm/ocr/handler.rs
View file

@ -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<C: BaseOcrConfig>(
) -> Result<LiteLLMOcrResponse, Error> {
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<C: BaseOcrConfig>(
.await
}
fn request_headers(request: &reqwest::Request) -> Result<Vec<(String, String)>, 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<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
@ -222,13 +208,13 @@ pub fn transport_error(error: reqwest::Error) -> Error {
pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
config: &C,
client: &OcrClient,
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: B,
signer: Option<&dyn RequestSigner>,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
let composed = litellm_core_utils::call_arguments::compose_body(
&request.optional_params,
&body,
@ -244,7 +230,17 @@ pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
});
}
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<B: Serialize>(
client: &OcrClient,
pub fn build_http_request(
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: &B,
) -> Result<reqwest::Request, Error> {
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<OutboundRequest, Error> {
Ok(OutboundRequest::json(
url,
headers,
body,
Some(request.connection.timeout),
)?)
}
pub async fn guardrail_document(

18
litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs
View file

@ -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<Error>,
) -> impl Future<Output = Result<reqwest::Request, Error>> + Send {
) -> impl Future<Output = Result<OutboundRequest, Error>> + 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
}
}
}

8
litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs
View file

@ -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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AudioTranscriptionAuth, Error> {
) -> Result<RequestAuth, Error> {
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,
})

13
litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs
View file

@ -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<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
) -> Result<RequestAuth, Error> {
// 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,
})
}

10
litellm-rust/crates/llms/src/bedrock/chat/tests.rs
View file

@ -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.

1
litellm-rust/crates/llms/src/lib.rs
View file

@ -1,4 +1,5 @@
pub mod anthropic;
pub mod aws_textract;
pub mod azure_ai;
pub mod base_llm;
pub mod bedrock;

9
litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

@ -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<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
prepare_upload_request(self, request, client, hooks).await
}
}
@ -251,7 +252,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
request: &PreparedOcrRequest,
client: &OcrClient,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
prepare_upload_request(self, request, client, hooks).await
}
}
@ -264,7 +265,7 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
request: &PreparedOcrRequest,
client: &OcrClient,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
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, &params, &headers)?;
@ -286,7 +287,7 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
&body,
config.get_supported_ocr_params(&request.model),
)?;
build_http_request(client, request, &url, &headers, &body)
build_http_request(request, url, headers, &body)
}
fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, Error> {

3
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -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())
}

4
litellm/ocr/dispatch.py
View file

@ -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(

88
litellm/proxy/public_endpoints/provider_create_fields.json
View file

@ -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",

1
litellm/rust_bridge/catalog.py
View file

@ -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"})),

1
litellm/types/utils.py
View file

@ -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"

36
tests/test_litellm/ocr/test_dispatch.py
View file

@ -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

12
tests/test_litellm/rust_bridge/test_catalog.py
View file

@ -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