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

Merge remote-tracking branch 'origin/main' into litellm_transcribe_passthrough

Browse source
This commit is contained in:
yassin 2026-09-17 23:07:57 +00:00
commit 912edaa8cc
249 changed files with 13175 additions and 5496 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cookbook/litellm_proxy_server/cli_token_usage.py
View file

@ -3,7 +3,7 @@
Example: Using CLI token with LiteLLM SDK
This example shows how to use the CLI authentication token
in your Python scripts after running `litellm-proxy login`.
in your Python scripts after running `lite login`.
"""
from textwrap import indent
@ -22,7 +22,7 @@ def main():
api_key = litellm.get_litellm_gateway_api_key()
if not api_key:
print("❌ No CLI token found. Please run 'litellm-proxy login' first.")
print("❌ No CLI token found. Please run 'lite login' first.")
return
print("✅ Found CLI token.")
@ -58,6 +58,6 @@ if __name__ == "__main__":
main()
print("\n💡 Tips:")
print("1. Run 'litellm-proxy login' to authenticate first")
print("1. Run 'lite login' to authenticate first")
print("2. Replace 'https://your-proxy.com' with your actual proxy URL")
print("3. The token is stored in your OS keychain, or in ~/.litellm/token.json when there is none")

5
db_scripts/partition_spend_logs.sql
View file

@ -53,6 +53,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
@ -78,6 +80,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
ON "LiteLLM_SpendLogs" ("api_key", "startTime");
-- Safety net: any row whose startTime has no explicit partition lands here so
-- writes never fail. The cleanup job never drops the DEFAULT partition.
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"

5
db_scripts/unpartition_spend_logs.sql
View file

@ -40,6 +40,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_api_key_startTime_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
@ -60,6 +62,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
ON "LiteLLM_SpendLogs" ("api_key", "startTime");
INSERT INTO "LiteLLM_SpendLogs"
SELECT * FROM "LiteLLM_SpendLogs_partitioned"
ON CONFLICT ("request_id") DO NOTHING;

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -678,6 +678,7 @@ model LiteLLM_SpendLogs {
@@index([end_user])
@@index([session_id])
@@index([litellm_call_id])
@@index([api_key, startTime])
}
model LiteLLM_BudgetWindowSpend {

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.98"
version = "0.4.99"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.98"
version = "0.4.99"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

160
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",
]
@ -1964,6 +2016,7 @@ dependencies = [
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-framing",
"litellm-providers",
"mime_guess",
"moka",
"rand 0.8.7",
@ -1974,6 +2027,7 @@ dependencies = [
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"sha2 0.10.9",
"strum",
"subtle",
@ -1999,6 +2053,18 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-providers"
version = "0.1.0"
dependencies = [
"litellm-auth",
"litellm-auth-aws",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-python-bridge"
version = "0.1.0"
@ -2036,7 +2102,7 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"indexmap",
"indexmap 2.14.0",
"itoa",
"rand 0.8.7",
"rstest",
@ -2757,6 +2823,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"
@ -3079,6 +3165,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"
@ -3160,6 +3270,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",
@ -3190,6 +3301,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"
@ -3665,7 +3807,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",

2
litellm-rust/Cargo.toml
View file

@ -16,6 +16,7 @@ litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-providers = { path = "crates/providers" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-token-counter = { path = "crates/token-counter" }
@ -30,6 +31,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"

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

4
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-providers.workspace = true
litellm-framing.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
@ -23,7 +24,8 @@ 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

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

@ -24,3 +24,23 @@ pub enum Error {
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}
impl From<litellm_providers::audio_transcription::Error> for Error {
fn from(error: litellm_providers::audio_transcription::Error) -> Self {
match error {
litellm_providers::audio_transcription::Error::InvalidType { expected, actual } => {
Self::InvalidType { expected, actual }
}
litellm_providers::audio_transcription::Error::MissingField(field) => {
Self::MissingField(field)
}
litellm_providers::audio_transcription::Error::InvalidRequest(message) => {
Self::InvalidRequest(message)
}
litellm_providers::audio_transcription::Error::InvalidResponse(message) => {
Self::InvalidResponse(message)
}
litellm_providers::audio_transcription::Error::Auth(error) => Self::Auth(error),
}
}
}

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 litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
use litellm_providers::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());

6
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 litellm_providers::audio_transcription::types;
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 litellm_providers::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
};
use litellm_providers::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(),

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 litellm_providers::anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use litellm_providers::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,
&litellm_providers::bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
),
_ => None,
}

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

@ -24,3 +24,19 @@ pub enum Error {
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}
impl From<litellm_providers::chat::Error> for Error {
fn from(error: litellm_providers::chat::Error) -> Self {
match error {
litellm_providers::chat::Error::MissingField(field) => Self::MissingField(field),
litellm_providers::chat::Error::InvalidRequest(message) => {
Self::InvalidRequest(message)
}
litellm_providers::chat::Error::InvalidResponse(message) => {
Self::InvalidResponse(message)
}
litellm_providers::chat::Error::Unsupported(reason) => Self::Unsupported(reason),
litellm_providers::chat::Error::Auth(error) => Self::Auth(error),
}
}
}

8
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 litellm_providers::base_llm::chat::transformation::ChatCompletionsAuth;
pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
@ -60,6 +59,7 @@ pub(super) async fn execute_chat_completions_provider_call(
request
.config
.transform_response(&request.model, ProviderChatResponseData { body })
.map_err(Error::from)
.map_err(as_response_error)
}
@ -87,7 +87,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,
};

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

@ -10,18 +10,15 @@ mod error;
pub use error::Error;
mod client;
mod common_utils;
pub mod conversation;
pub use litellm_providers::chat::{conversation, response_utils};
pub(crate) mod handler;
mod prepare;
pub mod response_utils;
pub mod streaming;
pub mod transformation;
pub mod types;
use serde_json::{Map, Value};
pub use litellm_providers::chat::types;
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 litellm_providers::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,

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 litellm_providers::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 {

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;

34
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,36 +1,4 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CustomLlmProvider<'a> {
pub model: &'a str,
pub custom_llm_provider: &'a str,
}
pub fn get_custom_llm_provider<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Option<CustomLlmProvider<'a>> {
if let Some(custom_llm_provider) = custom_llm_provider.filter(|provider| !provider.is_empty()) {
return Some(CustomLlmProvider {
model: strip_custom_llm_provider_prefix(model, custom_llm_provider),
custom_llm_provider,
});
}
let (custom_llm_provider, model) = model.split_once('/')?;
if custom_llm_provider.is_empty() || model.is_empty() {
return None;
}
Some(CustomLlmProvider {
model,
custom_llm_provider,
})
}
fn strip_custom_llm_provider_prefix<'a>(model: &'a str, custom_llm_provider: &str) -> &'a str {
model
.strip_prefix(custom_llm_provider)
.and_then(|model| model.strip_prefix('/'))
.unwrap_or(model)
}
pub use litellm_providers::provider_resolution::{CustomLlmProvider, get_custom_llm_provider};
#[cfg(test)]
mod tests {

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/llms/anthropic/chat/mod.rs Normal file
View file

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

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

@ -2,16 +2,16 @@ use std::collections::HashMap;
use serde_json::Value;
use super::super::experimental_pass_through::messages::streaming::{
AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
AnthropicStreamUsage,
};
use crate::chat_completions::Error;
use crate::chat_completions::streaming::StreamTransformer;
use crate::chat_completions::types::{
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionsUsage,
};
use crate::providers::anthropic::messages::streaming::{
AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
AnthropicStreamUsage,
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AnthropicJsonChunkType {

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

@ -5,7 +5,7 @@ use url::Url;
use crate::messages::Error;
use crate::messages::types::AnthropicMessagesResponse;
use crate::providers::anthropic::messages::transformation::resolve_anthropic_api_base;
use litellm_providers::anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base;
const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches";

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

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

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

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

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;

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

@ -0,0 +1 @@
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"));
}
}

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

@ -0,0 +1 @@
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)?
})
}

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

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

@ -0,0 +1,8 @@
pub mod anthropic;
pub mod azure_ai;
pub mod base_llm;
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/azure_ai/messages/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 litellm_providers::anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use litellm_providers::azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use litellm_providers::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),

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

@ -28,6 +28,22 @@ pub enum Error {
InvalidBedrockBase64(String),
}
impl From<litellm_providers::messages::Error> for Error {
fn from(error: litellm_providers::messages::Error) -> Self {
match error {
litellm_providers::messages::Error::MissingField(field) => Self::MissingField(field),
litellm_providers::messages::Error::InvalidRequest(message) => {
Self::InvalidRequest(message)
}
litellm_providers::messages::Error::InvalidResponse(message) => {
Self::InvalidResponse(message)
}
litellm_providers::messages::Error::Unsupported(reason) => Self::Unsupported(reason),
litellm_providers::messages::Error::Auth(error) => Self::Auth(error),
}
}
}
impl Error {
pub fn is_request(&self) -> bool {
match self {

10
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,10 @@ 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)
.map_err(Error::from)
}
pub(super) async fn execute_messages_provider_stream(

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

@ -13,8 +13,7 @@ mod client;
mod common_utils;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
pub use litellm_providers::messages::types;
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
use types::{AnthropicMessagesResponse, MessagesRequest};

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 litellm_providers::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,
};

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

5
litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs
View file

@ -1,5 +0,0 @@
mod transformation;
mod types;
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{DeepSeekOcrParams, DeepSeekOcrResponse};

101
litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs
View file

@ -1,101 +0,0 @@
use serde::de::IntoDeserializer;
use serde_json::{Value, json};
use super::types::*;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
pub(crate) fn transform_ocr_request(
provider_model: &str,
document: OcrDocument,
params: &DeepSeekOcrParams,
) -> Result<DeepSeekOcrRequest, OcrRequestError> {
if document.source().is_empty() {
return Err(OcrRequestError::MissingDocumentUrl);
}
let content = OcrDocument::ImageUrl {
image_url: document.source().to_string(),
extra_fields: serde_json::Map::new(),
};
Ok(DeepSeekOcrRequest {
model: provider_model.to_string(),
messages: vec![DeepSeekOcrMessage {
role: UserRole::User,
content: vec![content],
}],
params: params.clone(),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: DeepSeekOcrResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let content = response
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content)
.ok_or(OcrResponseError::EmptyContent)?;
let decoded = decode_content(content)?;
let pages = match decoded.result.pages {
Some(pages) if !pages.is_empty() => pages
.into_iter()
.map(|page| serde_json::to_value(page).expect("DeepSeek page serializes"))
.collect(),
_ => vec![json!({
"index":0,
"markdown":decoded.fallback_markdown,
"images":null
})],
};
Ok(LiteLLMOcrResponse {
pages,
model: decoded.result.model.unwrap_or_else(|| model.to_string()),
document_annotation: decoded.result.document_annotation,
usage_info: decoded.result.usage_info.or(response.usage),
object: "ocr".into(),
extra_fields: decoded.result.extra_fields,
provider_native_response: None,
})
}
struct DecodedContent {
result: DeepSeekOcrResult,
fallback_markdown: String,
}
fn decode_content(content: DeepSeekContent) -> Result<DecodedContent, OcrResponseError> {
let (result, fallback_markdown) = match content {
DeepSeekContent::Text(text) if text.is_empty() => {
return Err(OcrResponseError::EmptyContent);
}
DeepSeekContent::Text(text) => (decode_json_content(&text)?, text),
DeepSeekContent::Object(object) => {
let fallback =
serde_json::to_string(&object).map_err(|_| OcrResponseError::ResponseField {
path: "choices[0].message.content".into(),
})?;
(Some(object), fallback)
}
};
Ok(DecodedContent {
result: result.unwrap_or_default(),
fallback_markdown,
})
}
fn decode_json_content(text: &str) -> Result<Option<DeepSeekOcrResult>, OcrResponseError> {
if !text.trim_start().starts_with('{') {
return Ok(None);
}
let value = match serde_json::from_str::<Value>(text) {
Ok(value) => value,
Err(_) => return Ok(None),
};
serde_path_to_error::deserialize(value.into_deserializer())
.map(Some)
.map_err(|error| OcrResponseError::ResponseField {
path: format!("choices[0].message.content.{}", error.path()),
})
}

95
litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs
View file

@ -1,95 +0,0 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub stream: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<StopSequences>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum StopSequences {
One(String),
Many(Vec<String>),
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrRequest {
pub model: String,
pub messages: Vec<DeepSeekOcrMessage>,
#[serde(flatten)]
pub params: DeepSeekOcrParams,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrMessage {
pub role: UserRole,
pub content: Vec<crate::ocr::types::OcrDocument>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum UserRole {
User,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekOcrResponse {
#[serde(default)]
pub choices: Vec<DeepSeekChoice>,
pub usage: Option<Value>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekChoice {
pub message: DeepSeekResponseMessage,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekResponseMessage {
pub content: Option<DeepSeekContent>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub(crate) enum DeepSeekContent {
Text(String),
Object(DeepSeekOcrResult),
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrResult {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<Vec<DeepSeekPage>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage_info: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekPage {
#[serde(default)]
pub index: i64,
#[serde(default)]
pub markdown: String,
pub images: Option<Value>,
pub dimensions: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}

9
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs
View file

@ -1,9 +0,0 @@
mod params;
mod transformation;
mod types;
pub(crate) use params::{decode_input_params, map_ocr_params};
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{
AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, OperationStatus,
};

219
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs
View file

@ -1,219 +0,0 @@
use std::collections::BTreeSet;
use serde_json::{Map, Value};
use super::types::{
DocumentIntelligenceInputParams, DocumentIntelligenceParams, FeaturesInput, PagesInput,
};
use crate::ocr::error::OcrRequestError;
use crate::ocr::prepare::ParsedProviderParams;
pub(crate) fn decode_input_params(
params: Map<String, Value>,
prefix: &str,
) -> Result<ParsedProviderParams<DocumentIntelligenceInputParams>, OcrRequestError> {
if let Some(Value::Array(pages)) = params.get("pages") {
if pages.iter().any(Value::is_boolean) {
return Err(OcrRequestError::Pages("boolean page index".into()));
}
if pages
.iter()
.any(|page| page.is_number() && page.as_i64().is_none())
{
return Err(OcrRequestError::Pages("page index is out of range".into()));
}
if !pages.iter().all(Value::is_i64) && !pages.iter().all(Value::is_string) {
return Err(OcrRequestError::Pages("mixed page element types".into()));
}
}
crate::ocr::wire::decode_request_value(Value::Object(params), prefix)
}
pub(crate) fn map_ocr_params(
params: DocumentIntelligenceInputParams,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
Ok(DocumentIntelligenceParams {
pages: params.pages.map(normalize_pages).transpose()?.flatten(),
features: params
.features
.map(normalize_features)
.transpose()?
.flatten(),
})
}
fn normalize_pages(pages: PagesInput) -> Result<Option<String>, OcrRequestError> {
let normalized = match pages {
PagesInput::ZeroBasedIndices(indices) => {
if indices.is_empty() {
return Ok(None);
}
indices
.into_iter()
.map(|page| {
if page < 0 {
return Err(OcrRequestError::Pages("negative page index".into()));
}
page.checked_add(1)
.ok_or_else(|| OcrRequestError::Pages("page index is out of range".into()))
})
.collect::<Result<BTreeSet<_>, _>>()?
.into_iter()
.map(|page| page.to_string())
.collect::<Vec<_>>()
.join(",")
}
PagesInput::NativeTokens(tokens) => {
if tokens.is_empty() {
return Ok(None);
}
tokens
.iter()
.map(|token| token.trim())
.collect::<Vec<_>>()
.join(",")
}
PagesInput::NativeRange(range) => range
.split(',')
.map(str::trim)
.collect::<Vec<_>>()
.join(","),
};
if !normalized.split(',').all(valid_page_token) {
return Err(OcrRequestError::Pages("invalid native page range".into()));
}
Ok(Some(normalized))
}
fn valid_page_token(token: &str) -> bool {
let mut parts = token.split('-');
let start = parts.next().unwrap_or_default();
if start.is_empty() || !start.chars().all(|character| character.is_ascii_digit()) {
return false;
}
match parts.next() {
None => true,
Some(end) => {
!end.is_empty()
&& end.chars().all(|character| character.is_ascii_digit())
&& parts.next().is_none()
}
}
}
fn normalize_features(features: FeaturesInput) -> Result<Option<String>, OcrRequestError> {
let tokens = match features {
FeaturesInput::Names(names) => names,
FeaturesInput::CommaSeparated(names) => names.split(',').map(str::to_string).collect(),
};
if tokens.is_empty() {
return Ok(None);
}
let normalized = tokens.iter().map(|token| token.trim()).collect::<Vec<_>>();
if !normalized.iter().all(|token| {
let Some((first, rest)) = token.as_bytes().split_first() else {
return false;
};
first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric)
}) {
return Err(OcrRequestError::Features);
}
Ok(Some(normalized.join(",")))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
fn map(value: Value) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let fields = value.as_object().unwrap().clone();
map_ocr_params(decode_input_params(fields, "optional_params")?.known)
}
#[test]
fn input_params_retain_unknown_fields() {
let parsed = decode_input_params(
json!({
"pages": [0],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
})
.as_object()
.unwrap()
.clone(),
"optional_params",
)
.unwrap();
assert_eq!(
parsed.known.pages,
Some(PagesInput::ZeroBasedIndices(vec![0]))
);
assert_eq!(parsed.extra_params["future_ocr_option"], true);
assert_eq!(
parsed.extra_params["extra_body"],
json!({"provider_option": "value"})
);
assert_eq!(
serde_json::to_value(map_ocr_params(parsed.known).unwrap()).unwrap(),
json!({"pages": "1", "features": null})
);
}
#[rstest]
#[case(json!([0, 1, 2]), Some("1,2,3"))]
#[case(json!([2, 0, 0, 1]), Some("1,2,3"))]
#[case(json!([]), None)]
#[case(json!("3-9"), Some("3-9"))]
#[case(json!("1-3, 5"), Some("1-3,5"))]
#[case(json!(["1", "3-5"]), Some("1,3-5"))]
fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) {
assert_eq!(
map(json!({"pages": input})).unwrap().pages.as_deref(),
expected
);
}
#[rstest]
#[case(json!("a,b"))]
#[case(json!([-1]))]
#[case(json!([true, false]))]
#[case(json!([1, "2"]))]
#[case(json!(5))]
fn invalid_page_mapping_matches_python(#[case] input: Value) {
assert!(map(json!({"pages": input})).is_err());
}
#[rstest]
#[case(json!(["keyValuePairs"]), "keyValuePairs")]
#[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")]
#[case(json!("keyValuePairs"), "keyValuePairs")]
#[case(json!("keyValuePairs,languages"), "keyValuePairs,languages")]
#[case(json!("keyValuePairs, languages"), "keyValuePairs,languages")]
fn feature_mapping_matches_python(#[case] input: Value, #[case] expected: &str) {
assert_eq!(
map(json!({"features": input})).unwrap().features.as_deref(),
Some(expected)
);
}
#[rstest]
#[case(json!("keyValuePairs&pages=9"))]
#[case(json!("key value pairs"))]
#[case(json!(""))]
#[case(json!([1, 2]))]
#[case(json!([["keyValuePairs"]]))]
#[case(json!({"feature":"keyValuePairs"}))]
#[case(json!(5))]
fn invalid_feature_mapping_matches_python(#[case] input: Value) {
assert!(map(json!({"features": input})).is_err());
}
#[test]
fn empty_feature_list_is_omitted() {
assert_eq!(map(json!({"features": []})).unwrap().features, None);
}
}

107
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs
View file

@ -1,107 +0,0 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use serde_json::{Map, Value, json};
use super::types::*;
use crate::constants::{AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
pub(crate) fn transform_ocr_request(
document: OcrDocument,
) -> Result<DocumentIntelligenceRequest, OcrRequestError> {
let source = document.source();
if source.is_empty() {
return Err(OcrRequestError::MissingDocumentUrl);
}
Ok(if let Some(document) = InlineDocument::parse(source)? {
DocumentIntelligenceRequest::Base64Source(
STANDARD.encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?),
)
} else {
DocumentIntelligenceRequest::UrlSource(source.to_string())
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: AzureDocumentIntelligenceOperation,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
if response.status != Some(OperationStatus::Succeeded) {
return Err(OcrResponseError::OperationStatus(
response
.status
.map(|status| status.to_string())
.unwrap_or_else(|| "None".into()),
));
}
let result = response.analyze_result.unwrap_or_default();
let pages = result
.pages
.into_iter()
.map(normalize_page)
.collect::<Result<Vec<_>, _>>()?;
let pages_processed = pages.len();
let mut extra_fields = Map::new();
extra_fields.insert("content".into(), option_value(result.content));
extra_fields.insert("tables".into(), option_value(result.tables));
extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs));
Ok(LiteLLMOcrResponse {
pages,
model: model.into(),
document_annotation: None,
usage_info: Some(json!({"pages_processed":pages_processed})),
object: "ocr".into(),
extra_fields,
provider_native_response: None,
})
}
fn normalize_page(page: AzureDocumentIntelligencePage) -> Result<Value, OcrResponseError> {
let index = page
.page_number
.unwrap_or(1)
.checked_sub(1)
.ok_or(OcrResponseError::NumericRange("page.pageNumber"))?;
let scale = if page.unit.as_deref().unwrap_or("inch") == "inch" {
AZURE_DI_DEFAULT_DPI as f64
} else {
1.0
};
let width = pixel_dimension(
page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH),
scale,
"page.width",
)?;
let height = pixel_dimension(
page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT),
scale,
"page.height",
)?;
let markdown = page
.lines
.iter()
.map(|line| line.content.as_deref().unwrap_or_default())
.collect::<Vec<_>>()
.join("\n");
Ok(json!({
"index":index,
"markdown":markdown,
"images":null,
"dimensions":{"width":width,"height":height,"dpi":AZURE_DI_DEFAULT_DPI}
}))
}
fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result<i64, OcrResponseError> {
let value = value * scale;
if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 {
return Err(OcrResponseError::NumericRange(field));
}
Ok(value.trunc() as i64)
}
fn option_value<T: serde::Serialize>(value: Option<T>) -> Value {
value
.and_then(|value| serde_json::to_value(value).ok())
.unwrap_or(Value::Null)
}

138
litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs
View file

@ -1,138 +0,0 @@
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum PagesInput {
ZeroBasedIndices(Vec<i64>),
NativeTokens(Vec<String>),
NativeRange(String),
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum FeaturesInput {
Names(Vec<String>),
CommaSeparated(String),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct DocumentIntelligenceInputParams {
pub pages: Option<PagesInput>,
pub features: Option<FeaturesInput>,
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub(crate) struct DocumentIntelligenceParams {
pub pages: Option<String>,
pub features: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) enum DocumentIntelligenceRequest {
#[serde(rename = "urlSource")]
UrlSource(String),
#[serde(rename = "base64Source")]
Base64Source(String),
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum OperationStatus {
Succeeded,
Running,
NotStarted,
Failed,
Unknown(String),
}
impl<'de> Deserialize<'de> for OperationStatus {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
Ok(match String::deserialize(deserializer)?.as_str() {
"succeeded" => Self::Succeeded,
"running" => Self::Running,
"notStarted" => Self::NotStarted,
"failed" => Self::Failed,
value => Self::Unknown(value.to_string()),
})
}
}
impl std::fmt::Display for OperationStatus {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(match self {
Self::Succeeded => "succeeded",
Self::Running => "running",
Self::NotStarted => "notStarted",
Self::Failed => "failed",
Self::Unknown(value) => value,
})
}
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceOperation {
pub status: Option<OperationStatus>,
#[serde(rename = "analyzeResult")]
pub analyze_result: Option<AzureDocumentIntelligenceAnalyzeResult>,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceAnalyzeResult {
pub content: Option<String>,
#[serde(default)]
pub pages: Vec<AzureDocumentIntelligencePage>,
pub tables: Option<Vec<Map<String, Value>>>,
#[serde(rename = "keyValuePairs")]
pub key_value_pairs: Option<Vec<Map<String, Value>>>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligencePage {
#[serde(rename = "pageNumber", default, deserialize_with = "optional_i64")]
pub page_number: Option<i64>,
#[serde(default, deserialize_with = "optional_f64")]
pub width: Option<f64>,
#[serde(default, deserialize_with = "optional_f64")]
pub height: Option<f64>,
pub unit: Option<String>,
#[serde(default)]
pub lines: Vec<AzureDocumentIntelligenceLine>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceLine {
pub content: Option<String>,
}
fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<i64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_i64()
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected an integer")),
Some(Value::String(value)) => value
.parse::<i64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected an integer")),
Some(_) => Err(serde::de::Error::custom("expected an integer")),
}
}
fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<f64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_f64()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(Value::String(value)) => value
.parse::<f64>()
.ok()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(_) => Err(serde::de::Error::custom("expected a number")),
}
}

5
litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs
View file

@ -1,5 +0,0 @@
mod transformation;
mod types;
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse};

250
litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs
View file

@ -1,250 +0,0 @@
use super::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse};
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
pub(crate) fn transform_ocr_request(
model: &str,
document: OcrDocument,
params: &MistralOcrParams,
) -> Result<MistralOcrRequest, OcrRequestError> {
Ok(MistralOcrRequest {
model: model.to_string(),
document,
params: params.clone(),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: MistralOcrResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
Ok(LiteLLMOcrResponse {
pages: response.pages,
model: response.model.unwrap_or_else(|| model.to_string()),
document_annotation: response.document_annotation,
usage_info: response.usage_info,
object: "ocr".to_string(),
extra_fields: response.extra_fields,
provider_native_response: None,
})
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use serde_json::{Value, json};
fn mapped_params(value: Value) -> Value {
serde_json::to_value(serde_json::from_value::<MistralOcrParams>(value).unwrap()).unwrap()
}
fn document() -> OcrDocument {
serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.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_drops_unknown_params() {
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]
#[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("confidence_scores_granularity", json!("word"))]
#[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("include_image_base64", json!(true))]
#[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("confidence_scores_granularity", json!("word"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
let params: MistralOcrParams =
serde_json::from_value(json!({name: value.clone()})).unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document(), &params).unwrap())
.unwrap();
assert_eq!(result["model"], "model");
assert_eq!(result[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(
#[case] name: &str,
#[case] value: Value,
) {
let params: MistralOcrParams = serde_json::from_value(json!({name:value.clone()})).unwrap();
let result = serde_json::to_value(
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() {
let params: MistralOcrParams = serde_json::from_value(json!({
"table_format":"html",
"confidence_scores_granularity":"page",
"extract_header":true
}))
.unwrap();
let result = serde_json::to_value(
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 response: MistralOcrResponse = serde_json::from_value(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}
}))
.unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["blocks"][0]["type"], "title");
assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1);
assert_eq!(
result["pages"][0]["blocks"][0]["confidence_scores"]["mean"],
0.98
);
assert_eq!(
result["pages"][0]["confidence_scores"]["average_page_confidence_score"],
0.99
);
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 = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0], page);
}
}

60
litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs
View file

@ -1,60 +0,0 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::ocr::types::OcrDocument;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum MistralOcrPages {
Range(String),
Indices(Vec<i64>),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct MistralOcrParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<MistralOcrPages>,
#[serde(skip_serializing_if = "Option::is_none")]
pub include_image_base64: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image_limit: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image_min_size: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub bbox_annotation_format: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation_format: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation_prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub extract_header: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub extract_footer: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub table_format: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub confidence_scores_granularity: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub include_blocks: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct MistralOcrRequest {
pub model: String,
pub document: OcrDocument,
#[serde(flatten)]
pub params: MistralOcrParams,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct MistralOcrResponse {
#[serde(default)]
pub pages: Vec<Value>,
pub model: Option<String>,
pub document_annotation: Option<Value>,
pub usage_info: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}

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

@ -1,5 +0,0 @@
pub(crate) mod cohere;
pub(crate) mod deepseek;
pub(crate) mod document_intelligence;
pub(crate) mod mistral;
pub(crate) mod reducto;

9
litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs
View file

@ -1,9 +0,0 @@
mod transformation;
mod types;
pub(crate) use transformation::{
transform_legacy_ocr_request, transform_ocr_response, transform_v3_ocr_request,
};
pub(crate) use types::{
ReductoLegacyParams, ReductoResponse, ReductoUploadResponse, ReductoV3Params,
};

103
litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs
View file

@ -1,103 +0,0 @@
use std::collections::BTreeMap;
use serde_json::{Value, json};
use super::types::*;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
pub(crate) fn transform_v3_ocr_request(
_model: &str,
document: OcrDocument,
params: &ReductoV3Params,
) -> Result<ReductoV3Request, OcrRequestError> {
Ok(ReductoV3Request {
input: document.source().to_string(),
params: params.clone(),
})
}
pub(crate) fn transform_legacy_ocr_request(
_model: &str,
document: OcrDocument,
params: &ReductoLegacyParams,
) -> Result<ReductoLegacyRequest, OcrRequestError> {
Ok(ReductoLegacyRequest {
document_url: document.source().to_string(),
options: params.enhance.as_ref().map(|_| params.clone()),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: ReductoResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let result = match response.result {
Some(result) => result.unwrap_or_default(),
None => ReductoResult {
chunks: response.chunks,
},
};
let usage = response.usage.unwrap_or_default();
Ok(LiteLLMOcrResponse {
pages: build_pages(result.chunks.unwrap_or_default()),
model: model.to_string(),
document_annotation: None,
usage_info: Some(json!({
"pages_processed": usage.num_pages,
"credits": usage.credits,
})),
object: "ocr".to_string(),
extra_fields: serde_json::Map::new(),
provider_native_response: None,
})
}
fn build_pages(chunks: Vec<ReductoChunk>) -> Vec<Value> {
let blocks_by_page = chunks
.iter()
.flat_map(|chunk| chunk.blocks.iter().flatten())
.filter_map(|block| block.bbox.as_ref()?.page.map(|page| (page, block)))
.fold(
BTreeMap::<i64, Vec<&ReductoBlock>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = join_content(chunks.iter().map(|chunk| chunk.content.as_deref()));
return if markdown.is_empty() {
Vec::new()
} else {
vec![page(0, markdown, None)]
};
}
blocks_by_page
.into_iter()
.map(|(index, blocks)| {
let markdown = join_content(blocks.iter().map(|block| block.content.as_deref()));
page(
index.saturating_sub(1).max(0),
markdown,
Some(json!(blocks)),
)
})
.collect()
}
fn join_content<'a>(content: impl Iterator<Item = Option<&'a str>>) -> String {
content
.flatten()
.filter(|text| !text.is_empty())
.collect::<Vec<_>>()
.join("\n\n")
}
fn page(index: i64, markdown: String, blocks: Option<Value>) -> Value {
let mut result = json!({"index":index,"markdown":markdown,"images":null});
if let (Value::Object(fields), Some(blocks)) = (&mut result, blocks) {
fields.insert("blocks".into(), blocks);
}
result
}

128
litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs
View file

@ -1,128 +0,0 @@
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct ReductoV3Params {
#[serde(skip_serializing_if = "Option::is_none")]
pub formatting: Option<Map<String, Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub retrieval: Option<Map<String, Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub settings: Option<Map<String, Value>>,
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub enhance: Option<Map<String, Value>>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoV3Request {
pub input: String,
#[serde(flatten)]
pub params: ReductoV3Params,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyRequest {
pub document_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub options: Option<ReductoLegacyParams>,
}
#[derive(Deserialize)]
pub(crate) struct ReductoUploadResponse {
pub file_id: Option<String>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct ReductoResponse {
#[serde(default, deserialize_with = "present_nullable")]
pub result: Option<Option<ReductoResult>>,
pub usage: Option<ReductoUsage>,
#[serde(default)]
pub chunks: Option<Vec<ReductoChunk>>,
}
fn present_nullable<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
deserializer: D,
) -> Result<Option<Option<T>>, D::Error> {
Option::<T>::deserialize(deserializer).map(Some)
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct ReductoResult {
pub chunks: Option<Vec<ReductoChunk>>,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct ReductoUsage {
#[serde(default, deserialize_with = "optional_i64")]
pub num_pages: Option<i64>,
#[serde(default, deserialize_with = "optional_f64")]
pub credits: Option<f64>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct ReductoChunk {
pub content: Option<String>,
pub blocks: Option<Vec<ReductoBlock>>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoBlock {
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub bbox: Option<ReductoBoundingBox>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoBoundingBox {
#[serde(default, deserialize_with = "optional_i64")]
pub page: Option<i64>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<i64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_i64()
.or_else(|| number.as_f64().and_then(checked_truncated_i64))
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected an integer")),
Some(Value::String(value)) => value
.trim()
.parse::<i64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected an integer")),
Some(Value::Bool(value)) => Ok(Some(i64::from(value))),
Some(_) => Ok(None),
}
}
fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<f64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_f64()
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a number")),
Some(Value::String(value)) => value
.trim()
.parse::<f64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected a number")),
Some(_) => Ok(None),
}
}
fn checked_truncated_i64(value: f64) -> Option<i64> {
(value.is_finite() && value >= i64::MIN as f64 && value <= i64::MAX as f64)
.then(|| value.trunc() as i64)
}

50
litellm-rust/crates/core/src/ocr/document.rs
View file

@ -1,3 +1,4 @@
use std::collections::BTreeMap as Map;
use std::io::Read;
use std::path::Path;
@ -5,9 +6,10 @@ use base64::{Engine, engine::general_purpose::STANDARD};
use data_url::mime::Mime;
use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
use reqwest::Url;
use serde_json::Map;
use super::error::{OcrError, OcrRequestError, OcrResponseError};
use super::Error as OcrError;
use super::Error as OcrRequestError;
use super::Error as OcrResponseError;
use super::types::{OcrConnection, OcrDocument, OcrDocumentInput};
use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS};
use crate::media::Error as MediaError;
@ -47,11 +49,10 @@ pub fn read_path_document(
})
.map_err(|source| super::Error::FileRead {
path: path.to_owned(),
kind: source.kind(),
message: source.to_string(),
source: std::sync::Arc::new(source),
})?;
let name = path.file_name().map(|name| name.to_string_lossy());
Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?)
encode_file_document(&bytes, name.as_deref(), mime_type)
}
pub fn encode_file_document(
@ -164,7 +165,7 @@ pub(crate) async fn inline_remote_document(
connection: &OcrConnection,
) -> Result<OcrDocument, OcrError> {
let source = document.source();
if !source.starts_with("http://") && !source.starts_with("https://") {
if !document.is_remote() {
validate_inline_document(&document)?;
return Ok(document);
}
@ -193,12 +194,12 @@ pub(crate) async fn inline_remote_document(
fn map_media_error(error: MediaError) -> OcrError {
match error {
MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl.into(),
MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled.into(),
MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge.into(),
MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects.into(),
MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation.into(),
MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect.into(),
MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl,
MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled,
MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge,
MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects,
MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation,
MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect,
MediaError::Http(status) => TransportError::Http {
status,
body: "OCR document download failed".into(),
@ -215,8 +216,9 @@ fn map_media_error(error: MediaError) -> OcrError {
#[cfg(test)]
mod tests {
use std::collections::BTreeMap as Map;
use super::*;
use serde_json::Map;
fn document(source: &str) -> OcrDocument {
OcrDocument::DocumentUrl {
@ -286,17 +288,17 @@ mod tests {
document("data:application/pdf;base64,YWJj")
);
std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap();
assert_eq!(
assert!(matches!(
prepare_document(OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
}),
Err(OcrRequestError::InlineDocumentTooLarge.into())
);
Err(OcrRequestError::InlineDocumentTooLarge)
));
std::fs::remove_dir_all(&dir).unwrap();
let missing = dir.join("missing.pdf");
let Err(super::super::Error::FileRead { path, kind, .. }) =
let Err(super::super::Error::FileRead { path, source, .. }) =
prepare_document(OcrDocumentInput::Path {
path: missing.clone(),
mime_type: None,
@ -305,7 +307,7 @@ mod tests {
panic!("missing paths must surface a file read error");
};
assert_eq!(path, missing);
assert_eq!(kind, std::io::ErrorKind::NotFound);
assert_eq!(source.kind(), std::io::ErrorKind::NotFound);
}
#[test]
@ -325,10 +327,10 @@ mod tests {
#[test]
fn file_encoding_enforces_decoded_size_limit() {
let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1];
assert_eq!(
assert!(matches!(
encode_file_document(&bytes, None, None),
Err(OcrRequestError::InlineDocumentTooLarge)
);
));
let document = encode_file_document(&bytes[..OCR_INLINE_MAX_BYTES], None, None).unwrap();
let inline = InlineDocument::parse(document.source()).unwrap().unwrap();
assert_eq!(
@ -359,10 +361,10 @@ mod tests {
] {
let inline = InlineDocument::parse(source).unwrap().unwrap();
assert_eq!(inline.decode(expected.len()).unwrap(), expected);
assert_eq!(
assert!(matches!(
inline.decode(expected.len() - 1),
Err(OcrRequestError::InlineDocumentTooLarge)
);
));
}
}
@ -427,7 +429,7 @@ mod tests {
client.document_fetcher(),
OcrDocument::ImageUrl {
image_url: format!("http://{address}/image"),
extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
},
&OcrConnection::default(),
)
@ -439,7 +441,7 @@ mod tests {
converted,
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,YWJj".into(),
extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
}
);
assert!(!request.to_ascii_lowercase().contains("authorization"));

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