mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
feat: port OCR providers to Rust gateway
This commit is contained in:
Ishaan Jaff
parent
d8ef1da49d
commit
48b5c73f9a
29 changed files with 3241 additions and 474 deletions
162
litellm-rust/Cargo.lock
generated
162
litellm-rust/Cargo.lock
generated
|
|
@ -206,12 +206,24 @@ dependencies = [
|
|||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.2.2"
|
||||
|
|
@ -221,6 +233,21 @@ dependencies = [
|
|||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-executor",
|
||||
"futures-io",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-channel"
|
||||
version = "0.3.32"
|
||||
|
|
@ -237,12 +264,34 @@ version = "0.3.32"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
|
||||
|
||||
[[package]]
|
||||
name = "futures-executor"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-io"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.32"
|
||||
|
|
@ -261,8 +310,10 @@ version = "0.3.32"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"futures-macro",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"memchr",
|
||||
|
|
@ -307,6 +358,31 @@ dependencies = [
|
|||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
|
||||
dependencies = [
|
||||
"atomic-waker",
|
||||
"bytes",
|
||||
"fnv",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"http",
|
||||
"indexmap",
|
||||
"slab",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashbrown"
|
||||
version = "0.17.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.5.0"
|
||||
|
|
@ -368,6 +444,7 @@ dependencies = [
|
|||
"bytes",
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"httparse",
|
||||
|
|
@ -521,6 +598,16 @@ dependencies = [
|
|||
"icu_properties",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indexmap"
|
||||
version = "2.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
|
||||
dependencies = [
|
||||
"equivalent",
|
||||
"hashbrown",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indoc"
|
||||
version = "2.0.7"
|
||||
|
|
@ -544,9 +631,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
|
|||
|
||||
[[package]]
|
||||
name = "js-sys"
|
||||
version = "0.3.102"
|
||||
version = "0.3.103"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31"
|
||||
checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"futures-util",
|
||||
|
|
@ -564,6 +651,7 @@ name = "litellm-ai-gateway"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"axum",
|
||||
"base64",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"litellm-core",
|
||||
|
|
@ -594,7 +682,9 @@ dependencies = [
|
|||
"litellm-ai-gateway",
|
||||
"litellm-core",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -728,6 +818,19 @@ dependencies = [
|
|||
"unindent",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-async-runtimes"
|
||||
version = "0.23.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"once_cell",
|
||||
"pin-project-lite",
|
||||
"pyo3",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-build-config"
|
||||
version = "0.23.5"
|
||||
|
|
@ -913,6 +1016,7 @@ dependencies = [
|
|||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
|
|
@ -932,12 +1036,14 @@ dependencies = [
|
|||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams",
|
||||
"web-sys",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
|
@ -1335,6 +1441,19 @@ dependencies = [
|
|||
"tungstenite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-util"
|
||||
version = "0.7.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.3"
|
||||
|
|
@ -1507,9 +1626,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a"
|
||||
checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"once_cell",
|
||||
|
|
@ -1520,9 +1639,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-futures"
|
||||
version = "0.4.75"
|
||||
version = "0.4.76"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280"
|
||||
checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
|
|
@ -1530,9 +1649,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d"
|
||||
checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1"
|
||||
dependencies = [
|
||||
"quote",
|
||||
"wasm-bindgen-macro-support",
|
||||
|
|
@ -1540,9 +1659,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-macro-support"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd"
|
||||
checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e"
|
||||
dependencies = [
|
||||
"bumpalo",
|
||||
"proc-macro2",
|
||||
|
|
@ -1553,18 +1672,31 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "wasm-bindgen-shared"
|
||||
version = "0.2.125"
|
||||
version = "0.2.126"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f"
|
||||
checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.102"
|
||||
name = "wasm-streams"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d"
|
||||
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.103"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141"
|
||||
dependencies = [
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
|
|
|
|||
|
|
@ -16,13 +16,15 @@ litellm-core = { path = "crates/core" }
|
|||
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
|
||||
axum = "0.7"
|
||||
pyo3 = "0.23.5"
|
||||
pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] }
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
sha2 = "0.10"
|
||||
subtle = "2"
|
||||
thiserror = "2.0"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
base64 = "0.22"
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "tim
|
|||
tokio-tungstenite.workspace = true
|
||||
futures-util.workspace = true
|
||||
serde_json.workspace = true
|
||||
base64.workspace = true
|
||||
axum = { workspace = true, features = ["ws"], optional = true }
|
||||
serde = { workspace = true, optional = true }
|
||||
subtle = { workspace = true, optional = true }
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
//! End-to-end OCR orchestration.
|
||||
//!
|
||||
//! Owns the whole Mistral OCR call so the Python side stays a thin bridge:
|
||||
//! Owns supported OCR provider calls so the Python side stays a thin bridge:
|
||||
//! resolve the API key, build the URL + body via the pure transforms, POST it,
|
||||
//! and normalize the response. The HTTP client is built once and reused.
|
||||
|
||||
|
|
@ -8,75 +8,135 @@ use std::sync::OnceLock;
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::ocr::transformation::{OcrAuthStrategy, OcrResponseHandling};
|
||||
use litellm_core::CoreResult;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use litellm_core::providers::mistral::ocr::transformation as mistral;
|
||||
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
mod common_utils;
|
||||
|
||||
use common_utils::{
|
||||
convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence,
|
||||
string_headers, truncate_error_body,
|
||||
};
|
||||
|
||||
/// OCR over large documents can take a while; bound it generously rather than
|
||||
/// hanging forever on an unresponsive upstream. The client-level limit is the
|
||||
/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``.
|
||||
const OCR_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Maximum upstream body characters retained in error messages. OCR responses
|
||||
/// can echo document contents and prompts; keep enough for debugging without
|
||||
/// forwarding sensitive payloads across the host boundary.
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
/// Process-wide blocking HTTP client (connection pool + TLS reused across calls).
|
||||
fn http_client() -> &'static reqwest::blocking::Client {
|
||||
static CLIENT: OnceLock<reqwest::blocking::Client> = OnceLock::new();
|
||||
/// Process-wide async HTTP client (connection pool + TLS reused across calls).
|
||||
///
|
||||
/// The Python fallback path uses LiteLLM's standard `BaseLLMHTTPHandler`. This
|
||||
/// Rust path is opt-in and owns end-to-end OCR I/O, so it cannot call the
|
||||
/// Python handler directly; keep this route-scoped until litellm-rust has a
|
||||
/// shared HTTP abstraction.
|
||||
fn http_client() -> &'static reqwest::Client {
|
||||
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::blocking::Client::builder()
|
||||
reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
|
||||
.build()
|
||||
.expect("failed to build reqwest client")
|
||||
})
|
||||
}
|
||||
|
||||
fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
fn upstream_headers(
|
||||
headers: &[(String, String)],
|
||||
auth_strategy: OcrAuthStrategy,
|
||||
api_key: Option<&str>,
|
||||
) -> Vec<(String, String)> {
|
||||
let auth_header = api_key.map(|api_key| match auth_strategy {
|
||||
OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
|
||||
OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
|
||||
});
|
||||
auth_header
|
||||
.into_iter()
|
||||
.chain(headers.iter().cloned())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Perform a Mistral OCR call end to end and return the normalized response as
|
||||
pub struct OcrRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub document: Value,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: &'a str,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
/// Perform an OCR call end to end and return the normalized response as
|
||||
/// JSON (the shape the Python `OCRResponse` model expects).
|
||||
///
|
||||
/// Blocking: intended to be called with the GIL released from the Python bridge.
|
||||
pub fn run_ocr(
|
||||
model: &str,
|
||||
document: Value,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
optional_params: Map<String, Value>,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
let config = &MISTRAL_OCR_CONFIG;
|
||||
/// Async: intended to be awaited directly by the Python bridge's async entrypoint.
|
||||
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
|
||||
let model = request.model;
|
||||
let config = ocr_provider_config(request.custom_llm_provider, model)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.to_string()))?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
|
||||
let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?;
|
||||
let url = mistral::complete_url(api_base);
|
||||
let filtered_params = config.map_ocr_params(&optional_params);
|
||||
let headers = string_headers(request.extra_headers)?;
|
||||
let auth_strategy = config.auth_strategy();
|
||||
let api_key = (!has_header(&headers, auth_strategy.header_name()))
|
||||
.then(|| config.resolve_api_key(request.api_key, &env_lookup))
|
||||
.transpose()?;
|
||||
let url = config.complete_url(
|
||||
request.api_base,
|
||||
model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let filtered_params = config.map_ocr_params(&request.optional_params);
|
||||
let document = if config.requires_data_uri_document() {
|
||||
convert_document_url_to_data_uri(request.document).await?
|
||||
} else {
|
||||
request.document
|
||||
};
|
||||
let body = config
|
||||
.transform_ocr_request(model, document, filtered_params)?
|
||||
.data;
|
||||
let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
|
||||
|
||||
let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body);
|
||||
if let Some(duration) = timeout {
|
||||
request = request.timeout(duration);
|
||||
let mut request_builder = http_client().post(&url).json(&body);
|
||||
for (key, value) in &upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
if config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
&& status.as_u16() == 202
|
||||
{
|
||||
let operation_url = response
|
||||
.headers()
|
||||
.get("operation-location")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse(
|
||||
"Azure Document Intelligence returned 202 but no Operation-Location header found"
|
||||
.to_string(),
|
||||
)
|
||||
})?;
|
||||
let response_json =
|
||||
poll_document_intelligence(&operation_url, &url, &upstream_headers, request.timeout)
|
||||
.await?;
|
||||
return Ok(config
|
||||
.transform_ocr_response(model, response_json)?
|
||||
.into_json());
|
||||
}
|
||||
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
|
|
@ -97,6 +157,25 @@ pub fn run_ocr(
|
|||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
async fn read_http_headers(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
loop {
|
||||
let n = socket.read(&mut buffer).await.expect("reads request");
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
if request.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_passes_short_strings_through() {
|
||||
|
|
@ -106,7 +185,7 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn truncate_error_body_caps_long_payloads() {
|
||||
let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50);
|
||||
let body = "x".repeat(306);
|
||||
let truncated = truncate_error_body(&body);
|
||||
|
||||
assert!(truncated.ends_with("... (truncated)"));
|
||||
|
|
@ -115,13 +194,213 @@ mod tests {
|
|||
.expect("truncated marker present")
|
||||
.chars()
|
||||
.count();
|
||||
assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS);
|
||||
assert_eq!(prefix_chars, 256);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_does_not_split_multibyte_chars() {
|
||||
let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10);
|
||||
let body = "é".repeat(266);
|
||||
let truncated = truncate_error_body(&body);
|
||||
assert!(truncated.is_char_boundary(truncated.len()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ocr_dispatch_supports_migrated_providers() {
|
||||
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
|
||||
assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409")
|
||||
.expect("azure ai config resolves")
|
||||
.requires_data_uri_document());
|
||||
assert_eq!(
|
||||
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
|
||||
.expect("document intelligence config resolves")
|
||||
.response_handling(),
|
||||
OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
);
|
||||
assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
|
||||
.expect("vertex deepseek config resolves")
|
||||
.supported_ocr_params()
|
||||
.contains(&"temperature"));
|
||||
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_accepts_string_values() {
|
||||
let headers = json!({
|
||||
"x-trace-id": "trace-1"
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
|
||||
assert_eq!(
|
||||
string_headers(Some(headers)).expect("string headers accepted"),
|
||||
vec![("x-trace-id".to_string(), "trace-1".to_string())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_header_detection_is_case_insensitive() {
|
||||
let headers = vec![
|
||||
("x-trace-id".to_string(), "trace-1".to_string()),
|
||||
("authorization".to_string(), "Bearer sk-test".to_string()),
|
||||
];
|
||||
|
||||
assert!(has_header(&headers, "authorization"));
|
||||
|
||||
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
|
||||
assert!(has_header(&headers, "authorization"));
|
||||
|
||||
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
|
||||
assert!(!has_header(&headers, "authorization"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts one request");
|
||||
let request = read_http_headers(&mut socket).await;
|
||||
|
||||
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
response_body.len(),
|
||||
response_body
|
||||
);
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
request
|
||||
});
|
||||
|
||||
let mut headers = Map::new();
|
||||
headers.insert(
|
||||
"Authorization".to_string(),
|
||||
Value::String("Bearer sk-from-python".to_string()),
|
||||
);
|
||||
headers.insert(
|
||||
"x-trace-id".to_string(),
|
||||
Value::String("trace-1".to_string()),
|
||||
);
|
||||
|
||||
let response = ocr(OcrRequest {
|
||||
model: "mistral-ocr-latest",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("sk-for-rust-fallback"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: "mistral",
|
||||
extra_headers: Some(headers),
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
})
|
||||
.await
|
||||
.expect("ocr request succeeds");
|
||||
|
||||
assert_eq!(response["pages"][0]["markdown"], "ok");
|
||||
|
||||
let request = server.await.expect("server task completes");
|
||||
let authorization_count = request
|
||||
.lines()
|
||||
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
|
||||
.count();
|
||||
assert_eq!(authorization_count, 1, "{request}");
|
||||
assert!(
|
||||
request.contains("authorization: Bearer sk-from-python")
|
||||
|| request.contains("Authorization: Bearer sk-from-python"),
|
||||
"{request}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn document_intelligence_poll_uses_resolved_subscription_key() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("test listener binds");
|
||||
let addr = listener.local_addr().expect("listener has local addr");
|
||||
let operation_url = format!("http://{addr}/operations/1");
|
||||
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
|
||||
let post_request = read_http_headers(&mut post_socket).await;
|
||||
let post_response = format!(
|
||||
"HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
|
||||
);
|
||||
post_socket
|
||||
.write_all(post_response.as_bytes())
|
||||
.await
|
||||
.expect("writes post response");
|
||||
|
||||
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
|
||||
let poll_request = read_http_headers(&mut poll_socket).await;
|
||||
let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
|
||||
let poll_response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
response_body.len(),
|
||||
response_body
|
||||
);
|
||||
poll_socket
|
||||
.write_all(poll_response.as_bytes())
|
||||
.await
|
||||
.expect("writes poll response");
|
||||
(post_request, poll_request)
|
||||
});
|
||||
|
||||
let response = ocr(OcrRequest {
|
||||
model: "prebuilt-read",
|
||||
document: json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
}),
|
||||
api_key: Some("di-key"),
|
||||
api_base: Some(&format!("http://{addr}")),
|
||||
custom_llm_provider: "azure_ai/doc-intelligence",
|
||||
extra_headers: None,
|
||||
optional_params: Map::new(),
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
})
|
||||
.await
|
||||
.expect("document intelligence request succeeds");
|
||||
|
||||
assert_eq!(response["pages"][0]["markdown"], "ok");
|
||||
|
||||
let (post_request, poll_request) = server.await.expect("server task completes");
|
||||
assert!(
|
||||
post_request
|
||||
.to_ascii_lowercase()
|
||||
.contains("ocp-apim-subscription-key: di-key"),
|
||||
"{post_request}"
|
||||
);
|
||||
assert!(
|
||||
poll_request
|
||||
.to_ascii_lowercase()
|
||||
.contains("ocp-apim-subscription-key: di-key"),
|
||||
"{poll_request}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_rejects_non_string_values() {
|
||||
let headers = json!({
|
||||
"x-retry-count": 3
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
|
||||
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidRequest(
|
||||
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
440
litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs
Normal file
440
litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs
Normal file
|
|
@ -0,0 +1,440 @@
|
|||
use std::net::IpAddr;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use base64::Engine;
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::CoreResult;
|
||||
use reqwest::Url;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use litellm_core::providers::azure_ai::ocr::transformation::{
|
||||
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
|
||||
};
|
||||
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
|
||||
use litellm_core::providers::vertex_ai::ocr::transformation::{
|
||||
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
|
||||
};
|
||||
|
||||
use super::http_client;
|
||||
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
|
||||
const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
|
||||
const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
pub(super) fn ocr_provider_config(
|
||||
provider: &str,
|
||||
model: &str,
|
||||
) -> Option<&'static dyn OcrProviderConfig> {
|
||||
match provider {
|
||||
"mistral" => Some(&MISTRAL_OCR_CONFIG),
|
||||
"azure_ai" if is_azure_document_intelligence_model(model) => {
|
||||
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
|
||||
}
|
||||
"azure_ai/doc-intelligence" => Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG),
|
||||
"azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
|
||||
"vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
|
||||
"vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_azure_document_intelligence_model(model: &str) -> bool {
|
||||
let model = model.to_ascii_lowercase();
|
||||
model.contains("doc-intelligence") || model.contains("documentintelligence")
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"OCR extra_headers.{key} must be a string, got {}",
|
||||
litellm_core::error::json_type_name(&value)
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
fn document_url_field(document: &Value) -> CoreResult<Option<(&str, &str)>> {
|
||||
let Some(object) = document.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(doc_type) = object.get("type").and_then(Value::as_str) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let field = match doc_type {
|
||||
"document_url" => "document_url",
|
||||
"image_url" => "image_url",
|
||||
_ => return Ok(None),
|
||||
};
|
||||
let Some(url) = object.get(field).and_then(Value::as_str) else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some((field, url)))
|
||||
}
|
||||
|
||||
fn is_url_requiring_fetch(url: &str) -> bool {
|
||||
!url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://"))
|
||||
}
|
||||
|
||||
fn max_document_download_bytes() -> u64 {
|
||||
let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
|
||||
.ok()
|
||||
.and_then(|value| value.parse::<f64>().ok())
|
||||
.unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB);
|
||||
(max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64
|
||||
}
|
||||
|
||||
fn is_blocked_ip(ip: IpAddr) -> bool {
|
||||
match ip {
|
||||
IpAddr::V4(ip) => {
|
||||
ip.is_private()
|
||||
|| ip.is_loopback()
|
||||
|| ip.is_link_local()
|
||||
|| ip.is_broadcast()
|
||||
|| ip.is_multicast()
|
||||
|| ip.is_unspecified()
|
||||
}
|
||||
IpAddr::V6(ip) => {
|
||||
let first_segment = ip.segments()[0];
|
||||
let is_unique_local = (first_segment & 0xfe00) == 0xfc00;
|
||||
let is_link_local = (first_segment & 0xffc0) == 0xfe80;
|
||||
ip.is_loopback()
|
||||
|| ip.is_unspecified()
|
||||
|| ip.is_multicast()
|
||||
|| is_unique_local
|
||||
|| is_link_local
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn blocked_url_error(url: &Url) -> CoreError {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"OCR document URL rejected by SSRF protection: {url}"
|
||||
))
|
||||
}
|
||||
|
||||
async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
|
||||
if !matches!(url.scheme(), "http" | "https") {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
|
||||
let host = url.host_str().ok_or_else(|| blocked_url_error(url))?;
|
||||
if let Ok(ip) = host.parse::<IpAddr>() {
|
||||
if is_blocked_ip(ip) {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let port = url
|
||||
.port_or_known_default()
|
||||
.ok_or_else(|| blocked_url_error(url))?;
|
||||
let addresses = tokio::net::lookup_host((host, port))
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let mut saw_address = false;
|
||||
for address in addresses {
|
||||
saw_address = true;
|
||||
if is_blocked_ip(address.ip()) {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
}
|
||||
if !saw_address {
|
||||
return Err(blocked_url_error(url));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult<Url> {
|
||||
let location = response
|
||||
.headers()
|
||||
.get(reqwest::header::LOCATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("OCR document redirect missing Location header".to_string())
|
||||
})?;
|
||||
url.join(location)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}")))
|
||||
}
|
||||
|
||||
async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> {
|
||||
let client = reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let mut current_url = Url::parse(url)
|
||||
.map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
|
||||
|
||||
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
|
||||
validate_safe_fetch_url(¤t_url).await?;
|
||||
let response = client
|
||||
.get(current_url.clone())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
if !response.status().is_redirection() {
|
||||
return Ok((current_url, response));
|
||||
}
|
||||
current_url = redirect_location(&response, ¤t_url)?;
|
||||
}
|
||||
|
||||
Err(CoreError::InvalidRequest(
|
||||
"Too many redirects while fetching OCR document URL".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> {
|
||||
if max_bytes == 0 {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
|
||||
)));
|
||||
}
|
||||
if content_length > max_bytes {
|
||||
let size_mb = content_length as f64 / (1024.0 * 1024.0);
|
||||
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_response_with_limit(
|
||||
mut response: reqwest::Response,
|
||||
url: &Url,
|
||||
) -> CoreResult<Vec<u8>> {
|
||||
let max_bytes = max_document_download_bytes();
|
||||
if let Some(content_length) = response.content_length() {
|
||||
enforce_download_size(content_length, max_bytes, url)?;
|
||||
} else {
|
||||
enforce_download_size(0, max_bytes, url)?;
|
||||
}
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
let mut bytes_downloaded: u64 = 0;
|
||||
while let Some(chunk) = response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?
|
||||
{
|
||||
bytes_downloaded += chunk.len() as u64;
|
||||
enforce_download_size(bytes_downloaded, max_bytes, url)?;
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult<Value> {
|
||||
let Some((field, url)) = document_url_field(&document)? else {
|
||||
return Ok(document);
|
||||
};
|
||||
if !is_url_requiring_fetch(url) {
|
||||
return Ok(document);
|
||||
}
|
||||
|
||||
let (final_url, response) = safe_get_document_url(url).await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&body),
|
||||
});
|
||||
}
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.split(';').next())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("application/octet-stream")
|
||||
.to_string();
|
||||
let bytes = read_response_with_limit(response, &final_url).await?;
|
||||
let data_uri = format!(
|
||||
"data:{content_type};base64,{}",
|
||||
BASE64_STANDARD.encode(bytes)
|
||||
);
|
||||
|
||||
let mut transformed = document
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?;
|
||||
transformed.insert(field.to_string(), Value::String(data_uri));
|
||||
Ok(Value::Object(transformed))
|
||||
}
|
||||
|
||||
fn same_origin(left: &str, right: &str) -> bool {
|
||||
let Ok(left) = reqwest::Url::parse(left) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(right) = reqwest::Url::parse(right) else {
|
||||
return false;
|
||||
};
|
||||
left.scheme() == right.scheme()
|
||||
&& left.host_str() == right.host_str()
|
||||
&& left.port_or_known_default() == right.port_or_known_default()
|
||||
}
|
||||
|
||||
fn retry_after_secs(response: &reqwest::Response) -> u64 {
|
||||
response
|
||||
.headers()
|
||||
.get(reqwest::header::RETRY_AFTER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.unwrap_or(2)
|
||||
}
|
||||
|
||||
fn operation_status(response_json: &Value) -> CoreResult<&str> {
|
||||
let status = response_json
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("status"))?;
|
||||
match status {
|
||||
"succeeded" => Ok("succeeded"),
|
||||
"running" | "notStarted" => Ok("running"),
|
||||
"failed" => {
|
||||
let message = response_json
|
||||
.get("error")
|
||||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("Unknown error");
|
||||
Err(CoreError::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed: {message}"
|
||||
)))
|
||||
}
|
||||
other => Err(CoreError::InvalidResponse(format!(
|
||||
"Unknown operation status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn poll_document_intelligence(
|
||||
operation_url: &str,
|
||||
original_url: &str,
|
||||
headers: &[(String, String)],
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
if !same_origin(operation_url, original_url) {
|
||||
return Err(CoreError::InvalidResponse(
|
||||
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
let timeout = timeout.unwrap_or(Duration::from_secs(
|
||||
AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
|
||||
));
|
||||
loop {
|
||||
if start.elapsed() > timeout {
|
||||
return Err(CoreError::Network(format!(
|
||||
"Azure Document Intelligence operation polling timed out after {} seconds",
|
||||
timeout.as_secs()
|
||||
)));
|
||||
}
|
||||
|
||||
let mut request_builder = http_client().get(operation_url);
|
||||
for (key, value) in headers {
|
||||
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
}
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
let retry_after = retry_after_secs(&response);
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
|
||||
})?;
|
||||
if operation_status(&response_json)? == "succeeded" {
|
||||
return Ok(response_json);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs(retry_after)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn blocks_private_and_metadata_ips() {
|
||||
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("10.0.0.1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("169.254.169.254".parse().unwrap()));
|
||||
assert!(is_blocked_ip("::1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("fd00::1".parse().unwrap()));
|
||||
assert!(is_blocked_ip("fe80::1".parse().unwrap()));
|
||||
assert!(!is_blocked_ip("8.8.8.8".parse().unwrap()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn convert_document_url_rejects_loopback_fetch() {
|
||||
let error = convert_document_url_to_data_uri(json!({
|
||||
"type": "image_url",
|
||||
"image_url": "http://127.0.0.1/image.png"
|
||||
}))
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CoreError::InvalidRequest(message)
|
||||
if message.contains("SSRF protection")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn convert_document_url_leaves_data_uri_untouched() {
|
||||
let document = json!({
|
||||
"type": "image_url",
|
||||
"image_url": "data:image/png;base64,abcd"
|
||||
});
|
||||
|
||||
let transformed = convert_document_url_to_data_uri(document.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(transformed, document);
|
||||
}
|
||||
}
|
||||
|
|
@ -13,6 +13,10 @@ pub enum CoreError {
|
|||
MissingField(&'static str),
|
||||
#[error("invalid response: {0}")]
|
||||
InvalidResponse(String),
|
||||
#[error("invalid provider: {0}")]
|
||||
InvalidProvider(String),
|
||||
#[error("invalid request: {0}")]
|
||||
InvalidRequest(String),
|
||||
#[error("{0}")]
|
||||
Auth(String),
|
||||
#[error("OCR request failed with status {status}: {body}")]
|
||||
|
|
|
|||
|
|
@ -4,7 +4,28 @@ use crate::CoreResult;
|
|||
|
||||
use super::types::{OcrRequestData, OcrResponseData};
|
||||
|
||||
pub trait OcrProviderConfig {
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrAuthStrategy {
|
||||
Bearer,
|
||||
Header(&'static str),
|
||||
}
|
||||
|
||||
impl OcrAuthStrategy {
|
||||
pub fn header_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Bearer => "authorization",
|
||||
Self::Header(header_name) => header_name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OcrResponseHandling {
|
||||
Json,
|
||||
AzureDocumentIntelligencePoll,
|
||||
}
|
||||
|
||||
pub trait OcrProviderConfig: Sync {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str];
|
||||
|
||||
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
|
|
@ -29,4 +50,30 @@ pub trait OcrProviderConfig {
|
|||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData>;
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn auth_strategy(&self) -> OcrAuthStrategy {
|
||||
OcrAuthStrategy::Bearer
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn response_handling(&self) -> OcrResponseHandling {
|
||||
OcrResponseHandling::Json
|
||||
}
|
||||
}
|
||||
|
|
|
|||
1
litellm-rust/crates/core/src/providers/azure_ai/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/azure_ai/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,520 @@
|
|||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
|
||||
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30";
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96;
|
||||
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages"];
|
||||
|
||||
pub struct AzureAiOcrConfig;
|
||||
pub struct AzureDocumentIntelligenceOcrConfig;
|
||||
|
||||
pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig;
|
||||
pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig =
|
||||
AzureDocumentIntelligenceOcrConfig;
|
||||
|
||||
fn non_empty(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn resolve_value(
|
||||
explicit: Option<&str>,
|
||||
env_name: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
missing_message: &str,
|
||||
) -> CoreResult<String> {
|
||||
non_empty(explicit)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| CoreError::Auth(missing_message.to_string()))
|
||||
}
|
||||
|
||||
pub fn resolve_azure_ai_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_key,
|
||||
AZURE_AI_API_KEY_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params",
|
||||
)
|
||||
}
|
||||
|
||||
pub fn resolve_azure_ai_api_base(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_base,
|
||||
AZURE_AI_API_BASE_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter",
|
||||
)
|
||||
}
|
||||
|
||||
pub fn complete_azure_ai_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let base = resolve_azure_ai_api_base(api_base, env_lookup)?;
|
||||
Ok(format!(
|
||||
"{}/providers/mistral/azure/ocr",
|
||||
base.trim_end_matches('/')
|
||||
))
|
||||
}
|
||||
|
||||
pub fn resolve_document_intelligence_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_key,
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter",
|
||||
)
|
||||
}
|
||||
|
||||
pub fn resolve_document_intelligence_endpoint(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_value(
|
||||
api_base,
|
||||
AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV,
|
||||
env_lookup,
|
||||
"Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter",
|
||||
)
|
||||
}
|
||||
|
||||
fn encode_model_id(model: &str) -> String {
|
||||
let model_id = model.rsplit('/').next().unwrap_or(model);
|
||||
model_id
|
||||
.bytes()
|
||||
.flat_map(|byte| match byte {
|
||||
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
|
||||
vec![byte as char]
|
||||
}
|
||||
_ => format!("%{byte:02X}").chars().collect(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn pages_token_is_valid(token: &str) -> bool {
|
||||
let mut parts = token.split('-');
|
||||
let Some(start) = parts.next() else {
|
||||
return false;
|
||||
};
|
||||
if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) {
|
||||
return false;
|
||||
}
|
||||
match parts.next() {
|
||||
None => true,
|
||||
Some(end) => {
|
||||
!end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
|
||||
match pages {
|
||||
Value::String(value) => {
|
||||
let normalized = value
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
if normalized.split(',').all(pages_token_is_valid) {
|
||||
Ok(Some(normalized))
|
||||
} else {
|
||||
Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'."
|
||||
)))
|
||||
}
|
||||
}
|
||||
Value::Array(values) => {
|
||||
if values.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
if values.iter().all(Value::is_i64) {
|
||||
let mut pages = BTreeSet::new();
|
||||
for value in values {
|
||||
let page = value.as_i64().expect("checked is_i64");
|
||||
if page < 0 {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(),
|
||||
));
|
||||
}
|
||||
pages.insert(page + 1);
|
||||
}
|
||||
return Ok(Some(
|
||||
pages
|
||||
.into_iter()
|
||||
.map(|page| page.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(","),
|
||||
));
|
||||
}
|
||||
if values.iter().all(Value::is_string) {
|
||||
let normalized = values
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
if normalized.split(',').all(pages_token_is_valid) {
|
||||
return Ok(Some(normalized));
|
||||
}
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'."
|
||||
)));
|
||||
}
|
||||
Err(CoreError::InvalidRequest(
|
||||
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
|
||||
.to_string(),
|
||||
))
|
||||
}
|
||||
_ => Err(CoreError::InvalidRequest(
|
||||
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
|
||||
.to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn complete_document_intelligence_url(
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?;
|
||||
let mut url = format!(
|
||||
"{}/documentintelligence/documentModels/{}:analyze?api-version={}",
|
||||
endpoint.trim_end_matches('/'),
|
||||
encode_model_id(model),
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION
|
||||
);
|
||||
|
||||
if let Some(pages) = optional_params.get("pages") {
|
||||
if let Some(normalized) = normalize_pages_param(pages)? {
|
||||
url.push_str("&pages=");
|
||||
url.push_str(&normalized);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
|
||||
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("document.type"))?;
|
||||
let field_name = match doc_type {
|
||||
"document_url" => "document_url",
|
||||
"image_url" => "image_url",
|
||||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Invalid document type: {other}. Must be 'document_url' or 'image_url'"
|
||||
)))
|
||||
}
|
||||
};
|
||||
object
|
||||
.get(field_name)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField(field_name))
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_uri(data_uri: &str) -> &str {
|
||||
data_uri
|
||||
.split_once(',')
|
||||
.map(|(_, data)| data)
|
||||
.unwrap_or(data_uri)
|
||||
}
|
||||
|
||||
fn page_markdown(page: &Map<String, Value>) -> String {
|
||||
page.get("lines")
|
||||
.and_then(Value::as_array)
|
||||
.map(|lines| {
|
||||
lines
|
||||
.iter()
|
||||
.filter_map(|line| line.get("content").and_then(Value::as_str))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn page_dimensions(page: &Map<String, Value>) -> Value {
|
||||
let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5);
|
||||
let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0);
|
||||
let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch");
|
||||
let (width, height) = if unit == "inch" {
|
||||
(
|
||||
(width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64,
|
||||
(height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64,
|
||||
)
|
||||
} else {
|
||||
(width as i64, height as i64)
|
||||
};
|
||||
json!({
|
||||
"width": width,
|
||||
"height": height,
|
||||
"dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI,
|
||||
})
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for AzureAiOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_azure_ai_url(api_base, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_azure_ai_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
_model: &str,
|
||||
document: Value,
|
||||
_optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
let document_url = document_url_from_mistral_document(&document)?;
|
||||
let mut data = Map::new();
|
||||
if document_url.starts_with("data:") {
|
||||
data.insert(
|
||||
"base64Source".to_string(),
|
||||
Value::String(extract_base64_from_data_uri(document_url).to_string()),
|
||||
);
|
||||
} else {
|
||||
data.insert(
|
||||
"urlSource".to_string(),
|
||||
Value::String(document_url.to_string()),
|
||||
);
|
||||
}
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
let response = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
let status = response
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("status"))?;
|
||||
if status != "succeeded" {
|
||||
return Err(CoreError::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed with status: {status}"
|
||||
)));
|
||||
}
|
||||
|
||||
let azure_pages = response
|
||||
.get("analyzeResult")
|
||||
.and_then(|result| result.get("pages"))
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
let pages = azure_pages
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.map(|page| {
|
||||
let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1);
|
||||
json!({
|
||||
"index": page_number - 1,
|
||||
"markdown": page_markdown(page),
|
||||
"dimensions": page_dimensions(page),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(OcrResponseData {
|
||||
usage_info: Some(json!({
|
||||
"pages_processed": pages.len(),
|
||||
"doc_size_bytes": null,
|
||||
})),
|
||||
pages,
|
||||
model: model.to_string(),
|
||||
document_annotation: None,
|
||||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_document_intelligence_url(api_base, model, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_document_intelligence_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn auth_strategy(&self) -> OcrAuthStrategy {
|
||||
OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key")
|
||||
}
|
||||
|
||||
fn response_handling(&self) -> OcrResponseHandling {
|
||||
OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn azure_ai_reuses_mistral_body_transform() {
|
||||
let body = AZURE_AI_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"pixtral-12b-2409",
|
||||
json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}),
|
||||
serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body["model"], "pixtral-12b-2409");
|
||||
assert_eq!(body["include_image_base64"], true);
|
||||
assert_eq!(
|
||||
body["document"]["document_url"],
|
||||
"data:application/pdf;base64,abc"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_url_normalizes_zero_based_pages() {
|
||||
let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]);
|
||||
let url = complete_document_intelligence_url(
|
||||
Some("https://example.cognitiveservices.azure.com/"),
|
||||
"azure_ai/doc-intelligence/prebuilt-layout",
|
||||
¶ms,
|
||||
&|_| None,
|
||||
)
|
||||
.expect("url builds");
|
||||
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_request_uses_base64_source_for_data_uri() {
|
||||
let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"prebuilt-read",
|
||||
json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}),
|
||||
Map::new(),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body, json!({"base64Source": "abc123"}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_response_normalizes_pages() {
|
||||
let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
|
||||
.transform_ocr_response(
|
||||
"prebuilt-layout",
|
||||
json!({
|
||||
"status": "succeeded",
|
||||
"analyzeResult": {
|
||||
"pages": [{
|
||||
"pageNumber": 2,
|
||||
"width": 8.5,
|
||||
"height": 11,
|
||||
"unit": "inch",
|
||||
"lines": [{"content": "hello"}, {"content": "world"}]
|
||||
}]
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(response.pages[0]["index"], 1);
|
||||
assert_eq!(response.pages[0]["markdown"], "hello\nworld");
|
||||
assert_eq!(response.pages[0]["dimensions"]["width"], 816);
|
||||
assert_eq!(
|
||||
response.usage_info,
|
||||
Some(json!({"pages_processed": 1, "doc_size_bytes": null}))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -132,6 +132,24 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
Ok(complete_url(api_base))
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_api_key(api_key, env_lookup)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn supported_ocr_params() -> &'static [&'static str] {
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
pub mod azure_ai;
|
||||
pub mod mistral;
|
||||
pub mod openai;
|
||||
pub mod vertex_ai;
|
||||
|
|
|
|||
1
litellm-rust/crates/core/src/providers/vertex_ai/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/vertex_ai/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,435 @@
|
|||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
const VERTEX_DEFAULT_LOCATION: &str = "us-central1";
|
||||
const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com";
|
||||
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
|
||||
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
|
||||
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
|
||||
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
|
||||
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
|
||||
|
||||
#[rustfmt::skip]
|
||||
const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[
|
||||
"stream",
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"n",
|
||||
"stop",
|
||||
];
|
||||
|
||||
pub struct VertexAiOcrConfig;
|
||||
pub struct VertexAiDeepSeekOcrConfig;
|
||||
|
||||
pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig;
|
||||
pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig;
|
||||
|
||||
fn string_param<'a>(params: &'a Map<String, Value>, keys: &[&str]) -> Option<&'a str> {
|
||||
keys.iter()
|
||||
.find_map(|key| params.get(*key).and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub fn is_deepseek_model(model: &str) -> bool {
|
||||
model.to_ascii_lowercase().contains("deepseek")
|
||||
}
|
||||
|
||||
pub fn resolve_vertex_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::Auth(
|
||||
"Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn vertex_project(
|
||||
params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
string_param(params, &["vertex_project", "vertex_ai_project"])
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(
|
||||
"Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter"
|
||||
.to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn vertex_location(
|
||||
params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
string_param(params, &["vertex_location", "vertex_ai_location"])
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string())
|
||||
}
|
||||
|
||||
fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String {
|
||||
api_base
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com"))
|
||||
.trim_end_matches('/')
|
||||
.to_string()
|
||||
}
|
||||
|
||||
pub fn complete_vertex_mistral_url(
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let project = vertex_project(optional_params, env_lookup)?;
|
||||
let location = vertex_location(optional_params, env_lookup);
|
||||
let base = vertex_mistral_api_base(api_base, &location);
|
||||
Ok(format!(
|
||||
"{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict"
|
||||
))
|
||||
}
|
||||
|
||||
pub fn complete_vertex_deepseek_url(
|
||||
api_base: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let project = vertex_project(optional_params, env_lookup)?;
|
||||
let location = vertex_location(optional_params, env_lookup);
|
||||
let base = api_base
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE)
|
||||
.trim_end_matches('/');
|
||||
Ok(format!(
|
||||
"{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions"
|
||||
))
|
||||
}
|
||||
|
||||
fn document_content_item(document: &Value) -> CoreResult<Value> {
|
||||
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("document.type"))?;
|
||||
let url_field = match doc_type {
|
||||
"image_url" => "image_url",
|
||||
"document_url" => "document_url",
|
||||
other => {
|
||||
return Err(CoreError::InvalidRequest(format!(
|
||||
"Unsupported document type: {other}. Expected 'image_url' or 'document_url'"
|
||||
)))
|
||||
}
|
||||
};
|
||||
let url = object
|
||||
.get(url_field)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(CoreError::MissingField(url_field))?;
|
||||
|
||||
Ok(json!({
|
||||
"type": "image_url",
|
||||
"image_url": url,
|
||||
}))
|
||||
}
|
||||
|
||||
fn deepseek_model_name(model: &str) -> String {
|
||||
if model.starts_with("deepseek-ai/") {
|
||||
model.to_string()
|
||||
} else {
|
||||
format!("deepseek-ai/{model}")
|
||||
}
|
||||
}
|
||||
|
||||
fn first_choice_content(response: &Value) -> CoreResult<Value> {
|
||||
response
|
||||
.get("choices")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|choices| choices.first())
|
||||
.and_then(|choice| choice.get("message"))
|
||||
.and_then(|message| message.get("content"))
|
||||
.cloned()
|
||||
.filter(|content| match content {
|
||||
Value::String(value) => !value.is_empty(),
|
||||
Value::Object(_) => true,
|
||||
_ => false,
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> Value {
|
||||
match content {
|
||||
Value::String(content) => {
|
||||
if content.trim_start().starts_with('{') {
|
||||
serde_json::from_str(&content).unwrap_or_else(|_| {
|
||||
json!({
|
||||
"pages": [{"index": 0, "markdown": content}],
|
||||
"model": model,
|
||||
"usage_info": usage.unwrap_or_else(|| json!({})),
|
||||
})
|
||||
})
|
||||
} else {
|
||||
json!({
|
||||
"pages": [{"index": 0, "markdown": content}],
|
||||
"model": model,
|
||||
"usage_info": usage.unwrap_or_else(|| json!({})),
|
||||
})
|
||||
}
|
||||
}
|
||||
Value::Object(_) => content,
|
||||
other => json!({
|
||||
"pages": [{"index": 0, "markdown": other.to_string()}],
|
||||
"model": model,
|
||||
"usage_info": usage.unwrap_or_else(|| json!({})),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for VertexAiOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_vertex_mistral_url(api_base, model, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_vertex_api_key(api_key, env_lookup)
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
DEEPSEEK_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
let mut data = Map::new();
|
||||
data.insert(
|
||||
"model".to_string(),
|
||||
Value::String(deepseek_model_name(model)),
|
||||
);
|
||||
data.insert(
|
||||
"messages".to_string(),
|
||||
json!([{"role": "user", "content": [document_content_item(&document)?]}]),
|
||||
);
|
||||
for (key, value) in optional_params {
|
||||
if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) {
|
||||
data.insert(key, value);
|
||||
}
|
||||
}
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
let response = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
let usage = response.get("usage").cloned();
|
||||
let content = first_choice_content(&response_json)?;
|
||||
let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model);
|
||||
|
||||
if !ocr_data.get("pages").is_some_and(Value::is_array) {
|
||||
ocr_data = json!({
|
||||
"pages": [{
|
||||
"index": 0,
|
||||
"markdown": match content {
|
||||
Value::String(value) => value,
|
||||
other => other.to_string(),
|
||||
}
|
||||
}],
|
||||
"model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model),
|
||||
"usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})),
|
||||
});
|
||||
}
|
||||
|
||||
let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&ocr_data),
|
||||
})?;
|
||||
let pages = object
|
||||
.get("pages")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let usage_info = object
|
||||
.get("usage_info")
|
||||
.cloned()
|
||||
.or_else(|| response.get("usage").cloned());
|
||||
Ok(OcrResponseData {
|
||||
pages,
|
||||
model: object
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(model)
|
||||
.to_string(),
|
||||
document_annotation: object.get("document_annotation").cloned(),
|
||||
usage_info,
|
||||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_vertex_deepseek_url(api_base, optional_params, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
resolve_vertex_api_key(api_key, env_lookup)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn vertex_mistral_url_uses_project_location_and_model() {
|
||||
let params = Map::from_iter([
|
||||
("vertex_project".to_string(), json!("proj-1")),
|
||||
("vertex_location".to_string(), json!("europe-west4")),
|
||||
]);
|
||||
|
||||
let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None)
|
||||
.expect("url builds");
|
||||
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_mistral_reuses_mistral_body_transform() {
|
||||
let body = VERTEX_AI_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"mistral-ocr-maas",
|
||||
json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}),
|
||||
Map::new(),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body["model"], "mistral-ocr-maas");
|
||||
assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_deepseek_request_uses_ocr_endpoint_shape() {
|
||||
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
|
||||
.transform_ocr_request(
|
||||
"deepseek-ocr-maas",
|
||||
json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}),
|
||||
Map::from_iter([("temperature".to_string(), json!(0.1))]),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.data;
|
||||
|
||||
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
|
||||
assert_eq!(body["temperature"], 0.1);
|
||||
assert_eq!(
|
||||
body["messages"][0]["content"][0],
|
||||
json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_deepseek_response_wraps_markdown_content() {
|
||||
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG
|
||||
.transform_ocr_response(
|
||||
"deepseek-ocr-maas",
|
||||
json!({
|
||||
"choices": [{"message": {"content": "# OCR text"}}],
|
||||
"usage": {"prompt_tokens": 1}
|
||||
}),
|
||||
)
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(
|
||||
response.pages,
|
||||
vec![json!({"index": 0, "markdown": "# OCR text"})]
|
||||
);
|
||||
assert_eq!(response.model, "deepseek-ocr-maas");
|
||||
assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1})));
|
||||
}
|
||||
}
|
||||
|
|
@ -13,4 +13,6 @@ crate-type = ["cdylib"]
|
|||
litellm-core.workspace = true
|
||||
litellm-ai-gateway = { workspace = true, default-features = false }
|
||||
pyo3 = { workspace = true, features = ["extension-module"] }
|
||||
pyo3-async-runtimes.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_ai_gateway::io::ocr::run_ocr;
|
||||
use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest};
|
||||
use litellm_core::error::CoreError;
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
|
|
@ -9,6 +9,13 @@ use serde_json::{Map, Value};
|
|||
|
||||
mod gil;
|
||||
|
||||
type MarshaledOcrInputs = (
|
||||
Value,
|
||||
Option<Map<String, Value>>,
|
||||
Map<String, Value>,
|
||||
Option<Duration>,
|
||||
);
|
||||
|
||||
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
|
||||
let json = py.import("json")?;
|
||||
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
|
||||
|
|
@ -22,59 +29,93 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
|
|||
Ok(json.call_method1("loads", (encoded,))?.unbind())
|
||||
}
|
||||
|
||||
/// Map a core error to the closest Python exception. Caller-input problems
|
||||
/// (auth, bad types, missing fields) -> `ValueError`; everything else
|
||||
/// (network, upstream status, parse failures) -> `RuntimeError`.
|
||||
fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Auth(message) => PyValueError::new_err(message),
|
||||
CoreError::InvalidType { .. } | CoreError::MissingField(_) => {
|
||||
PyValueError::new_err(err.to_string())
|
||||
}
|
||||
CoreError::InvalidProvider(_)
|
||||
| CoreError::InvalidRequest(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_) => PyValueError::new_err(err.to_string()),
|
||||
other => PyRuntimeError::new_err(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a Mistral OCR call end to end and return the response as a dict.
|
||||
fn optional_object_to_map(
|
||||
py: Python<'_>,
|
||||
name: &'static str,
|
||||
value: Option<Py<PyAny>>,
|
||||
) -> PyResult<Map<String, Value>> {
|
||||
match value {
|
||||
Some(value) => match py_to_json(py, value.bind(py))? {
|
||||
Value::Object(map) => Ok(map),
|
||||
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
|
||||
},
|
||||
None => Ok(Map::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
|
||||
timeout_seconds.and_then(|secs| {
|
||||
if secs.is_finite() && secs > 0.0 {
|
||||
Some(Duration::from_secs_f64(secs))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn marshal_inputs(
|
||||
py: Python<'_>,
|
||||
document: Py<PyAny>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<MarshaledOcrInputs> {
|
||||
let document = py_to_json(py, document.bind(py))?;
|
||||
let extra_headers = match extra_headers {
|
||||
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
|
||||
None => None,
|
||||
};
|
||||
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
|
||||
let timeout = optional_timeout(timeout_seconds);
|
||||
|
||||
Ok((document, extra_headers, optional_params, timeout))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))]
|
||||
#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn ocr(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
document: Py<PyAny>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let document = py_to_json(py, document.bind(py))?;
|
||||
let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string());
|
||||
let (document, extra_headers, optional_params, timeout) = marshal_inputs(
|
||||
py,
|
||||
document,
|
||||
extra_headers,
|
||||
optional_params,
|
||||
timeout_seconds,
|
||||
)?;
|
||||
|
||||
let optional_params = match optional_params {
|
||||
Some(params) => match py_to_json(py, params.bind(py))? {
|
||||
Value::Object(map) => map,
|
||||
_ => return Err(PyValueError::new_err("optional_params must be a dict")),
|
||||
},
|
||||
None => Map::new(),
|
||||
};
|
||||
|
||||
let timeout = timeout_seconds.and_then(|secs| {
|
||||
if secs.is_finite() && secs > 0.0 {
|
||||
Some(Duration::from_secs_f64(secs))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
|
||||
// Release the GIL during the blocking HTTP call (counted for observability).
|
||||
let result = gil::release_gil(py, || {
|
||||
run_ocr(
|
||||
&model,
|
||||
pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest {
|
||||
model: &model,
|
||||
document,
|
||||
api_key.as_deref(),
|
||||
api_base.as_deref(),
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: &custom_llm_provider,
|
||||
extra_headers,
|
||||
optional_params,
|
||||
timeout,
|
||||
)
|
||||
}))
|
||||
});
|
||||
|
||||
match result {
|
||||
|
|
@ -83,8 +124,47 @@ fn ocr(
|
|||
}
|
||||
}
|
||||
|
||||
/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe
|
||||
/// how often the bridge has dropped the GIL for blocking work.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn aocr(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
document: Py<PyAny>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string());
|
||||
let (document, extra_headers, optional_params, timeout) = marshal_inputs(
|
||||
py,
|
||||
document,
|
||||
extra_headers,
|
||||
optional_params,
|
||||
timeout_seconds,
|
||||
)?;
|
||||
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let value = run_ocr(OcrRequest {
|
||||
model: &model,
|
||||
document,
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: &custom_llm_provider,
|
||||
extra_headers,
|
||||
optional_params,
|
||||
timeout,
|
||||
})
|
||||
.await
|
||||
.map_err(core_error_to_pyerr)?;
|
||||
|
||||
Python::with_gil(|py| json_to_py(py, value))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let stats = PyDict::new(py);
|
||||
|
|
@ -95,6 +175,7 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
|||
#[pymodule]
|
||||
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
module.add_function(wrap_pyfunction!(ocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(aocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -35,6 +35,8 @@ from litellm.llms.base_llm.ocr.transformation import (
|
|||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"
|
||||
|
||||
|
||||
class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
|
||||
"""
|
||||
|
|
@ -54,6 +56,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
|
|||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def get_api_key_env_var(self) -> Optional[str]:
|
||||
return AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR
|
||||
|
||||
def get_supported_ocr_params(self, model: str) -> list:
|
||||
"""
|
||||
Get supported OCR parameters for Azure Document Intelligence.
|
||||
|
|
@ -156,7 +161,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
|
|||
"""
|
||||
# Get API key from environment if not provided
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_API_KEY")
|
||||
api_key = get_secret_str(AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR)
|
||||
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestDat
|
|||
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
AZURE_AI_OCR_API_KEY_ENV_VAR = "AZURE_AI_API_KEY"
|
||||
|
||||
|
||||
class AzureAIOCRConfig(MistralOCRConfig):
|
||||
"""
|
||||
|
|
@ -30,6 +32,9 @@ class AzureAIOCRConfig(MistralOCRConfig):
|
|||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def get_api_key_env_var(self) -> Optional[str]:
|
||||
return AZURE_AI_OCR_API_KEY_ENV_VAR
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
|
|
@ -46,7 +51,7 @@ class AzureAIOCRConfig(MistralOCRConfig):
|
|||
"""
|
||||
# Get API key from environment if not provided
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("AZURE_AI_API_KEY")
|
||||
api_key = get_secret_str(AZURE_AI_OCR_API_KEY_ENV_VAR)
|
||||
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -101,6 +101,12 @@ class BaseOCRConfig:
|
|||
"""
|
||||
return []
|
||||
|
||||
def get_api_key_env_var(self) -> Optional[str]:
|
||||
"""
|
||||
Return the provider-specific API key environment variable name, if any.
|
||||
"""
|
||||
return None
|
||||
|
||||
def map_ocr_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ from litellm.llms.base_llm.ocr.transformation import (
|
|||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
MISTRAL_OCR_API_KEY_ENV_VAR = "MISTRAL_API_KEY"
|
||||
|
||||
|
||||
class MistralOCRConfig(BaseOCRConfig):
|
||||
"""
|
||||
|
|
@ -59,6 +61,9 @@ class MistralOCRConfig(BaseOCRConfig):
|
|||
"id",
|
||||
]
|
||||
|
||||
def get_api_key_env_var(self) -> Optional[str]:
|
||||
return MISTRAL_OCR_API_KEY_ENV_VAR
|
||||
|
||||
def map_ocr_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
|
|
@ -95,7 +100,7 @@ class MistralOCRConfig(BaseOCRConfig):
|
|||
"""
|
||||
# Get API key from environment if not provided
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("MISTRAL_API_KEY")
|
||||
api_key = get_secret_str(MISTRAL_OCR_API_KEY_ENV_VAR)
|
||||
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,8 @@ from litellm.llms.base_llm.ocr.transformation import (
|
|||
)
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY"
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
else:
|
||||
|
|
@ -28,14 +30,17 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
"""
|
||||
Vertex AI DeepSeek OCR transformation configuration.
|
||||
|
||||
Vertex AI DeepSeek OCR uses the chat completion API format through the openapi endpoint.
|
||||
This transformation converts OCR requests to chat completion format and vice versa.
|
||||
This transformation converts standard LiteLLM OCR requests to the
|
||||
Vertex AI DeepSeek OCR OpenAPI endpoint shape and normalizes the response.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.vertex_base = VertexBase()
|
||||
|
||||
def get_api_key_env_var(self) -> Optional[str]:
|
||||
return VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
|
|
@ -50,6 +55,13 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
|
||||
Vertex AI uses Bearer token authentication with access token from credentials.
|
||||
"""
|
||||
if api_key is not None:
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
**headers,
|
||||
}
|
||||
|
||||
# Extract Vertex AI parameters using safe helpers from VertexBase
|
||||
# Use safe_get_* methods that don't mutate litellm_params dict
|
||||
litellm_params = litellm_params or {}
|
||||
|
|
@ -86,9 +98,6 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
"""
|
||||
Get complete URL for Vertex AI DeepSeek OCR endpoint.
|
||||
|
||||
Vertex AI endpoint format:
|
||||
https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions
|
||||
|
||||
Args:
|
||||
api_base: Vertex AI API base URL (optional)
|
||||
model: Model name (e.g., "deepseek-ai/deepseek-ocr-maas")
|
||||
|
|
@ -123,8 +132,6 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
# Ensure no trailing slash
|
||||
api_base = api_base.rstrip("/")
|
||||
|
||||
# Vertex AI DeepSeek OCR endpoint format
|
||||
# Format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/endpoints/openapi/chat/completions
|
||||
return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions"
|
||||
|
||||
def transform_ocr_request(
|
||||
|
|
@ -136,9 +143,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
**kwargs,
|
||||
) -> OCRRequestData:
|
||||
"""
|
||||
Transform OCR request to chat completion format for Vertex AI DeepSeek OCR.
|
||||
Transform OCR request for Vertex AI DeepSeek OCR.
|
||||
|
||||
Converts OCR document format to chat completion messages format:
|
||||
Converts OCR document format to the Vertex AI DeepSeek OCR payload:
|
||||
- Input: {"type": "image_url", "image_url": "gs://..."}
|
||||
- Output: {"model": "deepseek-ai/deepseek-ocr-maas", "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "gs://..."}]}]}
|
||||
|
||||
|
|
@ -150,7 +157,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
OCRRequestData with JSON data in chat completion format
|
||||
OCRRequestData with JSON data for the DeepSeek OCR endpoint
|
||||
"""
|
||||
verbose_logger.debug(
|
||||
"Vertex AI DeepSeek OCR transform_ocr_request (sync) called"
|
||||
|
|
@ -173,7 +180,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'"
|
||||
)
|
||||
|
||||
# Build chat completion message content
|
||||
# Build DeepSeek OCR message content
|
||||
content_item = {}
|
||||
if image_url:
|
||||
content_item = {"type": "image_url", "image_url": image_url}
|
||||
|
|
@ -181,25 +188,21 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
# For document URLs, we use image_url type as well (Vertex AI supports both)
|
||||
content_item = {"type": "image_url", "image_url": document_url}
|
||||
|
||||
# Build chat completion request
|
||||
# Build DeepSeek OCR request
|
||||
data = {
|
||||
"model": "deepseek-ai/" + model,
|
||||
"messages": [{"role": "user", "content": [content_item]}],
|
||||
}
|
||||
|
||||
# Add optional parameters (stream, temperature, etc.)
|
||||
# Filter out OCR-specific params that don't apply to chat completion
|
||||
chat_completion_params = {}
|
||||
deepseek_ocr_params = {}
|
||||
for key, value in optional_params.items():
|
||||
# Include common chat completion params
|
||||
if key in ["stream", "temperature", "max_tokens", "top_p", "n", "stop"]:
|
||||
chat_completion_params[key] = value
|
||||
deepseek_ocr_params[key] = value
|
||||
|
||||
data.update(chat_completion_params)
|
||||
data.update(deepseek_ocr_params)
|
||||
|
||||
verbose_logger.debug(
|
||||
"Vertex AI DeepSeek OCR: Transformed request to chat completion format"
|
||||
)
|
||||
verbose_logger.debug("Vertex AI DeepSeek OCR: Transformed request")
|
||||
|
||||
return OCRRequestData(data=data, files=None)
|
||||
|
||||
|
|
@ -212,7 +215,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
**kwargs,
|
||||
) -> OCRRequestData:
|
||||
"""
|
||||
Transform OCR request to chat completion format for Vertex AI DeepSeek OCR (async).
|
||||
Transform OCR request for Vertex AI DeepSeek OCR (async).
|
||||
|
||||
Same as sync version - no async-specific logic needed.
|
||||
|
||||
|
|
@ -224,7 +227,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
OCRRequestData with JSON data in chat completion format
|
||||
OCRRequestData with JSON data for the DeepSeek OCR endpoint
|
||||
"""
|
||||
return self.transform_ocr_request(
|
||||
model=model,
|
||||
|
|
@ -242,12 +245,11 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
**kwargs,
|
||||
) -> OCRResponse:
|
||||
"""
|
||||
Transform chat completion response to OCR format.
|
||||
Transform Vertex AI DeepSeek OCR response to OCR format.
|
||||
|
||||
Vertex AI DeepSeek OCR returns chat completion format:
|
||||
Vertex AI DeepSeek OCR returns an OpenAPI response:
|
||||
{
|
||||
"id": "...",
|
||||
"object": "chat.completion",
|
||||
"choices": [{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
|
|
@ -274,16 +276,16 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
try:
|
||||
response_json = raw_response.json()
|
||||
|
||||
# Extract content from chat completion response
|
||||
# Extract OCR content from provider response
|
||||
choices = response_json.get("choices", [])
|
||||
if not choices:
|
||||
raise ValueError("No choices in chat completion response")
|
||||
raise ValueError("No choices in DeepSeek OCR response")
|
||||
|
||||
message = choices[0].get("message", {})
|
||||
content = message.get("content", "")
|
||||
|
||||
if not content:
|
||||
raise ValueError("No content in chat completion response")
|
||||
raise ValueError("No content in DeepSeek OCR response")
|
||||
|
||||
# Try to parse content as JSON (OCR result might be JSON string)
|
||||
ocr_data = None
|
||||
|
|
@ -376,7 +378,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
|
|||
**kwargs,
|
||||
) -> OCRResponse:
|
||||
"""
|
||||
Async transform chat completion response to OCR format.
|
||||
Async transform Vertex AI DeepSeek OCR response to OCR format.
|
||||
|
||||
Same as sync version - no async-specific logic needed.
|
||||
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
|
|||
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
VERTEX_AI_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY"
|
||||
|
||||
|
||||
class VertexAIOCRConfig(MistralOCRConfig):
|
||||
"""
|
||||
|
|
@ -32,6 +34,9 @@ class VertexAIOCRConfig(MistralOCRConfig):
|
|||
super().__init__()
|
||||
self.vertex_base = VertexBase()
|
||||
|
||||
def get_api_key_env_var(self) -> Optional[str]:
|
||||
return VERTEX_AI_OCR_API_KEY_ENV_VAR
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
|
|
@ -46,6 +51,13 @@ class VertexAIOCRConfig(MistralOCRConfig):
|
|||
|
||||
Vertex AI uses Bearer token authentication with access token from credentials.
|
||||
"""
|
||||
if api_key is not None:
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
**headers,
|
||||
}
|
||||
|
||||
# Extract Vertex AI parameters using safe helpers from VertexBase
|
||||
# Use safe_get_* methods that don't mutate litellm_params dict
|
||||
litellm_params = litellm_params or {}
|
||||
|
|
|
|||
|
|
@ -4,13 +4,12 @@ Main OCR function for LiteLLM.
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextvars
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
from functools import partial
|
||||
from dataclasses import dataclass
|
||||
from io import IOBase
|
||||
from typing import Any, Callable, Coroutine, Dict, Optional, Union, cast
|
||||
from typing import Any, Callable, Coroutine, Optional, Union, cast
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -20,7 +19,13 @@ from litellm.constants import request_timeout
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.ocr.rust_bridge import RustOcr, load_rust_ocr, rust_ocr_enabled
|
||||
from litellm.ocr.rust_bridge import (
|
||||
RustAocr,
|
||||
RustOcr,
|
||||
load_rust_aocr,
|
||||
load_rust_ocr,
|
||||
rust_ocr_enabled,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.utils import ProviderConfigManager, client
|
||||
|
||||
|
|
@ -29,6 +34,37 @@ base_llm_http_handler = BaseLLMHTTPHandler()
|
|||
#################################################
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PreparedOCRRequest:
|
||||
model: str
|
||||
document: dict[str, Any]
|
||||
api_key: Optional[str]
|
||||
api_base: Optional[str]
|
||||
custom_llm_provider: str
|
||||
extra_headers: Optional[dict[str, object]]
|
||||
provider_config: BaseOCRConfig
|
||||
optional_params: dict[str, object]
|
||||
litellm_params: dict[str, object]
|
||||
effective_timeout: Union[float, httpx.Timeout]
|
||||
litellm_logging_obj: LiteLLMLoggingObj
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PreparedRustOCRCall:
|
||||
api_key: Optional[str]
|
||||
api_base: Optional[str]
|
||||
headers: dict[str, object]
|
||||
optional_params: dict[str, object]
|
||||
|
||||
|
||||
_RUST_OCR_PROVIDERS = {
|
||||
"mistral",
|
||||
"azure_ai",
|
||||
"azure_ai/doc-intelligence",
|
||||
"vertex_ai",
|
||||
}
|
||||
|
||||
|
||||
def _timeout_to_seconds(
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
) -> Optional[float]:
|
||||
|
|
@ -45,18 +81,206 @@ def _timeout_to_seconds(
|
|||
return float(timeout)
|
||||
|
||||
|
||||
def _run_rust_ocr(
|
||||
rust_ocr: RustOcr,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
provider_config: BaseOCRConfig,
|
||||
resolve_api_key: Callable[[str], Optional[str]],
|
||||
def _prepare_ocr_request(
|
||||
model: str,
|
||||
document: dict[str, object],
|
||||
document: dict[str, Any],
|
||||
api_key: Optional[str],
|
||||
api_base: Optional[str],
|
||||
optional_params: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
timeout_seconds: Optional[float],
|
||||
timeout: Optional[Union[float, httpx.Timeout]],
|
||||
custom_llm_provider: Optional[str],
|
||||
extra_headers: Optional[dict[str, Any]],
|
||||
kwargs: dict[str, Any],
|
||||
) -> _PreparedOCRRequest:
|
||||
litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj"))
|
||||
litellm_call_id = cast(Optional[str], kwargs.get("litellm_call_id", None))
|
||||
|
||||
if not isinstance(document, dict):
|
||||
raise ValueError(
|
||||
f"document must be a dict with 'type' and URL/file field, got {type(document)}"
|
||||
)
|
||||
|
||||
doc_type = document.get("type")
|
||||
|
||||
if doc_type == "file":
|
||||
document = convert_file_document_to_url_document(document)
|
||||
doc_type = document.get("type")
|
||||
|
||||
if doc_type not in ["document_url", "image_url"]:
|
||||
raise ValueError(
|
||||
f"Invalid document type: {doc_type}. "
|
||||
"Must be 'document_url', 'image_url', or 'file'"
|
||||
)
|
||||
|
||||
(
|
||||
model,
|
||||
custom_llm_provider,
|
||||
dynamic_api_key,
|
||||
dynamic_api_base,
|
||||
) = litellm.get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
if dynamic_api_key:
|
||||
api_key = dynamic_api_key
|
||||
if dynamic_api_base:
|
||||
api_base = dynamic_api_base
|
||||
|
||||
ocr_provider_config = ProviderConfigManager.get_provider_ocr_config(
|
||||
model=model,
|
||||
provider=litellm.LlmProviders(custom_llm_provider),
|
||||
)
|
||||
|
||||
if ocr_provider_config is None:
|
||||
raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}")
|
||||
|
||||
verbose_logger.debug(f"OCR call - model: {model}, provider: {custom_llm_provider}")
|
||||
|
||||
litellm_params = GenericLiteLLMParams(**kwargs)
|
||||
|
||||
supported_params = ocr_provider_config.get_supported_ocr_params(model=model)
|
||||
non_default_params = {}
|
||||
for param in supported_params:
|
||||
if param in kwargs:
|
||||
non_default_params[param] = kwargs.pop(param)
|
||||
|
||||
optional_params = ocr_provider_config.map_ocr_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
model=model,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}")
|
||||
|
||||
effective_timeout = timeout or request_timeout
|
||||
|
||||
litellm_logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params={
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"api_base": api_base,
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
return _PreparedOCRRequest(
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=cast(Optional[dict[str, object]], extra_headers),
|
||||
provider_config=ocr_provider_config,
|
||||
optional_params=cast(dict[str, object], optional_params),
|
||||
litellm_params=dict(litellm_params),
|
||||
effective_timeout=effective_timeout,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
|
||||
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
|
||||
|
||||
|
||||
def _rust_bridge_optional_params(
|
||||
prepared_request: _PreparedOCRRequest,
|
||||
resolve_secret: Callable[[str], Optional[str]],
|
||||
) -> dict[str, object]:
|
||||
optional_params = dict(prepared_request.optional_params)
|
||||
if prepared_request.custom_llm_provider == "vertex_ai":
|
||||
vertex_project = (
|
||||
prepared_request.litellm_params.get("vertex_project")
|
||||
or prepared_request.litellm_params.get("vertex_ai_project")
|
||||
or litellm.vertex_project
|
||||
or resolve_secret("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_location = (
|
||||
prepared_request.litellm_params.get("vertex_location")
|
||||
or prepared_request.litellm_params.get("vertex_ai_location")
|
||||
or litellm.vertex_location
|
||||
or resolve_secret("VERTEXAI_LOCATION")
|
||||
or resolve_secret("VERTEX_LOCATION")
|
||||
)
|
||||
if vertex_project is not None:
|
||||
optional_params["vertex_project"] = vertex_project
|
||||
if vertex_location is not None:
|
||||
optional_params["vertex_location"] = vertex_location
|
||||
return optional_params
|
||||
|
||||
|
||||
def _rust_bridge_api_base(
|
||||
prepared_request: _PreparedOCRRequest,
|
||||
resolve_secret: Callable[[str], Optional[str]],
|
||||
) -> Optional[str]:
|
||||
if prepared_request.api_base is not None:
|
||||
return prepared_request.api_base
|
||||
if prepared_request.custom_llm_provider == "azure_ai/doc-intelligence":
|
||||
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
|
||||
if prepared_request.custom_llm_provider == "azure_ai":
|
||||
if (
|
||||
"doc-intelligence" in prepared_request.model
|
||||
or "documentintelligence" in prepared_request.model
|
||||
):
|
||||
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
|
||||
return resolve_secret("AZURE_AI_API_BASE")
|
||||
return None
|
||||
|
||||
|
||||
def _prepare_rust_ocr_call(
|
||||
prepared_request: _PreparedOCRRequest,
|
||||
resolve_api_key: Callable[[str], Optional[str]],
|
||||
) -> _PreparedRustOCRCall:
|
||||
provider_config = prepared_request.provider_config
|
||||
api_key_env_var = provider_config.get_api_key_env_var()
|
||||
resolved_api_key = prepared_request.api_key or (
|
||||
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
|
||||
)
|
||||
resolved_headers = provider_config.validate_environment(
|
||||
headers=prepared_request.extra_headers or {},
|
||||
model=prepared_request.model,
|
||||
api_key=resolved_api_key,
|
||||
api_base=prepared_request.api_base,
|
||||
litellm_params=prepared_request.litellm_params,
|
||||
)
|
||||
resolved_complete_url = provider_config.get_complete_url(
|
||||
api_base=prepared_request.api_base,
|
||||
model=prepared_request.model,
|
||||
optional_params=prepared_request.optional_params,
|
||||
litellm_params=prepared_request.litellm_params,
|
||||
)
|
||||
rust_api_base = _rust_bridge_api_base(prepared_request, resolve_api_key)
|
||||
rust_optional_params = _rust_bridge_optional_params(
|
||||
prepared_request, resolve_api_key
|
||||
)
|
||||
prepared_request.litellm_logging_obj.pre_call(
|
||||
input="OCR document processing",
|
||||
api_key=resolved_api_key,
|
||||
additional_args={
|
||||
"complete_input_dict": {
|
||||
"model": prepared_request.model,
|
||||
"document": prepared_request.document,
|
||||
**rust_optional_params,
|
||||
},
|
||||
"api_base": resolved_complete_url,
|
||||
"headers": resolved_headers,
|
||||
},
|
||||
)
|
||||
return _PreparedRustOCRCall(
|
||||
api_key=resolved_api_key,
|
||||
api_base=rust_api_base,
|
||||
headers=cast(dict[str, object], resolved_headers),
|
||||
optional_params=rust_optional_params,
|
||||
)
|
||||
|
||||
|
||||
def _run_rust_ocr(
|
||||
rust_ocr: RustOcr,
|
||||
prepared_request: _PreparedOCRRequest,
|
||||
resolve_api_key: Callable[[str], Optional[str]],
|
||||
) -> OCRResponse:
|
||||
"""Run the Mistral OCR call through the Rust bridge and wrap the result.
|
||||
|
||||
|
|
@ -66,41 +290,43 @@ def _run_rust_ocr(
|
|||
headers) is mirrored into pre_call so logs match the wire. Dependencies are
|
||||
injected so this stays unit-testable without patching module globals.
|
||||
"""
|
||||
resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY")
|
||||
resolved_headers = provider_config.validate_environment(
|
||||
headers={},
|
||||
model=model,
|
||||
api_key=resolved_api_key,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
resolved_complete_url = provider_config.get_complete_url(
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
logging_obj.pre_call(
|
||||
input="OCR document processing",
|
||||
api_key=resolved_api_key,
|
||||
additional_args={
|
||||
"complete_input_dict": {
|
||||
"model": model,
|
||||
"document": document,
|
||||
**optional_params,
|
||||
},
|
||||
"api_base": resolved_complete_url,
|
||||
"headers": resolved_headers,
|
||||
},
|
||||
prepared = _prepare_rust_ocr_call(
|
||||
prepared_request=prepared_request,
|
||||
resolve_api_key=resolve_api_key,
|
||||
)
|
||||
return OCRResponse.model_validate(
|
||||
rust_ocr(
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=resolved_api_key,
|
||||
api_base=api_base,
|
||||
optional_params=optional_params,
|
||||
timeout_seconds=timeout_seconds,
|
||||
model=prepared_request.model,
|
||||
document=cast(dict[str, object], prepared_request.document),
|
||||
api_key=prepared.api_key,
|
||||
api_base=prepared.api_base,
|
||||
custom_llm_provider=prepared_request.custom_llm_provider,
|
||||
extra_headers=prepared.headers,
|
||||
optional_params=prepared.optional_params,
|
||||
timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _run_rust_aocr(
|
||||
rust_aocr: RustAocr,
|
||||
prepared_request: _PreparedOCRRequest,
|
||||
resolve_api_key: Callable[[str], Optional[str]],
|
||||
) -> OCRResponse:
|
||||
prepared = _prepare_rust_ocr_call(
|
||||
prepared_request=prepared_request,
|
||||
resolve_api_key=resolve_api_key,
|
||||
)
|
||||
return OCRResponse.model_validate(
|
||||
await rust_aocr(
|
||||
model=prepared_request.model,
|
||||
document=cast(dict[str, object], prepared_request.document),
|
||||
api_key=prepared.api_key,
|
||||
api_base=prepared.api_base,
|
||||
custom_llm_provider=prepared_request.custom_llm_provider,
|
||||
extra_headers=prepared.headers,
|
||||
optional_params=prepared.optional_params,
|
||||
timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -108,12 +334,12 @@ def _run_rust_ocr(
|
|||
@client
|
||||
async def aocr(
|
||||
model: str,
|
||||
document: Dict[str, Any],
|
||||
document: dict[str, Any],
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
extra_headers: Optional[Dict[str, Any]] = None,
|
||||
extra_headers: Optional[dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> OCRResponse:
|
||||
"""
|
||||
|
|
@ -174,19 +400,18 @@ async def aocr(
|
|||
)
|
||||
```
|
||||
"""
|
||||
local_vars = locals()
|
||||
completion_kwargs: dict[str, object] = {
|
||||
"model": model,
|
||||
"document": document,
|
||||
"api_key": api_key,
|
||||
"api_base": api_base,
|
||||
"timeout": timeout,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"extra_headers": extra_headers,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
kwargs["aocr"] = True
|
||||
|
||||
# Get custom llm provider
|
||||
if custom_llm_provider is None:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model=model, api_base=api_base
|
||||
)
|
||||
|
||||
func = partial(
|
||||
ocr,
|
||||
prepared = _prepare_ocr_request(
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=api_key,
|
||||
|
|
@ -194,17 +419,47 @@ async def aocr(
|
|||
timeout=timeout,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
**kwargs,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
model = prepared.model
|
||||
custom_llm_provider = prepared.custom_llm_provider
|
||||
completion_kwargs.update(
|
||||
{"model": model, "custom_llm_provider": custom_llm_provider}
|
||||
)
|
||||
|
||||
ctx = contextvars.copy_context()
|
||||
func_with_context = partial(ctx.run, func)
|
||||
init_response = await loop.run_in_executor(None, func_with_context)
|
||||
if _rust_ocr_supported(prepared) and rust_ocr_enabled():
|
||||
rust_aocr = load_rust_aocr()
|
||||
if rust_aocr is None:
|
||||
verbose_logger.debug(
|
||||
"Async Rust OCR bridge unavailable; falling back to Python path"
|
||||
)
|
||||
else:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
if asyncio.iscoroutine(init_response):
|
||||
response = await init_response
|
||||
else:
|
||||
response = init_response
|
||||
response = await _run_rust_aocr(
|
||||
rust_aocr=rust_aocr,
|
||||
prepared_request=prepared,
|
||||
resolve_api_key=get_secret_str,
|
||||
)
|
||||
return response
|
||||
|
||||
response = base_llm_http_handler.ocr(
|
||||
model=prepared.model,
|
||||
document=prepared.document,
|
||||
optional_params=prepared.optional_params,
|
||||
timeout=prepared.effective_timeout,
|
||||
logging_obj=prepared.litellm_logging_obj,
|
||||
api_key=prepared.api_key,
|
||||
api_base=prepared.api_base,
|
||||
custom_llm_provider=prepared.custom_llm_provider,
|
||||
aocr=True,
|
||||
headers=prepared.extra_headers,
|
||||
provider_config=prepared.provider_config,
|
||||
litellm_params=prepared.litellm_params,
|
||||
)
|
||||
|
||||
if asyncio.iscoroutine(response):
|
||||
response = await response
|
||||
|
||||
if response is None:
|
||||
raise ValueError(
|
||||
|
|
@ -217,218 +472,7 @@ async def aocr(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
extra_kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
@client
|
||||
def ocr(
|
||||
model: str,
|
||||
document: Dict[str, Any],
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
extra_headers: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]:
|
||||
"""
|
||||
Synchronous OCR function.
|
||||
|
||||
Args:
|
||||
model: Model name (e.g., "mistral/mistral-ocr-latest")
|
||||
document: Document to process in Mistral format:
|
||||
{"type": "document_url", "document_url": "https://..."} for PDFs/docs,
|
||||
{"type": "image_url", "image_url": "https://..."} for images, or
|
||||
{"type": "file", "file": <path/bytes/file-obj>} for local files
|
||||
api_key: Optional API key
|
||||
api_base: Optional API base URL
|
||||
timeout: Optional timeout
|
||||
custom_llm_provider: Optional custom LLM provider
|
||||
extra_headers: Optional extra headers
|
||||
**kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit)
|
||||
|
||||
Returns:
|
||||
OCRResponse in Mistral OCR format with pages, model, usage_info, etc.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import litellm
|
||||
|
||||
# OCR with PDF
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": "https://arxiv.org/pdf/2201.04234"
|
||||
},
|
||||
include_image_base64=True
|
||||
)
|
||||
|
||||
# OCR with image
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={
|
||||
"type": "image_url",
|
||||
"image_url": "https://example.com/image.png"
|
||||
}
|
||||
)
|
||||
|
||||
# OCR with base64 encoded PDF
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": f"data:application/pdf;base64,{base64_pdf}"
|
||||
}
|
||||
)
|
||||
|
||||
# OCR with local file
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={"type": "file", "file": "/path/to/document.pdf"}
|
||||
)
|
||||
|
||||
# Access pages
|
||||
for page in response.pages:
|
||||
print(f"Page {page.index}: {page.markdown}")
|
||||
```
|
||||
"""
|
||||
local_vars = locals()
|
||||
try:
|
||||
litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj"))
|
||||
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
|
||||
_is_async = kwargs.pop("aocr", False) is True
|
||||
|
||||
# Validate document parameter format
|
||||
if not isinstance(document, dict):
|
||||
raise ValueError(
|
||||
f"document must be a dict with 'type' and URL/file field, got {type(document)}"
|
||||
)
|
||||
|
||||
doc_type = document.get("type")
|
||||
|
||||
# Handle file type: convert to document_url/image_url with base64 data URI
|
||||
if doc_type == "file":
|
||||
document = convert_file_document_to_url_document(document)
|
||||
doc_type = document.get("type")
|
||||
|
||||
if doc_type not in ["document_url", "image_url"]:
|
||||
raise ValueError(
|
||||
f"Invalid document type: {doc_type}. "
|
||||
"Must be 'document_url', 'image_url', or 'file'"
|
||||
)
|
||||
|
||||
(
|
||||
model,
|
||||
custom_llm_provider,
|
||||
dynamic_api_key,
|
||||
dynamic_api_base,
|
||||
) = litellm.get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
# Update with dynamic values if available
|
||||
if dynamic_api_key:
|
||||
api_key = dynamic_api_key
|
||||
if dynamic_api_base:
|
||||
api_base = dynamic_api_base
|
||||
|
||||
ocr_provider_config: Optional[BaseOCRConfig] = (
|
||||
ProviderConfigManager.get_provider_ocr_config(
|
||||
model=model,
|
||||
provider=litellm.LlmProviders(custom_llm_provider),
|
||||
)
|
||||
)
|
||||
|
||||
if ocr_provider_config is None:
|
||||
raise ValueError(
|
||||
f"OCR is not supported for provider: {custom_llm_provider}"
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"OCR call - model: {model}, provider: {custom_llm_provider}"
|
||||
)
|
||||
|
||||
litellm_params = GenericLiteLLMParams(**kwargs)
|
||||
|
||||
supported_params = ocr_provider_config.get_supported_ocr_params(model=model)
|
||||
non_default_params = {}
|
||||
for param in supported_params:
|
||||
if param in kwargs:
|
||||
non_default_params[param] = kwargs.pop(param)
|
||||
|
||||
optional_params = ocr_provider_config.map_ocr_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
model=model,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}")
|
||||
|
||||
effective_timeout = timeout or request_timeout
|
||||
|
||||
litellm_logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params={
|
||||
"litellm_call_id": litellm_call_id,
|
||||
"api_base": api_base,
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
# Optional Rust path: hand the whole Mistral OCR call to the Rust bridge.
|
||||
if custom_llm_provider == "mistral" and rust_ocr_enabled():
|
||||
rust_ocr = load_rust_ocr()
|
||||
if rust_ocr is None:
|
||||
verbose_logger.debug(
|
||||
"Rust OCR bridge unavailable; falling back to Python path"
|
||||
)
|
||||
else:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return _run_rust_ocr(
|
||||
rust_ocr=rust_ocr,
|
||||
logging_obj=litellm_logging_obj,
|
||||
provider_config=ocr_provider_config,
|
||||
resolve_api_key=get_secret_str,
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
optional_params=optional_params,
|
||||
litellm_params=dict(litellm_params),
|
||||
timeout_seconds=_timeout_to_seconds(effective_timeout),
|
||||
)
|
||||
|
||||
response = base_llm_http_handler.ocr(
|
||||
model=model,
|
||||
document=document,
|
||||
optional_params=optional_params,
|
||||
timeout=effective_timeout,
|
||||
logging_obj=litellm_logging_obj,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
aocr=_is_async,
|
||||
headers=extra_headers,
|
||||
provider_config=ocr_provider_config,
|
||||
litellm_params=dict(litellm_params),
|
||||
)
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=local_vars,
|
||||
completion_kwargs=completion_kwargs,
|
||||
extra_kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -466,7 +510,7 @@ def get_mime_type(file_path: str) -> str:
|
|||
return guessed or "application/octet-stream"
|
||||
|
||||
|
||||
def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str, str]:
|
||||
def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]:
|
||||
"""
|
||||
Convert a file-type document dict to a document_url-type document dict
|
||||
with an inline base64 data URI.
|
||||
|
|
@ -550,9 +594,153 @@ def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str,
|
|||
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
|
||||
)
|
||||
return {"type": "image_url", "image_url": data_uri}
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
f"OCR file input: Converted file to document_url data URI "
|
||||
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
|
||||
|
||||
verbose_logger.debug(
|
||||
f"OCR file input: Converted file to document_url data URI "
|
||||
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
|
||||
)
|
||||
return {"type": "document_url", "document_url": data_uri}
|
||||
|
||||
|
||||
@client
|
||||
def ocr(
|
||||
model: str,
|
||||
document: dict[str, Any],
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
extra_headers: Optional[dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]:
|
||||
"""
|
||||
Synchronous OCR function.
|
||||
|
||||
Args:
|
||||
model: Model name (e.g., "mistral/mistral-ocr-latest")
|
||||
document: Document to process in Mistral format:
|
||||
{"type": "document_url", "document_url": "https://..."} for PDFs/docs,
|
||||
{"type": "image_url", "image_url": "https://..."} for images, or
|
||||
{"type": "file", "file": <path/bytes/file-obj>} for local files
|
||||
api_key: Optional API key
|
||||
api_base: Optional API base URL
|
||||
timeout: Optional timeout
|
||||
custom_llm_provider: Optional custom LLM provider
|
||||
extra_headers: Optional extra headers
|
||||
**kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit)
|
||||
|
||||
Returns:
|
||||
OCRResponse in Mistral OCR format with pages, model, usage_info, etc.
|
||||
|
||||
Example:
|
||||
```python
|
||||
import litellm
|
||||
|
||||
# OCR with PDF
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": "https://arxiv.org/pdf/2201.04234"
|
||||
},
|
||||
include_image_base64=True
|
||||
)
|
||||
|
||||
# OCR with image
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={
|
||||
"type": "image_url",
|
||||
"image_url": "https://example.com/image.png"
|
||||
}
|
||||
)
|
||||
|
||||
# OCR with base64 encoded PDF
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": f"data:application/pdf;base64,{base64_pdf}"
|
||||
}
|
||||
)
|
||||
|
||||
# OCR with local file
|
||||
response = litellm.ocr(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
document={"type": "file", "file": "/path/to/document.pdf"}
|
||||
)
|
||||
|
||||
# Access pages
|
||||
for page in response.pages:
|
||||
print(f"Page {page.index}: {page.markdown}")
|
||||
```
|
||||
"""
|
||||
completion_kwargs: dict[str, object] = {
|
||||
"model": model,
|
||||
"document": document,
|
||||
"api_key": api_key,
|
||||
"api_base": api_base,
|
||||
"timeout": timeout,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"extra_headers": extra_headers,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
try:
|
||||
_is_async = kwargs.pop("aocr", False) is True
|
||||
completion_kwargs["aocr"] = _is_async
|
||||
prepared = _prepare_ocr_request(
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
kwargs=kwargs,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
model = prepared.model
|
||||
custom_llm_provider = prepared.custom_llm_provider
|
||||
completion_kwargs.update(
|
||||
{"model": model, "custom_llm_provider": custom_llm_provider}
|
||||
)
|
||||
|
||||
# Optional Rust path: hand supported OCR calls to the Rust bridge.
|
||||
if _rust_ocr_supported(prepared) and rust_ocr_enabled():
|
||||
rust_ocr = load_rust_ocr()
|
||||
if rust_ocr is None:
|
||||
verbose_logger.debug(
|
||||
"Rust OCR bridge unavailable; falling back to Python path"
|
||||
)
|
||||
else:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return _run_rust_ocr(
|
||||
rust_ocr=rust_ocr,
|
||||
prepared_request=prepared,
|
||||
resolve_api_key=get_secret_str,
|
||||
)
|
||||
|
||||
response = base_llm_http_handler.ocr(
|
||||
model=prepared.model,
|
||||
document=prepared.document,
|
||||
optional_params=prepared.optional_params,
|
||||
timeout=prepared.effective_timeout,
|
||||
logging_obj=prepared.litellm_logging_obj,
|
||||
api_key=prepared.api_key,
|
||||
api_base=prepared.api_base,
|
||||
custom_llm_provider=prepared.custom_llm_provider,
|
||||
aocr=_is_async,
|
||||
headers=prepared.extra_headers,
|
||||
provider_config=prepared.provider_config,
|
||||
litellm_params=prepared.litellm_params,
|
||||
)
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs=completion_kwargs,
|
||||
extra_kwargs=kwargs,
|
||||
)
|
||||
return {"type": "document_url", "document_url": data_uri}
|
||||
|
|
|
|||
|
|
@ -11,7 +11,8 @@ can import it statically without forming an import cycle.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final, Protocol, cast
|
||||
import os
|
||||
from typing import Awaitable, Final, Protocol, cast
|
||||
|
||||
|
||||
class RustOcr(Protocol):
|
||||
|
|
@ -23,9 +24,29 @@ class RustOcr(Protocol):
|
|||
document: dict[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout_seconds: float | None,
|
||||
) -> dict[str, object]: ...
|
||||
) -> dict[str, object]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class RustAocr(Protocol):
|
||||
"""Signature of the compiled ``litellm_python_bridge.aocr`` entrypoint."""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
model: str,
|
||||
document: dict[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout_seconds: float | None,
|
||||
) -> Awaitable[dict[str, object]]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class _Unset:
|
||||
|
|
@ -34,23 +55,39 @@ class _Unset:
|
|||
|
||||
_UNSET: Final[_Unset] = _Unset()
|
||||
|
||||
_rust_ocr_enabled = False
|
||||
|
||||
def _env_enables_rust_ocr() -> bool:
|
||||
return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
|
||||
|
||||
_rust_ocr_enabled = _env_enables_rust_ocr()
|
||||
_rust_ocr_impl: RustOcr | None = None
|
||||
_rust_aocr_impl: RustAocr | None = None
|
||||
|
||||
|
||||
def use_litellm_rust(
|
||||
enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET
|
||||
enabled: bool = True,
|
||||
*,
|
||||
ocr: RustOcr | None | _Unset = _UNSET,
|
||||
aocr: RustAocr | None | _Unset = _UNSET,
|
||||
) -> None:
|
||||
"""Route supported OCR calls through the packaged Rust extension.
|
||||
|
||||
``ocr`` injects the bridge callable; when omitted the compiled extension is
|
||||
loaded on demand and any previously injected bridge is preserved. Pass
|
||||
``ocr=None`` explicitly to clear a prior injection.
|
||||
``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled
|
||||
extension is loaded on demand and any previously injected bridge is
|
||||
preserved. Pass ``None`` explicitly to clear a prior injection.
|
||||
"""
|
||||
global _rust_ocr_enabled, _rust_ocr_impl
|
||||
global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl
|
||||
_rust_ocr_enabled = enabled
|
||||
if not isinstance(ocr, _Unset):
|
||||
_rust_ocr_impl = ocr
|
||||
if not isinstance(aocr, _Unset):
|
||||
_rust_aocr_impl = aocr
|
||||
|
||||
|
||||
def rust_ocr_enabled() -> bool:
|
||||
|
|
@ -73,3 +110,15 @@ def load_rust_ocr() -> RustOcr | None:
|
|||
if native_bridge is None:
|
||||
return None
|
||||
return cast(RustOcr, native_bridge.ocr)
|
||||
|
||||
|
||||
def load_rust_aocr() -> RustAocr | None:
|
||||
"""Return the async Rust OCR callable, or ``None`` when unavailable."""
|
||||
if _rust_aocr_impl is not None:
|
||||
return _rust_aocr_impl
|
||||
from litellm.rust_bridge import get_native_bridge
|
||||
|
||||
native_bridge = get_native_bridge()
|
||||
if native_bridge is None:
|
||||
return None
|
||||
return cast(RustAocr, getattr(native_bridge, "aocr", None))
|
||||
|
|
|
|||
|
|
@ -24,6 +24,12 @@ EXCLUDED_GUARD_ONLY_VARS = {
|
|||
"MAVVRIK_FOCUS_FREQUENCY",
|
||||
}
|
||||
|
||||
# Temporary/internal rollout flags are intentionally not added to the public
|
||||
# environment settings docs until the feature is ready for broad use.
|
||||
EXCLUDED_ROLLOUT_FLAGS = {
|
||||
"LITELLM_USE_RUST_OCR",
|
||||
}
|
||||
|
||||
EXCLUDED_TERMINAL_VARS = {
|
||||
"TERM",
|
||||
"TERM_PROGRAM",
|
||||
|
|
@ -71,6 +77,7 @@ for root, dirs, files in os.walk(repo_base):
|
|||
for match in getenv_matches
|
||||
if match not in EXCLUDED_TERMINAL_VARS
|
||||
and match not in EXCLUDED_GUARD_ONLY_VARS
|
||||
and match not in EXCLUDED_ROLLOUT_FLAGS
|
||||
) # Extract only the key part, excluding terminal vars
|
||||
|
||||
# Find all keys using litellm.get_secret()
|
||||
|
|
|
|||
|
|
@ -140,6 +140,35 @@ model_list:
|
|||
model_info:
|
||||
mode: realtime
|
||||
|
||||
- model_name: rust-ocr-mistral
|
||||
litellm_params:
|
||||
model: mistral/mistral-ocr-latest
|
||||
api_key: os.environ/MISTRAL_API_KEY
|
||||
|
||||
- model_name: rust-ocr-azure-ai
|
||||
litellm_params:
|
||||
model: azure_ai/mistral-document-ai-2505
|
||||
api_base: os.environ/AZURE_API_BASE
|
||||
api_key: os.environ/AZURE_API_KEY
|
||||
|
||||
- model_name: rust-ocr-azure-document-intelligence
|
||||
litellm_params:
|
||||
model: azure_ai/doc-intelligence/prebuilt-layout
|
||||
api_base: os.environ/AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT
|
||||
api_key: os.environ/AZURE_DOCUMENT_INTELLIGENCE_API_KEY
|
||||
|
||||
- model_name: rust-ocr-vertex-mistral
|
||||
litellm_params:
|
||||
model: vertex_ai/mistral-ocr-2505
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: us-central1
|
||||
|
||||
- model_name: rust-ocr-vertex-deepseek
|
||||
litellm_params:
|
||||
model: vertex_ai/deepseek-ocr-maas
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: us-central1
|
||||
|
||||
|
||||
mcp_servers:
|
||||
deepwiki_mcp:
|
||||
|
|
@ -167,4 +196,3 @@ guardrails:
|
|||
US_SSN: BLOCK
|
||||
PHONE_NUMBER: BLOCK
|
||||
|
||||
|
||||
|
|
|
|||
148
tests/e2e/gateway/test_ocr_rust_e2e.py
Normal file
148
tests/e2e/gateway/test_ocr_rust_e2e.py
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
"""
|
||||
Gateway E2E smoke for Rust-backed OCR.
|
||||
|
||||
Start the proxy with:
|
||||
|
||||
LITELLM_USE_RUST_OCR=1 litellm --config tests/e2e/gateway/litellm-config.yml --port 4000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
TEST_PDF_URL = (
|
||||
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
|
||||
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
|
||||
"/tests/llm_translation/fixtures/dummy.pdf"
|
||||
)
|
||||
TEST_IMAGE_URL = (
|
||||
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
|
||||
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
|
||||
"/tests/image_gen_tests/test_image.png"
|
||||
)
|
||||
|
||||
RUST_OCR_GATEWAY_CASES = [
|
||||
pytest.param(
|
||||
"rust-ocr-mistral",
|
||||
{"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
id="mistral",
|
||||
),
|
||||
pytest.param(
|
||||
"rust-ocr-azure-ai",
|
||||
{"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
id="azure_ai",
|
||||
),
|
||||
pytest.param(
|
||||
"rust-ocr-azure-document-intelligence",
|
||||
{"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
id="azure_document_intelligence",
|
||||
),
|
||||
pytest.param(
|
||||
"rust-ocr-vertex-mistral",
|
||||
{"type": "document_url", "document_url": TEST_PDF_URL},
|
||||
id="vertex_mistral",
|
||||
),
|
||||
pytest.param(
|
||||
"rust-ocr-vertex-deepseek",
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": os.getenv("RUST_OCR_IMAGE_URL", TEST_IMAGE_URL),
|
||||
},
|
||||
id="vertex_deepseek",
|
||||
),
|
||||
]
|
||||
|
||||
CONFIG_PATH = Path(__file__).with_name("litellm-config.yml")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OcrGateway:
|
||||
base_url: str
|
||||
master_key: str
|
||||
|
||||
def model_names(self) -> set[str]:
|
||||
with httpx.Client(
|
||||
timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120"))
|
||||
) as client:
|
||||
response = client.get(
|
||||
f"{self.base_url.rstrip('/')}/model/info",
|
||||
headers={"Authorization": f"Bearer {self.master_key}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return {
|
||||
model["model_name"]
|
||||
for model in response.json().get("data", [])
|
||||
if "model_name" in model
|
||||
}
|
||||
|
||||
def ocr(self, model: str, document: dict[str, str]) -> httpx.Response:
|
||||
with httpx.Client(
|
||||
timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120"))
|
||||
) as client:
|
||||
return client.post(
|
||||
f"{self.base_url.rstrip('/')}/v1/ocr",
|
||||
headers={"Authorization": f"Bearer {self.master_key}"},
|
||||
json={"model": model, "document": document},
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OcrResources:
|
||||
gateway: OcrGateway
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def resources() -> OcrResources:
|
||||
proxy_url = os.getenv("LITELLM_PROXY_URL")
|
||||
if not proxy_url:
|
||||
pytest.skip(
|
||||
"Start a Rust OCR proxy and set LITELLM_PROXY_URL, e.g. http://localhost:4000"
|
||||
)
|
||||
return OcrResources(
|
||||
gateway=OcrGateway(
|
||||
base_url=proxy_url,
|
||||
master_key=os.getenv("LITELLM_MASTER_KEY", "sk-1234"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _assert_ocr_response_shape(response_json: dict[str, Any]) -> None:
|
||||
assert response_json["object"] == "ocr"
|
||||
assert response_json["model"]
|
||||
assert isinstance(response_json["pages"], list)
|
||||
assert len(response_json["pages"]) > 0
|
||||
assert "index" in response_json["pages"][0]
|
||||
assert "markdown" in response_json["pages"][0]
|
||||
|
||||
|
||||
class TestRustOcrGateway:
|
||||
def test_rust_ocr_models_are_on_gateway_config(self) -> None:
|
||||
config = yaml.safe_load(CONFIG_PATH.read_text())
|
||||
configured_models = {
|
||||
model_config["model_name"] for model_config in config["model_list"]
|
||||
}
|
||||
|
||||
expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES}
|
||||
assert expected_models.issubset(configured_models)
|
||||
|
||||
def test_running_gateway_loaded_rust_ocr_models(
|
||||
self, resources: OcrResources
|
||||
) -> None:
|
||||
expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES}
|
||||
assert expected_models.issubset(resources.gateway.model_names())
|
||||
|
||||
@pytest.mark.parametrize(("model", "document"), RUST_OCR_GATEWAY_CASES)
|
||||
def test_rust_ocr_model_gateway_response(
|
||||
self, resources: OcrResources, model: str, document: dict[str, str]
|
||||
) -> None:
|
||||
response = resources.gateway.ocr(model, document)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
_assert_ocr_response_shape(response.json())
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
import importlib
|
||||
import builtins
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -18,9 +19,12 @@ rust_bridge = importlib.import_module("litellm.ocr.rust_bridge")
|
|||
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
|
||||
|
||||
MODEL = "mistral/mistral-ocr-latest"
|
||||
DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
|
||||
DOCUMENT: dict[str, object] = {
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf",
|
||||
}
|
||||
|
||||
FAKE_OCR_RESPONSE = {
|
||||
FAKE_OCR_RESPONSE: dict[str, object] = {
|
||||
"pages": [{"index": 0, "markdown": "hello world"}],
|
||||
"model": "mistral-ocr-2505-completion",
|
||||
"document_annotation": None,
|
||||
|
|
@ -29,21 +33,35 @@ FAKE_OCR_RESPONSE = {
|
|||
}
|
||||
|
||||
|
||||
class CapturedException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class RecordingBridge:
|
||||
"""A fake ``RustOcr`` callable that records the args it was handed."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
|
||||
def __call__(
|
||||
self, model, document, api_key, api_base, optional_params, timeout_seconds
|
||||
):
|
||||
self,
|
||||
model: str,
|
||||
document: dict[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout_seconds: float | None,
|
||||
) -> dict[str, object]:
|
||||
self.calls.append(
|
||||
{
|
||||
"model": model,
|
||||
"document": document,
|
||||
"api_key": api_key,
|
||||
"api_base": api_base,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"extra_headers": extra_headers,
|
||||
"optional_params": optional_params,
|
||||
"timeout_seconds": timeout_seconds,
|
||||
}
|
||||
|
|
@ -51,13 +69,81 @@ class RecordingBridge:
|
|||
return dict(FAKE_OCR_RESPONSE)
|
||||
|
||||
|
||||
class RecordingAsyncBridge:
|
||||
"""A fake async ``RustAocr`` callable that records the args it was handed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, object]] = []
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
model: str,
|
||||
document: dict[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout_seconds: float | None,
|
||||
) -> dict[str, object]:
|
||||
self.calls.append(
|
||||
{
|
||||
"model": model,
|
||||
"document": document,
|
||||
"api_key": api_key,
|
||||
"api_base": api_base,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"extra_headers": extra_headers,
|
||||
"optional_params": optional_params,
|
||||
"timeout_seconds": timeout_seconds,
|
||||
}
|
||||
)
|
||||
return dict(FAKE_OCR_RESPONSE)
|
||||
|
||||
|
||||
class RaisingBridge:
|
||||
def __call__(
|
||||
self,
|
||||
model: str,
|
||||
document: dict[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout_seconds: float | None,
|
||||
) -> dict[str, object]:
|
||||
raise RuntimeError("bridge failed")
|
||||
|
||||
|
||||
class RaisingAsyncBridge:
|
||||
async def __call__(
|
||||
self,
|
||||
model: str,
|
||||
document: dict[str, object],
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
custom_llm_provider: str,
|
||||
extra_headers: dict[str, object] | None,
|
||||
optional_params: dict[str, object],
|
||||
timeout_seconds: float | None,
|
||||
) -> dict[str, object]:
|
||||
raise RuntimeError("bridge failed")
|
||||
|
||||
|
||||
class RecordingLogging:
|
||||
"""A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``."""
|
||||
|
||||
def __init__(self):
|
||||
self.pre_call_kwargs = None
|
||||
def __init__(self) -> None:
|
||||
self.pre_call_kwargs: dict[str, object] | None = None
|
||||
|
||||
def pre_call(self, *, input, api_key, additional_args):
|
||||
def pre_call(
|
||||
self,
|
||||
*,
|
||||
input: str,
|
||||
api_key: str | None,
|
||||
additional_args: dict[str, object],
|
||||
) -> None:
|
||||
self.pre_call_kwargs = {
|
||||
"input": input,
|
||||
"api_key": api_key,
|
||||
|
|
@ -68,22 +154,70 @@ class RecordingLogging:
|
|||
class FakeOCRConfig:
|
||||
"""A stand-in ``BaseOCRConfig`` that echoes the request it would build."""
|
||||
|
||||
def validate_environment(
|
||||
self, *, headers, model, api_key, api_base, litellm_params
|
||||
):
|
||||
return {"authorization": f"Bearer {api_key}"}
|
||||
def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None:
|
||||
self.api_key_env_var = api_key_env_var
|
||||
|
||||
def get_complete_url(self, *, api_base, model, optional_params, litellm_params):
|
||||
def get_api_key_env_var(self) -> str:
|
||||
return self.api_key_env_var
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
*,
|
||||
headers: dict[str, object],
|
||||
model: str,
|
||||
api_key: str | None,
|
||||
api_base: str | None,
|
||||
litellm_params: dict[str, object],
|
||||
) -> dict[str, object]:
|
||||
return {"Authorization": f"Bearer {api_key}", **headers}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
*,
|
||||
api_base: str | None,
|
||||
model: str,
|
||||
optional_params: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
) -> str:
|
||||
return f"{api_base or 'https://api.mistral.ai/v1'}/ocr"
|
||||
|
||||
|
||||
def build_prepared_request(
|
||||
*,
|
||||
logging_obj: RecordingLogging | None = None,
|
||||
provider_config: FakeOCRConfig | None = None,
|
||||
model: str = "mistral-ocr-latest",
|
||||
document: dict[str, object] = DOCUMENT,
|
||||
api_key: str | None = "sk-test",
|
||||
api_base: str | None = None,
|
||||
custom_llm_provider: str = "mistral",
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
optional_params: dict[str, object] | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = 12.5,
|
||||
) -> Any:
|
||||
return ocr_main._PreparedOCRRequest(
|
||||
model=model,
|
||||
document=document,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
provider_config=provider_config or FakeOCRConfig(),
|
||||
optional_params=optional_params or {},
|
||||
litellm_params=litellm_params or {},
|
||||
effective_timeout=timeout,
|
||||
litellm_logging_obj=logging_obj or RecordingLogging(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_rust_flag():
|
||||
"""Keep the global toggle isolated between tests."""
|
||||
rust_bridge.use_litellm_rust(False, ocr=None)
|
||||
rust_bridge.use_litellm_rust(False, ocr=None, aocr=None)
|
||||
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
|
||||
yield
|
||||
rust_bridge.use_litellm_rust(False, ocr=None)
|
||||
rust_bridge.use_litellm_rust(False, ocr=None, aocr=None)
|
||||
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
|
||||
|
||||
|
||||
|
|
@ -95,6 +229,14 @@ def fake_bridge():
|
|||
return bridge
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_async_bridge():
|
||||
"""Enable the async Rust path with an injected recording bridge."""
|
||||
bridge = RecordingAsyncBridge()
|
||||
litellm.use_litellm_rust(True, aocr=bridge)
|
||||
return bridge
|
||||
|
||||
|
||||
def test_use_litellm_rust_toggles_flag():
|
||||
assert rust_bridge.rust_ocr_enabled() is False
|
||||
litellm.use_litellm_rust()
|
||||
|
|
@ -103,6 +245,11 @@ def test_use_litellm_rust_toggles_flag():
|
|||
assert rust_bridge.rust_ocr_enabled() is False
|
||||
|
||||
|
||||
def test_env_var_enables_rust_ocr(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
|
||||
assert rust_bridge._env_enables_rust_ocr() is True
|
||||
|
||||
|
||||
def test_load_rust_ocr_returns_injected_impl():
|
||||
bridge = RecordingBridge()
|
||||
litellm.use_litellm_rust(True, ocr=bridge)
|
||||
|
|
@ -147,6 +294,12 @@ def test_native_bridge_available_reflects_loader(monkeypatch):
|
|||
assert rust_bridge_loader.native_bridge_available() is True
|
||||
|
||||
|
||||
def test_load_rust_aocr_returns_injected_impl():
|
||||
bridge = RecordingAsyncBridge()
|
||||
litellm.use_litellm_rust(True, aocr=bridge)
|
||||
assert rust_bridge.load_rust_aocr() is bridge
|
||||
|
||||
|
||||
def test_toggle_without_ocr_arg_preserves_injected_impl():
|
||||
"""Regression: routine enable/disable calls must not clobber a prior injection.
|
||||
|
||||
|
|
@ -155,12 +308,15 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
|
|||
a caller toggled the flag without re-passing ``ocr=``.
|
||||
"""
|
||||
bridge = RecordingBridge()
|
||||
litellm.use_litellm_rust(True, ocr=bridge)
|
||||
async_bridge = RecordingAsyncBridge()
|
||||
litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge)
|
||||
|
||||
litellm.use_litellm_rust(False)
|
||||
assert rust_bridge.load_rust_ocr() is bridge
|
||||
assert rust_bridge.load_rust_aocr() is async_bridge
|
||||
litellm.use_litellm_rust(True)
|
||||
assert rust_bridge.load_rust_ocr() is bridge
|
||||
assert rust_bridge.load_rust_aocr() is async_bridge
|
||||
|
||||
|
||||
def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
|
||||
|
|
@ -170,10 +326,12 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
|
|||
lambda: None,
|
||||
)
|
||||
bridge = RecordingBridge()
|
||||
litellm.use_litellm_rust(True, ocr=bridge)
|
||||
async_bridge = RecordingAsyncBridge()
|
||||
litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge)
|
||||
|
||||
litellm.use_litellm_rust(True, ocr=None)
|
||||
litellm.use_litellm_rust(True, ocr=None, aocr=None)
|
||||
assert rust_bridge.load_rust_ocr() is None
|
||||
assert rust_bridge.load_rust_aocr() is None
|
||||
|
||||
|
||||
def test_load_rust_ocr_none_when_extension_absent(monkeypatch):
|
||||
|
|
@ -186,6 +344,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch):
|
|||
)
|
||||
litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI
|
||||
assert rust_bridge.load_rust_ocr() is None
|
||||
assert rust_bridge.load_rust_aocr() is None
|
||||
|
||||
|
||||
def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
|
||||
|
|
@ -194,6 +353,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
|
|||
built in CI, so stand in a fake module via the bridge loader."""
|
||||
fake_module = types.ModuleType("litellm.rust_bridge._native")
|
||||
fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
|
||||
fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
|
||||
monkeypatch.setattr(
|
||||
importlib.import_module("litellm.rust_bridge"),
|
||||
"get_native_bridge",
|
||||
|
|
@ -202,6 +362,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
|
|||
|
||||
litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension
|
||||
assert rust_bridge.load_rust_ocr() is fake_module.ocr
|
||||
assert rust_bridge.load_rust_aocr() is fake_module.aocr
|
||||
|
||||
|
||||
def test_timeout_to_seconds_handles_float_timeout_and_none():
|
||||
|
|
@ -216,16 +377,14 @@ def test_run_rust_ocr_forwards_args_and_wraps_response():
|
|||
|
||||
response = ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
logging_obj=logging_obj,
|
||||
provider_config=FakeOCRConfig(),
|
||||
prepared_request=build_prepared_request(
|
||||
logging_obj=logging_obj,
|
||||
api_base="https://proxy.internal",
|
||||
extra_headers={"x-trace-id": "trace-1"},
|
||||
optional_params={"include_image_base64": True},
|
||||
timeout=12.5,
|
||||
),
|
||||
resolve_api_key=lambda _name: None,
|
||||
model="mistral-ocr-latest",
|
||||
document=DOCUMENT,
|
||||
api_key="sk-test",
|
||||
api_base="https://proxy.internal",
|
||||
optional_params={"include_image_base64": True},
|
||||
litellm_params={},
|
||||
timeout_seconds=12.5,
|
||||
)
|
||||
|
||||
assert isinstance(response, OCRResponse)
|
||||
|
|
@ -236,6 +395,11 @@ def test_run_rust_ocr_forwards_args_and_wraps_response():
|
|||
"document": DOCUMENT,
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://proxy.internal",
|
||||
"custom_llm_provider": "mistral",
|
||||
"extra_headers": {
|
||||
"Authorization": "Bearer sk-test",
|
||||
"x-trace-id": "trace-1",
|
||||
},
|
||||
"optional_params": {"include_image_base64": True},
|
||||
"timeout_seconds": 12.5,
|
||||
}
|
||||
|
|
@ -248,23 +412,127 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
|
|||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
logging_obj=RecordingLogging(),
|
||||
provider_config=FakeOCRConfig(),
|
||||
prepared_request=build_prepared_request(api_key=None, timeout=None),
|
||||
resolve_api_key=lambda name: (
|
||||
"sk-from-vault" if name == "MISTRAL_API_KEY" else None
|
||||
),
|
||||
model="mistral-ocr-latest",
|
||||
document=DOCUMENT,
|
||||
api_key=None,
|
||||
api_base=None,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
timeout_seconds=None,
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["api_key"] == "sk-from-vault"
|
||||
|
||||
|
||||
def test_run_rust_ocr_uses_provider_api_key_env_var():
|
||||
bridge = RecordingBridge()
|
||||
resolver_calls = []
|
||||
|
||||
def _resolver(name):
|
||||
resolver_calls.append(name)
|
||||
return "sk-provider-env"
|
||||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
prepared_request=build_prepared_request(
|
||||
provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"),
|
||||
model="provider-ocr-model",
|
||||
api_key=None,
|
||||
timeout=None,
|
||||
),
|
||||
resolve_api_key=_resolver,
|
||||
)
|
||||
|
||||
assert resolver_calls == ["PROVIDER_OCR_API_KEY"]
|
||||
assert bridge.calls[0]["api_key"] == "sk-provider-env"
|
||||
|
||||
|
||||
def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
|
||||
bridge = RecordingBridge()
|
||||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="vertex_ai",
|
||||
model="mistral-ocr-maas",
|
||||
litellm_params={
|
||||
"vertex_project": "project-1",
|
||||
"vertex_location": "us-central1",
|
||||
"vertex_credentials": "redacted",
|
||||
},
|
||||
optional_params={"include_image_base64": True},
|
||||
timeout=None,
|
||||
),
|
||||
resolve_api_key=lambda _name: None,
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["optional_params"] == {
|
||||
"include_image_base64": True,
|
||||
"vertex_project": "project-1",
|
||||
"vertex_location": "us-central1",
|
||||
}
|
||||
|
||||
|
||||
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
|
||||
bridge = RecordingBridge()
|
||||
|
||||
def _resolver(name: str) -> str | None:
|
||||
return {
|
||||
"VERTEXAI_PROJECT": "project-from-secret",
|
||||
"VERTEXAI_LOCATION": "us-east5",
|
||||
}.get(name)
|
||||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="vertex_ai",
|
||||
model="mistral-ocr-maas",
|
||||
timeout=None,
|
||||
),
|
||||
resolve_api_key=_resolver,
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret"
|
||||
assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5"
|
||||
|
||||
|
||||
def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
|
||||
bridge = RecordingBridge()
|
||||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="azure_ai",
|
||||
model="pixtral-12b-2409",
|
||||
api_base=None,
|
||||
timeout=None,
|
||||
),
|
||||
resolve_api_key=lambda name: (
|
||||
"https://azure.example.com" if name == "AZURE_AI_API_BASE" else None
|
||||
),
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["api_base"] == "https://azure.example.com"
|
||||
|
||||
|
||||
def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
|
||||
bridge = RecordingBridge()
|
||||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
prepared_request=build_prepared_request(
|
||||
custom_llm_provider="azure_ai/doc-intelligence",
|
||||
model="prebuilt-layout",
|
||||
api_base=None,
|
||||
timeout=None,
|
||||
),
|
||||
resolve_api_key=lambda name: (
|
||||
"https://document-intelligence.example.com"
|
||||
if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com"
|
||||
|
||||
|
||||
def test_run_rust_ocr_prefers_explicit_key_over_resolver():
|
||||
bridge = RecordingBridge()
|
||||
resolver_calls = []
|
||||
|
|
@ -275,16 +543,8 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver():
|
|||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=bridge,
|
||||
logging_obj=RecordingLogging(),
|
||||
provider_config=FakeOCRConfig(),
|
||||
prepared_request=build_prepared_request(api_key="sk-explicit", timeout=None),
|
||||
resolve_api_key=_resolver,
|
||||
model="mistral-ocr-latest",
|
||||
document=DOCUMENT,
|
||||
api_key="sk-explicit",
|
||||
api_base=None,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
timeout_seconds=None,
|
||||
)
|
||||
|
||||
assert bridge.calls[0]["api_key"] == "sk-explicit"
|
||||
|
|
@ -297,16 +557,14 @@ def test_run_rust_ocr_runs_pre_call_logging():
|
|||
|
||||
ocr_main._run_rust_ocr(
|
||||
rust_ocr=RecordingBridge(),
|
||||
logging_obj=logging_obj,
|
||||
provider_config=FakeOCRConfig(),
|
||||
prepared_request=build_prepared_request(
|
||||
logging_obj=logging_obj,
|
||||
api_base="https://api.mistral.ai/v1",
|
||||
extra_headers={"x-trace-id": "trace-1"},
|
||||
optional_params={"include_image_base64": True},
|
||||
timeout=None,
|
||||
),
|
||||
resolve_api_key=lambda _name: None,
|
||||
model="mistral-ocr-latest",
|
||||
document=DOCUMENT,
|
||||
api_key="sk-test",
|
||||
api_base="https://api.mistral.ai/v1",
|
||||
optional_params={"include_image_base64": True},
|
||||
litellm_params={},
|
||||
timeout_seconds=None,
|
||||
)
|
||||
|
||||
assert logging_obj.pre_call_kwargs is not None
|
||||
|
|
@ -317,7 +575,10 @@ def test_run_rust_ocr_runs_pre_call_logging():
|
|||
assert complete_input["include_image_base64"] is True
|
||||
# The logged request mirrors what Rust sends: resolved URL + headers.
|
||||
assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr"
|
||||
assert additional_args["headers"] == {"authorization": "Bearer sk-test"}
|
||||
assert additional_args["headers"] == {
|
||||
"Authorization": "Bearer sk-test",
|
||||
"x-trace-id": "trace-1",
|
||||
}
|
||||
|
||||
|
||||
def test_ocr_routes_to_rust_when_enabled(fake_bridge):
|
||||
|
|
@ -325,6 +586,7 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge):
|
|||
model=MODEL,
|
||||
document=DOCUMENT,
|
||||
api_key="sk-test",
|
||||
extra_headers={"x-trace-id": "trace-1"},
|
||||
include_image_base64=True,
|
||||
)
|
||||
|
||||
|
|
@ -336,10 +598,93 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge):
|
|||
assert call["model"] == "mistral-ocr-latest"
|
||||
assert call["document"] == DOCUMENT
|
||||
assert call["api_key"] == "sk-test"
|
||||
assert call["custom_llm_provider"] == "mistral"
|
||||
assert call["extra_headers"] == {
|
||||
"Authorization": "Bearer sk-test",
|
||||
"x-trace-id": "trace-1",
|
||||
}
|
||||
# Raw OCR params ride along in optional_params; Rust filters to supported keys.
|
||||
assert call["optional_params"].get("include_image_base64") is True
|
||||
|
||||
|
||||
def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge):
|
||||
response = litellm.ocr(
|
||||
model="azure_ai/pixtral-12b-2409",
|
||||
document=DOCUMENT,
|
||||
api_key="sk-test",
|
||||
api_base="https://example.services.ai.azure.com",
|
||||
)
|
||||
|
||||
assert isinstance(response, OCRResponse)
|
||||
assert len(fake_bridge.calls) == 1
|
||||
assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409"
|
||||
assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai"
|
||||
|
||||
|
||||
def test_ocr_exception_type_uses_resolved_provider_context(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_exception_type(**kwargs: object) -> CapturedException:
|
||||
captured.update(kwargs)
|
||||
return CapturedException("wrapped")
|
||||
|
||||
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
|
||||
litellm.use_litellm_rust(True, ocr=RaisingBridge())
|
||||
|
||||
with pytest.raises(CapturedException):
|
||||
litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
|
||||
|
||||
assert captured["model"] == "mistral-ocr-latest"
|
||||
assert captured["custom_llm_provider"] == "mistral"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge):
|
||||
response = await litellm.aocr(
|
||||
model=MODEL,
|
||||
document=DOCUMENT,
|
||||
api_key="sk-test",
|
||||
extra_headers={"x-trace-id": "trace-1"},
|
||||
include_image_base64=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, OCRResponse)
|
||||
assert response.pages[0].markdown == "hello world"
|
||||
assert len(fake_async_bridge.calls) == 1
|
||||
call = fake_async_bridge.calls[0]
|
||||
assert call["model"] == "mistral-ocr-latest"
|
||||
assert call["document"] == DOCUMENT
|
||||
assert call["api_key"] == "sk-test"
|
||||
assert call["custom_llm_provider"] == "mistral"
|
||||
assert call["extra_headers"] == {
|
||||
"Authorization": "Bearer sk-test",
|
||||
"x-trace-id": "trace-1",
|
||||
}
|
||||
assert call["optional_params"].get("include_image_base64") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aocr_exception_type_uses_resolved_provider_context(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_exception_type(**kwargs: object) -> CapturedException:
|
||||
captured.update(kwargs)
|
||||
return CapturedException("wrapped")
|
||||
|
||||
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
|
||||
litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge())
|
||||
|
||||
with pytest.raises(CapturedException):
|
||||
await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
|
||||
|
||||
assert captured["model"] == "mistral-ocr-latest"
|
||||
assert captured["custom_llm_provider"] == "mistral"
|
||||
|
||||
|
||||
def test_ocr_forwards_timeout_to_rust(fake_bridge):
|
||||
"""Caller-supplied timeout must flow into the Rust bridge so the fixed 600s
|
||||
client ceiling doesn't silently override shorter deadlines."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue