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

refactor(rust): move OCR provider code into litellm-llms and delete core/src/llms

Browse source 
OCR transformations, BaseOcrConfig with its response and connection types,
the OCR error, and the HTTP pieces (custom_httpx: http_handler, transport,
media, llm_http_handler with OcrClient and the request/response handler)
now live in litellm-llms at their Python paths. Provider code no longer
reaches into the route: it gets the caller's hooks through a route-neutral
CallHooks trait that core implements over its host, and core dispatches to
llm_http_handler::ocr with the concrete config, the way Python calls
base_llm_http_handler.ocr(provider_config=...).

Core keeps the route: entrypoint, request types, credential fallback,
provider dispatch, the machine and hook glue. ocr/mod.rs no longer
re-exports anything, provider constants moved next to their only users,
and provider tests that drive the whole route moved to core's route test
files. Twenty-two of those were exact copies of tests already there and
were dropped; every one still runs once.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-09-17 22:28:49 -07:00
parent 49c50739d7
commit 5a76346047
102 changed files with 3753 additions and 4167 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm-rust/Cargo.lock generated
View file

@ -2027,12 +2027,9 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"bytes",
"data-url",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-callbacks",
"litellm-core-utils",
"litellm-llms",
@ -2047,8 +2044,6 @@ dependencies = [
"rustls-native-certs",
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"sha2 0.10.9",
"strum",
"subtle",
@ -2065,7 +2060,6 @@ name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"litellm-types",
"rstest",
"serde",
"serde_json",
"serde_path_to_error",
@ -2111,15 +2105,22 @@ dependencies = [
"aws-smithy-types",
"base64 0.22.1",
"bytes",
"data-url",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-callbacks",
"litellm-core-utils",
"litellm-framing",
"litellm-types",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"thiserror 2.0.19",
"time",
"tokio",
@ -2137,6 +2138,7 @@ dependencies = [
"litellm-callbacks-legacy",
"litellm-core",
"litellm-host-python",
"litellm-llms",
"litellm-token-counter",
"litellm-types",
"pyo3",

3
litellm-rust/crates/core-utils/Cargo.toml
View file

@ -13,6 +13,3 @@ serde_path_to_error = "0.1"
serde_with.workspace = true
thiserror.workspace = true
url.workspace = true
[dev-dependencies]
rstest.workspace = true

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

@ -1,27 +1,14 @@
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 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
## Crate layering
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business. Env reads are limited to credential fallback in a route's `prepare.rs`.
Each crate mirrors one top-level Python package, so a Rust path reads as its Python path with the crate name in place of the package directory. Dependencies only point down:
Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.
- `litellm-types` mirrors `litellm/types/`: pure serde data, no I/O
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments), no network I/O
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `custom_httpx/` (HTTP helpers and the OCR request handler)
- `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks
## Python/Rust transformation pairs
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::custom_httpx::llm_http_handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
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
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business.

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

@ -13,11 +13,8 @@ litellm-callbacks.workspace = true
bytes.workspace = true
futures-util.workspace = true
base64.workspace = true
data-url = "0.3.2"
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-llms.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
@ -27,8 +24,6 @@ rustls.workspace = true
rustls-native-certs.workspace = true
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_with.workspace = true
serde_path_to_error = "0.1"
strum.workspace = true
subtle.workspace = true
tokio = { workspace = true, features = ["sync"] }
@ -40,5 +35,6 @@ url.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true

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

@ -20,9 +20,9 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,10 +1,8 @@
use litellm_llms::custom_httpx::http_handler::{http_request, truncate_error_body};
use serde_json::Value;
use super::{Error, client::http_client};
use crate::{
audio_transcription::types::ProviderAudioTranscriptionRequest,
http_utils::{http_request, truncate_error_body},
};
use crate::audio_transcription::types::ProviderAudioTranscriptionRequest;
pub async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
@ -19,19 +17,24 @@ pub async fn execute_audio_transcription_provider_call(
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))?;
let response = http_request(request_builder).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
})?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))?;
let text = response.text().await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
})?;
if !status.is_success() {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
let response_json = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;

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

@ -4,12 +4,12 @@ use litellm_llms::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
},
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
custom_httpx::http_handler::{has_header, string_headers},
};
use super::Error;
use crate::{
audio_transcription::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest},
http_utils::{has_header, string_headers},
use crate::audio_transcription::types::{
AudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
};
fn provider_config(provider: &str) -> Option<&'static dyn BaseAudioTranscriptionConfig> {

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

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

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

@ -20,9 +20,9 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,11 +1,13 @@
use litellm_llms::base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData};
use litellm_llms::{
base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData},
custom_httpx::http_handler::{http_request, truncate_error_body},
};
use litellm_types::utils::ChatCompletionsResponse;
use serde_json::Value;
use super::{Error, client::http_client, prepare::prepare_provider_request};
use crate::{
chat_completions::types::{ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest},
http_utils::{http_request, truncate_error_body},
use crate::chat_completions::types::{
ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest,
};
pub(super) async fn execute_chat_completions_provider_call(
@ -32,23 +34,30 @@ pub(super) async fn execute_chat_completions_provider_call(
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
Error::Transport(crate::transport::Error::Connect(err.to_string()))
Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(
err.to_string(),
))
} else {
Error::Transport(crate::transport::Error::Network(err.to_string()))
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
}
})?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
if !status.is_success() {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
@ -73,7 +82,9 @@ pub(super) async fn execute_chat_completions_provider_call(
pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ (Error::InvalidResponse(_)
| Error::Transport(crate::transport::Error::Http { .. })) => already,
| Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
..
})) => already,
other => Error::InvalidResponse(other.to_string()),
}
}

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

@ -1,5 +1,8 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
use litellm_llms::{
base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth},
custom_httpx::http_handler::has_header,
};
use litellm_types::llms::openai::ChatMessage;
use serde_json::Value;
@ -7,11 +10,8 @@ use super::{
Error,
common_utils::{chat_completions_provider_config, string_headers},
};
use crate::{
chat_completions::types::{
ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest,
},
http_utils::has_header,
use crate::chat_completions::types::{
ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest,
};
pub(super) fn resolve_provider_config<'a>(

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

@ -265,7 +265,7 @@ fn rejects_non_string_extra_headers() {
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
decline(call),
Error::Headers(crate::http_utils::HeaderError {
Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
@ -771,7 +771,10 @@ mod round_trip {
assert!(
matches!(
err,
Error::Transport(crate::transport::Error::Http { status: 429, .. })
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status: 429,
..
})
),
"expected a 429, got {err:?}"
);
@ -796,7 +799,10 @@ mod round_trip {
.await
.expect_err("nothing is listening");
assert!(
matches!(err, Error::Transport(crate::transport::Error::Connect(_))),
matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_))
),
"expected a pre-send connect failure, got {err:?}"
);
}
@ -819,11 +825,16 @@ mod round_trip {
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(Error::Transport(crate::transport::Error::Http {
as_response_error(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: 500,
body: "boom".to_string()
}
)),
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status: 500,
body: "boom".to_string()
})),
Error::Transport(crate::transport::Error::Http { status: 500, .. })
..
})
));
}
}

35
litellm-rust/crates/core/src/constants.rs
View file

@ -8,10 +8,6 @@ pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
/// Connect timeout for Anthropic Messages provider calls, in seconds.
pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
/// Max characters of an upstream error body echoed across the call boundary
/// before truncation, so provider bodies are bounded and data-minimized.
pub(crate) const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256;
/// Provider name used for Anthropic Messages when a deployment's provider model
/// does not carry an explicit provider prefix.
pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
@ -27,34 +23,3 @@ pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600;
/// `object` field every non-streaming chat completion response carries.
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
/// Placeholder Python substitutes for empty or whitespace-only message text,
/// which Anthropic and Bedrock both reject. Must match
/// `_EMPTY_TEXT_PLACEHOLDER` in
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10;
pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120;
pub(crate) const OCR_POLL_RETRY_SECS: u64 = 2;
pub(crate) const AZURE_DI_API_VERSION: &str = "2024-11-30";
pub(crate) const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
pub(crate) const AZURE_DI_DEFAULT_DPI: i64 = 96;
pub(crate) const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
pub(crate) const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
pub(crate) const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://";
pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr";
pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";
pub(crate) const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com";
pub(crate) const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY";

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

@ -1,7 +1,9 @@
use litellm_llms::base_llm::ocr::error::Error as OcrError;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error(transparent)]
Ocr(#[from] crate::ocr::Error),
Ocr(#[from] OcrError),
#[error(transparent)]
Messages(#[from] crate::messages::Error),
#[error(transparent)]

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

@ -2,13 +2,9 @@ pub mod audio_transcription;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
pub mod llms;
pub mod machine;
mod media;
pub mod messages;
pub mod ocr;
pub mod responses;
pub mod transport;
pub use error::Error;

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

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

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

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

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

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

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

@ -1,619 +0,0 @@
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl};
use serde_json::Value;
use crate::{
constants::AZURE_AI_OCR_PATH,
llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext},
mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest},
},
ocr::{
OcrClient,
document::{inline_remote_document, validate_inline_document},
prepare::credential_env,
types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest},
},
};
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 serde_json::json;
use crate::ocr::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, 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")
);
}
#[tokio::test]
async fn rejects_non_inline_body_after_guardrails() {
let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| {
wire.body["document"] = json!({
"type":"document_url",
"document_url":"https://example.com/not-inline.pdf"
});
Ok(wire)
});
let error = perform_ocr_with(host).await.unwrap_err();
assert!(error.to_string().contains("data URI"));
}
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use litellm_auth::{
ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle,
};
use crate::ocr::{LiteLLMOcrRequest, test_support::header, wire::decode_request};
#[derive(Debug)]
struct CountingToken {
token: fn(usize) -> String,
calls: AtomicUsize,
}
impl CountingToken {
fn new(token: fn(usize) -> String) -> Arc<Self> {
Arc::new(Self {
token,
calls: AtomicUsize::new(0),
})
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
}
impl TokenProvider for CountingToken {
fn acquire(&self) -> TokenFuture<'_> {
let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
let token = SecretValue::new((self.token)(call));
Box::pin(async move {
Ok(ResolvedCredential::AccessToken {
token,
expires_on: None,
})
})
}
}
fn numbered_token(call: usize) -> String {
format!("callback-{call}")
}
fn azure_request(
provider: &Arc<CountingToken>,
api_base: Option<&str>,
api_key: Option<&str>,
extra_headers: Value,
optional_params: Value,
) -> LiteLLMOcrRequest {
let wire = serde_json::from_value(json!({
"model": "azure_ai/mistral-ocr-latest",
"document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": null,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": 2.0
}))
.unwrap();
LiteLLMOcrRequest {
azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())),
..decode_request(wire).unwrap()
}
}
fn ocr_page() -> MockResponse {
MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]}))
}
#[tokio::test]
async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() {
let provider = CountingToken::new(numbered_token);
let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await;
for _ in 0..2 {
perform_ocr(azure_request(
&provider,
Some(&base),
None,
Value::Null,
json!({}),
))
.await
.unwrap();
}
server.await.unwrap();
assert_eq!(provider.calls(), 2);
let requests = seen.lock().unwrap();
assert_eq!(
requests
.iter()
.map(|request| header(request, "authorization"))
.collect::<Vec<_>>(),
[Some("Bearer callback-1"), Some("Bearer callback-2")]
);
}
#[rstest]
#[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)]
#[case::provider_beats_static_token(
None,
Value::Null,
json!({"azure_ad_token":"static-token"}),
"Bearer callback-1",
1
)]
#[case::header_wins_on_the_wire_but_provider_still_runs(
None,
json!({"Authorization":"Bearer override"}),
json!({}),
"Bearer override",
1
)]
#[tokio::test]
async fn credential_precedence(
#[case] api_key: Option<&str>,
#[case] extra_headers: Value,
#[case] optional_params: Value,
#[case] expected_authorization: &str,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(numbered_token);
let (base, seen, server) = mock_server(vec![ocr_page()]).await;
perform_ocr(azure_request(
&provider,
Some(&base),
api_key,
extra_headers,
optional_params,
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(provider.calls(), expected_calls);
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(
header(&requests[0], "authorization"),
Some(expected_authorization)
);
}
#[rstest]
#[case::missing_api_base(
false,
json!({}),
numbered_token,
|error: &crate::ocr::Error| matches!(error, crate::ocr::Error::Auth(litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: AZURE_AI_API_BASE_ENV,
})),
0
)]
#[case::unsupported_oidc_reference(
true,
json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}),
numbered_token,
|error: &crate::ocr::Error| matches!(error, crate::ocr::Error::Auth(litellm_auth::Error::UnsupportedOidcReference)),
0
)]
#[case::empty_provider_token_ignores_static_token(
true,
json!({"azure_ad_token":"static-token"}),
|_| String::new(),
|error: &crate::ocr::Error| matches!(error, crate::ocr::Error::MissingAzureAiCredentials),
1
)]
#[tokio::test]
async fn credential_failures_send_no_provider_request(
#[case] with_api_base: bool,
#[case] optional_params: Value,
#[case] token: fn(usize) -> String,
#[case] expected: fn(&crate::ocr::Error) -> bool,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(token);
let (base, seen, server) = mock_server(vec![ocr_page()]).await;
let error = perform_ocr(azure_request(
&provider,
with_api_base.then_some(base.as_str()),
None,
Value::Null,
optional_params,
))
.await
.unwrap_err();
server.abort();
assert!(expected(&error), "unexpected error: {error:?}");
assert_eq!(provider.calls(), expected_calls);
assert!(seen.lock().unwrap().is_empty());
}
#[tokio::test]
async fn environment_supplies_api_base_and_bearer_key() {
let env = |name: &str| match name {
AZURE_AI_API_BASE_ENV => Some("https://env.example".to_string()),
AZURE_AI_API_KEY_ENV => Some("env-key".to_string()),
_ => None,
};
let connection = OcrConnection::default();
let headers = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &env)
.await
.unwrap();
let url = AzureAiOcrConfig.build_ocr_url(None, &env).unwrap();
assert_eq!(
headers,
[("Authorization".to_string(), "Bearer env-key".to_string())]
);
assert_eq!(url, "https://env.example/providers/mistral/azure/ocr");
}
}

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

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

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

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

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

@ -1,211 +0,0 @@
use std::future::Future;
use litellm_core_utils::call_arguments::CallArguments;
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use crate::ocr::{
OcrClient,
route::OcrHost,
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 host: &'a OcrHost,
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::emit_response_received(context.host, &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
View file

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

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

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

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

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

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

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

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

@ -1,6 +0,0 @@
pub mod azure_ai;
pub mod base_llm;
pub(crate) mod cohere;
pub(crate) mod mistral;
pub(crate) mod reducto;
pub(crate) mod vertex_ai;

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

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

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

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

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

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

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

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

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

@ -1,414 +0,0 @@
use litellm_auth_gcp::{self as vertex, VertexConfig};
use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl};
use serde_json::Value;
use super::common_utils::validate_destination;
use crate::{
llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrEnvironment, OcrRequestContext},
mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest},
},
ocr::{
OcrClient,
document::{inline_remote_document, validate_inline_document},
prepare::credential_env,
types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest},
},
};
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 rstest::rstest;
use super::VertexAiOcrConfig;
#[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,
mistral::ocr::transformation::MistralOcrConfig,
vertex_ai::ocr::transformation::VertexAiOcrConfig,
},
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_for_test(
crate::ocr::test_support::resolved_request(direct),
);
let vertex = crate::ocr::prepare::prepare_request_for_test(
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");
}
}

24
litellm-rust/crates/core/src/machine/mod.rs
View file

@ -37,7 +37,7 @@ struct PendingOp<R: Route> {
/// The provider side of the machine: how the in-flight call reaches its host.
pub struct HostChannel<R: Route> {
ops: Option<mpsc::UnboundedSender<PendingOp<R>>>,
ops: mpsc::UnboundedSender<PendingOp<R>>,
}
impl<R: Route> Clone for HostChannel<R> {
@ -48,24 +48,14 @@ impl<R: Route> Clone for HostChannel<R> {
}
}
impl<R: Route> HostChannel<R> {
/// A channel with no host behind it: the wire request goes out unchanged, events go
/// nowhere, and route operations fail. For tests that prepare a request without
/// driving it.
#[cfg(test)]
pub(crate) fn detached() -> Self {
Self { ops: None }
}
}
impl<R: Route> HostChannel<R>
where
R::Error: From<MachineFault>,
{
async fn invoke(&self, op: HostOp<R>) -> Result<HostResult<R>, R::Error> {
let ops = self.ops.as_ref().ok_or(MachineFault::Abandoned)?;
let (reply, answer) = oneshot::channel();
ops.send(PendingOp { op, reply })
self.ops
.send(PendingOp { op, reply })
.map_err(|_| MachineFault::Abandoned)?;
answer.await.map_err(|_| MachineFault::Abandoned.into())
}
@ -82,9 +72,6 @@ where
wire: WireRequest,
context: RequestContext,
) -> Result<WireRequest, R::Error> {
if self.ops.is_none() {
return Ok(wire);
}
let op = HostOp::BeforeSend {
wire: Box::new(wire),
context: Box::new(context),
@ -96,9 +83,6 @@ where
}
pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> {
if self.ops.is_none() {
return Ok(());
}
match self.invoke(HostOp::Emit(event)).await? {
HostResult::Emitted => Ok(()),
_ => Err(MachineFault::Mismatch.into()),
@ -128,7 +112,7 @@ where
Self {
execution: Execution::Unstarted(Box::new(execute)),
ops,
channel: HostChannel { ops: Some(ops_tx) },
channel: HostChannel { ops: ops_tx },
reply: None,
}
}

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

@ -1,13 +1,15 @@
pub(super) use litellm_llms::custom_httpx::http_handler::{
has_bearer_auth, has_header, truncate_error_body,
};
use litellm_llms::{
anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
custom_httpx::http_handler::string_headers as shared_string_headers,
};
use serde_json::{Map, Value};
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};
const HEADER_CONTEXT: &str = "messages";

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

@ -15,9 +15,9 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
}
impl From<LlmError> for Error {

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

@ -1,13 +1,11 @@
use litellm_llms::custom_httpx::http_handler::http_request;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use super::{
Error, client::http_client, common_utils::truncate_error_body,
prepare::prepare_provider_request,
};
use crate::{
constants::ANTHROPIC_MESSAGES_PROVIDER, http_utils::http_request,
messages::types::MessagesRequest,
};
use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, messages::types::MessagesRequest};
pub(super) async fn execute_messages_provider_call(
request: MessagesRequest<'_>,
@ -21,21 +19,26 @@ pub(super) async fn execute_messages_provider_call(
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
let response = http_request(request_builder).await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
if !status.is_success() {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
let response = serde_json::from_str(&text)
@ -62,19 +65,24 @@ pub(super) async fn execute_messages_provider_stream(
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
let response = http_request(request_builder).await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
let status = response.status();
if !status.is_success() {
let text = response
.text()
.await
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
Ok(response)
}

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

@ -82,7 +82,7 @@ fn string_headers_rejects_non_string_values() {
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::Headers(crate::http_utils::HeaderError {
Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
@ -432,7 +432,7 @@ async fn messages_maps_provider_error_status_to_http_error() {
assert!(matches!(
err,
Error::Transport(crate::transport::Error::Http { status: 401, .. })
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { status: 401, .. })
));
}

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

@ -1,4 +1,5 @@
use litellm_core_utils::call_arguments::ArgumentSpec;
use litellm_llms::base_llm::ocr::error::Error;
use super::provider_config::{OcrConfigKind, resolve_provider_config};
@ -30,7 +31,7 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b
pub fn consumed_optional_param_names(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Vec<&'static str>, super::Error> {
) -> Result<Vec<&'static str>, 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 {
@ -62,7 +63,7 @@ pub(crate) fn is_secret_param(name: &str) -> bool {
pub fn consumed_optional_params(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Vec<ArgumentSpec>, super::Error> {
) -> Result<Vec<ArgumentSpec>, Error> {
consumed_optional_param_names(model, custom_llm_provider).map(|names| {
names
.into_iter()

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

@ -1,176 +1,20 @@
use std::{sync::OnceLock, time::Duration};
use bytes::{Bytes, BytesMut};
use litellm_auth_gcp::VertexAuth;
use serde::de::DeserializeOwned;
use super::{
json::{DecodedOcrResponse, decode_response},
types::{LiteLLMOcrRequest, LiteLLMOcrResponse},
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::OcrClient,
};
use crate::{constants::OCR_CONNECT_TIMEOUT_SECS, media::MediaFetcher};
#[derive(Clone)]
pub struct OcrClient {
provider_http: reqwest::Client,
polling_http: reqwest::Client,
document_fetcher: MediaFetcher,
vertex_auth: VertexAuth,
use crate::ocr::{
route::{LocalOcrHost, ocr_machine},
types::LiteLLMOcrRequest,
};
pub async fn perform(
client: &OcrClient,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
}
impl OcrClient {
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()?,
document_fetcher,
vertex_auth: VertexAuth::default(),
})
}
pub fn shared() -> Result<Self, crate::ocr::Error> {
shared_client()
}
pub async fn perform(
&self,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
litellm_callbacks::run::run(
super::ocr_machine(self.clone()),
&super::LocalOcrHost::new(request),
)
.await
}
pub(crate) fn provider_http(&self) -> &reqwest::Client {
&self.provider_http
}
pub(crate) fn polling_http(&self) -> &reqwest::Client {
&self.polling_http
}
pub(crate) fn document_fetcher(&self) -> &MediaFetcher {
&self.document_fetcher
}
pub(crate) fn vertex_auth(&self) -> &VertexAuth {
&self.vertex_auth
}
#[cfg(test)]
pub(crate) fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
Self {
provider_http,
polling_http: no_redirect_http().expect("test polling client builds"),
document_fetcher: MediaFetcher::for_test(document_http),
vertex_auth: VertexAuth::default(),
}
}
}
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(crate::transport::Error::from)
}
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(crate::transport::Error::from)
.and_then(OcrClient::new)
})
.clone()?;
Ok(client)
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
shared_client()?.perform(request).await
}
pub async fn read_json_response<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
max_response_bytes: usize,
) -> Result<DecodedOcrResponse<T>, crate::ocr::Error> {
let bytes = read_response_bytes(response, max_response_bytes).await?;
decode_response(&bytes, native)
}
pub(crate) async fn read_response_bytes(
mut response: reqwest::Response,
limit: usize,
) -> Result<Bytes, crate::ocr::Error> {
let status = response.status();
if status.is_success()
&& response
.content_length()
.is_some_and(|length| length > limit as u64)
{
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(crate::ocr::Error::TooLarge { limit });
}
bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
if !status.is_success() && bytes.len() == limit {
break;
}
}
if !status.is_success() {
return Err(crate::transport::Error::Http {
status: status.as_u16(),
body: String::from_utf8_lossy(&bytes).into_owned(),
}
.into());
}
Ok(bytes.freeze())
}
pub(crate) fn transport_error(error: reqwest::Error) -> crate::ocr::Error {
if error.is_timeout() {
return crate::ocr::Error::Transport(crate::transport::Error::Http {
status: 408,
body: "OCR request timed out".into(),
});
}
crate::transport::Error::from(error).into()
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_timeout_has_an_http_408_status() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let _connection = listener.accept().await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
});
let error = reqwest::Client::new()
.get(format!("http://{address}"))
.timeout(Duration::from_millis(10))
.send()
.await
.unwrap_err();
assert!(matches!(
transport_error(error),
crate::ocr::Error::Transport(crate::transport::Error::Http { status: 408, .. })
));
server.abort();
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
perform(&OcrClient::shared()?, request).await
}

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

@ -1,20 +1,14 @@
use std::{collections::BTreeMap as Map, io::Read, path::Path};
use base64::{Engine, engine::general_purpose::STANDARD};
use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime};
use reqwest::Url;
use super::{
Error as OcrError, Error as OcrRequestError, Error as OcrResponseError,
types::{OcrConnection, OcrDocument, OcrDocumentInput},
};
use crate::{
constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS},
media::{DownloadPolicy, Error as MediaError, MediaFetcher},
transport::Error as TransportError,
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{OCR_INLINE_MAX_BYTES, OcrDocument},
};
pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, super::Error> {
use crate::ocr::types::OcrDocumentInput;
pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
match input {
OcrDocumentInput::Document(document) => Ok(document),
OcrDocumentInput::Path { path, mime_type } => {
@ -29,23 +23,20 @@ pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, super::E
file_name.as_deref(),
mime_type.as_deref(),
)?),
OcrDocumentInput::HostReader { .. } => Err(super::Error::InvalidRequest(
OcrDocumentInput::HostReader { .. } => Err(Error::InvalidRequest(
"OCR file reader was not read by the host".into(),
)),
}
}
pub fn read_path_document(
path: &Path,
mime_type: Option<&str>,
) -> Result<OcrDocument, super::Error> {
pub fn read_path_document(path: &Path, mime_type: Option<&str>) -> Result<OcrDocument, Error> {
let mut bytes = Vec::new();
std::fs::File::open(path)
.and_then(|file| {
file.take(OCR_INLINE_MAX_BYTES as u64 + 1)
.read_to_end(&mut bytes)
})
.map_err(|source| super::Error::FileRead {
.map_err(|source| Error::FileRead {
path: path.to_owned(),
source: std::sync::Arc::new(source),
})?;
@ -57,17 +48,17 @@ pub fn encode_file_document(
bytes: &[u8],
file_name: Option<&str>,
mime_type: Option<&str>,
) -> Result<OcrDocument, OcrRequestError> {
) -> Result<OcrDocument, Error> {
if bytes.is_empty() {
return Err(OcrRequestError::EmptyFile);
return Err(Error::EmptyFile);
}
if bytes.len() > OCR_INLINE_MAX_BYTES {
return Err(OcrRequestError::InlineDocumentTooLarge);
return Err(Error::InlineDocumentTooLarge);
}
if let Some(value) = mime_type
&& !valid_mime_type(value)
{
return Err(OcrRequestError::InvalidMimeType(value.into()));
return Err(Error::InvalidMimeType(value.into()));
}
let mime_type = mime_type
.map(str::to_string)
@ -117,105 +108,12 @@ pub fn mime_type_for_name(name: &str) -> &'static str {
}
}
pub(crate) struct InlineDocument<'a>(DataUrl<'a>);
impl<'a> InlineDocument<'a> {
pub(crate) fn parse(source: &'a str) -> Result<Option<Self>, OcrRequestError> {
match DataUrl::process(source) {
Ok(url) => Ok(Some(Self(url))),
Err(DataUrlError::NotADataUrl) => Ok(None),
Err(DataUrlError::NoComma) => Err(OcrRequestError::InvalidDataUri),
}
}
pub(crate) fn mime_type(&self) -> &Mime {
self.0.mime_type()
}
pub(crate) fn decode(&self, max_bytes: usize) -> Result<Vec<u8>, OcrRequestError> {
let mut body = Vec::new();
self.0
.decode(|bytes| {
if bytes.len() > max_bytes.saturating_sub(body.len()) {
return Err(OcrRequestError::InlineDocumentTooLarge);
}
body.extend_from_slice(bytes);
Ok(())
})
.map_err(|error| match error {
DecodeError::InvalidBase64(_) => OcrRequestError::InvalidDataUri,
DecodeError::WriteError(error) => error,
})?;
Ok(body)
}
}
pub(crate) fn validate_inline_document(document: &OcrDocument) -> Result<(), OcrRequestError> {
let inline =
InlineDocument::parse(document.source())?.ok_or(OcrRequestError::InvalidDataUri)?;
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
Ok(())
}
pub(crate) async fn inline_remote_document(
fetcher: &MediaFetcher,
document: OcrDocument,
connection: &OcrConnection,
) -> Result<OcrDocument, OcrError> {
let source = document.source();
if !document.is_remote() {
validate_inline_document(&document)?;
return Ok(document);
}
let url = Url::parse(source).map_err(|_| OcrRequestError::RequestField {
path: "document URL".into(),
})?;
let downloaded = fetcher
.fetch(
url,
DownloadPolicy {
timeout: connection.timeout,
max_bytes: connection.max_download_bytes,
max_redirects: OCR_MAX_FETCH_REDIRECTS,
},
)
.await
.map_err(map_media_error)?;
let result = document.with_source(format!(
"data:{};base64,{}",
downloaded.content_type,
STANDARD.encode(downloaded.bytes)
));
validate_inline_document(&result)?;
Ok(result)
}
fn map_media_error(error: MediaError) -> OcrError {
match error {
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(),
}
.into(),
MediaError::Timeout => TransportError::Http {
status: 408,
body: "OCR document download timed out".into(),
}
.into(),
MediaError::Transport(error) => error.into(),
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap as Map;
use litellm_llms::base_llm::ocr::document::InlineDocument;
use super::*;
fn document(source: &str) -> OcrDocument {
@ -291,12 +189,12 @@ mod tests {
path: path.clone(),
mime_type: None,
}),
Err(OcrRequestError::InlineDocumentTooLarge)
Err(Error::InlineDocumentTooLarge)
));
std::fs::remove_dir_all(&dir).unwrap();
let missing = dir.join("missing.pdf");
let Err(super::super::Error::FileRead { path, source, .. }) =
let Err(super::Error::FileRead { path, source, .. }) =
prepare_document(OcrDocumentInput::Path {
path: missing.clone(),
mime_type: None,
@ -327,7 +225,7 @@ mod tests {
let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1];
assert!(matches!(
encode_file_document(&bytes, None, None),
Err(OcrRequestError::InlineDocumentTooLarge)
Err(Error::InlineDocumentTooLarge)
));
let document = encode_file_document(&bytes[..OCR_INLINE_MAX_BYTES], None, None).unwrap();
let inline = InlineDocument::parse(document.source()).unwrap().unwrap();
@ -349,102 +247,4 @@ mod tests {
assert!(encode_file_document(b"abc", None, Some(mime)).is_err());
}
}
#[test]
fn decodes_data_urls_and_limits_decoded_size() {
for (source, expected) in [
("data:application/pdf;base64,YWJj", b"abc".as_slice()),
("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()),
("data:,a%20b%00%FF", b"a b\0\xff".as_slice()),
] {
let inline = InlineDocument::parse(source).unwrap().unwrap();
assert_eq!(inline.decode(expected.len()).unwrap(), expected);
assert!(matches!(
inline.decode(expected.len() - 1),
Err(OcrRequestError::InlineDocumentTooLarge)
));
}
}
#[test]
fn preserves_mime_parameters_and_standard_default() {
let inline = InlineDocument::parse("data:application/pdf;version=1.7;base64,YQ==")
.unwrap()
.unwrap();
assert!(inline.mime_type().matches("application", "pdf"));
assert_eq!(inline.mime_type().get_parameter("version"), Some("1.7"));
let default = InlineDocument::parse("data:,a").unwrap().unwrap();
assert!(default.mime_type().matches("text", "plain"));
assert_eq!(
default.mime_type().get_parameter("charset"),
Some("US-ASCII")
);
}
#[test]
fn rejects_invalid_inline_documents() {
for source in [
"https://example.com/document.pdf",
"data:application/pdf;base64",
"data:application/pdf;base64,INVALID!",
] {
assert!(validate_inline_document(&document(source)).is_err());
}
}
#[tokio::test]
async fn remote_conversion_preserves_kind_and_isolates_provider_credentials() {
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = vec![0_u8; 2048];
let count = socket.read(&mut request).await.unwrap();
socket
.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: image/png; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc")
.await
.unwrap();
String::from_utf8_lossy(&request[..count]).into_owned()
});
let mut provider_headers = reqwest::header::HeaderMap::new();
provider_headers.insert(
reqwest::header::AUTHORIZATION,
reqwest::header::HeaderValue::from_static("Bearer provider-secret"),
);
let provider_http = reqwest::Client::builder()
.default_headers(provider_headers)
.build()
.unwrap();
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.unwrap();
let client = super::super::OcrClient::for_test(provider_http, document_http);
let converted = inline_remote_document(
client.document_fetcher(),
OcrDocument::ImageUrl {
image_url: format!("http://{address}/image"),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
},
&OcrConnection::default(),
)
.await
.unwrap();
let request = server.await.unwrap();
assert_eq!(
converted,
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,YWJj".into(),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
}
);
assert!(!request.to_ascii_lowercase().contains("authorization"));
assert!(!request.contains("provider-secret"));
}
}

156
litellm-rust/crates/core/src/ocr/handler.rs
View file

@ -1,117 +1,81 @@
use litellm_callbacks::event::{CallEvent, RawResponse};
use futures_util::future::BoxFuture;
use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest};
use litellm_llms::{
base_llm::ocr::{
error::Error,
transformation::{LiteLLMOcrResponse, PreparedOcrRequest},
},
custom_httpx::llm_http_handler::{CallHooks, OcrClient},
};
use serde_json::Value;
use super::{
OcrClient,
arguments::is_secret_param, prepare::prepare_request, provider_config::OcrConfigKind,
route::OcrHost,
types::{LiteLLMOcrResponse, PreparedOcrRequest, ResolvedOcrRequest},
};
use crate::llms::base_llm::ocr::transformation::OcrResponseContext;
use crate::ocr::types::ResolvedOcrRequest;
pub(crate) async fn perform_ocr_request(
client: &OcrClient,
request: ResolvedOcrRequest,
host: &OcrHost,
caller_document: bool,
) -> Result<LiteLLMOcrResponse, super::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
request.response_format()?;
PreparedOcrCall::prepare(client.clone(), request, host, caller_document)
.await?
.execute()
.await
let config = request.config;
let request = prepare_request(request, caller_document);
let hooks = OcrCallHooks::new(host.clone(), &request, config);
config.ocr(client, &request, &hooks).await
}
pub(crate) struct PreparedOcrCall {
client: OcrClient,
request: PreparedOcrRequest,
http: reqwest::Request,
/// Lets provider code reach the host mid-call, filling in the request context only the
/// route knows.
pub(crate) struct OcrCallHooks {
host: OcrHost,
model: String,
custom_llm_provider: &'static str,
optional_params: Value,
secret_fields: Vec<String>,
}
impl PreparedOcrCall {
pub(crate) async fn prepare(
client: OcrClient,
request: ResolvedOcrRequest,
host: &OcrHost,
caller_document: bool,
) -> Result<Self, super::Error> {
let request = super::prepare::prepare_request(request, host.clone(), caller_document);
let http = request.config.prepare_request(&request, &client).await?;
Ok(Self {
client,
request,
http,
})
}
pub(crate) async fn execute(self) -> Result<LiteLLMOcrResponse, super::Error> {
let url = self.http.url().to_string();
let headers = request_headers(&self.http)?;
let response =
crate::http_utils::execute_http_request(self.client.provider_http(), self.http)
.await
.map_err(super::client::transport_error)?;
if !response.status().is_success() {
let headers = response
.headers()
.iter()
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.to_string(), value.to_string()))
})
.collect();
return match super::client::read_response_bytes(
response,
self.request.connection.max_response_bytes,
)
.await
{
Err(super::Error::Transport(crate::transport::Error::Http { status, body })) => {
Err(self.request.config.get_error_class(body, status, headers))
}
Err(error) => Err(error),
Ok(_) => unreachable!("non-success response produces an HTTP error"),
};
impl OcrCallHooks {
pub(crate) fn new(host: OcrHost, request: &PreparedOcrRequest, config: OcrConfigKind) -> Self {
Self {
host,
model: request.model.clone(),
custom_llm_provider: config.provider().into(),
optional_params: Value::Object(request.optional_params.clone().into()),
secret_fields: request
.optional_params
.keys()
.filter(|name| is_secret_param(name))
.cloned()
.collect(),
}
let model = &self.request.model;
let context = OcrResponseContext {
client: &self.client,
connection: &self.request.connection,
host: &self.request.host,
request_format: self.request.response_format()?,
url: &url,
headers: &headers,
};
self.request
.config
.async_transform_ocr_response(model, response, context)
.await
}
}
fn request_headers(request: &reqwest::Request) -> Result<Vec<(String, String)>, super::Error> {
request
.headers()
.iter()
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.to_string(), value.to_string()))
.map_err(|_| super::Error::RequestField {
path: "headers".into(),
})
})
.collect()
}
impl CallHooks<Error> for OcrCallHooks {
fn before_send(
&self,
wire: WireRequest,
passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
let context = RequestContext {
model: self.model.clone(),
custom_llm_provider: self.custom_llm_provider.into(),
optional_params: self.optional_params.clone(),
passthrough_fields,
secret_fields: self.secret_fields.clone(),
};
Box::pin(self.host.before_send(wire, context))
}
pub(crate) async fn emit_response_received(
host: &OcrHost,
bytes: &[u8],
) -> Result<(), super::Error> {
host.emit(CallEvent::ResponseReceived {
raw: RawResponse {
body: String::from_utf8_lossy(bytes).into_owned(),
},
})
.await
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(self.host.emit(CallEvent::ResponseReceived {
raw: RawResponse {
body: String::from_utf8_lossy(body).into_owned(),
},
}))
}
}

62
litellm-rust/crates/core/src/ocr/json.rs
View file

@ -1,62 +0,0 @@
use serde::de::{DeserializeOwned, IntoDeserializer};
use serde_json::{Map, Value};
#[derive(Debug)]
pub struct DecodedOcrResponse<T> {
pub data: T,
pub native: Option<Map<String, Value>>,
pub text: String,
}
pub(crate) fn decode_request_value<T: DeserializeOwned>(
value: Value,
prefix: &str,
) -> Result<T, crate::ocr::Error> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
crate::ocr::Error::RequestField {
path: format!("{prefix}.{}", error.path()),
}
})
}
pub(crate) fn decode_response_value<T: DeserializeOwned>(
value: Value,
prefix: &str,
) -> Result<T, crate::ocr::Error> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
crate::ocr::Error::ResponseField {
path: format!("{prefix}.{}", error.path()),
}
})
}
pub(crate) fn decode_response<T: DeserializeOwned>(
bytes: &[u8],
native: bool,
) -> Result<DecodedOcrResponse<T>, crate::ocr::Error> {
let mut deserializer = serde_json::Deserializer::from_slice(bytes);
let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| {
crate::ocr::Error::ResponseField {
path: error.path().to_string(),
}
})?;
deserializer
.end()
.map_err(|_| crate::ocr::Error::ResponseField {
path: "response".into(),
})?;
let native = if native {
Some(
serde_json::from_slice(bytes).map_err(|_| crate::ocr::Error::ResponseField {
path: "response".into(),
})?,
)
} else {
None
};
Ok(DecodedOcrResponse {
data,
native,
text: String::from_utf8_lossy(bytes).into_owned(),
})
}

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

@ -1,29 +1,13 @@
mod arguments;
pub mod arguments;
pub mod client;
pub(crate) mod document;
pub mod error;
pub use error::Error;
pub mod document;
pub(crate) mod handler;
pub(crate) mod json;
pub(crate) mod prepare;
mod provider_config;
pub mod provider_config;
pub mod route;
pub mod types;
pub mod wire;
pub use arguments::{
consumed_optional_param_names, consumed_optional_params, is_supported_request,
};
pub use client::{OcrClient, ocr};
pub use document::{encode_file_document, mime_type_for_name, read_path_document};
pub use provider_config::{get_api_key_env_var, get_health_check_document};
pub use route::{LocalOcrHost, Ocr, OcrHost, OcrMachine, OcrOp, OcrOpResult, ocr_machine};
pub use types::{
LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrConnectionInputs, OcrCredentialInputs,
OcrDocument, OcrDocumentInput, OcrFileContent, OcrPage, OcrPageDimensions, OcrPageImage,
OcrTransportConfig, OcrUsageInfo,
};
#[cfg(test)]
#[path = "../../tests/azure_ai_ocr.rs"]
mod azure_ai_tests;
@ -31,6 +15,9 @@ mod azure_ai_tests;
#[path = "../../tests/azure_document_intelligence_ocr.rs"]
mod azure_document_intelligence_tests;
#[cfg(test)]
#[path = "../../tests/cohere_ocr.rs"]
mod cohere_tests;
#[cfg(test)]
#[path = "../../tests/deepseek_ocr.rs"]
mod deepseek_tests;
#[cfg(test)]

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

@ -1,158 +1,20 @@
use litellm_callbacks::event::{Passthrough, RequestContext, WireRequest};
use serde::Serialize;
use serde_json::{Map, Value};
use super::{
OcrClient,
route::OcrHost,
types::{OcrConnection, OcrDocument, PreparedOcrRequest, ResolvedOcrRequest},
use litellm_auth::{InputSource, Sourced};
use litellm_llms::base_llm::ocr::transformation::{
OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env,
};
pub(crate) async fn transform_request_body<B>(
client: &OcrClient,
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: B,
validate: impl Fn(&Value) -> Result<(), super::Error>,
) -> Result<reqwest::Request, super::Error>
where
B: Serialize,
{
let composed = litellm_core_utils::call_arguments::compose_body(
&request.optional_params,
&body,
request.config.get_supported_ocr_params(&request.model),
)?;
validate(&composed)?;
let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed);
let changed = request
.host
.before_send(
wire_request(url, headers, composed),
request_context(request, passthrough_fields),
)
.await?;
if !changed.body.is_object() {
return Err(super::Error::RequestField {
path: "guardrail.body".into(),
});
}
validate(&changed.body)?;
build_http_request(client, request, url, &changed.headers, &changed.body)
}
fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest {
WireRequest {
url: url.into(),
headers: headers.to_vec(),
body,
}
}
fn caller_inputs(request: &PreparedOcrRequest) -> Result<Map<String, Value>, super::Error> {
let document = request
.caller_document
.then(|| serde_json::to_value(&request.document))
.transpose()
.map_err(|_| super::Error::RequestField {
path: "document".into(),
})?;
let params: Map<String, Value> = request.optional_params.clone().into();
Ok(params
.into_iter()
.chain(document.map(|document| ("document".to_string(), document)))
.collect())
}
fn request_context(
request: &PreparedOcrRequest,
passthrough_fields: Passthrough,
) -> RequestContext {
RequestContext {
model: request.model.clone(),
custom_llm_provider: request.provider_name().into(),
optional_params: Value::Object(request.optional_params.clone().into()),
passthrough_fields,
secret_fields: request
.optional_params
.keys()
.filter(|name| super::arguments::is_secret_param(name))
.cloned()
.collect(),
}
}
pub(crate) fn build_http_request<B: Serialize>(
client: &OcrClient,
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: &B,
) -> Result<reqwest::Request, super::Error> {
let builder = client
.provider_http()
.post(url)
.json(body)
.timeout(request.connection.timeout);
crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All)
.build()
.map_err(crate::transport::Error::from)
.map_err(super::Error::from)
}
pub(crate) async fn guardrail_document(
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
) -> Result<(OcrDocument, Vec<(String, String)>), super::Error> {
let body = serde_json::to_value(&request.document).map_err(|_| super::Error::RequestField {
path: "document".into(),
})?;
let changed = request
.host
.before_send(
wire_request(url, headers, body),
request_context(request, Passthrough::default()),
)
.await?;
let document = super::json::decode_request_value(changed.body, "guardrail.document")?;
Ok((document, changed.headers))
}
pub(crate) fn body_document(body: &Value) -> Result<OcrDocument, super::Error> {
let document = body
.get("document")
.and_then(Value::as_object)
.ok_or_else(|| super::Error::RequestField {
path: "body.document".into(),
})?;
let source = document
.iter()
.filter(|(name, _)| matches!(name.as_str(), "type" | "image_url" | "document_url"))
.map(|(name, value)| (name.clone(), value.clone()))
.collect();
super::json::decode_request_value(Value::Object(source), "body.document")
}
pub(crate) fn credential_env(name: &str) -> Option<String> {
std::env::var(name).ok()
}
use super::provider_config::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) fn prepare_request(
request: ResolvedOcrRequest,
host: OcrHost,
caller_document: bool,
) -> PreparedOcrRequest {
use litellm_auth::{InputSource, Sourced};
let credentials = request.credentials.clone();
let api_base_env = match request.config.provider() {
super::provider_config::OcrProvider::Mistral => Some("MISTRAL_API_BASE"),
super::provider_config::OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"),
super::provider_config::OcrProvider::Cohere
| super::provider_config::OcrProvider::Reducto
| super::provider_config::OcrProvider::VertexAi => None,
OcrProvider::Mistral => Some("MISTRAL_API_BASE"),
OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"),
OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => None,
};
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
@ -172,23 +34,34 @@ pub(crate) fn prepare_request(
});
let resolved = request
.config
.resolve_connection_params(super::types::OcrCredentialInputs {
.resolve_connection_params(OcrCredentialInputs {
dynamic_api_key,
dynamic_api_base,
..credentials
});
let transport = request.transport.clone();
PreparedOcrRequest::new(
request,
OcrConnection::new(resolved, transport),
host,
let LiteLLMOcrRequest {
model,
document,
transport,
optional_params,
input_sources,
azure_ad_token_provider,
..
} = request;
PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport),
caller_document,
)
optional_params,
input_sources,
azure_ad_token_provider,
}
}
#[cfg(test)]
pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest {
prepare_request(request, OcrHost::detached(), true)
prepare_request(request, true)
}
#[cfg(test)]

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

@ -1,46 +1,66 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use strum::{EnumString, IntoStaticStr};
use super::{
OcrClient,
types::{
LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest,
ResolvedOcrCredentials,
},
};
use crate::llms::{
use litellm_llms::{
azure_ai::ocr::{
cohere_parse_transformation::AzureAICohereParseConfig,
document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig,
transformation::AzureAiOcrConfig,
},
base_llm::ocr::transformation::{BaseOcrConfig, OcrResponseContext},
base_llm::ocr::{
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument,
PreparedOcrRequest, ResolvedOcrCredentials,
},
},
cohere::ocr::transformation::CohereParseConfig,
custom_httpx::llm_http_handler::{self, CallHooks, OcrClient},
mistral::ocr::transformation::MistralOcrConfig,
reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config},
vertex_ai::ocr::{
deepseek_transformation::VertexAIDeepSeekOCRConfig, transformation::VertexAiOcrConfig,
},
};
use strum::{EnumString, IntoStaticStr};
macro_rules! dispatch_config {
($config:expr, $method:ident($($argument:expr),* $(,)?)) => {
dispatch_config!(@arms $config, $method($($argument),*), )
};
($config:expr, $method:ident($($argument:expr),* $(,)?).await) => {
dispatch_config!(@arms $config, $method($($argument),*), .await)
};
(@arms $config:expr, $method:ident($($argument:expr),*), $($suffix:tt)*) => {
match $config {
OcrConfigKind::Cohere => CohereParseConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::Mistral => MistralOcrConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::AzureAi => AzureAiOcrConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::AzureCohere => AzureAICohereParseConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::AzureDocumentIntelligence => AzureDocumentIntelligenceOcrConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::ReductoLegacy => ReductoParseLegacyConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::ReductoV3 => ReductoParseV3Config.$method($($argument),*)$($suffix)*,
OcrConfigKind::VertexAi => VertexAiOcrConfig.$method($($argument),*)$($suffix)*,
OcrConfigKind::VertexDeepSeek => VertexAIDeepSeekOCRConfig.$method($($argument),*)$($suffix)*,
macro_rules! with_config {
($kind:expr, $config:ident => $body:expr) => {
match $kind {
OcrConfigKind::Cohere => {
let $config = CohereParseConfig;
$body
}
OcrConfigKind::Mistral => {
let $config = MistralOcrConfig;
$body
}
OcrConfigKind::AzureAi => {
let $config = AzureAiOcrConfig;
$body
}
OcrConfigKind::AzureCohere => {
let $config = AzureAICohereParseConfig;
$body
}
OcrConfigKind::AzureDocumentIntelligence => {
let $config = AzureDocumentIntelligenceOcrConfig;
$body
}
OcrConfigKind::ReductoLegacy => {
let $config = ReductoParseLegacyConfig;
$body
}
OcrConfigKind::ReductoV3 => {
let $config = ReductoParseV3Config;
$body
}
OcrConfigKind::VertexAi => {
let $config = VertexAiOcrConfig;
$body
}
OcrConfigKind::VertexDeepSeek => {
let $config = VertexAIDeepSeekOCRConfig;
$body
}
}
};
}
@ -72,58 +92,38 @@ impl OcrConfigKind {
}
pub(crate) fn get_supported_ocr_params(self, model: &str) -> &'static [&'static str] {
dispatch_config!(self, get_supported_ocr_params(model))
with_config!(self, config => config.get_supported_ocr_params(model))
}
pub(crate) fn get_api_key_env_var(self) -> Option<&'static str> {
dispatch_config!(self, get_api_key_env_var())
with_config!(self, config => config.get_api_key_env_var())
}
pub(crate) fn get_health_check_document(self) -> OcrDocument {
dispatch_config!(self, get_health_check_document())
with_config!(self, config => config.get_health_check_document())
}
pub(crate) fn resolve_connection_params(
self,
inputs: OcrCredentialInputs,
) -> ResolvedOcrCredentials {
dispatch_config!(self, resolve_connection_params(inputs))
with_config!(self, config => config.resolve_connection_params(inputs))
}
pub(crate) fn get_error_class(
pub(crate) async fn ocr(
self,
message: String,
status: u16,
headers: Vec<(String, String)>,
) -> super::Error {
dispatch_config!(self, get_error_class(message, status, headers))
}
pub(crate) async fn prepare_request(
self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, super::Error> {
dispatch_config!(self, prepare_request(request, client).await)
}
pub(crate) async fn async_transform_ocr_response(
self,
model: &str,
raw_response: reqwest::Response,
context: OcrResponseContext<'_>,
) -> Result<LiteLLMOcrResponse, super::Error> {
dispatch_config!(
self,
async_transform_ocr_response(model, raw_response, context).await
)
request: &PreparedOcrRequest,
hooks: &dyn CallHooks<Error>,
) -> Result<LiteLLMOcrResponse, Error> {
with_config!(self, config => llm_http_handler::ocr(&config, client, request, hooks).await)
}
}
pub fn get_api_key_env_var(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Option<&'static str>, super::Error> {
) -> Result<Option<&'static str>, Error> {
Ok(resolve_provider_config(model, custom_llm_provider)?
.1
.get_api_key_env_var())
@ -132,7 +132,7 @@ pub fn get_api_key_env_var(
pub fn get_health_check_document(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<OcrDocument, super::Error> {
) -> Result<OcrDocument, Error> {
Ok(resolve_provider_config(model, custom_llm_provider)?
.1
.get_health_check_document())
@ -151,7 +151,7 @@ pub(crate) enum OcrProvider {
pub(crate) fn resolve_provider_config(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<(String, OcrConfigKind), super::Error> {
) -> Result<(String, OcrConfigKind), Error> {
let provider =
get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider {
model,
@ -160,7 +160,7 @@ pub(crate) fn resolve_provider_config(
let ocr_provider = provider
.custom_llm_provider
.parse::<OcrProvider>()
.map_err(|_| super::Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
.map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
let config = match ocr_provider {
OcrProvider::Cohere => OcrConfigKind::Cohere,
OcrProvider::Mistral => OcrConfigKind::Mistral,
@ -194,6 +194,9 @@ fn is_document_intelligence_model(model: &str) -> bool {
#[cfg(test)]
mod tests {
use litellm_auth::{InputSource, Sourced};
use litellm_llms::{
base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document,
};
use rstest::rstest;
use super::*;
@ -216,7 +219,7 @@ mod tests {
fn invalid_provider_names_are_rejected(#[case] provider: &str) {
assert!(matches!(
resolve_provider_config("model", Some(provider)),
Err(crate::ocr::Error::InvalidProvider(value)) if value == provider
Err(Error::InvalidProvider(value)) if value == provider
));
}
@ -230,9 +233,7 @@ mod tests {
fn pdf_health_check_documents_are_valid(#[case] model: &str) {
let document = get_health_check_document(model, None).unwrap();
assert!(matches!(document, OcrDocument::DocumentUrl { .. }));
let inline = crate::ocr::document::InlineDocument::parse(document.source())
.unwrap()
.unwrap();
let inline = InlineDocument::parse(document.source()).unwrap().unwrap();
assert_eq!(inline.mime_type().to_string(), "application/pdf");
assert!(inline.decode(4096).unwrap().starts_with(b"%PDF-"));
}
@ -242,10 +243,8 @@ mod tests {
#[case("azure_ai/cohere-parse")]
fn png_health_check_documents_are_valid(#[case] model: &str) {
let document = get_health_check_document(model, None).unwrap();
crate::llms::cohere::ocr::validate_document(&document).unwrap();
let inline = crate::ocr::document::InlineDocument::parse(document.source())
.unwrap()
.unwrap();
validate_document(&document).unwrap();
let inline = InlineDocument::parse(document.source()).unwrap().unwrap();
assert_eq!(inline.mime_type().to_string(), "image/png");
assert!(
inline
@ -433,9 +432,7 @@ mod tests {
#[case] provider: Option<&str>,
) {
let error = resolve_provider_config(model, provider).unwrap_err();
assert!(
matches!(&error, crate::ocr::Error::InvalidProvider(provider) if provider == "not_a_provider")
);
assert!(matches!(&error, Error::InvalidProvider(provider) if provider == "not_a_provider"));
assert_eq!(error.http_status_code(), Some(400));
}
}

15
litellm-rust/crates/core/src/ocr/route.rs
View file

@ -5,13 +5,16 @@ use litellm_callbacks::{
event::{CallEvent, RequestContext, WireRequest},
route::Route,
};
use super::{
Error, LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient,
handler::perform_ocr_request,
types::{OcrDocumentInput, OcrFileContent, ResolvedOcrRequest},
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::OcrClient,
};
use super::handler::perform_ocr_request;
use crate::{
machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute},
ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest},
};
use crate::machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrOp {

461
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -1,69 +1,17 @@
use std::{collections::BTreeMap, path::PathBuf, time::Duration};
use bytes::Bytes;
use litellm_auth::{InputSource, Sourced, TokenProviderHandle};
use litellm_core_utils::{
call_arguments::CallArguments,
serde_compat::{FiniteF64, LaxI64},
use litellm_auth::{InputSource, TokenProviderHandle};
use litellm_core_utils::call_arguments::CallArguments;
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{
OcrCredentialInputs, OcrDocument, OcrResponseFormat, OcrTransportConfig, response_format,
},
};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use serde_with::serde_as;
use super::provider_config::{OcrConfigKind, resolve_provider_config};
use crate::constants::OCR_HTTP_TIMEOUT_SECS;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum OcrDocument {
#[serde(rename = "document_url")]
DocumentUrl {
document_url: String,
#[serde(flatten)]
extra_fields: BTreeMap<String, Option<String>>,
},
#[serde(rename = "image_url")]
ImageUrl {
image_url: String,
#[serde(flatten)]
extra_fields: BTreeMap<String, Option<String>>,
},
}
impl OcrDocument {
pub(crate) fn source(&self) -> &str {
match self {
Self::DocumentUrl { document_url, .. } => document_url,
Self::ImageUrl { image_url, .. } => image_url,
}
}
pub(crate) fn is_remote(&self) -> bool {
let source = self.source();
source.starts_with("http://") || source.starts_with("https://")
}
pub(crate) fn with_source(self, source: String) -> Self {
match self {
Self::DocumentUrl { extra_fields, .. } => Self::DocumentUrl {
document_url: source,
extra_fields,
},
Self::ImageUrl { extra_fields, .. } => Self::ImageUrl {
image_url: source,
extra_fields,
},
}
}
}
impl TryFrom<Value> for OcrDocument {
type Error = super::Error;
fn try_from(value: Value) -> Result<Self, Self::Error> {
super::json::decode_request_value(value, "document")
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum OcrDocumentInput {
@ -103,83 +51,6 @@ pub struct OcrFileContent {
pub file_name: Option<String>,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum OcrResponseFormat {
#[default]
Litellm,
Native,
}
#[derive(Clone, Default)]
pub struct OcrCredentialInputs {
pub api_key: Option<Sourced<String>>,
pub dynamic_api_key: Option<Sourced<String>>,
pub api_base: Option<Sourced<String>>,
pub dynamic_api_base: Option<Sourced<String>>,
}
impl OcrCredentialInputs {
pub fn new(
api_key: Option<String>,
api_key_source: InputSource,
api_base: Option<String>,
api_base_source: InputSource,
) -> Self {
Self {
api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)),
dynamic_api_key: None,
api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)),
dynamic_api_base: None,
}
}
}
#[derive(Clone)]
pub struct OcrTransportConfig {
pub extra_headers: Vec<(String, String)>,
pub extra_headers_source: InputSource,
pub timeout: Duration,
pub max_download_bytes: u64,
pub max_response_bytes: usize,
pub poll_timeout: Duration,
}
impl Default for OcrTransportConfig {
fn default() -> Self {
Self {
extra_headers: Vec::new(),
extra_headers_source: InputSource::Deployment,
timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES,
max_response_bytes: crate::constants::OCR_RESPONSE_MAX_BYTES,
poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS),
}
}
}
impl OcrTransportConfig {
pub fn with_overrides(
self,
extra_headers: Vec<(String, String)>,
extra_headers_source: InputSource,
timeout: Option<Duration>,
) -> Self {
Self {
extra_headers,
extra_headers_source,
timeout: timeout.unwrap_or(self.timeout),
..self
}
}
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
/// Caller-supplied connection overrides for a [`LiteLLMOcrRequest`], in the
/// shape hosts receive them: JSON-ish headers, optional timeout, optional
/// credentials, and per-field provenance in `input_sources`.
@ -197,14 +68,14 @@ impl OcrConnectionInputs {
self.input_sources.get(name).copied().unwrap_or_default()
}
fn header_pairs(&self) -> Result<Vec<(String, String)>, super::Error> {
fn header_pairs(&self) -> Result<Vec<(String, String)>, Error> {
self.extra_headers
.iter()
.map(|(name, value)| {
value
.as_str()
.map(|value| (name.clone(), value.to_string()))
.ok_or_else(|| super::Error::RequestField {
.ok_or_else(|| Error::RequestField {
path: format!("extra_headers.{name}"),
})
})
@ -212,62 +83,6 @@ impl OcrConnectionInputs {
}
}
#[derive(Clone)]
pub struct OcrConnection {
pub api_key: Option<String>,
pub api_key_source: InputSource,
pub api_base: Option<String>,
pub api_base_source: InputSource,
pub extra_headers: Vec<(String, String)>,
pub extra_headers_source: InputSource,
pub timeout: Duration,
pub max_download_bytes: u64,
pub max_response_bytes: usize,
pub poll_timeout: Duration,
}
impl OcrConnection {
pub(crate) fn new(credentials: ResolvedOcrCredentials, transport: OcrTransportConfig) -> Self {
let api_key_source = credentials
.api_key
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Deployment);
let api_base_source = credentials
.api_base
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Deployment);
Self {
api_key: credentials.api_key.map(Sourced::into_value),
api_key_source,
api_base: credentials.api_base.map(Sourced::into_value),
api_base_source,
extra_headers: transport.extra_headers,
extra_headers_source: transport.extra_headers_source,
timeout: transport.timeout,
max_download_bytes: transport.max_download_bytes,
max_response_bytes: transport.max_response_bytes,
poll_timeout: transport.poll_timeout,
}
}
}
impl Default for OcrConnection {
fn default() -> Self {
Self::new(
ResolvedOcrCredentials::default(),
OcrTransportConfig::default(),
)
}
}
#[derive(Clone, Default)]
pub(crate) struct ResolvedOcrCredentials {
pub api_key: Option<Sourced<String>>,
pub api_base: Option<Sourced<String>>,
}
pub struct LiteLLMOcrRequest<D = OcrDocumentInput> {
pub model: String,
pub document: D,
@ -285,7 +100,7 @@ impl LiteLLMOcrRequest {
document: impl Into<OcrDocumentInput>,
custom_llm_provider: Option<&str>,
optional_params: CallArguments,
) -> Result<Self, super::Error> {
) -> Result<Self, Error> {
let (model, config) = resolve_provider_config(&model, custom_llm_provider)?;
let default_transport = OcrTransportConfig::default();
let max_response_bytes = optional_params
@ -295,7 +110,7 @@ impl LiteLLMOcrRequest {
.as_u64()
.and_then(|value| usize::try_from(value).ok())
.filter(|value| *value > 0 && *value <= default_transport.max_response_bytes)
.ok_or_else(|| super::Error::RequestField {
.ok_or_else(|| Error::RequestField {
path: "max_response_bytes".into(),
})
})
@ -353,15 +168,8 @@ impl<D> LiteLLMOcrRequest<D> {
}
}
pub(crate) fn response_format(&self) -> Result<OcrResponseFormat, super::Error> {
self.optional_params
.get("req_format")
.filter(|value| !value.is_null())
.map(|value| {
serde_json::from_value(value.clone()).map_err(|_| super::Error::RequestFormat)
})
.transpose()
.map(|format| format.unwrap_or_default())
pub(crate) fn response_format(&self) -> Result<OcrResponseFormat, Error> {
response_format(&self.optional_params)
}
pub fn provider_name(&self) -> &'static str {
@ -395,7 +203,7 @@ impl LiteLLMOcrRequest {
custom_llm_provider: Option<&str>,
optional_params: CallArguments,
connection: OcrConnectionInputs,
) -> Result<Self, super::Error> {
) -> Result<Self, Error> {
let request = Self::new(model, document, custom_llm_provider, optional_params)?;
let transport = request.transport.clone().with_overrides(
connection.header_pairs()?,
@ -416,155 +224,6 @@ impl LiteLLMOcrRequest {
pub(crate) type ResolvedOcrRequest = LiteLLMOcrRequest<OcrDocument>;
pub(crate) struct PreparedOcrRequest {
pub model: String,
pub document: OcrDocument,
pub connection: OcrConnection,
pub host: super::route::OcrHost,
/// Whether the caller handed over the document as is, so the wire body's document
/// is the caller's own input rather than something the route prepared.
pub caller_document: bool,
pub optional_params: CallArguments,
pub input_sources: BTreeMap<String, InputSource>,
pub azure_ad_token_provider: Option<TokenProviderHandle>,
pub(crate) config: OcrConfigKind,
}
impl PreparedOcrRequest {
pub(crate) fn new(
request: ResolvedOcrRequest,
connection: OcrConnection,
host: super::route::OcrHost,
caller_document: bool,
) -> Self {
let LiteLLMOcrRequest {
model,
document,
credentials: _,
transport: _,
optional_params,
input_sources,
azure_ad_token_provider,
config,
} = request;
Self {
model,
document,
connection,
host,
caller_document,
optional_params,
input_sources,
azure_ad_token_provider,
config,
}
}
pub(crate) fn response_format(&self) -> Result<OcrResponseFormat, super::Error> {
self.optional_params
.get("req_format")
.filter(|value| !value.is_null())
.map(|value| {
serde_json::from_value(value.clone()).map_err(|_| super::Error::RequestFormat)
})
.transpose()
.map(|format| format.unwrap_or_default())
}
pub(crate) fn provider_name(&self) -> &'static str {
self.config.provider().into()
}
}
#[serde_as]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrPageDimensions {
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub dpi: Option<i64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub height: Option<i64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub width: Option<i64>,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrPageImage {
pub image_base64: Option<String>,
pub bbox: Option<Map<String, Value>>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[serde_as]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrPage {
#[serde_as(deserialize_as = "LaxI64")]
pub index: i64,
pub markdown: String,
pub images: Option<Vec<OcrPageImage>>,
pub dimensions: Option<OcrPageDimensions>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[serde_as]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrUsageInfo {
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub pages_processed: Option<i64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub pages_processed_annotation: Option<i64>,
#[serde_as(deserialize_as = "Option<FiniteF64>")]
pub credits: Option<f64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub doc_size_bytes: Option<i64>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct LiteLLMOcrResponse {
pub pages: Vec<OcrPage>,
pub model: String,
pub document_annotation: Option<Value>,
pub usage_info: Option<OcrUsageInfo>,
pub content: Option<String>,
pub tables: Option<Vec<Map<String, Value>>>,
#[serde(rename = "keyValuePairs")]
pub key_value_pairs: Option<Vec<Map<String, Value>>>,
#[serde(default = "ocr_object")]
pub object: String,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub provider_native_response: Option<Map<String, Value>>,
}
impl LiteLLMOcrResponse {
pub fn new(model: impl Into<String>, pages: Vec<OcrPage>) -> Self {
Self {
pages,
model: model.into(),
document_annotation: None,
usage_info: None,
content: None,
tables: None,
key_value_pairs: None,
object: ocr_object(),
extra_fields: Map::new(),
provider_native_response: None,
}
}
pub fn into_json(self) -> Value {
serde_json::to_value(self).expect("OCR response fields are JSON-compatible")
}
}
fn ocr_object() -> String {
"ocr".into()
}
#[cfg(test)]
mod tests {
use serde_json::json;
@ -645,97 +304,7 @@ mod tests {
};
assert!(matches!(
error,
super::super::Error::RequestField { ref path } if path == "extra_headers.x-a"
Error::RequestField { ref path } if path == "extra_headers.x-a"
));
}
#[test]
fn normalized_response_rejects_invalid_shared_fields() {
for fields in [
json!({"pages":[{}]}),
json!({"pages":[{"index":0,"markdown":false}]}),
json!({"pages":[{"index":0,"markdown":"","images":[{"bbox":[]}]}]}),
json!({"usage_info":{"pages_processed":1.5}}),
json!({"tables":[false]}),
json!({"keyValuePairs":[[]]}),
json!({"provider_native_response":[]}),
] {
let payload: Map<String, Value> = json!({"model":"model", "pages":[]})
.as_object()
.unwrap()
.iter()
.chain(fields.as_object().unwrap())
.map(|(key, value)| (key.clone(), value.clone()))
.collect();
assert!(serde_json::from_value::<LiteLLMOcrResponse>(Value::Object(payload)).is_err());
}
assert!(
serde_json::from_value::<OcrDocument>(json!({
"type":"image_url", "image_url":"https://example.com/image", "detail":42
}))
.is_err()
);
}
#[test]
fn numeric_coercion_preserves_integer_precision_and_rejects_fractional_values() {
for (value, expected) in [
(json!("9007199254740993.0"), 9_007_199_254_740_993),
(json!("+2.000"), 2),
(json!("1_000"), 1000),
(json!(true), 1),
(json!(2.0), 2),
] {
let page: OcrPage =
serde_json::from_value(json!({"index":value,"markdown":""})).unwrap();
assert_eq!(page.index, expected);
}
for value in [
json!("1e2"),
json!(".0"),
json!("2."),
json!("_2"),
json!("2__0"),
json!(2.5),
json!(null),
] {
assert!(
serde_json::from_value::<OcrPage>(json!({"index":value,"markdown":""})).is_err()
);
}
}
#[rstest::rstest]
#[case::document_url("document_url", "document_name", "application/pdf")]
#[case::image_url("image_url", "detail", "image/png")]
fn document_variants_preserve_provider_fields_when_rewriting_sources(
#[case] kind: &str,
#[case] field: &str,
#[case] mime_type: &str,
#[values(json!("kept"), Value::Null)] extra: Value,
) {
let original = "https://example.com/input";
let replacement = format!("data:{mime_type};base64,AA==");
let document: OcrDocument =
serde_json::from_value(json!({"type": kind, kind: original, field: extra})).unwrap();
assert_eq!(document.source(), original);
assert_eq!(
serde_json::to_value(document.with_source(replacement.clone())).unwrap(),
json!({"type": kind, kind: replacement, field: extra})
);
}
#[test]
fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() {
let response = LiteLLMOcrResponse {
extra_fields: json!({"provider_field":"kept"})
.as_object()
.unwrap()
.clone(),
..LiteLLMOcrResponse::new("model", vec![])
};
let serialized = response.into_json();
assert_eq!(serialized["provider_field"], "kept");
assert!(serialized.get("provider_native_response").is_none());
}
}

14
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -1,17 +1,20 @@
use std::{collections::BTreeMap, time::Duration};
use litellm_auth::InputSource;
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{OcrDocument, decode_request_value},
};
use serde::Deserialize;
use serde_json::{Map, Value};
pub use super::is_supported_request;
use super::{Error, LiteLLMOcrRequest, OcrConnectionInputs, OcrDocument, OcrDocumentInput};
use crate::ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput};
pub fn consumed_optional_params(
model: &str,
provider: Option<&str>,
) -> Result<Vec<litellm_core_utils::call_arguments::ArgumentSpec>, Error> {
let specs = super::consumed_optional_params(model, provider)?;
let specs = crate::ocr::arguments::consumed_optional_params(model, provider)?;
Ok(consumed_optional_param_names(model, provider)?
.into_iter()
.map(|name| litellm_core_utils::call_arguments::ArgumentSpec {
@ -25,7 +28,7 @@ pub fn consumed_optional_param_names(
model: &str,
provider: Option<&str>,
) -> Result<Vec<&'static str>, Error> {
let names = super::consumed_optional_param_names(model, provider)?;
let names = crate::ocr::arguments::consumed_optional_param_names(model, provider)?;
let (_, config) = super::provider_config::resolve_provider_config(model, provider)?;
if config == super::provider_config::OcrConfigKind::VertexDeepSeek {
return Ok(names
@ -99,7 +102,7 @@ pub fn decode_document(value: Value) -> Result<OcrDocument, Error> {
{
return Err(Error::MissingDocumentUrl);
}
super::json::decode_request_value(value, "document")
decode_request_value(value, "document")
}
#[cfg(test)]
@ -108,6 +111,7 @@ mod tests {
use serde_json::json;
use super::*;
use crate::ocr::arguments::is_supported_request;
#[rstest]
#[case::omitted(json!({"type":"document_url", "document_url":"https://example.com/a.pdf"}))]

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

@ -11,7 +11,7 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
}

39
litellm-rust/crates/core/src/responses/websocket.rs
View file

@ -95,7 +95,9 @@ impl ResponsesWebSocketConnection {
timeout: Option<Duration>,
) -> Result<Self, Error> {
let mut request = url.into_client_request().map_err(|error| {
Error::Transport(crate::transport::Error::Network(error.to_string()))
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
})?;
for (name, value) in headers {
let header_name = name
@ -108,7 +110,7 @@ impl ResponsesWebSocketConnection {
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Transport(crate::transport::Error::Network(
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
"Responses WebSocket connection timed out".into(),
))
})?,
@ -116,12 +118,14 @@ impl ResponsesWebSocketConnection {
};
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => {
Error::Transport(crate::transport::Error::Http {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status: response.status().as_u16(),
body: String::new(),
})
}
other => Error::Transport(crate::transport::Error::Network(other.to_string())),
other => Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
other.to_string(),
)),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
@ -131,14 +135,17 @@ impl ResponsesWebSocketConnection {
pub async fn send_text(&self, text: String) -> Result<(), Error> {
let mut socket = self.socket.lock().await;
let Some(socket) = socket.as_mut() else {
return Err(Error::Transport(crate::transport::Error::Network(
"Responses WebSocket is closed".into(),
)));
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Network(
"Responses WebSocket is closed".into(),
),
));
};
socket
.send(Message::Text(text))
.await
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))
socket.send(Message::Text(text)).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
})
}
pub async fn recv_text(&self) -> Result<Option<String>, Error> {
@ -153,9 +160,9 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(Error::Transport(crate::transport::Error::Network(
error.to_string(),
))),
Some(Err(error)) => Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Network(error.to_string()),
)),
}
}
@ -163,7 +170,9 @@ impl ResponsesWebSocketConnection {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket.close(None).await.map_err(|error| {
Error::Transport(crate::transport::Error::Network(error.to_string()))
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
})?;
}
*socket = None;

2
litellm-rust/crates/core/src/transport/mod.rs
View file

@ -1,2 +0,0 @@
mod error;
pub use error::Error;

219
litellm-rust/crates/core/tests/azure_ai_ocr.rs
View file

@ -1,9 +1,8 @@
use litellm_llms::base_llm::ocr::error::Error;
use serde_json::{Value, json};
use super::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request};
use crate::ocr::route::LocalOcrHost;
#[tokio::test]
async fn facade_executes_azure_mistral_with_prepared_auth() {
@ -80,3 +79,215 @@ async fn rejects_non_inline_body_after_guardrails() {
let error = perform_ocr_with(host).await.unwrap_err();
assert!(error.to_string().contains("data URI"));
}
mod transformation {
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use litellm_auth::{
ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle,
};
use rstest::rstest;
use serde_json::json;
use super::*;
use crate::ocr::{
test_support::{MockResponse, header, mock_server, perform_ocr},
types::LiteLLMOcrRequest,
wire::decode_request,
};
#[derive(Debug)]
struct CountingToken {
token: fn(usize) -> String,
calls: AtomicUsize,
}
impl CountingToken {
fn new(token: fn(usize) -> String) -> Arc<Self> {
Arc::new(Self {
token,
calls: AtomicUsize::new(0),
})
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
}
impl TokenProvider for CountingToken {
fn acquire(&self) -> TokenFuture<'_> {
let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1;
let token = SecretValue::new((self.token)(call));
Box::pin(async move {
Ok(ResolvedCredential::AccessToken {
token,
expires_on: None,
})
})
}
}
fn numbered_token(call: usize) -> String {
format!("callback-{call}")
}
fn azure_request(
provider: &Arc<CountingToken>,
api_base: Option<&str>,
api_key: Option<&str>,
extra_headers: Value,
optional_params: Value,
) -> LiteLLMOcrRequest {
let wire = serde_json::from_value(json!({
"model": "azure_ai/mistral-ocr-latest",
"document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": null,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": 2.0
}))
.unwrap();
LiteLLMOcrRequest {
azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())),
..decode_request(wire).unwrap()
}
}
fn ocr_page() -> MockResponse {
MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]}))
}
#[tokio::test]
async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() {
let provider = CountingToken::new(numbered_token);
let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await;
for _ in 0..2 {
perform_ocr(azure_request(
&provider,
Some(&base),
None,
Value::Null,
json!({}),
))
.await
.unwrap();
}
server.await.unwrap();
assert_eq!(provider.calls(), 2);
let requests = seen.lock().unwrap();
assert_eq!(
requests
.iter()
.map(|request| header(request, "authorization"))
.collect::<Vec<_>>(),
[Some("Bearer callback-1"), Some("Bearer callback-2")]
);
}
#[rstest]
#[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)]
#[case::provider_beats_static_token(
None,
Value::Null,
json!({"azure_ad_token":"static-token"}),
"Bearer callback-1",
1
)]
#[case::header_wins_on_the_wire_but_provider_still_runs(
None,
json!({"Authorization":"Bearer override"}),
json!({}),
"Bearer override",
1
)]
#[tokio::test]
async fn credential_precedence(
#[case] api_key: Option<&str>,
#[case] extra_headers: Value,
#[case] optional_params: Value,
#[case] expected_authorization: &str,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(numbered_token);
let (base, seen, server) = mock_server(vec![ocr_page()]).await;
perform_ocr(azure_request(
&provider,
Some(&base),
api_key,
extra_headers,
optional_params,
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(provider.calls(), expected_calls);
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(
header(&requests[0], "authorization"),
Some(expected_authorization)
);
}
#[rstest]
#[case::missing_api_base(
false,
json!({}),
numbered_token,
|error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase {
provider: "Azure AI",
environment_variable: "AZURE_AI_API_BASE",
})),
0
)]
#[case::unsupported_oidc_reference(
true,
json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}),
numbered_token,
|error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)),
0
)]
#[case::empty_provider_token_ignores_static_token(
true,
json!({"azure_ad_token":"static-token"}),
|_| String::new(),
|error: &Error| matches!(error, Error::MissingAzureAiCredentials),
1
)]
#[tokio::test]
async fn credential_failures_send_no_provider_request(
#[case] with_api_base: bool,
#[case] optional_params: Value,
#[case] token: fn(usize) -> String,
#[case] expected: fn(&Error) -> bool,
#[case] expected_calls: usize,
) {
let provider = CountingToken::new(token);
let (base, seen, server) = mock_server(vec![ocr_page()]).await;
let error = perform_ocr(azure_request(
&provider,
with_api_base.then_some(base.as_str()),
None,
Value::Null,
optional_params,
))
.await
.unwrap_err();
server.abort();
assert!(expected(&error), "unexpected error: {error:?}");
assert_eq!(provider.calls(), expected_calls);
assert!(seen.lock().unwrap().is_empty());
}
}

234
litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
View file

@ -1,12 +1,13 @@
use litellm_callbacks::event::CallEvent;
use litellm_llms::base_llm::ocr::error::Error;
use rstest::rstest;
use serde_json::{Value, json};
use super::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
wire::{OcrWireRequest, decode_request},
};
use crate::ocr::route::LocalOcrHost;
fn query_value(url: &str, key: &str) -> Option<String> {
url::Url::parse(url)
@ -27,12 +28,13 @@ async fn facade_maps_pages_features_and_url_document() {
&base,
json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}),
);
request.document = serde_json::from_value::<super::OcrDocument>(json!({
"type":"document_url",
"document_url":"https://example.com/document.pdf"
}))
.unwrap()
.into();
request.document =
serde_json::from_value::<litellm_llms::base_llm::ocr::transformation::OcrDocument>(json!({
"type":"document_url",
"document_url":"https://example.com/document.pdf"
}))
.unwrap()
.into();
perform_ocr(request).await.unwrap();
server.await.unwrap();
@ -52,16 +54,16 @@ async fn facade_maps_pages_features_and_url_document() {
}
#[rstest]
#[case(json!({"pages":[true]}), crate::ocr::Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages":[1,"2"]}), crate::ocr::Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages":[-1]}), crate::ocr::Error::Pages("negative page index".into()))]
#[case(json!({"pages":"1&&features=bad"}), crate::ocr::Error::Pages("invalid native page range".into()))]
#[case(json!({"features":"languages&pages=1"}), crate::ocr::Error::Features)]
#[case(json!({"req_format":"azure"}), crate::ocr::Error::RequestFormat)]
#[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))]
#[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))]
#[case(json!({"features":"languages&pages=1"}), Error::Features)]
#[case(json!({"req_format":"azure"}), Error::RequestFormat)]
#[tokio::test]
async fn rejects_invalid_pages_features_and_format(
#[case] options: Value,
#[case] expected: super::Error,
#[case] expected: Error,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await;
let result = decode_request(OcrWireRequest {
@ -460,3 +462,207 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() {
assert!(error.to_string().contains("dot segment"));
}
}
mod transformation {
use std::sync::{Arc, Mutex};
use litellm_callbacks::event::CallEvent;
use litellm_llms::base_llm::ocr::transformation::OcrDocument;
use serde_json::{Value, json};
use super::*;
use crate::ocr::{
route::LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
#[tokio::test]
async fn facade_maps_pages_features_and_url_document() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":{"pages":[]}
}))])
.await;
let mut request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}),
);
request.document = serde_json::from_value::<OcrDocument>(json!({
"type":"document_url",
"document_url":"https://example.com/document.pdf"
}))
.unwrap()
.into();
perform_ocr(request).await.unwrap();
server.await.unwrap();
let request = &seen.lock().unwrap()[0];
let target = request.split_whitespace().nth(1).unwrap();
let url = format!("{base}{target}");
assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3"));
assert_eq!(
query_value(&url, "features").as_deref(),
Some("keyValuePairs,languages")
);
let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false})
);
}
#[tokio::test]
async fn rejects_invalid_pages_features_and_format() {
for options in [
json!({"pages":[true]}),
json!({"pages":[1,"2"]}),
json!({"pages":[-1]}),
json!({"pages":"1&&features=bad"}),
json!({"features":"languages&pages=1"}),
json!({"req_format":"azure"}),
] {
let request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
"http://127.0.0.1:1",
options.clone(),
);
let rejected = perform_ocr(request).await.is_err();
assert!(rejected, "accepted {options}");
}
}
#[tokio::test]
async fn immediate_response_normalizes_pages_and_preserves_native() {
let operation = json!({
"status":"succeeded",
"operationExtension":42,
"analyzeResult":{
"content":"A\n\nB",
"tables":[{"cells":[]}],
"keyValuePairs":[{"key":{"content":"A"}}],
"pages":[{
"pageNumber":"2",
"width":"8.5",
"height":11,
"unit":"inch",
"lines":[{"content":"A"},{"content":null},{"content":"B"}]
}]
}
});
let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await;
let result = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"req_format":"native"}),
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0].index, 1);
assert_eq!(result.pages[0].markdown, "A\n\nB");
assert_eq!(
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
json!({"width":816,"height":1056,"dpi":96})
);
assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1));
let serialized = result.clone().into_json();
assert_eq!(serialized["content"], "A\n\nB");
assert_eq!(serialized["tables"], json!([{"cells":[]}]));
assert_eq!(
serialized["keyValuePairs"],
json!([{"key":{"content":"A"}}])
);
assert!(serialized.get("key_value_pairs").is_none());
assert_eq!(
result.provider_native_response.as_ref(),
operation.as_object()
);
}
#[tokio::test]
async fn accepted_response_polls_to_success_with_only_credentials() {
let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}});
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 200,
headers: vec![("Retry-After", "0".into())],
body: json!({"status":"running"}),
},
MockResponse::json(operation.clone()),
])
.await;
let mut request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"req_format":"native"}),
);
request
.transport
.extra_headers
.push(("X-Trace".into(), "initial-only".into()));
let result = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(
result.provider_native_response.as_ref(),
operation.as_object()
);
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 3);
assert!(requests[0].to_ascii_lowercase().contains("x-trace:"));
for poll in &requests[1..] {
assert!(!poll.to_ascii_lowercase().contains("x-trace:"));
assert!(
poll.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: test-key")
);
}
}
#[tokio::test]
async fn accepted_response_emits_response_received_for_submission_and_completed_poll() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({"submitted": true}),
},
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let responses_received = Arc::new(Mutex::new(Vec::new()));
let request_count = seen.clone();
let observed = responses_received.clone();
let host = LocalOcrHost::new(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
observed
.lock()
.unwrap()
.push((request_count.lock().unwrap().len(), raw.body.clone()));
}
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
assert_eq!(
*responses_received.lock().unwrap(),
[
(1, r#"{"submitted":true}"#.to_string()),
(2, r#"{"status":"succeeded"}"#.to_string()),
]
);
}
}

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

@ -0,0 +1,136 @@
mod transformation {
use litellm_llms::{
base_llm::ocr::{
error::Error,
transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat},
},
cohere::ocr::transformation::*,
};
use rstest::rstest;
use serde_json::{Value, json};
#[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", "timeout":30,
"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_for_test(request);
let http = CohereParseConfig
.prepare_request(
&request,
&crate::ocr::test_support::ocr_client(),
&crate::ocr::test_support::NoHooks,
)
.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]}}
})
);
}
#[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(),
OcrResponseFormat::Litellm
);
let request = crate::ocr::prepare::prepare_request_for_test(request);
let http = CohereParseConfig
.prepare_request(
&request,
&crate::ocr::test_support::ocr_client(),
&crate::ocr::test_support::NoHooks,
)
.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]
#[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")]
#[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")]
#[tokio::test]
async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key(
#[case] model: &str,
#[case] request_line: &str,
) {
use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr};
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let request = crate::ocr::test_support::wire_request(model, &base, json!({}))
.with_document(
serde_json::from_value::<OcrDocument>(
json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}),
)
.unwrap()
.into(),
);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with(request_line), "{}", requests[0]);
assert_eq!(
header(&requests[0], "authorization"),
Some("Bearer test-key")
);
}
#[rstest]
#[tokio::test]
async fn route_rejects_non_image_document_without_a_request(
#[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str,
) {
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr};
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let error = perform_ocr(crate::ocr::test_support::wire_request(
model,
&base,
json!({}),
))
.await
.unwrap_err();
server.abort();
assert!(matches!(error, Error::CohereImageOnly), "{error:?}");
assert!(seen.lock().unwrap().is_empty());
}
}

18
litellm-rust/crates/core/tests/deepseek_ocr.rs
View file

@ -1,17 +1,13 @@
use litellm_llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument},
vertex_ai::ocr::deepseek_transformation::{
DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig,
normalize_response as transform_ocr_response,
},
};
use rstest::rstest;
use serde_json::{Value, json};
use crate::{
llms::{
base_llm::ocr::transformation::BaseOcrConfig,
vertex_ai::ocr::deepseek_transformation::{
DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig,
normalize_response as transform_ocr_response,
},
},
ocr::types::OcrDocument,
};
fn document() -> OcrDocument {
serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap()
}

126
litellm-rust/crates/core/tests/ocr.rs
View file

@ -5,16 +5,23 @@ use litellm_callbacks::{
host::{Host, HostOp, HostResult},
machine::{HostFailure, Machine, MachineStep},
};
use litellm_llms::{
base_llm::ocr::{
error::Error as OcrError,
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
},
custom_httpx::llm_http_handler::OcrClient,
};
use rstest::rstest;
use serde_json::{Value, json};
use super::{
LocalOcrHost, OcrClient, OcrOp, OcrOpResult, ocr_machine,
test_support::{
MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request,
},
wire::{OcrWireRequest, decode_request},
};
use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine};
#[rstest]
#[case::mistral("mistral/model", json!({}))]
@ -41,7 +48,7 @@ async fn ocr_contract_upstream_error_preserves_status_body_and_headers(
.unwrap_err();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 1);
let super::Error::Provider {
let OcrError::Provider {
status,
body,
headers,
@ -175,11 +182,12 @@ async fn facade_uses_the_injected_http_client() {
.default_headers(default_headers)
.build()
.unwrap();
OcrClient::new(provider_http)
.unwrap()
.perform(wire_request("mistral/model", &base, json!({})))
.await
.unwrap();
crate::ocr::client::perform(
&OcrClient::new(provider_http).unwrap(),
wire_request("mistral/model", &base, json!({})),
)
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host"));
}
@ -193,7 +201,7 @@ fn event_name(event: &CallEvent) -> &'static str {
}
fn recording_host(
request: super::LiteLLMOcrRequest,
request: crate::ocr::types::LiteLLMOcrRequest,
events: Arc<Mutex<Vec<&'static str>>>,
block: bool,
) -> LocalOcrHost {
@ -202,7 +210,7 @@ fn recording_host(
.with_before_send(move |wire, _| {
before_send_events.lock().unwrap().push("before_send");
if block {
return Err(crate::ocr::Error::InvalidRequest("blocked".into()));
return Err(OcrError::InvalidRequest("blocked".into()));
}
Ok(wire)
})
@ -259,7 +267,7 @@ async fn before_send_context_names_passthrough_fields_and_secrets() {
&base,
json!({"client_secret": "shh", "tenant_id": "t"}),
);
let request = request.with_document(super::OcrDocumentInput::Bytes {
let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes {
bytes: b"abc".as_slice().into(),
file_name: None,
mime_type: Some("application/pdf".into()),
@ -302,7 +310,7 @@ async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() {
true,
);
let error = perform_ocr_with(host).await.unwrap_err();
assert!(matches!(error, crate::ocr::Error::InvalidRequest(message) if message == "blocked"));
assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked"));
assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]);
}
@ -331,11 +339,11 @@ async fn upstream_failure_emits_one_terminal_failure() {
async fn drive_until(
client: OcrClient,
host: &LocalOcrHost,
mut intercept: impl FnMut(WireRequest) -> Result<WireRequest, HostFailure<crate::ocr::Error>>,
mut intercept: impl FnMut(WireRequest) -> Result<WireRequest, HostFailure<OcrError>>,
) -> (
Result<super::LiteLLMOcrResponse, crate::ocr::Error>,
Result<LiteLLMOcrResponse, OcrError>,
Vec<&'static str>,
super::OcrMachine,
crate::ocr::route::OcrMachine,
) {
let mut machine = ocr_machine(client);
let mut result = None;
@ -386,13 +394,13 @@ async fn failed_before_send_does_not_replay_or_reach_transport() {
json!({}),
));
let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| {
Err(HostFailure::Error(crate::ocr::Error::InvalidRequest(
Err(HostFailure::Error(OcrError::InvalidRequest(
"before_send failed".into(),
)))
})
.await;
assert!(
matches!(outcome, Err(crate::ocr::Error::InvalidRequest(message)) if message == "before_send failed")
matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed")
);
assert_eq!(ops, ["ProjectRequest", "BeforeSend"]);
assert!(machine.resume(None).await.is_err());
@ -413,7 +421,7 @@ async fn invalid_provider_response_emits_response_received_before_normalization_
);
let error = perform_ocr_with(host).await.unwrap_err();
server.await.unwrap();
assert!(matches!(error, crate::ocr::Error::ResponseField { .. }));
assert!(matches!(error, OcrError::ResponseField { .. }));
assert_eq!(seen.lock().unwrap().len(), 1);
assert_eq!(
*responses_received.lock().unwrap(),
@ -435,14 +443,14 @@ async fn direct_native_host_drives_the_same_state_machine() {
assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]);
assert!(matches!(
machine.resume(None).await,
Err(crate::ocr::Error::InvalidRequest(_))
Err(OcrError::InvalidRequest(_))
));
}
async fn drive_native_file_call(
request: super::LiteLLMOcrRequest<super::OcrDocumentInput>,
content: Result<super::OcrFileContent, crate::ocr::Error>,
) -> (Result<super::LiteLLMOcrResponse, crate::ocr::Error>, usize) {
request: crate::ocr::types::LiteLLMOcrRequest<crate::ocr::types::OcrDocumentInput>,
content: Result<crate::ocr::types::OcrFileContent, OcrError>,
) -> (Result<LiteLLMOcrResponse, OcrError>, usize) {
let reads = Arc::new(Mutex::new(0));
let counted = reads.clone();
let content = Mutex::new(Some(content));
@ -462,13 +470,13 @@ async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_enco
}))])
.await;
let request = wire_request("mistral/model", &base, json!({})).with_document(
super::OcrDocumentInput::HostReader {
crate::ocr::types::OcrDocumentInput::HostReader {
mime_type: Some("application/pdf".into()),
},
);
let (response, reads) = drive_native_file_call(
request,
Ok(super::OcrFileContent {
Ok(crate::ocr::types::OcrFileContent {
bytes: b"abc".as_slice().into(),
file_name: Some("scan.png".into()),
}),
@ -484,30 +492,27 @@ async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_enco
async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() {
let (base, seen, _server) = mock_server(vec![]).await;
let request = wire_request("mistral/model", &base, json!({}));
let failure = crate::ocr::Error::InvalidRequest("reader exploded".into());
let failure = OcrError::InvalidRequest("reader exploded".into());
let (response, reads) = drive_native_file_call(
request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }),
Err(failure.clone()),
)
.await;
assert!(
matches!(response.unwrap_err(), crate::ocr::Error::InvalidRequest(message) if message == "reader exploded")
matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded")
);
assert_eq!(reads, 1);
let request = wire_request("mistral/model", &base, json!({}));
let (response, _) = drive_native_file_call(
request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }),
Ok(super::OcrFileContent {
request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }),
Ok(crate::ocr::types::OcrFileContent {
bytes: Default::default(),
file_name: None,
}),
)
.await;
assert!(matches!(
response.unwrap_err(),
crate::ocr::Error::EmptyFile
));
assert!(matches!(response.unwrap_err(), OcrError::EmptyFile));
assert!(seen.lock().unwrap().is_empty());
}
@ -522,16 +527,13 @@ async fn path_documents_are_read_by_core_without_a_host_operation() {
let path = dir.join("scan.png");
std::fs::write(&path, b"abc").unwrap();
let request = wire_request("mistral/model", &base, json!({})).with_document(
super::OcrDocumentInput::Path {
crate::ocr::types::OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
},
);
let (response, reads) = drive_native_file_call(
request,
Err(crate::ocr::Error::InvalidRequest("unused".into())),
)
.await;
let (response, reads) =
drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await;
server.await.unwrap();
std::fs::remove_dir_all(&dir).unwrap();
assert_eq!(response.unwrap().pages[0].markdown, "path");
@ -541,16 +543,16 @@ async fn path_documents_are_read_by_core_without_a_host_operation() {
let (base, seen, _server) = mock_server(vec![]).await;
let request = wire_request("mistral/model", &base, json!({}));
let (response, _) = drive_native_file_call(
request.with_document(super::OcrDocumentInput::Path {
request.with_document(crate::ocr::types::OcrDocumentInput::Path {
path: path.clone(),
mime_type: None,
}),
Err(crate::ocr::Error::InvalidRequest("unused".into())),
Err(OcrError::InvalidRequest("unused".into())),
)
.await;
assert!(matches!(
response.unwrap_err(),
crate::ocr::Error::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound
OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound
));
assert!(seen.lock().unwrap().is_empty());
}
@ -563,14 +565,12 @@ async fn cancellation_at_before_send_prevents_execution_and_further_resumption()
json!({}),
));
let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| {
Err(HostFailure::Cancelled(crate::ocr::Error::InvalidRequest(
Err(HostFailure::Cancelled(OcrError::InvalidRequest(
"cancelled".into(),
)))
})
.await;
assert!(
matches!(outcome, Err(crate::ocr::Error::InvalidRequest(message)) if message == "cancelled")
);
assert!(matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled"));
assert_eq!(ops, ["ProjectRequest", "BeforeSend"]);
assert!(machine.resume(Some(HostResult::Emitted)).await.is_err());
}
@ -596,10 +596,7 @@ async fn missing_host_result_preserves_pending_operation() {
));
}
async fn read_bounded_response(
response: Vec<u8>,
limit: usize,
) -> Result<bytes::Bytes, super::Error> {
async fn read_bounded_response(response: Vec<u8>, limit: usize) -> Result<bytes::Bytes, OcrError> {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@ -618,7 +615,7 @@ async fn read_bounded_response(
.unwrap();
let result = tokio::time::timeout(
std::time::Duration::from_secs(2),
super::client::read_response_bytes(response, limit),
litellm_llms::custom_httpx::llm_http_handler::read_response_bytes(response, limit),
)
.await;
server.abort();
@ -628,7 +625,7 @@ async fn read_bounded_response(
#[tokio::test]
async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() {
use super::Error;
use litellm_llms::base_llm::ocr::error::Error;
for response in [
"HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh",
@ -670,7 +667,10 @@ async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_dra
.await
.unwrap_err();
match error {
super::Error::Transport(crate::transport::Error::Http { status, body }) => {
OcrError::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status,
body,
}) => {
assert_eq!(status, 429);
assert_eq!(body, prefix);
}
@ -693,7 +693,7 @@ fn response_limit_is_validated_and_not_forwarded_to_the_provider() {
json!(true),
json!("123"),
json!(1.5),
json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1),
json!(OCR_RESPONSE_MAX_BYTES + 1),
Value::Null,
] {
let wire = serde_json::from_value(json!({
@ -738,8 +738,8 @@ async fn interrupt_drops_provider_captures_before_returning() {
let entered = Arc::new(tokio::sync::Notify::new());
let dropped = Arc::new(AtomicBool::new(false));
let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({}));
let request = super::LiteLLMOcrRequest {
transport: super::OcrTransportConfig {
let request = crate::ocr::types::LiteLLMOcrRequest {
transport: OcrTransportConfig {
extra_headers: vec![("authorization".into(), "Bearer test-key".into())],
..request.transport
},
@ -774,24 +774,24 @@ async fn interrupt_drops_provider_captures_before_returning() {
.await
.unwrap();
assert!(!dropped.load(Ordering::SeqCst));
let selected = crate::ocr::Error::InvalidRequest("cancelled".into());
let selected = OcrError::InvalidRequest("cancelled".into());
let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone()));
assert!(
dropped.load(Ordering::SeqCst),
"interrupt returned while provider captures were still alive"
);
assert!(
matches!(acknowledgement.await, Err(crate::ocr::Error::InvalidRequest(message)) if message == "cancelled")
matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled")
);
}
struct CallerTokenHost {
request: Mutex<Option<super::LiteLLMOcrRequest>>,
request: Mutex<Option<crate::ocr::types::LiteLLMOcrRequest>>,
trace: Mutex<Vec<String>>,
}
impl Host<super::Ocr> for CallerTokenHost {
async fn route(&self, op: OcrOp) -> Result<OcrOpResult, crate::ocr::Error> {
impl Host<crate::ocr::route::Ocr> for CallerTokenHost {
async fn route(&self, op: OcrOp) -> Result<OcrOpResult, OcrError> {
match op {
OcrOp::ProjectRequest => {
self.trace.lock().unwrap().push("project".into());
@ -808,7 +808,7 @@ impl Host<super::Ocr> for CallerTokenHost {
)),
))
}
OcrOp::ReadDocument => Err(crate::ocr::Error::InvalidRequest("no reader".into())),
OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())),
}
}
@ -816,7 +816,7 @@ impl Host<super::Ocr> for CallerTokenHost {
&self,
wire: WireRequest,
_: &litellm_callbacks::event::RequestContext,
) -> Result<WireRequest, crate::ocr::Error> {
) -> Result<WireRequest, OcrError> {
let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization");
let authorization = wire
.headers
@ -910,7 +910,7 @@ async fn interrupting_an_in_flight_provider_request_closes_its_connection() {
.await
.unwrap();
let cancelled = crate::ocr::Error::InvalidRequest("cancelled".into());
let cancelled = OcrError::InvalidRequest("cancelled".into());
assert!(
machine
.interrupt(HostFailure::Cancelled(cancelled))

13
litellm-rust/crates/core/tests/ocr/passthrough.rs
View file

@ -4,17 +4,16 @@ use std::{
};
use litellm_callbacks::event::{RequestContext, WireRequest};
use litellm_llms::base_llm::ocr::error::Error;
use rstest::rstest;
use rstest_reuse::{self, apply, template};
use serde_json::{Map, Value, json};
use super::{
LocalOcrHost,
test_support::{
MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with,
request_body, wire_request_with_document,
},
use super::test_support::{
MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body,
wire_request_with_document,
};
use crate::ocr::route::LocalOcrHost;
#[derive(Clone, Copy, Debug)]
enum Route {
@ -111,7 +110,7 @@ impl Host {
struct Sent {
caller: Map<String, Value>,
result: Result<(), crate::ocr::Error>,
result: Result<(), Error>,
before_send: Option<(WireRequest, RequestContext)>,
provider_body: Option<Value>,
}

37
litellm-rust/crates/core/tests/ocr/support.rs
View file

@ -1,5 +1,11 @@
use std::sync::{Arc, Mutex};
use futures_util::future::BoxFuture;
use litellm_callbacks::event::{Passthrough, WireRequest};
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::{CallHooks, OcrClient},
};
use serde_json::{Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
@ -7,10 +13,29 @@ use tokio::{
};
use crate::ocr::{
LiteLLMOcrRequest, LiteLLMOcrResponse, LocalOcrHost, OcrClient, ocr_machine,
route::{LocalOcrHost, ocr_machine},
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
/// Stands in for a host with no hooks registered: the wire request goes out unchanged
/// and response events go nowhere.
pub(crate) struct NoHooks;
impl CallHooks<Error> for NoHooks {
fn before_send(
&self,
wire: WireRequest,
_passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
Box::pin(async move { Ok(wire) })
}
fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(async { Ok(()) })
}
}
pub(crate) fn ocr_client() -> OcrClient {
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
@ -19,15 +44,11 @@ pub(crate) fn ocr_client() -> OcrClient {
OcrClient::for_test(reqwest::Client::new(), document_http)
}
pub(crate) async fn perform_ocr(
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
ocr_client().perform(request).await
pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
crate::ocr::client::perform(&ocr_client(), request).await
}
pub(crate) async fn perform_ocr_with(
host: LocalOcrHost,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
litellm_callbacks::run::run(ocr_machine(ocr_client()), &host).await
}

259
litellm-rust/crates/core/tests/reducto_ocr.rs
View file

@ -1,11 +1,10 @@
use litellm_callbacks::event::{CallEvent, WireRequest};
use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument};
use rstest::rstest;
use serde_json::{Value, json};
use super::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request};
use crate::ocr::route::LocalOcrHost;
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
@ -85,8 +84,8 @@ async fn data_uri_upload_preserves_multipart_headers(
} else {
json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")})
};
let mut request = super::LiteLLMOcrRequest {
document: serde_json::from_value::<super::OcrDocument>(document)
let mut request = crate::ocr::types::LiteLLMOcrRequest {
document: serde_json::from_value::<OcrDocument>(document)
.unwrap()
.into(),
..wire_request(&format!("reducto/{model}"), &base, json!({}))
@ -185,17 +184,14 @@ async fn upload_failure_stops_before_parse() {
}
#[rstest]
#[case("https://example.com/a.pdf", crate::ocr::Error::ReductoSource)]
#[case("reducto://", crate::ocr::Error::RequestField { path: "document file id".into() })]
#[case("data:application/pdf;base64", crate::ocr::Error::InvalidDataUri)]
#[case(
"data:application/pdf;base64,INVALID!",
crate::ocr::Error::InvalidDataUri
)]
#[case("https://example.com/a.pdf", Error::ReductoSource)]
#[case("reducto://", Error::RequestField { path: "document file id".into() })]
#[case("data:application/pdf;base64", Error::InvalidDataUri)]
#[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)]
#[tokio::test]
async fn rejects_invalid_document_sources_before_network(
#[case] source: &str,
#[case] expected: super::Error,
#[case] expected: Error,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await;
let request = super::test_support::with_source(
@ -220,7 +216,7 @@ async fn rejects_invalid_document_sources_before_network(
#[test]
fn response_normalization_groups_blocks_and_distinguishes_null_result() {
use crate::llms::reducto::ocr::transformation::{
use litellm_llms::reducto::ocr::transformation::{
ReductoResponse, normalize_response as transform_ocr_response,
};
@ -353,3 +349,236 @@ async fn guardrail_rewrites_document_before_upload() {
assert!(requests[0].starts_with("POST /parse "));
assert!(requests[0].contains("reducto://guarded.pdf"));
}
mod transformation {
use litellm_callbacks::event::{CallEvent, WireRequest};
use litellm_llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext},
reducto::ocr::transformation::*,
};
use rstest::rstest;
use super::*;
use crate::ocr::{
route::LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
#[tokio::test]
async fn v3_options_preserve_explicit_null() {
let overrides =
serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true}))
.unwrap();
let params = ReductoParseV3Config
.map_ocr_params(&overrides, "parse-v3")
.unwrap();
let client = crate::ocr::test_support::ocr_client();
let connection = OcrConnection::default();
let document = serde_json::from_value(
json!({"type":"document_url","document_url":"reducto://ready.pdf"}),
)
.unwrap();
let body = ReductoParseV3Config
.async_transform_ocr_request(
"parse-v3",
document,
&params,
&[],
OcrRequestContext {
client: &client,
connection: &connection,
},
)
.await
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap(),
json!({
"input":"reducto://ready.pdf", "formatting":null, "settings":{}
})
);
let absent = ReductoParseV3Config
.map_ocr_params(
&litellm_core_utils::call_arguments::CallArguments::default(),
"parse-v3",
)
.unwrap();
assert_eq!(serde_json::to_value(absent).unwrap(), json!({}));
}
#[rstest]
#[case(
"reducto/parse-v3",
json!({
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://already.pdf",
json!({
"input":"reducto://already.pdf",
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[case(
"reducto/parse-legacy",
json!({
"enhance":{"agentic":[{"type":"table"}]},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://legacy.pdf",
json!({
"document_url":"reducto://legacy.pdf",
"options":{"enhance":{"agentic":[{"type":"table"}]}},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[tokio::test]
async fn request_mapping_matches_python(
#[case] model: &str,
#[case] options: Value,
#[case] source: &str,
#[case] expected: Value,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"result":{"chunks":[]}
}))])
.await;
let request =
crate::ocr::test_support::with_source(wire_request(model, &base, options), source);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert_eq!(request_body(&requests[0]), expected);
}
#[rstest]
#[case("parse-v3")]
#[case("parse-legacy")]
#[tokio::test]
async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})),
])
.await;
let mut request = wire_request(&format!("reducto/{model}"), &base, json!({}));
request.transport.extra_headers = vec![
("Content-Type".into(), "application/json".into()),
("X-Trace".into(), "upload-test".into()),
];
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(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("content-type: multipart/form-data; boundary=")
);
assert!(requests[0].contains("x-trace: upload-test"));
assert!(requests[0].contains("application/pdf"));
assert!(requests[0].contains("abc"));
assert!(requests[1].starts_with("POST /parse "));
}
#[tokio::test]
async fn response_received_stays_after_reducto_upload_and_parse() {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[]}})),
])
.await;
let request_count = seen.clone();
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
assert_eq!(request_count.lock().unwrap().len(), 2);
assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
}
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
}
#[rstest]
#[case("https://example.com/a.pdf")]
#[case("reducto://")]
#[case("data:application/pdf;base64")]
#[case("data:application/pdf;base64,INVALID!")]
#[tokio::test]
async fn rejects_invalid_document_sources_before_network(#[case] source: &str) {
let request = crate::ocr::test_support::with_source(
wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})),
source,
);
assert!(perform_ocr(request).await.is_err());
}
#[tokio::test]
async fn facade_omits_native_response_by_default_and_preserves_auth_priority() {
let raw = json!({"job_id":"job-1","result":{"chunks":[]}});
let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await;
let mut request = crate::ocr::test_support::with_source(
wire_request("reducto/parse-v3", &base, json!({})),
"reducto://ready.pdf",
);
request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())];
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.provider_native_response, None);
assert!(
seen.lock().unwrap()[0]
.to_ascii_lowercase()
.contains("authorization: bearer existing")
);
}
#[rstest]
#[case("reducto/parse-v3")]
#[case("reducto/parse-legacy")]
#[tokio::test]
async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[]}})),
])
.await;
let mut request = wire_request(model, &base, json!({}));
request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())];
let host = LocalOcrHost::new(request).with_before_send(|wire, _| {
Ok(WireRequest {
headers: vec![("authorization".into(), "Bearer guarded".into())],
..wire
})
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(requests[1].starts_with("POST /parse "));
for request in requests.iter() {
assert!(request.contains("authorization: Bearer guarded"));
assert!(!request.contains("Bearer original"));
}
}
}

60
litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs
View file

@ -56,11 +56,11 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() {
#[test]
fn host_registration_selects_deepseek_without_affecting_mistral() {
assert!(crate::ocr::wire::is_supported_request(
assert!(crate::ocr::arguments::is_supported_request(
"deepseek-ocr-maas",
Some("vertex_ai")
));
assert!(crate::ocr::wire::is_supported_request(
assert!(crate::ocr::arguments::is_supported_request(
"mistral-ocr-maas",
Some("vertex_ai")
));
@ -85,3 +85,59 @@ async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
.contains("request-controlled Vertex AI endpoint")
);
}
mod deepseek_transformation {
use serde_json::json;
use super::*;
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
#[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"})
);
}
}

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

@ -1,4 +1,5 @@
use litellm_auth::InputSource;
use litellm_llms::base_llm::ocr::transformation::OcrResponseFormat;
use serde_json::{Value, json};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
@ -102,15 +103,14 @@ async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
async fn adapters_build_complete_requests_and_share_mistral_normalization() {
use std::time::Duration;
use crate::{
llms::{
base_llm::ocr::transformation::BaseOcrConfig,
mistral::ocr::transformation::MistralOcrConfig,
vertex_ai::ocr::transformation::VertexAiOcrConfig,
},
ocr::test_support::ocr_client,
use litellm_llms::{
base_llm::ocr::transformation::BaseOcrConfig,
mistral::ocr::transformation::MistralOcrConfig,
vertex_ai::ocr::transformation::VertexAiOcrConfig,
};
use crate::ocr::test_support::ocr_client;
let client = ocr_client();
let options = json!({
"pages": [0, 2],
@ -132,11 +132,11 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() {
super::test_support::resolved_request(vertex),
);
let direct_http = MistralOcrConfig
.prepare_request(&direct, &client)
.prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks)
.await
.unwrap();
let vertex_http = VertexAiOcrConfig
.prepare_request(&vertex, &client)
.prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr");
@ -164,19 +164,11 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() {
let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"});
let raw = serde_json::to_vec(&payload).unwrap();
let direct_response = MistralOcrConfig
.transform_ocr_response(
&direct.model,
&raw,
crate::ocr::types::OcrResponseFormat::Litellm,
)
.transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
let vertex_response = VertexAiOcrConfig
.transform_ocr_response(
&vertex.model,
&raw,
crate::ocr::types::OcrResponseFormat::Litellm,
)
.transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
assert_eq!(direct_response, vertex_response);
@ -184,3 +176,99 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() {
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
mod transformation {
use rstest::rstest;
use serde_json::{Value, json};
use crate::ocr::test_support::wire_request;
#[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 litellm_llms::{
base_llm::ocr::transformation::BaseOcrConfig,
mistral::ocr::transformation::MistralOcrConfig,
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_for_test(
crate::ocr::test_support::resolved_request(direct),
);
let vertex = crate::ocr::prepare::prepare_request_for_test(
crate::ocr::test_support::resolved_request(vertex),
);
let direct_http = MistralOcrConfig
.prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks)
.await
.unwrap();
let vertex_http = VertexAiOcrConfig
.prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr");
assert_eq!(
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");
}
}

21
litellm-rust/crates/llms/AGENTS.md Normal file
View file

@ -0,0 +1,21 @@
litellm-llms mirrors `litellm/llms/`: base config traits, provider transformations, and the `custom_httpx` handlers. See `../core/AGENTS.md` for how the crates layer.
## 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/llms/<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 live next to `BaseOcrConfig` in `src/base_llm/ocr/transformation.rs`, as they do in Python; 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. `src/base_llm/ocr/error.rs` and `src/base_llm/ocr/document.rs` are Rust-only: the OCR error taxonomy shared with the route, and inline-document helpers shared by several providers
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 in litellm-core. 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

13
litellm-rust/crates/llms/Cargo.toml
View file

@ -5,19 +5,30 @@ edition.workspace = true
license.workspace = true
repository.workspace = true
[features]
test-support = []
[dependencies]
litellm-types.workspace = true
litellm-core-utils.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-callbacks.workspace = true
litellm-framing.workspace = true
base64.workspace = true
bytes.workspace = true
data-url = "0.3.2"
futures-util.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_path_to_error = "0.1"
serde_with.workspace = true
thiserror.workspace = true
time.workspace = true
tokio = { workspace = true, features = ["sync"] }
url.workspace = true
[dev-dependencies]

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

@ -1 +1,2 @@
pub mod anthropic;
pub mod ocr;

47
litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs → litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs
View file

@ -2,23 +2,22 @@ use litellm_core_utils::{call_arguments::CallArguments, url_utils::ApiUrl};
use serde_json::Value;
use crate::{
llms::{
base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext},
cohere::ocr::{
CohereOptions,
transformation::{CohereParseConfig, CohereRequest},
validate_document,
base_llm::ocr::{
document::{inline_remote_document, validate_inline_document},
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat,
PreparedOcrRequest,
},
},
ocr::{
OcrClient,
document::{inline_remote_document, validate_inline_document},
types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest},
cohere::ocr::transformation::{
CohereOptions, CohereParseConfig, CohereRequest, validate_document,
},
custom_httpx::llm_http_handler::OcrClient,
};
#[derive(Default)]
pub(crate) struct AzureAICohereParseConfig;
pub struct AzureAICohereParseConfig;
impl BaseOcrConfig for AzureAICohereParseConfig {
type OcrParams = CohereOptions;
@ -37,7 +36,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
BaseOcrConfig::validate_environment(
&super::transformation::AzureAiOcrConfig,
request,
@ -51,10 +50,10 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
request: &PreparedOcrRequest,
_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
let base = super::transformation::AzureAiOcrConfig::resolve_api_base(
request.connection.api_base.as_deref(),
&crate::ocr::prepare::credential_env,
&crate::base_llm::ocr::transformation::credential_env,
)?;
self.get_complete_url(&base)
}
@ -65,7 +64,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
document: OcrDocument,
params: &CohereOptions,
headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
) -> Result<CohereRequest, Error> {
CohereParseConfig.transform_ocr_request(model, document, params, headers)
}
@ -77,7 +76,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
&self,
arguments: &CallArguments,
model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
) -> Result<CohereOptions, Error> {
CohereParseConfig.map_ocr_params(arguments, model)
}
@ -88,7 +87,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
optional_params: &CohereOptions,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<CohereRequest, crate::ocr::Error> {
) -> Result<CohereRequest, Error> {
validate_document(&document)?;
let document = inline_remote_document(
context.client.document_fetcher(),
@ -103,20 +102,20 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, 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)?;
fn validate_request_body(&self, body: &Value) -> Result<(), Error> {
let document = crate::custom_httpx::llm_http_handler::body_document(body)?;
validate_document(&document)?;
validate_inline_document(&document)
}
}
impl AzureAICohereParseConfig {
fn get_complete_url(&self, base: &str) -> Result<String, crate::ocr::Error> {
fn get_complete_url(&self, base: &str) -> Result<String, Error> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base());
@ -134,8 +133,8 @@ impl AzureAICohereParseConfig {
}
}
fn invalid_api_base() -> crate::ocr::Error {
crate::ocr::Error::RequestField {
fn invalid_api_base() -> Error {
Error::RequestField {
path: "api_base".into(),
}
}

8
litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs → litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs
View file

@ -3,12 +3,12 @@ use std::sync::OnceLock;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
use crate::ocr::types::OcrConnection;
use crate::base_llm::ocr::{error::Error, transformation::OcrConnection};
pub(super) async fn resolve_entra(
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<String>>, crate::ocr::Error> {
) -> Result<Option<Sourced<String>>, Error> {
static SERVICE: OnceLock<AzureAuthService> = OnceLock::new();
SERVICE
.get_or_init(AzureAuthService::default)
@ -25,13 +25,13 @@ pub(super) async fn resolve_entra(
Sourced::new(value, source)
})
})
.map_err(crate::ocr::Error::from)
.map_err(Error::from)
}
pub(super) fn validate_destination(
connection: &OcrConnection,
credential_source: InputSource,
) -> Result<(), crate::ocr::Error> {
) -> Result<(), Error> {
if connection.api_base.is_some()
&& connection.api_base_source == InputSource::Request
&& credential_source != InputSource::Request

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

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

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

@ -15,33 +15,31 @@ use serde_with::serde_as;
use tokio::time::Instant;
use crate::{
constants::{
AZURE_DI_API_VERSION, AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT,
AZURE_DI_DEFAULT_WIDTH, AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS,
},
llms::base_llm::ocr::transformation::{
BaseOcrConfig, OcrResponseContext, decode_and_normalize_response,
},
ocr::{
OcrClient,
client::read_json_response,
base_llm::ocr::{
document::InlineDocument,
json::DecodedOcrResponse,
prepare::credential_env,
route::OcrHost,
types::{
LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
OcrPageDimensions, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
ResolvedOcrCredentials,
error::Error,
transformation::{
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES,
OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo,
PreparedOcrRequest, ResolvedOcrCredentials, credential_env,
decode_and_normalize_response, decode_response,
},
},
custom_httpx::llm_http_handler::{CallHooks, OcrClient, read_json_response},
};
const AZURE_DI_API_VERSION: &str = "2024-11-30";
const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
const AZURE_DI_DEFAULT_DPI: i64 = 96;
const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
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, PartialEq, Serialize)]
pub(crate) struct DocumentIntelligenceParams {
pub struct DocumentIntelligenceParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
@ -50,7 +48,7 @@ pub(crate) struct DocumentIntelligenceParams {
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum DocumentIntelligenceRequest {
pub enum DocumentIntelligenceRequest {
UrlSource {
#[serde(rename = "urlSource")]
url_source: String,
@ -95,7 +93,7 @@ impl std::fmt::Display for OperationStatus {
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceOperation {
pub struct AzureDocumentIntelligenceOperation {
status: Option<OperationStatus>,
#[serde(rename = "analyzeResult")]
analyze_result: Option<AzureDocumentIntelligenceAnalyzeResult>,
@ -132,7 +130,7 @@ struct AzureDocumentIntelligenceLine {
}
#[derive(Clone, Debug)]
pub(crate) struct AzureDocumentIntelligenceOcrConfig;
pub struct AzureDocumentIntelligenceOcrConfig;
impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
type OcrParams = DocumentIntelligenceParams;
@ -168,7 +166,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
&self,
non_default_params: &CallArguments,
_model: &str,
) -> Result<DocumentIntelligenceParams, crate::ocr::Error> {
) -> Result<DocumentIntelligenceParams, Error> {
Ok(DocumentIntelligenceParams {
pages: normalize_pages_param(non_default_params.get("pages"))?,
features: normalize_features_param(non_default_params.get("features"))?,
@ -179,7 +177,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
let config = AzureAuthInputs {
azure_ad_token_provider: request.azure_ad_token_provider.clone(),
..AzureAuthInputs::from_sourced_optional_params(
@ -196,10 +194,10 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
request: &PreparedOcrRequest,
optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
let endpoint = nonblank(request.connection.api_base.clone())
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
.ok_or_else(|| crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?;
.ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?;
self.build_ocr_url(&endpoint, &request.model, optional_params)
}
@ -209,7 +207,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
document: OcrDocument,
_optional_params: &DocumentIntelligenceParams,
_headers: &[(String, String)],
) -> Result<DocumentIntelligenceRequest, crate::ocr::Error> {
) -> Result<DocumentIntelligenceRequest, Error> {
build_request(document)
}
@ -218,7 +216,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(
model,
raw_response,
@ -232,7 +230,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
model: &str,
raw_response: reqwest::Response,
context: OcrResponseContext<'_>,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let decoded = read_operation_response(
context.client.polling_http(),
raw_response,
@ -240,7 +238,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
context.headers,
context.connection,
context.request_format == OcrResponseFormat::Native,
context.host,
context.hooks,
)
.await?;
Ok(LiteLLMOcrResponse {
@ -250,7 +248,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
}
}
fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, crate::ocr::Error> {
fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, Error> {
let normalized = match pages {
None | Some(Value::Null) => return Ok(None),
Some(Value::Array(pages)) if pages.is_empty() => return Ok(None),
@ -259,12 +257,12 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, crate:
.map(|page| {
let page = page
.as_i64()
.ok_or_else(|| crate::ocr::Error::Pages("page index is out of range".into()))?;
.ok_or_else(|| Error::Pages("page index is out of range".into()))?;
if page < 0 {
return Err(crate::ocr::Error::Pages("negative page index".into()));
return Err(Error::Pages("negative page index".into()));
}
page.checked_add(1)
.ok_or_else(|| crate::ocr::Error::Pages("page index is out of range".into()))
.ok_or_else(|| Error::Pages("page index is out of range".into()))
})
.collect::<Result<BTreeSet<_>, _>>()?
.into_iter()
@ -274,9 +272,10 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, crate:
Some(Value::Array(tokens)) => tokens
.iter()
.map(|token| {
token.as_str().map(str::trim).ok_or_else(|| {
crate::ocr::Error::Pages("expected only integers or only strings".into())
})
token
.as_str()
.map(str::trim)
.ok_or_else(|| Error::Pages("expected only integers or only strings".into()))
})
.collect::<Result<Vec<_>, _>>()?
.join(","),
@ -286,13 +285,13 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, crate:
.collect::<Vec<_>>()
.join(","),
Some(_) => {
return Err(crate::ocr::Error::Pages(
return Err(Error::Pages(
"expected an array of integers or strings, or a native page range".into(),
));
}
};
if !normalized.split(',').all(valid_page_token) {
return Err(crate::ocr::Error::Pages("invalid native page range".into()));
return Err(Error::Pages("invalid native page range".into()));
}
Ok(Some(normalized))
}
@ -313,15 +312,15 @@ fn valid_page_token(token: &str) -> bool {
}
}
fn normalize_features_param(features: Option<&Value>) -> Result<Option<String>, crate::ocr::Error> {
fn normalize_features_param(features: Option<&Value>) -> Result<Option<String>, Error> {
let tokens = match features {
None | Some(Value::Null) => return Ok(None),
Some(Value::Array(names)) => names
.iter()
.map(|name| name.as_str().ok_or(crate::ocr::Error::Features))
.map(|name| name.as_str().ok_or(Error::Features))
.collect::<Result<Vec<_>, _>>()?,
Some(Value::String(names)) => names.split(',').collect(),
Some(_) => return Err(crate::ocr::Error::Features),
Some(_) => return Err(Error::Features),
};
if tokens.is_empty() {
return Ok(None);
@ -333,20 +332,19 @@ fn normalize_features_param(features: Option<&Value>) -> Result<Option<String>,
};
first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric)
}) {
return Err(crate::ocr::Error::Features);
return Err(Error::Features);
}
Ok(Some(normalized.join(",")))
}
fn build_request(document: OcrDocument) -> Result<DocumentIntelligenceRequest, crate::ocr::Error> {
fn build_request(document: OcrDocument) -> Result<DocumentIntelligenceRequest, Error> {
let source = document.source();
if source.is_empty() {
return Err(crate::ocr::Error::MissingDocumentUrl);
return Err(Error::MissingDocumentUrl);
}
Ok(if let Some(document) = InlineDocument::parse(source)? {
DocumentIntelligenceRequest::Base64Source {
base64_source: STANDARD
.encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?),
base64_source: STANDARD.encode(document.decode(OCR_INLINE_MAX_BYTES)?),
}
} else {
DocumentIntelligenceRequest::UrlSource {
@ -358,9 +356,9 @@ fn build_request(document: OcrDocument) -> Result<DocumentIntelligenceRequest, c
fn transform_completed_response(
model: &str,
response: AzureDocumentIntelligenceOperation,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
if response.status != Some(OperationStatus::Succeeded) {
return Err(crate::ocr::Error::OperationStatus(
return Err(Error::OperationStatus(
response
.status
.map(|status| status.to_string())
@ -373,8 +371,7 @@ fn transform_completed_response(
.into_iter()
.map(transform_azure_page)
.collect::<Result<Vec<_>, _>>()?;
let pages_processed =
i64::try_from(pages.len()).map_err(|_| crate::ocr::Error::NumericRange("pages"))?;
let pages_processed = i64::try_from(pages.len()).map_err(|_| Error::NumericRange("pages"))?;
Ok(LiteLLMOcrResponse {
content: result.content,
tables: result.tables,
@ -387,12 +384,12 @@ fn transform_completed_response(
})
}
fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage, crate::ocr::Error> {
fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage, Error> {
let index = page
.page_number
.unwrap_or(1)
.checked_sub(1)
.ok_or(crate::ocr::Error::NumericRange("page.pageNumber"))?;
.ok_or(Error::NumericRange("page.pageNumber"))?;
let dimensions = convert_dimensions(
page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH),
page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT),
@ -412,11 +409,7 @@ fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage,
})
}
fn convert_dimensions(
width: f64,
height: f64,
unit: &str,
) -> Result<OcrPageDimensions, crate::ocr::Error> {
fn convert_dimensions(width: f64, height: f64, unit: &str) -> Result<OcrPageDimensions, Error> {
let scale = if unit == "inch" {
AZURE_DI_DEFAULT_DPI as f64
} else {
@ -429,10 +422,10 @@ fn convert_dimensions(
})
}
fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result<i64, crate::ocr::Error> {
fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result<i64, Error> {
let value = value * scale;
if !value.is_finite() || value < i64::MIN as f64 || value >= -(i64::MIN as f64) {
return Err(crate::ocr::Error::NumericRange(field));
return Err(Error::NumericRange(field));
}
Ok(value.trunc() as i64)
}
@ -444,33 +437,38 @@ async fn read_operation_response(
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
host: &OcrHost,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, crate::ocr::Error> {
hooks: &dyn CallHooks<Error>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, Error> {
if response.status() != reqwest::StatusCode::ACCEPTED {
let bytes =
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes)
.await?;
crate::ocr::handler::emit_response_received(host, &bytes).await?;
return crate::ocr::json::decode_response(&bytes, native);
let bytes = crate::custom_httpx::llm_http_handler::read_response_bytes(
response,
connection.max_response_bytes,
)
.await?;
hooks.response_received(&bytes).await?;
return decode_response(&bytes, native);
}
let location = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.ok_or(crate::ocr::Error::PollLocation)?
.ok_or(Error::PollLocation)?
.to_string();
let original = Url::parse(original_url).map_err(|_| crate::ocr::Error::PollOrigin)?;
let operation = Url::parse(&location).map_err(|_| crate::ocr::Error::PollOrigin)?;
let original = Url::parse(original_url).map_err(|_| Error::PollOrigin)?;
let operation = Url::parse(&location).map_err(|_| Error::PollOrigin)?;
if original.origin() != operation.origin()
|| !operation.username().is_empty()
|| operation.password().is_some()
{
return Err(crate::ocr::Error::PollOrigin);
return Err(Error::PollOrigin);
}
let bytes =
crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?;
crate::ocr::handler::emit_response_received(host, &bytes).await?;
poll_operation(http_client, operation, headers, connection, native, host).await
let bytes = crate::custom_httpx::llm_http_handler::read_response_bytes(
response,
connection.max_response_bytes,
)
.await?;
hooks.response_received(&bytes).await?;
poll_operation(http_client, operation, headers, connection, native, hooks).await
}
async fn poll_operation(
@ -479,29 +477,35 @@ async fn poll_operation(
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
host: &OcrHost,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, crate::ocr::Error> {
hooks: &dyn CallHooks<Error>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, Error> {
let deadline = Instant::now()
.checked_add(connection.poll_timeout)
.ok_or(crate::ocr::Error::PollTimeout)?;
.ok_or(Error::PollTimeout)?;
loop {
let remaining = deadline
.checked_duration_since(Instant::now())
.filter(|remaining| !remaining.is_zero())
.ok_or(crate::ocr::Error::PollTimeout)?;
.ok_or(Error::PollTimeout)?;
let builder = http_client
.get(url.clone())
.timeout(remaining.min(connection.timeout));
let builder = crate::http_utils::with_headers(
let builder = crate::custom_httpx::http_handler::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]),
crate::custom_httpx::http_handler::HeaderPolicy::Only(&[
AZURE_DI_SUBSCRIPTION_HEADER,
"authorization",
]),
);
let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder))
.await
.map_err(|_| crate::ocr::Error::PollTimeout)?
.map_err(crate::transport::Error::from)?;
let response = tokio::time::timeout_at(
deadline,
crate::custom_httpx::http_handler::http_request(builder),
)
.await
.map_err(|_| Error::PollTimeout)?
.map_err(crate::custom_httpx::transport::Error::from)?;
let retry = response
.headers()
.get(reqwest::header::RETRY_AFTER)
@ -518,19 +522,19 @@ async fn poll_operation(
),
)
.await
.map_err(|_| crate::ocr::Error::PollTimeout)??;
.map_err(|_| Error::PollTimeout)??;
match &decoded.data.status {
Some(OperationStatus::Succeeded) => {
crate::ocr::handler::emit_response_received(host, decoded.text.as_bytes()).await?;
hooks.response_received(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(|_| crate::ocr::Error::PollTimeout)?;
.map_err(|_| Error::PollTimeout)?;
}
status => {
return Err(crate::ocr::Error::OperationStatus(
return Err(Error::OperationStatus(
status
.as_ref()
.map(ToString::to_string)
@ -547,7 +551,7 @@ impl AzureDocumentIntelligenceOcrConfig {
endpoint: &str,
model: &str,
params: &DocumentIntelligenceParams,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
let model = format!("{}:analyze", model_id(model)?);
ApiUrl::parse(endpoint)
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
@ -565,7 +569,7 @@ impl AzureDocumentIntelligenceOcrConfig {
)
.into_string()
})
.map_err(|_| crate::ocr::Error::RequestField {
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
}
@ -575,9 +579,9 @@ impl AzureDocumentIntelligenceOcrConfig {
connection: &OcrConnection,
config: &AzureAuthInputs,
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")
|| crate::http_utils::has_header(
) -> Result<Vec<(String, String)>, Error> {
if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization")
|| crate::custom_httpx::http_handler::has_header(
&connection.extra_headers,
AZURE_DI_SUBSCRIPTION_HEADER,
)
@ -604,7 +608,7 @@ impl AzureDocumentIntelligenceOcrConfig {
}
let token = super::super::common_utils::resolve_entra(config, env_lookup)
.await?
.ok_or(crate::ocr::Error::MissingAzureDocumentIntelligenceCredentials)?;
.ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?;
super::super::common_utils::validate_destination(connection, token.source())?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {}", token.value())))
@ -614,10 +618,10 @@ impl AzureDocumentIntelligenceOcrConfig {
}
}
fn model_id(model: &str) -> Result<&str, crate::ocr::Error> {
fn model_id(model: &str) -> Result<&str, Error> {
let model = model.rsplit('/').next().unwrap_or(model);
if matches!(model, "." | "..") {
return Err(crate::ocr::Error::DotModel);
return Err(Error::DotModel);
}
Ok(model)
}
@ -635,7 +639,7 @@ mod tests {
use super::*;
fn map(value: Value) -> Result<DocumentIntelligenceParams, crate::ocr::Error> {
fn map(value: Value) -> Result<DocumentIntelligenceParams, Error> {
let arguments = serde_json::from_value(value).unwrap();
AzureDocumentIntelligenceOcrConfig.map_ocr_params(&arguments, "model")
}
@ -809,410 +813,4 @@ mod tests {
(AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into())
);
}
use std::sync::{Arc, Mutex};
use litellm_callbacks::event::CallEvent;
use crate::ocr::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
fn query_value(url: &str, key: &str) -> Option<String> {
url::Url::parse(url)
.unwrap()
.query_pairs()
.find_map(|(name, value)| (name == key).then(|| value.into_owned()))
}
#[tokio::test]
async fn facade_maps_pages_features_and_url_document() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":{"pages":[]}
}))])
.await;
let mut request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}),
);
request.document = serde_json::from_value::<OcrDocument>(json!({
"type":"document_url",
"document_url":"https://example.com/document.pdf"
}))
.unwrap()
.into();
perform_ocr(request).await.unwrap();
server.await.unwrap();
let request = &seen.lock().unwrap()[0];
let target = request.split_whitespace().nth(1).unwrap();
let url = format!("{base}{target}");
assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3"));
assert_eq!(
query_value(&url, "features").as_deref(),
Some("keyValuePairs,languages")
);
let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false})
);
}
#[tokio::test]
async fn rejects_invalid_pages_features_and_format() {
for options in [
json!({"pages":[true]}),
json!({"pages":[1,"2"]}),
json!({"pages":[-1]}),
json!({"pages":"1&&features=bad"}),
json!({"features":"languages&pages=1"}),
json!({"req_format":"azure"}),
] {
let request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
"http://127.0.0.1:1",
options.clone(),
);
let rejected = perform_ocr(request).await.is_err();
assert!(rejected, "accepted {options}");
}
}
#[tokio::test]
async fn inline_document_decodes_to_base64_source() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded"
}))])
.await;
let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
perform_ocr(request).await.unwrap();
server.await.unwrap();
let request = &seen.lock().unwrap()[0];
let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(body, json!({"base64Source":"YWJj"}));
}
#[tokio::test]
async fn immediate_response_normalizes_pages_and_preserves_native() {
let operation = json!({
"status":"succeeded",
"operationExtension":42,
"analyzeResult":{
"content":"A\n\nB",
"tables":[{"cells":[]}],
"keyValuePairs":[{"key":{"content":"A"}}],
"pages":[{
"pageNumber":"2",
"width":"8.5",
"height":11,
"unit":"inch",
"lines":[{"content":"A"},{"content":null},{"content":"B"}]
}]
}
});
let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await;
let result = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"req_format":"native"}),
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0].index, 1);
assert_eq!(result.pages[0].markdown, "A\n\nB");
assert_eq!(
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
json!({"width":816,"height":1056,"dpi":96})
);
assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1));
let serialized = result.clone().into_json();
assert_eq!(serialized["content"], "A\n\nB");
assert_eq!(serialized["tables"], json!([{"cells":[]}]));
assert_eq!(
serialized["keyValuePairs"],
json!([{"key":{"content":"A"}}])
);
assert!(serialized.get("key_value_pairs").is_none());
assert_eq!(
result.provider_native_response.as_ref(),
operation.as_object()
);
}
#[tokio::test]
async fn accepted_response_polls_to_success_with_only_credentials() {
let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}});
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 200,
headers: vec![("Retry-After", "0".into())],
body: json!({"status":"running"}),
},
MockResponse::json(operation.clone()),
])
.await;
let mut request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"req_format":"native"}),
);
request
.transport
.extra_headers
.push(("X-Trace".into(), "initial-only".into()));
let result = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(
result.provider_native_response.as_ref(),
operation.as_object()
);
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 3);
assert!(requests[0].to_ascii_lowercase().contains("x-trace:"));
for poll in &requests[1..] {
assert!(!poll.to_ascii_lowercase().contains("x-trace:"));
assert!(
poll.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: test-key")
);
}
}
#[tokio::test]
async fn accepted_response_emits_response_received_for_submission_and_completed_poll() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({"submitted": true}),
},
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let responses_received = Arc::new(Mutex::new(Vec::new()));
let request_count = seen.clone();
let observed = responses_received.clone();
let host = LocalOcrHost::new(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
observed
.lock()
.unwrap()
.push((request_count.lock().unwrap().len(), raw.body.clone()));
}
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
assert_eq!(
*responses_received.lock().unwrap(),
[
(1, r#"{"submitted":true}"#.to_string()),
(2, r#"{"status":"succeeded"}"#.to_string()),
]
);
}
#[tokio::test]
async fn polling_forwards_bearer_credentials() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
request.credentials.api_key = None;
request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())];
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert!(
requests[1]
.to_ascii_lowercase()
.contains("authorization: bearer token")
);
}
#[tokio::test]
async fn polling_does_not_follow_redirects() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 302,
headers: vec![("Location", "{base}/redirected".into())],
body: json!({}),
},
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
assert!(error.to_string().contains("status 302"), "{error}");
assert_eq!(seen.lock().unwrap().len(), 2);
server.abort();
}
#[tokio::test]
async fn polling_rejects_terminal_failure() {
let (base, _, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse::json(json!({"status":"failed"})),
])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("status failed"));
}
#[tokio::test]
async fn malformed_provider_pages_report_response_paths() {
for (analysis, path) in [
(json!({"pages":null}), "pages"),
(json!({"pages":[null]}), "pages[0]"),
(json!({"pages":[{"lines":null}]}), "lines"),
(json!({"pages":[{"width":"bad"}]}), "width"),
] {
let (base, _, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":analysis
}))])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains(path), "{error}");
}
}
#[tokio::test]
async fn rejects_missing_invalid_and_cross_origin_operation_locations() {
for headers in [
Vec::new(),
vec![("Operation-Location", "/relative".into())],
vec![("Operation-Location", "http://example.com/operation".into())],
vec![(
"Operation-Location",
"http://user:password@127.0.0.1/operation".into(),
)],
] {
let (base, _, server) = mock_server(vec![MockResponse {
status: 202,
headers,
body: json!({}),
}])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("operation-location"));
}
}
#[tokio::test]
async fn polling_deadline_bounds_retry_delay() {
let (base, _, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 200,
headers: vec![("Retry-After", "9999".into())],
body: json!({"status":"notStarted"}),
},
])
.await;
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
request.transport.poll_timeout = std::time::Duration::from_millis(100);
let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request))
.await
.unwrap()
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("timed out"));
}
#[tokio::test]
async fn model_id_is_encoded_and_dot_segments_are_rejected() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded"
}))])
.await;
perform_ocr(wire_request(
"azure_ai/doc-intelligence/a ?#é",
&base,
json!({}),
))
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze"));
for model in [
"azure_ai/doc-intelligence/.",
"azure_ai/doc-intelligence/..",
] {
let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains("dot segment"));
}
}
}

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

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

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

@ -0,0 +1,330 @@
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::AzureAuthInputs;
use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl};
use serde_json::Value;
use crate::{
base_llm::ocr::{
document::{inline_remote_document, validate_inline_document},
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrRequestContext,
OcrResponseFormat, PreparedOcrRequest, credential_env,
},
},
custom_httpx::llm_http_handler::OcrClient,
mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest},
};
const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr";
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 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, Error> {
MistralOcrConfig.map_ocr_params(non_default_params, model)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, 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, 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, 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, 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: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, Error> {
MistralOcrConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), Error> {
validate_inline_document(&crate::custom_httpx::llm_http_handler::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, Error> {
nonblank(api_base.map(str::to_string))
.or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV)))
.ok_or(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)>, Error> {
Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?;
if crate::custom_httpx::http_handler::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(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, 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(|_| 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(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())
);
}
#[tokio::test]
async fn environment_supplies_api_base_and_bearer_key() {
let env = |name: &str| match name {
AZURE_AI_API_BASE_ENV => Some("https://env.example".to_string()),
AZURE_AI_API_KEY_ENV => Some("env-key".to_string()),
_ => None,
};
let connection = OcrConnection::default();
let headers = AzureAiOcrConfig
.resolve_headers(&connection, &Default::default(), &env)
.await
.unwrap();
let url = AzureAiOcrConfig.build_ocr_url(None, &env).unwrap();
assert_eq!(
headers,
[("Authorization".to_string(), "Bearer env-key".to_string())]
);
assert_eq!(url, "https://env.example/providers/mistral/azure/ocr");
}
}

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

@ -2,4 +2,5 @@ pub mod anthropic_messages;
pub mod audio_transcription;
pub mod base_model_iterator;
pub mod chat;
pub mod ocr;
pub mod responses;

225
litellm-rust/crates/llms/src/base_llm/ocr/document.rs Normal file
View file

@ -0,0 +1,225 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime};
use reqwest::Url;
use crate::{
base_llm::ocr::{
error::Error,
transformation::{
OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS, OcrConnection, OcrDocument,
},
},
custom_httpx::{
media::{DownloadPolicy, Error as MediaError, MediaFetcher},
transport::Error as TransportError,
},
};
pub struct InlineDocument<'a>(DataUrl<'a>);
impl<'a> InlineDocument<'a> {
pub fn parse(source: &'a str) -> Result<Option<Self>, Error> {
match DataUrl::process(source) {
Ok(url) => Ok(Some(Self(url))),
Err(DataUrlError::NotADataUrl) => Ok(None),
Err(DataUrlError::NoComma) => Err(Error::InvalidDataUri),
}
}
pub fn mime_type(&self) -> &Mime {
self.0.mime_type()
}
pub fn decode(&self, max_bytes: usize) -> Result<Vec<u8>, Error> {
let mut body = Vec::new();
self.0
.decode(|bytes| {
if bytes.len() > max_bytes.saturating_sub(body.len()) {
return Err(Error::InlineDocumentTooLarge);
}
body.extend_from_slice(bytes);
Ok(())
})
.map_err(|error| match error {
DecodeError::InvalidBase64(_) => Error::InvalidDataUri,
DecodeError::WriteError(error) => error,
})?;
Ok(body)
}
}
pub fn validate_inline_document(document: &OcrDocument) -> Result<(), Error> {
let inline = InlineDocument::parse(document.source())?.ok_or(Error::InvalidDataUri)?;
inline.decode(OCR_INLINE_MAX_BYTES)?;
Ok(())
}
pub async fn inline_remote_document(
fetcher: &MediaFetcher,
document: OcrDocument,
connection: &OcrConnection,
) -> Result<OcrDocument, Error> {
let source = document.source();
if !document.is_remote() {
validate_inline_document(&document)?;
return Ok(document);
}
let url = Url::parse(source).map_err(|_| Error::RequestField {
path: "document URL".into(),
})?;
let downloaded = fetcher
.fetch(
url,
DownloadPolicy {
timeout: connection.timeout,
max_bytes: connection.max_download_bytes,
max_redirects: OCR_MAX_FETCH_REDIRECTS,
},
)
.await
.map_err(map_media_error)?;
let result = document.with_source(format!(
"data:{};base64,{}",
downloaded.content_type,
STANDARD.encode(downloaded.bytes)
));
validate_inline_document(&result)?;
Ok(result)
}
fn map_media_error(error: MediaError) -> Error {
match error {
MediaError::BlockedUrl => Error::BlockedDocumentUrl,
MediaError::DownloadDisabled => Error::DownloadDisabled,
MediaError::DownloadTooLarge => Error::DownloadTooLarge,
MediaError::TooManyRedirects => Error::TooManyRedirects,
MediaError::MissingRedirectLocation => Error::MissingRedirectLocation,
MediaError::InvalidRedirect => Error::InvalidRedirect,
MediaError::Http(status) => TransportError::Http {
status,
body: "OCR document download failed".into(),
}
.into(),
MediaError::Timeout => TransportError::Http {
status: 408,
body: "OCR document download timed out".into(),
}
.into(),
MediaError::Transport(error) => error.into(),
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap as Map;
use super::*;
fn document(source: &str) -> OcrDocument {
OcrDocument::DocumentUrl {
document_url: source.into(),
extra_fields: Map::new(),
}
}
#[test]
fn decodes_data_urls_and_limits_decoded_size() {
for (source, expected) in [
("data:application/pdf;base64,YWJj", b"abc".as_slice()),
("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()),
("data:,a%20b%00%FF", b"a b\0\xff".as_slice()),
] {
let inline = InlineDocument::parse(source).unwrap().unwrap();
assert_eq!(inline.decode(expected.len()).unwrap(), expected);
assert!(matches!(
inline.decode(expected.len() - 1),
Err(Error::InlineDocumentTooLarge)
));
}
}
#[test]
fn preserves_mime_parameters_and_standard_default() {
let inline = InlineDocument::parse("data:application/pdf;version=1.7;base64,YQ==")
.unwrap()
.unwrap();
assert!(inline.mime_type().matches("application", "pdf"));
assert_eq!(inline.mime_type().get_parameter("version"), Some("1.7"));
let default = InlineDocument::parse("data:,a").unwrap().unwrap();
assert!(default.mime_type().matches("text", "plain"));
assert_eq!(
default.mime_type().get_parameter("charset"),
Some("US-ASCII")
);
}
#[test]
fn rejects_invalid_inline_documents() {
for source in [
"https://example.com/document.pdf",
"data:application/pdf;base64",
"data:application/pdf;base64,INVALID!",
] {
assert!(validate_inline_document(&document(source)).is_err());
}
}
#[tokio::test]
async fn remote_conversion_preserves_kind_and_isolates_provider_credentials() {
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = vec![0_u8; 2048];
let count = socket.read(&mut request).await.unwrap();
socket
.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: image/png; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc")
.await
.unwrap();
String::from_utf8_lossy(&request[..count]).into_owned()
});
let mut provider_headers = reqwest::header::HeaderMap::new();
provider_headers.insert(
reqwest::header::AUTHORIZATION,
reqwest::header::HeaderValue::from_static("Bearer provider-secret"),
);
let provider_http = reqwest::Client::builder()
.default_headers(provider_headers)
.build()
.unwrap();
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.unwrap();
let client = crate::custom_httpx::llm_http_handler::OcrClient::for_test(
provider_http,
document_http,
);
let converted = inline_remote_document(
client.document_fetcher(),
OcrDocument::ImageUrl {
image_url: format!("http://{address}/image"),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
},
&OcrConnection::default(),
)
.await
.unwrap();
let request = server.await.unwrap();
assert_eq!(
converted,
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,YWJj".into(),
extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]),
}
);
assert!(!request.to_ascii_lowercase().contains("authorization"));
assert!(!request.contains("provider-secret"));
}
}

8
litellm-rust/crates/core/src/ocr/error.rs → litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -95,11 +95,11 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
Transport(#[from] crate::custom_httpx::transport::Error),
#[error(transparent)]
Params(#[from] litellm_core_utils::params::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
Headers(#[from] crate::custom_httpx::http_handler::HeaderError),
}
impl From<litellm_core_utils::call_arguments::ArgumentError> for Error {
@ -114,7 +114,9 @@ impl Error {
pub fn http_status_code(&self) -> Option<u16> {
match self {
Self::Provider { status, .. }
| Self::Transport(crate::transport::Error::Http { status, .. }) => Some(*status),
| Self::Transport(crate::custom_httpx::transport::Error::Http { status, .. }) => {
Some(*status)
}
error if error.is_request() => Some(400),
_ => None,
}

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

@ -0,0 +1,3 @@
pub mod document;
pub mod error;
pub mod transformation;

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

@ -0,0 +1,667 @@
use std::{collections::BTreeMap, future::Future, time::Duration};
use litellm_auth::{InputSource, Sourced, TokenProviderHandle};
use litellm_core_utils::{
call_arguments::CallArguments,
serde_compat::{FiniteF64, LaxI64},
};
use serde::{
Deserialize, Serialize,
de::{DeserializeOwned, IntoDeserializer},
};
use serde_json::{Map, Value};
use serde_with::serde_as;
use crate::{
base_llm::ocr::error::Error,
custom_httpx::llm_http_handler::{
CallHooks, OcrClient, read_response_bytes, transform_request_body,
},
};
pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
pub const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
pub const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
pub const OCR_MAX_FETCH_REDIRECTS: usize = 10;
pub const OCR_POLL_TIMEOUT_SECS: u64 = 120;
pub const OCR_POLL_RETRY_SECS: u64 = 2;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum OcrDocument {
#[serde(rename = "document_url")]
DocumentUrl {
document_url: String,
#[serde(flatten)]
extra_fields: BTreeMap<String, Option<String>>,
},
#[serde(rename = "image_url")]
ImageUrl {
image_url: String,
#[serde(flatten)]
extra_fields: BTreeMap<String, Option<String>>,
},
}
impl OcrDocument {
pub fn source(&self) -> &str {
match self {
Self::DocumentUrl { document_url, .. } => document_url,
Self::ImageUrl { image_url, .. } => image_url,
}
}
pub fn is_remote(&self) -> bool {
let source = self.source();
source.starts_with("http://") || source.starts_with("https://")
}
pub fn with_source(self, source: String) -> Self {
match self {
Self::DocumentUrl { extra_fields, .. } => Self::DocumentUrl {
document_url: source,
extra_fields,
},
Self::ImageUrl { extra_fields, .. } => Self::ImageUrl {
image_url: source,
extra_fields,
},
}
}
}
impl TryFrom<Value> for OcrDocument {
type Error = Error;
fn try_from(value: Value) -> Result<Self, Self::Error> {
decode_request_value(value, "document")
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum OcrResponseFormat {
#[default]
Litellm,
Native,
}
#[derive(Clone, Default)]
pub struct OcrCredentialInputs {
pub api_key: Option<Sourced<String>>,
pub dynamic_api_key: Option<Sourced<String>>,
pub api_base: Option<Sourced<String>>,
pub dynamic_api_base: Option<Sourced<String>>,
}
impl OcrCredentialInputs {
pub fn new(
api_key: Option<String>,
api_key_source: InputSource,
api_base: Option<String>,
api_base_source: InputSource,
) -> Self {
Self {
api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)),
dynamic_api_key: None,
api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)),
dynamic_api_base: None,
}
}
}
#[derive(Clone)]
pub struct OcrTransportConfig {
pub extra_headers: Vec<(String, String)>,
pub extra_headers_source: InputSource,
pub timeout: Duration,
pub max_download_bytes: u64,
pub max_response_bytes: usize,
pub poll_timeout: Duration,
}
impl Default for OcrTransportConfig {
fn default() -> Self {
Self {
extra_headers: Vec::new(),
extra_headers_source: InputSource::Deployment,
timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
max_download_bytes: OCR_DOWNLOAD_MAX_BYTES,
max_response_bytes: OCR_RESPONSE_MAX_BYTES,
poll_timeout: Duration::from_secs(OCR_POLL_TIMEOUT_SECS),
}
}
}
impl OcrTransportConfig {
pub fn with_overrides(
self,
extra_headers: Vec<(String, String)>,
extra_headers_source: InputSource,
timeout: Option<Duration>,
) -> Self {
Self {
extra_headers,
extra_headers_source,
timeout: timeout.unwrap_or(self.timeout),
..self
}
}
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[derive(Clone)]
pub struct OcrConnection {
pub api_key: Option<String>,
pub api_key_source: InputSource,
pub api_base: Option<String>,
pub api_base_source: InputSource,
pub extra_headers: Vec<(String, String)>,
pub extra_headers_source: InputSource,
pub timeout: Duration,
pub max_download_bytes: u64,
pub max_response_bytes: usize,
pub poll_timeout: Duration,
}
impl OcrConnection {
pub fn new(credentials: ResolvedOcrCredentials, transport: OcrTransportConfig) -> Self {
let api_key_source = credentials
.api_key
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Deployment);
let api_base_source = credentials
.api_base
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Deployment);
Self {
api_key: credentials.api_key.map(Sourced::into_value),
api_key_source,
api_base: credentials.api_base.map(Sourced::into_value),
api_base_source,
extra_headers: transport.extra_headers,
extra_headers_source: transport.extra_headers_source,
timeout: transport.timeout,
max_download_bytes: transport.max_download_bytes,
max_response_bytes: transport.max_response_bytes,
poll_timeout: transport.poll_timeout,
}
}
}
impl Default for OcrConnection {
fn default() -> Self {
Self::new(
ResolvedOcrCredentials::default(),
OcrTransportConfig::default(),
)
}
}
#[derive(Clone, Default)]
pub struct ResolvedOcrCredentials {
pub api_key: Option<Sourced<String>>,
pub api_base: Option<Sourced<String>>,
}
pub struct PreparedOcrRequest {
pub model: String,
pub document: OcrDocument,
pub connection: OcrConnection,
/// Whether the caller handed over the document as is, so the wire body's document
/// is the caller's own input rather than something the route prepared.
pub caller_document: bool,
pub optional_params: CallArguments,
pub input_sources: BTreeMap<String, InputSource>,
pub azure_ad_token_provider: Option<TokenProviderHandle>,
}
impl PreparedOcrRequest {
pub fn response_format(&self) -> Result<OcrResponseFormat, Error> {
response_format(&self.optional_params)
}
}
pub fn response_format(optional_params: &CallArguments) -> Result<OcrResponseFormat, Error> {
optional_params
.get("req_format")
.filter(|value| !value.is_null())
.map(|value| serde_json::from_value(value.clone()).map_err(|_| Error::RequestFormat))
.transpose()
.map(|format| format.unwrap_or_default())
}
#[serde_as]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrPageDimensions {
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub dpi: Option<i64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub height: Option<i64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub width: Option<i64>,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrPageImage {
pub image_base64: Option<String>,
pub bbox: Option<Map<String, Value>>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[serde_as]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrPage {
#[serde_as(deserialize_as = "LaxI64")]
pub index: i64,
pub markdown: String,
pub images: Option<Vec<OcrPageImage>>,
pub dimensions: Option<OcrPageDimensions>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[serde_as]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct OcrUsageInfo {
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub pages_processed: Option<i64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub pages_processed_annotation: Option<i64>,
#[serde_as(deserialize_as = "Option<FiniteF64>")]
pub credits: Option<f64>,
#[serde_as(deserialize_as = "Option<LaxI64>")]
pub doc_size_bytes: Option<i64>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct LiteLLMOcrResponse {
pub pages: Vec<OcrPage>,
pub model: String,
pub document_annotation: Option<Value>,
pub usage_info: Option<OcrUsageInfo>,
pub content: Option<String>,
pub tables: Option<Vec<Map<String, Value>>>,
#[serde(rename = "keyValuePairs")]
pub key_value_pairs: Option<Vec<Map<String, Value>>>,
#[serde(default = "ocr_object")]
pub object: String,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub provider_native_response: Option<Map<String, Value>>,
}
impl LiteLLMOcrResponse {
pub fn new(model: impl Into<String>, pages: Vec<OcrPage>) -> Self {
Self {
pages,
model: model.into(),
document_annotation: None,
usage_info: None,
content: None,
tables: None,
key_value_pairs: None,
object: ocr_object(),
extra_fields: Map::new(),
provider_native_response: None,
}
}
pub fn into_json(self) -> Value {
serde_json::to_value(self).expect("OCR response fields are JSON-compatible")
}
}
fn ocr_object() -> String {
"ocr".into()
}
#[derive(Debug)]
pub struct DecodedOcrResponse<T> {
pub data: T,
pub native: Option<Map<String, Value>>,
pub text: String,
}
pub fn decode_request_value<T: DeserializeOwned>(value: Value, prefix: &str) -> Result<T, Error> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
Error::RequestField {
path: format!("{prefix}.{}", error.path()),
}
})
}
pub fn decode_response_value<T: DeserializeOwned>(value: Value, prefix: &str) -> Result<T, Error> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
Error::ResponseField {
path: format!("{prefix}.{}", error.path()),
}
})
}
pub fn decode_response<T: DeserializeOwned>(
bytes: &[u8],
native: bool,
) -> Result<DecodedOcrResponse<T>, Error> {
let mut deserializer = serde_json::Deserializer::from_slice(bytes);
let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| {
Error::ResponseField {
path: error.path().to_string(),
}
})?;
deserializer.end().map_err(|_| Error::ResponseField {
path: "response".into(),
})?;
let native = if native {
Some(
serde_json::from_slice(bytes).map_err(|_| Error::ResponseField {
path: "response".into(),
})?,
)
} else {
None
};
Ok(DecodedOcrResponse {
data,
native,
text: String::from_utf8_lossy(bytes).into_owned(),
})
}
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 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 struct OcrRequestContext<'a> {
pub client: &'a OcrClient,
pub connection: &'a OcrConnection,
}
#[derive(Clone, Copy)]
pub struct OcrResponseContext<'a> {
pub client: &'a OcrClient,
pub connection: &'a OcrConnection,
pub hooks: &'a dyn CallHooks<Error>,
pub request_format: OcrResponseFormat,
pub url: &'a str,
pub headers: &'a [(String, String)],
}
pub 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, Error>;
fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> impl Future<Output = Result<Self::Environment, Error>> + Send;
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
optional_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, Error>;
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &Self::OcrParams,
headers: &[(String, String)],
) -> Result<Self::ProviderRequest, 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, 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, Error>;
fn async_transform_ocr_response(
&self,
model: &str,
raw_response: reqwest::Response,
context: OcrResponseContext<'_>,
) -> impl Future<Output = Result<LiteLLMOcrResponse, Error>> + Send {
async move {
let bytes =
read_response_bytes(raw_response, context.connection.max_response_bytes).await?;
context.hooks.response_received(&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)>,
) -> Error {
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<(), 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,
hooks: &dyn CallHooks<Error>,
) -> impl Future<Output = Result<reqwest::Request, 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?;
transform_request_body(self, client, request, &url, headers, body, hooks).await
}
}
}
pub fn decode_and_normalize_response<T: DeserializeOwned>(
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
normalize: impl FnOnce(&str, T) -> Result<LiteLLMOcrResponse, Error>,
) -> Result<LiteLLMOcrResponse, Error> {
let decoded = decode_response(raw_response, request_format == OcrResponseFormat::Native)?;
Ok(LiteLLMOcrResponse {
provider_native_response: decoded.native,
..normalize(model, decoded.data)?
})
}
pub fn credential_env(name: &str) -> Option<String> {
std::env::var(name).ok()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn normalized_response_rejects_invalid_shared_fields() {
for fields in [
json!({"pages":[{}]}),
json!({"pages":[{"index":0,"markdown":false}]}),
json!({"pages":[{"index":0,"markdown":"","images":[{"bbox":[]}]}]}),
json!({"usage_info":{"pages_processed":1.5}}),
json!({"tables":[false]}),
json!({"keyValuePairs":[[]]}),
json!({"provider_native_response":[]}),
] {
let payload: Map<String, Value> = json!({"model":"model", "pages":[]})
.as_object()
.unwrap()
.iter()
.chain(fields.as_object().unwrap())
.map(|(key, value)| (key.clone(), value.clone()))
.collect();
assert!(serde_json::from_value::<LiteLLMOcrResponse>(Value::Object(payload)).is_err());
}
assert!(
serde_json::from_value::<OcrDocument>(json!({
"type":"image_url", "image_url":"https://example.com/image", "detail":42
}))
.is_err()
);
}
#[test]
fn numeric_coercion_preserves_integer_precision_and_rejects_fractional_values() {
for (value, expected) in [
(json!("9007199254740993.0"), 9_007_199_254_740_993),
(json!("+2.000"), 2),
(json!("1_000"), 1000),
(json!(true), 1),
(json!(2.0), 2),
] {
let page: OcrPage =
serde_json::from_value(json!({"index":value,"markdown":""})).unwrap();
assert_eq!(page.index, expected);
}
for value in [
json!("1e2"),
json!(".0"),
json!("2."),
json!("_2"),
json!("2__0"),
json!(2.5),
json!(null),
] {
assert!(
serde_json::from_value::<OcrPage>(json!({"index":value,"markdown":""})).is_err()
);
}
}
#[rstest::rstest]
#[case::document_url("document_url", "document_name", "application/pdf")]
#[case::image_url("image_url", "detail", "image/png")]
fn document_variants_preserve_provider_fields_when_rewriting_sources(
#[case] kind: &str,
#[case] field: &str,
#[case] mime_type: &str,
#[values(json!("kept"), Value::Null)] extra: Value,
) {
let original = "https://example.com/input";
let replacement = format!("data:{mime_type};base64,AA==");
let document: OcrDocument =
serde_json::from_value(json!({"type": kind, kind: original, field: extra})).unwrap();
assert_eq!(document.source(), original);
assert_eq!(
serde_json::to_value(document.with_source(replacement.clone())).unwrap(),
json!({"type": kind, kind: replacement, field: extra})
);
}
#[test]
fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() {
let response = LiteLLMOcrResponse {
extra_fields: json!({"provider_field":"kept"})
.as_object()
.unwrap()
.clone(),
..LiteLLMOcrResponse::new("model", vec![])
};
let serialized = response.into_json();
assert_eq!(serialized["provider_field"], "kept");
assert!(serialized.get("provider_native_response").is_none());
}
}

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

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

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

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

232
litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs → litellm-rust/crates/llms/src/cohere/ocr/transformation.rs
View file

@ -8,37 +8,39 @@ use serde_json::{Map, Value};
use serde_with::serde_as;
use crate::{
constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE},
llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response},
ocr::{
OcrClient,
base_llm::ocr::{
document::InlineDocument,
prepare::credential_env,
types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage,
OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument,
OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
credential_env, decode_and_normalize_response, decode_response_value,
},
},
custom_httpx::llm_http_handler::OcrClient,
};
const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com";
const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY";
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 {
pub enum OutputFormat {
#[default]
Markdown,
Blocks,
}
#[derive(Default, Deserialize, Serialize)]
pub(crate) struct CohereOptions {
pub struct CohereOptions {
#[serde(skip_serializing_if = "Option::is_none")]
pub output_format: Option<OutputFormat>,
}
#[derive(Deserialize, Serialize)]
pub(crate) struct CohereRequest {
pub struct CohereRequest {
pub model: String,
pub document: CohereParseDocument,
pub output_format: String,
@ -46,13 +48,13 @@ pub(crate) struct CohereRequest {
#[derive(Deserialize, Serialize)]
#[serde(tag = "type")]
pub(crate) enum CohereParseDocument {
pub enum CohereParseDocument {
#[serde(rename = "image_url")]
ImageUrl { image_url: String },
}
#[derive(Deserialize)]
pub(crate) struct CohereResponse {
pub struct CohereResponse {
#[serde(default)]
pages: Vec<CoherePage>,
meta: Option<CohereMeta>,
@ -87,7 +89,7 @@ struct CohereBilledUnits {
}
#[derive(Default)]
pub(crate) struct CohereParseConfig;
pub struct CohereParseConfig;
impl BaseOcrConfig for CohereParseConfig {
type OcrParams = CohereOptions;
@ -113,7 +115,7 @@ impl BaseOcrConfig for CohereParseConfig {
&self,
non_default_params: &CallArguments,
_model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
) -> Result<CohereOptions, Error> {
Ok(parse_options(non_default_params)?)
}
@ -121,7 +123,7 @@ impl BaseOcrConfig for CohereParseConfig {
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
self.resolve_headers(&request.connection, &credential_env)
}
@ -130,7 +132,7 @@ impl BaseOcrConfig for CohereParseConfig {
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
self.build_ocr_url(
request
.connection
@ -146,7 +148,7 @@ impl BaseOcrConfig for CohereParseConfig {
document: OcrDocument,
optional_params: &CohereOptions,
_headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
) -> Result<CohereRequest, Error> {
let image_url = image_url(document)?;
Ok(build_request(model, image_url, optional_params))
}
@ -156,12 +158,12 @@ impl BaseOcrConfig for CohereParseConfig {
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, 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)?)
fn validate_request_body(&self, body: &Value) -> Result<(), Error> {
validate_document(&crate::custom_httpx::llm_http_handler::body_document(body)?)
}
}
@ -170,8 +172,9 @@ impl CohereParseConfig {
&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") {
) -> Result<Vec<(String, String)>, Error> {
if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization")
{
return Ok(connection.extra_headers.clone());
}
let key = connection
@ -186,7 +189,7 @@ impl CohereParseConfig {
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| {
crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication(
Error::Auth(litellm_auth::Error::ProviderAuthentication(
"Missing COHERE_API_KEY - set it in the environment or pass api_key".into(),
))
})?;
@ -197,7 +200,7 @@ impl CohereParseConfig {
)
}
fn build_ocr_url(&self, api_base: &str) -> Result<String, crate::ocr::Error> {
fn build_ocr_url(&self, api_base: &str) -> Result<String, Error> {
let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(invalid_api_base());
@ -209,35 +212,35 @@ impl CohereParseConfig {
}
}
pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), crate::ocr::Error> {
pub fn validate_document(document: &OcrDocument) -> Result<(), Error> {
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(crate::ocr::Error::CohereImageOnly);
return Err(Error::CohereImageOnly);
};
if image_url.is_empty() {
return Err(crate::ocr::Error::CohereImageOnly);
return Err(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);
return Err(Error::CohereImageOnly);
}
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
inline.decode(OCR_INLINE_MAX_BYTES)?;
}
Ok(())
}
pub(crate) fn normalize_response(
pub fn normalize_response(
model: &str,
response: CohereResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, 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"))
i64::try_from(response.pages.len()).map_err(|_| Error::NumericRange("pages"))
})?;
let pages = response
.pages
.into_iter()
.enumerate()
.map(|(position, page)| normalize_page(page, position))
.collect::<Result<Vec<_>, crate::ocr::Error>>()?;
.collect::<Result<Vec<_>, Error>>()?;
Ok(LiteLLMOcrResponse {
usage_info: Some(OcrUsageInfo {
pages_processed: Some(pages_processed),
@ -247,10 +250,10 @@ pub(crate) fn normalize_response(
})
}
fn image_url(document: OcrDocument) -> Result<String, crate::ocr::Error> {
fn image_url(document: OcrDocument) -> Result<String, Error> {
validate_document(&document)?;
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(crate::ocr::Error::CohereImageOnly);
return Err(Error::CohereImageOnly);
};
Ok(image_url)
}
@ -267,19 +270,16 @@ fn build_request(model: &str, image_url: String, params: &CohereOptions) -> Cohe
}
}
fn page_image(
mut image: Map<String, Value>,
path: &str,
) -> Result<OcrPageImage, crate::ocr::Error> {
fn page_image(mut image: Map<String, Value>, path: &str) -> Result<OcrPageImage, 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)
decode_response_value(Value::Object(image), path)
}
fn normalize_page(page: CoherePage, position: usize) -> Result<OcrPage, crate::ocr::Error> {
fn normalize_page(page: CoherePage, position: usize) -> Result<OcrPage, Error> {
let index = page.index.map(Ok).unwrap_or_else(|| {
i64::try_from(position).map_err(|_| crate::ocr::Error::NumericRange("page index"))
i64::try_from(position).map_err(|_| Error::NumericRange("page index"))
})?;
let (markdown, images) = match page.markdown {
Some(markdown) => {
@ -326,8 +326,8 @@ 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 {
fn invalid_api_base() -> Error {
Error::RequestField {
path: "api_base".into(),
}
}
@ -338,42 +338,7 @@ mod tests {
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", "timeout":30,
"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_for_test(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]}}
})
);
}
use crate::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig;
#[rstest]
#[case::cohere(false)]
@ -384,8 +349,7 @@ mod tests {
}))
.unwrap();
let mapped = if azure {
crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig
.map_ocr_params(&arguments, "parse")
AzureAICohereParseConfig.map_ocr_params(&arguments, "parse")
} else {
CohereParseConfig.map_ocr_params(&arguments, "parse")
}
@ -403,7 +367,7 @@ mod tests {
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 })
Err(Error::RequestField { path })
if path == "optional_params.output_format"
));
}
@ -465,7 +429,7 @@ mod tests {
.unwrap();
assert!(matches!(
normalize_response("parse", response).unwrap_err(),
crate::ocr::Error::ResponseField { path }
Error::ResponseField { path }
if path == "pages[0].markdown.images[0].image_base64"
));
}
@ -501,33 +465,6 @@ mod tests {
);
}
#[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_for_test(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!({
@ -621,10 +558,10 @@ mod tests {
#[rstest]
fn response_types_documented_block_variants(
#[values(
crate::ocr::types::OcrResponseFormat::Litellm,
crate::ocr::types::OcrResponseFormat::Native
crate::base_llm::ocr::transformation::OcrResponseFormat::Litellm,
crate::base_llm::ocr::transformation::OcrResponseFormat::Native
)]
response_format: crate::ocr::types::OcrResponseFormat,
response_format: crate::base_llm::ocr::transformation::OcrResponseFormat,
) {
let payload = json!({
"pages": [{
@ -694,10 +631,10 @@ mod tests {
Some(1)
);
match response_format {
crate::ocr::types::OcrResponseFormat::Litellm => {
crate::base_llm::ocr::transformation::OcrResponseFormat::Litellm => {
assert!(normalized.provider_native_response.is_none());
}
crate::ocr::types::OcrResponseFormat::Native => {
crate::base_llm::ocr::transformation::OcrResponseFormat::Native => {
assert_eq!(
normalized.provider_native_response.as_ref(),
payload.as_object()
@ -717,7 +654,7 @@ mod tests {
fn request_requires_image(#[case] value: Value) {
assert!(matches!(
validate_document(&serde_json::from_value(value).unwrap()),
Err(crate::ocr::Error::CohereImageOnly)
Err(Error::CohereImageOnly)
));
}
@ -786,7 +723,7 @@ mod tests {
},
&|_| None,
),
Err(crate::ocr::Error::Auth(_))
Err(Error::Auth(_))
));
}
@ -812,61 +749,4 @@ mod tests {
assert!(error.to_string().contains(COHERE_API_KEY_ENV), "{error}");
}
#[rstest]
#[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")]
#[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")]
#[tokio::test]
async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key(
#[case] model: &str,
#[case] request_line: &str,
) {
use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr};
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let request = crate::ocr::test_support::wire_request(model, &base, json!({}))
.with_document(
serde_json::from_value::<OcrDocument>(
json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}),
)
.unwrap()
.into(),
);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with(request_line), "{}", requests[0]);
assert_eq!(
header(&requests[0], "authorization"),
Some("Bearer test-key")
);
}
#[rstest]
#[tokio::test]
async fn route_rejects_non_image_document_without_a_request(
#[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str,
) {
use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr};
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let error = perform_ocr(crate::ocr::test_support::wire_request(
model,
&base,
json!({}),
))
.await
.unwrap_err();
server.abort();
assert!(
matches!(error, crate::ocr::Error::CohereImageOnly),
"{error:?}"
);
assert!(seen.lock().unwrap().is_empty());
}
}

12
litellm-rust/crates/core/src/http_utils.rs → litellm-rust/crates/llms/src/custom_httpx/http_handler.rs
View file

@ -9,13 +9,15 @@ pub struct HeaderError {
use litellm_core_utils::core_helpers::json_type_name;
use serde_json::{Map, Value};
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
/// Max characters of an upstream error body echoed across the call boundary
/// before truncation, so provider bodies are bounded and data-minimized.
const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256;
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) enum HeaderPolicy<'a> {
pub enum HeaderPolicy<'a> {
All,
Only(&'a [&'a str]),
Except(&'a [&'a str]),
@ -25,7 +27,7 @@ pub(crate) enum HeaderPolicy<'a> {
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) fn with_headers(
pub fn with_headers(
builder: reqwest::RequestBuilder,
headers: &[(String, String)],
policy: HeaderPolicy<'_>,
@ -109,9 +111,7 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) fn deserialize_optional_param<'de, D, T>(
deserializer: D,
) -> Result<Option<Option<T>>, D::Error>
pub fn deserialize_optional_param<'de, D, T>(deserializer: D) -> Result<Option<Option<T>>, D::Error>
where
D: serde::Deserializer<'de>,
T: serde::Deserialize<'de>,

343
litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs Normal file
View file

@ -0,0 +1,343 @@
use std::{sync::OnceLock, time::Duration};
use bytes::{Bytes, BytesMut};
use futures_util::future::BoxFuture;
use litellm_auth_gcp::VertexAuth;
use litellm_callbacks::event::{Passthrough, WireRequest};
use serde::{Serialize, de::DeserializeOwned};
use serde_json::{Map, Value};
use crate::{
base_llm::ocr::{
error::Error,
transformation::{
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_CONNECT_TIMEOUT_SECS,
OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value,
decode_response,
},
},
custom_httpx::{
http_handler::{HeaderPolicy, execute_http_request, with_headers},
media::MediaFetcher,
transport,
},
};
/// The route's view of one call, handed to provider code that has to reach the
/// caller's hooks mid-flight (guardrails on the outgoing body, raw response events).
pub trait CallHooks<E>: Send + Sync {
fn before_send(
&self,
wire: WireRequest,
passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, E>>;
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), E>>;
}
#[derive(Clone)]
pub struct OcrClient {
provider_http: reqwest::Client,
polling_http: reqwest::Client,
document_fetcher: MediaFetcher,
vertex_auth: VertexAuth,
}
impl OcrClient {
pub fn new(provider_http: reqwest::Client) -> Result<Self, transport::Error> {
let document_fetcher = MediaFetcher::new().map_err(transport::Error::from)?;
Ok(Self {
provider_http,
polling_http: no_redirect_http()?,
document_fetcher,
vertex_auth: VertexAuth::default(),
})
}
pub fn shared() -> Result<Self, Error> {
static CLIENT: OnceLock<Result<OcrClient, 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(transport::Error::from)
.and_then(OcrClient::new)
})
.clone()?;
Ok(client)
}
pub fn provider_http(&self) -> &reqwest::Client {
&self.provider_http
}
pub fn polling_http(&self) -> &reqwest::Client {
&self.polling_http
}
pub fn document_fetcher(&self) -> &MediaFetcher {
&self.document_fetcher
}
pub fn vertex_auth(&self) -> &VertexAuth {
&self.vertex_auth
}
#[cfg(any(test, feature = "test-support"))]
pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
Self {
provider_http,
polling_http: no_redirect_http().expect("test polling client builds"),
document_fetcher: MediaFetcher::for_test(document_http),
vertex_auth: VertexAuth::default(),
}
}
}
fn no_redirect_http() -> Result<reqwest::Client, transport::Error> {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(transport::Error::from)
}
/// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request,
/// send it, and hand the response to the config for normalization.
pub async fn ocr<C: BaseOcrConfig>(
config: &C,
client: &OcrClient,
request: &PreparedOcrRequest,
hooks: &dyn CallHooks<Error>,
) -> Result<LiteLLMOcrResponse, Error> {
let http = config.prepare_request(request, client, hooks).await?;
let url = http.url().to_string();
let headers = request_headers(&http)?;
let response = execute_http_request(client.provider_http(), http)
.await
.map_err(transport_error)?;
if !response.status().is_success() {
let headers = response
.headers()
.iter()
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.to_string(), value.to_string()))
})
.collect();
return match read_response_bytes(response, request.connection.max_response_bytes).await {
Err(Error::Transport(transport::Error::Http { status, body })) => {
Err(config.get_error_class(body, status, headers))
}
Err(error) => Err(error),
Ok(_) => unreachable!("non-success response produces an HTTP error"),
};
}
let context = OcrResponseContext {
client,
connection: &request.connection,
hooks,
request_format: request.response_format()?,
url: &url,
headers: &headers,
};
config
.async_transform_ocr_response(&request.model, response, context)
.await
}
fn request_headers(request: &reqwest::Request) -> Result<Vec<(String, String)>, Error> {
request
.headers()
.iter()
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.to_string(), value.to_string()))
.map_err(|_| Error::RequestField {
path: "headers".into(),
})
})
.collect()
}
pub async fn read_json_response<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
max_response_bytes: usize,
) -> Result<DecodedOcrResponse<T>, Error> {
let bytes = read_response_bytes(response, max_response_bytes).await?;
decode_response(&bytes, native)
}
pub async fn read_response_bytes(
mut response: reqwest::Response,
limit: usize,
) -> Result<Bytes, Error> {
let status = response.status();
if status.is_success()
&& response
.content_length()
.is_some_and(|length| length > limit as u64)
{
return Err(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(Error::TooLarge { limit });
}
bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
if !status.is_success() && bytes.len() == limit {
break;
}
}
if !status.is_success() {
return Err(transport::Error::Http {
status: status.as_u16(),
body: String::from_utf8_lossy(&bytes).into_owned(),
}
.into());
}
Ok(bytes.freeze())
}
pub fn transport_error(error: reqwest::Error) -> Error {
if error.is_timeout() {
return Error::Transport(transport::Error::Http {
status: 408,
body: "OCR request timed out".into(),
});
}
transport::Error::from(error).into()
}
pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
config: &C,
client: &OcrClient,
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: B,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
let composed = litellm_core_utils::call_arguments::compose_body(
&request.optional_params,
&body,
config.get_supported_ocr_params(&request.model),
)?;
config.validate_request_body(&composed)?;
let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed);
let changed = hooks
.before_send(wire_request(url, headers, composed), passthrough_fields)
.await?;
if !changed.body.is_object() {
return Err(Error::RequestField {
path: "guardrail.body".into(),
});
}
config.validate_request_body(&changed.body)?;
build_http_request(client, request, url, &changed.headers, &changed.body)
}
fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest {
WireRequest {
url: url.into(),
headers: headers.to_vec(),
body,
}
}
fn caller_inputs(request: &PreparedOcrRequest) -> Result<Map<String, Value>, Error> {
let document = request
.caller_document
.then(|| serde_json::to_value(&request.document))
.transpose()
.map_err(|_| Error::RequestField {
path: "document".into(),
})?;
let params: Map<String, Value> = request.optional_params.clone().into();
Ok(params
.into_iter()
.chain(document.map(|document| ("document".to_string(), document)))
.collect())
}
pub fn build_http_request<B: Serialize>(
client: &OcrClient,
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: &B,
) -> Result<reqwest::Request, Error> {
let builder = client
.provider_http()
.post(url)
.json(body)
.timeout(request.connection.timeout);
with_headers(builder, headers, HeaderPolicy::All)
.build()
.map_err(transport::Error::from)
.map_err(Error::from)
}
pub async fn guardrail_document(
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
hooks: &dyn CallHooks<Error>,
) -> Result<(OcrDocument, Vec<(String, String)>), Error> {
let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField {
path: "document".into(),
})?;
let changed = hooks
.before_send(wire_request(url, headers, body), Passthrough::default())
.await?;
let document = decode_request_value(changed.body, "guardrail.document")?;
Ok((document, changed.headers))
}
pub fn body_document(body: &Value) -> Result<OcrDocument, Error> {
let document = body
.get("document")
.and_then(Value::as_object)
.ok_or_else(|| Error::RequestField {
path: "body.document".into(),
})?;
let source = document
.iter()
.filter(|(name, _)| matches!(name.as_str(), "type" | "image_url" | "document_url"))
.map(|(name, value)| (name.clone(), value.clone()))
.collect();
decode_request_value(Value::Object(source), "body.document")
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_timeout_has_an_http_408_status() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let _connection = listener.accept().await.unwrap();
tokio::time::sleep(Duration::from_secs(1)).await;
});
let error = reqwest::Client::new()
.get(format!("http://{address}"))
.timeout(Duration::from_millis(10))
.send()
.await
.unwrap_err();
assert!(matches!(
transport_error(error),
Error::Transport(transport::Error::Http { status: 408, .. })
));
server.abort();
}
}

44
litellm-rust/crates/core/src/media.rs → litellm-rust/crates/llms/src/custom_httpx/media.rs
View file

@ -12,10 +12,10 @@ use reqwest::{
dns::{Addrs, Name, Resolve, Resolving},
};
use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS;
const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
#[derive(Debug, thiserror::Error)]
pub(crate) enum Error {
pub enum Error {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
@ -33,11 +33,11 @@ pub(crate) enum Error {
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] crate::transport::Error),
Transport(#[from] crate::custom_httpx::transport::Error),
}
#[derive(Clone)]
pub(crate) struct MediaFetcher {
pub struct MediaFetcher {
client: reqwest::Client,
address_resolver: Arc<dyn AddressResolver>,
allow_private_network: bool,
@ -50,20 +50,20 @@ trait AddressResolver: Send + Sync {
}
#[derive(Clone, Copy)]
pub(crate) struct DownloadPolicy {
pub(crate) timeout: Duration,
pub(crate) max_bytes: u64,
pub(crate) max_redirects: usize,
pub struct DownloadPolicy {
pub timeout: Duration,
pub max_bytes: u64,
pub max_redirects: usize,
}
#[derive(Debug)]
pub(crate) struct DownloadedMedia {
pub(crate) bytes: Vec<u8>,
pub(crate) content_type: String,
pub struct DownloadedMedia {
pub bytes: Vec<u8>,
pub content_type: String,
}
impl MediaFetcher {
pub(crate) fn new() -> Result<Self, reqwest::Error> {
pub fn new() -> Result<Self, reqwest::Error> {
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
}
@ -87,8 +87,8 @@ impl MediaFetcher {
})
}
#[cfg(test)]
pub(crate) fn for_test(client: reqwest::Client) -> Self {
#[cfg(any(test, feature = "test-support"))]
pub fn for_test(client: reqwest::Client) -> Self {
Self {
client,
address_resolver: Arc::new(AllowPrivateResolver),
@ -96,11 +96,7 @@ impl MediaFetcher {
}
}
pub(crate) async fn fetch(
&self,
url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, Error> {
pub async fn fetch(&self, url: Url, policy: DownloadPolicy) -> Result<DownloadedMedia, Error> {
if policy.max_bytes == 0 {
return Err(Error::DownloadDisabled);
}
@ -122,7 +118,7 @@ impl MediaFetcher {
.get(url.clone())
.send()
.await
.map_err(crate::transport::Error::from)?;
.map_err(crate::custom_httpx::transport::Error::from)?;
if response.status().is_redirection() {
if redirects_followed == policy.max_redirects {
return Err(Error::TooManyRedirects);
@ -153,7 +149,7 @@ impl MediaFetcher {
while let Some(chunk) = response
.chunk()
.await
.map_err(crate::transport::Error::from)?
.map_err(crate::custom_httpx::transport::Error::from)?
{
enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?;
bytes.extend_from_slice(&chunk);
@ -184,7 +180,7 @@ impl MediaFetcher {
.address_resolver
.resolve(host, port)
.await
.map_err(|error| crate::transport::Error::Network(error.to_string()))?;
.map_err(|error| crate::custom_httpx::transport::Error::Network(error.to_string()))?;
validate_addresses(&addresses)
}
}
@ -254,10 +250,10 @@ impl AddressResolver for SystemAddressResolver {
}
}
#[cfg(test)]
#[cfg(any(test, feature = "test-support"))]
struct AllowPrivateResolver;
#[cfg(test)]
#[cfg(any(test, feature = "test-support"))]
impl AddressResolver for AllowPrivateResolver {
fn resolve<'a>(&'a self, _host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) })

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

@ -0,0 +1,4 @@
pub mod http_handler;
pub mod llm_http_handler;
pub mod media;
pub mod transport;

11
litellm-rust/crates/core/src/transport/error.rs → litellm-rust/crates/llms/src/custom_httpx/transport.rs
View file

@ -38,8 +38,11 @@ mod tests {
.send()
.await
.expect_err("invalid port");
let error = crate::transport::Error::from_reqwest_before_dispatch(error);
assert!(matches!(error, crate::transport::Error::Connect(_)));
let error = crate::custom_httpx::transport::Error::from_reqwest_before_dispatch(error);
assert!(matches!(
error,
crate::custom_httpx::transport::Error::Connect(_)
));
assert!(!error.to_string().contains("secret"));
assert!(!error.to_string().contains("private"));
}
@ -68,8 +71,8 @@ mod tests {
let error = response.expect_err("server does not respond");
assert!(error.is_timeout());
assert!(matches!(
crate::transport::Error::from_reqwest_before_dispatch(error),
crate::transport::Error::Network(_)
crate::custom_httpx::transport::Error::from_reqwest_before_dispatch(error),
crate::custom_httpx::transport::Error::Network(_)
));
}
}

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

@ -2,4 +2,9 @@ pub mod anthropic;
pub mod azure_ai;
pub mod base_llm;
pub mod bedrock;
pub mod cohere;
pub mod custom_httpx;
pub mod mistral;
pub mod openai;
pub mod reducto;
pub mod vertex_ai;

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

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

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

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

77
litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs → litellm-rust/crates/llms/src/mistral/ocr/transformation.rs
View file

@ -3,22 +3,23 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
constants::MISTRAL_OCR_API_BASE,
llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response},
ocr::{
OcrClient,
prepare::credential_env,
types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat,
OcrUsageInfo, PreparedOcrRequest,
base_llm::ocr::{
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage,
OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env,
decode_and_normalize_response,
},
},
custom_httpx::llm_http_handler::OcrClient,
};
const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";
const MISTRAL_OCR_API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY";
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct MistralOcrRequest {
pub struct MistralOcrRequest {
pub model: String,
pub document: OcrDocument,
#[serde(flatten)]
@ -26,7 +27,7 @@ pub(crate) struct MistralOcrRequest {
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct MistralOcrResponse {
pub struct MistralOcrResponse {
#[serde(default)]
pub pages: Vec<OcrPage>,
#[serde(
@ -42,7 +43,7 @@ pub(crate) struct MistralOcrResponse {
}
#[derive(Clone, Debug, Default)]
pub(crate) struct MistralOcrConfig;
pub struct MistralOcrConfig;
impl BaseOcrConfig for MistralOcrConfig {
type OcrParams = OpaqueParams;
@ -75,7 +76,7 @@ impl BaseOcrConfig for MistralOcrConfig {
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<OpaqueParams, crate::ocr::Error> {
) -> Result<OpaqueParams, Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
@ -85,7 +86,7 @@ impl BaseOcrConfig for MistralOcrConfig {
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
self.resolve_headers(&request.connection, &credential_env)
}
@ -94,7 +95,7 @@ impl BaseOcrConfig for MistralOcrConfig {
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
self.build_ocr_url(request.connection.api_base.as_deref())
}
@ -104,7 +105,7 @@ impl BaseOcrConfig for MistralOcrConfig {
document: OcrDocument,
optional_params: &OpaqueParams,
_headers: &[(String, String)],
) -> Result<MistralOcrRequest, crate::ocr::Error> {
) -> Result<MistralOcrRequest, Error> {
Ok(MistralOcrRequest {
model: model.to_string(),
document,
@ -117,7 +118,7 @@ impl BaseOcrConfig for MistralOcrConfig {
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
}
@ -127,8 +128,9 @@ impl MistralOcrConfig {
&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") {
) -> Result<Vec<(String, String)>, Error> {
if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization")
{
return Ok(connection.extra_headers.clone());
}
let api_key = connection
@ -153,7 +155,7 @@ impl MistralOcrConfig {
)
}
fn build_ocr_url(&self, api_base: Option<&str>) -> Result<String, crate::ocr::Error> {
fn build_ocr_url(&self, api_base: Option<&str>) -> Result<String, Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
@ -161,20 +163,20 @@ impl MistralOcrConfig {
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
}
}
pub(crate) fn normalize_response(
pub fn normalize_response(
model: &str,
response: MistralOcrResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let model = match response.model {
Some(Some(model)) => model,
Some(None) => {
return Err(crate::ocr::Error::ResponseField {
return Err(Error::ResponseField {
path: "model".into(),
});
}
@ -194,6 +196,7 @@ mod tests {
use serde_json::{Value, json};
use super::*;
use crate::base_llm::ocr::transformation::decode_response;
#[fixture]
fn document() -> OcrDocument {
@ -220,7 +223,7 @@ mod tests {
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"
Error::ResponseField { path } if path == "model"
));
}
@ -247,14 +250,12 @@ mod tests {
#[case] payload: Value,
#[case] path: &str,
) {
let error = crate::ocr::json::decode_response::<MistralOcrResponse>(
&serde_json::to_vec(&payload).unwrap(),
false,
)
.unwrap_err();
let error =
decode_response::<MistralOcrResponse>(&serde_json::to_vec(&payload).unwrap(), false)
.unwrap_err();
assert!(matches!(
error,
crate::ocr::Error::ResponseField { path: actual } if actual == path
Error::ResponseField { path: actual } if actual == path
));
}
@ -316,7 +317,11 @@ mod tests {
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)
.transform_ocr_response(
"model",
raw,
crate::base_llm::ocr::transformation::OcrResponseFormat::Native,
)
.unwrap();
assert_eq!(response.pages[0].index, 2);
let native = response.provider_native_response.unwrap();
@ -640,12 +645,10 @@ mod tests {
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,
}
))
Err(Error::Auth(litellm_auth::Error::MissingApiKey {
provider: "Mistral",
environment_variable: MISTRAL_OCR_API_KEY_ENV_VAR,
}))
));
}
}

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

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

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

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

435
litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs → litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

@ -9,44 +9,47 @@ use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value, json};
use crate::{
constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX},
llms::base_llm::ocr::transformation::{
BaseOcrConfig, OcrRequestContext, decode_and_normalize_response,
},
ocr::{
OcrClient,
base_llm::ocr::{
document::InlineDocument,
prepare::{build_http_request, credential_env, guardrail_document},
types::{
LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat,
OcrUsageInfo, PreparedOcrRequest,
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument,
OcrPage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
credential_env, decode_and_normalize_response,
},
},
custom_httpx::llm_http_handler::{
CallHooks, OcrClient, build_http_request, guardrail_document,
},
};
const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
const REDUCTO_ID_PREFIX: &str = "reducto://";
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(transparent)]
pub(crate) struct ReductoFileId(String);
pub struct ReductoFileId(String);
pub(crate) type ReductoV3Params = OpaqueParams;
pub(crate) type ReductoLegacyParams = OpaqueParams;
pub type ReductoV3Params = OpaqueParams;
pub type ReductoLegacyParams = OpaqueParams;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoV3Request {
pub struct ReductoV3Request {
pub input: ReductoFileId,
#[serde(flatten)]
pub params: ReductoV3Params,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyRequest {
pub struct ReductoLegacyRequest {
pub document_url: ReductoFileId,
#[serde(skip_serializing_if = "Option::is_none")]
pub options: Option<ReductoLegacyOptions>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyOptions {
pub struct ReductoLegacyOptions {
pub enhance: Value,
}
@ -56,7 +59,7 @@ struct ReductoUploadResponse {
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct ReductoResponse {
pub struct ReductoResponse {
#[serde(default, deserialize_with = "present_nullable")]
result: Option<Option<ReductoResult>>,
usage: Option<ReductoUsage>,
@ -85,7 +88,7 @@ struct ReductoChunk {
}
#[derive(Clone, Debug)]
pub(crate) struct ReductoParseV3Config;
pub struct ReductoParseV3Config;
impl BaseOcrConfig for ReductoParseV3Config {
type OcrParams = ReductoV3Params;
@ -100,7 +103,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<ReductoV3Params, crate::ocr::Error> {
) -> Result<ReductoV3Params, Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
@ -110,7 +113,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
resolve_headers(&request.connection, &credential_env)
}
@ -119,7 +122,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
request: &PreparedOcrRequest,
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
build_ocr_url(request.connection.api_base.as_deref())
}
@ -129,7 +132,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
document: OcrDocument,
optional_params: &Self::OcrParams,
_headers: &[(String, String)],
) -> Result<Self::ProviderRequest, crate::ocr::Error> {
) -> Result<Self::ProviderRequest, Error> {
Ok(ReductoV3Request {
input: uploaded_file_id(document)?,
params: optional_params.clone(),
@ -143,7 +146,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
optional_params: &ReductoV3Params,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoV3Request, crate::ocr::Error> {
) -> Result<ReductoV3Request, Error> {
let file_id = ensure_file_id_async(document, headers, context).await?;
Ok(ReductoV3Request {
input: file_id,
@ -156,7 +159,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
@ -164,13 +167,14 @@ impl BaseOcrConfig for ReductoParseV3Config {
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, crate::ocr::Error> {
prepare_upload_request(self, request, client).await
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
prepare_upload_request(self, request, client, hooks).await
}
}
#[derive(Clone, Debug)]
pub(crate) struct ReductoParseLegacyConfig;
pub struct ReductoParseLegacyConfig;
impl BaseOcrConfig for ReductoParseLegacyConfig {
type OcrParams = ReductoLegacyParams;
@ -185,7 +189,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
&self,
non_default_params: &CallArguments,
model: &str,
) -> Result<ReductoLegacyParams, crate::ocr::Error> {
) -> Result<ReductoLegacyParams, Error> {
Ok(non_default_params
.select(self.get_supported_ocr_params(model))
.into())
@ -195,7 +199,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
ReductoParseV3Config
.validate_environment(request, client)
.await
@ -206,7 +210,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
request: &PreparedOcrRequest,
optional_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
ReductoParseV3Config.get_complete_url(request, optional_params, environment)
}
@ -216,7 +220,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
document: OcrDocument,
optional_params: &Self::OcrParams,
_headers: &[(String, String)],
) -> Result<Self::ProviderRequest, crate::ocr::Error> {
) -> Result<Self::ProviderRequest, Error> {
Ok(build_legacy_body(
uploaded_file_id(document)?,
optional_params,
@ -230,7 +234,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
optional_params: &ReductoLegacyParams,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoLegacyRequest, crate::ocr::Error> {
) -> Result<ReductoLegacyRequest, Error> {
let file_id = ensure_file_id_async(document, headers, context).await?;
Ok(build_legacy_body(file_id, optional_params))
}
@ -240,7 +244,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
ReductoParseV3Config.transform_ocr_response(model, raw_response, request_format)
}
@ -248,8 +252,9 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, crate::ocr::Error> {
prepare_upload_request(self, request, client).await
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
prepare_upload_request(self, request, client, hooks).await
}
}
@ -260,11 +265,12 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
config: &C,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, crate::ocr::Error> {
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
let params = config.map_ocr_params(&request.optional_params, &request.model)?;
let headers = config.validate_environment(request, client).await?;
let url = config.get_complete_url(request, &params, &headers)?;
let (document, headers) = guardrail_document(request, &url, &headers).await?;
let (document, headers) = guardrail_document(request, &url, &headers, hooks).await?;
let body = config
.async_transform_ocr_request(
&request.model,
@ -285,15 +291,15 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
build_http_request(client, request, &url, &headers, &body)
}
fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, crate::ocr::Error> {
fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, Error> {
if !document.source().starts_with(REDUCTO_ID_PREFIX) {
return Err(crate::ocr::Error::ReductoSource);
return Err(Error::ReductoSource);
}
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(crate::ocr::Error::RequestField {
return Err(Error::RequestField {
path: "document file id".into(),
});
}
@ -322,10 +328,10 @@ fn checked_truncated_i64(value: f64) -> Option<i64> {
.then(|| value.trunc() as i64)
}
pub(crate) fn normalize_response(
pub fn normalize_response(
model: &str,
response: ReductoResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let result = match response.result {
Some(result) => result.unwrap_or_default(),
None => ReductoResult {
@ -346,7 +352,7 @@ pub(crate) fn normalize_response(
})
}
fn build_pages_from_reducto(chunks: Vec<ReductoChunk>) -> Result<Vec<OcrPage>, crate::ocr::Error> {
fn build_pages_from_reducto(chunks: Vec<ReductoChunk>) -> Result<Vec<OcrPage>, Error> {
let blocks_by_page = chunks
.iter()
.flat_map(|chunk| chunk.blocks.iter().flatten())
@ -376,7 +382,7 @@ fn build_pages_from_reducto(chunks: Vec<ReductoChunk>) -> Result<Vec<OcrPage>, c
.map(|block| match block.get("content") {
None | Some(Value::Null) => Ok(None),
Some(Value::String(content)) => Ok(Some(content.as_str())),
Some(_) => Err(crate::ocr::Error::ResponseField {
Some(_) => Err(Error::ResponseField {
path: "result.chunks.blocks.content".into(),
}),
})
@ -410,11 +416,11 @@ fn page(index: i64, markdown: String, blocks: Option<Value>) -> OcrPage {
..Default::default()
}
}
fn build_ocr_url(api_base: Option<&str>) -> Result<String, crate::ocr::Error> {
fn build_ocr_url(api_base: Option<&str>) -> Result<String, Error> {
complete_endpoint_url(api_base, "parse")
}
fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result<String, crate::ocr::Error> {
fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result<String, Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
@ -422,7 +428,7 @@ fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result<String, c
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&[path]))
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
}
@ -430,8 +436,8 @@ fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result<String, c
fn resolve_headers(
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") {
) -> Result<Vec<(String, String)>, Error> {
if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
@ -445,7 +451,7 @@ fn resolve_headers(
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or(crate::ocr::Error::MissingReductoApiKey)?;
.ok_or(Error::MissingReductoApiKey)?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
@ -472,22 +478,21 @@ async fn ensure_file_id_async(
document: OcrDocument,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoFileId, crate::ocr::Error> {
) -> Result<ReductoFileId, Error> {
if document.source().starts_with(REDUCTO_ID_PREFIX) {
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(crate::ocr::Error::RequestField {
return Err(Error::RequestField {
path: "document file id".into(),
});
}
return Ok(ReductoFileId(document.source().to_string()));
}
let inline =
InlineDocument::parse(document.source())?.ok_or(crate::ocr::Error::ReductoSource)?;
let inline = InlineDocument::parse(document.source())?.ok_or(Error::ReductoSource)?;
let mime = inline.mime_type().to_string();
let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
let bytes = inline.decode(OCR_INLINE_MAX_BYTES)?;
upload_bytes_async(bytes, &mime, headers, context).await
}
@ -496,12 +501,12 @@ async fn upload_bytes_async(
mime: &str,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<ReductoFileId, crate::ocr::Error> {
) -> Result<ReductoFileId, Error> {
let OcrRequestContext { client, connection } = context;
let part = reqwest::multipart::Part::bytes(bytes)
.file_name("document")
.mime_str(mime)
.map_err(|_| crate::ocr::Error::InvalidDataUri)?;
.map_err(|_| Error::InvalidDataUri)?;
let builder = client
.provider_http()
.post(complete_endpoint_url(
@ -510,28 +515,32 @@ async fn upload_bytes_async(
)?)
.multipart(reqwest::multipart::Form::new().part("file", part))
.timeout(connection.timeout);
let builder = crate::http_utils::with_headers(
let builder = crate::custom_httpx::http_handler::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]),
crate::custom_httpx::http_handler::HeaderPolicy::Except(&[
"content-type",
"content-length",
]),
);
let response = crate::http_utils::http_request(builder)
let response = crate::custom_httpx::http_handler::http_request(builder)
.await
.map_err(crate::transport::Error::from)?;
let uploaded = crate::ocr::client::read_json_response::<ReductoUploadResponse>(
response,
false,
connection.max_response_bytes,
)
.await?
.data;
.map_err(crate::custom_httpx::transport::Error::from)?;
let uploaded =
crate::custom_httpx::llm_http_handler::read_json_response::<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(crate::ocr::Error::ResponseField {
return Err(Error::ResponseField {
path: "file_id".into(),
});
};
@ -592,48 +601,6 @@ mod tests {
assert_eq!(normalized.usage_info.unwrap().credits, Some(1.0));
}
#[tokio::test]
async fn v3_options_preserve_explicit_null() {
let overrides =
serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true}))
.unwrap();
let params = ReductoParseV3Config
.map_ocr_params(&overrides, "parse-v3")
.unwrap();
let client = crate::ocr::test_support::ocr_client();
let connection = OcrConnection::default();
let document = serde_json::from_value(
json!({"type":"document_url","document_url":"reducto://ready.pdf"}),
)
.unwrap();
let body = ReductoParseV3Config
.async_transform_ocr_request(
"parse-v3",
document,
&params,
&[],
OcrRequestContext {
client: &client,
connection: &connection,
},
)
.await
.unwrap();
assert_eq!(
serde_json::to_value(body).unwrap(),
json!({
"input":"reducto://ready.pdf", "formatting":null, "settings":{}
})
);
let absent = ReductoParseV3Config
.map_ocr_params(
&litellm_core_utils::call_arguments::CallArguments::default(),
"parse-v3",
)
.unwrap();
assert_eq!(serde_json::to_value(absent).unwrap(), json!({}));
}
#[test]
fn legacy_body_omits_null_enhance_and_wraps_mapped_options() {
for (value, expected) in [
@ -691,178 +658,9 @@ mod tests {
);
}
use litellm_callbacks::event::{CallEvent, WireRequest};
use rstest::rstest;
use crate::ocr::{
LocalOcrHost,
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
};
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[rstest]
#[case(
"reducto/parse-v3",
json!({
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://already.pdf",
json!({
"input":"reducto://already.pdf",
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[case(
"reducto/parse-legacy",
json!({
"enhance":{"agentic":[{"type":"table"}]},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://legacy.pdf",
json!({
"document_url":"reducto://legacy.pdf",
"options":{"enhance":{"agentic":[{"type":"table"}]}},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[tokio::test]
async fn request_mapping_matches_python(
#[case] model: &str,
#[case] options: Value,
#[case] source: &str,
#[case] expected: Value,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"result":{"chunks":[]}
}))])
.await;
let request =
crate::ocr::test_support::with_source(wire_request(model, &base, options), source);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert_eq!(request_body(&requests[0]), expected);
}
#[rstest]
#[case("parse-v3")]
#[case("parse-legacy")]
#[tokio::test]
async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})),
])
.await;
let mut request = wire_request(&format!("reducto/{model}"), &base, json!({}));
request.transport.extra_headers = vec![
("Content-Type".into(), "application/json".into()),
("X-Trace".into(), "upload-test".into()),
];
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(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("content-type: multipart/form-data; boundary=")
);
assert!(requests[0].contains("x-trace: upload-test"));
assert!(requests[0].contains("application/pdf"));
assert!(requests[0].contains("abc"));
assert!(requests[1].starts_with("POST /parse "));
}
#[tokio::test]
async fn response_received_stays_after_reducto_upload_and_parse() {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[]}})),
])
.await;
let request_count = seen.clone();
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
assert_eq!(request_count.lock().unwrap().len(), 2);
assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#);
}
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 2);
}
#[rstest]
#[case(json!({"file_id":""}))]
#[case(json!({}))]
#[case(json!({"file_id":null}))]
#[tokio::test]
async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) {
let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await;
let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({})))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("file_id"));
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn upload_failure_stops_before_parse() {
let (base, seen, server) = mock_server(vec![MockResponse {
status: 503,
headers: vec![],
body: json!({"error":"unavailable"}),
}])
.await;
assert!(
perform_ocr(wire_request("reducto/parse-v3", &base, json!({})))
.await
.is_err()
);
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[rstest]
#[case("https://example.com/a.pdf")]
#[case("reducto://")]
#[case("data:application/pdf;base64")]
#[case("data:application/pdf;base64,INVALID!")]
#[tokio::test]
async fn rejects_invalid_document_sources_before_network(#[case] source: &str) {
let request = crate::ocr::test_support::with_source(
wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})),
source,
);
assert!(perform_ocr(request).await.is_err());
}
#[test]
fn response_normalization_groups_blocks_and_distinguishes_null_result() {
use crate::llms::reducto::ocr::transformation::{ReductoResponse, normalize_response};
use crate::reducto::ocr::transformation::{ReductoResponse, normalize_response};
let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[
{"blocks":[{
@ -906,79 +704,4 @@ mod tests {
let null = normalize_response("parse-v3", null).unwrap();
assert!(null.pages.is_empty());
}
#[tokio::test]
async fn facade_omits_native_response_by_default_and_preserves_auth_priority() {
let raw = json!({"job_id":"job-1","result":{"chunks":[]}});
let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await;
let mut request = crate::ocr::test_support::with_source(
wire_request("reducto/parse-v3", &base, json!({})),
"reducto://ready.pdf",
);
request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())];
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.provider_native_response, None);
assert!(
seen.lock().unwrap()[0]
.to_ascii_lowercase()
.contains("authorization: bearer existing")
);
}
#[rstest]
#[case("reducto/parse-v3")]
#[case("reducto/parse-legacy")]
#[tokio::test]
async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[]}})),
])
.await;
let mut request = wire_request(model, &base, json!({}));
request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())];
let host = LocalOcrHost::new(request).with_before_send(|wire, _| {
Ok(WireRequest {
headers: vec![("authorization".into(), "Bearer guarded".into())],
..wire
})
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(requests[1].starts_with("POST /parse "));
for request in requests.iter() {
assert!(request.contains("authorization: Bearer guarded"));
assert!(!request.contains("Bearer original"));
}
}
#[tokio::test]
async fn guardrail_rewrites_document_before_upload() {
let (base, seen, server) =
mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await;
let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({})))
.with_before_send(|wire, _| {
assert_eq!(
wire.body["document_url"],
"data:application/pdf;base64,YWJj"
);
Ok(WireRequest {
body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
..wire
})
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert!(requests[0].contains("reducto://guarded.pdf"));
}
}

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

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

4
litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs → litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs
View file

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

184
litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs → litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs
View file

@ -5,15 +5,15 @@ use serde_json::{Map, Value};
use super::transformation::VertexAiOcrConfig;
use crate::{
llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext},
ocr::{
OcrClient,
prepare::credential_env,
types::{
LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage,
OcrUsageInfo, PreparedOcrRequest,
base_llm::ocr::{
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions,
OcrPageImage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
credential_env, decode_and_normalize_response, decode_response_value,
},
},
custom_httpx::llm_http_handler::OcrClient,
};
const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
@ -21,10 +21,10 @@ 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;
pub type DeepSeekOcrParams = OpaqueParams;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrRequest {
pub struct DeepSeekOcrRequest {
pub model: String,
pub messages: Vec<DeepSeekOcrMessage>,
#[serde(flatten)]
@ -32,26 +32,26 @@ pub(crate) struct DeepSeekOcrRequest {
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrMessage {
pub struct DeepSeekOcrMessage {
pub role: UserRole,
pub content: Vec<DeepSeekDocument>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
pub(crate) enum DeepSeekDocument {
pub enum DeepSeekDocument {
#[serde(rename = "image_url")]
ImageUrl { image_url: String },
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum UserRole {
pub enum UserRole {
User,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekOcrResponse {
pub struct DeepSeekOcrResponse {
#[serde(default)]
choices: Vec<DeepSeekChoice>,
#[serde(default = "empty_object")]
@ -89,7 +89,7 @@ struct DeepSeekPage {
}
#[derive(Clone, Debug)]
pub(crate) struct VertexAIDeepSeekOCRConfig;
pub struct VertexAIDeepSeekOCRConfig;
impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
type OcrParams = DeepSeekOcrParams;
@ -104,7 +104,7 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
&self,
_arguments: &CallArguments,
_model: &str,
) -> Result<DeepSeekOcrParams, crate::ocr::Error> {
) -> Result<DeepSeekOcrParams, Error> {
Ok(DeepSeekOcrParams::default())
}
@ -112,7 +112,7 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
) -> Result<Self::Environment, Error> {
VertexAiOcrConfig
.validate_environment(request, client)
.await
@ -123,7 +123,7 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
request: &PreparedOcrRequest,
_params: &Self::OcrParams,
environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
@ -144,7 +144,7 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
optional_params: &DeepSeekOcrParams,
headers: &[(String, String)],
_context: OcrRequestContext<'_>,
) -> Result<DeepSeekOcrRequest, crate::ocr::Error> {
) -> Result<DeepSeekOcrRequest, Error> {
self.transform_ocr_request(model, document, optional_params, headers)
}
@ -152,14 +152,9 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
&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,
)
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
fn transform_ocr_request(
@ -168,9 +163,9 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
document: OcrDocument,
optional_params: &DeepSeekOcrParams,
_headers: &[(String, String)],
) -> Result<DeepSeekOcrRequest, crate::ocr::Error> {
) -> Result<DeepSeekOcrRequest, Error> {
if document.source().is_empty() {
return Err(crate::ocr::Error::MissingDocumentUrl);
return Err(Error::MissingDocumentUrl);
}
Ok(DeepSeekOcrRequest {
model: provider_model(model)?,
@ -189,19 +184,19 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
}
}
pub(crate) fn normalize_response(
pub fn normalize_response(
model: &str,
response: DeepSeekOcrResponse,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
) -> Result<LiteLLMOcrResponse, Error> {
let content = response
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content)
.ok_or(crate::ocr::Error::EmptyContent)?;
.ok_or(Error::EmptyContent)?;
let (ocr_data, fallback_markdown) = match content {
DeepSeekContent::Text(text) if text.is_empty() => {
return Err(crate::ocr::Error::EmptyContent);
return Err(Error::EmptyContent);
}
DeepSeekContent::Text(text) => {
let parsed = text
@ -212,7 +207,7 @@ pub(crate) fn normalize_response(
(parsed.unwrap_or_default(), text)
}
DeepSeekContent::Object(data) if data.is_empty() => {
return Err(crate::ocr::Error::EmptyContent);
return Err(Error::EmptyContent);
}
DeepSeekContent::Object(data) => {
let fallback = if data.contains_key("pages") {
@ -236,7 +231,7 @@ pub(crate) fn normalize_response(
.enumerate()
.filter(|(_, page)| page.is_object())
.map(|(position, page)| {
let page: DeepSeekPage = crate::ocr::json::decode_response_value(
let page: DeepSeekPage = decode_response_value(
page.clone(),
&format!("choices[0].message.content.pages[{position}]"),
)?;
@ -248,7 +243,7 @@ pub(crate) fn normalize_response(
..Default::default()
})
})
.collect::<Result<Vec<_>, crate::ocr::Error>>()?,
.collect::<Result<Vec<_>, Error>>()?,
Some(_) => return Err(response_field("pages")),
None => Vec::new(),
};
@ -257,7 +252,7 @@ pub(crate) fn normalize_response(
.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"))
.map(|usage| decode_response_value(usage.clone(), "usage_info"))
.transpose()?;
let model = match ocr_data.get("model") {
Some(Value::String(model)) => model.clone(),
@ -357,16 +352,16 @@ impl serde_json::ser::Formatter for PythonJsonFormatter {
}
}
fn response_field(field: &str) -> crate::ocr::Error {
crate::ocr::Error::ResponseField {
fn response_field(field: &str) -> Error {
Error::ResponseField {
path: format!("choices[0].message.content.{field}"),
}
}
pub(crate) fn provider_model(model: &str) -> Result<String, crate::ocr::Error> {
pub fn provider_model(model: &str) -> Result<String, Error> {
let local_model = model.trim_start_matches(MODEL_PREFIX);
if local_model.is_empty() {
return Err(crate::ocr::Error::RequestField {
return Err(Error::RequestField {
path: "model".into(),
});
}
@ -379,7 +374,7 @@ impl VertexAIDeepSeekOCRConfig {
api_base: Option<&str>,
project: &str,
location: &str,
) -> Result<String, crate::ocr::Error> {
) -> Result<String, Error> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
@ -399,7 +394,7 @@ impl VertexAIDeepSeekOCRConfig {
])
})
.map(|url| url.into_string())
.map_err(|_| crate::ocr::Error::RequestField {
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
}
@ -407,18 +402,24 @@ impl VertexAIDeepSeekOCRConfig {
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::{
DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, normalize_response,
provider_model,
};
use crate::base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument};
fn document() -> OcrDocument {
serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap()
}
#[test]
fn unconsumed_options_remain_available_for_body_composition() {
use serde_json::json;
use crate::llms::base_llm::ocr::transformation::BaseOcrConfig;
use crate::base_llm::ocr::transformation::BaseOcrConfig;
let arguments =
serde_json::from_value(json!({"temperature":0.5,"extension":null})).unwrap();
@ -460,14 +461,6 @@ mod tests {
);
}
use rstest::rstest;
use crate::{llms::base_llm::ocr::transformation::BaseOcrConfig, 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))]
@ -614,93 +607,4 @@ mod tests {
);
}
}
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/llms/src/vertex_ai/ocr/mod.rs Normal file
View file

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

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

@ -0,0 +1,224 @@
use litellm_auth_gcp::{self as vertex, VertexConfig};
use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl};
use serde_json::Value;
use super::common_utils::validate_destination;
use crate::{
base_llm::ocr::{
document::{inline_remote_document, validate_inline_document},
error::Error,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrEnvironment,
OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, credential_env,
},
},
custom_httpx::llm_http_handler::OcrClient,
mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest},
};
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug, Default)]
pub 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, Error> {
MistralOcrConfig.map_ocr_params(non_default_params, model)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, 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, 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, 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, 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: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, Error> {
MistralOcrConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), Error> {
validate_inline_document(&crate::custom_httpx::llm_http_handler::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, Error> {
validate_destination(connection)?;
client
.vertex_auth()
.validate_environment(
connection.extra_headers.clone(),
connection.api_key.as_deref(),
config,
&credential_env,
)
.await
.map_err(Error::from)
}
fn build_ocr_url(
&self,
api_base: Option<&str>,
project: &str,
location: &str,
model: &str,
) -> Result<String, 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(|_| Error::RequestField {
path: "api_base".into(),
})
}
}
fn validate_location(location: &str) -> Result<(), 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(Error::RequestField {
path: "vertex_location".into(),
})
}
#[cfg(test)]
mod tests {
use super::VertexAiOcrConfig;
#[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()
);
}
}

1
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -20,6 +20,7 @@ bytes.workspace = true
litellm-auth.workspace = true
litellm-callbacks-legacy.workspace = true
litellm-core.workspace = true
litellm-llms.workspace = true
litellm-types.workspace = true
litellm-host-python.workspace = true
litellm-token-counter.workspace = true

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

@ -1,6 +1,6 @@
use litellm_core::{
Error, audio_transcription, chat_completions, messages, ocr, responses,
transport::Error as TransportError,
use litellm_core::{Error, audio_transcription, chat_completions, messages, responses};
use litellm_llms::{
base_llm::ocr::error::Error as OcrError, custom_httpx::transport::Error as TransportError,
};
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
@ -43,11 +43,11 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr {
error.is_request()
|| matches!(
error,
ocr::Error::Auth(_)
| ocr::Error::InvalidProvider(_)
| ocr::Error::InvalidRequest(_)
| ocr::Error::MissingField(_)
| ocr::Error::MissingDocumentUrl
OcrError::Auth(_)
| OcrError::InvalidProvider(_)
| OcrError::InvalidRequest(_)
| OcrError::MissingField(_)
| OcrError::MissingDocumentUrl
)
}
Error::Messages(error) => match error {

2
litellm-rust/crates/python-bridge/src/routes/ocr/document.rs
View file

@ -1,7 +1,7 @@
use std::path::PathBuf;
use bytes::Bytes;
use litellm_core::ocr::{OcrDocumentInput, OcrFileContent};
use litellm_core::ocr::types::{OcrDocumentInput, OcrFileContent};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
gc::{PyTraverseError, PyVisit},

9
litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs
View file

@ -1,4 +1,4 @@
use litellm_core::ocr::Error;
use litellm_llms::base_llm::ocr::error::Error;
use pyo3::{
exceptions::{PyFileNotFoundError, PyOSError},
prelude::*,
@ -15,9 +15,10 @@ pub(super) fn to_pyerr(error: Error) -> PyErr {
body,
headers,
} => upstream_error(py, status, body, headers)?,
Error::Transport(litellm_core::transport::Error::Http { status, body }) => {
upstream_error(py, status, body, Vec::new())?
}
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status,
body,
}) => upstream_error(py, status, body, Vec::new())?,
Error::RequestFormat => {
let error = core_error_to_pyerr(Error::RequestFormat.into());
error

11
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -1,6 +1,7 @@
use litellm_auth::ResolvedCredential;
use litellm_core::ocr::{LiteLLMOcrResponse, Ocr, OcrOp, OcrOpResult};
use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult};
use litellm_host_python::{RouteHost, missing_state, to_py};
use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse};
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
@ -41,7 +42,7 @@ impl OcrRouteHost {
}
}
fn read_document(&self, py: Python<'_>) -> PyResult<litellm_core::ocr::OcrFileContent> {
fn read_document(&self, py: Python<'_>) -> PyResult<litellm_core::ocr::types::OcrFileContent> {
self.handles()?
.reader
.as_ref()
@ -94,12 +95,12 @@ impl RouteHost for OcrRouteHost {
.map(Bound::unbind)
}
fn native_error(error: litellm_core::ocr::Error) -> PyErr {
fn native_error(error: Error) -> PyErr {
ocr_error_to_pyerr(error)
}
fn host_error(error: &PyErr) -> litellm_core::ocr::Error {
litellm_core::ocr::Error::InvalidRequest(error.to_string())
fn host_error(error: &PyErr) -> Error {
Error::InvalidRequest(error.to_string())
}
fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult<PyErr> {

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