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

Merge remote-tracking branch 'origin/main' into litellm-providers/price-sync

Browse source
This commit is contained in:
kerry 2026-09-17 19:29:18 +00:00
commit 5e5e5086f0
318 changed files with 17942 additions and 9117 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -3009,7 +3009,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, browser]
suite: [management, accounting, database, providers, extensions, sdk, browser]
filters:
branches:
only:

2
.github/workflows/test-unit.yml vendored
View file

@ -100,6 +100,7 @@ jobs:
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/chat_completions
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
@ -109,6 +110,7 @@ jobs:
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/messages
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag

151
litellm-rust/Cargo.lock generated
View file

@ -948,8 +948,18 @@ version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
dependencies = [
"darling_core",
"darling_macro",
"darling_core 0.20.11",
"darling_macro 0.20.11",
]
[[package]]
name = "darling"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0"
dependencies = [
"darling_core 0.21.3",
"darling_macro 0.21.3",
]
[[package]]
@ -966,13 +976,38 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "darling_core"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4"
dependencies = [
"fnv",
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
dependencies = [
"darling_core",
"darling_core 0.20.11",
"quote",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81"
dependencies = [
"darling_core 0.21.3",
"quote",
"syn 2.0.119",
]
@ -1022,7 +1057,7 @@ version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
dependencies = [
"darling",
"darling 0.20.11",
"proc-macro2",
"quote",
"syn 2.0.119",
@ -1363,7 +1398,7 @@ dependencies = [
"futures-sink",
"futures-util",
"http 0.2.12",
"indexmap",
"indexmap 2.14.0",
"slab",
"tokio",
"tokio-util",
@ -1382,7 +1417,7 @@ dependencies = [
"futures-core",
"futures-sink",
"http 1.4.2",
"indexmap",
"indexmap 2.14.0",
"slab",
"tokio",
"tokio-util",
@ -1400,6 +1435,12 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.17.1"
@ -1736,6 +1777,17 @@ dependencies = [
"icu_properties",
]
[[package]]
name = "indexmap"
version = "1.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
dependencies = [
"autocfg",
"hashbrown 0.12.3",
"serde",
]
[[package]]
name = "indexmap"
version = "2.14.0"
@ -1743,7 +1795,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
dependencies = [
"equivalent",
"hashbrown",
"hashbrown 0.17.1",
"serde",
"serde_core",
]
@ -1953,6 +2005,8 @@ dependencies = [
name = "litellm-core"
version = "0.1.0"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-types",
"base64 0.22.1",
"bytes",
"data-url",
@ -1961,6 +2015,7 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-framing",
"mime_guess",
"moka",
"rand 0.8.7",
@ -1971,10 +2026,12 @@ dependencies = [
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"sha2 0.10.9",
"strum",
"subtle",
"thiserror 2.0.19",
"time",
"tokio",
"tokio-tungstenite",
"url",
@ -2032,7 +2089,7 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"indexmap",
"indexmap 2.14.0",
"itoa",
"rand 0.8.7",
"rstest",
@ -2753,6 +2810,26 @@ dependencies = [
"bitflags",
]
[[package]]
name = "ref-cast"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3"
dependencies = [
"ref-cast-impl",
]
[[package]]
name = "ref-cast-impl"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.0",
]
[[package]]
name = "regex"
version = "1.13.1"
@ -3075,6 +3152,30 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "schemars"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f"
dependencies = [
"dyn-clone",
"ref-cast",
"serde",
"serde_json",
]
[[package]]
name = "schemars"
version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
dependencies = [
"dyn-clone",
"ref-cast",
"serde",
"serde_json",
]
[[package]]
name = "scopeguard"
version = "1.2.0"
@ -3156,6 +3257,7 @@ version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"indexmap 2.14.0",
"itoa",
"memchr",
"serde",
@ -3186,6 +3288,37 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_with"
version = "3.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4fa237f2807440d238e0364a218270b98f767a00d3dada77b1c53ae88940e2e7"
dependencies = [
"base64 0.22.1",
"chrono",
"hex",
"indexmap 1.9.3",
"indexmap 2.14.0",
"schemars 0.9.0",
"schemars 1.2.2",
"serde_core",
"serde_json",
"serde_with_macros",
"time",
]
[[package]]
name = "serde_with_macros"
version = "3.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52a8e3ca0ca629121f70ab50f95249e5a6f925cc0f6ffe8256c45b728875706c"
dependencies = [
"darling 0.21.3",
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "sha1"
version = "0.10.7"
@ -3661,7 +3794,7 @@ version = "0.25.13+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b"
dependencies = [
"indexmap",
"indexmap 2.14.0",
"toml_datetime",
"toml_parser",
"winnow",

3
litellm-rust/Cargo.toml
View file

@ -11,6 +11,7 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
bytes = "1"
litellm-core = { path = "crates/core" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
@ -29,6 +30,7 @@ rustls = { version = "0.23", default-features = false, features = ["ring", "std"
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.16.1", default-features = false, features = ["std", "macros"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"
@ -39,6 +41,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
veil = "0.3.0"

22
litellm-rust/crates/core/AGENTS.md
View file

@ -1,7 +1,27 @@
litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src/<route>/` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back.
A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate.
A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. Provider code and base config traits live under `src/llms/`, mirroring their Python source paths. This applies to every API surface: shared orchestration stays in its route module (`ocr/`, `chat_completions/`, `messages/`, `audio_transcription/`, or `responses/`), while provider transformations live under the corresponding Python-mirrored `llms/<provider>/` path. Import implementations directly from their canonical paths; do not add a `src/providers/` layer or compatibility re-exports. Shared provider resolution lives under `src/litellm_core_utils/get_llm_provider_logic.rs`. Handlers belong in core, never in a host crate
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`.
Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.
## Python/Rust transformation pairs
Use the base OCR and Mistral OCR pairs as the reference when aligning transformations. Derive `src/<relative_path>.rs` from `litellm/<relative_path>.py`, preserving meaningful basenames such as `messages_transformation`
Keep corresponding operation names and parameter names when their responsibilities match. Rust types retain the Python semantic name with Rust acronym casing (`BaseOCRConfig` / `BaseOcrConfig`, `MistralOCRConfig` / `MistralOcrConfig`). Private Python helpers can drop their leading underscore. Give Rust adapter helpers distinct responsibility names rather than duplicating trait method names
Order OCR config methods as supported parameters, credential metadata and connection resolution, health-check input, parameter mapping, environment validation, URL construction, request transformation, async request transformation, response transformation, async response transformation, and error conversion. Put constants and data types before the config, private helpers after it in operation order, and tests last. Rust-only trait hooks follow the corresponding Python methods
Use trait defaults for unchanged inherited behavior and explicit delegation for shared provider behavior. Keep typed inputs, ownership, `Result`, and async I/O idiomatic. A matching path or symbol identifies the counterpart, not a claim of full behavioral parity
Use named `#[rstest]` cases for independent input/output scenarios instead of loops or repeated calls in one test. Inject reusable setup with `#[fixture]` arguments and use `#[with(...)]` for fixture overrides. Keep assertions about the same result together
For base OCR, Python response models correspond to `src/ocr/types.rs`; Rust context/environment types support the runtime. `BaseOcrConfig::prepare_request` corresponds to Python's HTTP-handler preparation rather than a `BaseOCRConfig` method, and `validate_request_body` is a Rust-only hook
For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages; layout refactors must preserve those behaviors and verify them with the existing tests
For non-OCR pairs, order corresponding methods as parameter support/mapping, environment validation, URL construction, request transformation, and response transformation, followed by Rust-only runtime hooks. Auth resolution remains split between configs and route preparation. Chat `supported_openai_param_mappings` describes accepted OpenAI/provider name pairs, unlike Python's `get_supported_openai_params` name list. Audio `map_transcription_params` remains a Rust filtering helper
Azure Messages maps to `llms/azure_ai/anthropic/messages_transformation.py`; Bedrock Converse maps to `llms/bedrock/chat/converse_transformation.py`. `AnthropicConfig`, `AmazonConverseConfig`, and the non-OCR base traits are partial ports. `OpenAiResponsesApiConfig` currently implements only the WebSocket surface. Preserve their acceptance gates, passthrough behavior, and host fallback contracts when aligning layout

7
litellm-rust/crates/core/Cargo.toml
View file

@ -15,6 +15,7 @@ litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-framing.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
rand.workspace = true
@ -22,16 +23,20 @@ reqwest.workspace = true
rustls.workspace = true
rustls-native-certs.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_with.workspace = true
serde_path_to_error = "0.1"
strum.workspace = true
subtle.workspace = true
tokio = { workspace = true, features = ["sync"] }
tokio-tungstenite.workspace = true
thiserror.workspace = true
time.workspace = true
sha2.workspace = true
url.workspace = true
veil.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-types = "1.6.1"
rstest.workspace = true

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

@ -1,10 +1,9 @@
use serde_json::Value;
use super::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::types::ProviderAudioTranscriptionRequest;
use crate::http_utils::{http_request, truncate_error_body};
pub async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
@ -37,7 +36,7 @@ pub async fn execute_audio_transcription_provider_call(
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.transform_audio_transcription_response(&request.model, response_json)?
.into_json())
}
@ -48,9 +47,8 @@ async fn signed_headers(
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
use crate::providers::bedrock::audio_transcription::aws_auth_config;
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use crate::llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());

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

@ -3,13 +3,11 @@ pub use error::Error;
mod client;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
use serde_json::Value;
pub use handler::execute_audio_transcription_provider_call;
pub use prepare::prepare_audio_transcription_provider_call;
use serde_json::Value;
pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {

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

@ -1,12 +1,15 @@
use super::Error;
use crate::http_utils::{has_header, string_headers};
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use crate::http_utils::{has_header, string_headers};
use crate::litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, get_custom_llm_provider,
};
use crate::llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
};
use crate::llms::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
fn provider_config(provider: &str) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
if provider == "bedrock" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
}
@ -46,7 +49,7 @@ pub fn prepare_audio_transcription_provider_call(
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));
}
let url = config.complete_url(
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
@ -54,7 +57,7 @@ pub fn prepare_audio_transcription_provider_call(
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_transcription_request(&model, request.audio, filtered_params)?;
config.transform_audio_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: provider_info.custom_llm_provider.to_string(),

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

@ -3,7 +3,9 @@ use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
use crate::llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
};
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
@ -20,7 +22,7 @@ pub struct AudioTranscriptionRequest<'a> {
pub struct ProviderAudioTranscriptionRequest {
pub(super) model: String,
pub(super) custom_llm_provider: String,
pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
pub(super) config: &'static dyn BaseAudioTranscriptionConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,

468
litellm-rust/crates/core/src/call_arguments.rs Normal file
View file

@ -0,0 +1,468 @@
use std::ops::Deref;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct CallArguments(Map<String, Value>);
impl CallArguments {
pub(crate) fn select(&self, names: &[&str]) -> Map<String, Value> {
self.iter()
.filter(|(name, _)| names.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("invalid argument: {path}")]
pub struct ArgumentError {
pub path: String,
}
pub fn parse_options<T: DeserializeOwned>(arguments: &CallArguments) -> Result<T, ArgumentError> {
let deserializer = serde::de::value::MapDeserializer::new(
arguments.iter().map(|(name, value)| (name.as_str(), value)),
);
serde_path_to_error::deserialize(deserializer).map_err(|error| ArgumentError {
path: error.path().to_string(),
})
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ArgumentSpec {
pub name: &'static str,
pub secret: bool,
}
pub fn should_project(name: &str, consumed: &[ArgumentSpec], bound_fields: &[&str]) -> bool {
consumed.iter().any(|field| field.name == name)
|| (!bound_fields.contains(&name) && !is_control(name))
}
pub fn is_control(name: &str) -> bool {
crate::params::is_control_param(name) || HOST_CONTROLS.contains(&name)
}
const HOST_CONTROLS: &[&str] = &[
"_agentic_loop_api_surface",
"_agentic_loop_depth",
"_agentic_loop_fingerprints",
"_code_interpreter_interception_active",
"_code_interpreter_interception_converted_stream",
"_code_interpreter_interception_sandbox_key",
"_code_interpreter_interception_session_scoped",
"_headroom_interception_converted_stream",
"_litellm_strip_stream_usage",
"_router_weights",
"_websearch_interception_converted_stream",
"_websearch_interception_emit_native_blocks",
"acompletion",
"adaptive_router_config",
"adaptive_router_default_model",
"aembedding",
"aimg_generation",
"allm_passthrough_route",
"allow_client_keepalive_override",
"allowed_model_region",
"allowed_openai_params",
"annotation_cost_per_page",
"api_version",
"arize_api_key",
"arize_space_id",
"arize_space_key",
"assistant_continue_message",
"async_call",
"atext_completion",
"attempted_targets",
"auto_router_config",
"auto_router_config_path",
"auto_router_default_model",
"auto_router_embedding_model",
"auto_router_max_input_chars",
"auto_router_model_compression",
"auto_router_routing_compression",
"aws_batch_role_arn",
"azure",
"azure_password",
"azure_username",
"base_model",
"bedrock_tags",
"bos_token",
"budget_duration",
"cache",
"cache_creation_input_audio_token_cost",
"cache_creation_input_token_cost",
"cache_creation_input_token_cost_above_1hr",
"cache_creation_input_token_cost_above_200k_tokens",
"cache_creation_input_token_cost_above_272k_tokens",
"cache_creation_input_token_cost_above_272k_tokens_flex",
"cache_creation_input_token_cost_above_272k_tokens_priority",
"cache_creation_input_token_cost_flex",
"cache_creation_input_token_cost_priority",
"cache_creation_input_token_cost_ultrafast",
"cache_key",
"cache_read_input_audio_token_cost",
"cache_read_input_token_cost",
"cache_read_input_token_cost_above_200k_tokens",
"cache_read_input_token_cost_above_200k_tokens_priority",
"cache_read_input_token_cost_above_272k_tokens",
"cache_read_input_token_cost_above_272k_tokens_flex",
"cache_read_input_token_cost_above_272k_tokens_priority",
"cache_read_input_token_cost_above_512k_tokens",
"cache_read_input_token_cost_flex",
"cache_read_input_token_cost_priority",
"cache_read_input_token_cost_ultrafast",
"caching",
"caching_groups",
"citation_cost_per_token",
"client",
"client_side_timeout",
"complete_response",
"completion_call_id",
"complexity_router_config",
"complexity_router_default_model",
"configurable_clientside_auth_params",
"context_window_fallback_dict",
"cooldown_time",
"cost_per_query",
"custom_prompt_dict",
"data_residency",
"dd_agent_host",
"dd_agent_port",
"dd_api_key",
"dd_site",
"default_api_key_rpm_limit",
"default_api_key_tpm_limit",
"disable_add_transform_inline_image_block",
"enable_json_schema_validation",
"enable_prompt_caching",
"enable_tag_filtering",
"ensure_alternating_roles",
"eos_token",
"fallback_depth",
"fallbacks",
"fastest_response",
"final_prompt_value",
"force_timeout",
"gcs_bucket_name",
"gcs_path_service_account",
"google_maps_grounding_cost_per_query",
"headers",
"hf_model_name",
"humanloop_api_key",
"id",
"input_cost_per_audio_per_second",
"input_cost_per_audio_per_second_above_128k_tokens",
"input_cost_per_audio_token",
"input_cost_per_audio_token_batches",
"input_cost_per_character",
"input_cost_per_character_above_128k_tokens",
"input_cost_per_image",
"input_cost_per_image_above_128k_tokens",
"input_cost_per_image_token",
"input_cost_per_image_token_batches",
"input_cost_per_pixel",
"input_cost_per_query",
"input_cost_per_second",
"input_cost_per_token",
"input_cost_per_token_above_128k_tokens",
"input_cost_per_token_above_200k_tokens",
"input_cost_per_token_above_200k_tokens_priority",
"input_cost_per_token_above_272k_tokens",
"input_cost_per_token_above_272k_tokens_flex",
"input_cost_per_token_above_272k_tokens_priority",
"input_cost_per_token_above_512k_tokens",
"input_cost_per_token_batches",
"input_cost_per_token_cache_hit",
"input_cost_per_token_flex",
"input_cost_per_token_priority",
"input_cost_per_token_ultrafast",
"input_cost_per_video_per_second",
"input_cost_per_video_per_second_above_128k_tokens",
"input_cost_per_video_per_second_above_15s_interval",
"input_cost_per_video_per_second_above_8s_interval",
"input_cost_per_video_token",
"input_cost_per_video_token_batches",
"itpm",
"keepalive_seconds",
"langfuse_environment",
"langfuse_host",
"langfuse_prompt_version",
"langfuse_public_key",
"langfuse_secret",
"langfuse_secret_key",
"langsmith_api_key",
"langsmith_base_url",
"langsmith_project",
"langsmith_sampling_rate",
"langsmith_tenant_id",
"litellm_credential_name",
"litellm_disabled_callbacks",
"litellm_request_debug",
"litellm_session_id",
"litellm_system_prompt",
"litellm_trace_id",
"litellm_trusted_callback_vars",
"logger_fn",
"max_agentic_loops",
"max_budget",
"max_fallbacks",
"max_parallel_requests",
"merge_reasoning_content_in_choices",
"metadata",
"mock_response",
"mock_timeout",
"model_alias_map",
"model_config",
"model_file_id_mapping",
"model_info",
"model_list",
"newrelic_api_key",
"newrelic_region",
"no-log",
"num_retries",
"ocr_cost_per_credit",
"ocr_cost_per_page",
"order",
"otpm",
"output_cost_per_audio_per_second",
"output_cost_per_audio_token",
"output_cost_per_character",
"output_cost_per_character_above_128k_tokens",
"output_cost_per_image",
"output_cost_per_image_token",
"output_cost_per_pixel",
"output_cost_per_reasoning_token",
"output_cost_per_reasoning_token_flex",
"output_cost_per_reasoning_token_priority",
"output_cost_per_second",
"output_cost_per_second_1080p",
"output_cost_per_second_480p",
"output_cost_per_second_4k",
"output_cost_per_second_720p",
"output_cost_per_token",
"output_cost_per_token_above_128k_tokens",
"output_cost_per_token_above_200k_tokens",
"output_cost_per_token_above_200k_tokens_priority",
"output_cost_per_token_above_272k_tokens",
"output_cost_per_token_above_272k_tokens_flex",
"output_cost_per_token_above_272k_tokens_priority",
"output_cost_per_token_above_512k_tokens",
"output_cost_per_token_batches",
"output_cost_per_token_flex",
"output_cost_per_token_priority",
"output_cost_per_token_ultrafast",
"output_cost_per_video_per_second",
"output_cost_per_video_token",
"output_vector_size",
"posthog_api_key",
"posthog_api_url",
"preset_cache_key",
"prompt_environment",
"prompt_id",
"prompt_label",
"prompt_variables",
"prompt_version",
"provider_specific_header",
"quality_router_config",
"quality_router_default_model",
"region_name",
"regional_endpoint_uplift_multiplier",
"regional_processing_uplift_multiplier_eu",
"regional_processing_uplift_multiplier_us",
"retry_policy",
"retry_strategy",
"roles",
"routing_strategy",
"rpm",
"rust",
"s3_bucket_name",
"s3_output_bucket_name",
"s3_region_name",
"search_context_cost_per_query",
"search_tool_name",
"secret_fields",
"self",
"shared_session",
"ssl_verify",
"stream_response",
"stream_timeout",
"supports_system_message",
"tags",
"text_completion",
"tiered_pricing",
"tpm",
"ttl",
"turn_off_message_logging",
"use_chat_completions_api",
"use_client",
"use_in_pass_through",
"use_litellm_proxy",
"use_xai_oauth",
"user_continue_message",
"verbose",
"wandb_api_key",
"weave_project_id",
"weight",
];
pub fn compose_body<B: Serialize>(
arguments: &CallArguments,
body: &B,
consumed: &[&str],
) -> Result<Value, crate::params::Error> {
let Value::Object(fields) =
serde_json::to_value(body).map_err(|_| crate::params::Error::Body)?
else {
return Err(crate::params::Error::Body);
};
let overrides = match arguments.get("extra_body") {
None | Some(Value::Null) => None,
Some(Value::Object(fields)) => Some(fields),
Some(_) => return Err(crate::params::Error::ExtraBody),
};
let extensions = arguments.iter().filter(|(name, _)| {
!consumed.contains(&name.as_str()) && name.as_str() != "extra_body" && !is_control(name)
});
Ok(Value::Object(
fields
.into_iter()
.chain(
extensions
.chain(overrides.into_iter().flatten())
.filter(|(name, _)| {
name.as_str() != "model"
&& name.as_str() != "extra_body"
&& !crate::params::is_control_param(name)
})
.map(|(name, value)| (name.clone(), value.clone())),
)
.collect(),
))
}
impl Deref for CallArguments {
type Target = Map<String, Value>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl From<Map<String, Value>> for CallArguments {
fn from(values: Map<String, Value>) -> Self {
Self(values)
}
}
impl From<CallArguments> for Map<String, Value> {
fn from(arguments: CallArguments) -> Self {
arguments.0
}
}
impl FromIterator<(String, Value)> for CallArguments {
fn from_iter<T: IntoIterator<Item = (String, Value)>>(iter: T) -> Self {
Self(iter.into_iter().collect())
}
}
impl IntoIterator for CallArguments {
type Item = (String, Value);
type IntoIter = serde_json::map::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn composition_preserves_extensions_and_applies_shallow_explicit_overrides() {
let original = json!({
"known": false, "future": {"old": 1}, "null": null, "zero": 0,
"metadata": {"host": true}, "shared_session": "host", "api_key": "secret",
"extra_body": {
"known": null, "future": {"new": [false, 0, null]},
"metadata": {"provider": true}, "model": "ignored", "api_key": "ignored"
}
});
let arguments = serde_json::from_value(original.clone()).unwrap();
let body = compose_body(
&arguments,
&json!({"model":"resolved", "known":false}),
&["known"],
)
.unwrap();
assert_eq!(
body,
json!({
"model":"resolved", "known":null, "future":{"new":[false,0,null]},
"null":null, "zero":0, "metadata":{"provider":true}
})
);
assert_eq!(serde_json::to_value(arguments).unwrap(), original);
}
#[test]
fn projection_prioritizes_consumed_fields_and_keeps_unknown_names() {
let fields = [ArgumentSpec {
name: "id",
secret: false,
}];
assert!(should_project("id", &fields, &[]));
assert!(!should_project("id", &[], &[]));
assert!(should_project("future_option", &[], &[]));
assert!(!should_project("document", &fields, &["document"]));
assert!(!should_project("metadata", &fields, &[]));
assert!(!should_project("callbacks", &fields, &[]));
assert!(!should_project("ocr_cost_per_page", &fields, &[]));
}
#[test]
fn invalid_extra_body_is_rejected_without_coercing_it_to_empty() {
for value in [json!(false), json!(0), json!([]), json!("")] {
let arguments = serde_json::from_value(json!({"extra_body":value})).unwrap();
assert_eq!(
compose_body(&arguments, &json!({}), &[]),
Err(crate::params::Error::ExtraBody)
);
}
let arguments = serde_json::from_value(json!({"extra_body":null})).unwrap();
assert_eq!(
compose_body(&arguments, &json!({}), &[]).unwrap(),
json!({})
);
}
#[test]
fn typed_views_preserve_missing_and_explicit_null_in_the_source() {
#[derive(Deserialize)]
struct Options {
enabled: Option<bool>,
}
let arguments: CallArguments =
serde_json::from_value(json!({"enabled":null,"future":0})).unwrap();
assert!(
parse_options::<Options>(&arguments)
.unwrap()
.enabled
.is_none()
);
assert_eq!(arguments.get("enabled"), Some(&Value::Null));
assert_eq!(arguments.get("missing"), None);
let invalid = serde_json::from_value(json!({"enabled":0})).unwrap();
assert_eq!(
parse_options::<Options>(&invalid).err().unwrap().path,
"enabled"
);
}
}

3
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -228,10 +228,11 @@ fn epoch_seconds() -> f64 {
#[cfg(test)]
mod tests {
use super::*;
use std::pin::Pin;
use std::sync::Mutex;
use super::*;
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
#[derive(Default)]

14
litellm-rust/crates/core/src/chat_completions/common_utils.rs
View file

@ -1,19 +1,17 @@
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use serde_json::{Map, Value};
use super::transformation::ChatCompletionsProviderConfig;
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::llms::anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use crate::llms::base_llm::chat::transformation::BaseConfig;
const HEADER_CONTEXT: &str = "chat completions";
pub(super) fn chat_completions_provider_config(
provider: &str,
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
pub(super) fn chat_completions_provider_config(provider: &str) -> Option<&'static dyn BaseConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
"bedrock" => Some(
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
&crate::llms::bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
),
_ => None,
}

6
litellm-rust/crates/core/src/chat_completions/conversation.rs
View file

@ -10,9 +10,8 @@
//! `_bedrock_converse_messages_pt` for the text-only surface this route
//! accepts; anything richer is declined upstream by the capability gate.
use crate::constants::EMPTY_TEXT_PLACEHOLDER;
use super::types::{ChatMessage, ChatMessageContent};
use crate::constants::EMPTY_TEXT_PLACEHOLDER;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TurnRole {
@ -132,9 +131,10 @@ pub fn build_conversation(messages: &[ChatMessage]) -> Conversation {
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
fn messages(value: serde_json::Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
}

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

@ -1,15 +1,14 @@
use serde_json::Value;
use super::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::prepare::prepare_provider_request;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
ResolvedChatCompletionsRequest,
};
use crate::http_utils::{http_request, truncate_error_body};
use crate::llms::base_llm::chat::transformation::ChatCompletionsAuth;
pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
@ -87,7 +86,7 @@ pub(super) async fn signed_headers(
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::providers::bedrock::aws_base::{
use litellm_auth_aws::{
aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
};

5
litellm-rust/crates/core/src/chat_completions/mod.rs
View file

@ -14,13 +14,12 @@ pub mod conversation;
pub(crate) mod handler;
mod prepare;
pub mod response_utils;
pub mod transformation;
pub mod streaming;
pub mod types;
use serde_json::{Map, Value};
use handler::execute_chat_completions_provider_call;
use prepare::{parse_messages, resolve_provider_config, resolve_request};
use serde_json::{Map, Value};
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
pub async fn chat_completions(

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

@ -1,20 +1,21 @@
use serde_json::Value;
use super::Error;
use crate::http_utils::has_header;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{chat_completions_provider_config, string_headers};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
use super::types::{
ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
ResolvedChatCompletionsRequest,
};
use crate::http_utils::has_header;
use crate::litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, get_custom_llm_provider,
};
use crate::llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
pub(super) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> {
) -> Result<(String, &'static dyn BaseConfig), Error> {
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
.or_else(|| {
custom_llm_provider.map(|provider| CustomLlmProvider {
@ -65,7 +66,7 @@ pub(super) fn resolve_request(
fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn ChatCompletionsProviderConfig,
config: &dyn BaseConfig,
) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> {
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers(request.extra_headers.clone())?;
@ -122,7 +123,7 @@ pub(super) fn prepare_provider_request(
let model = request.model;
let config = request.config;
let env_lookup = |key: &str| std::env::var(key).ok();
let url = config.complete_url(
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,

9
litellm-rust/crates/core/src/chat_completions/streaming.rs Normal file
View file

@ -0,0 +1,9 @@
pub trait StreamTransformer {
type Input;
type Output;
type Error;
fn transform(&mut self, input: Self::Input) -> Result<Vec<Self::Output>, Self::Error>;
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error>;
}

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

@ -1,10 +1,9 @@
use serde_json::{Map, Value, json};
use super::Error;
use super::prepare::{prepare_provider_request, resolve_request};
use super::transformation::ChatCompletionsAuth;
use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
use crate::llms::base_llm::chat::transformation::ChatCompletionsAuth;
fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
@ -588,10 +587,10 @@ fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
}
mod round_trip {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::*;
use crate::chat_completions::chat_completions;
async fn read_http_request(socket: &mut TcpStream) -> String {

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

@ -3,7 +3,7 @@ use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
use crate::llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
/// A `/chat/completions` call as it crosses into the core.
///
@ -24,7 +24,7 @@ pub struct ChatCompletionsRequest<'a> {
pub(super) struct ResolvedChatCompletionsRequest<'a> {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) config: &'static dyn BaseConfig,
pub(super) messages: Vec<ChatMessage>,
pub(super) optional_params: Map<String, Value>,
pub(super) api_key: Option<&'a str>,
@ -35,7 +35,7 @@ pub(super) struct ResolvedChatCompletionsRequest<'a> {
pub(super) struct ProviderChatCompletionsRequest {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) config: &'static dyn BaseConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
@ -120,3 +120,83 @@ pub struct ChatCompletionsResponse {
pub choices: Vec<ChatCompletionsChoice>,
pub usage: ChatCompletionsUsage,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionToolCallFunctionChunk {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
pub arguments: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionToolCallChunk {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(rename = "type")]
pub tool_type: String,
pub function: ChatCompletionToolCallFunctionChunk,
pub index: i64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatCompletionThinkingBlock {
Thinking {
#[serde(default, skip_serializing_if = "Option::is_none")]
thinking: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<Value>,
},
RedactedThinking {
#[serde(default, skip_serializing_if = "Option::is_none")]
data: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<Value>,
},
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionDelta {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ChatCompletionToolCallChunk>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thinking_blocks: Option<Vec<ChatCompletionThinkingBlock>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionStreamingChoice {
pub index: u64,
pub delta: ChatCompletionDelta,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub finish_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub logprobs: Option<Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionChunk {
pub id: String,
pub created: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
pub object: String,
pub choices: Vec<ChatCompletionStreamingChoice>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<ChatCompletionsUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
}

3
litellm-rust/crates/core/src/http_utils.rs
View file

@ -131,9 +131,10 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
#[rstest::rstest]
#[case(HeaderPolicy::All, true, true)]
#[case(HeaderPolicy::Only(&["authorization"]), true, false)]

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

@ -1,14 +1,18 @@
pub mod audio_transcription;
pub mod call_arguments;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
pub mod litellm_core_utils;
pub mod llms;
mod media;
pub mod messages;
pub mod ocr;
pub mod providers;
pub mod params;
pub mod responses;
mod serde_compat;
pub mod transport;
mod url_utils;

0
litellm-rust/crates/core/src/providers/custom_llm_provider.rs → litellm-rust/crates/core/src/litellm_core_utils/get_llm_provider_logic.rs
View file

1
litellm-rust/crates/core/src/litellm_core_utils/mod.rs Normal file
View file

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

1
litellm-rust/crates/core/src/providers/openai/responses/mod.rs → litellm-rust/crates/core/src/llms/anthropic/chat/mod.rs
View file

@ -1 +1,2 @@
pub mod streaming;
pub mod transformation;

164
litellm-rust/crates/core/src/llms/anthropic/chat/streaming.rs Normal file
View file

@ -0,0 +1,164 @@
use std::collections::HashMap;
use serde_json::Value;
use crate::chat_completions::Error;
use crate::chat_completions::streaming::StreamTransformer;
use crate::chat_completions::types::{
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionsUsage,
};
use crate::llms::anthropic::experimental_pass_through::messages::streaming::{
AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
AnthropicStreamUsage,
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AnthropicJsonChunkType {
ValidJson,
AccumulatedJson,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AnthropicContentBlockType {
Text,
ToolUse,
ServerToolUse,
Thinking,
RedactedThinking,
Compaction,
ToolResult(String),
Other(String),
}
#[derive(Clone, Debug, PartialEq)]
pub struct AnthropicContentBlockDeltaEvent {
pub index: u64,
pub delta: AnthropicContentBlockDelta,
}
pub struct AnthropicChatCompletionsStreamTransformer {
pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
pub tool_index: i64,
pub json_mode: bool,
pub speed: Option<String>,
pub tool_name_reverse_map: HashMap<String, String>,
pub response_id: String,
pub served_model: Option<String>,
pub is_response_format_tool: bool,
pub converted_response_format_tool: bool,
pub accumulated_json: String,
pub chunk_type: AnthropicJsonChunkType,
pub current_content_block_type: Option<AnthropicContentBlockType>,
pub web_search_results: Vec<Value>,
pub web_search_calls: HashMap<String, Value>,
pub compaction_blocks: Vec<Value>,
pub reasoning_content_chunks: Vec<String>,
pub server_tool_inputs: HashMap<String, Value>,
pub tool_results: Vec<Value>,
pub current_server_tool_id: Option<String>,
pub container_id: Option<String>,
}
impl AnthropicChatCompletionsStreamTransformer {
pub fn new(
_json_mode: bool,
_speed: Option<String>,
_tool_name_reverse_map: HashMap<String, String>,
) -> Self {
todo!()
}
pub fn check_empty_tool_call_args(&self) -> bool {
todo!()
}
pub fn handle_usage(&mut self, _usage: AnthropicStreamUsage) -> ChatCompletionsUsage {
todo!()
}
pub fn handle_content_block_delta(
&mut self,
_index: u64,
_delta: AnthropicContentBlockDelta,
) -> (
String,
Option<ChatCompletionToolCallChunk>,
Vec<ChatCompletionThinkingBlock>,
Option<Value>,
Option<String>,
) {
todo!()
}
pub fn handle_content_block_start(
&mut self,
_index: u64,
_content_block: AnthropicContentBlock,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
pub fn handle_json_mode_chunk(
&mut self,
_text: String,
_tool_use: Option<ChatCompletionToolCallChunk>,
) -> (String, Option<ChatCompletionToolCallChunk>) {
todo!()
}
pub fn handle_accumulated_json_chunk(
&mut self,
_data: &str,
_is_final: bool,
) -> Result<Option<ChatCompletionChunk>, Error> {
todo!()
}
pub fn handle_redacted_thinking_content(
&mut self,
_content_block: &AnthropicContentBlock,
) -> Vec<ChatCompletionThinkingBlock> {
todo!()
}
pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
todo!()
}
pub fn complete_web_search_call(&mut self, _result: Value) {
todo!()
}
pub fn build_code_interpreter_results(&self) -> Vec<Value> {
todo!()
}
pub fn handle_message_delta(
&mut self,
_event: AnthropicMessagesStreamEvent,
) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
todo!()
}
pub fn chunk_parser(
&mut self,
_event: AnthropicMessagesStreamEvent,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
}
impl StreamTransformer for AnthropicChatCompletionsStreamTransformer {
type Input = AnthropicMessagesStreamEvent;
type Output = ChatCompletionChunk;
type Error = Error;
fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
}

5
litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs → litellm-rust/crates/core/src/llms/anthropic/chat/tests.rs
View file

@ -1,6 +1,7 @@
use serde_json::json;
use super::*;
use crate::chat_completions::Error;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
@ -419,7 +420,7 @@ fn resolves_the_messages_url_and_x_api_key_auth() {
let config = &ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
assert_eq!(
config
.complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None)
.get_complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None)
.expect("url builds"),
"https://api.anthropic.com/v1/messages"
);

197
litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs → litellm-rust/crates/core/src/llms/anthropic/chat/transformation.rs
View file

@ -2,20 +2,18 @@ use serde_json::{Map, Value, json};
use crate::chat_completions::Error;
use crate::chat_completions::conversation::{Conversation, build_conversation};
use crate::chat_completions::transformation::{
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
unsupported_param,
};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::chat_completions::types::{
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage,
ProviderChatRequestData, ProviderChatResponseData,
};
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
use crate::providers::anthropic::messages::transformation::{
use crate::llms::anthropic::experimental_pass_through::messages::transformation::{
complete_anthropic_url, resolve_anthropic_api_key,
};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::llms::base_llm::chat::transformation::{
BaseConfig, ChatCompletionsAuth, Unsupported, unsupported_message, unsupported_param,
};
/// Anthropic parameter names, post `map_openai_params`, that the Rust path can
/// place verbatim in the Messages body.
@ -34,46 +32,16 @@ const SUPPORTED_PARAMS: &[(&str, &str)] = &[
("stop", "stop_sequences"),
];
pub struct AnthropicChatCompletionsConfig;
pub struct AnthropicConfig;
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig =
AnthropicChatCompletionsConfig;
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicConfig = AnthropicConfig;
fn text_block(text: &str) -> Value {
json!({"type": "text", "text": text})
}
impl BaseConfig for AnthropicConfig {
fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] {
SUPPORTED_PARAMS
}
fn anthropic_body(model: &str, conversation: &Conversation, params: Map<String, Value>) -> Value {
let messages: Vec<Value> = conversation
.turns
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
})
})
.collect();
let system: Vec<Value> = conversation.system.iter().map(|s| text_block(s)).collect();
let body = Map::from_iter(
[
("model".to_string(), json!(model)),
("messages".to_string(), json!(messages)),
]
.into_iter()
// Python builds `{"model", "messages", **optional_params}` with
// `system` already folded into optional_params, so a caller-supplied
// key of the same name wins here too.
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system))))
.chain(params),
);
Value::Object(body)
}
impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
fn complete_url(
fn get_complete_url(
&self,
api_base: Option<&str>,
_model: &str,
@ -83,60 +51,6 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
Ok(complete_anthropic_url(api_base, env_lookup))
}
fn auth(
&self,
api_key: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
Ok(ChatCompletionsAuth::Header {
name: "x-api-key",
value: resolve_anthropic_api_key(api_key, env_lookup)?,
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[
("anthropic-version", "2023-06-01"),
("content-type", "application/json"),
]
}
/// An OAuth bearer is the whole credential: Python's `validate_environment`
/// authenticates with it and drops `x-api-key` rather than resolving one, so
/// the resolved key must not be applied over the top. Any other forwarded
/// `authorization` is unrelated to this header and does not defer, which is
/// also what Python does: it sends the deployment's `x-api-key` alongside.
fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("authorization")
&& value
.strip_prefix("Bearer ")
.is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX))
})
}
fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] {
SUPPORTED_PARAMS
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(self.supported_openai_params(), &[], optional_params)
.or_else(|| messages.iter().find_map(unsupported_message))
// Anthropic rejects a request whose first turn is not a user turn.
// Python only repairs that under `litellm.modify_params`, which the
// core cannot observe, so decline instead of guessing.
.or_else(|| {
(!build_conversation(messages).opens_on_user_turn())
.then_some(Unsupported("conversation does not open on a user turn"))
})
}
fn transform_request(
&self,
model: &str,
@ -210,6 +124,93 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
),
})
}
fn auth(
&self,
api_key: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
Ok(ChatCompletionsAuth::Header {
name: "x-api-key",
value: resolve_anthropic_api_key(api_key, env_lookup)?,
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[
("anthropic-version", "2023-06-01"),
("content-type", "application/json"),
]
}
/// An OAuth bearer is the whole credential: Python's `validate_environment`
/// authenticates with it and drops `x-api-key` rather than resolving one, so
/// the resolved key must not be applied over the top. Any other forwarded
/// `authorization` is unrelated to this header and does not defer, which is
/// also what Python does: it sends the deployment's `x-api-key` alongside.
fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("authorization")
&& value
.strip_prefix("Bearer ")
.is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX))
})
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(self.supported_openai_param_mappings(), &[], optional_params)
.or_else(|| messages.iter().find_map(unsupported_message))
// Anthropic rejects a request whose first turn is not a user turn.
// Python only repairs that under `litellm.modify_params`, which the
// core cannot observe, so decline instead of guessing.
.or_else(|| {
(!build_conversation(messages).opens_on_user_turn())
.then_some(Unsupported("conversation does not open on a user turn"))
})
}
}
fn text_block(text: &str) -> Value {
json!({"type": "text", "text": text})
}
fn anthropic_body(
model: &str,
conversation: &Conversation,
optional_params: Map<String, Value>,
) -> Value {
let messages: Vec<Value> = conversation
.turns
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
})
})
.collect();
let system: Vec<Value> = conversation.system.iter().map(|s| text_block(s)).collect();
let body = Map::from_iter(
[
("model".to_string(), json!(model)),
("messages".to_string(), json!(messages)),
]
.into_iter()
// Python builds `{"model", "messages", **optional_params}` with
// `system` already folded into optional_params, so a caller-supplied
// key of the same name wins here too.
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system))))
.chain(optional_params),
);
Value::Object(body)
}
#[cfg(test)]

338
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/batches.rs Normal file
View file

@ -0,0 +1,338 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use url::Url;
use crate::llms::anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base;
use crate::messages::Error;
use crate::messages::types::AnthropicMessagesResponse;
const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches";
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicBatchRequestCounts {
#[serde(default)]
pub processing: u64,
#[serde(default)]
pub succeeded: u64,
#[serde(default)]
pub errored: u64,
#[serde(default)]
pub canceled: u64,
#[serde(default)]
pub expired: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicMessageBatch {
#[serde(default)]
pub id: String,
#[serde(default = "default_processing_status")]
pub processing_status: String,
pub created_at: Option<String>,
pub ended_at: Option<String>,
pub expires_at: Option<String>,
pub cancel_initiated_at: Option<String>,
pub archived_at: Option<String>,
#[serde(default)]
pub request_counts: AnthropicBatchRequestCounts,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BatchStatus {
InProgress,
Cancelling,
Completed,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BatchRequestCounts {
pub total: u64,
pub completed: u64,
pub failed: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiteLlmMessageBatch {
pub id: String,
pub object: String,
pub endpoint: String,
pub input_file_id: String,
pub completion_window: String,
pub status: BatchStatus,
pub output_file_id: String,
pub created_at: i64,
pub in_progress_at: Option<i64>,
pub expires_at: Option<i64>,
pub completed_at: Option<i64>,
pub expired_at: Option<i64>,
pub cancelling_at: Option<i64>,
pub cancelled_at: Option<i64>,
pub request_counts: BatchRequestCounts,
}
pub trait AnthropicBatchesConfig {
fn create_batch_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_create_batch_request(&self) -> Result<Value, Error>;
fn transform_create_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> Result<LiteLlmMessageBatch, Error>;
fn retrieve_batch_url(
&self,
api_base: Option<&str>,
batch_id: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_retrieve_batch_request(&self) -> Value;
fn transform_retrieve_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> LiteLlmMessageBatch;
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error>;
}
pub struct AnthropicBatchesTransformation;
pub const ANTHROPIC_BATCHES_TRANSFORMATION: AnthropicBatchesTransformation =
AnthropicBatchesTransformation;
fn default_processing_status() -> String {
"in_progress".into()
}
fn timestamp(value: Option<&str>) -> Option<i64> {
value
.and_then(|value| {
OffsetDateTime::parse(value, &time::format_description::well_known::Rfc3339).ok()
})
.map(OffsetDateTime::unix_timestamp)
}
fn batches_base_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Url, Error> {
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
let api_base = api_base.trim_end_matches('/');
let complete_url = if api_base.ends_with(BATCHES_PATH_SUFFIX) {
api_base.to_string()
} else if let Some(base) = api_base.strip_suffix("/v1/messages") {
format!("{base}{BATCHES_PATH_SUFFIX}")
} else {
format!("{api_base}{BATCHES_PATH_SUFFIX}")
};
Url::parse(&complete_url)
.map_err(|error| Error::InvalidRequest(format!("invalid Anthropic API base: {error}")))
}
impl AnthropicBatchesConfig for AnthropicBatchesTransformation {
fn create_batch_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(batches_base_url(api_base, env_lookup)?.into())
}
fn transform_create_batch_request(&self) -> Result<Value, Error> {
Err(Error::Unsupported("Anthropic message batch creation"))
}
fn transform_create_batch_response(
&self,
_response: AnthropicMessageBatch,
_now: i64,
) -> Result<LiteLlmMessageBatch, Error> {
Err(Error::Unsupported("Anthropic message batch creation"))
}
fn retrieve_batch_url(
&self,
api_base: Option<&str>,
batch_id: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
if batch_id.is_empty() {
return Err(Error::MissingField("batch_id"));
}
let mut url = batches_base_url(api_base, env_lookup)?;
url.path_segments_mut()
.map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))?
.push(batch_id);
Ok(url.into())
}
fn transform_retrieve_batch_request(&self) -> Value {
Value::Object(Default::default())
}
fn transform_retrieve_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> LiteLlmMessageBatch {
let created_at = timestamp(response.created_at.as_deref());
let ended_at = timestamp(response.ended_at.as_deref());
let expires_at = timestamp(response.expires_at.as_deref());
let cancel_initiated_at = timestamp(response.cancel_initiated_at.as_deref());
let archived_at = timestamp(response.archived_at.as_deref());
let status = match response.processing_status.as_str() {
"canceling" => BatchStatus::Cancelling,
"ended" => BatchStatus::Completed,
_ => BatchStatus::InProgress,
};
let request_counts = BatchRequestCounts {
total: response.request_counts.processing
+ response.request_counts.succeeded
+ response.request_counts.errored
+ response.request_counts.canceled
+ response.request_counts.expired,
completed: response.request_counts.succeeded,
failed: response.request_counts.errored,
};
LiteLlmMessageBatch {
id: response.id.clone(),
object: "batch".into(),
endpoint: "/v1/messages".into(),
input_file_id: "None".into(),
completion_window: "24h".into(),
status,
output_file_id: response.id,
created_at: created_at.unwrap_or(now),
in_progress_at: (response.processing_status == "in_progress")
.then_some(created_at)
.flatten(),
expires_at,
completed_at: (response.processing_status == "ended")
.then_some(ended_at)
.flatten(),
expired_at: archived_at,
cancelling_at: (response.processing_status == "canceling")
.then_some(cancel_initiated_at)
.flatten(),
cancelled_at: (response.processing_status == "canceling")
.then_some(ended_at)
.flatten(),
request_counts,
}
}
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error> {
body.lines()
.filter(|line| !line.trim().is_empty())
.filter_map(|line| serde_json::from_str::<Value>(line.trim()).ok())
.map(|record| {
serde_json::from_value(record["result"]["message"].clone()).map_err(|error| {
Error::InvalidResponse(format!("invalid Anthropic batch result: {error}"))
})
})
.collect()
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn builds_and_encodes_message_batch_urls() {
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.create_batch_url(None, &|_| None)
.unwrap(),
"https://api.anthropic.com/v1/messages/batches"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.create_batch_url(Some("https://proxy.test/v1/messages/batches"), &|_| None)
.unwrap(),
"https://proxy.test/v1/messages/batches"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.retrieve_batch_url(Some("https://proxy.test"), "batch/id ?", &|_| None)
.unwrap(),
"https://proxy.test/v1/messages/batches/batch%2Fid%20%3F"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_request(),
json!({})
);
}
#[test]
fn maps_retrieved_batch_status_counts_and_timestamps_like_python() {
let response: AnthropicMessageBatch = serde_json::from_value(json!({
"id": "msgbatch_1",
"processing_status": "ended",
"created_at": "2025-01-01T00:00:00Z",
"ended_at": "2025-01-01T00:01:00Z",
"expires_at": "not-a-timestamp",
"request_counts": {
"processing": 1,
"succeeded": 2,
"errored": 3,
"canceled": 4,
"expired": 5
}
}))
.unwrap();
let batch = ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_response(response, 7);
assert_eq!(batch.status, BatchStatus::Completed);
assert_eq!(batch.created_at, 1_735_689_600);
assert_eq!(batch.completed_at, Some(1_735_689_660));
assert_eq!(batch.expires_at, None);
assert_eq!(
batch.request_counts,
BatchRequestCounts {
total: 15,
completed: 2,
failed: 3
}
);
}
#[test]
fn extracts_message_responses_from_ndjson_and_skips_non_json_lines() {
let body = r#"not-json
{"result":{"message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":"end_turn","stop_sequence":null}}}
"#;
let messages = ANTHROPIC_BATCHES_TRANSFORMATION
.transform_batch_results(body)
.unwrap();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].id, "msg_1");
}
#[test]
fn preserves_python_placeholder_for_batch_creation() {
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_request(),
Err(Error::Unsupported("Anthropic message batch creation"))
));
let response: AnthropicMessageBatch = serde_json::from_value(json!({})).unwrap();
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_response(response, 0),
Err(Error::Unsupported("Anthropic message batch creation"))
));
}
}

168
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/count_tokens.rs Normal file
View file

@ -0,0 +1,168 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
use crate::messages::Error;
use crate::messages::types::{AnthropicMessage, SystemPrompt};
const COUNT_TOKENS_ENDPOINT: &str = "https://api.anthropic.com/v1/messages/count_tokens";
const TOKEN_COUNTING_BETA: &str = "token-counting-2024-11-01";
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicCountTokensRequest {
pub model: String,
pub messages: Vec<AnthropicMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<SystemPrompt>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicCountTokensResponse {
pub input_tokens: u64,
}
pub trait AnthropicCountTokensConfig {
fn endpoint(&self) -> &'static str;
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error>;
fn transform_request(
&self,
model: &str,
messages: Vec<AnthropicMessage>,
tools: Option<Vec<Value>>,
system: Option<SystemPrompt>,
) -> Result<AnthropicCountTokensRequest, Error>;
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)>;
}
pub struct AnthropicCountTokensTransformation;
pub const ANTHROPIC_COUNT_TOKENS_TRANSFORMATION: AnthropicCountTokensTransformation =
AnthropicCountTokensTransformation;
impl AnthropicCountTokensConfig for AnthropicCountTokensTransformation {
fn endpoint(&self) -> &'static str {
COUNT_TOKENS_ENDPOINT
}
fn transform_request(
&self,
model: &str,
messages: Vec<AnthropicMessage>,
tools: Option<Vec<Value>>,
system: Option<SystemPrompt>,
) -> Result<AnthropicCountTokensRequest, Error> {
self.validate_request(model, &messages)?;
Ok(AnthropicCountTokensRequest {
model: model.to_string(),
messages,
tools,
system,
})
}
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error> {
if model.is_empty() {
return Err(Error::MissingField("model"));
}
if messages.is_empty() {
return Err(Error::MissingField("messages"));
}
Ok(())
}
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)> {
let auth = if api_key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) {
("authorization", format!("Bearer {api_key}"))
} else {
("x-api-key", api_key.to_string())
};
vec![
("content-type", "application/json".to_string()),
auth,
("anthropic-version", "2023-06-01".to_string()),
("anthropic-beta", TOKEN_COUNTING_BETA.to_string()),
]
}
}
#[cfg(test)]
mod tests {
use serde_json::{Map, json};
use super::*;
use crate::messages::types::MessageContent;
fn message() -> AnthropicMessage {
AnthropicMessage {
role: "user".into(),
content: MessageContent::Text("hello".into()),
extra: Map::new(),
}
}
#[test]
fn maps_the_python_count_tokens_contract() {
let request = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION
.transform_request(
"claude-test",
vec![message()],
Some(vec![json!({"name": "lookup"})]),
Some(SystemPrompt::Text("system".into())),
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap(),
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hello"}],
"tools": [{"name": "lookup"}],
"system": "system"
})
);
assert_eq!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.endpoint(),
COUNT_TOKENS_ENDPOINT
);
}
#[test]
fn rejects_the_invalid_requests_python_rejects() {
assert!(matches!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
"",
vec![message()],
None,
None
),
Err(Error::MissingField("model"))
));
assert!(matches!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
"claude-test",
vec![],
None,
None
),
Err(Error::MissingField("messages"))
));
}
#[test]
fn uses_api_key_or_oauth_headers_without_combining_credentials() {
let api_key = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-api");
assert!(api_key.contains(&("x-api-key", "sk-ant-api".into())));
assert!(!api_key.iter().any(|(name, _)| *name == "authorization"));
let oauth = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-oat-test");
assert!(oauth.contains(&("authorization", "Bearer sk-ant-oat-test".into())));
assert!(!oauth.iter().any(|(name, _)| *name == "x-api-key"));
assert!(oauth.contains(&("anthropic-beta", TOKEN_COUNTING_BETA.into())));
}
}

4
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub mod batches;
pub mod count_tokens;
pub mod streaming;
pub mod transformation;

282
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/streaming.rs Normal file
View file

@ -0,0 +1,282 @@
use base64::Engine;
use bytes::Buf;
use futures_util::{Stream, StreamExt};
use litellm_framing::Framer;
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
use litellm_framing::sse::{SseFrame, SseFramer};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::messages::Error;
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamUsage {
#[serde(default)]
pub input_tokens: u64,
#[serde(default)]
pub output_tokens: u64,
#[serde(default)]
pub cache_creation_input_tokens: u64,
#[serde(default)]
pub cache_read_input_tokens: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub server_tool_use: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamMessage {
pub id: String,
#[serde(rename = "type")]
pub message_type: String,
pub role: String,
pub model: String,
pub content: Vec<Value>,
pub stop_reason: Option<String>,
pub stop_sequence: Option<String>,
pub usage: AnthropicStreamUsage,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AnthropicContentBlockDelta {
TextDelta {
text: String,
},
InputJsonDelta {
partial_json: String,
},
#[serde(rename = "citations_delta")]
Citations {
citation: Value,
},
ThinkingDelta {
thinking: String,
},
SignatureDelta {
signature: String,
},
CompactionDelta {
content: String,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicContentBlock {
#[serde(rename = "type")]
pub block_type: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thinking: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub data: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub caller: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AnthropicMessageDelta {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_sequence: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_details: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub container: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamError {
#[serde(rename = "type")]
pub error_type: String,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub details: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AnthropicMessagesStreamEvent {
MessageStart {
message: AnthropicStreamMessage,
},
ContentBlockStart {
index: u64,
content_block: AnthropicContentBlock,
},
ContentBlockDelta {
index: u64,
delta: AnthropicContentBlockDelta,
},
ContentBlockStop {
index: u64,
},
MessageDelta {
delta: AnthropicMessageDelta,
#[serde(default, skip_serializing_if = "Option::is_none")]
usage: Option<AnthropicStreamUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
context_management: Option<Value>,
},
MessageStop,
Ping,
Error {
error: AnthropicStreamError,
},
}
#[derive(Deserialize)]
struct BedrockChunkPayload {
bytes: String,
}
pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result<AnthropicMessagesStreamEvent, Error> {
let data = frame.data.ok_or(Error::MissingStreamData)?;
serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn decode_bedrock_anthropic_frame(
frame: AwsEventStreamFrame,
) -> Result<AnthropicMessagesStreamEvent, Error> {
let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload)
.map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?;
let event = base64::engine::general_purpose::STANDARD
.decode(payload.bytes)
.map_err(|error| Error::InvalidBedrockBase64(error.to_string()))?;
serde_json::from_slice(&event).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn direct_anthropic_event_stream<S, B, E>(
input: S,
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
SseFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_anthropic_sse_frame(frame)
})
}
pub fn bedrock_anthropic_event_stream<S, B, E>(
input: S,
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
AwsEventStreamFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_bedrock_anthropic_frame(frame)
})
}
#[cfg(test)]
mod tests {
use std::io;
use aws_smithy_eventstream::frame::write_message_to;
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
use base64::engine::general_purpose::STANDARD;
use bytes::Bytes;
use futures_util::TryStreamExt;
use super::*;
const TEXT_DELTA: &str =
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
#[tokio::test]
async fn direct_anthropic_sse_frames_into_typed_events() {
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
let events = direct_anthropic_event_stream(futures_util::stream::iter(
wire.as_bytes().chunks(3).map(Ok::<_, io::Error>),
))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: AnthropicContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
#[test]
fn decodes_citations_delta_events() {
let event = decode_anthropic_sse_frame(SseFrame {
event: Some("content_block_delta".into()),
data: Some(
r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#
.into(),
),
id: None,
retry: None,
})
.unwrap();
assert!(matches!(
event,
AnthropicMessagesStreamEvent::ContentBlockDelta {
delta: AnthropicContentBlockDelta::Citations { .. },
..
}
));
}
#[tokio::test]
async fn bedrock_aws_frames_into_the_same_typed_events() {
let payload = serde_json::json!({"bytes": STANDARD.encode(TEXT_DELTA)});
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(
Header::new(":event-type", HeaderValue::String("chunk".into())),
);
let mut wire = Vec::new();
write_message_to(&message, &mut wire).unwrap();
let events = bedrock_anthropic_event_stream(futures_util::stream::iter(
wire.chunks(3).map(Ok::<_, io::Error>),
))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: AnthropicContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
}

55
litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs → litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/transformation.rs
View file

@ -1,5 +1,5 @@
use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
use crate::messages::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE";
@ -10,6 +10,25 @@ pub struct AnthropicMessagesConfig;
pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig;
impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig {
fn get_complete_url(
&self,
api_base: Option<&str>,
_model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(complete_anthropic_url(api_base, env_lookup))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from)
}
}
pub fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
@ -31,10 +50,7 @@ pub fn complete_anthropic_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
let api_base = non_empty(api_base)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string());
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
let api_base = api_base.trim_end_matches('/');
if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
@ -43,27 +59,14 @@ pub fn complete_anthropic_url(
format!("{api_base}{MESSAGES_PATH_SUFFIX}")
}
impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(complete_anthropic_url(api_base, env_lookup))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
MessagesAuthStrategy::Header("x-api-key")
}
pub fn resolve_anthropic_api_base(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
non_empty(api_base)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string())
}
#[cfg(test)]

0
litellm-rust/crates/core/src/providers/azure_ai/mod.rs → litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/mod.rs
View file

2
litellm-rust/crates/core/src/llms/anthropic/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod chat;
pub mod experimental_pass_through;

141
litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs → litellm-rust/crates/core/src/llms/azure_ai/anthropic/messages_transformation.rs
View file

@ -1,13 +1,16 @@
use serde_json::{Map, Value};
use crate::llms::anthropic::experimental_pass_through::messages::transformation::{
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
};
use crate::llms::base_llm::anthropic_messages::transformation::{
BaseAnthropicMessagesConfig, MessagesAuthStrategy,
};
use crate::messages::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::messages::types::{
AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
MessageContent, SystemPrompt,
};
use crate::providers::anthropic::messages::transformation::{
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
};
use serde_json::{Map, Value};
const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY";
const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE";
@ -25,6 +28,61 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig =
anthropic: ANTHROPIC_MESSAGES_CONFIG,
};
impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig {
fn get_complete_url(
&self,
api_base: Option<&str>,
_model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
complete_azure_anthropic_url(api_base, env_lookup)
}
fn transform_anthropic_messages_request(
&self,
request: AnthropicMessagesRequest,
) -> Result<AnthropicMessagesRequest, Error> {
let mut request = fold_system_role_messages(request);
if let Some(system) = request.system.as_mut() {
strip_scope_from_system(system);
}
request
.messages
.iter_mut()
.for_each(strip_scope_from_message);
self.anthropic.transform_anthropic_messages_request(request)
}
fn transform_anthropic_messages_response(
&self,
model: &str,
response: AnthropicMessagesResponse,
) -> Result<AnthropicMessagesResponse, Error> {
self.anthropic
.transform_anthropic_messages_response(model, response)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_azure_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
self.anthropic.auth_strategy()
}
fn accepts_bearer_auth(&self) -> bool {
true
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
self.anthropic.default_headers()
}
}
pub fn resolve_azure_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
@ -135,65 +193,12 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess
}
}
impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
complete_azure_anthropic_url(api_base, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_azure_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
self.anthropic.auth_strategy()
}
fn accepts_bearer_auth(&self) -> bool {
true
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
self.anthropic.default_headers()
}
fn transform_request(
&self,
request: AnthropicMessagesRequest,
) -> Result<AnthropicMessagesRequest, Error> {
let mut request = fold_system_role_messages(request);
if let Some(system) = request.system.as_mut() {
strip_scope_from_system(system);
}
request
.messages
.iter_mut()
.for_each(strip_scope_from_message);
self.anthropic.transform_request(request)
}
fn transform_response(
&self,
model: &str,
response: AnthropicMessagesResponse,
) -> Result<AnthropicMessagesResponse, Error> {
self.anthropic.transform_response(model, response)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest {
serde_json::from_value(value).expect("valid request")
}
@ -337,7 +342,7 @@ mod tests {
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms"),
);
@ -364,10 +369,10 @@ mod tests {
"messages": [{"role": "user", "content": "hi"}]
}));
let once = AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms");
let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(once.clone())
.transform_anthropic_messages_request(once.clone())
.expect("request transforms");
assert_eq!(once, twice);
assert_eq!(to_value(once)["system"], json!("plain string system"));
@ -401,7 +406,7 @@ mod tests {
});
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request_from(body.clone()))
.transform_anthropic_messages_request(request_from(body.clone()))
.expect("request transforms"),
);
assert_eq!(transformed, body);
@ -421,7 +426,7 @@ mod tests {
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms"),
);
@ -451,7 +456,7 @@ mod tests {
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms"),
);
@ -478,7 +483,7 @@ mod tests {
});
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request_from(body.clone()))
.transform_anthropic_messages_request(request_from(body.clone()))
.expect("request transforms"),
);
assert_eq!(transformed, body);
@ -505,7 +510,7 @@ mod tests {
}))
.expect("valid response");
let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_response("claude-sonnet-4-5", response)
.transform_anthropic_messages_response("claude-sonnet-4-5", response)
.expect("response transforms");
let value = serde_json::to_value(transformed).expect("serializable");
assert_eq!(value["stop_reason"], json!("end_turn"));

1
litellm-rust/crates/core/src/llms/azure_ai/anthropic/mod.rs Normal file
View file

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

2
litellm-rust/crates/core/src/llms/azure_ai/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod anthropic;
pub(crate) mod ocr;

166
litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs Normal file
View file

@ -0,0 +1,166 @@
use serde_json::Value;
use crate::call_arguments::CallArguments;
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
use crate::llms::cohere::ocr::transformation::{CohereParseConfig, CohereRequest};
use crate::llms::cohere::ocr::{CohereOptions, validate_document};
use crate::ocr::OcrClient;
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest};
use crate::url_utils::ApiUrl;
#[derive(Default)]
pub(crate) struct AzureAICohereParseConfig;
impl BaseOcrConfig for AzureAICohereParseConfig {
type OcrParams = CohereOptions;
type ProviderRequest = CohereRequest;
type Environment = Vec<(String, String)>;
fn get_api_key_env_var(&self) -> Option<&'static str> {
super::transformation::AzureAiOcrConfig.get_api_key_env_var()
}
fn get_health_check_document(&self) -> OcrDocument {
CohereParseConfig.get_health_check_document()
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
BaseOcrConfig::validate_environment(
&super::transformation::AzureAiOcrConfig,
request,
client,
)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
let base = super::transformation::AzureAiOcrConfig::resolve_api_base(
request.connection.api_base.as_deref(),
&crate::ocr::prepare::credential_env,
)?;
self.get_complete_url(&base)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
params: &CohereOptions,
headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
CohereParseConfig.transform_ocr_request(model, document, params, headers)
}
fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] {
CohereParseConfig.get_supported_ocr_params(model)
}
fn map_ocr_params(
&self,
arguments: &CallArguments,
model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
CohereParseConfig.map_ocr_params(arguments, model)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &CohereOptions,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<CohereRequest, crate::ocr::Error> {
validate_document(&document)?;
let document = inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
CohereParseConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
let document = crate::ocr::prepare::body_document(body)?;
validate_document(&document)?;
validate_inline_document(&document)
}
}
impl AzureAICohereParseConfig {
fn get_complete_url(&self, base: &str) -> Result<String, crate::ocr::Error> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
let path = url.path().trim_end_matches('/').to_string();
if path.ends_with("/v2/parse") {
url.set_path(&path);
return Ok(url.into());
}
url.set_path(path.strip_suffix("/models").unwrap_or(&path));
ApiUrl::parse(url.as_str())
.and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
}
}
fn invalid_api_base() -> crate::ocr::Error {
crate::ocr::Error::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in [
"",
"/models",
"/providers/cohere/v2",
"/providers/cohere/v2/parse",
] {
assert_eq!(
AzureAICohereParseConfig
.get_complete_url(&format!("https://example.com{suffix}?tenant=a"))
.unwrap(),
"https://example.com/providers/cohere/v2/parse?tenant=a"
);
}
assert_eq!(
AzureAICohereParseConfig
.get_complete_url("https://example.com/v2/parse?tenant=a")
.unwrap(),
"https://example.com/v2/parse?tenant=a"
);
assert!(
AzureAICohereParseConfig
.get_complete_url("relative/path")
.is_err()
);
}
}

25
litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs → litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs
View file

@ -1,25 +1,14 @@
mod cohere;
mod document_intelligence;
mod mistral;
use std::sync::OnceLock;
use crate::ocr::Error;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
pub(crate) use cohere::AzureCohereAdapter;
pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
pub(crate) use mistral::AzureMistralAdapter;
pub(super) use mistral::validate_environment as validate_ai_environment;
use crate::ocr::types::OcrConnection;
async fn resolve_entra(
pub(super) async fn resolve_entra(
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<String>>, Error> {
) -> Result<Option<Sourced<String>>, crate::ocr::Error> {
static SERVICE: OnceLock<AzureAuthService> = OnceLock::new();
SERVICE
.get_or_init(AzureAuthService::default)
@ -36,18 +25,18 @@ async fn resolve_entra(
Sourced::new(value, source)
})
})
.map_err(Error::from)
.map_err(crate::ocr::Error::from)
}
fn validate_destination(
pub(super) fn validate_destination(
connection: &OcrConnection,
credential_source: InputSource,
) -> Result<(), OcrError> {
) -> Result<(), crate::ocr::Error> {
if connection.api_base.is_some()
&& connection.api_base_source == InputSource::Request
&& credential_source != InputSource::Request
{
return Err(Error::from(litellm_auth::Error::RequestAzureCredentialDestination).into());
return Err(litellm_auth::Error::RequestAzureCredentialDestination.into());
}
Ok(())
}

1
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

1261
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -0,0 +1,4 @@
pub(crate) mod cohere_parse_transformation;
pub(crate) mod common_utils;
pub(crate) mod document_intelligence;
pub(crate) mod transformation;

405
litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs Normal file
View file

@ -0,0 +1,405 @@
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
use serde_json::Value;
use crate::call_arguments::CallArguments;
use crate::constants::AZURE_AI_OCR_PATH;
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
use crate::llms::mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest};
use crate::ocr::OcrClient;
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::prepare::credential_env;
use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest};
use crate::params::OpaqueParams;
use crate::url_utils::ApiUrl;
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
#[derive(Clone, Debug, Default)]
pub(crate) struct AzureAiOcrConfig;
impl BaseOcrConfig for AzureAiOcrConfig {
type OcrParams = OpaqueParams;
type ProviderRequest = MistralOcrRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] {
MistralOcrConfig.get_supported_ocr_params(model)
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some(AZURE_AI_API_KEY_ENV)
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<OpaqueParams, crate::ocr::Error> {
MistralOcrConfig.map_ocr_params(non_default_params, model)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
let config = AzureAuthInputs {
azure_ad_token_provider: request.azure_ad_token_provider.clone(),
..AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?
};
self.resolve_headers(&request.connection, &config, &credential_env)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
self.build_ocr_url(request.connection.api_base.as_deref(), &credential_env)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
headers: &[(String, String)],
) -> Result<MistralOcrRequest, crate::ocr::Error> {
MistralOcrConfig.transform_ocr_request(model, document, optional_params, headers)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<MistralOcrRequest, crate::ocr::Error> {
let document = inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
MistralOcrConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
validate_inline_document(&crate::ocr::prepare::body_document(body)?)
}
}
impl AzureAiOcrConfig {
/// Python `AzureAIOCRConfig.validate_environment` requires the endpoint
/// before it resolves credentials; keep that order so a missing base is
/// reported without invoking any token provider.
pub(super) fn resolve_api_base(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, crate::ocr::Error> {
nonblank(api_base.map(str::to_string))
.or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV)))
.ok_or(crate::ocr::Error::Auth(
litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: AZURE_AI_API_BASE_ENV,
},
))
}
async fn resolve_headers(
&self,
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?;
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
if config.azure_ad_token_provider.is_some() {
super::common_utils::resolve_entra(config, env_lookup).await?;
}
super::common_utils::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(self.get_api_key_env_var().and_then(env_lookup))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::common_utils::validate_destination(connection, key.source())?;
return Ok(bearer_headers(connection, key.value()));
}
let key = super::common_utils::resolve_entra(config, env_lookup)
.await?
.ok_or(crate::ocr::Error::MissingAzureAiCredentials)?;
super::common_utils::validate_destination(connection, key.source())?;
Ok(bearer_headers(connection, key.value()))
}
fn build_ocr_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, crate::ocr::Error> {
let base = Self::resolve_api_base(api_base, env_lookup)?;
let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect();
ApiUrl::parse(&base)
.and_then(|url| url.complete_path(&path))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
}
fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> {
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use rstest::{fixture, rstest};
use super::*;
#[fixture]
fn connection() -> OcrConnection {
OcrConnection {
api_key: Some("request-key".into()),
api_base: Some("https://example.com".into()),
..Default::default()
}
}
#[rstest]
#[case::base_with_query(
"https://example.com/?tenant=a",
"https://example.com/providers/mistral/azure/ocr?tenant=a"
)]
#[case::complete_endpoint(
"https://example.com/providers/mistral/azure/ocr",
"https://example.com/providers/mistral/azure/ocr"
)]
fn completes_azure_path_and_preserves_query(#[case] api_base: &str, #[case] expected: &str) {
assert_eq!(
AzureAiOcrConfig
.build_ocr_url(Some(api_base), &|_| None)
.unwrap(),
expected
);
}
#[test]
fn missing_api_base_is_structured() {
assert!(matches!(
AzureAiOcrConfig::resolve_api_base(None, &|_| None),
Err(crate::ocr::Error::Auth(
litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: AZURE_AI_API_BASE_ENV,
}
))
));
}
#[rstest]
#[tokio::test]
async fn supplied_authorization_precedes_keys(connection: OcrConnection) {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer prepared".into())],
..connection
};
assert_eq!(
AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap(),
connection.extra_headers
);
}
#[rstest]
#[tokio::test]
async fn request_key_precedes_environment_key(connection: OcrConnection) {
assert_eq!(
AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap()[0],
("Authorization".into(), "Bearer request-key".into())
);
}
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|name| {
(name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
("Authorization".into(), "Bearer request-key".into())
);
}
use std::sync::Arc;
use serde_json::json;
use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
#[tokio::test]
async fn facade_executes_azure_mistral_with_prepared_auth() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"hello"}],
"usage_info":{"pages_processed":1}
}))])
.await;
let mut request = wire_request(
"azure_ai/model",
&base,
json!({"include_image_base64":true}),
);
request.credentials.api_key = None;
request.transport.extra_headers = vec![(
"Authorization".into(),
"Bearer python-prepared-token".into(),
)];
let result = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0].markdown, "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer python-prepared-token\r\n")
);
let body: Value =
serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({
"model":"model",
"document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"include_image_base64":true
})
);
}
#[tokio::test]
async fn facade_acquires_supplied_entra_token_for_final_request() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut request = wire_request(
"azure_ai/model",
&base,
json!({"azure_ad_token":"rust-owned-token"}),
);
request.credentials.api_key = None;
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer rust-owned-token\r\n")
);
}
struct ReplaceBodyDocument;
impl OcrHooks for ReplaceBodyDocument {
fn intercepts_requests(&self) -> bool {
true
}
fn during_call(
&self,
mut request: OcrDuringCallRequest,
) -> OcrHookFuture<'_, OcrDuringCallRequest> {
Box::pin(async move {
request.body["document"] = json!({
"type":"document_url",
"document_url":"https://example.com/not-inline.pdf"
});
Ok(request)
})
}
}
#[tokio::test]
async fn rejects_non_inline_body_after_guardrails() {
let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
request.hooks = Arc::new(ReplaceBodyDocument);
let error = perform_ocr(request).await.unwrap_err();
assert!(error.to_string().contains("data URI"));
}
}

0
litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs → litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/mod.rs
View file

38
litellm-rust/crates/core/src/messages/transformation.rs → litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/transformation.rs
View file

