mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge pull request #31253 from BerriAI/litellm_hotfix_ocr_async_rust
Some checks are pending
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
LiteLLM Rust / rustfmt, clippy, test (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Some checks are pending
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
LiteLLM Rust / rustfmt, clippy, test (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
feat: make rust OCR async-first
This commit is contained in:
ishaan-berri
GitHub
commit
bd2a1653bd
22 changed files with 1639 additions and 415 deletions
|
|
@ -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/<route>/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/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
|
||||
3. **HTTP / transport (the host)** — `crates/providers/src/<route>.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/<route>/`: 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/<provider>/<route>/transformation.rs`, mirroring the Python provider tree and exposing a `const <PROVIDER>_<ROUTE>_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/<route>.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`.
|
||||
|
|
|
|||
118
litellm-rust/Cargo.lock
generated
118
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
357
litellm-rust/crates/core/build.rs
Normal file
357
litellm-rust/crates/core/build.rs
Normal file
|
|
@ -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<String, ProviderInput>,
|
||||
#[serde(default)]
|
||||
default_creds: BTreeMap<String, ProviderDefaultCredsInput>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ProviderInput {
|
||||
display_name: String,
|
||||
url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ProviderDefaultCredsInput {
|
||||
default_api_base: Option<String>,
|
||||
api_key_env_var: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ProviderMetadataInput {
|
||||
routing_name: String,
|
||||
display_name: String,
|
||||
docs_url: String,
|
||||
default_api_base: Option<String>,
|
||||
api_key_env_var: Option<String>,
|
||||
}
|
||||
|
||||
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<ProviderMetadataInput> {
|
||||
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<String> = providers
|
||||
.iter()
|
||||
.map(|provider| variant_name(&provider.routing_name))
|
||||
.collect();
|
||||
|
||||
let enum_variants = variants
|
||||
.iter()
|
||||
.map(|variant| format!(" {variant},"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let all_values = variants
|
||||
.iter()
|
||||
.map(|variant| format!(" LlmProvider::{variant},"))
|
||||
.collect::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let metadata_match_arms = variants
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, variant)| {
|
||||
format!(" LlmProvider::{variant} => &Self::METADATA[{index}],")
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.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<Self, Self::Err> {{
|
||||
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<String> {{
|
||||
let raw = include_str!(concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
"/../../../provider_endpoints_support.json"
|
||||
));
|
||||
let registry = serde_json::from_str::<serde_json::Value>(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::<Vec<_>>();
|
||||
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::<Vec<_>>(),
|
||||
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");
|
||||
}
|
||||
|
|
@ -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}")]
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<String, Value>) -> Map<String, Value> {
|
||||
|
|
|
|||
19
litellm-rust/crates/core/src/providers/README.md
Normal file
19
litellm-rust/crates/core/src/providers/README.md
Normal file
|
|
@ -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/<provider>/<route>/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`.
|
||||
3
litellm-rust/crates/core/src/providers/mod.rs
Normal file
3
litellm-rust/crates/core/src/providers/mod.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
pub mod provider;
|
||||
|
||||
pub use provider::LlmProvider;
|
||||
1
litellm-rust/crates/core/src/providers/provider.rs
Normal file
1
litellm-rust/crates/core/src/providers/provider.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
include!(concat!(env!("OUT_DIR"), "/provider_generated.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<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")
|
||||
|
|
@ -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<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()
|
||||
}
|
||||
|
||||
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<Map<String, Value>>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
/// 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<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 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()
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<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()?;
|
||||
|
|
@ -28,53 +34,94 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
|
|||
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<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))
|
||||
}
|
||||
|
||||
/// 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<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).
|
||||
// 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<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 = 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<Py<PyAny>> {
|
||||
let stats = PyDict::new(py);
|
||||
|
|
@ -95,6 +184,7 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
|||
#[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(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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": <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 +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": <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 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}
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue