diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md index 2fa81798605..113d27e84e5 100644 --- a/litellm-rust/ADDING_A_PROVIDER.md +++ b/litellm-rust/ADDING_A_PROVIDER.md @@ -1,9 +1,10 @@ # Adding a provider / route to litellm-rust -Three layers, same for every route (see `ocr` and `realtime` as references): - -1. **Transform contract (pure)** — `crates/core/src//transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth. -2. **Provider config (pure)** — `crates/providers/src///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. -3. **HTTP / transport (the host)** — `crates/providers/src/.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O. - -**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`. +- Keep the route contract pure in `crates/core/src//`: define the typed request/response structs and a provider config trait with no network, env, auth, or logging. +- Add provider identity to the repo-root `provider_endpoints_support.json`: use the LiteLLM provider slug, display name, docs URL, and endpoint support flags. Put optional stable provider-level base URL defaults under the top-level `default_creds` map; keep route-specific API key env vars in provider config/transform code so key resolution has one owner. +- Put provider-specific transforms in `crates/providers/src///transformation.rs`, mirroring the Python provider tree and exposing a `const __CONFIG`. +- The provider config owns three pure steps: map LiteLLM params, transform the LiteLLM request into the provider request, and transform the provider response back into the LiteLLM response. +- If the provider has a reverse or normalization step, keep it pure and explicit next to the transforms; do not hide reverse mapping inside the HTTP transport. +- Route host functions in `crates/providers/src/.rs` must be async: resolve auth/base URL, call the transforms, send with async transport, then call the response transform. Use Tokio/async all the way through Rust route I/O; only the PyO3 sync compatibility wrapper should `block_on` the async route, and it must release the GIL while waiting. +- Do not add per-provider HTTP clients casually. Today Rust cannot call Python's `BaseLLMHTTPHandler`; if a route needs end-to-end Rust I/O, keep the async transport route-scoped, opt-in from Python, and do not broaden it to more providers until there is a shared Rust HTTP abstraction. +- Register modules in `lib.rs` / `mod.rs`, add parity tests for params/request/response behavior, then run `cargo fmt && cargo clippy --workspace --all-targets --locked -- -D warnings && cargo test --workspace --locked`. diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2ec86fc00e7..5f9c3a0b270 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -137,12 +137,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" @@ -152,6 +164,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" @@ -168,12 +195,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" @@ -192,8 +241,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", @@ -238,6 +289,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" @@ -293,6 +369,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", @@ -445,6 +522,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" @@ -511,7 +598,9 @@ dependencies = [ "litellm-core", "litellm-providers", "pyo3", + "pyo3-async-runtimes", "serde_json", + "tokio", ] [[package]] @@ -633,6 +722,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" @@ -815,9 +917,8 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64", "bytes", - "futures-channel", "futures-core", - "futures-util", + "h2", "http", "http-body", "http-body-util", @@ -1218,6 +1319,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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 0b1528672a7..63c7d6e8b9d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -15,7 +15,8 @@ repository = "https://github.com/BerriAI/litellm" litellm-core = { path = "crates/core" } litellm-providers = { path = "crates/providers" } pyo3 = "0.23.5" -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "http2"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" thiserror = "2.0" diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index e54002fe5e8..c3b16431b36 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -9,3 +9,7 @@ repository.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true + +[build-dependencies] +serde.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/core/build.rs b/litellm-rust/crates/core/build.rs new file mode 100644 index 00000000000..efc1da37f97 --- /dev/null +++ b/litellm-rust/crates/core/build.rs @@ -0,0 +1,357 @@ +use std::collections::{BTreeMap, HashSet}; +use std::env; +use std::fs; +use std::path::PathBuf; + +use serde::Deserialize; + +#[derive(Debug, Deserialize)] +struct ProviderEndpointSupportInput { + providers: BTreeMap, + #[serde(default)] + default_creds: BTreeMap, +} + +#[derive(Debug, Deserialize)] +struct ProviderInput { + display_name: String, + url: String, +} + +#[derive(Debug, Deserialize)] +struct ProviderDefaultCredsInput { + default_api_base: Option, + api_key_env_var: Option, +} + +#[derive(Debug)] +struct ProviderMetadataInput { + routing_name: String, + display_name: String, + docs_url: String, + default_api_base: Option, + api_key_env_var: Option, +} + +fn rust_string(value: &str) -> String { + format!("{value:?}") +} + +fn rust_option(value: Option<&str>) -> String { + value + .map(|value| format!("Some({})", rust_string(value))) + .unwrap_or_else(|| "None".to_string()) +} + +fn variant_name(routing_name: &str) -> String { + routing_name + .split(['_', '-', '.', '/']) + .filter(|part| !part.is_empty()) + .map(|part| { + let mut chars = part.chars(); + match chars.next() { + Some(first) => { + let mut out = String::new(); + out.extend(first.to_uppercase()); + out.push_str(chars.as_str()); + out + } + None => String::new(), + } + }) + .collect() +} + +fn provider_metadata_from_endpoint_support( + registry: ProviderEndpointSupportInput, +) -> Vec { + let ProviderEndpointSupportInput { + providers, + default_creds, + } = registry; + + providers + .into_iter() + .map(|(routing_name, provider)| { + let default_creds = default_creds.get(&routing_name); + ProviderMetadataInput { + routing_name, + display_name: provider.display_name, + docs_url: provider.url, + default_api_base: default_creds.and_then(|creds| creds.default_api_base.clone()), + api_key_env_var: default_creds.and_then(|creds| creds.api_key_env_var.clone()), + } + }) + .collect() +} + +fn validate_providers(providers: &[ProviderMetadataInput]) { + let mut routing_names = HashSet::new(); + let mut variants = HashSet::new(); + + for provider in providers { + if provider.routing_name.trim().is_empty() { + panic!("provider routing_name cannot be empty"); + } + if provider.display_name.trim().is_empty() { + panic!( + "provider {} has an empty display_name", + provider.routing_name + ); + } + if provider.docs_url.trim().is_empty() { + panic!("provider {} has an empty docs_url", provider.routing_name); + } + if !routing_names.insert(provider.routing_name.as_str()) { + panic!("duplicate provider routing_name: {}", provider.routing_name); + } + + let variant = variant_name(&provider.routing_name); + if variant.is_empty() { + panic!( + "provider {} generated an empty Rust variant", + provider.routing_name + ); + } + if variant.chars().next().is_some_and(|ch| ch.is_ascii_digit()) { + panic!( + "provider {} generated Rust variant {variant} starting with a digit", + provider.routing_name + ); + } + if !variants.insert(variant.clone()) { + panic!( + "provider {} generated duplicate Rust variant {variant}", + provider.routing_name + ); + } + } +} + +fn generate_provider_code(providers: &[ProviderMetadataInput]) -> String { + let variants: Vec = providers + .iter() + .map(|provider| variant_name(&provider.routing_name)) + .collect(); + + let enum_variants = variants + .iter() + .map(|variant| format!(" {variant},")) + .collect::>() + .join("\n"); + + let all_values = variants + .iter() + .map(|variant| format!(" LlmProvider::{variant},")) + .collect::>() + .join("\n"); + + let metadata_values = providers + .iter() + .zip(variants.iter()) + .map(|(provider, variant)| { + format!( + " ProviderMetadata {{ provider: LlmProvider::{variant}, routing_name: {}, display_name: {}, docs_url: {}, default_api_base: {}, api_key_env_var: {} }},", + rust_string(&provider.routing_name), + rust_string(&provider.display_name), + rust_string(&provider.docs_url), + rust_option(provider.default_api_base.as_deref()), + rust_option(provider.api_key_env_var.as_deref()), + ) + }) + .collect::>() + .join("\n"); + + let metadata_match_arms = variants + .iter() + .enumerate() + .map(|(index, variant)| { + format!(" LlmProvider::{variant} => &Self::METADATA[{index}],") + }) + .collect::>() + .join("\n"); + + format!( + r#"// @generated by crates/core/build.rs from provider_endpoints_support.json. +// Do not edit this file by hand. + +use std::fmt; +use std::str::FromStr; + +use crate::error::CoreError; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ProviderMetadata {{ + pub provider: LlmProvider, + pub routing_name: &'static str, + pub display_name: &'static str, + pub docs_url: &'static str, + pub default_api_base: Option<&'static str>, + pub api_key_env_var: Option<&'static str>, +}} + +#[allow(clippy::enum_variant_names)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LlmProvider {{ +{enum_variants} +}} + +impl LlmProvider {{ + pub const ALL: [LlmProvider; {provider_count}] = [ +{all_values} + ]; + + pub const METADATA: [ProviderMetadata; {provider_count}] = [ +{metadata_values} + ]; + + pub fn metadata(self) -> &'static ProviderMetadata {{ + match self {{ +{metadata_match_arms} + }} + }} + + pub fn as_str(self) -> &'static str {{ + self.metadata().routing_name + }} + + pub fn display_name(self) -> &'static str {{ + self.metadata().display_name + }} + + pub fn docs_url(self) -> &'static str {{ + self.metadata().docs_url + }} + + pub fn default_api_base(self) -> Option<&'static str> {{ + self.metadata().default_api_base + }} + + pub fn api_key_env_var(self) -> Option<&'static str> {{ + self.metadata().api_key_env_var + }} +}} + +impl fmt::Display for LlmProvider {{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {{ + f.write_str(self.as_str()) + }} +}} + +impl FromStr for LlmProvider {{ + type Err = CoreError; + + fn from_str(value: &str) -> Result {{ + LlmProvider::ALL + .iter() + .copied() + .find(|provider| provider.as_str() == value) + .ok_or_else(|| CoreError::InvalidProvider(value.to_string())) + }} +}} + +#[cfg(test)] +mod tests {{ + use super::*; + + fn endpoint_support_provider_values() -> Vec {{ + let raw = include_str!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/../../../provider_endpoints_support.json" + )); + let registry = serde_json::from_str::(raw) + .expect("provider_endpoints_support.json parses"); + let mut values = registry + .get("providers") + .and_then(|providers| providers.as_object()) + .expect("provider_endpoints_support.json has providers object") + .keys() + .cloned() + .collect::>(); + values.sort(); + values + }} + + #[test] + fn provider_values_match_endpoint_support_registry() {{ + let registry_values = endpoint_support_provider_values(); + assert_eq!(LlmProvider::ALL.len(), registry_values.len()); + assert_eq!( + LlmProvider::ALL + .iter() + .map(|provider| provider.as_str().to_string()) + .collect::>(), + registry_values + ); + }} + + #[test] + fn from_str_round_trips_all_providers() {{ + for provider in LlmProvider::ALL {{ + assert_eq!(LlmProvider::from_str(provider.as_str()), Ok(provider)); + assert_eq!(provider.to_string(), provider.as_str()); + assert!(!provider.docs_url().is_empty()); + }} + }} + + #[test] + fn from_str_rejects_unknown_provider() {{ + assert_eq!( + LlmProvider::from_str("not-a-provider"), + Err(CoreError::InvalidProvider("not-a-provider".to_string())) + ); + }} + + #[test] + fn provider_metadata_exposes_optional_defaults() {{ + assert_eq!( + LlmProvider::Mistral.display_name(), + "Mistral AI API (`mistral`)" + ); + assert_eq!( + LlmProvider::Mistral.docs_url(), + "https://docs.litellm.ai/docs/providers/mistral" + ); + assert_eq!( + LlmProvider::Mistral.default_api_base(), + Some("https://api.mistral.ai/v1") + ); + assert_eq!(LlmProvider::Mistral.api_key_env_var(), None); + }} +}} +"#, + enum_variants = enum_variants, + provider_count = providers.len(), + all_values = all_values, + metadata_values = metadata_values, + metadata_match_arms = metadata_match_arms, + ) +} + +fn main() { + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); + let providers_path = manifest_dir.join("../../../provider_endpoints_support.json"); + println!("cargo:rerun-if-changed={}", providers_path.display()); + println!( + "cargo:rerun-if-changed={}", + manifest_dir.join("build.rs").display() + ); + + let raw = fs::read_to_string(&providers_path).unwrap_or_else(|err| { + panic!( + "failed to read provider endpoint support registry {}: {err}", + providers_path.display() + ) + }); + let registry: ProviderEndpointSupportInput = serde_json::from_str(&raw) + .unwrap_or_else(|err| panic!("failed to parse {}: {err}", providers_path.display())); + let providers = provider_metadata_from_endpoint_support(registry); + validate_providers(&providers); + + let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); + fs::write( + out_dir.join("provider_generated.rs"), + generate_provider_code(&providers), + ) + .expect("failed to write generated provider code"); +} diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 645e261f76d..c05bd6103e9 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -13,6 +13,10 @@ pub enum CoreError { MissingField(&'static str), #[error("invalid response: {0}")] InvalidResponse(String), + #[error("invalid request: {0}")] + InvalidRequest(String), + #[error("invalid provider: {0}")] + InvalidProvider(String), #[error("{0}")] Auth(String), #[error("OCR request failed with status {status}: {body}")] diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 5d4d5bfd142..6142cece16f 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,5 +1,7 @@ pub mod error; pub mod ocr; +pub mod providers; pub mod realtime; pub use error::{CoreError, CoreResult}; +pub use providers::LlmProvider; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 7353d9d22c4..049959c7bc1 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -4,7 +4,12 @@ use crate::CoreResult; use super::types::{OcrRequestData, OcrResponseData}; -pub trait OcrProviderConfig { +/// Provider-specific OCR transforms. +/// +/// Implementations should stay pure and non-blocking: map supported params, +/// build the provider request body, and normalize the provider response. The +/// route layer owns async HTTP I/O. +pub trait OcrProviderConfig: Send + Sync { fn supported_ocr_params(&self) -> &'static [&'static str]; fn map_ocr_params(&self, non_default_params: &Map) -> Map { diff --git a/litellm-rust/crates/core/src/providers/README.md b/litellm-rust/crates/core/src/providers/README.md new file mode 100644 index 00000000000..5c00534fd6e --- /dev/null +++ b/litellm-rust/crates/core/src/providers/README.md @@ -0,0 +1,19 @@ +# Rust Provider Metadata + +The repo-root `provider_endpoints_support.json` is the shared source of truth +for provider identity and docs metadata in `litellm-rust`. `crates/core/build.rs` +reads it at compile time and generates the typed `LlmProvider` enum plus static +provider metadata. Runtime code does not parse this JSON. + +To add a provider: + +- Add a `provider_endpoints_support.json` provider entry using the LiteLLM + provider slug, display name, docs URL, and endpoint support flags. +- Add optional defaults under the top-level `default_creds` map only when there + is a stable provider-level base URL. Keep route-specific key env var names in + the provider transform/config so auth resolution has one owner. +- Put request/response logic under + `crates/providers/src///transformation.rs`; do not put + transforms, signing logic, or secrets in provider metadata. +- Run `cargo test -p litellm-core --locked`; it verifies the Rust registry stays + in parity with `provider_endpoints_support.json`. diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs new file mode 100644 index 00000000000..3d6f90356a1 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -0,0 +1,3 @@ +pub mod provider; + +pub use provider::LlmProvider; diff --git a/litellm-rust/crates/core/src/providers/provider.rs b/litellm-rust/crates/core/src/providers/provider.rs new file mode 100644 index 00000000000..c0a511d10bb --- /dev/null +++ b/litellm-rust/crates/core/src/providers/provider.rs @@ -0,0 +1 @@ +include!(concat!(env!("OUT_DIR"), "/provider_generated.rs")); diff --git a/litellm-rust/crates/providers/src/ocr.rs b/litellm-rust/crates/providers/src/ocr.rs index dcd56a5f0b4..4f45c9a6dd1 100644 --- a/litellm-rust/crates/providers/src/ocr.rs +++ b/litellm-rust/crates/providers/src/ocr.rs @@ -4,12 +4,14 @@ //! 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. +use std::str::FromStr; use std::sync::OnceLock; use std::time::Duration; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrProviderConfig; use litellm_core::CoreResult; +use litellm_core::LlmProvider; use serde_json::{Map, Value}; use crate::mistral::ocr::transformation as mistral; @@ -25,11 +27,16 @@ const OCR_TIMEOUT_SECS: u64 = 600; /// 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 = 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 = 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") @@ -44,39 +51,89 @@ fn truncate_error_body(body: &str) -> String { format!("{truncated}... (truncated)") } +fn ocr_config_for(provider: LlmProvider) -> Option<&'static dyn OcrProviderConfig> { + match provider { + LlmProvider::Mistral => Some(&MISTRAL_OCR_CONFIG), + _ => None, + } +} + +fn string_headers(extra_headers: Option>) -> CoreResult> { + 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() +} + +fn has_authorization_header(headers: &[(String, String)]) -> bool { + headers + .iter() + .any(|(key, _)| key.eq_ignore_ascii_case("authorization")) +} + +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>, + pub optional_params: Map, + pub timeout: Option, +} + /// Perform a Mistral 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, - timeout: Option, -) -> CoreResult { - 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 { + let model = request.model; + let provider = LlmProvider::from_str(request.custom_llm_provider)?; + let config = + ocr_config_for(provider).ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?; - 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); + // TODO: key and URL resolution are still Mistral-specific while Mistral is + // the only Rust OCR provider. Move these onto the trait when another OCR + // provider is added here. + let api_key = mistral::resolve_api_key(request.api_key, &|key| std::env::var(key).ok())?; + let url = mistral::complete_url(request.api_base); + let filtered_params = config.map_ocr_params(&request.optional_params); let body = config - .transform_ocr_request(model, document, filtered_params)? + .transform_ocr_request(model, request.document, filtered_params)? .data; - let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body); - if let Some(duration) = timeout { - request = request.timeout(duration); + let headers = string_headers(request.extra_headers)?; + let mut request_builder = http_client().post(&url).json(&body); + if !has_authorization_header(&headers) { + request_builder = request_builder.bearer_auth(&api_key); + } + for (key, value) in headers { + request_builder = request_builder.header(&key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); } - let response = request + let response = request_builder .send() + .await .map_err(|err| CoreError::Network(err.to_string()))?; let status = response.status(); let text = response .text() + .await .map_err(|err| CoreError::Network(err.to_string()))?; if !status.is_success() { @@ -97,6 +154,9 @@ pub fn run_ocr( #[cfg(test)] mod tests { use super::*; + use serde_json::json; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; #[test] fn truncate_error_body_passes_short_strings_through() { @@ -124,4 +184,135 @@ mod tests { let truncated = truncate_error_body(&body); assert!(truncated.is_char_boundary(truncated.len())); } + + #[test] + fn ocr_registry_supports_only_mistral() { + assert!(ocr_config_for(LlmProvider::Mistral).is_some()); + assert!(ocr_config_for(LlmProvider::Openai).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 has_authorization_header_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_authorization_header(&headers)); + + let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; + assert!(has_authorization_header(&headers)); + + let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; + assert!(!has_authorization_header(&headers)); + } + + #[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 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; + } + } + + 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"); + String::from_utf8(request).expect("request is utf8") + }); + + 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}" + ); + } + + #[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() + ) + ); + } } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 80b6478daac..d0a7cd56cc6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -13,4 +13,6 @@ crate-type = ["cdylib"] litellm-core.workspace = true litellm-providers.workspace = true pyo3 = { workspace = true, features = ["extension-module"] } +pyo3-async-runtimes.workspace = true serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 15e93f7b00c..aa1fb2e0127 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,7 +1,6 @@ use std::time::Duration; use litellm_core::error::CoreError; -use litellm_providers::ocr::run_ocr; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; @@ -9,6 +8,13 @@ use serde_json::{Map, Value}; mod gil; +type MarshaledOcrInputs = ( + Value, + Option>, + Map, + Option, +); + fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { let json = py.import("json")?; let encoded: String = json.call_method1("dumps", (value,))?.extract()?; @@ -28,53 +34,94 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { 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::InvalidType { .. } + | CoreError::InvalidRequest(_) + | CoreError::MissingField(_) => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } +fn optional_object_to_map( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult> { + 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) -> Option { + 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, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + 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)) +} + /// Perform a Mistral OCR call end to end and return the response as a dict. #[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, api_key: Option, api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - 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). + // Release the GIL while the sync API waits on async Rust work. let result = gil::release_gil(py, || { - run_ocr( - &model, - document, - api_key.as_deref(), - api_base.as_deref(), - optional_params, - timeout, - ) + pyo3_async_runtimes::tokio::get_runtime().block_on(litellm_providers::ocr::ocr( + litellm_providers::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, + }, + )) }); match result { @@ -83,8 +130,50 @@ fn ocr( } } +/// Perform an OCR call end to end and return an asyncio awaitable. +#[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, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + 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 = litellm_providers::ocr::ocr(litellm_providers::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)) + }) +} + /// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe -/// how often the bridge has dropped the GIL for blocking work. +/// how often the sync bridge has dropped the GIL while awaiting Rust work. #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); @@ -95,6 +184,7 @@ fn gil_stats(py: Python<'_>) -> PyResult> { #[pymodule] fn litellm_python_bridge(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(()) } diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 263e0c094ce..de6bc2471ec 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -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, diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index 21e0e27a314..49555baaf70 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -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( diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 3a9ef8db804..80fde8ff9f3 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -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,27 @@ 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] + headers: dict[str, object] + + def _timeout_to_seconds( timeout: Optional[Union[float, httpx.Timeout]], ) -> Optional[float]: @@ -45,18 +71,152 @@ 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 _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, + ) + 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, + **prepared_request.optional_params, + }, + "api_base": resolved_complete_url, + "headers": resolved_headers, + }, + ) + return _PreparedRustOCRCall( + api_key=resolved_api_key, + headers=cast(dict[str, object], resolved_headers), + ) + + +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 +226,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_request.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared_request.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_request.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared_request.optional_params, + timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), ) ) @@ -108,12 +270,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 +336,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 +355,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 prepared.custom_llm_provider == "mistral" 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 +408,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": } 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 +446,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 +530,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": } 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 the whole Mistral OCR call to the Rust bridge. + if prepared.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, + 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} diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 61f9e8ca69a..1e3312c1473 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -11,7 +11,7 @@ can import it statically without forming an import cycle. from __future__ import annotations -from typing import Final, Protocol, cast +from typing import Awaitable, Final, Protocol, cast class RustOcr(Protocol): @@ -23,9 +23,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: @@ -36,21 +56,27 @@ _UNSET: Final[_Unset] = _Unset() _rust_ocr_enabled = False _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 Rust ``litellm_python_bridge`` 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: @@ -72,3 +98,14 @@ def load_rust_ocr() -> RustOcr | None: except ImportError: return None return cast(RustOcr, litellm_python_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 + try: + import litellm_python_bridge + except ImportError: + return None + return cast(RustAocr, getattr(litellm_python_bridge, "aocr", None)) diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dd7712aabca..bf4ca71043d 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -31,6 +31,11 @@ } } }, + "default_creds": { + "mistral": { + "default_api_base": "https://api.mistral.ai/v1" + } + }, "providers": { "a2a": { "display_name": "A2A (Agent-to-Agent) (`a2a`)", @@ -179,6 +184,7 @@ }, "apertis": { "display_name": "Apertis (`apertis`)", + "url": "https://docs.litellm.ai/docs/providers/apertis", "endpoints": { "chat_completions": true, "messages": false, @@ -460,6 +466,7 @@ }, "chutes": { "display_name": "Chutes (`chutes`)", + "url": "https://docs.litellm.ai/docs/providers/chutes", "endpoints": { "chat_completions": true, "messages": false, @@ -1470,6 +1477,7 @@ }, "nanogpt": { "display_name": "NanoGPT (`nanogpt`)", + "url": "https://docs.litellm.ai/docs/providers/nanogpt", "endpoints": { "chat_completions": true, "messages": false, @@ -1803,6 +1811,7 @@ }, "poe": { "display_name": "Poe (`poe`)", + "url": "https://docs.litellm.ai/docs/providers/poe", "endpoints": { "chat_completions": true, "messages": false, @@ -2030,6 +2039,7 @@ }, "synthetic": { "display_name": "Synthetic (`synthetic`)", + "url": "https://docs.litellm.ai/docs/providers/synthetic", "endpoints": { "chat_completions": true, "messages": true, diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index e32fee6afc5..823b4a5e86c 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, Request, Response, UploadFile from fastapi.responses import ORJSONResponse from litellm._logging import verbose_proxy_logger -from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -27,6 +26,8 @@ def _build_document_from_upload( Delegates to convert_file_document_to_url_document after resolving MIME type from the upload's content_type header or filename. """ + from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type + mime_type = content_type.split(";")[0].strip() if content_type else None if not mime_type or mime_type == "application/octet-stream": if filename: diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index b137ec59a1f..eb34389bcf9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -31,6 +31,11 @@ } } }, + "default_creds": { + "mistral": { + "default_api_base": "https://api.mistral.ai/v1" + } + }, "providers": { "a2a": { "display_name": "A2A (Agent-to-Agent) (`a2a`)", @@ -179,6 +184,7 @@ }, "apertis": { "display_name": "Apertis (`apertis`)", + "url": "https://docs.litellm.ai/docs/providers/apertis", "endpoints": { "chat_completions": true, "messages": false, @@ -495,6 +501,7 @@ }, "chutes": { "display_name": "Chutes (`chutes`)", + "url": "https://docs.litellm.ai/docs/providers/chutes", "endpoints": { "chat_completions": true, "messages": false, @@ -1593,6 +1600,7 @@ }, "nanogpt": { "display_name": "NanoGPT (`nanogpt`)", + "url": "https://docs.litellm.ai/docs/providers/nanogpt", "endpoints": { "chat_completions": true, "messages": false, @@ -1993,6 +2001,7 @@ }, "poe": { "display_name": "Poe (`poe`)", + "url": "https://docs.litellm.ai/docs/providers/poe", "endpoints": { "chat_completions": true, "messages": false, @@ -2253,6 +2262,7 @@ }, "synthetic": { "display_name": "Synthetic (`synthetic`)", + "url": "https://docs.litellm.ai/docs/providers/synthetic", "endpoints": { "chat_completions": true, "messages": true, diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e028064e4c..aad5a4d6f42 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,6 +3,7 @@ import importlib import sys import types +from typing import Any import httpx import pytest @@ -17,9 +18,12 @@ ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") 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, @@ -28,21 +32,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, } @@ -50,13 +68,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, @@ -67,21 +153,69 @@ 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) yield - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) @pytest.fixture @@ -92,6 +226,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() @@ -106,6 +248,12 @@ def test_load_rust_ocr_returns_injected_impl(): assert rust_bridge.load_rust_ocr() is bridge +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. @@ -114,20 +262,25 @@ 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(): 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(): @@ -135,6 +288,7 @@ def test_load_rust_ocr_none_when_extension_absent(): caller degrades to the Python path instead of raising ImportError.""" 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): @@ -143,10 +297,12 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): built in CI, so stand in a fake module via ``sys.modules``.""" fake_module = types.ModuleType("litellm_python_bridge") 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.setitem(sys.modules, "litellm_python_bridge", fake_module) 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(): @@ -161,16 +317,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) @@ -181,6 +335,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, } @@ -193,23 +352,38 @@ 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_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() resolver_calls = [] @@ -220,16 +394,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" @@ -242,16 +408,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 @@ -262,7 +426,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): @@ -270,6 +437,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, ) @@ -281,10 +449,79 @@ 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_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."""