@ -1,5 +1,5 @@
use super::Error;
use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
use crate::messages::Error;
use crate::messages::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MessagesAuthStrategy {
@ -16,14 +16,29 @@ impl MessagesAuthStrategy {
}
}
pub trait AnthropicMessagesProviderConfig: Sync {
fn complete_url(
pub trait BaseAnthropicMessagesConfig: Sync {
fn get_complete_url(
&self,
api_base: Option<&str>,
model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_anthropic_messages_request(
&self,
request: AnthropicMessagesRequest,
) -> Result<AnthropicMessagesRequest, Error> {
Ok(request)
}
fn transform_anthropic_messages_response(
&self,
_model: &str,
response: AnthropicMessagesResponse,
) -> Result<AnthropicMessagesResponse, Error> {
Ok(response)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
@ -44,19 +59,4 @@ pub trait AnthropicMessagesProviderConfig: Sync {
("content-type", "application/json"),
]
}
fn transform_request(
&self,
request: AnthropicMessagesRequest,
) -> Result<AnthropicMessagesRequest, Error> {
Ok(request)
}
fn transform_response(
&self,
_model: &str,
response: AnthropicMessagesResponse,
) -> Result<AnthropicMessagesResponse, Error> {
Ok(response)
}
}

0
litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs → litellm-rust/crates/core/src/llms/base_llm/audio_transcription/mod.rs
View file

50
litellm-rust/crates/core/src/audio_transcription/transformation.rs → litellm-rust/crates/core/src/llms/base_llm/audio_transcription/transformation.rs
View file

@ -1,7 +1,9 @@
use super::Error;
use serde_json::{Map, Value};
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};
use crate::audio_transcription::Error;
use crate::audio_transcription::types::{
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AudioTranscriptionAuth {
@ -12,34 +14,21 @@ pub enum AudioTranscriptionAuth {
},
}
pub trait AudioTranscriptionProviderConfig: Sync {
fn supported_transcription_params(&self) -> &'static [&'static str];
pub trait BaseAudioTranscriptionConfig: Sync {
fn get_supported_openai_params(&self) -> &'static [&'static str];
fn map_transcription_params(&self, params: &Map<String, Value>) -> Map<String, Value> {
params
fn map_transcription_params(
&self,
non_default_params: &Map<String, Value>,
) -> Map<String, Value> {
non_default_params
.iter()
.filter(|(key, _)| {
self.supported_transcription_params()
.contains(&key.as_str())
})
.filter(|(key, _)| self.get_supported_openai_params().contains(&key.as_str()))
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
fn transform_transcription_request(
&self,
model: &str,
audio: Value,
optional_params: Map<String, Value>,
) -> Result<AudioTranscriptionRequestData, Error>;
fn transform_transcription_response(
&self,
model: &str,
response_json: Value,
) -> Result<AudioTranscriptionResponseData, Error>;
fn complete_url(
fn get_complete_url(
&self,
api_base: Option<&str>,
model: &str,
@ -47,6 +36,19 @@ pub trait AudioTranscriptionProviderConfig: Sync {
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_audio_transcription_request(
&self,
model: &str,
audio: Value,
optional_params: Map<String, Value>,
) -> Result<AudioTranscriptionRequestData, Error>;
fn transform_audio_transcription_response(
&self,
model: &str,
response_json: Value,
) -> Result<AudioTranscriptionResponseData, Error>;
fn auth_strategy(
&self,
model: &str,

0
litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs → litellm-rust/crates/core/src/llms/base_llm/chat/mod.rs
View file

52
litellm-rust/crates/core/src/chat_completions/transformation.rs → litellm-rust/crates/core/src/llms/base_llm/chat/transformation.rs
View file

@ -1,11 +1,17 @@
use super::Error;
use serde_json::{Map, Value};
use super::types::{
use crate::chat_completions::Error;
use crate::chat_completions::types::{
ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData,
ProviderChatResponseData,
};
pub const STREAM_PARAM: &str = "stream";
/// Message fields that carry no meaning for the upstream body, so their
/// 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)]
@ -25,14 +31,11 @@ pub enum ChatCompletionsAuth {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Unsupported(pub &'static str);
pub const STREAM_PARAM: &str = "stream";
pub trait BaseConfig: Sync {
/// Supported OpenAI parameter names paired with their provider names.
fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)];
/// Message fields that carry no meaning for the upstream body, so their
/// presence does not make a request untranslatable.
const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"];
pub trait ChatCompletionsProviderConfig: Sync {
fn complete_url(
fn get_complete_url(
&self,
api_base: Option<&str>,
model: &str,
@ -40,6 +43,19 @@ pub trait ChatCompletionsProviderConfig: Sync {
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_request(
&self,
model: &str,
messages: Vec<ChatMessage>,
optional_params: Map<String, Value>,
) -> Result<ProviderChatRequestData, Error>;
fn transform_response(
&self,
model: &str,
response: ProviderChatResponseData,
) -> Result<ChatCompletionsResponse, Error>;
fn auth(
&self,
api_key: Option<&str>,
@ -62,9 +78,6 @@ pub trait ChatCompletionsProviderConfig: Sync {
false
}
/// Supported OpenAI parameter names paired with their provider names.
fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)];
/// Parameters consumed as call configuration (credentials, endpoints)
/// rather than placed in the body. Accepted, never serialized.
fn config_params(&self) -> &'static [&'static str] {
@ -77,25 +90,12 @@ pub trait ChatCompletionsProviderConfig: Sync {
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(
self.supported_openai_params(),
self.supported_openai_param_mappings(),
self.config_params(),
optional_params,
)
.or_else(|| messages.iter().find_map(unsupported_message))
}
fn transform_request(
&self,
model: &str,
messages: Vec<ChatMessage>,
optional_params: Map<String, Value>,
) -> Result<ProviderChatRequestData, Error>;
fn transform_response(
&self,
model: &str,
response: ProviderChatResponseData,
) -> Result<ChatCompletionsResponse, Error>;
}
pub fn unsupported_param(

4
litellm-rust/crates/core/src/llms/base_llm/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub mod anthropic_messages;
pub mod audio_transcription;
pub mod chat;
pub(crate) mod ocr;

1
litellm-rust/crates/core/src/llms/base_llm/ocr/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

211
litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs Normal file
View file

@ -0,0 +1,211 @@
use std::future::Future;
use std::sync::Arc;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use crate::call_arguments::CallArguments;
use crate::ocr::OcrClient;
use crate::ocr::hooks::OcrHooks;
use crate::ocr::types::{
LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrResponseFormat,
PreparedOcrRequest, ResolvedOcrCredentials,
};
const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y=";
/// Output of `validate_environment`: whatever a provider resolves up front
/// (headers at minimum; Vertex also carries the project id).
pub(crate) trait OcrEnvironment: Send + Sync {
fn headers(&self) -> &[(String, String)];
}
impl OcrEnvironment for Vec<(String, String)> {
fn headers(&self) -> &[(String, String)] {
self
}
}
#[derive(Clone, Copy)]
pub(crate) struct OcrRequestContext<'a> {
pub client: &'a OcrClient,
pub connection: &'a OcrConnection,
}
#[derive(Clone, Copy)]
pub(crate) struct OcrResponseContext<'a> {
pub client: &'a OcrClient,
pub connection: &'a OcrConnection,
pub hooks: &'a Arc<dyn OcrHooks>,
pub request_format: OcrResponseFormat,
pub url: &'a str,
pub headers: &'a [(String, String)],
}
pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static {
type OcrParams: Send + Sync;
type ProviderRequest: Serialize + Send;
type Environment: OcrEnvironment;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&[]
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
None
}
fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials {
ResolvedOcrCredentials {
api_key: inputs
.dynamic_api_key
.filter(|value| !value.value().is_empty())
.or(inputs.api_key),
api_base: inputs
.dynamic_api_base
.filter(|value| !value.value().is_empty())
.or(inputs.api_base),
}
}
fn get_health_check_document(&self) -> OcrDocument {
OcrDocument::DocumentUrl {
document_url: HEALTH_CHECK_PDF_DATA_URI.into(),
extra_fields: Default::default(),
}
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<Self::OcrParams, crate::ocr::Error>;
fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<Self::Environment, crate::ocr::Error>> + Send;
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
optional_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error>;
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
headers: &[(String, String)],
) -> Result<Self::ProviderRequest, crate::ocr::Error>;
fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
headers: &[(String, String)],
_context: OcrRequestContext<'_>,
) -> impl Future<Output = Result<Self::ProviderRequest, crate::ocr::Error>> + Send {
async move { self.transform_ocr_request(model, document, optional_params, headers) }
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error>;
fn async_transform_ocr_response(
&self,
model: &str,
raw_response: reqwest::Response,
context: OcrResponseContext<'_>,
) -> impl Future<Output = Result<LiteLLMOcrResponse, crate::ocr::Error>> + Send {
async move {
let bytes = crate::ocr::client::read_response_bytes(
raw_response,
context.connection.max_response_bytes,
)
.await?;
crate::ocr::handler::post_call(context.hooks, &bytes).await?;
self.transform_ocr_response(model, &bytes, context.request_format)
}
}
fn get_error_class(
&self,
error_message: String,
status_code: u16,
headers: Vec<(String, String)>,
) -> crate::ocr::Error {
crate::ocr::Error::Provider {
status: status_code,
body: error_message,
headers,
}
}
/// Provider-specific check applied to the composed body, both before and
/// after guardrail hooks. Defaults to accepting any body.
fn validate_request_body(&self, _body: &Value) -> Result<(), crate::ocr::Error> {
Ok(())
}
/// Rust counterpart of `BaseLLMHTTPHandler._async_prepare_ocr_request`:
/// map params, validate environment, build URL, transform, compose body.
fn prepare_request(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<reqwest::Request, crate::ocr::Error>> + Send {
async move {
let params = self.map_ocr_params(&request.optional_params, &request.model)?;
let environment = self.validate_environment(request, client).await?;
let url = self.get_complete_url(request, &params, &environment)?;
let headers = environment.headers();
let body = self
.async_transform_ocr_request(
&request.model,
request.document.clone(),
&params,
headers,
OcrRequestContext {
client,
connection: &request.connection,
},
)
.await?;
crate::ocr::prepare::transform_request_body(
client,
request,
&url,
headers,
body,
|body| self.validate_request_body(body),
)
.await
}
}
}
pub(crate) fn decode_and_normalize_response<T: DeserializeOwned>(
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
normalize: impl FnOnce(&str, T) -> Result<LiteLLMOcrResponse, crate::ocr::Error>,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let decoded = crate::ocr::json::decode_response(
raw_response,
request_format == OcrResponseFormat::Native,
)?;
Ok(LiteLLMOcrResponse {
provider_native_response: decoded.native,
..normalize(model, decoded.data)?
})
}

29
litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs → litellm-rust/crates/core/src/llms/bedrock/audio_transcription/mod.rs
View file

@ -1,16 +1,15 @@
use serde_json::{Map, Value, json};
use crate::audio_transcription::Error;
use crate::audio_transcription::transformation::{
AudioTranscriptionAuth, AudioTranscriptionProviderConfig,
};
use crate::audio_transcription::types::{
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
};
use crate::http_utils::json_type_name;
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
use crate::llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
};
use litellm_auth_aws::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
use litellm_auth_aws::{bedrock_model_id_and_region, resolve_bedrock_region};
const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"];
@ -46,12 +45,12 @@ fn optional_string<'a>(params: &'a Map<String, Value>, key: &str) -> Option<&'a
.filter(|value| !value.is_empty())
}
impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
fn supported_transcription_params(&self) -> &'static [&'static str] {
impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig {
fn get_supported_openai_params(&self) -> &'static [&'static str] {
SUPPORTED_PARAMS
}
fn transform_transcription_request(
fn transform_audio_transcription_request(
&self,
_model: &str,
audio: Value,
@ -84,7 +83,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
})
}
fn transform_transcription_response(
fn transform_audio_transcription_response(
&self,
_model: &str,
response_json: Value,
@ -106,7 +105,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
Ok(AudioTranscriptionResponseData { text })
}
fn complete_url(
fn get_complete_url(
&self,
api_base: Option<&str>,
model: &str,
@ -161,7 +160,7 @@ mod tests {
]);
let params = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.map_transcription_params(&params);
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
.transform_transcription_request(
.transform_audio_transcription_request(
"mistral.voxtral-mini-3b-2507",
json!({"data": "AQI=", "format": "wav", "filename": "sample.wav"}),
params,
@ -186,7 +185,7 @@ mod tests {
#[test]
fn response_concatenates_content_blocks() {
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
.transform_transcription_response(
.transform_audio_transcription_response(
"model",
json!({"output": {"message": {"content": [{"text": "hello "}, {"text": "world"}]}}}),
)
@ -197,7 +196,7 @@ mod tests {
#[test]
fn invalid_audio_is_rejected() {
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_transcription_request(
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_audio_transcription_request(
"model",
json!({"data": "AQI="}),
Map::new(),
@ -209,7 +208,7 @@ mod tests {
fn region_and_url_precedence_match_python() {
let params = Map::from_iter([("aws_region_name".to_string(), json!("eu-west-1"))]);
let url = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
.complete_url(
.get_complete_url(
None,
"bedrock/us-east-1/mistral.voxtral-mini-3b-2507",
&params,

267
litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs → litellm-rust/crates/core/src/llms/bedrock/chat/converse_transformation.rs
View file

@ -3,18 +3,16 @@ use serde_json::{Map, Value, json};
use crate::chat_completions::Error;
use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::chat_completions::transformation::{
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
unsupported_param,
};
use crate::chat_completions::types::{
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse,
ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData,
ProviderChatResponseData,
};
use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region};
use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};
use crate::llms::base_llm::chat::transformation::{
BaseConfig, ChatCompletionsAuth, Unsupported, unsupported_message, unsupported_param,
};
use litellm_auth_aws::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};
use litellm_auth_aws::{bedrock_model_id_and_region, resolve_bedrock_region};
/// Converse parameter names, post `map_openai_params`, that the Rust path can
/// place verbatim in `inferenceConfig`.
@ -50,62 +48,16 @@ const CONFIG_PARAMS: &[&str] = &[
const CONVERSE_PATH_SUFFIX: &str = "/converse";
pub struct BedrockChatCompletionsConfig;
pub struct AmazonConverseConfig;
pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: BedrockChatCompletionsConfig =
BedrockChatCompletionsConfig;
pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: AmazonConverseConfig = AmazonConverseConfig;
fn converse_body(conversation: &Conversation, params: &Map<String, Value>) -> Value {
let messages: Vec<Value> = conversation
.turns
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"content": turn.texts.iter().map(|text| json!({"text": text})).collect::<Vec<_>>(),
})
})
.collect();
let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| {
params
.get(*name)
.map(|value| ((*name).to_string(), value.clone()))
}));
let system: Vec<Value> = conversation
.system
.iter()
.map(|text| json!({"text": text}))
.collect();
Value::Object(Map::from_iter(
[
(
"inferenceConfig".to_string(),
Value::Object(inference_config),
),
("messages".to_string(), json!(messages)),
]
.into_iter()
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))),
))
}
fn has_blank_text(message: &ChatMessage) -> bool {
match &message.content {
None => false,
Some(ChatMessageContent::Text(text)) => text.trim().is_empty(),
Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| {
part.get("text")
.and_then(Value::as_str)
.is_none_or(|text| text.trim().is_empty())
}),
impl BaseConfig for AmazonConverseConfig {
fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] {
SUPPORTED_PARAMS
}
}
impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
fn complete_url(
fn get_complete_url(
&self,
api_base: Option<&str>,
model: &str,
@ -132,82 +84,6 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
Ok(format!("{endpoint}/model/{model_id}{CONVERSE_PATH_SUFFIX}"))
}
fn auth(
&self,
api_key: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, 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
// all-whitespace token stays a bearer token here because Python sends
// it too: treating it as absent would sign as the host principal
// instead, which is the identity swap this branch exists to prevent.
let bearer = match api_key {
Some(key) => Some(key.to_string()),
None => env_lookup(AWS_BEARER_TOKEN_BEDROCK),
}
.filter(|token| !token.is_empty());
if let Some(token) = bearer {
return Ok(ChatCompletionsAuth::Bearer { token });
}
let (_, model_region) = bedrock_model_id_and_region(model);
Ok(ChatCompletionsAuth::AwsSigV4 {
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[("Content-Type", "application/json")]
}
fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] {
SUPPORTED_PARAMS
}
fn config_params(&self) -> &'static [&'static str] {
CONFIG_PARAMS
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(
self.supported_openai_params(),
CONFIG_PARAMS,
optional_params,
)
.or_else(|| messages.iter().find_map(unsupported_message))
// Python's Converse translation drops blank text blocks instead of
// substituting the placeholder the shared conversation builder
// applies, so decline blank text rather than diverge.
.or_else(|| {
messages
.iter()
.any(has_blank_text)
.then_some(Unsupported("blank message text"))
})
// Converse has no assistant prefill: Python inserts a continue turn
// when a conversation opens or closes on an assistant message, and
// only under `litellm.modify_params`, which the core cannot see.
// Declining both ends also keeps the shared builder's final
// assistant right-strip (an Anthropic rule) unreachable here.
.or_else(|| {
let conversation = build_conversation(messages);
let ends_on_assistant = conversation
.turns
.last()
.is_some_and(|turn| turn.role == TurnRole::Assistant);
(!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
"conversation does not run user turn to user turn",
))
})
}
fn transform_request(
&self,
_model: &str,
@ -296,6 +172,127 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
usage,
})
}
fn auth(
&self,
api_key: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, 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
// all-whitespace token stays a bearer token here because Python sends
// it too: treating it as absent would sign as the host principal
// instead, which is the identity swap this branch exists to prevent.
let bearer = match api_key {
Some(key) => Some(key.to_string()),
None => env_lookup(AWS_BEARER_TOKEN_BEDROCK),
}
.filter(|token| !token.is_empty());
if let Some(token) = bearer {
return Ok(ChatCompletionsAuth::Bearer { token });
}
let (_, model_region) = bedrock_model_id_and_region(model);
Ok(ChatCompletionsAuth::AwsSigV4 {
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[("Content-Type", "application/json")]
}
fn config_params(&self) -> &'static [&'static str] {
CONFIG_PARAMS
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(
self.supported_openai_param_mappings(),
CONFIG_PARAMS,
optional_params,
)
.or_else(|| messages.iter().find_map(unsupported_message))
// Python's Converse translation drops blank text blocks instead of
// substituting the placeholder the shared conversation builder
// applies, so decline blank text rather than diverge.
.or_else(|| {
messages
.iter()
.any(has_blank_text)
.then_some(Unsupported("blank message text"))
})
// Converse has no assistant prefill: Python inserts a continue turn
// when a conversation opens or closes on an assistant message, and
// only under `litellm.modify_params`, which the core cannot see.
// Declining both ends also keeps the shared builder's final
// assistant right-strip (an Anthropic rule) unreachable here.
.or_else(|| {
let conversation = build_conversation(messages);
let ends_on_assistant = conversation
.turns
.last()
.is_some_and(|turn| turn.role == TurnRole::Assistant);
(!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
"conversation does not run user turn to user turn",
))
})
}
}
fn converse_body(conversation: &Conversation, optional_params: &Map<String, Value>) -> Value {
let messages: Vec<Value> = conversation
.turns
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"content": turn.texts.iter().map(|text| json!({"text": text})).collect::<Vec<_>>(),
})
})
.collect();
let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| {
optional_params
.get(*name)
.map(|value| ((*name).to_string(), value.clone()))
}));
let system: Vec<Value> = conversation
.system
.iter()
.map(|text| json!({"text": text}))
.collect();
Value::Object(Map::from_iter(
[
(
"inferenceConfig".to_string(),
Value::Object(inference_config),
),
("messages".to_string(), json!(messages)),
]
.into_iter()
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))),
))
}
fn has_blank_text(message: &ChatMessage) -> bool {
match &message.content {
None => false,
Some(ChatMessageContent::Text(text)) => text.trim().is_empty(),
Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| {
part.get("text")
.and_then(Value::as_str)
.is_none_or(|text| text.trim().is_empty())
}),
}
}
#[cfg(test)]

1
litellm-rust/crates/core/src/llms/bedrock/chat/mod.rs Normal file
View file

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

15
litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs → litellm-rust/crates/core/src/llms/bedrock/chat/tests.rs
View file

@ -1,6 +1,7 @@
use serde_json::json;
use super::*;
use crate::chat_completions::Error;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
@ -225,7 +226,7 @@ fn builds_the_converse_url_from_the_region_in_the_model_id() {
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
assert_eq!(
config
.complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| {
.get_complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| {
None
})
.expect("url builds"),
@ -239,13 +240,13 @@ fn falls_back_to_the_region_env_then_the_default_region() {
let with_env = |key: &str| (key == "AWS_REGION_NAME").then(|| "eu-west-1".to_string());
assert_eq!(
config
.complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env)
.get_complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env)
.expect("url builds"),
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
assert_eq!(
config
.complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None)
.get_complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None)
.expect("url builds"),
"https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-v2/converse"
);
@ -257,7 +258,7 @@ fn prefers_an_explicit_runtime_endpoint_over_the_api_base() {
let overrides = params(json!({"aws_bedrock_runtime_endpoint": "https://vpce.internal/"}));
assert_eq!(
config
.complete_url(
.get_complete_url(
Some("https://ignored.example"),
"anthropic.claude-v2",
&overrides,
@ -539,7 +540,7 @@ fn leaves_a_complete_converse_url_untouched() {
"https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-v2%3A0/converse";
assert_eq!(
config
.complete_url(
.get_complete_url(
Some(already_built),
"anthropic.claude-v2",
&Map::new(),
@ -553,7 +554,7 @@ fn leaves_a_complete_converse_url_untouched() {
#[test]
fn host_supplied_credentials_outrank_ambient_profile_and_role_state() {
use crate::providers::bedrock::aws_base::host_supplied_credentials;
use litellm_auth_aws::host_supplied_credentials;
let supplied = params(json!({
"aws_access_key_id": "AKIAHOST",

2
litellm-rust/crates/core/src/llms/bedrock/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod audio_transcription;
pub mod chat;

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

@ -0,0 +1 @@
pub(crate) mod ocr;

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

@ -0,0 +1,3 @@
pub(crate) mod transformation;
pub(crate) use transformation::{CohereOptions, validate_document};

782
litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs Normal file
View file

@ -0,0 +1,782 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use serde_with::serde_as;
use crate::call_arguments::{CallArguments, parse_options};
use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE};
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response};
use crate::ocr::OcrClient;
use crate::ocr::document::InlineDocument;
use crate::ocr::prepare::credential_env;
use crate::ocr::types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage, OcrResponseFormat,
OcrUsageInfo, PreparedOcrRequest,
};
use crate::serde_compat::LaxI64;
use crate::url_utils::ApiUrl;
const COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC";
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum OutputFormat {
#[default]
Markdown,
Blocks,
}
#[derive(Default, Deserialize, Serialize)]
pub(crate) struct CohereOptions {
#[serde(skip_serializing_if = "Option::is_none")]
pub output_format: Option<OutputFormat>,
}
#[derive(Deserialize, Serialize)]
pub(crate) struct CohereRequest {
pub model: String,
pub document: CohereParseDocument,
pub output_format: String,
}
#[derive(Deserialize, Serialize)]
#[serde(tag = "type")]
pub(crate) enum CohereParseDocument {
#[serde(rename = "image_url")]
ImageUrl { image_url: String },
}
#[derive(Deserialize)]
pub(crate) struct CohereResponse {
#[serde(default)]
pages: Vec<CoherePage>,
meta: Option<CohereMeta>,
}
#[serde_as]
#[derive(Deserialize)]
struct CoherePage {
#[serde_as(deserialize_as = "Option<LaxI64>")]
index: Option<i64>,
markdown: Option<CohereMarkdown>,
blocks: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize, Serialize)]
struct CohereMarkdown {
#[serde(default)]
content: String,
images: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize)]
struct CohereMeta {
billed_units: Option<CohereBilledUnits>,
}
#[serde_as]
#[derive(Deserialize)]
struct CohereBilledUnits {
#[serde_as(deserialize_as = "Option<LaxI64>")]
pages: Option<i64>,
}
#[derive(Default)]
pub(crate) struct CohereParseConfig;
impl BaseOcrConfig for CohereParseConfig {
type OcrParams = CohereOptions;
type ProviderRequest = CohereRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&["output_format", "req_format"]
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some(COHERE_API_KEY_ENV)
}
fn get_health_check_document(&self) -> OcrDocument {
OcrDocument::ImageUrl {
image_url: COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI.into(),
extra_fields: Default::default(),
}
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
_model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
Ok(parse_options(non_default_params)?)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
self.resolve_headers(&request.connection, &credential_env)
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
self.build_ocr_url(
request
.connection
.api_base
.as_deref()
.unwrap_or(COHERE_PARSE_API_BASE),
)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &CohereOptions,
_headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
let image_url = image_url(document)?;
Ok(build_request(model, image_url, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
validate_document(&crate::ocr::prepare::body_document(body)?)
}
}
impl CohereParseConfig {
fn resolve_headers(
&self,
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
self.get_api_key_env_var()
.and_then(env_lookup)
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| {
crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication(
"Missing COHERE_API_KEY - set it in the environment or pass api_key".into(),
))
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn build_ocr_url(&self, api_base: &str) -> Result<String, crate::ocr::Error> {
let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
ApiUrl::parse(api_base)
.and_then(|url| url.complete_path(&["v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
}
}
pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), crate::ocr::Error> {
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(crate::ocr::Error::CohereImageOnly);
};
if image_url.is_empty() {
return Err(crate::ocr::Error::CohereImageOnly);
}
if let Some(inline) = InlineDocument::parse(image_url)? {
if !inline.mime_type().type_.eq_ignore_ascii_case("image") {
return Err(crate::ocr::Error::CohereImageOnly);
}
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
}
Ok(())
}
pub(crate) fn normalize_response(
model: &str,
response: CohereResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let pages_processed = billed_pages(&response).map(Ok).unwrap_or_else(|| {
i64::try_from(response.pages.len()).map_err(|_| crate::ocr::Error::NumericRange("pages"))
})?;
let pages = response
.pages
.into_iter()
.enumerate()
.map(|(position, page)| normalize_page(page, position))
.collect::<Result<Vec<_>, crate::ocr::Error>>()?;
Ok(LiteLLMOcrResponse {
usage_info: Some(OcrUsageInfo {
pages_processed: Some(pages_processed),
..Default::default()
}),
..LiteLLMOcrResponse::new(model, pages)
})
}
fn image_url(document: OcrDocument) -> Result<String, crate::ocr::Error> {
validate_document(&document)?;
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(crate::ocr::Error::CohereImageOnly);
};
Ok(image_url)
}
fn build_request(model: &str, image_url: String, params: &CohereOptions) -> CohereRequest {
CohereRequest {
model: model.into(),
document: CohereParseDocument::ImageUrl { image_url },
output_format: match params.output_format.unwrap_or_default() {
OutputFormat::Markdown => "markdown",
OutputFormat::Blocks => "blocks",
}
.into(),
}
}
fn page_image(
mut image: Map<String, Value>,
path: &str,
) -> Result<OcrPageImage, crate::ocr::Error> {
if let Some(Value::Object(bbox)) = image.get("bounding_box") {
image.insert("bbox".into(), Value::Object(bbox.clone()));
}
crate::ocr::json::decode_response_value(Value::Object(image), path)
}
fn normalize_page(page: CoherePage, position: usize) -> Result<OcrPage, crate::ocr::Error> {
let index = page.index.map(Ok).unwrap_or_else(|| {
i64::try_from(position).map_err(|_| crate::ocr::Error::NumericRange("page index"))
})?;
let (markdown, images) = match page.markdown {
Some(markdown) => {
let images = markdown
.images
.filter(|images| !images.is_empty())
.map(|images| {
images
.into_iter()
.enumerate()
.map(|(image_index, image)| {
page_image(
image,
&format!("pages[{position}].markdown.images[{image_index}]"),
)
})
.collect::<Result<Vec<_>, _>>()
})
.transpose()?;
(markdown.content, images)
}
None => (String::new(), None),
};
let extra_fields = page
.blocks
.map(|blocks| {
(
"blocks".into(),
Value::Array(blocks.into_iter().map(Value::Object).collect()),
)
})
.into_iter()
.collect();
Ok(OcrPage {
index,
markdown,
images,
extra_fields,
..Default::default()
})
}
fn billed_pages(response: &CohereResponse) -> Option<i64> {
response.meta.as_ref()?.billed_units.as_ref()?.pages
}
fn invalid_api_base() -> crate::ocr::Error {
crate::ocr::Error::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[tokio::test]
async fn composed_body_preserves_native_document_fields_and_untyped_overrides() {
let request = crate::ocr::test_support::wire_request(
"cohere/parse",
"https://example.com",
json!({
"output_format":"markdown", "metadata":{"host":true},
"extra_body":{
"output_format": {"future":true},
"document":{"type":"image_url","image_url":"https://example.com/a.png",
"provider_options":{"nested":[false,0,null]}}
}
}),
);
let request = request.with_document(
serde_json::from_value(json!({
"type":"image_url","image_url":"https://example.com/original.png"
}))
.unwrap(),
);
let request = crate::ocr::prepare::prepare_request(request);
let http = CohereParseConfig
.prepare_request(&request, &crate::ocr::test_support::ocr_client())
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(
body,
json!({
"model":"parse", "output_format":{"future":true},
"document":{"type":"image_url","image_url":"https://example.com/a.png",
"provider_options":{"nested":[false,0,null]}}
})
);
}
#[rstest]
#[case::cohere(false)]
#[case::azure(true)]
fn options_read_known_fields_without_changing_arguments(#[case] azure: bool) {
let arguments = serde_json::from_value(json!({
"output_format":"blocks", "req_format":"native", "extension":false
}))
.unwrap();
let mapped = if azure {
crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig
.map_ocr_params(&arguments, "parse")
} else {
CohereParseConfig.map_ocr_params(&arguments, "parse")
}
.unwrap();
assert_eq!(
serde_json::to_value(mapped).unwrap(),
json!({"output_format":"blocks"})
);
assert_eq!(arguments["req_format"], "native");
assert_eq!(arguments["extension"], false);
}
#[test]
fn options_reject_invalid_output_format() {
let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap();
assert!(matches!(
CohereParseConfig.map_ocr_params(&invalid, "parse"),
Err(crate::ocr::Error::RequestField { path })
if path == "optional_params.output_format"
));
}
#[test]
fn billed_pages_accept_integral_doubles() {
let response = serde_json::from_str::<CohereResponse>(
r#"{"pages":[],"meta":{"billed_units":{"pages":1.0}}}"#,
)
.unwrap();
let normalized = normalize_response("parse", response).unwrap();
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1));
}
#[test]
fn billed_pages_reject_fractional_counts() {
assert!(
serde_json::from_str::<CohereResponse>(
r#"{"pages":[],"meta":{"billed_units":{"pages":1.5}}}"#,
)
.is_err()
);
}
#[test]
fn response_preserves_python_mapping_shapes_and_extensions() {
let blocks = json!([
{"type":"text", "text":"Total Due: $4.00"},
{"type":"future", "payload":{"nested":[null,false,0]}}
]);
let response = serde_json::from_value(json!({
"pages":[{
"index":"2",
"markdown":{"content":"receipt", "images":[
{"bounding_box":{"x":1}, "bbox":"replaced", "category":"future", "extension":null},
{"image_base64":"encoded"}
]},
"blocks":blocks
}],
"meta":{"billed_units":{"pages":0}}
})).unwrap();
let response = normalize_response("parse", response).unwrap();
assert_eq!(response.pages[0].index, 2);
assert_eq!(response.usage_info.unwrap().pages_processed, Some(0));
assert_eq!(response.pages[0].extra_fields["blocks"], blocks);
let images = response.pages[0].images.as_ref().unwrap();
assert_eq!(images[0].bbox.as_ref().unwrap()["x"], 1);
assert_eq!(images[0].extra_fields["category"], "future");
assert_eq!(images[0].extra_fields.get("extension"), Some(&Value::Null));
assert_eq!(images[1].image_base64.as_deref(), Some("encoded"));
assert!(images[1].bbox.is_none());
}
#[test]
fn malformed_normalized_image_fields_report_the_original_path() {
let response = serde_json::from_value(json!({
"pages":[{"markdown":{"images":[{"image_base64":42}]}}]
}))
.unwrap();
assert!(matches!(
normalize_response("parse", response).unwrap_err(),
crate::ocr::Error::ResponseField { path }
if path == "pages[0].markdown.images[0].image_base64"
));
}
#[rstest]
fn provider_options_exclude_response_controls_and_extensions(
#[values("markdown", "blocks")] output_format: &str,
#[values("https://example.com/a.png", "data:image/png;base64,YWJj")] source: &str,
) {
let arguments = serde_json::from_value(
json!({"output_format":output_format,"req_format":"native","unknown":true}),
)
.unwrap();
let params = CohereParseConfig
.map_ocr_params(&arguments, "parse")
.unwrap();
assert_eq!(
serde_json::to_value(&params).unwrap(),
json!({"output_format":output_format})
);
let document = serde_json::from_value(
json!({"type":"image_url","image_url":source,"ignored":"field"}),
)
.unwrap();
let body = CohereParseConfig
.transform_ocr_request("parse", document, &params, &[])
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap(),
json!({
"model":"parse", "document":{"type":"image_url","image_url":source}, "output_format":output_format
})
);
}
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let request = crate::ocr::test_support::wire_request(
"cohere/parse",
"https://example.com",
json!({"output_format":null,"req_format":null}),
);
let request = request.with_document(
serde_json::from_value(
json!({"type":"image_url","image_url":"https://example.com/a.png"}),
)
.unwrap(),
);
assert_eq!(
request.response_format().unwrap(),
crate::ocr::types::OcrResponseFormat::Litellm
);
let request = crate::ocr::prepare::prepare_request(request);
let http = CohereParseConfig
.prepare_request(&request, &crate::ocr::test_support::ocr_client())
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[rstest]
fn response_normalizes_markdown_images_blocks_and_billed_pages() {
let payload = json!({
"pages": [
{
"type":"markdown",
"index":4,
"markdown":{
"content":"receipt",
"images":[{
"id":"image",
"bounding_box":{
"top_left_x":1,
"top_left_y":2,
"bottom_right_x":48,
"bottom_right_y":49
},
"bounding_box_normalized":{
"top_left_x":0.04,
"top_left_y":0.05,
"bottom_right_x":0.15,
"bottom_right_y":0.16
},
"description":"scan",
"category":"logo",
"provider_extension":"preserved"
}]
}
},
{"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]}
],
"meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}}
});
let response = serde_json::from_value(payload.clone()).unwrap();
let normalized = normalize_response("parse-v5.0", response).unwrap();
assert_eq!(normalized.pages[0].index, 4);
assert_eq!(normalized.pages[0].markdown, "receipt");
let image = &normalized.pages[0].images.as_ref().unwrap()[0];
let original_image = &payload["pages"][0]["markdown"]["images"][0];
assert_eq!(
serde_json::to_value(&image.bbox).unwrap(),
original_image["bounding_box"]
);
assert_eq!(
image.extra_fields["bounding_box_normalized"],
original_image["bounding_box_normalized"]
);
assert_eq!(image.extra_fields["id"], original_image["id"]);
assert_eq!(image.extra_fields["description"], "scan");
assert_eq!(image.extra_fields["category"], "logo");
assert_eq!(image.extra_fields["provider_extension"], "preserved");
assert_eq!(normalized.pages[1].index, 1);
assert_eq!(normalized.pages[1].markdown, "");
assert_eq!(
normalized.pages[1].extra_fields["blocks"][0]["text"]["content"],
"total"
);
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(3));
}
#[rstest]
#[case::empty(json!({}))]
#[case::null_meta(json!({"meta":null}))]
#[case::null_billed_units(json!({"pages":[],"meta":{"billed_units":null}}))]
fn response_defaults(#[case] value: Value) {
let normalized =
normalize_response("parse", serde_json::from_value(value).unwrap()).unwrap();
assert!(normalized.pages.is_empty());
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(0));
}
#[rstest]
#[case::null_pages(json!({"pages":null}))]
#[case::invalid_markdown(json!({"pages":[{"markdown":"text"}]}))]
#[case::invalid_index(json!({"pages":[{"index":"bad"}]}))]
fn response_rejects_invalid_fields(#[case] value: Value) {
assert!(serde_json::from_value::<CohereResponse>(value).is_err());
}
#[test]
fn null_markdown_uses_page_defaults() {
let normalized = normalize_response(
"parse",
serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(),
)
.unwrap();
assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1));
assert!(normalized.pages[0].images.is_none());
}
#[rstest]
fn response_types_documented_block_variants(
#[values(
crate::ocr::types::OcrResponseFormat::Litellm,
crate::ocr::types::OcrResponseFormat::Native
)]
response_format: crate::ocr::types::OcrResponseFormat,
) {
let payload = json!({
"pages": [{
"type": "blocks",
"index": 0,
"blocks": [
{"type": "text", "text": {"content": "hello"}},
{
"type": "image",
"image": {
"id": "img-0",
"description": "logo",
"category": "logo",
"bounding_box": {
"top_left_x": 1,
"top_left_y": 2,
"bottom_right_x": 3,
"bottom_right_y": 4
},
"bounding_box_normalized": {
"top_left_x": 0.1,
"top_left_y": 0.2,
"bottom_right_x": 0.3,
"bottom_right_y": 0.4
}
}
},
{
"type": "table",
"table": {
"type": "html",
"html": "<table></table>",
"bounding_box": {
"top_left_x": 5,
"top_left_y": 6,
"bottom_right_x": 7,
"bottom_right_y": 8
},
"bounding_box_normalized": {
"top_left_x": 0.5,
"top_left_y": 0.6,
"bottom_right_x": 0.7,
"bottom_right_y": 0.8
},
"title": "Totals",
"description": "Invoice totals"
}
}
]
}]
});
let normalized = CohereParseConfig
.transform_ocr_response(
"parse-v5.0",
&serde_json::to_vec(&payload).unwrap(),
response_format,
)
.unwrap();
assert_eq!(
normalized.pages[0].extra_fields["blocks"],
payload["pages"][0]["blocks"]
);
assert_eq!(normalized.pages[0].markdown, "");
assert_eq!(normalized.pages[0].index, 0);
assert_eq!(
normalized.usage_info.as_ref().unwrap().pages_processed,
Some(1)
);
match response_format {
crate::ocr::types::OcrResponseFormat::Litellm => {
assert!(normalized.provider_native_response.is_none());
}
crate::ocr::types::OcrResponseFormat::Native => {
assert_eq!(
normalized.provider_native_response.as_ref(),
payload.as_object()
);
}
}
assert_eq!(
normalized.into_json()["pages"][0]["blocks"],
payload["pages"][0]["blocks"]
);
}
#[rstest]
#[case::document_url(json!({"type":"document_url","document_url":"https://example.com/a.pdf"}))]
#[case::empty_image_url(json!({"type":"image_url","image_url":""}))]
#[case::pdf_data_uri(json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}))]
fn request_requires_image(#[case] value: Value) {
assert!(matches!(
validate_document(&serde_json::from_value(value).unwrap()),
Err(crate::ocr::Error::CohereImageOnly)
));
}
#[rstest]
#[case::markdown("markdown", true)]
#[case::blocks("blocks", true)]
#[case::unsupported("html", false)]
fn request_requires_supported_output_format(#[case] format: &str, #[case] valid: bool) {
assert_eq!(
serde_json::from_value::<CohereOptions>(json!({"output_format":format})).is_ok(),
valid
);
}
#[test]
fn request_defaults_to_markdown() {
let request = CohereParseConfig
.transform_ocr_request(
"parse-v5.0",
serde_json::from_value(json!({
"type":"image_url",
"image_url":"https://example.com/image.png"
}))
.unwrap(),
&serde_json::from_value(json!({})).unwrap(),
&[],
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap()["output_format"],
"markdown"
);
}
#[rstest]
#[case::base("")]
#[case::version("/v2")]
#[case::complete("/v2/parse")]
fn completes_provider_urls_without_duplicate_paths_and_preserves_queries(#[case] suffix: &str) {
assert_eq!(
CohereParseConfig
.build_ocr_url(&format!("https://example.com{suffix}?tenant=a"))
.unwrap(),
"https://example.com/v2/parse?tenant=a"
);
}
#[rstest]
#[case::relative("relative/path")]
#[case::unsupported_scheme("ftp://example.com")]
fn rejects_invalid_urls(#[case] api_base: &str) {
assert!(CohereParseConfig.build_ocr_url(api_base).is_err());
}
#[test]
fn rejects_blank_keys() {
assert!(matches!(
CohereParseConfig.resolve_headers(
&OcrConnection {
api_key: Some(" ".into()),
..Default::default()
},
&|_| None,
),
Err(crate::ocr::Error::Auth(_))
));
}
}

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

@ -0,0 +1 @@
pub(crate) mod ocr;

1
litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

633
litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs Normal file
View file

@ -0,0 +1,633 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::call_arguments::CallArguments;
use crate::constants::MISTRAL_OCR_API_BASE;
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response};
use crate::ocr::OcrClient;
use crate::ocr::prepare::credential_env;
use crate::ocr::types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo,
PreparedOcrRequest,
};
use crate::params::OpaqueParams;
use crate::url_utils::ApiUrl;
const MISTRAL_OCR_API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY";
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct MistralOcrRequest {
pub model: String,
pub document: OcrDocument,
#[serde(flatten)]
pub params: OpaqueParams,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct MistralOcrResponse {
#[serde(default)]
pub pages: Vec<OcrPage>,
#[serde(
default,
deserialize_with = "serde_with::rust::double_option::deserialize"
)]
pub model: Option<Option<String>>,
pub document_annotation: Option<Value>,
pub usage_info: Option<OcrUsageInfo>,
#[serde(flatten)]
pub extra_fields: serde_json::Map<String, Value>,
}
#[derive(Clone, Debug, Default)]
pub(crate) struct MistralOcrConfig;
impl BaseOcrConfig for MistralOcrConfig {
type OcrParams = OpaqueParams;
type ProviderRequest = MistralOcrRequest;
type Environment = Vec<(String, String)>;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&[
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
"document_annotation_prompt",
"extract_header",
"extract_footer",
"table_format",
"confidence_scores_granularity",
"include_blocks",
"id",
]
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some(MISTRAL_OCR_API_KEY_ENV_VAR)
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<OpaqueParams, crate::ocr::Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
self.resolve_headers(&request.connection, &credential_env)
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
self.build_ocr_url(request.connection.api_base.as_deref())
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
_headers: &[(String, String)],
) -> Result<MistralOcrRequest, crate::ocr::Error> {
Ok(MistralOcrRequest {
model: model.to_string(),
document,
params: optional_params.clone(),
})
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
}
impl MistralOcrConfig {
fn resolve_headers(
&self,
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, crate::ocr::Error> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
self.get_api_key_env_var()
.and_then(env_lookup)
.filter(|key| !key.trim().is_empty())
})
.ok_or(litellm_auth::Error::MissingApiKey {
provider: "Mistral",
environment_variable: MISTRAL_OCR_API_KEY_ENV_VAR,
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn build_ocr_url(&self, api_base: Option<&str>) -> Result<String, crate::ocr::Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(MISTRAL_OCR_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
}
pub(crate) fn normalize_response(
model: &str,
response: MistralOcrResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let model = match response.model {
Some(Some(model)) => model,
Some(None) => {
return Err(crate::ocr::Error::ResponseField {
path: "model".into(),
});
}
None => model.to_string(),
};
Ok(LiteLLMOcrResponse {
extra_fields: response.extra_fields,
document_annotation: response.document_annotation,
usage_info: response.usage_info,
..LiteLLMOcrResponse::new(model, response.pages)
})
}
#[cfg(test)]
mod tests {
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::*;
#[fixture]
fn document() -> OcrDocument {
serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap()
}
#[fixture]
fn connection(
#[default(None)] api_key: Option<&str>,
#[default(vec![])] extra_headers: Vec<(String, String)>,
) -> OcrConnection {
OcrConnection {
api_key: api_key.map(str::to_string),
extra_headers,
..OcrConnection::default()
}
}
#[test]
fn explicit_null_model_does_not_use_the_missing_model_default() {
let response = serde_json::from_value(json!({"model":null})).unwrap();
assert!(matches!(
normalize_response("fallback", response).unwrap_err(),
crate::ocr::Error::ResponseField { path } if path == "model"
));
}
#[rstest]
#[case::non_object_page(json!({"pages":[42]}), "pages[0]")]
#[case::missing_markdown(json!({"pages":[{"index":0}]}), "pages[0]")]
#[case::non_string_markdown(
json!({"pages":[{"index":0,"markdown":42}]}),
"pages[0].markdown"
)]
#[case::non_object_image(
json!({"pages":[{"index":0,"markdown":"","images":[42]}]}),
"pages[0].images[0]"
)]
#[case::fractional_width(
json!({"pages":[{"index":0,"markdown":"","dimensions":{"width":1.5}}]}),
"pages[0].dimensions.width"
)]
#[case::invalid_page_count(
json!({"usage_info":{"pages_processed":"bad"}}),
"usage_info.pages_processed"
)]
fn response_validates_normalized_shapes_at_the_provider_boundary(
#[case] payload: Value,
#[case] path: &str,
) {
let error = crate::ocr::json::decode_response::<MistralOcrResponse>(
&serde_json::to_vec(&payload).unwrap(),
false,
)
.unwrap_err();
assert!(matches!(
error,
crate::ocr::Error::ResponseField { path: actual } if actual == path
));
}
#[test]
fn response_normalizes_python_numeric_inputs_and_shared_defaults() {
let response = serde_json::from_value(json!({
"pages":[{"index":"2","markdown":"text","dimensions":{"width":1.0},"extension":false}],
"usage_info":{"pages_processed":true,"credits":"1.5","custom":0},
"extra":"ignored"
}))
.unwrap();
let response = normalize_response("model", response).unwrap();
assert_eq!(response.pages[0].index, 2);
assert_eq!(
response.pages[0].dimensions.as_ref().unwrap().width,
Some(1)
);
assert_eq!(
response.usage_info.as_ref().unwrap().pages_processed,
Some(1)
);
assert_eq!(response.usage_info.as_ref().unwrap().credits, Some(1.5));
let serialized = response.into_json();
assert_eq!(serialized["pages"][0]["extension"], false);
assert!(serialized["pages"][0]["images"].is_null());
assert!(serialized["usage_info"]["doc_size_bytes"].is_null());
assert_eq!(serialized["usage_info"]["custom"], 0);
assert!(serialized["content"].is_null());
assert_eq!(serialized["extra"], "ignored");
}
#[test]
fn map_ocr_params_selects_known_fields_without_changing_arguments() {
let input =
serde_json::from_value(json!({"pages":null,"extract_header":false,"unknown":true}))
.unwrap();
let params = MistralOcrConfig.map_ocr_params(&input, "model").unwrap();
assert_eq!(
serde_json::to_value(params).unwrap(),
json!({"pages":null,"extract_header":false})
);
assert_eq!(input["unknown"], true);
assert_eq!(input.get("pages"), Some(&Value::Null));
}
#[rstest]
fn request_transform_uses_already_mapped_params_without_filtering_again(document: OcrDocument) {
let params = serde_json::from_value(json!({"extension":{"nested":null}})).unwrap();
let body = MistralOcrConfig
.transform_ocr_request("model", document, &params, &[])
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap()["extension"],
json!({"nested":null})
);
}
#[test]
fn raw_response_transform_keeps_native_payload_separate_from_typed_normalization() {
let raw = br#"{"pages":[{"index":"2","markdown":"text"}],"provider_extension":false}"#;
let response = MistralOcrConfig
.transform_ocr_response("model", raw, crate::ocr::types::OcrResponseFormat::Native)
.unwrap();
assert_eq!(response.pages[0].index, 2);
let native = response.provider_native_response.unwrap();
assert_eq!(native["pages"][0]["index"], "2");
assert_eq!(native["provider_extension"], false);
assert_eq!(response.extra_fields["provider_extension"], false);
}
#[rstest]
fn raw_response_transform_rejects_invalid_page(
#[values(OcrResponseFormat::Litellm, OcrResponseFormat::Native)]
request_format: OcrResponseFormat,
) {
assert!(
MistralOcrConfig
.transform_ocr_response("model", br#"{"pages":[{"index":0}]}"#, request_format)
.is_err()
);
}
fn mapped_params(value: Value) -> Value {
let params = serde_json::from_value(value).unwrap();
serde_json::to_value(MistralOcrConfig.map_ocr_params(&params, "model").unwrap()).unwrap()
}
#[rstest]
fn extract_header_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn extract_footer_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_footer":false}))["extract_footer"],
false
);
}
#[rstest]
fn existing_ocr_params_remain_supported() {
let mapped = mapped_params(json!({
"pages":[0,2],
"include_image_base64":true,
"image_limit":2,
"image_min_size":100,
"bbox_annotation_format":{"type":"json_schema"},
"document_annotation_format":{"type":"json_schema"}
}));
assert_eq!(mapped["pages"], json!([0, 2]));
assert_eq!(mapped["include_image_base64"], true);
assert_eq!(mapped["image_limit"], 2);
assert_eq!(mapped["image_min_size"], 100);
assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema");
assert_eq!(mapped["document_annotation_format"]["type"], "json_schema");
}
#[rstest]
fn map_ocr_params_forwards_extract_header() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_footer() {
assert_eq!(
mapped_params(json!({"extract_footer":true}))["extract_footer"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_header_and_footer() {
let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false}));
assert_eq!(mapped["extract_header"], true);
assert_eq!(mapped["extract_footer"], false);
}
#[rstest]
fn map_ocr_params_excludes_extensions_from_the_provider_options() {
let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"}));
assert_eq!(mapped["extract_header"], true);
assert!(mapped.get("unsupported_param").is_none());
}
#[rstest]
fn map_ocr_params_preserves_unvalidated_values_and_explicit_null() {
let mapped = mapped_params(json!({
"pages":{"future":"shape"},
"include_image_base64":null
}));
assert_eq!(mapped["pages"], json!({"future":"shape"}));
assert!(mapped.get("include_image_base64").unwrap().is_null());
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("block"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("table_format", json!("markdown"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("page"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("pages", json!([0, 2]))]
#[case("pages", json!("0,2-4"))]
#[case("pages", Value::Null)]
#[case("include_image_base64", json!(true))]
#[case("include_image_base64", json!(false))]
#[case("image_limit", json!(2))]
#[case("image_min_size", json!(100))]
#[case("bbox_annotation_format", json!({"type":"json_schema"}))]
#[case("document_annotation_format", json!({"type":"json_schema"}))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("extract_header", json!(true))]
#[case("extract_footer", json!(false))]
#[case("table_format", json!("html"))]
#[case("table_format", json!("markdown"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("confidence_scores_granularity", json!("page"))]
#[case("confidence_scores_granularity", json!("block"))]
#[case("include_blocks", json!(true))]
#[case("include_blocks", json!(false))]
#[case("id", json!("req-123"))]
fn request_mapping_preserves_supplied_options(
document: OcrDocument,
#[case] name: &str,
#[case] value: Value,
) {
let arguments = serde_json::from_value(json!({name: value.clone()})).unwrap();
let params = MistralOcrConfig
.map_ocr_params(&arguments, "model")
.unwrap();
let result = serde_json::to_value(
MistralOcrConfig
.transform_ocr_request("model", document.clone(), &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(
result,
json!({"model":"model", "document":document, name:value})
);
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("id", json!("req-123"))]
#[case("extract_header", json!(true))]
#[case("include_blocks", json!(true))]
#[case("pages", json!([0,1]))]
fn transform_ocr_request_includes_each_optional_param(
document: OcrDocument,
#[case] name: &str,
#[case] value: Value,
) {
let params: OpaqueParams = serde_json::from_value(json!({name:value.clone()})).unwrap();
let result = serde_json::to_value(
MistralOcrConfig
.transform_ocr_request("mistral-ocr-latest", document, &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(result[name], value);
assert_eq!(result["model"], "mistral-ocr-latest");
}
#[rstest]
fn transform_ocr_request_includes_multiple_new_params(document: OcrDocument) {
let params: OpaqueParams = serde_json::from_value(json!({
"table_format":"html",
"confidence_scores_granularity":"page",
"extract_header":true
}))
.unwrap();
let result = serde_json::to_value(
MistralOcrConfig
.transform_ocr_request("mistral-ocr-latest", document, &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(result["table_format"], "html");
assert_eq!(result["confidence_scores_granularity"], "page");
assert_eq!(result["extract_header"], true);
}
#[rstest]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let payload = json!({
"pages":[{
"index":0,
"markdown":"hello",
"images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}],
"dimensions":{"width":612,"height":792,"dpi":72},
"blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}],
"confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97}
}],
"model":"returned-model",
"document_annotation":"{\"language\":\"en\"}",
"usage_info":{"pages_processed":1}
});
let response: MistralOcrResponse = serde_json::from_value(payload.clone()).unwrap();
let result = normalize_response("model", response).unwrap().into_json();
assert_eq!(result["pages"][0]["blocks"], payload["pages"][0]["blocks"]);
assert_eq!(
result["pages"][0]["confidence_scores"],
payload["pages"][0]["confidence_scores"]
);
assert_eq!(result["pages"][0]["images"][0]["id"], "img-0");
assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72);
assert_eq!(result["model"], "returned-model");
assert_eq!(result["document_annotation"], "{\"language\":\"en\"}");
assert_eq!(result["usage_info"]["pages_processed"], 1);
}
#[rstest]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let page = json!({
"index":0,
"markdown":"table page",
"tables":[{"rows":2,"cols":3}],
"hyperlinks":["https://example.com"],
"header":"header",
"footer":"footer"
});
let response: MistralOcrResponse =
serde_json::from_value(json!({"pages":[page.clone()]})).unwrap();
let result = normalize_response("model", response).unwrap().into_json();
assert_eq!(result["pages"][0]["tables"], page["tables"]);
assert_eq!(result["pages"][0]["hyperlinks"], page["hyperlinks"]);
assert_eq!(result["pages"][0]["header"], page["header"]);
assert_eq!(result["pages"][0]["footer"], page["footer"]);
assert!(result["pages"][0]["images"].is_null());
assert!(result["pages"][0]["dimensions"].is_null());
}
#[rstest]
#[case::default_base(None, "https://api.mistral.ai/v1/ocr")]
#[case::versioned_base(
Some("https://example.com/v1?tenant=a"),
"https://example.com/v1/ocr?tenant=a"
)]
#[case::complete_endpoint(
Some("https://example.com/v1/ocr?tenant=a"),
"https://example.com/v1/ocr?tenant=a"
)]
fn complete_url_defaults_and_dedupes_v1(
#[case] api_base: Option<&str>,
#[case] expected: &str,
) {
assert_eq!(MistralOcrConfig.build_ocr_url(api_base).unwrap(), expected);
}
#[rstest]
#[case::explicit_key(Some("explicit"), "Bearer explicit")]
#[case::environment_fallback(None, "Bearer environment")]
fn environment_prefers_explicit_key_then_environment(
#[case] _api_key: Option<&str>,
#[case] expected: &str,
#[with(_api_key)] connection: OcrConnection,
) {
assert_eq!(
MistralOcrConfig
.resolve_headers(&connection, &|_| Some("environment".into()))
.unwrap()[0],
("Authorization".into(), expected.into())
);
}
#[rstest]
fn environment_preserves_forwarded_authorization(
#[with(None, vec![("authorization".into(), "Bearer forwarded".into())])]
connection: OcrConnection,
) {
assert_eq!(
MistralOcrConfig
.resolve_headers(&connection, &|_| None)
.unwrap(),
connection.extra_headers
);
}
#[rstest]
fn environment_rejects_missing_key(connection: OcrConnection) {
assert!(matches!(
MistralOcrConfig.resolve_headers(&connection, &|_| None),
Err(crate::ocr::Error::Auth(
litellm_auth::Error::MissingApiKey {
provider: "Mistral",
environment_variable: MISTRAL_OCR_API_KEY_ENV_VAR,
}
))
));
}
}

9
litellm-rust/crates/core/src/llms/mod.rs Normal file
View file

@ -0,0 +1,9 @@
pub mod anthropic;
pub mod azure_ai;
pub mod base_llm;
pub mod bedrock;
pub(crate) mod cohere;
pub(crate) mod mistral;
pub mod openai;
pub(crate) mod reducto;
pub(crate) mod vertex_ai;

0
litellm-rust/crates/core/src/providers/openai/mod.rs → litellm-rust/crates/core/src/llms/openai/mod.rs
View file

0
litellm-rust/crates/core/src/providers/bedrock/chat_completions/mod.rs → litellm-rust/crates/core/src/llms/openai/responses/mod.rs
View file

6
litellm-rust/crates/core/src/providers/openai/responses/transformation.rs → litellm-rust/crates/core/src/llms/openai/responses/transformation.rs
View file

@ -2,11 +2,11 @@ use crate::responses::Error;
use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult};
use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model};
pub struct OpenAIResponsesWsConfig;
pub struct OpenAiResponsesApiConfig;
pub const OPENAI_RESPONSES_WS_CONFIG: OpenAIResponsesWsConfig = OpenAIResponsesWsConfig;
pub const OPENAI_RESPONSES_WS_CONFIG: OpenAiResponsesApiConfig = OpenAiResponsesApiConfig;
impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig {
impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig {
fn supports_native_websocket(&self) -> bool {
true
}

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

@ -0,0 +1 @@
pub(crate) mod ocr;

1
litellm-rust/crates/core/src/llms/reducto/ocr/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

1019
litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -0,0 +1 @@
pub(crate) mod ocr;

10
litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs Normal file
View file

@ -0,0 +1,10 @@
use litellm_auth::InputSource;
use crate::ocr::types::OcrConnection;
pub(super) fn validate_destination(connection: &OcrConnection) -> Result<(), crate::ocr::Error> {
if connection.api_base.is_some() && connection.api_base_source == InputSource::Request {
return Err(litellm_auth::Error::RequestVertexCredentialDestination.into());
}
Ok(())
}

701
litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs Normal file
View file

@ -0,0 +1,701 @@
use litellm_auth_gcp::{self as vertex, VertexConfig};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::VertexAiOcrConfig;
use crate::call_arguments::CallArguments;
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
use crate::ocr::OcrClient;
use crate::ocr::prepare::credential_env;
use crate::ocr::types::{
LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, OcrUsageInfo,
PreparedOcrRequest,
};
use crate::params::OpaqueParams;
use crate::url_utils::ApiUrl;
const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
const MODEL_PREFIX: &str = "deepseek-ai/";
const DEFAULT_LOCATION: &str = "us-central1";
const DEEPSEEK_OCR_PARAMS: &[&str] = &["stream", "temperature", "max_tokens", "top_p", "n", "stop"];
pub(crate) type DeepSeekOcrParams = OpaqueParams;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrRequest {
pub model: String,
pub messages: Vec<DeepSeekOcrMessage>,
#[serde(flatten)]
pub params: OpaqueParams,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrMessage {
pub role: UserRole,
pub content: Vec<DeepSeekDocument>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
pub(crate) enum DeepSeekDocument {
#[serde(rename = "image_url")]
ImageUrl { image_url: String },
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum UserRole {
User,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekOcrResponse {
#[serde(default)]
choices: Vec<DeepSeekChoice>,
#[serde(default = "empty_object")]
usage: Value,
}
#[derive(Clone, Debug, Deserialize)]
struct DeepSeekChoice {
#[serde(default)]
message: DeepSeekResponseMessage,
}
#[derive(Clone, Debug, Default, Deserialize)]
struct DeepSeekResponseMessage {
content: Option<DeepSeekContent>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
enum DeepSeekContent {
Text(String),
Object(Map<String, Value>),
}
#[serde_with::serde_as]
#[derive(Deserialize)]
struct DeepSeekPage {
#[serde(default)]
#[serde_as(deserialize_as = "crate::serde_compat::LaxI64")]
index: i64,
#[serde(default)]
markdown: String,
images: Option<Vec<OcrPageImage>>,
dimensions: Option<OcrPageDimensions>,
}
#[derive(Clone, Debug)]
pub(crate) struct VertexAIDeepSeekOCRConfig;
impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
type OcrParams = DeepSeekOcrParams;
type ProviderRequest = DeepSeekOcrRequest;
type Environment = vertex::VertexEnvironment;
fn get_api_key_env_var(&self) -> Option<&'static str> {
VertexAiOcrConfig.get_api_key_env_var()
}
fn map_ocr_params(
&self,
_arguments: &CallArguments,
_model: &str,
) -> Result<DeepSeekOcrParams, crate::ocr::Error> {
Ok(DeepSeekOcrParams::default())
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
VertexAiOcrConfig
.validate_environment(request, client)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
self.get_complete_url(
request.connection.api_base.as_deref(),
&environment.project_id,
&location,
)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &DeepSeekOcrParams,
headers: &[(String, String)],
_context: OcrRequestContext<'_>,
) -> Result<DeepSeekOcrRequest, crate::ocr::Error> {
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
crate::llms::base_llm::ocr::transformation::decode_and_normalize_response(
model,
raw_response,
request_format,
normalize_response,
)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &DeepSeekOcrParams,
_headers: &[(String, String)],
) -> Result<DeepSeekOcrRequest, crate::ocr::Error> {
if document.source().is_empty() {
return Err(crate::ocr::Error::MissingDocumentUrl);
}
Ok(DeepSeekOcrRequest {
model: provider_model(model)?,
messages: vec![DeepSeekOcrMessage {
role: UserRole::User,
content: vec![DeepSeekDocument::ImageUrl {
image_url: document.source().to_string(),
}],
}],
params: optional_params
.iter()
.filter(|(name, _)| DEEPSEEK_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
})
}
}
pub(crate) fn normalize_response(
model: &str,
response: DeepSeekOcrResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
let content = response
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content)
.ok_or(crate::ocr::Error::EmptyContent)?;
let (ocr_data, fallback_markdown) = match content {
DeepSeekContent::Text(text) if text.is_empty() => {
return Err(crate::ocr::Error::EmptyContent);
}
DeepSeekContent::Text(text) => {
let parsed = text
.trim_start()
.starts_with('{')
.then(|| serde_json::from_str::<Map<String, Value>>(&text).ok())
.flatten();
(parsed.unwrap_or_default(), text)
}
DeepSeekContent::Object(data) if data.is_empty() => {
return Err(crate::ocr::Error::EmptyContent);
}
DeepSeekContent::Object(data) => {
let fallback = if data.contains_key("pages") {
String::new()
} else {
let mut output = Vec::new();
data.serialize(&mut serde_json::Serializer::with_formatter(
&mut output,
PythonJsonFormatter,
))
.map_err(|_| response_field("content"))?;
String::from_utf8(output).map_err(|_| response_field("content"))?
};
(data, fallback)
}
};
let has_pages = ocr_data.contains_key("pages");
let pages = match ocr_data.get("pages") {
Some(Value::Array(pages)) => pages
.iter()
.enumerate()
.filter(|(_, page)| page.is_object())
.map(|(position, page)| {
let page: DeepSeekPage = crate::ocr::json::decode_response_value(
page.clone(),
&format!("choices[0].message.content.pages[{position}]"),
)?;
Ok(OcrPage {
index: page.index,
markdown: page.markdown,
images: page.images,
dimensions: page.dimensions,
..Default::default()
})
})
.collect::<Result<Vec<_>, crate::ocr::Error>>()?,
Some(_) => return Err(response_field("pages")),
None => Vec::new(),
};
let usage = ocr_data
.get("usage_info")
.or_else(|| (!has_pages).then_some(&response.usage));
let usage_info: Option<OcrUsageInfo> = usage
.filter(|usage| usage.is_object())
.map(|usage| crate::ocr::json::decode_response_value(usage.clone(), "usage_info"))
.transpose()?;
let model = match ocr_data.get("model") {
Some(Value::String(model)) => model.clone(),
Some(_) => return Err(response_field("model")),
None => model.to_string(),
};
Ok(LiteLLMOcrResponse {
extra_fields: ocr_data
.iter()
.filter(|(name, _)| {
!matches!(
name.as_str(),
"pages"
| "model"
| "document_annotation"
| "usage_info"
| "object"
| "content"
| "tables"
| "keyValuePairs"
| "provider_native_response"
)
})
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
document_annotation: has_pages
.then(|| ocr_data.get("document_annotation").cloned())
.flatten(),
usage_info,
..LiteLLMOcrResponse::new(
model,
if pages.is_empty() {
vec![OcrPage {
markdown: fallback_markdown,
..Default::default()
}]
} else {
pages
},
)
})
}
fn empty_object() -> Value {
Value::Object(Map::new())
}
struct PythonJsonFormatter;
impl serde_json::ser::Formatter for PythonJsonFormatter {
fn begin_array_value<W: std::io::Write + ?Sized>(
&mut self,
writer: &mut W,
first: bool,
) -> std::io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_key<W: std::io::Write + ?Sized>(
&mut self,
writer: &mut W,
first: bool,
) -> std::io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_value<W: std::io::Write + ?Sized>(
&mut self,
writer: &mut W,
) -> std::io::Result<()> {
writer.write_all(b": ")
}
fn write_string_fragment<W: std::io::Write + ?Sized>(
&mut self,
writer: &mut W,
fragment: &str,
) -> std::io::Result<()> {
for character in fragment.chars() {
if character.is_ascii() && character != '\u{7f}' {
writer.write_all(&[character as u8])?;
} else {
for unit in character.encode_utf16(&mut [0; 2]) {
write!(writer, "\\u{unit:04x}")?;
}
}
}
Ok(())
}
}
fn response_field(field: &str) -> crate::ocr::Error {
crate::ocr::Error::ResponseField {
path: format!("choices[0].message.content.{field}"),
}
}
pub(crate) fn provider_model(model: &str) -> Result<String, crate::ocr::Error> {
let local_model = model.trim_start_matches(MODEL_PREFIX);
if local_model.is_empty() {
return Err(crate::ocr::Error::RequestField {
path: "model".into(),
});
}
Ok(format!("{MODEL_PREFIX}{local_model}"))
}
impl VertexAIDeepSeekOCRConfig {
fn get_complete_url(
&self,
api_base: Option<&str>,
project: &str,
location: &str,
) -> Result<String, crate::ocr::Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(DEFAULT_API_BASE);
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"endpoints",
"openapi",
"chat",
"completions",
])
})
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
}
#[cfg(test)]
mod tests {
use serde_json::{Value, json};
use super::{
DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, normalize_response,
provider_model,
};
#[test]
fn unconsumed_options_remain_available_for_body_composition() {
use serde_json::json;
use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
let arguments =
serde_json::from_value(json!({"temperature":0.5,"extension":null})).unwrap();
assert_eq!(
serde_json::to_value(
VertexAIDeepSeekOCRConfig
.map_ocr_params(&arguments, "deepseek-ocr")
.unwrap()
)
.unwrap(),
json!({})
);
assert_eq!(
crate::call_arguments::compose_body(&arguments, &json!({"model":"deepseek-ocr"}), &[])
.unwrap(),
json!({"model":"deepseek-ocr","temperature":0.5,"extension":null})
);
}
#[test]
fn config_owns_model_namespace_and_endpoint() {
assert_eq!(
provider_model("deepseek-ocr-maas").unwrap(),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
provider_model("deepseek-ai/deepseek-ocr-maas").unwrap(),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
VertexAIDeepSeekOCRConfig
.get_complete_url(None, "proj-1", "europe-west4")
.unwrap(),
"https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions"
);
}
use rstest::rstest;
use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
use crate::ocr::types::OcrDocument;
fn document() -> OcrDocument {
serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap()
}
#[rstest]
#[case("stream", json!(true))]
#[case("temperature", json!(0.1))]
#[case("max_tokens", json!(1024))]
#[case("top_p", json!(0.9))]
#[case("n", json!(2))]
#[case("stop", json!("done"))]
#[case("stop", json!(["done", "stop"]))]
#[case("temperature", json!(null))]
fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
let params: DeepSeekOcrParams =
serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap();
let result = serde_json::to_value(
VertexAIDeepSeekOCRConfig
.transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), &params, &[])
.unwrap(),
)
.unwrap();
assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(
result["messages"][0]["content"][0],
json!({"type":"image_url","image_url":"gs://bucket/a.png"})
);
assert_eq!(result[name], value);
assert!(result.get("ignored").is_none());
}
#[rstest]
#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))]
#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))]
fn request_maps_both_document_types_to_image_content(#[case] document: Value) {
let source = document
.get("image_url")
.or_else(|| document.get("document_url"))
.unwrap()
.clone();
let request = VertexAIDeepSeekOCRConfig
.transform_ocr_request(
"deepseek-ai/deepseek-ocr-maas",
serde_json::from_value(document).unwrap(),
&DeepSeekOcrParams::default(),
&[],
)
.unwrap();
let result = serde_json::to_value(request).unwrap();
assert_eq!(
result["messages"][0]["content"][0],
json!({"type":"image_url","image_url":source})
);
}
#[rstest]
#[case(json!("# hello"), "# hello")]
#[case(json!("{broken"), "{broken")]
#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")]
#[case(json!({"pages":[]}), "")]
#[case(json!("[]"), "[]")]
#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")]
#[case(json!({"pages":[{"markdown":"object"}]}), "object")]
fn response_transform_handles_text_json_and_objects(
#[case] content: Value,
#[case] expected: &str,
) {
let has_pages = content
.as_object()
.is_some_and(|data| data.contains_key("pages"))
|| content
.as_str()
.is_some_and(|text| text.contains("\"pages\""));
let response: DeepSeekOcrResponse = serde_json::from_value(
json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}),
)
.unwrap();
let result = normalize_response("model", response).unwrap().into_json();
assert_eq!(result["pages"][0]["markdown"], expected);
assert_eq!(result["pages"][0]["index"], 0);
if has_pages {
assert!(result["usage_info"].is_null());
} else {
assert_eq!(result["usage_info"]["prompt_tokens"], 1);
}
}
#[test]
fn structured_result_maps_pages_usage_model_and_annotation() {
let response: DeepSeekOcrResponse = serde_json::from_value(json!({
"choices":[{"message":{"content":{
"pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}],
"model":"provider-model",
"usage_info":{"pages_processed":1},
"document_annotation":{"language":"en"},
"future":"kept"
}}}]
}))
.unwrap();
let result = normalize_response("requested", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["index"], 2);
assert_eq!(result["pages"][0]["images"][0]["id"], "one");
assert_eq!(result["model"], "provider-model");
assert_eq!(result["usage_info"]["pages_processed"], 1);
assert_eq!(result["document_annotation"]["language"], "en");
assert_eq!(result["future"], "kept");
}
#[test]
fn response_transform_rejects_missing_empty_and_malformed_content() {
for value in [
json!({"choices":[]}),
json!({"choices":[{"message":{"content":{}}}]}),
json!({"choices":[{"message":{"content":""}}]}),
json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}),
json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}),
] {
let result = serde_json::from_value::<DeepSeekOcrResponse>(value)
.map_err(|_| ())
.and_then(|response| normalize_response("model", response).map_err(|_| ()));
assert!(result.is_err());
}
}
#[test]
fn structured_content_preserves_usage_presence_and_shared_page_defaults() {
for (usage, expected) in [(json!(null), None), (json!({"pages_processed":2}), Some(2))] {
let response = serde_json::from_value(json!({
"choices":[{"message":{"content":{
"pages":[42, {"index":"2", "images":[{"id":"kept"}], "ignored":true}],
"usage_info":usage
}}}],
"usage":{"pages_processed":99}
}))
.unwrap();
let normalized = normalize_response("model", response).unwrap();
assert_eq!(normalized.pages.len(), 1);
assert_eq!(normalized.pages[0].index, 2);
assert_eq!(normalized.pages[0].markdown, "");
assert!(normalized.pages[0].extra_fields.is_empty());
assert_eq!(
normalized
.usage_info
.and_then(|usage| usage.pages_processed),
expected
);
}
}
use litellm_auth::InputSource;
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[tokio::test]
async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"choices":[{"message":{"content":"recognized"}}],
"usage":{"prompt_tokens":1}
}))])
.await;
let request = wire_request(
"vertex_ai/deepseek-ocr-maas",
&base,
json!({
"vertex_project":"project-1",
"vertex_location":"europe-west4",
"temperature":0.1,
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
);
let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf");
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.pages[0].markdown, "recognized");
assert_eq!(
response.usage_info.unwrap().extra_fields["prompt_tokens"],
1
);
let requests = seen.lock().unwrap();
assert!(requests[0].starts_with(
"POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions "
));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
let body = request_body(&requests[0]);
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(body["temperature"], 0.1);
assert_eq!(body["future_ocr_option"], true);
assert_eq!(body["provider_option"], "value");
assert!(body.get("vertex_project").is_none());
assert!(body.get("extra_body").is_none());
assert_eq!(
body["messages"][0]["content"][0],
json!({"type":"image_url","image_url":"gs://bucket/document.pdf"})
);
}
#[test]
fn host_registration_selects_deepseek_without_affecting_mistral() {
assert!(crate::ocr::is_supported_request(
"deepseek-ocr-maas",
Some("vertex_ai")
));
assert!(crate::ocr::is_supported_request(
"mistral-ocr-maas",
Some("vertex_ai")
));
}
#[tokio::test]
async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
let mut request = wire_request(
"vertex_ai/deepseek-ocr-maas",
"https://caller.example",
json!({"vertex_project":"project-1"}),
);
request.credentials.api_base = Some(litellm_auth::Sourced::new(
"https://caller.example".into(),
InputSource::Request,
));
let error = perform_ocr(request).await.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Vertex AI endpoint")
);
}
}

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

@ -0,0 +1,3 @@
pub(crate) mod common_utils;
pub(crate) mod deepseek_transformation;
pub(crate) mod transformation;

407
litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs Normal file
View file

@ -0,0 +1,407 @@
use litellm_auth_gcp::{self as vertex, VertexConfig};
use serde_json::Value;
use super::common_utils::validate_destination;
use crate::call_arguments::CallArguments;
use crate::llms::base_llm::ocr::transformation::{
BaseOcrConfig, OcrEnvironment, OcrRequestContext,
};
use crate::llms::mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest};
use crate::ocr::OcrClient;
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::prepare::credential_env;
use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest};
use crate::params::OpaqueParams;
use crate::url_utils::ApiUrl;
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug, Default)]
pub(crate) struct VertexAiOcrConfig;
impl BaseOcrConfig for VertexAiOcrConfig {
type OcrParams = OpaqueParams;
type ProviderRequest = MistralOcrRequest;
type Environment = vertex::VertexEnvironment;
fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] {
MistralOcrConfig.get_supported_ocr_params(model)
}
fn get_api_key_env_var(&self) -> Option<&'static str> {
Some("VERTEX_AI_API_KEY")
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<OpaqueParams, crate::ocr::Error> {
MistralOcrConfig.map_ocr_params(non_default_params, model)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?;
self.resolve_environment(&request.connection, &config, client)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
self.build_ocr_url(
request.connection.api_base.as_deref(),
&environment.project_id,
&location,
&request.model,
)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
headers: &[(String, String)],
) -> Result<MistralOcrRequest, crate::ocr::Error> {
MistralOcrConfig.transform_ocr_request(model, document, optional_params, headers)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &OpaqueParams,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<MistralOcrRequest, crate::ocr::Error> {
let document = inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
MistralOcrConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
validate_inline_document(&crate::ocr::prepare::body_document(body)?)
}
}
impl OcrEnvironment for vertex::VertexEnvironment {
fn headers(&self) -> &[(String, String)] {
&self.headers
}
}
impl VertexAiOcrConfig {
async fn resolve_environment(
&self,
connection: &OcrConnection,
config: &VertexConfig,
client: &OcrClient,
) -> Result<vertex::VertexEnvironment, crate::ocr::Error> {
validate_destination(connection)?;
client
.vertex_auth()
.validate_environment(
connection.extra_headers.clone(),
connection.api_key.as_deref(),
config,
&credential_env,
)
.await
.map_err(crate::ocr::Error::from)
}
fn build_ocr_url(
&self,
api_base: Option<&str>,
project: &str,
location: &str,
model: &str,
) -> Result<String, crate::ocr::Error> {
validate_location(location)?;
let default_base = format!("https://{location}-aiplatform.googleapis.com");
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(&default_base);
let prediction = format!("{model}:rawPredict");
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"publishers",
"mistralai",
"models",
&prediction,
])
})
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
path: "api_base".into(),
})
}
}
fn validate_location(location: &str) -> Result<(), crate::ocr::Error> {
let valid = !location.is_empty()
&& location
.bytes()
.all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-')
&& location
.as_bytes()
.first()
.is_some_and(u8::is_ascii_alphanumeric)
&& location
.as_bytes()
.last()
.is_some_and(u8::is_ascii_alphanumeric);
if valid {
return Ok(());
}
Err(crate::ocr::Error::RequestField {
path: "vertex_location".into(),
})
}
#[cfg(test)]
mod tests {
use super::VertexAiOcrConfig;
use rstest::rstest;
#[test]
fn endpoint_uses_location_project_and_model() {
assert_eq!(
VertexAiOcrConfig
.build_ocr_url(None, "proj-1", "europe-west4", "mistral-ocr-maas")
.unwrap(),
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[test]
fn endpoint_rejects_invalid_location() {
assert!(
VertexAiOcrConfig
.build_ocr_url(None, "proj-1", "attacker.example/path", "model")
.is_err()
);
}
use litellm_auth::InputSource;
use serde_json::{Value, json};
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[tokio::test]
async fn facade_executes_vertex_mistral_with_resolved_project_and_location() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"hello"}],
"usage_info":{"pages_processed":1}
}))])
.await;
let request = wire_request(
"vertex_ai/mistral-ocr-maas",
&base,
json!({
"vertex_project":"project-1",
"vertex_location":"europe-west4",
"extract_footer":true
}),
);
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.pages[0].markdown, "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with(
"POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict "
));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert_eq!(
request_body(&requests[0]),
json!({
"model":"mistral-ocr-maas",
"document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"extract_footer":true
})
);
}
#[tokio::test]
async fn supplied_authorization_is_forwarded_without_a_static_token() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut request = wire_request(
"vertex_ai/model",
&base,
json!({"vertex_project":"project-1"}),
);
request.credentials.api_key = None;
request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())];
perform_ocr(request).await.unwrap();
server.await.unwrap();
assert!(
seen.lock().unwrap()[0]
.to_ascii_lowercase()
.contains("authorization: bearer supplied")
);
}
#[tokio::test]
async fn invalid_credentials_fail_before_provider_http() {
let request = wire_request(
"vertex_ai/model",
"http://127.0.0.1:1",
json!({"vertex_credentials": true}),
);
let error = perform_ocr(request).await.unwrap_err();
assert!(error.to_string().contains("vertex_credentials"));
}
#[tokio::test]
async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
let mut request = wire_request(
"vertex_ai/mistral-ocr-maas",
"https://caller.example",
json!({"vertex_project":"project-1"}),
);
request.credentials.api_base = Some(litellm_auth::Sourced::new(
"https://caller.example".into(),
InputSource::Request,
));
let error = perform_ocr(request).await.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Vertex AI endpoint")
);
}
#[rstest]
#[case::mistral(false)]
#[case::vertex(true)]
#[tokio::test]
async fn configs_build_complete_requests_and_share_mistral_normalization(
#[case] use_vertex: bool,
) {
use std::time::Duration;
use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
use crate::llms::mistral::ocr::transformation::MistralOcrConfig;
use crate::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig;
use crate::ocr::test_support::ocr_client;
let client = ocr_client();
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "preserved"
});
let direct = wire_request(
"mistral/mistral-ocr-maas",
"https://mistral.test",
options.clone(),
);
let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options);
let direct = crate::ocr::prepare::prepare_request(
crate::ocr::test_support::resolved_request(direct),
);
let vertex = crate::ocr::prepare::prepare_request(
crate::ocr::test_support::resolved_request(vertex),
);
let direct_http = MistralOcrConfig
.prepare_request(&direct, &client)
.await
.unwrap();
let vertex_http = VertexAiOcrConfig
.prepare_request(&vertex, &client)
.await
.unwrap();
assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url().as_str(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
let http = if use_vertex {
&vertex_http
} 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!(
body,
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true,
"unknown": "preserved"
})
);
let payload = serde_json::to_vec(
&json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}),
)
.unwrap();
let direct_response = MistralOcrConfig
.transform_ocr_response(&direct.model, &payload, Default::default())
.unwrap()
.into_json();
let vertex_response = VertexAiOcrConfig
.transform_ocr_response(&vertex.model, &payload, Default::default())
.unwrap()
.into_json();
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
}

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

@ -279,11 +279,13 @@ impl Resolve for PublicDnsResolver {
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use super::*;
async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
.await

13
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -1,18 +1,17 @@
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use serde_json::{Map, Value};
use super::transformation::AnthropicMessagesProviderConfig;
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body};
use crate::llms::anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use crate::llms::azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
const HEADER_CONTEXT: &str = "messages";
pub(super) fn messages_provider_config(
provider: &str,
) -> Option<&'static dyn AnthropicMessagesProviderConfig> {
) -> Option<&'static dyn BaseAnthropicMessagesConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG),
"azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG),

42
litellm-rust/crates/core/src/messages/error.rs
View file

@ -2,16 +2,54 @@
pub enum Error {
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("routing error: {0}")]
Routing(String),
#[error("unsupported by the Rust messages route: {0}")]
Unsupported(&'static str),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
#[error("stream framing failed: {0}")]
StreamFraming(String),
#[error("Anthropic SSE frame has no data")]
MissingStreamData,
#[error("Anthropic stream event is invalid: {0}")]
InvalidStreamEvent(String),
#[error("Bedrock event payload is invalid: {0}")]
InvalidBedrockPayload(String),
#[error("Bedrock event payload has invalid base64: {0}")]
InvalidBedrockBase64(String),
}
impl Error {
pub fn is_request(&self) -> bool {
match self {
Self::InvalidProvider(_)
| Self::MissingField(_)
| Self::InvalidRequest(_)
| Self::Unsupported(_)
| Self::Headers(_) => true,
Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }),
_ => false,
}
}
pub fn is_response(&self) -> bool {
matches!(
self,
Self::InvalidResponse(_)
| Self::StreamFraming(_)
| Self::MissingStreamData
| Self::InvalidStreamEvent(_)
| Self::InvalidBedrockPayload(_)
| Self::InvalidBedrockBase64(_)
)
}
}

13
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -1,11 +1,10 @@
use super::Error;
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
use crate::http_utils::http_request;
use super::client::http_client;
use super::common_utils::truncate_error_body;
use super::prepare::prepare_provider_request;
use super::types::{AnthropicMessagesResponse, MessagesRequest};
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
use crate::http_utils::http_request;
pub(super) async fn execute_messages_provider_call(
request: MessagesRequest<'_>,
@ -38,7 +37,9 @@ pub(super) async fn execute_messages_provider_call(
let response = serde_json::from_str(&text)
.map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?;
request.config.transform_response(&request.model, response)
request
.config
.transform_anthropic_messages_response(&request.model, response)
}
pub(super) async fn execute_messages_provider_stream(
@ -46,9 +47,7 @@ pub(super) async fn execute_messages_provider_stream(
) -> Result<reqwest::Response, Error> {
let request = prepare_provider_request(request)?;
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::InvalidRequest(
"streaming messages is not supported for this provider".to_string(),
));
return Err(Error::Unsupported("streaming messages for this provider"));
}
let mut request_builder = http_client().post(&request.url).json(&request.body);

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

@ -13,7 +13,6 @@ mod client;
mod common_utils;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
use handler::{execute_messages_provider_call, execute_messages_provider_stream};

22
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -1,11 +1,15 @@
use super::Error;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use super::types::{MessagesRequest, ProviderMessagesRequest};
use serde_json::{Map, Value};
use super::Error;
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
use super::types::{MessagesRequest, ProviderMessagesRequest};
use crate::litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, get_custom_llm_provider,
};
use crate::llms::base_llm::anthropic_messages::transformation::{
BaseAnthropicMessagesConfig, MessagesAuthStrategy,
};
pub(super) fn prepare_provider_request(
request: MessagesRequest<'_>,
) -> Result<ProviderMessagesRequest, Error> {
@ -36,14 +40,14 @@ pub(super) fn prepare_provider_request(
let typed_request = serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
})?;
let transformed = config.transform_request(typed_request)?;
let transformed = config.transform_anthropic_messages_request(typed_request)?;
let body = serde_json::to_value(transformed).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize Anthropic messages request: {err}"
))
})?;
let url = config.complete_url(request.api_base, &model, &env_lookup)?;
let url = config.get_complete_url(request.api_base, &model, &env_lookup)?;
Ok(ProviderMessagesRequest {
provider: provider.to_string(),
@ -57,7 +61,7 @@ pub(super) fn prepare_provider_request(
}
fn validate_environment(
config: &dyn AnthropicMessagesProviderConfig,
config: &dyn BaseAnthropicMessagesConfig,
extra_headers: Option<Map<String, Value>>,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,

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

@ -5,7 +5,6 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::Error;
use super::common_utils::{
has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
};

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

@ -3,7 +3,7 @@ use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::AnthropicMessagesProviderConfig;
use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
pub struct MessagesRequest<'a> {
pub model: &'a str,
@ -18,7 +18,7 @@ pub struct MessagesRequest<'a> {
pub(super) struct ProviderMessagesRequest {
pub(super) provider: String,
pub(super) model: String,
pub(super) config: &'static dyn AnthropicMessagesProviderConfig,
pub(super) config: &'static dyn BaseAnthropicMessagesConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,

131
litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs
View file

@ -1,131 +0,0 @@
use super::super::OcrAdapter;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::cohere::{
CohereParams, CohereResponse, transform_request, transform_response, validate_document,
};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
use litellm_auth_azure::AzureAuthInputs;
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
pub(crate) struct AzureCohereAdapter;
impl OcrAdapter for AzureCohereAdapter {
type ProviderResponse = CohereResponse;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = super::super::super::wire::decode_request_value::<CohereParams>(
serde_json::Value::Object(request.optional_params.clone()),
"optional_params",
)?;
let mut config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
let base = request
.connection
.api_base
.clone()
.or_else(|| credential_env(AZURE_AI_API_BASE_ENV))
.filter(|base| !base.trim().is_empty())
.ok_or_else(|| {
Error::Auth(
"Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(),
)
})?;
let headers =
super::validate_ai_environment(&request.connection, &config, &credential_env).await?;
validate_document(&request.document)?;
let remote = request.document.source().starts_with("http://")
|| request.document.source().starts_with("https://");
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = transform_request(&request.model, document, params)?;
transform_request_body(
client,
request,
&complete_url(&base)?,
&headers,
!remote,
body,
|body| {
validate_document(&body.document)?;
validate_inline_document(&body.document)
},
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
transform_response(&request.model, response)
}
}
fn complete_url(base: &str) -> Result<String, OcrError> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base().into());
}
let path = url.path().trim_end_matches('/').to_string();
if path.ends_with("/v2/parse") {
url.set_path(&path);
return Ok(url.into());
}
url.set_path(path.strip_suffix("/models").unwrap_or(&path));
ApiUrl::parse(url.as_str())
.and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base().into())
}
fn invalid_api_base() -> OcrRequestError {
OcrRequestError::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in [
"",
"/models",
"/providers/cohere/v2",
"/providers/cohere/v2/parse",
] {
assert_eq!(
complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(),
"https://example.com/providers/cohere/v2/parse?tenant=a"
);
}
assert_eq!(
complete_url("https://example.com/v2/parse?tenant=a").unwrap(),
"https://example.com/v2/parse?tenant=a"
);
assert!(complete_url("relative/path").is_err());
}
}

214
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs
View file

@ -1,214 +0,0 @@
use super::super::OcrAdapter;
use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER};
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::document_intelligence::{
self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams,
};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
use crate::url_utils::ApiUrl;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
mod polling;
const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
#[derive(Clone, Debug)]
pub(crate) struct AzureDocumentIntelligenceAdapter;
impl OcrAdapter for AzureDocumentIntelligenceAdapter {
type ProviderResponse = AzureDocumentIntelligenceOperation;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = map_ocr_params(request)?;
let mut config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let endpoint = nonblank(request.connection.api_base.clone())
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
.ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?;
let url = get_complete_url(&endpoint, &request.model, &params)?;
let body = document_intelligence::transform_ocr_request(request.document.clone())?;
transform_request_body(client, request, &url, &headers, false, body, |_| Ok(())).await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
document_intelligence::transform_ocr_response(&request.model, response)
}
async fn read_response(
&self,
client: &OcrClient,
response: reqwest::Response,
url: &str,
headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> Result<crate::ocr::wire::DecodedOcrResponse<Self::ProviderResponse>, OcrError> {
polling::read_operation_response(
client.polling_http(),
response,
url,
headers,
&request.connection,
request.response_format()? == OcrResponseFormat::Native,
&request.hooks,
)
.await
}
}
fn map_ocr_params(
request: &LiteLLMOcrRequest,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let params = document_intelligence::decode_input_params(
request.optional_params.clone(),
"optional_params",
)?;
let crate::ocr::prepare::ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = params;
document_intelligence::map_ocr_params(params)
}
fn get_complete_url(
endpoint: &str,
model: &str,
params: &DocumentIntelligenceParams,
) -> Result<String, OcrError> {
let model = format!("{}:analyze", model_id(model)?);
ApiUrl::parse(endpoint)
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
.map(|url| {
url.append_query_pairs(
[("api-version", AZURE_DI_API_VERSION)]
.into_iter()
.chain(params.pages.iter().map(|pages| ("pages", pages.as_str())))
.chain(
params
.features
.iter()
.map(|features| ("features", features.as_str())),
),
)
.into_string()
})
.map_err(|_| OcrRequestError::RequestField {
path: "api_base".into(),
})
.map_err(OcrError::from)
}
async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization")
|| crate::http_utils::has_header(&connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER)
{
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_DI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(
std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value()))
.chain(connection.extra_headers.clone())
.collect(),
);
}
let token = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?;
super::validate_destination(connection, token.source())?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {}", token.value())))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn model_id(model: &str) -> Result<&str, OcrRequestError> {
let model = model.rsplit('/').next().unwrap_or(model);
if matches!(model, "." | "..") {
return Err(OcrRequestError::DotModel);
}
Ok(model)
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
(AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into())
);
}
}

119
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs
View file

@ -1,119 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use reqwest::Url;
use tokio::time::Instant;
use crate::constants::{AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS};
use crate::ocr::client::read_json_response;
use crate::ocr::codecs::document_intelligence::{
AzureDocumentIntelligenceOperation, OperationStatus,
};
use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError};
use crate::ocr::hooks::OcrHooks;
use crate::ocr::types::OcrConnection;
use crate::ocr::wire::DecodedOcrResponse;
pub(super) async fn read_operation_response(
http_client: &reqwest::Client,
response: reqwest::Response,
original_url: &str,
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
hooks: &Arc<dyn OcrHooks>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
if response.status() != reqwest::StatusCode::ACCEPTED {
let bytes =
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
.await?;
crate::ocr::handler::post_call(hooks, &bytes).await?;
return Ok(crate::ocr::wire::decode_response(&bytes, native)?);
}
let location = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.ok_or(OcrPollingError::PollLocation)?
.to_string();
let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?;
let operation = Url::parse(&location).map_err(|_| OcrPollingError::PollOrigin)?;
if original.origin() != operation.origin()
|| !operation.username().is_empty()
|| operation.password().is_some()
{
return Err(OcrPollingError::PollOrigin.into());
}
let bytes =
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?;
crate::ocr::handler::post_call(hooks, &bytes).await?;
poll_operation(http_client, operation, headers, connection, native, hooks).await
}
async fn poll_operation(
http_client: &reqwest::Client,
url: Url,
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
hooks: &Arc<dyn OcrHooks>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
let deadline = Instant::now()
.checked_add(connection.poll_timeout)
.ok_or(OcrPollingError::PollTimeout)?;
loop {
let remaining = deadline
.checked_duration_since(Instant::now())
.filter(|remaining| !remaining.is_zero())
.ok_or(OcrPollingError::PollTimeout)?;
let builder = http_client
.get(url.clone())
.timeout(remaining.min(connection.timeout));
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]),
);
let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder))
.await
.map_err(|_| OcrPollingError::PollTimeout)?
.map_err(crate::transport::Error::from)?;
let retry = response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(OCR_POLL_RETRY_SECS)
.max(1);
let decoded = tokio::time::timeout_at(
deadline,
read_json_response::<AzureDocumentIntelligenceOperation>(
response,
native,
connection.max_response_bytes,
),
)
.await
.map_err(|_| OcrPollingError::PollTimeout)??;
match &decoded.data.status {
Some(OperationStatus::Succeeded) => {
crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?;
return Ok(decoded);
}
Some(OperationStatus::Running | OperationStatus::NotStarted) => {
tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry)))
.await
.map_err(|_| OcrPollingError::PollTimeout)?;
}
status => {
return Err(OcrResponseError::OperationStatus(
status
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| "None".into()),
)
.into());
}
}
}
}

229
litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs
View file

@ -1,229 +0,0 @@
use super::super::OcrAdapter;
use crate::constants::AZURE_AI_OCR_PATH;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::url_utils::ApiUrl;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
#[derive(Clone, Debug)]
pub(crate) struct AzureMistralAdapter;
impl OcrAdapter for AzureMistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let mut config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
config.azure_ad_token_provider = request.azure_ad_token_provider.clone();
let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?;
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let retains_document = !request.document.source().starts_with("http://")
&& !request.document.source().starts_with("https://");
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = mistral::transform_ocr_request(&request.model, document, &params)?;
transform_request_body(
client,
request,
&url,
&headers,
retains_document,
body,
|body| validate_inline_document(&body.document),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
fn get_complete_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, OcrError> {
let base = nonblank(api_base.map(str::to_string))
.or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV)))
.ok_or_else(|| Error::Auth(
"Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter".into(),
))?;
let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect();
ApiUrl::parse(&base)
.and_then(|url| url.complete_path(&path))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
pub(in crate::ocr::adapters) async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
if config.azure_ad_token_provider.is_some() {
super::resolve_entra(config, env_lookup).await?;
}
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_AI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(bearer_headers(connection, key.value()));
}
let key = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureAiCredentials)?;
super::validate_destination(connection, key.source())?;
Ok(bearer_headers(connection, key.value()))
}
fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> {
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_azure_path_and_preserves_query() {
assert_eq!(
get_complete_url(Some("https://example.com/?tenant=a"), &|_| None).unwrap(),
"https://example.com/providers/mistral/azure/ocr?tenant=a"
);
assert_eq!(
get_complete_url(
Some("https://example.com/providers/mistral/azure/ocr"),
&|_| None
)
.unwrap(),
"https://example.com/providers/mistral/azure/ocr"
);
}
#[tokio::test]
async fn supplied_authorization_precedes_keys() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
extra_headers: vec![("authorization".into(), "Bearer prepared".into())],
..Default::default()
};
assert_eq!(
validate_environment(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap(),
connection.extra_headers
);
}
#[tokio::test]
async fn request_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
..Default::default()
};
assert_eq!(
validate_environment(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap()[0],
("Authorization".into(), "Bearer request-key".into())
);
}
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
("Authorization".into(), "Bearer request-key".into())
);
}
}

123
litellm-rust/crates/core/src/ocr/adapters/cohere.rs
View file

@ -1,123 +0,0 @@
use super::OcrAdapter;
use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE};
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::cohere::{
CohereParams, CohereResponse, transform_request, transform_response, validate_document,
};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::url_utils::ApiUrl;
pub(crate) struct CohereAdapter;
impl OcrAdapter for CohereAdapter {
type ProviderResponse = CohereResponse;
const PROVIDER: OcrProvider = OcrProvider::Cohere;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = super::super::wire::decode_request_value::<CohereParams>(
serde_json::Value::Object(request.optional_params.clone()),
"optional_params",
)?;
let headers = validate_environment(&request.connection, &credential_env)?;
let url = complete_url(
request
.connection
.api_base
.as_deref()
.unwrap_or(COHERE_PARSE_API_BASE),
)?;
let body = transform_request(&request.model, request.document.clone(), params)?;
transform_request_body(client, request, &url, &headers, true, body, |body| {
validate_document(&body.document)
})
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
transform_response(&request.model, response)
}
}
fn complete_url(base: &str) -> Result<String, OcrError> {
let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(invalid_api_base().into());
}
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base().into())
}
fn invalid_api_base() -> OcrRequestError {
OcrRequestError::RequestField {
path: "api_base".into(),
}
}
fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(COHERE_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or_else(|| {
Error::Auth("Missing COHERE_API_KEY - set it in the environment or pass api_key".into())
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in ["", "/v2", "/v2/parse"] {
assert_eq!(
complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(),
"https://example.com/v2/parse?tenant=a"
);
}
}
#[test]
fn rejects_invalid_urls_and_blank_keys() {
assert!(complete_url("relative/path").is_err());
assert!(complete_url("ftp://example.com").is_err());
assert!(matches!(
validate_environment(
&OcrConnection {
api_key: Some(" ".into()),
..Default::default()
},
&|_| None,
),
Err(OcrError::Public(Error::Auth(_)))
));
}
}

147
litellm-rust/crates/core/src/ocr/adapters/mistral.rs
View file

@ -1,147 +0,0 @@
use super::OcrAdapter;
use crate::constants::MISTRAL_OCR_API_BASE;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::url_utils::ApiUrl;
const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY";
#[derive(Clone, Debug)]
pub(crate) struct MistralAdapter;
impl OcrAdapter for MistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::Mistral;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let headers = validate_environment(&request.connection, &credential_env)?;
let url = get_complete_url(request.connection.api_base.as_deref())?;
let body =
mistral::transform_ocr_request(&request.model, request.document.clone(), &params)?;
transform_request_body(client, request, &url, &headers, true, body, |_| Ok(())).await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
pub(crate) fn get_complete_url(api_base: Option<&str>) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(MISTRAL_OCR_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or(Error::MissingApiKey {
provider: "Mistral",
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn complete_url_defaults_and_dedupes_v1() {
assert_eq!(
get_complete_url(None).unwrap(),
"https://api.mistral.ai/v1/ocr"
);
assert_eq!(
get_complete_url(Some("https://example.com/v1?tenant=a")).unwrap(),
"https://example.com/v1/ocr?tenant=a"
);
assert_eq!(
get_complete_url(Some("https://example.com/v1/ocr?tenant=a")).unwrap(),
"https://example.com/v1/ocr?tenant=a"
);
}
#[test]
fn environment_prefers_explicit_key_then_environment() {
let explicit = OcrConnection {
api_key: Some("explicit".into()),
..OcrConnection::default()
};
assert_eq!(
validate_environment(&explicit, &|_| Some("environment".into())).unwrap()[0],
("Authorization".into(), "Bearer explicit".into())
);
assert_eq!(
validate_environment(&OcrConnection::default(), &|_| Some("environment".into()))
.unwrap()[0],
("Authorization".into(), "Bearer environment".into())
);
}
#[test]
fn environment_preserves_forwarded_authorization() {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer forwarded".into())],
..OcrConnection::default()
};
assert_eq!(
validate_environment(&connection, &|_| None).unwrap(),
connection.extra_headers
);
}
#[test]
fn environment_rejects_missing_key() {
assert!(matches!(
validate_environment(&OcrConnection::default(), &|_| None),
Err(OcrError::Public(Error::MissingApiKey {
provider: "Mistral"
}))
));
}
}

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

@ -1,91 +0,0 @@
use std::future::Future;
use serde::de::DeserializeOwned;
use super::OcrClient;
use super::error::{OcrError, OcrResponseError};
use super::registry::OcrProvider;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
mod azure;
mod cohere;
mod mistral;
mod reducto;
mod vertex;
pub(crate) use azure::{AzureCohereAdapter, AzureDocumentIntelligenceAdapter, AzureMistralAdapter};
pub(crate) use cohere::CohereAdapter;
pub(crate) use mistral::MistralAdapter;
pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter};
pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter};
/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response.
pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
/// Provider JSON schema; direct and Vertex Mistral share `MistralOcrResponse`.
type ProviderResponse: DeserializeOwned + Send;
const PROVIDER: OcrProvider;
/// Prepares the complete provider HTTP request.
/// `request` contains the model, document, connection, and unmapped caller options.
/// `client` supplies reusable provider and document HTTP clients.
/// Returns the complete HTTP request, whereas Python returns body data.
fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<reqwest::Request, OcrError>> + Send;
/// Python: `transform_ocr_response`.
/// `request` supplies caller context, including the fallback model.
/// `response` is the decoded provider payload; the output is the shared LiteLLM schema.
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError>;
/// Decodes provider HTTP; adapters may override this to poll asynchronous operations.
/// Python performs that polling inside `async_transform_ocr_response`.
/// `client` is reused for polling; `response` is the initial HTTP response.
/// `url` and `headers` describe the submitted call; `request` supplies limits and format.
fn read_response(
&self,
_client: &OcrClient,
response: reqwest::Response,
_url: &str,
_headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> impl Future<
Output = Result<super::wire::DecodedOcrResponse<Self::ProviderResponse>, OcrError>,
> + Send {
async move {
let bytes =
super::client::read_response_bytes(response, request.connection.max_response_bytes)
.await?;
super::handler::post_call(&request.hooks, &bytes).await?;
Ok(super::wire::decode_response(
&bytes,
request.response_format()? == super::types::OcrResponseFormat::Native,
)?)
}
}
}
macro_rules! for_each_ocr_adapter {
($callback:ident) => {
$callback! {
Cohere, $crate::ocr::adapters::CohereAdapter, $crate::ocr::adapters::CohereAdapter, Cohere;
AzureCohere, $crate::ocr::adapters::AzureCohereAdapter, $crate::ocr::adapters::AzureCohereAdapter, AzureAi;
Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral;
AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi;
AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi;
ReductoLegacy, $crate::ocr::adapters::ReductoLegacyAdapter, $crate::ocr::adapters::ReductoLegacyAdapter, Reducto;
ReductoV3, $crate::ocr::adapters::ReductoV3Adapter, $crate::ocr::adapters::ReductoV3Adapter, Reducto;
VertexMistral, $crate::ocr::adapters::VertexMistralAdapter, $crate::ocr::adapters::VertexMistralAdapter, VertexAi;
VertexDeepSeek, $crate::ocr::adapters::VertexDeepSeekAdapter, $crate::ocr::adapters::VertexDeepSeekAdapter, VertexAi;
}
};
}
pub(crate) use for_each_ocr_adapter;

45
litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs
View file

@ -1,45 +0,0 @@
use super::super::OcrAdapter;
use crate::ocr::OcrClient;
use crate::ocr::codecs::reducto::{self, ReductoLegacyParams, ReductoResponse};
use crate::ocr::error::{OcrError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env,
guardrail_document, merge_extra_params,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
#[derive(Clone, Debug)]
pub(crate) struct ReductoLegacyAdapter;
impl OcrAdapter for ReductoLegacyAdapter {
type ProviderResponse = ReductoResponse;
const PROVIDER: OcrProvider = OcrProvider::Reducto;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params,
} = _prepare_ocr_request::<ReductoLegacyParams>(request)?;
let headers = super::validate_environment(&request.connection, &credential_env)?;
let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
let (document, headers) = guardrail_document(request, &url, &headers).await?;
let document =
super::prepare_document(client, document, &request.connection, &headers).await?;
let body = reducto::transform_legacy_ocr_request(&request.model, document, &params)?;
let body = merge_extra_params(&body, extra_params)?;
build_http_request(client, request, &url, &headers, &body)
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
reducto::transform_ocr_response(&request.model, response)
}
}

148
litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs
View file

@ -1,148 +0,0 @@
mod legacy;
mod v3;
use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX};
use crate::ocr::Error;
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::types::{OcrConnection, OcrDocument};
use crate::url_utils::ApiUrl;
pub(crate) use legacy::ReductoLegacyAdapter;
pub(crate) use v3::ReductoV3Adapter;
pub(super) fn get_complete_url(api_base: Option<&str>, path: &str) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&[path]))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
pub(super) fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or(Error::MissingReductoApiKey)?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
pub(super) async fn prepare_document(
client: &crate::ocr::OcrClient,
document: OcrDocument,
connection: &OcrConnection,
headers: &[(String, String)],
) -> Result<OcrDocument, OcrError> {
if document.source().starts_with(REDUCTO_ID_PREFIX) {
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(OcrRequestError::RequestField {
path: "document file id".into(),
}
.into());
}
return Ok(document);
}
let inline = InlineDocument::parse(document.source())?.ok_or(OcrRequestError::ReductoSource)?;
let mime = inline.mime_type().to_string();
let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
let part = reqwest::multipart::Part::bytes(bytes)
.file_name("document")
.mime_str(&mime)
.map_err(|_| OcrRequestError::InvalidDataUri)?;
let builder = client
.provider_http()
.post(get_complete_url(connection.api_base.as_deref(), "upload")?)
.multipart(reqwest::multipart::Form::new().part("file", part))
.timeout(connection.timeout);
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]),
);
let response = crate::http_utils::http_request(builder)
.await
.map_err(crate::transport::Error::from)?;
let uploaded = crate::ocr::client::read_json_response::<
crate::ocr::codecs::reducto::ReductoUploadResponse,
>(response, false, connection.max_response_bytes)
.await?
.data;
let file_id = uploaded
.file_id
.as_deref()
.map(str::trim)
.filter(|id| !id.is_empty());
let Some(file_id) = file_id else {
return Err(OcrResponseError::ResponseField {
path: "file_id".into(),
}
.into());
};
Ok(document.with_source(file_id.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn explicit_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("passed-key".into()),
..Default::default()
};
let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap();
assert_eq!(headers[0].1, "Bearer passed-key");
}
#[test]
fn blank_explicit_key_uses_environment_key() {
let connection = OcrConnection {
api_key: Some(" ".into()),
..Default::default()
};
let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap();
assert_eq!(headers[0].1, "Bearer env-key");
}
#[test]
fn existing_authorization_skips_key_lookup() {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer existing".into())],
..Default::default()
};
assert_eq!(
validate_environment(&connection, &|_| None).unwrap(),
connection.extra_headers
);
}
}

45
litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs
View file

@ -1,45 +0,0 @@
use super::super::OcrAdapter;
use crate::ocr::OcrClient;
use crate::ocr::codecs::reducto::{self, ReductoResponse, ReductoV3Params};
use crate::ocr::error::{OcrError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env,
guardrail_document, merge_extra_params,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
#[derive(Clone, Debug)]
pub(crate) struct ReductoV3Adapter;
impl OcrAdapter for ReductoV3Adapter {
type ProviderResponse = ReductoResponse;
const PROVIDER: OcrProvider = OcrProvider::Reducto;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params,
} = _prepare_ocr_request::<ReductoV3Params>(request)?;
let headers = super::validate_environment(&request.connection, &credential_env)?;
let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
let (document, headers) = guardrail_document(request, &url, &headers).await?;
let document =
super::prepare_document(client, document, &request.connection, &headers).await?;
let body = reducto::transform_v3_ocr_request(&request.model, document, &params)?;
let body = merge_extra_params(&body, extra_params)?;
build_http_request(client, request, &url, &headers, &body)
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
reducto::transform_ocr_response(&request.model, response)
}
}

140
litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs
View file

@ -1,140 +0,0 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::deepseek::{self, DeepSeekOcrParams, DeepSeekOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
use litellm_auth_gcp::{self as vertex, VertexConfig};
const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
const MODEL_NAMESPACE: &str = "deepseek-ai";
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug)]
pub(crate) struct VertexDeepSeekAdapter;
impl OcrAdapter for VertexDeepSeekAdapter {
type ProviderResponse = DeepSeekOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::VertexAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
validate_destination(&request.connection)?;
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<DeepSeekOcrParams>(request)?;
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let authentication = client
.vertex_auth()
.validate_environment(
request.connection.extra_headers.clone(),
request.connection.api_key.as_deref(),
&config,
&credential_env,
)
.await
.map_err(Error::from)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
let url = get_complete_url(
request.connection.api_base.as_deref(),
&authentication.project_id,
&location,
)?;
let document = request.document.clone();
let body =
deepseek::transform_ocr_request(&provider_model(&request.model), document, &params)?;
transform_request_body(
client,
request,
&url,
&authentication.headers,
false,
body,
|_| Ok(()),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
deepseek::transform_ocr_response(&request.model, response)
}
}
fn provider_model(model: &str) -> String {
if model.starts_with(&format!("{MODEL_NAMESPACE}/")) {
model.to_string()
} else {
format!("{MODEL_NAMESPACE}/{model}")
}
}
fn get_complete_url(
api_base: Option<&str>,
project: &str,
location: &str,
) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(DEFAULT_API_BASE);
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"endpoints",
"openapi",
"chat",
"completions",
])
})
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
#[cfg(test)]
mod tests {
use super::{get_complete_url, provider_model};
#[test]
fn adapter_owns_model_namespace_and_endpoint() {
assert_eq!(
provider_model("deepseek-ocr-maas"),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
provider_model("deepseek-ai/deepseek-ocr-maas"),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
get_complete_url(None, "proj-1", "europe-west4").unwrap(),
"https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions"
);
}
}

157
litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs
View file

@ -1,157 +0,0 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::ocr::Error;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
use litellm_auth_gcp::{self as vertex, VertexConfig};
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug)]
pub(crate) struct VertexMistralAdapter;
impl OcrAdapter for VertexMistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::VertexAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
validate_destination(&request.connection)?;
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let authentication = client
.vertex_auth()
.validate_environment(
request.connection.extra_headers.clone(),
request.connection.api_key.as_deref(),
&config,
&credential_env,
)
.await
.map_err(Error::from)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
let url = get_complete_url(
request.connection.api_base.as_deref(),
&authentication.project_id,
&location,
&request.model,
)?;
let retains_document = !request.document.source().starts_with("http://")
&& !request.document.source().starts_with("https://");
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = mistral::transform_ocr_request(&request.model, document, &params)?;
transform_request_body(
client,
request,
&url,
&authentication.headers,
retains_document,
body,
|body| validate_inline_document(&body.document),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
fn get_complete_url(
api_base: Option<&str>,
project: &str,
location: &str,
model: &str,
) -> Result<String, OcrError> {
validate_location(location)?;
let default_base = format!("https://{location}-aiplatform.googleapis.com");
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(&default_base);
let prediction = format!("{model}:rawPredict");
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"publishers",
"mistralai",
"models",
&prediction,
])
})
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
fn validate_location(location: &str) -> Result<(), OcrError> {
let valid = !location.is_empty()
&& location
.bytes()
.all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-')
&& location
.as_bytes()
.first()
.is_some_and(u8::is_ascii_alphanumeric)
&& location
.as_bytes()
.last()
.is_some_and(u8::is_ascii_alphanumeric);
if valid {
return Ok(());
}
Err(OcrRequestError::RequestField {
path: "vertex_location".into(),
}
.into())
}
#[cfg(test)]
mod tests {
use super::get_complete_url;
#[test]
fn endpoint_uses_location_project_and_model() {
assert_eq!(
get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas").unwrap(),
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
assert!(get_complete_url(None, "proj-1", "attacker.example/path", "model").is_err());
}
}

18
litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs
View file

@ -1,18 +0,0 @@
mod deepseek;
mod mistral;
use crate::ocr::Error;
use litellm_auth::InputSource;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
pub(crate) use deepseek::VertexDeepSeekAdapter;
pub(crate) use mistral::VertexMistralAdapter;
fn validate_destination(connection: &OcrConnection) -> Result<(), OcrError> {
if connection.api_base.is_some() && connection.api_base_source == InputSource::Request {
return Err(Error::from(litellm_auth::Error::RequestVertexCredentialDestination).into());
}
Ok(())
}

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

@ -0,0 +1,100 @@
use super::provider_config::{OcrConfigKind, resolve_provider_config};
use crate::call_arguments::ArgumentSpec;
const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"];
const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[
"azure_ad_token",
"tenant_id",
"client_id",
"client_secret",
"azure_scope",
"azure_authority_host",
"azure_credential",
"azure_federated_token_file",
"enable_azure_ad_token_refresh",
];
const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[
"vertex_credentials",
"vertex_ai_credentials",
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
];
pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> bool {
resolve_provider_config(model, custom_llm_provider).is_ok()
}
pub fn consumed_optional_param_names(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Vec<&'static str>, super::Error> {
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::AzureAi
| OcrConfigKind::AzureDocumentIntelligence
| OcrConfigKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS,
OcrConfigKind::VertexAi | OcrConfigKind::VertexDeepSeek => VERTEX_AUTH_OPTION_FIELDS,
_ => &[],
};
Ok(COMMON_OPTION_FIELDS
.iter()
.chain(provider_fields)
.chain(auth_fields)
.copied()
.collect())
}
pub fn consumed_optional_params(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Vec<ArgumentSpec>, super::Error> {
consumed_optional_param_names(model, custom_llm_provider).map(|names| {
names
.into_iter()
.map(|name| ArgumentSpec {
name,
secret: matches!(
name,
"azure_ad_token"
| "client_secret"
| "azure_federated_token_file"
| "vertex_credentials"
| "vertex_ai_credentials"
),
})
.collect()
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn consumed_params_include_provider_options_and_mark_credentials() {
let mistral = consumed_optional_param_names("mistral/model", None).unwrap();
assert!(mistral.contains(&"pages"));
assert!(mistral.contains(&"req_format"));
assert!(!mistral.contains(&"vertex_project"));
let vertex = consumed_optional_param_names("vertex_ai/deepseek-ocr", None).unwrap();
assert!(!vertex.contains(&"temperature"));
assert!(vertex.contains(&"vertex_credentials"));
assert!(!vertex.contains(&"pages"));
let azure = consumed_optional_params("model", Some("azure_ai")).unwrap();
assert!(
azure
.iter()
.any(|spec| spec.name == "client_secret" && spec.secret)
);
assert!(
azure
.iter()
.any(|spec| spec.name == "tenant_id" && !spec.secret)
);
}
}

73
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -2,15 +2,13 @@ use std::sync::OnceLock;
use std::time::Duration;
use bytes::{Bytes, BytesMut};
use litellm_auth_gcp::VertexAuth;
use serde::de::DeserializeOwned;
use super::error::{Error, OcrError, OcrResponseError};
use super::json::{DecodedOcrResponse, decode_response};
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use super::wire::{DecodedOcrResponse, decode_response};
use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
use crate::media::MediaFetcher;
use crate::transport::Error as TransportError;
use litellm_auth_gcp::VertexAuth;
#[derive(Clone)]
pub struct OcrClient {
@ -21,8 +19,8 @@ pub struct OcrClient {
}
impl OcrClient {
pub fn new(provider_http: reqwest::Client) -> Result<Self, TransportError> {
let document_fetcher = MediaFetcher::new().map_err(TransportError::from)?;
pub fn new(provider_http: reqwest::Client) -> Result<Self, crate::transport::Error> {
let document_fetcher = MediaFetcher::new().map_err(crate::transport::Error::from)?;
Ok(Self {
provider_http,
polling_http: no_redirect_http()?,
@ -31,11 +29,14 @@ impl OcrClient {
})
}
pub fn shared() -> Result<Self, Error> {
pub fn shared() -> Result<Self, crate::ocr::Error> {
shared_client()
}
pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
pub async fn perform(
&self,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
use super::{
NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost,
OcrHostOperation, OcrHostResult,
@ -45,7 +46,7 @@ impl OcrClient {
let mut request = Some(request);
let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all())
else {
return Err(Error::InvalidRequest(
return Err(crate::ocr::Error::InvalidRequest(
"native OCR host admission declined".into(),
));
};
@ -54,16 +55,11 @@ impl OcrClient {
match call.resume(result.take()).await? {
OcrCallStep::Host(OcrHostOperation::ProjectRequest) => {
result = Some(OcrHostResult::Request(Ok((
Box::new(
request
.take()
.ok_or_else(|| {
Error::InvalidRequest(
"OCR request was already projected".into(),
)
})?
.into(),
),
Box::new(request.take().ok_or_else(|| {
crate::ocr::Error::InvalidRequest(
"OCR request was already projected".into(),
)
})?),
false,
))))
}
@ -100,29 +96,29 @@ impl OcrClient {
}
}
fn no_redirect_http() -> Result<reqwest::Client, TransportError> {
fn no_redirect_http() -> Result<reqwest::Client, crate::transport::Error> {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(TransportError::from)
.map_err(crate::transport::Error::from)
}
pub(crate) fn shared_client() -> Result<OcrClient, Error> {
static CLIENT: OnceLock<Result<OcrClient, TransportError>> = OnceLock::new();
pub(crate) fn shared_client() -> Result<OcrClient, crate::ocr::Error> {
static CLIENT: OnceLock<Result<OcrClient, crate::transport::Error>> = OnceLock::new();
let client = CLIENT
.get_or_init(|| {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.build()
.map_err(TransportError::from)
.map_err(crate::transport::Error::from)
.and_then(OcrClient::new)
})
.clone()?;
Ok(client)
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
shared_client()?.perform(request).await
}
@ -130,33 +126,28 @@ pub async fn read_json_response<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
max_response_bytes: usize,
) -> Result<DecodedOcrResponse<T>, OcrError> {
) -> Result<DecodedOcrResponse<T>, crate::ocr::Error> {
let bytes = read_response_bytes(response, max_response_bytes).await?;
Ok(decode_response(&bytes, native)?)
decode_response(&bytes, native)
}
pub(crate) async fn read_response_bytes(
mut response: reqwest::Response,
max_response_bytes: usize,
) -> Result<Bytes, OcrError> {
limit: usize,
) -> Result<Bytes, crate::ocr::Error> {
let status = response.status();
let limit = if status.is_success() {
max_response_bytes
} else {
max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1))
};
if status.is_success()
&& response
.content_length()
.is_some_and(|length| length > limit as u64)
{
return Err(OcrResponseError::TooLarge { limit }.into());
return Err(crate::ocr::Error::TooLarge { limit });
}
let mut bytes = BytesMut::new();
while let Some(chunk) = response.chunk().await.map_err(transport_error)? {
let remaining = limit.saturating_sub(bytes.len());
if status.is_success() && chunk.len() > remaining {
return Err(OcrResponseError::TooLarge { limit }.into());
return Err(crate::ocr::Error::TooLarge { limit });
}
bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
if !status.is_success() && bytes.len() == limit {
@ -166,19 +157,19 @@ pub(crate) async fn read_response_bytes(
if !status.is_success() {
return Err(crate::transport::Error::Http {
status: status.as_u16(),
body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)),
body: String::from_utf8_lossy(&bytes).into_owned(),
}
.into());
}
Ok(bytes.freeze())
}
pub(crate) fn transport_error(error: reqwest::Error) -> Error {
pub(crate) fn transport_error(error: reqwest::Error) -> crate::ocr::Error {
if error.is_timeout() {
return Error::Http {
return crate::ocr::Error::Transport(crate::transport::Error::Http {
status: 408,
body: "OCR request timed out".into(),
};
});
}
crate::transport::Error::from(error).into()
}
@ -203,7 +194,7 @@ mod tests {
.unwrap_err();
assert!(matches!(
transport_error(error),
Error::Http { status: 408, .. }
crate::ocr::Error::Transport(crate::transport::Error::Http { status: 408, .. })
));
server.abort();
}

254
litellm-rust/crates/core/src/ocr/codecs/cohere.rs
View file

@ -1,254 +0,0 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum OutputFormat {
#[default]
Markdown,
Blocks,
}
#[derive(Deserialize)]
pub(crate) struct CohereParams {
#[serde(default)]
pub output_format: OutputFormat,
}
#[derive(Deserialize, Serialize)]
pub(crate) struct CohereRequest {
pub model: String,
pub document: OcrDocument,
pub output_format: OutputFormat,
}
pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), OcrRequestError> {
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(OcrRequestError::CohereImageOnly);
};
if image_url.is_empty() {
return Err(OcrRequestError::CohereImageOnly);
}
if let Some(inline) = InlineDocument::parse(image_url)? {
if !inline.mime_type().type_.eq_ignore_ascii_case("image") {
return Err(OcrRequestError::CohereImageOnly);
}
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
}
Ok(())
}
#[derive(Deserialize)]
pub(crate) struct CohereResponse {
#[serde(default)]
pages: Vec<CoherePage>,
meta: Option<CohereMeta>,
}
#[derive(Deserialize)]
struct CoherePage {
index: Option<i64>,
markdown: Option<CohereMarkdown>,
blocks: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize)]
struct CohereMarkdown {
#[serde(default)]
content: String,
images: Option<Vec<Map<String, Value>>>,
}
#[derive(Deserialize)]
struct CohereMeta {
billed_units: Option<CohereBilledUnits>,
}
#[derive(Deserialize)]
struct CohereBilledUnits {
pages: Option<i64>,
}
pub(crate) fn transform_response(
model: &str,
response: CohereResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let pages_processed = response
.meta
.and_then(|meta| meta.billed_units)
.and_then(|units| units.pages)
.map(Ok)
.unwrap_or_else(|| {
i64::try_from(response.pages.len()).map_err(|_| OcrResponseError::NumericRange("pages"))
})?;
let pages = response
.pages
.into_iter()
.enumerate()
.map(|(position, page)| {
let index = page.index.map(Ok).unwrap_or_else(|| {
i64::try_from(position).map_err(|_| OcrResponseError::NumericRange("page index"))
})?;
let (content, images) = page
.markdown
.map(|markdown| {
let images =
markdown
.images
.filter(|images| !images.is_empty())
.map(|images| {
images
.into_iter()
.map(|mut image| {
if let Some(Value::Object(bbox)) =
image.get("bounding_box").cloned()
{
image.insert("bbox".into(), Value::Object(bbox));
}
Value::Object(image)
})
.collect::<Vec<_>>()
});
(markdown.content, images)
})
.unwrap_or_default();
let mut normalized = json!({"index": index, "markdown": content, "images": images});
if let Some(blocks) = page.blocks {
normalized["blocks"] = json!(blocks);
}
Ok(normalized)
})
.collect::<Result<Vec<_>, OcrResponseError>>()?;
Ok(LiteLLMOcrResponse {
pages,
model: model.into(),
document_annotation: None,
usage_info: Some(json!({"pages_processed": pages_processed})),
object: "ocr".into(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
pub(crate) fn transform_request(
model: &str,
document: OcrDocument,
params: CohereParams,
) -> Result<CohereRequest, OcrRequestError> {
validate_document(&document)?;
Ok(CohereRequest {
model: model.into(),
document,
output_format: params.output_format,
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn response_normalizes_markdown_images_blocks_and_billed_pages() {
let response = serde_json::from_value(json!({
"pages": [
{
"type":"markdown",
"index":4,
"markdown":{
"content":"receipt",
"images":[{
"id":"image",
"bounding_box":{"top_left_x":1,"bottom_right_x":48},
"bounding_box_normalized":{"top_left_x":0.04,"bottom_right_x":0.15},
"description":"scan",
"category":"logo"
}]
}
},
{"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]}
],
"meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}}
}))
.unwrap();
let normalized = transform_response("parse-v5.0", response).unwrap();
assert_eq!(normalized.pages[0]["index"], 4);
assert_eq!(normalized.pages[0]["markdown"], "receipt");
assert_eq!(normalized.pages[0]["images"][0]["bbox"]["top_left_x"], 1);
assert_eq!(
normalized.pages[0]["images"][0]["bounding_box_normalized"]["bottom_right_x"],
0.15
);
assert_eq!(normalized.pages[0]["images"][0]["description"], "scan");
assert_eq!(normalized.pages[0]["images"][0]["category"], "logo");
assert_eq!(normalized.pages[1]["index"], 1);
assert_eq!(normalized.pages[1]["markdown"], "");
assert_eq!(normalized.pages[1]["blocks"][0]["text"]["content"], "total");
assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 3);
}
#[test]
fn response_defaults_and_invalid_fields() {
for value in [
json!({}),
json!({"meta":null}),
json!({"pages":[],"meta":{"billed_units":null}}),
] {
let normalized =
transform_response("parse", serde_json::from_value(value).unwrap()).unwrap();
assert!(normalized.pages.is_empty());
assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 0);
}
for value in [
json!({"pages":null}),
json!({"pages":[{"markdown":"text"}]}),
json!({"pages":[{"index":"bad"}]}),
] {
assert!(serde_json::from_value::<CohereResponse>(value).is_err());
}
let normalized = transform_response(
"parse",
serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(),
)
.unwrap();
assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 1);
assert!(normalized.pages[0]["images"].is_null());
}
#[test]
fn request_requires_image_and_supported_output_format() {
for value in [
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
json!({"type":"image_url","image_url":""}),
json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}),
] {
assert_eq!(
validate_document(&serde_json::from_value(value).unwrap()),
Err(OcrRequestError::CohereImageOnly)
);
}
assert!(serde_json::from_value::<CohereParams>(json!({"output_format":"html"})).is_err());
for format in ["markdown", "blocks"] {
assert!(
serde_json::from_value::<CohereParams>(json!({"output_format":format})).is_ok()
);
}
let request = transform_request(
"parse-v5.0",
serde_json::from_value(json!({
"type":"image_url",
"image_url":"https://example.com/image.png"
}))
.unwrap(),
serde_json::from_value(json!({})).unwrap(),
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap()["output_format"],
"markdown"
);
}
}

Some files were not shown because too many files have changed in this diff Show more