From b1f74b360306c0a8da179c0872d7ac0db9c52d1f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 7 Sep 2026 12:09:54 -0700 Subject: [PATCH] wip --- .opencode/plugins/litellm.ts | 212 ++ litellm-rust/README.md | 43 +- .../crates/ai-gateway/src/ocr/common_utils.rs | 20 +- litellm-rust/crates/core/src/ocr/mod.rs | 87 + litellm-rust/crates/core/src/ocr/prepare.rs | 115 ++ .../crates/core/src/ocr/transformation.rs | 20 +- litellm-rust/crates/core/src/ocr/types.rs | 69 + .../providers/azure_ai/ocr/transformation.rs | 51 +- .../providers/vertex_ai/ocr/transformation.rs | 30 +- litellm-rust/crates/core/tests/ocr.rs | 632 ++++++ .../python-bridge/src/routes/definition.rs | 2 +- .../crates/python-bridge/src/routes/ocr.rs | 896 ++++++-- .../tests/callback_lifecycle.rs | 32 +- .../tests/fixtures/callback_components.py | 385 +++- litellm/ocr/main.py | 125 +- litellm/rust_bridge/ocr.py | 224 +- .../ocr/test_ocr_native_format.py | 46 +- tests/test_litellm/ocr/test_rust_bridge.py | 1823 ++++++++--------- .../rust_bridge/native_route_wheel_test.py | 196 +- 19 files changed, 3516 insertions(+), 1492 deletions(-) create mode 100644 .opencode/plugins/litellm.ts create mode 100644 litellm-rust/crates/core/src/ocr/prepare.rs create mode 100644 litellm-rust/crates/core/tests/ocr.rs diff --git a/.opencode/plugins/litellm.ts b/.opencode/plugins/litellm.ts new file mode 100644 index 00000000000..8dccec0acae --- /dev/null +++ b/.opencode/plugins/litellm.ts @@ -0,0 +1,212 @@ +import type { Plugin } from "@opencode-ai/plugin" + +type JsonObject = Record + +type LiteLLMModelGroup = { + model_group: string + max_input_tokens?: number + max_output_tokens?: number + input_cost_per_token?: number + output_cost_per_token?: number + mode?: string + supports_vision?: boolean + supports_reasoning?: boolean + supports_function_calling?: boolean + supported_reasoning_efforts?: string[] + supported_openai_params?: string[] +} + +type OpenCodeModel = { + name: string + attachment: boolean + reasoning: boolean + temperature: boolean + tool_call: boolean + cost?: { + input: number + output: number + } + limit?: { + context: number + input?: number + output: number + } + modalities: { + input: Array<"text" | "image"> + output: ["text"] + } + variants?: Record +} + +const PROVIDER_ID = "litellm" + +function isObject(value: unknown): value is JsonObject { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function optionalBoolean(value: unknown): boolean | undefined { + return typeof value === "boolean" ? value : undefined +} + +function optionalNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined +} + +function optionalString(value: unknown): string | undefined { + return typeof value === "string" && value.length > 0 ? value : undefined +} + +function optionalStrings(value: unknown): string[] | undefined { + if (!Array.isArray(value) || !value.every((item) => typeof item === "string")) return undefined + return value +} + +function parseModelIDs(payload: unknown): string[] { + if (!isObject(payload) || !Array.isArray(payload.data)) return [] + + return Array.from( + new Set( + payload.data.flatMap((item) => { + if (!isObject(item)) return [] + const id = optionalString(item.id) + return id && !id.includes("*") ? [id] : [] + }), + ), + ) +} + +function parseModelGroups(payload: unknown): Map { + if (!isObject(payload) || !Array.isArray(payload.data)) return new Map() + + return new Map( + payload.data.flatMap((item): Array<[string, LiteLLMModelGroup]> => { + if (!isObject(item)) return [] + const modelGroup = optionalString(item.model_group) + if (!modelGroup) return [] + + return [ + [ + modelGroup, + { + model_group: modelGroup, + max_input_tokens: optionalNumber(item.max_input_tokens), + max_output_tokens: optionalNumber(item.max_output_tokens), + input_cost_per_token: optionalNumber(item.input_cost_per_token), + output_cost_per_token: optionalNumber(item.output_cost_per_token), + mode: optionalString(item.mode), + supports_vision: optionalBoolean(item.supports_vision), + supports_reasoning: optionalBoolean(item.supports_reasoning), + supports_function_calling: optionalBoolean(item.supports_function_calling), + supported_reasoning_efforts: optionalStrings(item.supported_reasoning_efforts), + supported_openai_params: optionalStrings(item.supported_openai_params), + }, + ], + ] + }), + ) +} + +function modelConfig(id: string, group?: LiteLLMModelGroup): OpenCodeModel { + const context = group?.max_input_tokens + const output = group?.max_output_tokens + const inputCost = group?.input_cost_per_token + const outputCost = group?.output_cost_per_token + const reasoningEfforts = group?.supported_reasoning_efforts + + return { + name: id, + attachment: group?.supports_vision ?? false, + reasoning: group?.supports_reasoning ?? false, + temperature: group?.supported_openai_params?.includes("temperature") ?? true, + tool_call: group?.supports_function_calling ?? true, + ...(inputCost !== undefined && outputCost !== undefined + ? { cost: { input: inputCost * 1_000_000, output: outputCost * 1_000_000 } } + : {}), + ...(context !== undefined && output !== undefined + ? { limit: { context, output } } + : {}), + modalities: { + input: group?.supports_vision ? ["text", "image"] : ["text"], + output: ["text"], + }, + ...(reasoningEfforts?.length + ? { + variants: Object.fromEntries( + reasoningEfforts.map((effort) => [effort, { reasoningEffort: effort }]), + ), + } + : {}), + } +} + +async function getJSON(url: string, apiKey: string): Promise { + const response = await fetch(url, { + headers: { + Authorization: `Bearer ${apiKey}`, + Accept: "application/json", + }, + signal: AbortSignal.timeout(10_000), + }) + + if (!response.ok) throw new Error(`${url} returned ${response.status}`) + return response.json() +} + +export const LiteLLM: Plugin = async ({ client }) => { + const configuredURL = process.env.LITELLM_BASE_URL?.trim() + const apiKey = process.env.LITELLM_API_KEY?.trim() + if (!configuredURL || !apiKey) return {} + + const gatewayURL = configuredURL.replace(/\/+$/, "") + const apiRoot = gatewayURL.endsWith("/v1") ? gatewayURL.slice(0, -3) : gatewayURL + const inferenceURL = `${apiRoot}/v1` + + try { + const [modelsPayload, groupsPayload] = await Promise.all([ + getJSON(`${inferenceURL}/models`, apiKey), + getJSON(`${apiRoot}/model_group/info`, apiKey).catch(() => undefined), + ]) + const groups = parseModelGroups(groupsPayload) + const models = Object.fromEntries( + parseModelIDs(modelsPayload) + .filter((id) => { + const mode = groups.get(id)?.mode + return mode === undefined || mode === "chat" + }) + .map((id) => [id, modelConfig(id, groups.get(id))]), + ) + + if (Object.keys(models).length === 0) throw new Error("the gateway returned no selectable models") + + return { + async config(config) { + config.provider ??= {} + const existing = config.provider[PROVIDER_ID] + config.provider[PROVIDER_ID] = { + npm: "@ai-sdk/openai-compatible", + name: "LiteLLM Gateway", + ...existing, + options: { + baseURL: inferenceURL, + apiKey, + ...existing?.options, + }, + models: { + ...models, + ...existing?.models, + }, + } + }, + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + await client.app.log({ + body: { + service: "litellm-model-discovery", + level: "warn", + message: `LiteLLM model discovery skipped: ${message}`, + }, + }) + return {} + } +} diff --git a/litellm-rust/README.md b/litellm-rust/README.md index 68e02ab335d..5b7180d0148 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -51,6 +51,33 @@ function per top-level route, mirroring the core entrypoints. ## Checks +### Native OCR Boundary + +`LITELLM_RUST=1` selects native OCR before Python provider preparation. With Rust +disabled, the existing Python execution and authentication paths are unchanged + +The bridge retains the complete call argument dictionary as a Python object, +including opaque callback and metadata objects. It creates callback-visible +request dictionaries with shared parameter references, then reads the execution +roots after pre-call dispatch. Mistral retains the original document; Azure and +Vertex Mistral use a shallow document copy, and Vertex DeepSeek projects it into +chat messages using the existing Rust transform. Rust performs provider preparation, +encoding, HTTP and response normalization. Python continues to dispatch existing +logging operations and construct the public response object + +This is an opt-in implementation scaffold, not full OCR parity. Azure Mistral and +Vertex Mistral accept inline data URIs with supplied keys/tokens, native environment +keys or auth headers. Azure also accepts a supplied `azure_ad_token`. Vertex +DeepSeek uses its existing chat request and OCR response transforms. Cloud +credential acquisition fails explicitly only when no native credential is available. +HTTP document URL conversion fails only for configs requiring data URIs. Azure +Document Intelligence selects its own config but fails at the polling capability +check before sending a billable analyze request. Cohere transforms, file inputs, +streaming, native response format and compression remain unsupported. +An enabled but missing native extension also fails; +neither case falls back to Python execution. Transport failures currently use a +generic error rather than the SDK's timeout-specific exception + Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust changes. That list is the single source of truth and matches what GitHub Actions runs for changes under `litellm-rust/`. @@ -90,14 +117,16 @@ The gate checks that native `ocr` and `aocr` are importable, then runs `LITELLM_REQUIRE_NATIVE_OCR=1`, so unavailable native OCR fails instead of skipping. CI uses `make test-rust-ocr RUST_OCR_WHEEL=/absolute/path/to/current.whl` to test the release wheel it just built. The stdlib-only -`native_route_wheel_test.py` also exercises sync/async OCR through a small -boundary, including 429 handling in `finish`/`afinish`, alongside the other -native routes +`native_route_wheel_test.py` also exercises sync/async OCR through the retained +argument dictionary, including native request preparation and public 429 error +mapping, alongside the other native routes -The Python-integrated Cargo tests validate retained callback identity, mutation, invocation context, -and ownership against Python behavior, including existing LiteLLM components. -They do not wire retained callbacks into production routes or change provider -preparation, authentication, HTTP transport, or response transformation +The Python-integrated Cargo tests validate retained callback identity, mutation, +invocation context and ownership against Python behavior, including existing +LiteLLM components. Short synthetic pre-call contracts use Rust-owned table-driven +cases with inline Python callbacks; larger component scenarios share Python +fixtures. These generic proofs complement, rather than replace, native OCR +acceptance tests The callback lifecycle scenarios use `#[serial(python_interpreter)]` to isolate CPython GC and interpreter-wide diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index d2be17260a3..f2799244071 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -8,15 +8,10 @@ use litellm_core::ocr::transformation::OcrProviderConfig; use reqwest::Url; use serde_json::{Map, Value}; -use litellm_core::providers::azure_ai::ocr::transformation::{ - AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG, -}; +use litellm_core::providers::azure_ai::ocr::transformation as azure_ai; use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; use litellm_core::providers::reducto::ocr::transformation as reducto; use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; -use litellm_core::providers::vertex_ai::ocr::transformation::{ - VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, -}; use crate::client::http_client; @@ -41,21 +36,12 @@ pub(super) fn ocr_provider_config( match provider { "mistral" => Some(&MISTRAL_OCR_CONFIG), "reducto" => reducto::config_for_model(model), - "azure_ai" if is_azure_document_intelligence_model(model) => { - Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) - } - "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), - "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), - "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), + "azure_ai" => azure_ai::config_for_model(model).ok(), + "vertex_ai" => vertex_ai::config_for_model(model).ok(), _ => None, } } -fn is_azure_document_intelligence_model(model: &str) -> bool { - let model = model.to_ascii_lowercase(); - model.contains("doc-intelligence") || model.contains("documentintelligence") -} - pub(super) fn string_headers( extra_headers: Option>, ) -> Result, Error> { diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index ec2fbb969a6..420c4d222a0 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,2 +1,89 @@ +pub mod prepare; pub mod transformation; pub mod types; + +use serde_json::Value; + +use crate::Error; +use crate::error::json_type_name; +use crate::http_utils::{buffered_post, has_header}; + +use types::{OcrDocument, OcrDocumentProjection}; +pub use types::{OcrRequest, OcrResponseData, PreparedOcr}; + +pub fn terminal_callbacks(asynchronous: bool, success: bool) -> &'static [&'static str] { + match (asynchronous, success) { + (false, true) => &["sync_success"], + (true, true) => &["async_success", "sync_success_if_needed"], + (false, false) => &["sync_failure"], + (true, false) => &["sync_failure", "async_failure"], + } +} + +pub async fn ocr( + prepared: PreparedOcr, + headers: Vec<(String, String)>, + body: Value, +) -> Result { + let config = prepare::provider_config(&prepared.custom_llm_provider, &prepared.model)?; + prepare::validate_capabilities(config)?; + let object = body.as_object().ok_or_else(|| Error::InvalidType { + expected: "object", + actual: json_type_name(&body), + })?; + if config.document_projection() != OcrDocumentProjection::Transformed { + let document = object + .get("document") + .ok_or(Error::MissingField("document"))?; + let document: OcrDocument = serde_json::from_value(document.clone()) + .map_err(|_| Error::InvalidRequest("invalid OCR document".into()))?; + document.validate(config.requires_data_uri_document())?; + } + if object.get("stream").and_then(Value::as_bool) == Some(true) { + return Err(Error::Unsupported("OCR streaming response handling")); + } + if headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("content-encoding") && !value.eq_ignore_ascii_case("identity") + }) { + return Err(Error::Unsupported("compressed OCR request")); + } + let headers = headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_str())) + .chain( + [ + ("Content-Type", "application/json"), + ("Accept-Encoding", "identity"), + ] + .into_iter() + .filter(|(name, _)| !has_header(&headers, name)), + ) + .map(|(name, value)| (name.as_bytes().to_vec(), value.as_bytes().to_vec())) + .collect(); + let body = serde_json::to_vec(&body) + .map_err(|_| Error::InvalidRequest("could not encode OCR request".into()))?; + let response = buffered_post::send(buffered_post::Request { + url: prepared.url, + headers, + body, + timeout_seconds: prepared.timeout_seconds, + }) + .await?; + if !(200..300).contains(&response.status) { + return Err(Error::Http { + status: response.status, + body: "OCR provider request failed".into(), + }); + } + if response.headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case(b"content-encoding") + && value + .split(|byte| *byte == b',') + .any(|encoding| !encoding.trim_ascii().eq_ignore_ascii_case(b"identity")) + }) { + return Err(Error::Unsupported("compressed OCR response")); + } + let response_json = serde_json::from_slice(&response.content) + .map_err(|_| Error::InvalidResponse("invalid OCR JSON response".into()))?; + config.transform_ocr_response(&prepared.model, response_json) +} diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs new file mode 100644 index 00000000000..c2d86a3c01d --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -0,0 +1,115 @@ +use std::time::Duration; + +use serde_json::{Map, Value}; + +use crate::Error; +use crate::providers::azure_ai::ocr::transformation as azure_ai; +use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use crate::providers::vertex_ai::ocr::transformation as vertex_ai; +use crate::routing_utils::provider::get_custom_llm_provider; + +use super::transformation::{OcrProviderConfig, OcrResponseHandling}; +use super::types::OcrRequest; +pub use super::types::PreparedOcr; + +pub fn prepare(request: OcrRequest) -> Result { + match request.request_format.as_deref() { + None | Some("litellm") => {} + Some("native") => return Err(Error::Unsupported("native OCR request format")), + Some(_) => { + return Err(Error::Unsupported( + "OCR request format must be litellm or native", + )); + } + } + let provider = get_custom_llm_provider(&request.model, request.custom_llm_provider.as_deref()) + .ok_or_else(|| Error::InvalidProvider("unable to resolve OCR provider".into()))?; + let config = provider_config(provider.custom_llm_provider, provider.model)?; + if provider.model.trim().is_empty() { + return Err(Error::InvalidRequest("OCR model must not be empty".into())); + } + Duration::try_from_secs_f64(request.timeout_seconds) + .ok() + .filter(|timeout| !timeout.is_zero()) + .ok_or_else(|| Error::InvalidRequest("timeout must be positive and finite".into()))?; + + let env_lookup = |key: &str| std::env::var(key).ok(); + let headers = config + .validate_credentials( + request.extra_headers, + request.api_key.as_deref(), + request.azure_ad_token.as_deref(), + &env_lookup, + ) + .map_err( + |error| match (error, config.credential_acquisition_operation()) { + (Error::Auth(_), Some(operation)) => Error::Unsupported(operation), + (error, _) => error, + }, + )?; + validate_capabilities(config)?; + request + .document + .validate(config.requires_data_uri_document())?; + if request.stream { + return Err(Error::Unsupported("OCR streaming response handling")); + } + let url_params = [ + ("vertex_project", request.vertex_project), + ("vertex_location", request.vertex_location), + ] + .into_iter() + .filter_map(|(name, value)| value.map(|value| (name.into(), Value::String(value)))) + .collect(); + let url = config.complete_url( + request.api_base.as_deref(), + provider.model, + &url_params, + &env_lookup, + )?; + let parsed_url = reqwest::Url::parse(&url) + .map_err(|_| Error::InvalidRequest("invalid OCR API URL".into()))?; + if !matches!(parsed_url.scheme(), "http" | "https") || parsed_url.host_str().is_none() { + return Err(Error::InvalidRequest("invalid OCR API URL".into())); + } + let document = serde_json::to_value(request.document) + .map_err(|_| Error::InvalidRequest("could not project OCR document".into()))?; + let template = config.transform_ocr_request(provider.model, document, Map::new())?; + if template.files.is_some() { + return Err(Error::Unsupported("OCR multipart document preparation")); + } + let Value::Object(body) = template.data else { + return Err(Error::Unsupported("non-object OCR request template")); + }; + Ok(PreparedOcr { + model: provider.model.to_string(), + custom_llm_provider: provider.custom_llm_provider.to_string(), + url, + headers, + body, + document_projection: config.document_projection(), + parameter_fields: config.supported_ocr_params(), + timeout_seconds: request.timeout_seconds, + }) +} + +pub(super) fn validate_capabilities(config: &dyn OcrProviderConfig) -> Result<(), Error> { + match config.response_handling() { + OcrResponseHandling::Json => Ok(()), + OcrResponseHandling::AzureDocumentIntelligencePoll => Err(Error::Unsupported( + "Azure Document Intelligence OCR polling", + )), + } +} + +pub(super) fn provider_config( + provider: &str, + model: &str, +) -> Result<&'static dyn OcrProviderConfig, Error> { + match provider { + "mistral" => Ok(&MISTRAL_OCR_CONFIG), + "azure_ai" => azure_ai::config_for_model(model), + "vertex_ai" => vertex_ai::config_for_model(model), + _ => Err(Error::Unsupported("OCR provider")), + } +} diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 62299faf9ed..6bb13771396 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -1,7 +1,7 @@ use crate::Error; use serde_json::{Map, Value}; -use super::types::{OcrRequestData, OcrResponseData}; +use super::types::{OcrDocumentProjection, OcrRequestData, OcrResponseData}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrAuthStrategy { @@ -25,6 +25,24 @@ pub enum OcrResponseHandling { } pub trait OcrProviderConfig: Sync { + fn document_projection(&self) -> OcrDocumentProjection { + OcrDocumentProjection::RetainedDocument + } + + fn credential_acquisition_operation(&self) -> Option<&'static str> { + None + } + + fn validate_credentials( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + _azure_ad_token: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result, Error> { + self.validate_environment(headers, api_key, env_lookup) + } + fn supported_ocr_params(&self) -> &'static [&'static str]; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 71cdb232a87..78f8fab5b3e 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,6 +1,75 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +pub struct OcrRequest { + pub model: String, + pub custom_llm_provider: Option, + pub api_key: Option, + pub api_base: Option, + pub extra_headers: Vec<(String, String)>, + pub timeout_seconds: f64, + pub request_format: Option, + pub document: OcrDocument, + pub azure_ad_token: Option, + pub vertex_project: Option, + pub vertex_location: Option, + pub stream: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum OcrDocument { + DocumentUrl { + document_url: String, + }, + ImageUrl { + image_url: String, + }, + File, + #[serde(other)] + Unsupported, +} + +impl OcrDocument { + pub fn validate(&self, requires_data_uri: bool) -> Result<(), crate::Error> { + let url = match self { + Self::DocumentUrl { document_url } => document_url, + Self::ImageUrl { image_url } => image_url, + Self::File => return Err(crate::Error::Unsupported("OCR file document preparation")), + Self::Unsupported => return Err(crate::Error::Unsupported("OCR document type")), + }; + let parsed = reqwest::Url::parse(url) + .map_err(|_| crate::Error::Unsupported("OCR local or non-HTTP document preparation"))?; + match parsed.scheme() { + "http" | "https" if requires_data_uri => Err(crate::Error::Unsupported( + "OCR HTTP document URL to data URI conversion", + )), + "http" | "https" | "data" => Ok(()), + _ => Err(crate::Error::Unsupported( + "OCR local or non-HTTP document preparation", + )), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrDocumentProjection { + RetainedDocument, + ShallowCopyDocument, + Transformed, +} + +pub struct PreparedOcr { + pub model: String, + pub custom_llm_provider: String, + pub url: String, + pub headers: Vec<(String, String)>, + pub body: Map, + pub document_projection: OcrDocumentProjection, + pub parameter_fields: &'static [&'static str], + pub timeout_seconds: f64, +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct OcrRequestData { pub data: Value, diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index d15c032f0bc..6fa932e4f0c 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -2,7 +2,7 @@ use std::collections::BTreeSet; use crate::error::{Error, json_type_name}; use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use crate::ocr::types::{OcrDocumentProjection, OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -24,6 +24,19 @@ pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig; pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig = AzureDocumentIntelligenceOcrConfig; +pub fn config_for_model(model: &str) -> Result<&'static dyn OcrProviderConfig, Error> { + let model = model.to_ascii_lowercase(); + if model.contains("cohere") { + return Err(Error::Unsupported( + "Azure Cohere OCR request transformation", + )); + } + if model.contains("doc-intelligence") || model.contains("documentintelligence") { + return Ok(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG); + } + Ok(&AZURE_AI_OCR_CONFIG) +} + fn non_empty(value: Option<&str>) -> Option<&str> { value.map(str::trim).filter(|value| !value.is_empty()) } @@ -503,6 +516,24 @@ fn transform_document_intelligence_response( } impl OcrProviderConfig for AzureAiOcrConfig { + fn document_projection(&self) -> OcrDocumentProjection { + OcrDocumentProjection::ShallowCopyDocument + } + + fn credential_acquisition_operation(&self) -> Option<&'static str> { + Some("Azure OCR credential acquisition") + } + + fn validate_credentials( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + azure_ad_token: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result, Error> { + validate_azure_ai_environment(headers, api_key, azure_ad_token, env_lookup) + } + fn supported_ocr_params(&self) -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() } @@ -550,6 +581,24 @@ impl OcrProviderConfig for AzureAiOcrConfig { } impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { + fn document_projection(&self) -> OcrDocumentProjection { + OcrDocumentProjection::Transformed + } + + fn credential_acquisition_operation(&self) -> Option<&'static str> { + Some("Azure OCR credential acquisition") + } + + fn validate_credentials( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + azure_ad_token: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result, Error> { + validate_document_intelligence_environment(headers, api_key, azure_ad_token, env_lookup) + } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index c324de8cb45..2b4941f459a 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,6 +1,6 @@ use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; -use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use crate::ocr::types::{OcrDocumentProjection, OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -29,6 +29,18 @@ pub struct VertexAiDeepSeekOcrConfig; pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig; pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; +pub fn config_for_model(model: &str) -> Result<&'static dyn OcrProviderConfig, Error> { + if model.to_ascii_lowercase().contains("cohere") { + return Err(Error::Unsupported( + "Vertex Cohere OCR request transformation", + )); + } + if is_deepseek_model(model) { + return Ok(&VERTEX_AI_DEEPSEEK_OCR_CONFIG); + } + Ok(&VERTEX_AI_OCR_CONFIG) +} + fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { keys.iter() .find_map(|key| params.get(*key).and_then(Value::as_str)) @@ -208,6 +220,14 @@ fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> V } impl OcrProviderConfig for VertexAiOcrConfig { + fn document_projection(&self) -> OcrDocumentProjection { + OcrDocumentProjection::ShallowCopyDocument + } + + fn credential_acquisition_operation(&self) -> Option<&'static str> { + Some("Vertex OCR credential acquisition") + } + fn supported_ocr_params(&self) -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() } @@ -255,6 +275,14 @@ impl OcrProviderConfig for VertexAiOcrConfig { } impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { + fn document_projection(&self) -> OcrDocumentProjection { + OcrDocumentProjection::Transformed + } + + fn credential_acquisition_operation(&self) -> Option<&'static str> { + VERTEX_AI_OCR_CONFIG.credential_acquisition_operation() + } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { DEEPSEEK_SUPPORTED_OCR_PARAMS diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs new file mode 100644 index 00000000000..c200db6c951 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -0,0 +1,632 @@ +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::process::Command; +use std::thread; +use std::time::Duration; + +use litellm_core::Error; +use litellm_core::ocr::prepare::prepare; +use litellm_core::ocr::types::{OcrDocument, OcrDocumentProjection}; +use litellm_core::ocr::{OcrRequest, PreparedOcr, ocr}; +use serde_json::{Value, json}; + +fn request() -> OcrRequest { + OcrRequest { + model: "mistral/mistral-ocr-latest".into(), + custom_llm_provider: None, + api_key: Some(" test-key ".into()), + api_base: None, + extra_headers: vec![], + timeout_seconds: 2.0, + request_format: None, + document: OcrDocument::DocumentUrl { + document_url: "data:application/pdf;base64,cGRm".into(), + }, + azure_ad_token: None, + vertex_project: Some("test-project".into()), + vertex_location: Some("us-central1".into()), + stream: false, + } +} + +fn body(prepared: &PreparedOcr) -> Value { + let mut body = prepared.body.clone(); + body.insert( + "document".into(), + json!({"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"}), + ); + body.insert("include_image_base64".into(), json!(true)); + body.insert("pages".into(), json!([0, 2])); + Value::Object(body) +} + +#[test] +fn prepares_provider_template_auth_and_url() { + let prepared = prepare(OcrRequest { + api_base: Some(" https://ocr.example/v1/ ".into()), + extra_headers: vec![("X-Request-Id".into(), "request-1".into())], + request_format: Some("litellm".into()), + ..request() + }) + .unwrap(); + assert_eq!(prepared.model, "mistral-ocr-latest"); + assert_eq!(prepared.custom_llm_provider, "mistral"); + assert_eq!(prepared.url, "https://ocr.example/v1/ocr"); + assert_eq!(prepared.timeout_seconds, 2.0); + assert_eq!( + prepared.document_projection, + OcrDocumentProjection::RetainedDocument + ); + assert_eq!( + Value::Object(prepared.body), + json!({"model": "mistral-ocr-latest", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"}}) + ); + assert_eq!( + prepared.parameter_fields, + litellm_core::providers::mistral::ocr::transformation::supported_ocr_params() + ); + assert_eq!( + prepared.headers, + vec![ + ("Authorization".into(), "Bearer test-key".into()), + ("X-Request-Id".into(), "request-1".into()), + ] + ); + let explicit = prepare(OcrRequest { + model: "mistral-ocr-latest".into(), + custom_llm_provider: Some("mistral".into()), + api_key: None, + extra_headers: vec![("aUtHoRiZaTiOn".into(), "Bearer explicit".into())], + ..request() + }) + .unwrap(); + assert_eq!(explicit.model, "mistral-ocr-latest"); + assert_eq!(explicit.url, "https://api.mistral.ai/v1/ocr"); + assert_eq!( + explicit.headers, + vec![("aUtHoRiZaTiOn".into(), "Bearer explicit".into())] + ); +} + +#[test] +fn prepare_environment_credentials() { + if let Ok(case) = std::env::var("LITELLM_OCR_ENV_TEST") { + let result = prepare(OcrRequest { + api_key: Some(" ".into()), + ..request() + }); + if case == "present" { + assert_eq!( + result.unwrap().headers, + vec![("Authorization".into(), "Bearer env-key".into())] + ); + assert_eq!( + prepare(request()).unwrap().headers, + vec![("Authorization".into(), "Bearer test-key".into())] + ); + } else { + assert!(matches!(result, Err(Error::Auth(_)))); + } + return; + } + for (case, key) in [ + ("present", Some("env-key")), + ("absent", None), + ("blank", Some(" ")), + ] { + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args(["--exact", "prepare_environment_credentials"]) + .env("LITELLM_OCR_ENV_TEST", case) + .env_remove("MISTRAL_API_KEY"); + if let Some(key) = key { + command.env("MISTRAL_API_KEY", key); + } + assert!(command.status().unwrap().success()); + } +} + +#[test] +fn prepare_rejects_unsupported_providers_formats_and_invalid_metadata() { + for (provider, capability) in [("reducto", "OCR provider"), ("openai", "OCR provider")] { + let result = prepare(OcrRequest { + custom_llm_provider: Some(provider.into()), + api_key: None, + api_base: Some("not a URL".into()), + ..request() + }); + assert!(matches!(result, Err(Error::Unsupported(message)) if message.contains(capability))); + } + for format in ["native", "json", "", "LiteLLM"] { + assert!(matches!( + prepare(OcrRequest { + request_format: Some(format.into()), + ..request() + }), + Err(Error::Unsupported(_)) + )); + } + for timeout_seconds in [0.0, -1.0, f64::NAN, f64::INFINITY, f64::MAX] { + assert!(matches!( + prepare(OcrRequest { + timeout_seconds, + ..request() + }), + Err(Error::InvalidRequest(_)) + )); + } + assert!(matches!( + prepare(OcrRequest { + model: "mistral-ocr-latest".into(), + ..request() + }), + Err(Error::InvalidProvider(_)) + )); + assert!(matches!( + prepare(OcrRequest { + api_base: Some("file:///secret".into()), + ..request() + }), + Err(Error::InvalidRequest(_)) + )); +} + +#[test] +fn cloud_capabilities_fail_only_when_required() { + for model in [ + "azure_ai/mistral-ocr-latest", + "vertex_ai/mistral-ocr-latest", + ] { + assert!(matches!( + prepare(OcrRequest { + model: model.into(), + api_base: Some("http://127.0.0.1:1".into()), + document: OcrDocument::ImageUrl { + image_url: "https://example.test/image.png".into() + }, + ..request() + }), + Err(Error::Unsupported( + "OCR HTTP document URL to data URI conversion" + )) + )); + } + for model in [ + "azure_ai/doc-intelligence/prebuilt-read", + "azure_ai/documentintelligence/prebuilt-layout", + ] { + assert!(matches!( + prepare(OcrRequest { + model: model.into(), + api_base: Some("http://127.0.0.1:1".into()), + ..request() + }), + Err(Error::Unsupported( + "Azure Document Intelligence OCR polling" + )) + )); + } + for model in [ + "azure_ai/cohere/parse-v5.0", + "vertex_ai/cohere/parse-v5.0", + "cohere/parse-v5.0", + ] { + assert!(matches!( + prepare(OcrRequest { + model: model.into(), + ..request() + }), + Err(Error::Unsupported(_)) + )); + } +} + +#[test] +fn cloud_credentials_use_native_keys_headers_or_narrow_acquisition_stub() { + if std::env::var_os("LITELLM_CLOUD_OCR_ENV_TEST").is_some() { + for (model, operation, header, key) in [ + ( + "azure_ai/mistral-ocr-latest", + "Azure OCR credential acquisition", + "Api-Key", + "env-azure", + ), + ( + "vertex_ai/mistral-ocr-latest", + "Vertex OCR credential acquisition", + "Authorization", + "Bearer env-vertex", + ), + ] { + let make_request = || OcrRequest { + model: model.into(), + api_key: None, + api_base: Some("http://127.0.0.1:1".into()), + ..request() + }; + let result = prepare(make_request()); + if std::env::var("LITELLM_CLOUD_OCR_ENV_TEST").unwrap() == "present" { + assert_eq!(result.unwrap().headers, vec![(header.into(), key.into())]); + } else { + assert!(matches!(result, Err(Error::Unsupported(message)) if message == operation)); + } + let prepared = prepare(OcrRequest { + extra_headers: vec![("aUtHoRiZaTiOn".into(), "Bearer supplied".into())], + ..make_request() + }) + .unwrap(); + assert_eq!( + prepared.headers, + vec![("aUtHoRiZaTiOn".into(), "Bearer supplied".into())] + ); + } + return; + } + for case in ["present", "absent"] { + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args([ + "--exact", + "cloud_credentials_use_native_keys_headers_or_narrow_acquisition_stub", + ]) + .env("LITELLM_CLOUD_OCR_ENV_TEST", case) + .env_remove("AZURE_AI_API_KEY") + .env_remove("VERTEX_AI_API_KEY") + .env_remove("VERTEXAI_API_KEY"); + if case == "present" { + command + .env("AZURE_AI_API_KEY", "env-azure") + .env("VERTEX_AI_API_KEY", "env-vertex"); + } + assert!(command.status().unwrap().success()); + } +} + +#[tokio::test] +async fn cloud_providers_use_existing_auth_urls_and_request_response_transforms() { + for (model, path, auth_name, auth_value, deepseek) in [ + ( + "azure_ai/mistral-ocr-latest", + "/providers/mistral/azure/ocr", + "api-key", + "test-key", + false, + ), + ( + "vertex_ai/mistral-ocr-latest", + "/v1/projects/test-project/locations/us-central1/publishers/mistralai/models/mistral-ocr-latest:rawPredict", + "authorization", + "Bearer test-key", + false, + ), + ( + "vertex_ai/deepseek-ai/deepseek-ocr-maas", + "/v1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions", + "authorization", + "Bearer test-key", + true, + ), + ] { + let response = if deepseek { + json!({"choices": [{"message": {"content": "proof"}}], "usage": {"pages_processed": 1}}) + } else { + json!({"pages": [{"index": 0, "markdown": "proof"}], "usage_info": {"pages_processed": 1}}) + }; + let (base, handle) = server(200, "", &response.to_string(), Duration::ZERO); + let prepared = prepare(OcrRequest { + model: model.into(), + api_base: Some(base), + ..request() + }) + .unwrap(); + let mut body = Value::Object(prepared.body.clone()); + body[if deepseek { + "temperature" + } else { + "include_image_base64" + }] = if deepseek { json!(0.1) } else { json!(true) }; + assert_eq!( + prepared.headers, + vec![( + if auth_name == "api-key" { + "Api-Key" + } else { + "Authorization" + } + .into(), + auth_value.into() + )] + ); + let headers = prepared.headers.clone(); + let response = ocr(prepared, headers, body.clone()).await.unwrap(); + let (headers, sent) = handle.join().unwrap(); + assert!(headers.starts_with(&format!("POST {path} HTTP/1.1\r\n"))); + assert!(headers.contains(&format!("{auth_name}: {auth_value}\r\n"))); + assert_eq!(sent, body); + if deepseek { + assert_eq!(sent["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + sent["messages"][0]["content"][0], + json!({"type": "image_url", "image_url": "data:application/pdf;base64,cGRm"}) + ); + assert!(sent.get("document").is_none()); + } else { + assert_eq!( + sent["document"]["document_url"], + "data:application/pdf;base64,cGRm" + ); + } + assert_eq!(response.pages[0]["markdown"], "proof"); + } +} + +fn server( + status: u16, + headers: &str, + response_body: &str, + delay: Duration, +) -> (String, thread::JoinHandle<(String, Value)>) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let response = format!( + "HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n{headers}\r\n{response_body}", + response_body.len() + ); + let handle = thread::spawn(move || { + listener.set_nonblocking(true).unwrap(); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + let mut stream = loop { + match listener.accept() { + Ok((stream, _)) => break stream, + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + assert!( + std::time::Instant::now() < deadline, + "no OCR request received" + ); + thread::sleep(Duration::from_millis(5)); + } + Err(error) => panic!("{error}"), + } + }; + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + stream + .set_write_timeout(Some(Duration::from_secs(5))) + .unwrap(); + let mut received = Vec::new(); + let mut buffer = [0; 4096]; + let header_end = loop { + let count = stream.read(&mut buffer).unwrap(); + assert_ne!(count, 0); + received.extend_from_slice(&buffer[..count]); + if let Some(end) = received.windows(4).position(|part| part == b"\r\n\r\n") { + break end + 4; + } + }; + let headers = String::from_utf8(received[..header_end].to_vec()).unwrap(); + let length: usize = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse().unwrap()) + }) + .unwrap(); + while received.len() < header_end + length { + let count = stream.read(&mut buffer).unwrap(); + assert_ne!(count, 0); + received.extend_from_slice(&buffer[..count]); + } + let body = serde_json::from_slice(&received[header_end..header_end + length]).unwrap(); + thread::sleep(delay); + let _ = stream.write_all(response.as_bytes()); + (headers, body) + }); + (base, handle) +} + +#[tokio::test] +async fn posts_filled_body_and_normalizes_provider_response() { + let response_json = json!({ + "pages": [{"index": 0, "markdown": "result", "header": "title", "blocks": []}], + "model": "provider-model", + "document_annotation": {"title": "result"}, + "usage_info": {"pages_processed": 1}, + "private_provider_field": "not forwarded" + }); + let (base, handle) = server(200, "", &response_json.to_string(), Duration::ZERO); + let prepared = prepare(OcrRequest { + api_base: Some(base), + ..request() + }) + .unwrap(); + let body = body(&prepared); + let headers = prepared + .headers + .iter() + .cloned() + .chain([("X-Retained".into(), "header".into())]) + .collect(); + let response = ocr(prepared, headers, body.clone()) + .await + .unwrap() + .into_json(); + let (headers, sent_body) = handle.join().unwrap(); + let headers = headers.to_ascii_lowercase(); + assert!(headers.starts_with("post /v1/ocr http/1.1\r\n")); + assert!(headers.contains("authorization: bearer test-key\r\n")); + assert!(headers.contains("accept-encoding: identity\r\n")); + assert!(!headers.contains("gzip")); + assert!(headers.contains("content-type: application/json\r\n")); + assert!(headers.contains("x-retained: header\r\n")); + assert_eq!(sent_body, body); + assert_eq!( + response, + json!({ + "pages": response_json["pages"], + "model": "provider-model", + "document_annotation": response_json["document_annotation"], + "usage_info": response_json["usage_info"], + "object": "ocr" + }) + ); +} + +#[tokio::test] +async fn preserves_callback_body_and_header_changes() { + let (base, handle) = server(200, "", "{}", Duration::ZERO); + let prepared = prepare(OcrRequest { + api_base: Some(base), + ..request() + }) + .unwrap(); + let mut body = body(&prepared); + body["model"] = json!("callback-model"); + body["custom_provider_field"] = json!({"nested": [1, true, null]}); + let headers = prepared + .headers + .iter() + .cloned() + .chain([ + ("cOnTeNt-TyPe".into(), "application/vnd.ocr+json".into()), + ("aCcEpT-EnCoDiNg".into(), "gzip, br".into()), + ]) + .collect(); + ocr(prepared, headers, body.clone()).await.unwrap(); + let (headers, sent_body) = handle.join().unwrap(); + let headers = headers.to_ascii_lowercase(); + assert_eq!(sent_body, body); + assert_eq!( + headers + .lines() + .filter(|line| line.starts_with("content-type:")) + .collect::>(), + vec!["content-type: application/vnd.ocr+json"] + ); + assert_eq!( + headers + .lines() + .filter(|line| line.starts_with("accept-encoding:")) + .collect::>(), + vec!["accept-encoding: gzip, br"] + ); +} + +#[tokio::test] +async fn rejects_unsupported_inputs_before_io() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + for case in ["file", "local", "provider", "compression"] { + let mut prepared = prepare(OcrRequest { + api_base: Some(base.clone()), + ..request() + }) + .unwrap(); + let mut body = body(&prepared); + let mut headers = prepared.headers.clone(); + match case { + "file" => body["document"] = json!({"type": "file", "file": "private"}), + "local" => body["document"]["document_url"] = json!("file:///private.pdf"), + "provider" => prepared.custom_llm_provider = "cohere".into(), + "compression" => headers.push(("Content-Encoding".into(), "gzip".into())), + _ => unreachable!(), + } + assert!(matches!( + ocr(prepared, headers, body).await, + Err(Error::Unsupported(_)) + )); + assert_eq!( + listener.accept().unwrap_err().kind(), + std::io::ErrorKind::WouldBlock + ); + } +} + +#[tokio::test] +async fn handles_errors_compression_and_timeout_without_exposing_payloads() { + for (status, headers, response_body, delay, expected) in [ + ( + 401, + "", + "private-document api-key", + Duration::ZERO, + Error::Http { + status: 401, + body: "OCR provider request failed".into(), + }, + ), + ( + 302, + "Location: http://127.0.0.1:1/private\r\n", + "secret", + Duration::ZERO, + Error::Http { + status: 302, + body: "OCR provider request failed".into(), + }, + ), + ( + 200, + "", + "private-document invalid JSON", + Duration::ZERO, + Error::InvalidResponse("invalid OCR JSON response".into()), + ), + ( + 200, + "Content-Encoding: gzip\r\n", + "{}", + Duration::ZERO, + Error::Unsupported("compressed OCR response"), + ), + ( + 200, + "Content-Encoding: identity, br\r\n", + "{}", + Duration::ZERO, + Error::Unsupported("compressed OCR response"), + ), + ( + 200, + "", + "{}", + Duration::from_millis(200), + Error::Network("transport failed".into()), + ), + ] { + let (base, handle) = server(status, headers, response_body, delay); + let prepared = prepare(OcrRequest { + api_base: Some(base), + timeout_seconds: if delay.is_zero() { 2.0 } else { 0.05 }, + ..request() + }) + .unwrap(); + let body = body(&prepared); + let headers = prepared.headers.clone(); + assert_eq!(ocr(prepared, headers, body).await.unwrap_err(), expected); + handle.join().unwrap(); + } +} + +#[tokio::test] +async fn normalizes_missing_fields_and_accepts_identity_response() { + let (base, handle) = server(200, "Content-Encoding: Identity\r\n", "{}", Duration::ZERO); + let prepared = prepare(OcrRequest { + api_base: Some(base), + ..request() + }) + .unwrap(); + let body = body(&prepared); + let headers = prepared.headers.clone(); + let response = ocr(prepared, headers, body).await.unwrap().into_json(); + handle.join().unwrap(); + assert_eq!( + response, + json!({ + "pages": [], "model": "mistral-ocr-latest", "document_annotation": null, + "usage_info": null, "object": "ocr" + }) + ); +} diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 594de0220bd..1e66a207592 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -222,7 +222,7 @@ mod tests { let module = PyModule::new(py, "routes").expect("module should be created"); crate::routes::register(&module).expect("routes should register"); let routes = [ - ("ocr", "aocr", "(boundary)"), + ("ocr", "aocr", "(arguments)"), ( "transcription", "atranscription", diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 172b8de48a7..7ff84273fbb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -1,153 +1,352 @@ -//! Retained OCR route: Python owns request/response objects, Rust sequences -//! prepare -> encode -> POST -> finish through owning `Py` handles. +//! Native OCR retains the entire Python argument graph through completion. +//! Missing native operations raise NotImplementedError before +//! callbacks; no Python preparation, auth, encoding, or provider transforms run. -use litellm_core::http_utils::buffered_post::{self, Request, Response}; -use litellm_python_interop::{InvocationMode, InvocationOutcome, PreparedCall}; +use litellm_core::error::Error; +use litellm_core::ocr::types::{OcrDocumentProjection, OcrRequest, PreparedOcr}; +use litellm_core::routing_utils::provider::get_custom_llm_provider; +use litellm_python_interop::{Pythonized, from_py, to_py}; +use pyo3::exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}; use pyo3::prelude::*; +use pyo3::pyclass::{PyTraverseError, PyVisit}; use pyo3::sync::PyOnceLock; -use pyo3::types::{PyBytes, PyList, PyTuple}; +use pyo3::types::PyDict; +use serde_json::Value; use crate::errors::core_error_to_pyerr; use crate::execution::{run_async_value, run_sync_value}; -#[derive(Clone, Copy)] -struct BoundaryStep { - method: &'static str, - awaited: bool, +#[pyclass] +struct OcrState { + arguments: Option>, + body: Option>, + headers: Option>, + logging: Option>, + prepared: Option, } -const PREPARE_SYNC: BoundaryStep = BoundaryStep { - method: "prepare", - awaited: false, -}; -const PREPARE_ASYNC: BoundaryStep = BoundaryStep { - method: "aprepare", - awaited: true, -}; -const ENCODE: BoundaryStep = BoundaryStep { - method: "encode", - awaited: false, -}; -const FINISH_SYNC: BoundaryStep = BoundaryStep { - method: "finish", - awaited: false, -}; -const FINISH_ASYNC: BoundaryStep = BoundaryStep { - method: "afinish", - awaited: true, -}; +#[pymethods] +impl OcrState { + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.arguments)?; + visit.call(&self.body)?; + visit.call(&self.headers)?; + visit.call(&self.logging) + } -fn invoke( - boundary: &Bound<'_, PyAny>, - step: BoundaryStep, - args: Bound<'_, PyTuple>, -) -> PyResult> { - let call = PreparedCall::new( - if step.awaited { - InvocationMode::Await - } else { - InvocationMode::Direct - }, - boundary.getattr(step.method)?.unbind(), - args.unbind(), - None, - ); - match call.invoke(boundary.py())? { - InvocationOutcome::Returned(value) | InvocationOutcome::Awaitable(value) => Ok(value), + fn __clear__(slf: &Bound<'_, Self>) { + let roots = { + let mut state = slf.borrow_mut(); + ( + state.arguments.take(), + state.body.take(), + state.headers.take(), + state.logging.take(), + ) + }; + drop(roots); } } -#[pyfunction] -fn prepare(boundary: &Bound<'_, PyAny>, asynchronous: bool) -> PyResult> { - let step = if asynchronous { - PREPARE_ASYNC - } else { - PREPARE_SYNC +fn ocr_error_to_pyerr(py: Python<'_>, error: Error, model: &str, provider: &str) -> PyErr { + let status = match error { + Error::Unsupported(message) => return PyNotImplementedError::new_err(message), + Error::Auth(_) => 401, + Error::Http { status, .. } => status, + Error::Network(_) | Error::Connect(_) => { + return PyRuntimeError::new_err("OCR transport failed"); + } + Error::InvalidResponse(_) => { + return PyRuntimeError::new_err("Invalid OCR provider response"); + } + other => return core_error_to_pyerr(other), }; - invoke(boundary, step, PyTuple::empty(boundary.py())) -} - -fn request(boundary: &Bound<'_, PyAny>, roots: &Bound<'_, PyAny>) -> PyResult { - type ByteHeaders<'py> = Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>; - let py = boundary.py(); - let encoded = invoke(boundary, ENCODE, PyTuple::new(py, [roots])?)?; - let (url, headers, body, timeout_seconds): (String, ByteHeaders<'_>, Bound<'_, PyBytes>, f64) = - encoded.into_bound(py).extract()?; - Ok(Request { - url, - headers: headers - .into_iter() - .map(|(name, value)| (name.as_bytes().to_vec(), value.as_bytes().to_vec())) - .collect(), - body: body.as_bytes().to_vec(), - timeout_seconds, - }) -} - -struct Wire(Response); - -impl<'py> IntoPyObject<'py> for Wire { - type Target = PyTuple; - type Output = Bound<'py, PyTuple>; - type Error = PyErr; - - fn into_pyobject(self, py: Python<'py>) -> PyResult { - let headers = PyList::new( - py, - self.0 - .headers - .iter() - .map(|(name, value)| (PyBytes::new(py, name), PyBytes::new(py, value))), + let class = match status { + 400 => "BadRequestError", + 401 => "AuthenticationError", + 403 => "PermissionDeniedError", + 404 => "NotFoundError", + 422 => "UnprocessableEntityError", + 429 => "RateLimitError", + 500 => "InternalServerError", + 502 => "BadGatewayError", + 503 => "ServiceUnavailableError", + _ => "APIError", + }; + let exception = || -> PyResult { + let kwargs = PyDict::new(py); + kwargs.set_item( + "message", + format!("OCR provider request failed (HTTP {status})"), )?; - (self.0.status, headers, PyBytes::new(py, &self.0.content)).into_pyobject(py) - } + kwargs.set_item("model", model)?; + kwargs.set_item("llm_provider", provider)?; + if class == "APIError" { + kwargs.set_item("status_code", status)?; + } else { + let httpx = py.import("httpx")?; + let request = httpx + .getattr("Request")? + .call1(("POST", "https://litellm.ai"))?; + let response_kwargs = PyDict::new(py); + response_kwargs.set_item("request", request)?; + let response = httpx + .getattr("Response")? + .call((status,), Some(&response_kwargs))?; + kwargs.set_item("response", response)?; + } + let instance = py + .import("litellm.exceptions")? + .getattr(class)? + .call((), Some(&kwargs))?; + Ok(PyErr::from_value(instance)) + }; + exception().unwrap_or_else(|error| error) +} + +fn scalar(arguments: &Bound<'_, PyDict>, name: &str) -> PyResult> { + arguments + .get_item(name)? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose() + .map(|value| value.filter(|value| !value.trim().is_empty())) +} + +fn header_pairs(headers: &Bound<'_, PyDict>) -> PyResult> { + headers + .iter() + .map(|(name, value)| Ok((name.extract()?, value.extract()?))) + .collect() } #[pyfunction] -fn send<'a>( - boundary: &'a Bound<'a, PyAny>, - roots: &Bound<'_, PyAny>, -) -> PyResult> { - let request = request(boundary, roots)?; - pyo3_async_runtimes::tokio::future_into_py(boundary.py(), async move { - let response = run_async_value(buffered_post::send(request), core_error_to_pyerr).await?; - Ok(Wire(response)) +#[pyo3(signature = (arguments, asynchronous=false))] +fn prepare(py: Python<'_>, arguments: Py, asynchronous: bool) -> PyResult> { + let bag = arguments.bind(py); + let document = bag + .get_item("document")? + .ok_or_else(|| PyValueError::new_err("OCR requires document"))? + .cast_into::()?; + if scalar(&document, "type")?.as_deref() == Some("file") { + return Err(ocr_error_to_pyerr( + py, + Error::Unsupported( + "Native OCR does not support file documents; pass a document_url or image_url dict", + ), + "", + "", + )); + } + let timeout_seconds = match bag.get_item("timeout")?.filter(|value| !value.is_none()) { + None => py + .import("litellm.constants")? + .getattr("request_timeout")? + .extract::()?, + Some(timeout) => match timeout.extract::() { + Ok(seconds) => seconds, + Err(_) => timeout.getattr("read")?.extract::()?, + }, + }; + let extra_headers = match bag + .get_item("extra_headers")? + .filter(|value| !value.is_none()) + { + Some(headers) => header_pairs(headers.cast::()?)?, + None => Vec::new(), + }; + let model = scalar(bag, "model")?.ok_or_else(|| PyValueError::new_err("OCR requires model"))?; + let custom_llm_provider = scalar(bag, "custom_llm_provider")?; + let document_input = PyDict::new(py); + for name in ["type", "document_url", "image_url"] { + if let Some(value) = document.get_item(name)? { + document_input.set_item(name, value)?; + } + } + let prepared = litellm_core::ocr::prepare::prepare(OcrRequest { + model: model.clone(), + custom_llm_provider: custom_llm_provider.clone(), + api_key: scalar(bag, "api_key")?, + api_base: scalar(bag, "api_base")?, + extra_headers, + timeout_seconds, + request_format: scalar(bag, "req_format")?, + document: from_py(document_input.as_any())?, + azure_ad_token: scalar(bag, "azure_ad_token")?, + vertex_project: scalar(bag, "vertex_project")?.or(scalar(bag, "vertex_ai_project")?), + vertex_location: scalar(bag, "vertex_location")?.or(scalar(bag, "vertex_ai_location")?), + stream: bag + .get_item("stream")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()? + .unwrap_or(false), + }) + .map_err(|error| { + let resolved = get_custom_llm_provider(&model, custom_llm_provider.as_deref()); + ocr_error_to_pyerr( + py, + error, + resolved + .as_ref() + .map_or(model.as_str(), |value| value.model), + resolved + .as_ref() + .map_or("", |value| value.custom_llm_provider), + ) + })?; + + let body = to_py(py, &prepared.body)? + .into_bound(py) + .cast_into::()?; + match prepared.document_projection { + OcrDocumentProjection::RetainedDocument => body.set_item("document", &document)?, + OcrDocumentProjection::ShallowCopyDocument => { + body.set_item("document", document.copy()?)? + } + OcrDocumentProjection::Transformed => {} + } + let optional_params = PyDict::new(py); + for &name in prepared.parameter_fields { + if let Some(value) = bag.get_item(name)? { + body.set_item(name, &value)?; + optional_params.set_item(name, value)?; + } + } + let headers = PyDict::new(py); + for (name, value) in &prepared.headers { + headers.set_item(name, value)?; + } + let logging = py + .import("litellm.rust_bridge.ocr")? + .getattr("initialize_logging")? + .call1((bag, asynchronous))?; + let litellm_params = PyDict::new(py); + litellm_params.set_item("litellm_call_id", bag.get_item("litellm_call_id")?)?; + litellm_params.set_item("api_base", bag.get_item("api_base")?)?; + let update = PyDict::new(py); + update.set_item("kwargs", bag)?; + update.set_item("model", &prepared.model)?; + update.set_item("optional_params", optional_params)?; + update.set_item("litellm_params", litellm_params)?; + update.set_item("custom_llm_provider", &prepared.custom_llm_provider)?; + logging.call_method("update_from_kwargs", (), Some(&update))?; + + let additional_args = PyDict::new(py); + additional_args.set_item("complete_input_dict", &body)?; + additional_args.set_item("api_base", &prepared.url)?; + additional_args.set_item("headers", &headers)?; + let pre_call = PyDict::new(py); + pre_call.set_item("input", "OCR document processing")?; + pre_call.set_item("api_key", bag.get_item("api_key")?)?; + pre_call.set_item("additional_args", additional_args)?; + logging.call_method("pre_call", (), Some(&pre_call))?; + + let logging = logging.unbind(); + Py::new( + py, + OcrState { + arguments: Some(arguments), + body: Some(body.unbind()), + headers: Some(headers.unbind()), + logging: Some(logging), + prepared: Some(prepared), + }, + ) +} + +type OcrWireRequest = (PreparedOcr, Vec<(String, String)>, Value); + +fn request(py: Python<'_>, state: &Py) -> PyResult { + let (prepared, body, headers) = { + let mut state = state.borrow_mut(py); + let prepared = state + .prepared + .take() + .ok_or_else(|| PyRuntimeError::new_err("OCR request was already sent or cleared"))?; + let body = state + .body + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("OCR body was cleared"))? + .clone_ref(py); + let headers = state + .headers + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("OCR headers were cleared"))? + .clone_ref(py); + (prepared, body, headers) + }; + Ok(( + prepared, + header_pairs(headers.bind(py))?, + from_py(body.bind(py).as_any())?, + )) +} + +#[pyfunction] +fn send(py: Python<'_>, state: Py) -> PyResult> { + let (prepared, headers, body) = request(py, &state)?; + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let _state = state; + let model = prepared.model.clone(); + let provider = prepared.custom_llm_provider.clone(); + let response = run_async_value( + async move { Ok(litellm_core::ocr::ocr(prepared, headers, body).await) }, + core_error_to_pyerr, + ) + .await? + .map_err(|error| Python::attach(|py| ocr_error_to_pyerr(py, error, &model, &provider)))?; + Ok(Pythonized(response.into_json())) }) } #[pyfunction] -fn finish( - boundary: &Bound<'_, PyAny>, - wire: &Bound<'_, PyAny>, - asynchronous: bool, -) -> PyResult> { - let step = if asynchronous { - FINISH_ASYNC - } else { - FINISH_SYNC - }; - invoke(boundary, step, PyTuple::new(boundary.py(), [wire])?) +fn finish(py: Python<'_>, response: Py) -> PyResult> { + let fields = response.bind(py); + let native_response = fields.get_item("provider_native_response")?; + if native_response.is_some() { + fields.del_item("provider_native_response")?; + } + let response = py + .import("litellm.llms.base_llm.ocr.transformation")? + .getattr("OCRResponse")? + .call((), Some(fields))?; + if let Some(native_response) = native_response.filter(|value| !value.is_none()) { + response.call_method1("set_provider_native_response", (native_response,))?; + } + Ok(response.unbind()) } #[pyfunction] -fn ocr(boundary: &Bound<'_, PyAny>) -> PyResult> { - let py = boundary.py(); - let roots = prepare(boundary, false)?; - let request = request(boundary, roots.bind(py))?; - let response = run_sync_value(py, buffered_post::send(request), core_error_to_pyerr)?; - let wire = Wire(response).into_pyobject(py)?; - finish(boundary, &wire, false) +fn send_sync(py: Python<'_>, state: Py) -> PyResult> { + let (prepared, headers, body) = request(py, &state)?; + let model = prepared.model.clone(); + let provider = prepared.custom_llm_provider.clone(); + let response = run_sync_value( + py, + async move { Ok(litellm_core::ocr::ocr(prepared, headers, body).await) }, + core_error_to_pyerr, + )? + .map_err(|error| ocr_error_to_pyerr(py, error, &model, &provider))?; + let fields = to_py(py, &response.into_json())? + .into_bound(py) + .cast_into::()?; + let response = finish(py, fields.unbind()); + drop(state); + response } #[pyfunction] -fn aocr<'a>(boundary: &'a Bound<'a, PyAny>) -> PyResult> { - driver(boundary.py())?.getattr("drive")?.call1((boundary,)) +fn ocr(py: Python<'_>, arguments: Py) -> PyResult> { + driver(py)?.getattr("drive_sync")?.call1((arguments,)) } -/// The async route must await `aprepare`/`afinish` inline in the caller's -/// Python task, so a Python driver coroutine owns the roots between steps. -/// Compiling the driver runs Python (audit hooks can re-enter `aocr`), so -/// compile first and publish only a finished module into the once-lock. +#[pyfunction] +fn aocr(py: Python<'_>, arguments: Py) -> PyResult> { + driver(py)?.getattr("drive")?.call1((arguments,)) +} + +// Compilation can re-enter through audit hooks; publish only a finished module. fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { static DRIVER: PyOnceLock> = PyOnceLock::new(); if let Some(module) = DRIVER.get(py) { @@ -155,17 +354,65 @@ fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { } let module = PyModule::from_code( py, - c"async def drive(boundary): - roots = await _prepare(boundary, True) - wire = await _send(boundary, roots) - return await _finish(boundary, wire, True) + c"from datetime import datetime +from litellm.rust_bridge.ocr import invoke_terminal + +def drive_sync(arguments): + start = datetime.now() + state = None + try: + state = _prepare(arguments, False) + response = _send_sync(state) + except Exception as error: + logger = arguments.get('litellm_logging_obj') + if state is not None or (logger is not None and not isinstance(error, NotImplementedError)): + end = datetime.now() + for action in _sync_failure: + invoke_terminal(action, (arguments, state), logger, error, start, end) + raise + end = datetime.now() + for action in _sync_success: + invoke_terminal(action, (arguments, state), arguments['litellm_logging_obj'], response, start, end) + return response + +async def drive(arguments): + start = datetime.now() + state = None + try: + state = _prepare(arguments, True) + response = _finish(await _send(state)) + except Exception as error: + logger = arguments.get('litellm_logging_obj') + if state is not None or (logger is not None and not isinstance(error, NotImplementedError)): + end = datetime.now() + for action in _async_failure: + pending = invoke_terminal(action, (arguments, state), logger, error, start, end) + if pending is not None: + await pending + raise + end = datetime.now() + for action in _async_success: + invoke_terminal(action, (arguments, state), arguments['litellm_logging_obj'], response, start, end) + return response ", c"ocr_driver.py", c"_ocr_driver", )?; module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; module.add("_send", wrap_pyfunction!(send, &module)?)?; + module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?; module.add("_finish", wrap_pyfunction!(finish, &module)?)?; + for (name, asynchronous, success) in [ + ("_sync_success", false, true), + ("_async_success", true, true), + ("_sync_failure", false, false), + ("_async_failure", true, false), + ] { + module.add( + name, + litellm_core::ocr::terminal_callbacks(asynchronous, success).to_vec(), + )?; + } Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) } @@ -179,3 +426,404 @@ pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { register(module) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[ignore = "requires the Python SDK and its dependencies on PYTHONPATH"] + fn structured_errors_use_public_sdk_exceptions() { + Python::initialize(); + Python::attach(|py| { + let exceptions = py.import("litellm.exceptions").unwrap(); + for (status, class) in [ + (400, "BadRequestError"), + (401, "AuthenticationError"), + (403, "PermissionDeniedError"), + (404, "NotFoundError"), + (422, "UnprocessableEntityError"), + (429, "RateLimitError"), + (500, "InternalServerError"), + (502, "BadGatewayError"), + (503, "ServiceUnavailableError"), + (504, "APIError"), + ] { + let error = ocr_error_to_pyerr( + py, + Error::Http { + status, + body: "private upstream content".into(), + }, + "mistral-ocr-latest", + "mistral", + ); + let value = error.value(py); + assert!( + value + .is_instance(&exceptions.getattr(class).unwrap()) + .unwrap(), + "HTTP {status}: expected {class}, got {error}" + ); + assert_eq!( + value + .getattr("status_code") + .unwrap() + .extract::() + .unwrap(), + status + ); + assert_eq!( + value.getattr("model").unwrap().extract::().unwrap(), + "mistral-ocr-latest" + ); + assert_eq!( + value + .getattr("llm_provider") + .unwrap() + .extract::() + .unwrap(), + "mistral" + ); + assert!(!error.to_string().contains("private upstream content")); + } + let error = ocr_error_to_pyerr( + py, + Error::Auth("private credential".into()), + "model", + "mistral", + ); + assert!( + error + .value(py) + .is_instance(&exceptions.getattr("AuthenticationError").unwrap()) + .unwrap() + ); + assert!(!error.to_string().contains("private credential")); + }); + } + + #[test] + fn unstructured_transport_errors_are_not_guessed_from_strings() { + Python::initialize(); + Python::attach(|py| { + for error in [ + Error::Network("timeout secret".into()), + Error::Connect("401 secret".into()), + Error::InvalidResponse("429 secret".into()), + ] { + let error = ocr_error_to_pyerr(py, error, "model", "mistral"); + assert!(error.is_instance_of::(py)); + assert!(!error.to_string().contains("secret")); + } + }); + } + + #[test] + fn native_send_owns_state_without_the_python_driver() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "ocr_test").unwrap(); + module + .add_function(wrap_pyfunction!(prepare, &module).unwrap()) + .unwrap(); + module + .add_function(wrap_pyfunction!(send, &module).unwrap()) + .unwrap(); + let globals = PyDict::new(py); + globals.set_item("native", module).unwrap(); + py.run( + cr" +import asyncio +import gc +import weakref + +class Logger: + def update_from_kwargs(self, **values): + pass + def pre_call(self, **values): + pass + +async def exercise(): + received = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + + async def respond(reader, writer): + await reader.readuntil(b'\r\n\r\n') + received.set() + await release.wait() + writer.close() + await writer.wait_closed() + closed.set() + + server = await asyncio.start_server(respond, '127.0.0.1', 0) + async with server: + port = server.sockets[0].getsockname()[1] + logger = Logger() + alive = weakref.ref(logger) + state = native.prepare(dict( + model='mistral/mistral-ocr-latest', api_key='test-key', timeout=5.0, + api_base=f'http://127.0.0.1:{port}', litellm_logging_obj=logger, + document={'type': 'document_url', 'document_url': 'https://example.test/doc.pdf'}, + )) + pending = native.send(state) + del state, logger + try: + await asyncio.wait_for(received.wait(), 5) + gc.collect() + assert alive() is not None + pending.cancel() + try: + await pending + except asyncio.CancelledError: + pass + for _ in range(500): + await asyncio.sleep(0.01) + gc.collect() + if alive() is None: + break + assert alive() is None + finally: + release.set() + await asyncio.wait_for(closed.wait(), 5) + +asyncio.run(exercise()) +", + Some(&globals), + Some(&globals), + ) + .unwrap(); + }); + } + + #[pyfunction] + fn snapshot(py: Python<'_>, state: Py) -> PyResult> { + let (_, headers, body) = request(py, &state)?; + to_py(py, &(headers, body)) + } + + #[test] + fn retains_identity_independent_wire_roots_and_collects_cycles() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "ocr_test").unwrap(); + module + .add_function(wrap_pyfunction!(prepare, &module).unwrap()) + .unwrap(); + module + .add_function(wrap_pyfunction!(snapshot, &module).unwrap()) + .unwrap(); + let globals = PyDict::new(py); + globals.set_item("native", module).unwrap(); + py.run( + c" +import gc +import weakref + +class Opaque: + pass + +class Timeout: + read = 5.0 + +class Logger: + def update_from_kwargs(self, **values): + assert values['kwargs'] is arguments + assert values['kwargs']['metadata'] is metadata + assert values['kwargs']['opaque'] is opaque + assert values['optional_params']['pages'] is pages + self.calls = ['update'] + + def pre_call(self, **values): + self.calls.append('pre') + view = values['additional_args'] + self.body = view['complete_input_dict'] + self.headers = view['headers'] + assert self.body['document'] is document + assert self.body['pages'] is pages + document['document_url'] = 'https://example.test/changed.pdf' + pages.append(2) + self.headers['x-hook'] = 'changed' + view['complete_input_dict'] = {'replacement': True} + view['headers'] = {'replacement': 'true'} + +document = {'type': 'document_url', 'document_url': 'https://example.test/test.pdf'} +pages = [0] +metadata = {'nested': []} +opaque = Opaque() +logger = Logger() +arguments = dict(model='mistral/mistral-ocr-latest', document=document, + api_key='test-key', pages=pages, metadata=metadata, + opaque=opaque, litellm_logging_obj=logger, timeout=Timeout()) +state = native.prepare(arguments) +assert logger.calls == ['update', 'pre'] +roots = gc.get_referents(state) +assert any(root is arguments for root in roots) +assert any(root is logger for root in roots) +assert any(root is logger.body for root in roots) +assert any(root is logger.headers for root in roots) +headers, body = native.snapshot(state) +assert body['document']['document_url'] == 'https://example.test/changed.pdf' +assert body['pages'] == [0, 2] +assert dict(headers)['x-hook'] == 'changed' +assert 'replacement' not in body and 'replacement' not in dict(headers) + +arguments['cycle'] = state +logger.cycle = state +logger.body['cycle'] = state +logger.headers['cycle'] = state +alive = weakref.ref(opaque) +del roots, arguments, logger, opaque +gc.collect() +assert alive() is not None +del state +gc.collect() +assert alive() is None +", + Some(&globals), + Some(&globals), + ) + .unwrap(); + }); + } + + #[test] + fn async_callbacks_are_inline_and_unsupported_requests_never_call_them() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "ocr_test").unwrap(); + module + .add_function(wrap_pyfunction!(ocr, &module).unwrap()) + .unwrap(); + module + .add_function(wrap_pyfunction!(aocr, &module).unwrap()) + .unwrap(); + let globals = PyDict::new(py); + globals.set_item("native", module).unwrap(); + py.run( + cr" +import asyncio +import contextvars +import gc +import json +import threading +import weakref + +marker = contextvars.ContextVar('ocr_marker') + +class Opaque: + pass + +class Logger: + def __init__(self): + self.calls = [] + + def failure_handler(self, error, trace, start, end): + assert asyncio.current_task() is caller + assert marker.get() == 'pre' + self.error = error + + async def async_failure_handler(self, error, trace, start, end): + await asyncio.sleep(0) + assert asyncio.current_task() is caller + assert error is self.error + + def update_from_kwargs(self, **values): + assert asyncio.current_task() is caller + assert threading.get_ident() == caller_thread + assert marker.get() == 'caller' + assert values['kwargs'] is arguments + self.calls.append('update') + marker.set('updated') + + def pre_call(self, **values): + assert asyncio.current_task() is caller + assert threading.get_ident() == caller_thread + assert marker.get() == 'updated' + self.calls.append('pre') + values['additional_args']['complete_input_dict']['pages'].append(3) + marker.set('pre') + +async def exercise(): + global arguments, caller, caller_thread + caller = asyncio.current_task() + caller_thread = threading.get_ident() + marker.set('caller') + logger = Logger() + document = {'type': 'document_url', 'document_url': 'https://example.test/test.pdf'} + arguments = dict(model='mistral/mistral-ocr-latest', document=document, + api_key='test-key', pages=[0], opaque=Opaque(), + litellm_logging_obj=logger) + alive = weakref.ref(arguments['opaque']) + received = asyncio.Event() + errors = [] + + async def respond(reader, writer): + try: + header = await reader.readuntil(b'\r\n\r\n') + length = next(int(line.split(b':', 1)[1]) for line in header.split(b'\r\n') + if line.lower().startswith(b'content-length:')) + body = json.loads(await reader.readexactly(length)) + assert body['pages'] == [0, 3] + assert 'opaque' not in body + gc.collect() + assert alive() is not None + assert logger.calls == ['update', 'pre'] + globals().pop('arguments') + gc.collect() + assert alive() is not None + except BaseException as error: + errors.append(error) + finally: + writer.write(b'HTTP/1.1 200 OK\r\nContent-Length: 1\r\nConnection: close\r\n\r\nx') + await writer.drain() + writer.close() + await writer.wait_closed() + received.set() + + server = await asyncio.start_server(respond, '127.0.0.1', 0) + async with server: + port = server.sockets[0].getsockname()[1] + arguments['api_base'] = f'http://127.0.0.1:{port}' + arguments['timeout'] = 5.0 + try: + await native.aocr(arguments) + except RuntimeError as error: + assert error is logger.error + else: + raise AssertionError('expected upstream error') + await asyncio.wait_for(received.wait(), 5) + assert not errors, errors + assert marker.get() == 'pre' + assert logger.calls == ['update', 'pre'] + + for model, doc in [ + ('azure_ai/doc-intelligence/prebuilt-read', document), + ('vertex_ai/ocr', document), + ('mistral/mistral-ocr-latest', {'type': 'file', 'file': Opaque()}), + ]: + logger.calls.clear() + unsupported = dict(model=model, document=doc, api_key='test-key', timeout=5.0, + litellm_logging_obj=logger) + for asynchronous in (False, True): + try: + if asynchronous: + await native.aocr(unsupported) + else: + native.ocr(unsupported) + except NotImplementedError: + pass + else: + raise AssertionError('expected strict unsupported error') + assert logger.calls == [] + +asyncio.run(exercise()) +", + Some(&globals), + Some(&globals), + ) + .unwrap(); + }); + } +} diff --git a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs index 9a87d19c17f..5f50a68a2a2 100644 --- a/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs +++ b/litellm-rust/crates/python-interop/tests/callback_lifecycle.rs @@ -113,8 +113,38 @@ fn lifecycle_contract( } #[rstest] +#[case::identity_and_ignored_returns("pre_call_identity_and_ignored_returns")] +#[case::mutations_visible_to_later_callbacks("pre_call_mutations_visible_to_later_callbacks")] +#[case::mutation_survives_failure("pre_call_mutation_survives_failure")] +#[ignore = "requires the repository Python environment and LiteLLM on PYTHONPATH"] +#[serial(python_interpreter)] +fn pre_call_contract( + scenario_scope: Py, + #[case] scenario: &str, + #[values(false, true)] retained: bool, +) -> PyResult<()> { + run_scenario_fixture( + scenario_scope, + scenario, + retained, + Some(( + include_str!("fixtures/callback_components.py"), + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/callback_components.py" + ), + )), + ) +} + +#[rstest] +#[case::real_post_call_logging("real_post_call_logging")] +#[case::real_post_call_dict_response("real_post_call_dict_response")] +#[case::real_sync_logging("real_sync_logging")] +#[case::real_sync_logging_hook_failure("real_sync_logging_hook_failure")] +#[case::real_sync_failure_chain("real_sync_failure_chain")] +#[case::real_async_failure_chain("real_async_failure_chain")] #[case::real_async_logging("real_async_logging")] -#[case::real_pre_call_logging("real_pre_call_logging")] #[case::real_copy_boundaries("real_copy_boundaries")] #[case::real_logging_worker("real_logging_worker")] #[case::real_sync_stream_copies("real_sync_stream_copies")] diff --git a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py index d5379a9dc98..59469a6545e 100644 --- a/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py +++ b/litellm-rust/crates/python-interop/tests/fixtures/callback_components.py @@ -17,36 +17,68 @@ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import Delta, ModelResponse, ModelResponseStream, StreamingChoices, Usage -def logger_for(callbacks=(), stream=False, input_callbacks=(), sync_callbacks=()): +def logger_for( + callbacks=(), + stream=False, + input_callbacks=(), + sync_callbacks=(), + failure_callbacks=(), + async_failure_callbacks=(), + call_type="acompletion", +): return Logging( model="test", messages=[{"role": "user", "content": "test"}], stream=stream, - call_type="acompletion", + call_type=call_type, start_time=datetime.now(), litellm_call_id="retained-test", function_id="retained-test", dynamic_async_success_callbacks=list(callbacks), dynamic_input_callbacks=list(input_callbacks), dynamic_success_callbacks=list(sync_callbacks), + dynamic_failure_callbacks=list(failure_callbacks), + dynamic_async_failure_callbacks=list(async_failure_callbacks), ) -async def real_pre_call_logging(owners): - retained = [] - observed = [] - snapshots = [] - order = [] +def invoke_pre_call(owners, logger, additional): + owner = owners.prepare(logger.pre_call, (logger.messages, "test-key"), {"additional_args": additional}) + try: + return owner.invoke() + finally: + owner.close() + + +async def pre_call_identity_and_ignored_returns(owners): + saved = [] ignored = {"replacement": True} + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + saved.append((kwargs, messages)) + return ignored + + logger = logger_for(input_callbacks=[Retain(), Retain()]) + details = logger.model_call_details + additional = {"headers": {"test": "header"}} + assert invoke_pre_call(owners, logger, additional) is None + assert len(saved) == 2 + assert saved[0][0] is saved[1][0] is details + assert saved[0][1] is saved[1][1] is logger.messages is details["input"] + assert details["additional_args"] is additional + assert "replacement" not in details + + +async def pre_call_mutations_visible_to_later_callbacks(owners): + saved, observed, order = [], [], [] metadata = {"secret": "private", "keep": []} removed = object() - lock = threading.Lock() class Retain(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): order.append("retain") - retained.append(kwargs) - return ignored + saved.append(kwargs) class Mutate(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): @@ -54,7 +86,30 @@ async def real_pre_call_logging(owners): kwargs["normalized"] = "normalized" assert kwargs.pop("remove") is removed kwargs["retained_metadata"]["secret"] = "masked" - return ignored + return {"replacement": True} + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + order.append("observe") + observed.append((kwargs["normalized"], "remove" in kwargs, kwargs["retained_metadata"]["secret"])) + + logger = logger_for(input_callbacks=[Retain(), Mutate(), Observe()]) + details = logger.model_call_details + details.update(retained_metadata=metadata, normalized=None, remove=removed) + assert invoke_pre_call(owners, logger, {}) is None + assert order == ["retain", "mutate", "observe"] + assert observed == [("normalized", False, "masked")] + assert len(saved) == 1 and saved[0] is details + assert details["retained_metadata"] is metadata + assert metadata == {"secret": "masked", "keep": []} + assert details["normalized"] == "normalized" and "remove" not in details + assert "replacement" not in details + + +async def pre_call_mutation_survives_failure(owners): + observed, order = [], [] + metadata = {"keep": []} + lock = threading.Lock() class Fail(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): @@ -66,40 +121,292 @@ async def real_pre_call_logging(owners): class Observe(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): order.append("observe") - snapshots.append( - ( - kwargs.get("normalized"), - "remove" in kwargs, - kwargs["retained_metadata"]["secret"], - tuple(kwargs["retained_metadata"]["keep"]), - "lock" in kwargs, - "replacement" in kwargs, - ) - ) - observed.append((kwargs, messages)) + observed.append((kwargs, tuple(kwargs["retained_metadata"]["keep"]), kwargs["lock"])) - logger = logger_for(input_callbacks=[Retain(), Mutate(), Fail(), Observe()]) + logger = logger_for(input_callbacks=[Fail(), Observe()]) details = logger.model_call_details - details.update(retained_metadata=metadata, normalized=None, remove=removed) - messages = logger.messages - additional = {"headers": {"test": "header"}} - owner = owners.prepare(logger.pre_call, (messages, "test-key"), {"additional_args": additional}) + details["retained_metadata"] = metadata + assert invoke_pre_call(owners, logger, {}) is None + assert order == ["fail", "observe"] + assert len(observed) == 1 and observed[0][0] is details + assert observed[0][1] == ("before failure",) and observed[0][2] is lock + assert details["retained_metadata"] is metadata + assert metadata == {"keep": ["before failure"]} and details["lock"] is lock + with TestCase().assertRaises(TypeError): + json.dumps({"lock": details["lock"]}) + + +async def real_post_call_logging(owners): + saved, observed, order = [], [], [] + shared = {"values": []} + response = ModelResponse(model="test") + ignored = {"replacement": True} + error = RuntimeError("expected post-call callback failure") + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["stash"] = shared + saved.append(kwargs) + + def log_post_api_call(self, kwargs, response_obj, start_time, end_time): + order.append("retain") + saved.append(kwargs) + return ignored + + class MutateThenFail(CustomLogger): + def log_post_api_call(self, kwargs, response_obj, start_time, end_time): + order.append("fail") + kwargs["stash"]["values"].append("post") + kwargs["callback_error"] = error + kwargs["original_response"].choices[0].message.content = "mutated" + raise error + + class Observe(CustomLogger): + def log_post_api_call(self, kwargs, response_obj, start_time, end_time): + order.append("observe") + observed.append((kwargs, response_obj, start_time, end_time, tuple(kwargs["stash"]["values"]))) + return ignored + + logger = logger_for(input_callbacks=[Retain(), MutateThenFail(), Observe()]) + details = logger.model_call_details + assert invoke_pre_call(owners, logger, {}) is None + additional = {"headers": {"test": "post"}} + owner = owners.prepare(logger.post_call, (response, logger.messages, "test-key"), {"additional_args": additional}) try: assert owner.invoke() is None finally: owner.close() - assert order == ["retain", "mutate", "fail", "observe"] - assert snapshots == [("normalized", False, "masked", ("before failure",), True, False)] - assert len(retained) == len(observed) == 1 - assert retained[0] is details and observed[0][0] is details - assert observed[0][1] is messages and details["input"] is messages - assert details["additional_args"] is additional - assert details["retained_metadata"] is metadata - assert metadata == {"secret": "masked", "keep": ["before failure"]} - assert details["normalized"] == "normalized" and "remove" not in details - assert details["lock"] is lock and "replacement" not in details - with TestCase().assertRaises(TypeError): - json.dumps({"lock": details["lock"]}) + assert order == ["retain", "fail", "observe"] + assert len(saved) == 2 and saved[0] is saved[1] is details + assert len(observed) == 1 and observed[0][0] is details + assert observed[0][1] is None and observed[0][2] is logger.start_time and observed[0][3] is None + assert observed[0][4] == ("post",) + assert details["original_response"] is response and response.choices[0].message.content == "mutated" + assert details["input"] is logger.messages and details["additional_args"] is additional + assert details["log_event_type"] == "post_api_call" and details["api_key"] == "test-key" + assert details["stash"] is shared and details["callback_error"] is error + assert "replacement" not in details + shared["values"].append("later") + assert saved[0]["stash"]["values"] == ["post", "later"] + + +async def real_post_call_dict_response(owners): + observed = [] + + class Observe(CustomLogger): + def log_post_api_call(self, kwargs, response_obj, start_time, end_time): + observed.append(kwargs["original_response"]) + + response = {"content": ["original"], "timestamp": datetime(2026, 1, 1)} + logger = logger_for(input_callbacks=[Observe()]) + owner = owners.prepare(logger.post_call, (response,)) + try: + assert owner.invoke() is None + finally: + owner.close() + assert len(observed) == 1 and observed[0] is logger.model_call_details["original_response"] + assert isinstance(observed[0], str) + assert json.loads(observed[0]) == {"content": ["original"], "timestamp": "2026-01-01 00:00:00"} + response["content"].append("later") + assert json.loads(observed[0])["content"] == ["original"] + + +async def real_sync_logging(owners): + saved, observations, replacements = [], [], [] + shared = {"values": []} + result = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "original"}}]) + replacement = ModelResponse(model="test", choices=[{"message": {"role": "assistant", "content": "replacement"}}]) + ignored = {"ignored": True}, result + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["stash"] = shared + saved.append(kwargs) + + def logging_hook(self, kwargs, result, call_type): + observations.append(("retain", kwargs, result, call_type)) + kwargs["stash"]["values"].append("hook") + result.choices[0].message.content = "mutated" + return kwargs, result + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + kwargs["stash"]["values"].append("event") + observations.append(("event", kwargs, response_obj, tuple(kwargs["stash"]["values"]))) + return ignored + + class Replace(CustomLogger): + def logging_hook(self, kwargs, result, call_type): + observations.append(("replace", kwargs, result, result.choices[0].message.content)) + updated = {**kwargs, "adopted": True} + replacements.append(updated) + return updated, replacement + + class Observe(CustomLogger): + def logging_hook(self, kwargs, result, call_type): + observations.append(("observe", kwargs, result, tuple(kwargs["stash"]["values"]))) + return kwargs, result + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append(("success", kwargs, response_obj, tuple(kwargs["stash"]["values"]))) + + retain = Retain() + logger = logger_for(input_callbacks=[retain], sync_callbacks=[retain, Replace(), Observe()], call_type="completion") + assert invoke_pre_call(owners, logger, {}) is None + owner = owners.prepare(logger.success_handler, (result,)) + try: + assert owner.invoke() is None + finally: + owner.close() + assert [entry[0] for entry in observations] == ["retain", "replace", "observe", "event", "success"] + assert len(saved) == len(replacements) == 1 + assert observations[0][1] is observations[1][1] is saved[0] + assert observations[0][2] is observations[1][2] is result + assert observations[0][3] == "completion" and observations[1][3] == "mutated" + assert observations[2][3] == ("hook",) and observations[3][3] == observations[4][3] == ("hook", "event") + assert all(entry[1] is replacements[0] is logger.model_call_details for entry in observations[2:]) + assert all(entry[2] is replacement for entry in observations[2:]) + assert logger.model_call_details is not saved[0] + assert logger.model_call_details["adopted"] and "adopted" not in saved[0] + assert "ignored" not in logger.model_call_details + assert result.choices[0].message.content == "mutated" + assert replacement.choices[0].message.content == "replacement" + assert saved[0]["stash"] is logger.model_call_details["stash"] is shared + shared["values"].append("later") + assert observations[-1][1]["stash"]["values"] == ["hook", "event", "later"] + + +async def real_sync_logging_hook_failure(owners): + order, saved = [], [] + result = ModelResponse(model="test") + error = RuntimeError("expected sync logging hook failure") + + class MutateThenFail(CustomLogger): + def logging_hook(self, kwargs, result, call_type): + order.append("fail") + saved.append(kwargs) + kwargs["callback_error"] = error + result.choices[0].message.content = "before failure" + raise error + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + order.append("unexpected success") + + class Observe(CustomLogger): + def logging_hook(self, kwargs, result, call_type): + order.append("unexpected hook") + return kwargs, result + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + order.append("unexpected later success") + + logger = logger_for(sync_callbacks=[MutateThenFail(), Observe()], call_type="completion") + details = logger.model_call_details + owner = owners.prepare(logger.success_handler, (result,)) + try: + assert owner.invoke() is None + finally: + owner.close() + assert order == ["fail"] + assert len(saved) == 1 and saved[0] is logger.model_call_details is details + assert details["callback_error"] is error and error.__traceback__ is not None + assert result.choices[0].message.content == "before failure" + + +async def real_sync_failure_chain(owners): + await real_failure_chain(owners, awaited=False) + + +async def real_async_failure_chain(owners): + await real_failure_chain(owners, awaited=True) + + +async def real_failure_chain(owners, awaited): + saved, observations, hooks = [], [], [] + shared = {"values": []} + error = ValueError("provider failure") + callback_error = RuntimeError("expected failure callback error") + ignored = {"replacement": True}, object() + task = asyncio.current_task() + end = datetime.now() + + class Stage(CustomLogger): + def __init__(self, name): + super().__init__() + self.name = name + + def log_pre_api_call(self, model, messages, kwargs): + kwargs["stash"] = shared + saved.append(kwargs) + + def logging_hook(self, kwargs, result, call_type): + hooks.append("sync") + return ignored + + async def async_logging_hook(self, kwargs, result, call_type): + hooks.append("async") + return ignored + + def record(self, kwargs, response_obj, start_time, end_time): + observations.append( + ( + self.name, + kwargs, + response_obj, + kwargs["exception"], + tuple(kwargs["stash"]["values"]), + start_time, + end_time, + asyncio.current_task(), + ) + ) + if self.name == "fail": + kwargs["stash"]["values"].append("before failure") + kwargs["callback_error"] = callback_error + raise callback_error + return ignored + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + return self.record(kwargs, response_obj, start_time, end_time) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await asyncio.sleep(0) + return self.record(kwargs, response_obj, start_time, end_time) + + retain, fail, observe = Stage("retain"), Stage("fail"), Stage("observe") + callbacks = [retain, fail, observe] + logger = logger_for( + input_callbacks=[retain], + failure_callbacks=() if awaited else callbacks, + async_failure_callbacks=callbacks if awaited else (), + call_type="acompletion" if awaited else "completion", + ) + logger.model_call_details["litellm_params"]["acompletion"] = awaited + details = logger.model_call_details + assert invoke_pre_call(owners, logger, {}) is None + owner = owners.prepare( + logger.async_failure_handler if awaited else logger.failure_handler, + (error, "provider traceback"), + {"start_time": logger.start_time, "end_time": end}, + awaited=awaited, + ) + try: + if awaited: + assert await owner.invoke() is None + else: + assert owner.invoke() is None + finally: + owner.close() + assert [entry[0] for entry in observations] == ["retain", "fail", "observe"] + assert len(saved) == 1 and saved[0] is logger.model_call_details is details + assert all(entry[1] is details and entry[2] is None and entry[3] is error for entry in observations) + assert [entry[4] for entry in observations] == [(), (), ("before failure",)] + assert all(entry[5] is logger.start_time and entry[6] is end and entry[7] is task for entry in observations) + assert details["exception"] is error and details["callback_error"] is callback_error + assert callback_error.__traceback__ is not None + assert details["traceback_exception"] == "provider traceback" and details["log_event_type"] == "failed_api_call" + assert details["stash"] is shared and "replacement" not in details and hooks == [] + shared["values"].append("later") + assert saved[0]["stash"]["values"] == observations[-1][1]["stash"]["values"] == ["before failure", "later"] async def real_async_logging(owners): diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 0164d829bec..8ecb2724acb 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -53,13 +53,6 @@ class _PreparedOCRRequest: litellm_logging_obj: LiteLLMLoggingObj -_RUST_OCR_PROVIDERS: Final = { - "mistral", - "azure_ai", - "vertex_ai", -} - - def _prepare_ocr_request( model: str, document: Mapping[str, object], @@ -181,50 +174,6 @@ def _prepare_ocr_request( ) -def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: - if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": - return False - if not prepared_request.provider_config.supports_rust_bridge(): - return False - return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS - - -def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool: - raw_request_override: Final = prepared_request.litellm_params.get("rust") - request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None - return rust_enabled(request_override=request_override) - - -def _ocr_boundary(prepared_request: _PreparedOCRRequest) -> rust_ocr_bridge.OCRBoundary: - return rust_ocr_bridge.OCRBoundary( - handler=base_llm_http_handler, - model=prepared_request.model, - document=prepared_request.document, - optional_params=prepared_request.optional_params, - logging_obj=prepared_request.litellm_logging_obj, - api_key=prepared_request.api_key, - api_base=prepared_request.api_base, - headers=prepared_request.extra_headers, - provider_config=prepared_request.provider_config, - litellm_params=prepared_request.litellm_params, - custom_llm_provider=prepared_request.custom_llm_provider, - timeout=prepared_request.effective_timeout, - ) - - -def _run_rust_ocr(prepared_request: _PreparedOCRRequest) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_ocr() is None: - return None - return rust_ocr_bridge.ocr(_ocr_boundary(prepared_request)) - - -async def _run_rust_aocr(prepared_request: _PreparedOCRRequest) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_aocr() is None: - return None - return await rust_ocr_bridge.aocr(_ocr_boundary(prepared_request)) - - -@client async def aocr( model: str, document: Mapping[str, object], @@ -293,6 +242,32 @@ async def aocr( ) ``` """ + if rust_enabled(): + return await rust_ocr_bridge.aocr( + { + **kwargs, + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + } + ) + return await _legacy_aocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) + + +async def _legacy_aocr( + model: str, + document: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + **kwargs: object, +) -> OCRResponse: completion_kwargs: Final[dict[str, object]] = { "model": model, "document": document, @@ -318,13 +293,6 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): - rust_response: Final = await _run_rust_aocr(prepared_request=prepared) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response = base_llm_http_handler.ocr( model=prepared.model, document=prepared.document, @@ -485,7 +453,6 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, return {"type": "document_url", "document_url": data_uri} -@client def ocr( model: str, document: Mapping[str, object], @@ -558,6 +525,33 @@ def ocr( print(f"Page {page.index}: {page.markdown}") ``` """ + if rust_enabled(): + arguments: Final[dict[str, object]] = { + **kwargs, + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + } + if kwargs.get("aocr") is True: + return rust_ocr_bridge.aocr(arguments) + return rust_ocr_bridge.ocr(arguments) + return _legacy_ocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) + + +def _legacy_ocr( + model: str, + document: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + **kwargs: object, +) -> OCRResponse | Coroutine[object, object, OCRResponse]: completion_kwargs: Final[dict[str, object]] = { "model": model, "document": document, @@ -585,13 +579,6 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared): - rust_response: Final = _run_rust_ocr(prepared_request=prepared) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response: Final = base_llm_http_handler.ocr( model=prepared.model, document=prepared.document, @@ -616,3 +603,9 @@ def ocr( completion_kwargs=completion_kwargs, extra_kwargs=kwargs, ) + + +_legacy_aocr.__name__ = "aocr" +_legacy_ocr.__name__ = "ocr" +_legacy_aocr = client(_legacy_aocr) +_legacy_ocr = client(_legacy_ocr) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 2a412957119..9ced40eab8d 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -1,145 +1,25 @@ -"""Retained OCR bridge: Python owns the request/response objects, Rust drives -prepare -> encode -> POST -> finish against those same objects.""" +"""Whole-call argument boundary for native OCR execution.""" from __future__ import annotations -import math +import inspect +import traceback from collections.abc import Awaitable -from dataclasses import dataclass, field +from contextvars import copy_context +from datetime import datetime from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables +from uuid import uuid4 -import httpx - -import litellm -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, DocumentType, OCRResponse -from litellm.llms.custom_httpx.http_handler import ( - AsyncHTTPHandler, - HTTPHandler, - _get_httpx_client, - get_async_httpx_client, -) -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.rust_bridge import configuration as _configuration +from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds - -rust_ocr_enabled = _configuration.rust_enabled -rust = _configuration.rust - -OCRRoots = tuple[dict[str, object], str, dict[str, object], None] -OCRWire = tuple[int, list[tuple[bytes, bytes]], bytes] -OCREncoded = tuple[str, list[tuple[bytes, bytes]], bytes, float] - - -def _positive_timeout_seconds(timeout: float | httpx.Timeout) -> float: - seconds: Final = _timeout_to_seconds(timeout) - if seconds is None or not math.isfinite(seconds) or seconds <= 0: - raise ValueError("OCR bridge requires a positive finite timeout") - return seconds - - -@dataclass(kw_only=True, slots=True) -class OCRBoundary: - handler: BaseLLMHTTPHandler - model: str - document: DocumentType - optional_params: dict[str, object] - logging_obj: Logging - api_key: str | None - api_base: str | None - headers: dict[str, object] | None - provider_config: BaseOCRConfig - litellm_params: dict[str, object] - custom_llm_provider: str - timeout: float | httpx.Timeout - client: HTTPHandler | AsyncHTTPHandler | None = None - request: httpx.Request | None = field(default=None, init=False) - - def prepare(self) -> OCRRoots: - roots: Final = self.handler._prepare_ocr_request( - model=self.model, - document=self.document, - optional_params=self.optional_params, - logging_obj=self.logging_obj, - api_key=self.api_key, - api_base=self.api_base, - headers=self.headers, - provider_config=self.provider_config, - litellm_params=self.litellm_params, - ) - if not isinstance(self.client, HTTPHandler): - self.client = _get_httpx_client() - return roots - - async def aprepare(self) -> OCRRoots: - roots: Final = await self.handler._async_prepare_ocr_request( - model=self.model, - document=self.document, - optional_params=self.optional_params, - logging_obj=self.logging_obj, - api_key=self.api_key, - api_base=self.api_base, - headers=self.headers, - provider_config=self.provider_config, - litellm_params=self.litellm_params, - ) - if not isinstance(self.client, AsyncHTTPHandler): - self.client = get_async_httpx_client(llm_provider=litellm.LlmProviders(self.custom_llm_provider)) - return roots - - def encode(self, roots: OCRRoots) -> OCREncoded: - headers, url, data, _files = roots - seconds: Final = _positive_timeout_seconds(self.timeout) - if self.client is None: - raise RuntimeError("OCR boundary must be prepared before encoding") - try: - self.request = self.client.client.build_request( - "POST", - url, - headers=cast(dict[str, str], headers), - json=data, - timeout=self.timeout, - ) - return str(self.request.url), self.request.headers.raw, self.request.read(), seconds - except Exception as e: # noqa: BLE001 # match the Python OCR handler's encoding error mapping - raise self.handler._handle_error(e=e, provider_config=self.provider_config) - - def _response(self, wire: OCRWire) -> httpx.Response: - if self.request is None: - raise RuntimeError("OCR boundary must be encoded before finishing") - status, headers, content = wire - try: - response: Final = httpx.Response(status, headers=headers, content=content, request=self.request) - response.raise_for_status() - except Exception as e: # noqa: BLE001 # match the Python OCR handler's response error mapping - raise self.handler._handle_error(e=e, provider_config=self.provider_config) - return response - - def finish(self, wire: OCRWire) -> OCRResponse: - return self.handler._transform_ocr_response( - provider_config=self.provider_config, - model=self.model, - response=self._response(wire), - logging_obj=self.logging_obj, - optional_params=self.optional_params, - ) - - async def afinish(self, wire: OCRWire) -> OCRResponse: - return await self.provider_config.async_transform_ocr_response( - model=self.model, - raw_response=self._response(wire), - logging_obj=self.logging_obj, - optional_params=self.optional_params, - ) class RustOcr(Protocol): - def __call__(self, boundary: OCRBoundary) -> OCRResponse: ... + def __call__(self, arguments: dict[str, object]) -> OCRResponse: ... class RustAocr(Protocol): - def __call__(self, boundary: OCRBoundary) -> Awaitable[OCRResponse]: ... + def __call__(self, arguments: dict[str, object]) -> Awaitable[OCRResponse]: ... def _as_ocr(value: object) -> RustOcr | None: @@ -162,15 +42,83 @@ def load_rust_aocr() -> RustAocr | None: return _AOCR.load() -def ocr(boundary: OCRBoundary) -> OCRResponse | None: - rust_ocr: Final = load_rust_ocr() - if rust_ocr is None: - return None - return rust_ocr(boundary) +def ocr(arguments: dict[str, object]) -> OCRResponse: + implementation: Final = load_rust_ocr() + if implementation is None: + raise RuntimeError("Rust OCR is enabled but the native OCR extension is unavailable") + return implementation(arguments) -async def aocr(boundary: OCRBoundary) -> OCRResponse | None: - rust_aocr: Final = load_rust_aocr() - if rust_aocr is None: +async def aocr(arguments: dict[str, object]) -> OCRResponse: + implementation: Final = load_rust_aocr() + if implementation is None: + raise RuntimeError("Rust OCR is enabled but the native OCR extension is unavailable") + return await implementation(arguments) + + +def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> object: + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + + supplied: Final = arguments.get("litellm_logging_obj") + if supplied is not None: + return supplied + callbacks: Final = tuple(dict.fromkeys((*litellm.callbacks, *cast(list, arguments.get("callbacks") or [])))) + success: Final = tuple(dict.fromkeys((*callbacks, *cast(list, arguments.get("success_callback") or [])))) + failure: Final = tuple(dict.fromkeys((*callbacks, *cast(list, arguments.get("failure_callback") or [])))) + call_id: Final = str(arguments.get("litellm_call_id") or uuid4()) + logger: Final = Logging( + model=str(arguments["model"]), + messages="default-message-value", + stream=False, + call_type="aocr" if asynchronous else "ocr", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=str(arguments.get("id") or ""), + litellm_trace_id=cast(str | None, arguments.get("litellm_trace_id")), + dynamic_input_callbacks=[cb for cb in callbacks if cb not in litellm.input_callback], + dynamic_success_callbacks=[cb for cb in success if not inspect.iscoroutinefunction(cb)], + dynamic_async_success_callbacks=list(success), + dynamic_failure_callbacks=[cb for cb in failure if not inspect.iscoroutinefunction(cb)], + dynamic_async_failure_callbacks=list(failure), + kwargs=arguments, + supports_correlation_logging=asynchronous, + ) + arguments["litellm_call_id"] = call_id + arguments["litellm_logging_obj"] = logger + return logger + + +def invoke_terminal( + action: str, roots: object, logger: object, value: object, start_time: datetime, end_time: datetime +) -> object: + from litellm import utils + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging: Final = cast(Logging, logger) + if action == "sync_success": + + def run() -> None: + _retained: Final = roots + logging.success_handler(value, start_time, end_time) + + return utils.executor.submit(copy_context().run, run) + if action == "async_success": + + async def run_async() -> None: + _retained: Final = roots + await logging.async_success_handler(value, start_time, end_time) + + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=run_async()) return None - return await rust_aocr(boundary) + if action == "sync_success_if_needed": + if logging._should_run_sync_callbacks_for_async_calls(): + return invoke_terminal("sync_success", roots, logger, value, start_time, end_time) + return None + exception: Final = cast(Exception, value) + trace: Final = "".join(traceback.format_exception(type(exception), exception, exception.__traceback__)) + if action == "sync_failure": + logging.failure_handler(exception, trace, start_time, end_time) + return None + return logging.async_failure_handler(exception, trace, start_time, end_time) diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 249fbda713e..01f8c44ccc0 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -1,52 +1,20 @@ """ -Tests for the OCR `req_format` option in the SDK request path: -providers that don't support a native response must reject it, and the Rust -bridge (which only returns the normalized shape) must not serve native requests. +Tests for the OCR `req_format` option in the Python SDK request path. """ -import dataclasses -from unittest.mock import MagicMock - import pytest import litellm -from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig -from litellm.llms.cohere.ocr.transformation import CohereParseConfig -from litellm.ocr.main import _PreparedOCRRequest, _rust_ocr_supported +from litellm.rust_bridge import configuration DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -def _prepared(optional_params: dict[str, object]) -> _PreparedOCRRequest: - return _PreparedOCRRequest( - model="doc-intelligence/prebuilt-layout", - document=dict(DOCUMENT), - api_key="fake-key", - api_base="https://example.cognitiveservices.azure.com", - custom_llm_provider="azure_ai", - extra_headers=None, - provider_config=MagicMock(), - optional_params=optional_params, - litellm_params={}, - effective_timeout=60.0, - litellm_logging_obj=MagicMock(), - ) - - -@pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) -def test_rust_ocr_serves_default_format(optional_params): - assert _rust_ocr_supported(_prepared(optional_params)) is True - - -def test_rust_ocr_skipped_for_native_format(): - assert _rust_ocr_supported(_prepared({"req_format": "native"})) is False - - -@pytest.mark.parametrize("provider_config", [CohereParseConfig(), AzureAICohereParseConfig()]) -def test_rust_ocr_skipped_for_configs_without_bridge_support(provider_config): - prepared = dataclasses.replace(_prepared({}), provider_config=provider_config) - - assert _rust_ocr_supported(prepared) is False +@pytest.fixture(autouse=True) +def python_ocr_path(): + litellm.rust(False) + yield + configuration.reset_rust_configuration() @pytest.mark.asyncio diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 9a05b248078..3b4c1580b59 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,720 +1,427 @@ -"""Tests for the optional Rust-backed OCR path.""" +"""Strict whole-argument OCR dispatch and opt-in native transport contracts.""" import asyncio +import atexit import builtins import contextvars -import copy import gc import importlib import inspect +import json import os -import subprocess -import sys import threading -import types import weakref -from dataclasses import replace -from datetime import datetime -from typing import Any, Final, cast +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import ModuleType +from unittest.mock import AsyncMock, Mock import httpx import pytest import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.base_llm.ocr.transformation import OCRRequestData, OCRResponse -from litellm.llms.mistral.ocr.transformation import MistralOCRConfig +from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge import configuration -# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` -# function onto `litellm.ocr` and shadows the submodule, so import the modules -# explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" -DOCUMENT: dict[str, object] = { - "type": "document_url", - "document_url": "https://example.com/doc.pdf", -} - -FAKE_OCR_RESPONSE: dict[str, object] = { - "pages": [{"index": 0, "markdown": "hello world"}], - "model": "mistral-ocr-2505-completion", - "document_annotation": None, - "usage_info": {"pages_processed": 1}, +RESPONSE_DATA = { + "pages": [ + { + "index": 0, + "markdown": "proof", + "images": [{"id": "image-0", "image_base64": "aW1hZ2U="}], + "dimensions": {"dpi": 200, "height": 2200, "width": 1700}, + } + ], + "model": "mistral-ocr-latest", + "document_annotation": {"title": "Test document"}, + "usage_info": {"pages_processed": 1, "doc_size_bytes": 1234}, "object": "ocr", } -class CapturedException(Exception): - pass - - -class RecordingBridge: - """A fake ``RustOcr`` callable that records the boundary it was handed.""" - - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: - self.calls.append( - { - "model": boundary.model, - "document": boundary.document, - "api_key": boundary.api_key, - "api_base": boundary.api_base, - "custom_llm_provider": boundary.custom_llm_provider, - "extra_headers": boundary.headers, - "optional_params": boundary.optional_params, - "timeout": boundary.timeout, - } - ) - return OCRResponse.model_validate(FAKE_OCR_RESPONSE) - - -class RecordingAsyncBridge: - """A fake async ``RustAocr`` callable that records the boundary it was handed.""" - - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - async def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: - self.calls.append( - { - "model": boundary.model, - "document": boundary.document, - "api_key": boundary.api_key, - "api_base": boundary.api_base, - "custom_llm_provider": boundary.custom_llm_provider, - "extra_headers": boundary.headers, - "optional_params": boundary.optional_params, - "timeout": boundary.timeout, - } - ) - return OCRResponse.model_validate(FAKE_OCR_RESPONSE) - - -class RaisingBridge: - def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: - raise RuntimeError("bridge failed") - - -class RaisingAsyncBridge: - async def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: - raise RuntimeError("bridge failed") - - -class RecordingLogging: - """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" - - def __init__(self) -> None: - self.pre_call_kwargs: dict[str, object] | None = None - - 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, - "additional_args": additional_args, - } - - -class FakeOCRConfig: - """A stand-in ``BaseOCRConfig`` that echoes the request it would build.""" - - def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None: - self.api_key_env_var = api_key_env_var - self.seen_api_keys: list[str | None] = [] - - 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]: - self.seen_api_keys.append(api_key) - 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 transform_ocr_request( - self, - *, - model: str, - document: dict[str, object], - optional_params: dict[str, object], - headers: dict[str, object], - api_key: str | None, - api_base: str | None, - ) -> OCRRequestData: - return OCRRequestData(data={"model": model, "document": document, **optional_params}, files=None) - - async def async_transform_ocr_request( - self, - *, - model: str, - document: dict[str, object], - optional_params: dict[str, object], - headers: dict[str, object], - api_key: str | None, - api_base: str | None, - ) -> OCRRequestData: - return self.transform_ocr_request( - model=model, - document=document, - optional_params=optional_params, - headers=headers, - api_key=api_key, - api_base=api_base, - ) - - -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: - from litellm.constants import request_timeout - - 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 if timeout is not None else float(request_timeout), - litellm_logging_obj=logging_obj or RecordingLogging(), - ) - - -class BoundaryDriver: - """Stands in for the native route: drives the boundary's own methods.""" - - def __init__(self) -> None: - self.roots: rust_bridge.OCRRoots | None = None - self.encoded: rust_bridge.OCREncoded | None = None - - def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: - self.roots = boundary.prepare() - self.encoded = boundary.encode(self.roots) - return OCRResponse.model_validate(FAKE_OCR_RESPONSE) - - @pytest.fixture(autouse=True) -def _reset_rust_flag(): - """Keep the global toggle isolated between tests.""" +def reset_rust_state(monkeypatch): + monkeypatch.delenv("LITELLM_RUST", raising=False) + configuration.reset_rust_configuration() rust_bridge._OCR.reset() rust_bridge._AOCR.reset() - configuration.reset_rust_configuration() - rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - yield - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() - configuration.reset_rust_configuration() - rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - - -@pytest.fixture -def fake_bridge(): - """Enable the Rust path with an injected recording bridge (no native wheel).""" - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - return bridge - - -@pytest.fixture -def fake_async_bridge(): - """Enable the async Rust path with an injected recording bridge.""" - bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._AOCR.override(bridge) - return bridge - - -def test_load_rust_ocr_returns_injected_impl(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - assert rust_bridge.load_rust_ocr() is bridge - - -def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch): - real_import = builtins.__import__ - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - if name == "litellm.rust_bridge" and "_native" in fromlist: - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - - -def test_native_bridge_loader_caches_absent_extension(monkeypatch): - real_import = builtins.__import__ - attempts = 0 - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal attempts - if name == "litellm.rust_bridge" and "_native" in fromlist: - attempts += 1 - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - assert rust_bridge_loader.get_native_bridge() is None - assert attempts == 1 - - -def test_native_bridge_loader_reset_forces_relookup(monkeypatch): - real_import = builtins.__import__ - attempts = 0 - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal attempts - if name == "litellm.rust_bridge" and "_native" in fromlist: - attempts += 1 - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None rust_bridge_loader.reset_native_bridge_cache() - assert rust_bridge_loader.get_native_bridge() is None - assert attempts == 2 + try: + yield + finally: + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() + configuration.reset_rust_configuration() + rust_bridge_loader.reset_native_bridge_cache() -def test_native_bridge_available_reflects_loader(monkeypatch): - fake_module = types.ModuleType("litellm.rust_bridge._native") - monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module) - - assert rust_bridge_loader.native_bridge_available() is True +@pytest.fixture +def document(): + return {"type": "document_url", "document_url": "https://example.invalid/document.pdf"} -def test_load_rust_aocr_returns_injected_impl(): - bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._AOCR.override(bridge) - assert rust_bridge.load_rust_aocr() is bridge +@pytest.fixture +def response(): + return OCRResponse.model_validate(RESPONSE_DATA) -def test_toggle_without_ocr_arg_preserves_injected_impl(): - """The public flag must not clobber an internal test binding.""" - bridge = RecordingBridge() - async_bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) - - litellm.rust(False) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge - litellm.rust(True) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge +@pytest.fixture +def injected_native(response): + sync = Mock(return_value=response) + asynchronous = AsyncMock(return_value=response) + rust_bridge._OCR.override(sync) + rust_bridge._AOCR.override(asynchronous) + return sync, asynchronous -def test_explicit_ocr_none_clears_injected_impl(monkeypatch): - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: None, +@pytest.fixture +def no_python_ocr(monkeypatch): + prepare = Mock(side_effect=lambda **kwargs: pytest.fail("Python OCR preparation ran")) + handler = Mock(side_effect=lambda **kwargs: pytest.fail("Python OCR transport ran")) + mapping = Mock(side_effect=lambda **kwargs: pytest.fail("Python exception mapping ran")) + monkeypatch.setattr(ocr_main, "_prepare_ocr_request", prepare) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "_prepare_ocr_request", prepare) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "_async_prepare_ocr_request", prepare) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", handler) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "async_ocr", handler) + monkeypatch.setattr(litellm, "exception_type", mapping) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("enable_with", ["global", "environment"]) +async def test_unwrapped_dispatch_preserves_every_argument( + monkeypatch, no_python_ocr, injected_native, response, asynchronous, enable_with +): + if enable_with == "global": + litellm.rust(True) + else: + monkeypatch.setenv("LITELLM_RUST", "1") + opaque = object() + document = {"type": "file", "file": opaque, "mime_type": "application/pdf"} + metadata = {"opaque": opaque, "nested": []} + arguments = { + "model": "unresolved-provider/opaque-model", + "document": document, + "api_key": opaque, + "api_base": opaque, + "timeout": httpx.Timeout(30, read=12.5), + "custom_llm_provider": "unported-provider", + "extra_headers": {"x-opaque": opaque}, + "metadata": metadata, + "litellm_metadata": {"opaque": opaque}, + "litellm_logging_obj": opaque, + "litellm_call_id": "whole-arguments", + "pages": [0, 2], + "include_image_base64": True, + "document_annotation_format": {"schema": opaque}, + "req_format": "native", + "client": opaque, + "arbitrary_option": opaque, + "optional_params": {"opaque": opaque}, + "kwargs": {"caller_owned": opaque}, + "aocr": opaque, + } + result = ( + await inspect.unwrap(ocr_main.aocr)(**arguments) if asynchronous else inspect.unwrap(ocr_main.ocr)(**arguments) ) - bridge = RecordingBridge() - async_bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) + sync, async_native = injected_native + selected, unused = (async_native, sync) if asynchronous else (sync, async_native) + selected.assert_called_once() + unused.assert_not_called() + if asynchronous: + async_native.assert_awaited_once() + assert result is response + assert selected.call_args.kwargs == {} + (forwarded,) = selected.call_args.args + assert isinstance(forwarded, dict) + assert forwarded.keys() == arguments.keys() + for name, value in arguments.items(): + assert forwarded[name] is value, name + assert document == {"type": "file", "file": opaque, "mime_type": "application/pdf"} + assert metadata == {"opaque": opaque, "nested": []} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_unwrapped_dispatch_keeps_unresolved_defaults( + no_python_ocr, injected_native, document, response, asynchronous +): + litellm.rust(True) + result = ( + await inspect.unwrap(ocr_main.aocr)(MODEL, document) + if asynchronous + else inspect.unwrap(ocr_main.ocr)(MODEL, document) + ) + selected = injected_native[int(asynchronous)] + selected.assert_called_once_with( + { + "model": MODEL, + "document": document, + "api_key": None, + "api_base": None, + "timeout": None, + "custom_llm_provider": None, + "extra_headers": None, + } + ) + assert result is response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_public_decorator_routes_full_kwargs(no_python_ocr, injected_native, document, response, asynchronous): + litellm.rust(True) + metadata = {"test_tag": "whole-arguments"} + pages = [0, 2] + arguments = { + "model": MODEL, + "document": document, + "api_key": "sk-test", + "api_base": "https://example.invalid", + "timeout": 12.5, + "extra_headers": {"x-trace-id": "trace-1"}, + "pages": pages, + "include_image_base64": True, + "metadata": metadata, + "arbitrary_option": {"nested": ["preserved"]}, + "num_retries": 0, + } + result = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + + selected = injected_native[int(asynchronous)] + selected.assert_called_once() + injected_native[not asynchronous].assert_not_called() + (forwarded,) = selected.call_args.args + assert result is response + for name in ("model", "api_key", "api_base", "timeout", "extra_headers", "arbitrary_option", "num_retries"): + assert forwarded[name] == arguments[name] + assert forwarded["document"] is document + assert forwarded["pages"] is pages + assert forwarded["include_image_base64"] is True + assert forwarded["metadata"]["test_tag"] == "whole-arguments" + assert forwarded["custom_llm_provider"] is None + assert "litellm_logging_obj" not in forwarded + assert "litellm_call_id" not in forwarded + assert "kwargs" not in forwarded + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_bridge_passes_same_dictionary_and_response(injected_native, document, response, asynchronous): + arguments = {"model": MODEL, "document": document, "opaque": object()} + result = await rust_bridge.aocr(arguments) if asynchronous else rust_bridge.ocr(arguments) + selected = injected_native[int(asynchronous)] + selected.assert_called_once_with(arguments) + assert selected.call_args.args[0] is arguments + assert result is response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("public", [False, True], ids=["unwrapped", "decorated"]) +@pytest.mark.parametrize("failure", ["missing", "unsupported", "runtime"]) +async def test_native_failures_propagate_without_fallback( + monkeypatch, no_python_ocr, injected_native, document, asynchronous, public, failure +): + litellm.rust(True) + error = ( + NotImplementedError("native OCR provider is not implemented") + if failure == "unsupported" + else RuntimeError("native OCR failed") + ) + if failure == "missing": + rust_bridge._OCR.reset() + rust_bridge._AOCR.reset() + monkeypatch.setattr(rust_bridge_bindings, "get_native_bridge", lambda: None) + else: + injected_native[int(asynchronous)].side_effect = error + function = litellm.aocr if asynchronous else litellm.ocr + route = function if public else inspect.unwrap(function) + with pytest.raises(RuntimeError if failure == "missing" else type(error)) as caught: + result = route(model=MODEL, document=document, api_key="sk-test", num_retries=0) + if asynchronous: + await result + if failure == "missing": + assert "OCR" in str(caught.value).upper() + else: + assert caught.value is error + injected_native[int(asynchronous)].assert_called_once() + injected_native[not asynchronous].assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_bridge_missing_binding_is_strict(document, asynchronous): rust_bridge._OCR.override(None) rust_bridge._AOCR.override(None) - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None + with pytest.raises(RuntimeError, match="(?i)ocr"): + if asynchronous: + await rust_bridge.aocr({"model": MODEL, "document": document}) + else: + rust_bridge.ocr({"model": MODEL, "document": document}) -def test_load_rust_ocr_none_when_extension_absent(monkeypatch): - """With no injected impl and no compiled wheel, the loader returns None so the - caller degrades to the Python path instead of raising ImportError.""" - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: None, - ) - 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): - """With no injected impl but a packaged ``litellm.rust_bridge._native`` importable, - the loader returns the extension's ``ocr`` callable. The native wheel isn't - built in CI, so stand in a fake module via the bridge loader.""" - fake_module = types.ModuleType("litellm.rust_bridge._native") - fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: fake_module, - ) - - 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(): - assert rust_bridge._timeout_to_seconds(12.5) == 12.5 - assert rust_bridge._timeout_to_seconds(None) is None - assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 - - -def test_run_rust_ocr_forwards_boundary_fields(): - bridge = RecordingBridge() - logging_obj = RecordingLogging() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - response = ocr_main._run_rust_ocr( - 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, - ) - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert bridge.calls[0] == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("setting", ["default", "disabled", "overrides-environment"]) +async def test_rust_off_keeps_python_preparation_and_transport( + monkeypatch, injected_native, response, asynchronous, setting +): + if setting == "overrides-environment": + monkeypatch.setenv("LITELLM_RUST", "1") + if setting != "default": + litellm.rust(False) + prepare = Mock(wraps=ocr_main._prepare_ocr_request) + handler = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + lookup = Mock(side_effect=lambda: pytest.fail("disabled Rust binding was consulted")) + monkeypatch.setattr(ocr_main, "_prepare_ocr_request", prepare) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", handler) + monkeypatch.setattr(rust_bridge, "load_rust_ocr", lookup) + monkeypatch.setattr(rust_bridge, "load_rust_aocr", lookup) + arguments = { + "model": MODEL, + "document": {"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": {"x-trace-id": "trace-1"}, - "optional_params": {"include_image_base64": True}, "timeout": 12.5, + "extra_headers": {"x-trace-id": "python"}, + "pages": [0], + "include_image_base64": True, + "num_retries": 0, } - assert logging_obj.pre_call_kwargs is None # pre_call now runs inside the boundary + result = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - -def test_run_rust_ocr_passes_raw_api_key_to_provider_config(): - """Key resolution moved into provider ``validate_environment``; the boundary - forwards the caller's key unchanged.""" - provider_config = FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY") - driver = BoundaryDriver() - litellm.rust(True) - rust_bridge._OCR.override(driver) - - ocr_main._run_rust_ocr(build_prepared_request(provider_config=provider_config, api_key="sk-explicit", timeout=None)) - ocr_main._run_rust_ocr(build_prepared_request(provider_config=provider_config, api_key=None, timeout=None)) - - assert provider_config.seen_api_keys == ["sk-explicit", None] - - -def test_run_rust_ocr_preserves_native_response_identity(): - """The bridge returns the boundary's own finish() object, not a re-validated copy.""" - sentinel = OCRResponse.model_validate(FAKE_OCR_RESPONSE) - - class IdentityBridge: - def __call__(self, boundary: rust_bridge.OCRBoundary) -> OCRResponse: - return sentinel - - litellm.rust(True) - rust_bridge._OCR.override(IdentityBridge()) - - response = ocr_main._run_rust_ocr(build_prepared_request()) - - assert response is sentinel - - -def test_boundary_prepare_runs_pre_call_and_encodes_the_same_roots(): - """The logging view must alias the execution roots: mutating the headers the - callback received must surface in the encoded wire headers, while replacing - a view field must not.""" - logging_obj = RecordingLogging() - driver = BoundaryDriver() - litellm.rust(True) - rust_bridge._OCR.override(driver) - - seen: dict[str, object] = {} - - original_pre_call = logging_obj.pre_call - - def observing_pre_call(**kwargs: object) -> None: - original_pre_call(**kwargs) - view = cast(dict[str, object], kwargs["additional_args"]) - seen["headers"] = view["headers"] - cast(dict[str, object], view["headers"])["X-Proof"] = "mutated" - - logging_obj.pre_call = observing_pre_call # type: ignore[method-assign] - - ocr_main._run_rust_ocr(build_prepared_request(logging_obj=logging_obj, api_base="https://api.mistral.ai/v1")) - - assert logging_obj.pre_call_kwargs is not None - additional_args = logging_obj.pre_call_kwargs["additional_args"] - assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" - assert additional_args["headers"] == {"Authorization": "Bearer sk-test", "X-Proof": "mutated"} - assert driver.roots is not None - roots_headers, _url, _data, _files = driver.roots - assert roots_headers is seen["headers"] - assert driver.encoded is not None - encoded_headers = {name.lower(): value for name, value in driver.encoded[1]} - assert encoded_headers[b"x-proof"] == b"mutated" - - -def test_ocr_routes_to_rust_when_enabled(fake_bridge): - response = litellm.ocr( - 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_bridge.calls) == 1 - call = fake_bridge.calls[0] + assert result is response + prepare.assert_called_once() + handler.assert_called_once() + if asynchronous: + handler.assert_awaited_once() + call = handler.call_args.kwargs 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"] == {"x-trace-id": "trace-1"} - assert call["optional_params"].get("include_image_base64") is True - - -def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): - response = litellm.ocr( - model="azure_ai/pixtral-12b-2409", - document=DOCUMENT, - api_key="sk-test", - api_base="https://example.services.ai.azure.com", - ) - - assert isinstance(response, OCRResponse) - assert len(fake_bridge.calls) == 1 - assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409" - assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" - - -def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): - response = litellm.ocr( - model=MODEL, - document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, - api_key="sk-test", - ) - - assert isinstance(response, OCRResponse) - document = fake_bridge.calls[0]["document"] - assert document["type"] == "document_url" - assert document["document_url"].startswith("data:application/pdf;base64,") - - -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.rust(True) - rust_bridge._OCR.override(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" + assert call["document"] == { + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + } + assert call["optional_params"] == {"pages": [0], "include_image_base64": True} + assert call["timeout"] == 12.5 + assert call["headers"] == arguments["extra_headers"] + assert call["aocr"] is asynchronous + lookup.assert_not_called() + for native in injected_native: + native.assert_not_called() @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"] == {"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.rust(True) - rust_bridge._AOCR.override(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.""" - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test", timeout=12.5) - - assert fake_bridge.calls[0]["timeout"] == 12.5 - - -def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - from litellm.constants import request_timeout - - assert fake_bridge.calls[0]["timeout"] == float(request_timeout) - - -def test_ocr_does_not_route_to_rust_when_disabled(): - """With the flag off, the bridge must not be consulted even if an impl exists.""" - bridge = RecordingBridge() +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_rust_off_preserves_python_exception_mapping(monkeypatch, injected_native, document, asynchronous): litellm.rust(False) - rust_bridge.set_rust_ocr(ocr=bridge) - - assert rust_bridge.rust_ocr_enabled() is False - # The impl stays available for injection, but the disabled flag gates usage, - # so ocr() never reaches the Rust path (asserted via the enabled-path test). - assert bridge.calls == [] + original_error = ValueError("Python transport failed") + mapped_error = RuntimeError("mapped Python error") + mapping = Mock(return_value=mapped_error) + handler = AsyncMock(side_effect=original_error) if asynchronous else Mock(side_effect=original_error) + monkeypatch.setattr(litellm, "exception_type", mapping) + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", handler) + arguments = dict(model=MODEL, document=document, api_key="sk-test", litellm_logging_obj=Mock()) + with pytest.raises(RuntimeError) as caught: + if asynchronous: + await inspect.unwrap(ocr_main._legacy_aocr)(**arguments) + else: + inspect.unwrap(ocr_main._legacy_ocr)(**arguments) + assert caught.value is mapped_error + handler.assert_called_once() + mapping.assert_called_once() + assert mapping.call_args.kwargs["original_exception"] is original_error + assert mapping.call_args.kwargs["model"] == "mistral-ocr-latest" + assert mapping.call_args.kwargs["custom_llm_provider"] == "mistral" + for native in injected_native: + native.assert_not_called() -def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): - """Rust enabled but no bridge available (no injected impl, no compiled wheel): - ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) - litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI - - captured = {} - - def fake_handler_ocr(**kwargs): - captured["called"] = True - return OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") - - monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) - - response = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured.get("called") is True # Python path was used - assert isinstance(response, OCRResponse) +def test_global_toggle_preserves_injected_bindings(injected_native): + for enabled in (True, False, True): + litellm.rust(enabled) + assert configuration.rust_enabled() is enabled + assert rust_bridge.load_rust_ocr() is injected_native[0] + assert rust_bridge.load_rust_aocr() is injected_native[1] -def test_ocr_provider_configs_expose_api_key_env_vars(): - from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( - AzureDocumentIntelligenceOCRConfig, - ) - from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig - from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig - from litellm.llms.mistral.ocr.transformation import MistralOCRConfig - from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( - VertexAIDeepSeekOCRConfig, - ) - from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig - - assert BaseOCRConfig().get_api_key_env_var() is None - assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" - assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" - assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" - assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" - assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" +@pytest.mark.parametrize("name", ["ocr", "aocr"]) +def test_binding_override_none_and_reset(monkeypatch, name): + native = ModuleType("litellm.rust_bridge._native") + implementation = Mock() + setattr(native, name, implementation) + monkeypatch.setattr(rust_bridge_bindings, "get_native_bridge", lambda: native) + binding = rust_bridge._OCR if name == "ocr" else rust_bridge._AOCR + assert binding.load() is implementation + override = Mock() + binding.override(override) + assert binding.load() is override + binding.override(None) + assert binding.load() is None + binding.reset() + assert binding.load() is implementation + setattr(native, name, object()) + assert binding.load() is None -################################################# -# Proof: pre_call callbacks run against the same objects the wire request is -# built from, through the real native route, compared with the Python route. -################################################# +def test_loader_caches_missing_extension_until_reset(monkeypatch): + original_import = builtins.__import__ + attempts = Mock() -MISTRAL_OCR_RESPONSE_JSON: Final = ( - b'{"pages": [{"index": 0, "markdown": "proof"}], "model": "mistral-ocr-2505-completion",' - b' "document_annotation": null, "usage_info": {"pages_processed": 1}, "object": "ocr"}' -) + def import_without_native(name, globals=None, locals=None, fromlist=(), level=0): + if name == "litellm.rust_bridge" and "_native" in fromlist: + attempts() + raise ImportError("extension not installed") + return original_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", import_without_native) + assert rust_bridge_loader.get_native_bridge() is None + assert rust_bridge_loader.native_bridge_available() is False + attempts.assert_called_once() + rust_bridge_loader.reset_native_bridge_cache() + assert rust_bridge_loader.get_native_bridge() is None + assert attempts.call_count == 2 -class _WireRecorder: - """Loopback OCR server recording every request it serves.""" +@pytest.fixture +def native_ocr(): + native = rust_bridge_loader.get_native_bridge() + try: + available = native is not None and all( + tuple(inspect.signature(getattr(native, name)).parameters) == ("arguments",) for name in ("ocr", "aocr") + ) + except (AttributeError, TypeError, ValueError): + available = False + if not available: + message = "native whole-argument OCR extension unavailable or stale; rebuild the extension" + if os.environ.get("LITELLM_REQUIRE_NATIVE_OCR") == "1": + pytest.fail(message) + pytest.skip(message) + return native - def __init__(self) -> None: - import json - import threading - from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer - self.requests: list[dict[str, object]] = [] +class WireRecorder: + def __init__(self): + self.requests = [] self.received = threading.Event() self.release = threading.Event() self.finished = threading.Event() self.release.set() self.status = 200 + self.response_data = RESPONSE_DATA recorder = self class Handler(BaseHTTPRequestHandler): - def do_POST(self) -> None: # http.server API + def do_POST(self): body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) recorder.requests.append( { @@ -727,229 +434,451 @@ class _WireRecorder: try: if not recorder.release.wait(timeout=10): return + payload = json.dumps(recorder.response_data).encode() self.send_response(recorder.status) self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(MISTRAL_OCR_RESPONSE_JSON))) + self.send_header("Content-Length", str(len(payload))) self.end_headers() - self.wfile.write(MISTRAL_OCR_RESPONSE_JSON) + self.wfile.write(payload) except (BrokenPipeError, ConnectionResetError): pass finally: recorder.finished.set() - def log_message(self, *args: object) -> None: + def log_message(self, *args): pass - self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) - self._thread.start() + self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() @property - def api_base(self) -> str: - host, port = self._server.server_address[:2] + def api_base(self): + host, port = self.server.server_address[:2] return f"http://{host}:{port}" - def stop(self) -> None: + def stop(self): self.release.set() - self._server.shutdown() - self._server.server_close() - self._thread.join(timeout=5) + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) @pytest.fixture def wire_recorder(): - recorder = _WireRecorder() + recorder = WireRecorder() try: yield recorder finally: recorder.stop() -def _native_boundary_route_available() -> bool: - bridge = rust_bridge_loader.get_native_bridge() - if bridge is None: - return False - ocr_fn = getattr(bridge, "ocr", None) - if ocr_fn is None: - return False - try: - return str(inspect.signature(ocr_fn)) == "(boundary)" - except (TypeError, ValueError): - return False - - -@pytest.fixture -def native_ocr(): - if not _native_boundary_route_available(): - if os.environ.get("LITELLM_REQUIRE_NATIVE_OCR") == "1": - pytest.fail("native OCR boundary route not built") - pytest.skip("native OCR boundary route not built") - return rust_bridge_loader.get_native_bridge() - - -class ProofCallback(CustomLogger): - def __init__(self, document, context, events, fail=False) -> None: - self.document = document - self.context = context - self.events = events - self.fail = fail - self.headers: dict[str, object] | None = None - self.body: dict[str, object] | None = None - self.details = None - self.calls = 0 - - def log_pre_api_call(self, model, messages, kwargs): - self.calls += 1 - self.events.append(("mutate", self.context.get(), threading.get_ident(), asyncio.current_task())) - self.context.set("callback") - self.details = kwargs - view = kwargs["additional_args"] - self.headers = view["headers"] - self.body = view["complete_input_dict"] - self.headers["X-Proof"] = "mutated" - self.body["document"]["document_url"] = "https://example.invalid/mutated-by-callback.pdf" - self.document["document_name"] = "closure-mutation" - view["headers"] = {"X-Replacement": "must-not-reach-wire"} - view["complete_input_dict"] = {"model": "logging-only"} - if self.fail: - raise RuntimeError("expected pre-call failure") - return {"additional_args": {"headers": {"X-Return": "ignored"}}} - - -def _assert_retained_wire(request: dict[str, object]) -> None: - assert request["path"] == "/v1/ocr" - headers = request["headers"] - assert headers["authorization"] == "Bearer sk-test" - assert headers["x-proof"] == "mutated" - assert "x-replacement" not in headers - assert "x-return" not in headers - body = request["body"] - assert body["model"] == "mistral-ocr-latest" - assert body["document"]["document_url"] == "https://example.invalid/mutated-by-callback.pdf" - assert body["document"]["document_name"] == "closure-mutation" - - -async def _pre_call_contract(native, wire_recorder, monkeypatch, asynchronous, enabled, fail=False, control="original"): - document = { - "type": "document_url", - "document_url": "https://example.invalid/original.pdf", - } - context = contextvars.ContextVar("ocr-pre-call", default="caller") +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("callback_raises", [False, True], ids=["ignored-return", "caught-error"]) +async def test_native_mistral_wire_response_and_callback_identity( + native_ocr, wire_recorder, monkeypatch, no_python_ocr, document, asynchronous, callback_raises +): + litellm.rust(True) + context = contextvars.ContextVar("ocr-callback-context", default="caller") + caller = (threading.get_ident(), asyncio.current_task()) events = [] + retained = [] observations = [] - native_calls = [] - proof = ProofCallback(document, context, events, fail=fail) + pages = [0] + + class Mutate(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + events.append(("mutate", context.get(), threading.get_ident(), asyncio.current_task())) + context.set("callback") + view = kwargs["additional_args"] + body, headers = view["complete_input_dict"], view["headers"] + retained.append((kwargs, body, headers)) + headers["X-Proof"] = "mutated" + body["document"]["document_url"] = "https://example.invalid/mutated.pdf" + document["document_name"] = "closure-mutation" + pages.append(2) + view["headers"] = {"X-Replacement": "logging-only"} + view["complete_input_dict"] = {"model": "logging-only"} + if callback_raises: + raise RuntimeError("expected callback failure") + return {"additional_args": {"headers": {"X-Return": "ignored"}}} class Observe(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): events.append(("observe", context.get(), threading.get_ident(), asyncio.current_task())) - view = kwargs["additional_args"] observations.append( ( - kwargs is proof.details, - tuple(view["headers"].items()), - view["complete_input_dict"]["model"], - document["document_url"], + kwargs is retained[0][0], + dict(kwargs["additional_args"]["headers"]), + dict(kwargs["additional_args"]["complete_input_dict"]), ) ) - def selected(boundary): - if control == "copy-input": - return replace(boundary, document=copy.deepcopy(boundary.document)) - if control in {"logging-roots", "tuple-only"}: - - class EncodeControl: - def __getattr__(self, name): - return getattr(boundary, name) - - def encode(self, roots): - headers, url, body, files = roots - if control == "logging-roots": - view = boundary.logging_obj.model_call_details["additional_args"] - return boundary.encode((view["headers"], url, view["complete_input_dict"], files)) - return boundary.encode((headers, url, body, files)) - - return EncodeControl() - return boundary - - def sync_call(boundary): - native_calls.append("sync") - if control == "duplicate-prepare": - boundary.prepare() - return native.ocr(selected(boundary)) - - async def async_call(boundary): - native_calls.append("async") - if control == "duplicate-prepare": - await boundary.aprepare() - if control == "new-task": - return await asyncio.create_task(native.aocr(selected(boundary))) - return await native.aocr(selected(boundary)) - - rust_bridge._OCR.override(sync_call) - rust_bridge._AOCR.override(async_call) - monkeypatch.setattr(litellm, "input_callback", [proof, Observe()]) - litellm.rust(enabled) - caller = (threading.get_ident(), asyncio.current_task()) - arguments = dict(model=MODEL, document=document, api_key="sk-test", api_base=wire_recorder.api_base, num_retries=0) + monkeypatch.setattr(litellm, "input_callback", [Mutate(), Observe()]) + arguments = { + "model": MODEL, + "document": document, + "api_key": "sk-test", + "api_base": wire_recorder.api_base, + "timeout": httpx.Timeout(5, read=3), + "extra_headers": {"X-Trace": "caller"}, + "pages": pages, + "include_image_base64": True, + "metadata": {"opaque": object()}, + "arbitrary_option": object(), + "num_retries": 0, + } response = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - assert native_calls == (["async" if asynchronous else "sync"] if enabled else []), "native dispatch" - assert isinstance(response, OCRResponse) and response.pages[0].markdown == "proof" - assert proof.calls == 1, "callback count" - assert proof.body is not None and proof.body["document"] is document, "caller identity" - assert events == [("mutate", "caller", *caller), ("observe", "callback", *caller)], "callback context/order" - assert context.get() == "callback", "caller context write" - assert observations == [ - ( - True, - (("X-Replacement", "must-not-reach-wire"),), - "logging-only", - "https://example.invalid/mutated-by-callback.pdf", - ) - ], "observation-time values" - assert len(wire_recorder.requests) == 1, "POST count" - assert wire_recorder.requests[0]["headers"].get("x-proof") == "mutated", "execution roots" - _assert_retained_wire(wire_recorder.requests[0]) - proof.body["document"]["document_url"] = "https://example.invalid/after-encode.pdf" - assert document["document_url"] == "https://example.invalid/after-encode.pdf" - assert ( - wire_recorder.requests[0]["body"]["document"]["document_url"] - == "https://example.invalid/mutated-by-callback.pdf" - ) + + assert isinstance(response, OCRResponse) + assert response.object == "ocr" + assert response.model == RESPONSE_DATA["model"] + assert response.pages[0].index == 0 + assert response.pages[0].markdown == "proof" + assert response.pages[0].dimensions.dpi == 200 + assert response.pages[0].images[0].image_base64 == "aW1hZ2U=" + assert response.document_annotation == {"title": "Test document"} + assert response.usage_info.pages_processed == 1 + assert response.usage_info.doc_size_bytes == 1234 + assert response.get_provider_native_response() is None + assert events == [("mutate", "caller", *caller), ("observe", "callback", *caller)] + assert context.get() == "callback" + assert observations == [(True, {"X-Replacement": "logging-only"}, {"model": "logging-only"})] + assert len(retained) == 1 + assert retained[0][1]["document"] is document + assert retained[0][1]["pages"] is pages + assert len(wire_recorder.requests) == 1 + request = wire_recorder.requests[0] + assert request["path"] == "/v1/ocr" + assert request["headers"]["authorization"] == "Bearer sk-test" + assert request["headers"]["x-trace"] == "caller" + assert request["headers"]["x-proof"] == "mutated" + assert "x-replacement" not in request["headers"] + assert "x-return" not in request["headers"] + assert request["body"] == { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.invalid/mutated.pdf", + "document_name": "closure-mutation", + }, + "pages": [0, 2], + "include_image_base64": True, + } + document["document_url"] = "https://example.invalid/after-send.pdf" + assert request["body"]["document"]["document_url"] == "https://example.invalid/mutated.pdf" @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) -@pytest.mark.parametrize("fail", [False, True], ids=["return-ignored", "caught-error"]) -async def test_pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, enabled, fail): - await _pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, enabled, fail) +@pytest.mark.parametrize("failure", [False, True], ids=["success", "failure"]) +@pytest.mark.parametrize("callback_source", ["global", "per-call"]) +async def test_public_native_callback_lifecycle( + native_ocr, wire_recorder, monkeypatch, no_python_ocr, document, asynchronous, failure, callback_source +): + from litellm import utils + from litellm.litellm_core_utils import litellm_logging, logging_worker + + litellm.rust(True) + wire_recorder.status = 429 if failure else 200 + monkeypatch.setattr(utils, "function_setup", Mock(side_effect=AssertionError("native OCR entered function_setup"))) + context = contextvars.ContextVar("ocr-lifecycle-context", default="caller") + caller_thread, caller_task = threading.get_ident(), asyncio.current_task() + shared = {"phase": "pre_call"} + pre_calls = [] + terminal_calls = [] + finished = threading.Event() + executor = ThreadPoolExecutor(max_workers=1) + worker = logging_worker.LoggingWorker(timeout=5, concurrency=1) + monkeypatch.setattr(utils, "executor", executor) + monkeypatch.setattr(litellm_logging, "executor", executor) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", worker) + + class Lifecycle(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + pre_calls.append((kwargs, context.get(), threading.get_ident(), asyncio.current_task())) + kwargs["ocr_lifecycle_state"] = shared + context.set("pre_call") + + def record(self, event, kwargs, response_obj, start_time, end_time): + terminal_calls.append( + { + "event": event, + "kwargs": kwargs, + "response": response_obj, + "exception": kwargs.get("exception"), + "shared": kwargs.get("ocr_lifecycle_state"), + "phase": kwargs.get("ocr_lifecycle_state", {}).get("phase"), + "context": context.get(), + "thread": threading.get_ident(), + "task": asyncio.current_task() if threading.get_ident() == caller_thread else None, + "start_time": start_time, + "end_time": end_time, + } + ) + kwargs["ocr_lifecycle_state"]["phase"] = "terminal" + context.set("terminal") + finished.set() + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.record("sync_success", kwargs, response_obj, start_time, end_time) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.record("sync_failure", kwargs, response_obj, start_time, end_time) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await asyncio.sleep(0) + self.record("async_success", kwargs, response_obj, start_time, end_time) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await asyncio.sleep(0) + self.record("async_failure", kwargs, response_obj, start_time, end_time) + + callback = Lifecycle() + monkeypatch.setattr(litellm, "callbacks", [callback] if callback_source == "global" else []) + arguments = dict( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + num_retries=0, + **({"callbacks": [callback]} if callback_source == "per-call" else {}), + ) + try: + if failure: + with pytest.raises(litellm.RateLimitError) as caught: + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) + assert caught.value.status_code == 429 + assert caught.value.model == "mistral-ocr-latest" + assert caught.value.llm_provider == "mistral" + else: + response = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "proof" + assert await asyncio.to_thread(finished.wait, 5), "terminal callback was not delivered" + finally: + try: + await asyncio.wait_for(worker.flush(), timeout=5) + await asyncio.wait_for(asyncio.wrap_future(executor.submit(lambda: None)), timeout=5) + finally: + try: + await asyncio.wait_for(worker.stop(), timeout=5) + finally: + atexit.unregister(worker._flush_on_exit) + executor.shutdown(wait=False, cancel_futures=True) + + assert len(wire_recorder.requests) == 1 + assert len(pre_calls) == 1 + details, pre_context, pre_thread, pre_task = pre_calls[0] + assert (pre_context, pre_thread, pre_task) == ("caller", caller_thread, caller_task) + assert details["additional_args"]["complete_input_dict"]["document"] is document + expected_events = ( + {"sync_failure", "async_failure"} + if failure and asynchronous + else {f"{'async' if asynchronous else 'sync'}_{'failure' if failure else 'success'}"} + ) + assert len(terminal_calls) == len(expected_events) + assert {terminal["event"] for terminal in terminal_calls} == expected_events + assert shared == {"phase": "terminal"} + assert details["litellm_call_id"] + for terminal in terminal_calls: + assert terminal["kwargs"] is details + assert terminal["shared"] is shared + expected_phase = "terminal" if terminal["event"] == "async_failure" else "pre_call" + assert terminal["phase"] == expected_phase + assert terminal["context"] == expected_phase + assert terminal["start_time"] <= terminal["end_time"] + if failure: + assert terminal["response"] is None + assert terminal["exception"] is caught.value + assert "RateLimitError" in details["traceback_exception"] + assert (terminal["thread"], terminal["task"]) == (caller_thread, caller_task) + else: + assert terminal["response"] is response + assert terminal["exception"] is None + assert terminal["task"] is not caller_task + assert (terminal["thread"] == caller_thread) is asynchronous + if asynchronous: + assert terminal["task"] is not None + assert context.get() == ("terminal" if failure else "pre_call") @pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("public", [False, True], ids=["unwrapped", "public"]) @pytest.mark.parametrize( - "asynchronous, control, message", + "model,provider,options,message", [ - (False, "copy-input", "caller identity"), - (True, "copy-input", "caller identity"), - (False, "duplicate-prepare", "callback count"), - (True, "duplicate-prepare", "callback count"), - (True, "new-task", "callback context/order"), - (False, "logging-roots", "execution roots"), - (True, "logging-roots", "execution roots"), + ("azure_ai/doc-intelligence/prebuilt-read", None, {}, "Document Intelligence OCR polling"), + ("documentintelligence/prebuilt-layout", "azure_ai", {}, "Document Intelligence OCR polling"), + ("azure_ai/mistral-ocr-latest", None, {}, "HTTP document URL to data URI conversion"), + ("vertex_ai/mistral-ocr-latest", None, {}, "HTTP document URL to data URI conversion"), + ("mistral-ocr-latest", "vertex_ai", {}, "HTTP document URL to data URI conversion"), + ("azure_ai/mistral-ocr-latest", None, {"api_key": None}, "Azure OCR credential acquisition"), + ("vertex_ai/mistral-ocr-latest", None, {"api_key": None}, "Vertex OCR credential acquisition"), + ("vertex_ai/deepseek-ocr-maas", None, {"api_key": None}, "Vertex OCR credential acquisition"), + ("azure_ai/cohere/parse-v5.0", None, {}, "Cohere OCR request transformation"), + ("cohere/parse-v5.0", None, {}, "OCR provider"), + ("vertex_ai/deepseek-ocr-maas", None, {"stream": True}, "OCR streaming response handling"), + (MODEL, None, {"document": {"type": "file", "file": b"%PDF-1.4"}}, "file"), + (MODEL, None, {"req_format": "native"}, "native"), ], ) -async def test_pre_call_contract_rejects_boundary_mutants( - native_ocr, wire_recorder, monkeypatch, asynchronous, control, message +async def test_native_unsupported_requests_never_prepare_or_send( + native_ocr, + wire_recorder, + monkeypatch, + no_python_ocr, + document, + asynchronous, + public, + model, + provider, + options, + message, ): - with pytest.raises(AssertionError, match=message): - await _pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, True, control=control) + monkeypatch.setenv("LITELLM_RUST", "1") + for name in ("AZURE_AI_API_KEY", "VERTEX_AI_API_KEY", "VERTEXAI_API_KEY"): + monkeypatch.delenv(name, raising=False) + logger = Mock() + arguments = { + "model": model, + "document": document, + "custom_llm_provider": provider, + "api_key": "sk-test", + "api_base": wire_recorder.api_base, + "timeout": 3, + **({} if public else {"litellm_logging_obj": logger}), + "num_retries": 0, + **options, + } + with pytest.raises(NotImplementedError, match=message): + function = litellm.aocr if asynchronous else litellm.ocr + result = (function if public else inspect.unwrap(function))(**arguments) + if asynchronous: + await result + logger.assert_not_called() + assert logger.mock_calls == [] + assert wire_recorder.requests == [] @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -async def test_pre_call_contract_allows_root_tuple_reconstruction(native_ocr, wire_recorder, monkeypatch, asynchronous): - await _pre_call_contract(native_ocr, wire_recorder, monkeypatch, asynchronous, True, control="tuple-only") +@pytest.mark.parametrize("auth", ["key", "header", "environment"]) +@pytest.mark.parametrize("provider", ["azure_ai", "vertex_ai", "deepseek"]) +async def test_public_native_cloud_wire_and_shallow_boundaries( + native_ocr, wire_recorder, monkeypatch, no_python_ocr, asynchronous, auth, provider +): + monkeypatch.setenv("LITELLM_RUST", "1") + deepseek = provider == "deepseek" + model = "vertex_ai/deepseek-ai/deepseek-ocr-maas" if deepseek else f"{provider}/mistral-ocr-latest" + document = {"type": "image_url", "image_url": "data:image/png;base64,YWJj", "nested": []} + shared_param = ["stop"] if deepseek else [0] + param = "stop" if deepseek else "pages" + captured = [] + opaque = object() + + class Mutate(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + view = kwargs["additional_args"] + body = view["complete_input_dict"] + captured.append(body) + assert body[param] is shared_param + shared_param.append("changed" if deepseek else 2) + document["image_url"] = "https://example.invalid/original-only.png" + if deepseek: + assert "document" not in body + body["messages"][0]["content"][0]["image_url"] = "data:image/png;base64,ZGVm" + else: + assert body["document"] is not document + assert body["document"]["nested"] is document["nested"] + document["nested"].append("shared") + assert body["document"]["image_url"] == "data:image/png;base64,YWJj" + body["document"]["image_url"] = "data:image/png;base64,ZGVm" + view["headers"]["X-Callback"] = "native" + view["complete_input_dict"] = {"replacement": True} + + monkeypatch.setattr(litellm, "input_callback", [Mutate()]) + if deepseek: + wire_recorder.response_data = { + "choices": [{"message": {"content": "proof"}}], + "usage": {"pages_processed": 1}, + } + if auth == "environment": + monkeypatch.setenv("AZURE_AI_API_KEY" if provider == "azure_ai" else "VERTEX_AI_API_KEY", "native-token") + arguments = dict( + model=model, + document=document, + api_key="native-token" if auth == "key" else None, + extra_headers={"aUtHoRiZaTiOn": "Bearer native-token"} if auth == "header" else {}, + api_base=wire_recorder.api_base, + vertex_ai_project="test-project", + vertex_ai_location="europe-west4", + timeout=3, + num_retries=0, + metadata={"opaque": opaque}, + vertex_credentials=opaque, + arbitrary_option=opaque, + **{param: shared_param}, + ) + response = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "proof" + assert response.usage_info.pages_processed == 1 + assert len(captured) == len(wire_recorder.requests) == 1 + sent = wire_recorder.requests[0] + auth_header = "api-key" if provider == "azure_ai" and auth != "header" else "authorization" + assert sent["headers"][auth_header] == ("native-token" if auth_header == "api-key" else "Bearer native-token") + assert sent["headers"]["x-callback"] == "native" + assert sent["body"][param] == shared_param + assert "arbitrary_option" not in sent["body"] and "metadata" not in sent["body"] + assert "vertex_credentials" not in sent["body"] and "vertex_ai_project" not in sent["body"] + if deepseek: + assert sent["path"] == "/v1/projects/test-project/locations/europe-west4/endpoints/openapi/chat/completions" + assert sent["body"] == { + "model": "deepseek-ai/deepseek-ocr-maas", + "messages": [ + {"role": "user", "content": [{"type": "image_url", "image_url": "data:image/png;base64,ZGVm"}]} + ], + "stop": ["stop", "changed"], + } + else: + expected_path = ( + "/providers/mistral/azure/ocr" + if provider == "azure_ai" + else "/v1/projects/test-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-latest:rawPredict" + ) + assert sent["path"] == expected_path + assert sent["body"] == { + "model": "mistral-ocr-latest", + "document": {"type": "image_url", "image_url": "data:image/png;base64,ZGVm", "nested": ["shared"]}, + "pages": [0, 2], + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_public_native_azure_supplied_entra_token( + native_ocr, wire_recorder, monkeypatch, no_python_ocr, asynchronous +): + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + arguments = dict( + model="azure_ai/mistral-ocr-latest", + document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + azure_ad_token="entra-token", + api_base=wire_recorder.api_base, + timeout=3, + num_retries=0, + ) + response = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert response.pages[0].markdown == "proof" + assert wire_recorder.requests[0]["headers"]["authorization"] == "Bearer entra-token" class PreCallAbort(BaseException): @@ -958,247 +887,117 @@ class PreCallAbort(BaseException): @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) -@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) -async def test_pre_call_escape_never_sends_or_replays(native_ocr, wire_recorder, monkeypatch, asynchronous, enabled): - document = {"type": "document_url", "document_url": "https://example.invalid/original.pdf"} +async def test_native_callback_escape_never_sends_or_replays( + native_ocr, wire_recorder, monkeypatch, no_python_ocr, document, asynchronous +): + litellm.rust(True) error = PreCallAbort("stop before POST") calls = [] - native_calls = [] class Abort(CustomLogger): def log_pre_api_call(self, model, messages, kwargs): calls.append(kwargs["additional_args"]["complete_input_dict"]["document"]) - document["document_url"] = "https://example.invalid/aborted.pdf" raise error - def sync_call(boundary): - native_calls.append("sync") - return native_ocr.ocr(boundary) - - async def async_call(boundary): - native_calls.append("async") - return await native_ocr.aocr(boundary) - - rust_bridge._OCR.override(sync_call) - rust_bridge._AOCR.override(async_call) monkeypatch.setattr(litellm, "input_callback", [Abort()]) - litellm.rust(enabled) arguments = dict(model=MODEL, document=document, api_key="sk-test", api_base=wire_recorder.api_base, num_retries=0) - with pytest.raises(PreCallAbort, match="stop before POST") as caught: - await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + with pytest.raises(PreCallAbort) as caught: + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) assert caught.value is error assert len(calls) == 1 and calls[0] is document - assert native_calls == (["async" if asynchronous else "sync"] if enabled else []) - assert document["document_url"] == "https://example.invalid/aborted.pdf" assert wire_recorder.requests == [] -class RetainedDocument(dict): +def test_native_sync_callback_reentry_without_event_loop( + native_ocr, wire_recorder, monkeypatch, no_python_ocr, document +): + litellm.rust(True) + context = contextvars.ContextVar("ocr-reentry", default="caller") + events = [] + arguments = dict(model=MODEL, document=document, api_key="sk-test", api_base=wire_recorder.api_base, num_retries=0) + + class Reenter(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + with pytest.raises(RuntimeError, match="no running event loop"): + asyncio.get_running_loop() + events.append((context.get(), threading.get_ident())) + if len(events) == 1: + context.set("nested") + result = litellm.ocr(**arguments) + events.append((result.pages[0].markdown, threading.get_ident())) + + monkeypatch.setattr(litellm, "input_callback", [Reenter()]) + response = litellm.ocr(**arguments) + assert response.pages[0].markdown == "proof" + assert events == [(value, threading.get_ident()) for value in ("caller", "nested", "proof")] + assert context.get() == "nested" + assert len(wire_recorder.requests) == 2 + + +class Opaque: pass @pytest.mark.asyncio -@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) @pytest.mark.parametrize("outcome", ["success", "error", "cancel"]) -async def test_ocr_retention_during_post_and_terminal_cleanup(native_ocr, wire_recorder, enabled, outcome): - retained = [] - references = [] - calls = [] +async def test_native_retains_opaque_arguments_until_terminal_cleanup(native_ocr, wire_recorder, document, outcome): wire_recorder.release.clear() wire_recorder.status = 429 if outcome == "error" else 200 - - class Retain(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - body = kwargs["additional_args"]["complete_input_dict"] - retained.append(body) - references.append(weakref.ref(body["document"])) - calls.append("pre_call") - - async def request(): - document = RetainedDocument(type="document_url", document_url="https://example.invalid/original.pdf") - logging_obj = Logging( - model="mistral-ocr-latest", - messages=[], - stream=False, - call_type="aocr", - start_time=datetime.now(), - litellm_call_id="retention", - function_id="retention", - dynamic_input_callbacks=[Retain()], - ) - arguments = dict( - model="mistral-ocr-latest", - document=document, - optional_params={}, - logging_obj=logging_obj, - api_key="sk-test", - api_base=wire_recorder.api_base, - headers=None, - provider_config=MistralOCRConfig(), - litellm_params={}, - custom_llm_provider="mistral", - timeout=5.0, - ) - if enabled: - return await native_ocr.aocr(rust_bridge.OCRBoundary(handler=ocr_main.base_llm_http_handler, **arguments)) - return await ocr_main.base_llm_http_handler.async_ocr(**arguments) - - task = asyncio.create_task(request()) + opaque = Opaque() + reference = weakref.ref(opaque) + logger = Mock() + logger.async_success_handler = AsyncMock() + logger.async_failure_handler = AsyncMock() + logger._should_run_sync_callbacks_for_async_calls.return_value = False + arguments = dict( + model=MODEL, + document=document, + api_key="sk-test", + api_base=wire_recorder.api_base, + timeout=5, + litellm_logging_obj=logger, + opaque=opaque, + ) + task = asyncio.create_task(native_ocr.aocr(arguments)) + del arguments, opaque try: assert await asyncio.to_thread(wire_recorder.received.wait, 5), "POST never reached server" assert not task.done() - assert calls == ["pre_call"] - assert len(references) == 1 and references[0]() is not None - retained[0]["document"]["document_url"] = "https://example.invalid/after-consumption.pdf" - retained[0]["document"]["cycle"] = retained[0] - assert wire_recorder.requests[0]["body"]["document"]["document_url"] == "https://example.invalid/original.pdf" + logger.update_from_kwargs.assert_called_once() + logger.pre_call.assert_called_once() + logger.reset_mock() + gc.collect() + assert reference() is not None if outcome == "cancel": task.cancel() with pytest.raises(asyncio.CancelledError): await task elif outcome == "error": wire_recorder.release.set() - with pytest.raises(BaseLLMException) as caught: + with pytest.raises(litellm.RateLimitError) as caught: await task assert caught.value.status_code == 429 del caught else: wire_recorder.release.set() - response = await task - assert response.pages[0].markdown == "proof" + result = await task + assert result.pages[0].markdown == "proof" finally: wire_recorder.release.set() if not task.done(): task.cancel() await asyncio.gather(task, return_exceptions=True) - assert await asyncio.to_thread(wire_recorder.finished.wait, 5), "server did not finish" + assert await asyncio.to_thread(wire_recorder.finished.wait, 5) del task - gc.collect() - assert references[0]() is retained[0]["document"] - assert retained[0]["document"]["document_url"] == "https://example.invalid/after-consumption.pdf" - assert len(wire_recorder.requests) == 1 and calls == ["pre_call"] - retained.clear() + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + if outcome == "success": + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), 5) + logger.reset_mock() await asyncio.sleep(0) gc.collect() - assert references[0]() is None, "execution retained the caller graph after cleanup" - - -@pytest.mark.parametrize("filename", ["ocr_driver.py", "retained_callback.py"]) -def test_native_ocr_cold_cache_reentry(native_ocr, filename): - script = """ -import sys -from litellm.rust_bridge import _native - -events = [] -compilations = [] -error = LookupError('preparation stopped') - -class Boundary: - async def aprepare(self): - raise error - -def invoke(): - pending = _native.aocr(Boundary()) - try: - pending.send(None) - except LookupError as caught: - assert caught is error - events.append('raised') - else: - raise AssertionError('preparation did not raise') - finally: - pending.close() - error.__traceback__ = None - -def audit(event, args): - if event == 'compile' and args[1] == sys.argv[1]: - compilations.append(args[1]) - if len(compilations) == 1: - events.append('entered') - invoke() - events.append('returned') - -sys.addaudithook(audit) -invoke() -invoke() -assert events == ['entered', 'raised', 'returned', 'raised', 'raised'], events -assert len(compilations) == 2, compilations -print('cold reentry passed') -""" - isolation = ["-I"] if sys.flags.isolated else [] - result = subprocess.run( - [sys.executable, *isolation, "-c", script, filename], capture_output=True, text=True, timeout=30 - ) - assert result.returncode == 0, result.stdout + result.stderr - assert result.stdout.strip() == "cold reentry passed" - - -@pytest.mark.parametrize("enabled", [False, True], ids=["python", "native"]) -def test_sync_pre_call_reentry_without_event_loop(native_ocr, wire_recorder, monkeypatch, enabled): - context = contextvars.ContextVar("sync-ocr-context", default="caller") - events = [] - native_calls = [] - document = {"type": "document_url", "document_url": "https://example.invalid/original.pdf"} - - class Reenter(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - try: - asyncio.get_running_loop() - except RuntimeError: - loop_running = False - else: - loop_running = True - events.append((context.get(), threading.get_ident(), loop_running)) - if len(events) == 1: - context.set("nested") - response = litellm.ocr( - model=MODEL, - document=dict(document), - api_key="sk-test", - api_base=wire_recorder.api_base, - num_retries=0, - ) - events.append((response.pages[0].markdown, threading.get_ident(), loop_running)) - - def sync_call(boundary): - native_calls.append(boundary) - return native_ocr.ocr(boundary) - - rust_bridge._OCR.override(sync_call) - monkeypatch.setattr(litellm, "input_callback", [Reenter()]) - litellm.rust(enabled) - response = litellm.ocr( - model=MODEL, - document=document, - api_key="sk-test", - api_base=wire_recorder.api_base, - num_retries=0, - ) - assert response.pages[0].markdown == "proof" - assert events == [(value, threading.get_ident(), False) for value in ("caller", "nested", "proof")] - assert context.get() == "nested" - assert len(native_calls) == (2 if enabled else 0) - assert len(wire_recorder.requests) == 2 - - -@pytest.mark.parametrize("explicit_close", [False, True], ids=["abandoned", "closed"]) -def test_unstarted_native_ocr_driver_releases_cyclic_input(native_ocr, explicit_close): - document = RetainedDocument(type="document_url", document_url="https://example.invalid/original.pdf") - reference = weakref.ref(document) - logger = RecordingLogging() - boundary = ocr_main._ocr_boundary(build_prepared_request(document=document, logging_obj=logger)) - pending = native_ocr.aocr(boundary) - document["pending"] = pending - del boundary, document - assert reference() is not None - assert logger.pre_call_kwargs is None - if explicit_close: - pending.close() - del pending - gc.collect() - else: - del pending - with pytest.warns(RuntimeWarning, match="coroutine .* was never awaited"): - gc.collect() - assert reference() is None - assert logger.pre_call_kwargs is None + assert reference() is None, "native execution retained opaque kwargs after cleanup" + assert len(wire_recorder.requests) == 1 diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 4f5aed50985..7427eba5fa5 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -15,8 +15,9 @@ from http.client import HTTPMessage from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from socket import socket as Socket +from types import ModuleType from typing import Final -from urllib.error import HTTPError +from unittest.mock import patch REQUEST_STARTED: Final = threading.Event() REQUEST_CANCELLED: Final = threading.Event() @@ -87,6 +88,8 @@ def assert_native_request( assert body["model"] == "mistral-ocr-latest" assert body["document"]["document_url"] == "https://example.com/document.pdf" assert body["include_image_base64"] is True + assert headers.get("x-ocr-callback") == "inline" + assert set(body) == {"model", "document", "include_image_base64"} return if route == "transcription": assert path == "/model/mistral.voxtral-mini-3b-2507/converse" @@ -110,7 +113,7 @@ def native_response(status: int, route: str | None) -> bytes: if status == 429: return b'{"error":"native-rate-limit"}' if route == "ocr": - return b'{"pages":[{"index":0,"markdown":"native-ocr"}]}' + return b'{"pages":[{"index":0,"markdown":"native-ocr"}],"usage_info":{"pages_processed":1}}' if route == "transcription": return b'{"output":{"message":{"content":[{"text":"native-transcription"}]}}}' return ANTHROPIC_RESPONSE @@ -126,48 +129,97 @@ def load_native(native_path: Path) -> object: @dataclass(frozen=True) -class OCRBoundary: - api_base: str - outcome: str +class WheelOCRResponse: + pages: list[dict[str, object]] + model: str + document_annotation: object + usage_info: dict[str, object] | None + object: str - def prepare(self) -> dict[str, object]: - return { - "model": "mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, - "include_image_base64": True, - } - async def aprepare(self) -> dict[str, object]: - return self.prepare() +@dataclass(frozen=True) +class WheelHTTPRequest: + method: str + url: str - def encode(self, roots: dict[str, object]) -> tuple[str, list[tuple[bytes, bytes]], bytes, float]: - return ( - f"{self.api_base}/v1/ocr", - [ - (b"authorization", b"Bearer sk-native"), - (b"content-type", b"application/json"), - (b"x-test-route", b"ocr"), - (b"x-test-outcome", self.outcome.encode()), - ], - json.dumps(roots).encode(), - 3.0, - ) - def finish(self, wire: tuple[int, list[tuple[bytes, bytes]], bytes]) -> object: - status, headers, content = wire - assert (b"content-type", b"application/json") in headers - if status != 200: - assert content == b'{"error":"native-rate-limit"}' - raise HTTPError(f"{self.api_base}/v1/ocr", status, "native-rate-limit", HTTPMessage(), None) - return json.loads(content) +@dataclass(frozen=True) +class WheelHTTPResponse: + status_code: int + request: WheelHTTPRequest - async def afinish(self, wire: tuple[int, list[tuple[bytes, bytes]], bytes]) -> object: - return self.finish(wire) + +class WheelRateLimitError(Exception): + def __init__(self, *, message: str, model: str, llm_provider: str, response: WheelHTTPResponse) -> None: + super().__init__(message) + self.status_code = response.status_code + self.model = model + self.llm_provider = llm_provider + + +class OCRLogging: + def __init__(self, arguments: dict[str, object]) -> None: + self.arguments = arguments + self.calls: tuple[str, ...] = () + self.thread_id = threading.get_ident() + try: + self.task = asyncio.current_task() + except RuntimeError: + self.task = None + + def update_from_kwargs(self, **kwargs: object) -> None: + assert self.calls == () + assert threading.get_ident() == self.thread_id + if self.task is not None: + assert asyncio.current_task() is self.task + assert kwargs["kwargs"] is self.arguments + assert kwargs["model"] == "mistral-ocr-latest" + assert kwargs["custom_llm_provider"] == "mistral" + assert kwargs["optional_params"] == {"include_image_base64": True} + self.calls = ("update",) + + def pre_call(self, **kwargs: object) -> None: + assert self.calls == ("update",) + assert threading.get_ident() == self.thread_id + if self.task is not None: + assert asyncio.current_task() is self.task + assert kwargs["input"] == "OCR document processing" + assert kwargs["api_key"] == "sk-native" + additional_args: Final = kwargs["additional_args"] + assert isinstance(additional_args, dict) + assert additional_args["api_base"] == f"{self.arguments['api_base']}/v1/ocr" + assert additional_args["complete_input_dict"]["document"] is self.arguments["document"] + assert additional_args["headers"]["Authorization"] == "Bearer sk-native" + additional_args["headers"]["x-ocr-callback"] = "inline" + self.calls = ("update", "pre") + + +def initialize_ocr_logging(arguments: dict[str, object], asynchronous: bool) -> object: + return arguments["litellm_logging_obj"] + + +def invoke_ocr_terminal( + action: str, roots: object, logger: OCRLogging, value: object, start: object, end: object +) -> None: + assert action in {"sync_success", "async_success", "sync_success_if_needed", "sync_failure", "async_failure"} + assert isinstance(roots, tuple) and roots[0] is logger.arguments + assert logger.calls == ("update", "pre") def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: if route == "ocr": - return {"boundary": OCRBoundary(api_base, outcome)} + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/document.pdf"}, + "include_image_base64": True, + "api_base": api_base, + "api_key": "sk-native", + "extra_headers": {"x-test-outcome": outcome, "x-test-route": route}, + "timeout": 3.0, + "metadata": {"opaque": object()}, + } + arguments["litellm_logging_obj"] = OCRLogging(arguments) + return {"arguments": arguments} common: Final = { "api_base": api_base, "extra_headers": {"x-test-outcome": outcome, "x-test-route": route}, @@ -207,19 +259,23 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: def assert_success(route: str, response: object) -> None: + if route == "ocr": + assert isinstance(response, WheelOCRResponse), f"expected OCRResponse construction, got {response!r}" + assert response.pages == [{"index": 0, "markdown": "native-ocr"}] + assert response.model == "mistral-ocr-latest" + assert response.object == "ocr" + assert response.document_annotation is None + assert response.usage_info == {"pages_processed": 1} + return if not isinstance(response, dict): raise TypeError(f"{route} returned {type(response).__name__}, expected dict") actual: Final = success_value(route, response) - expected: Final = ( - "native-ocr" if route == "ocr" else "native-transcription" if route == "transcription" else "native-message" - ) + expected: Final = "native-transcription" if route == "transcription" else "native-message" if actual != expected: raise AssertionError(f"{route} returned {actual!r}, expected {expected!r}") def success_value(route: str, response: dict[object, object]) -> object: - if route == "ocr": - return response["pages"][0]["markdown"] if route == "transcription": return response["text"] if route == "messages": @@ -229,8 +285,11 @@ def success_value(route: str, response: dict[object, object]) -> object: def assert_rate_limit(native: object, route: str, error: BaseException) -> None: if route == "ocr": - if not isinstance(error, HTTPError) or error.code != 429: + if not isinstance(error, WheelRateLimitError) or error.status_code != 429: raise AssertionError(f"{route} returned the wrong 429 error: {error!r}") + assert error.model == "mistral-ocr-latest" + assert error.llm_provider == "mistral" + assert str(error) == "OCR provider request failed (HTTP 429)" return if route == "chat_completions": upstream_error: Final = native.RustUpstreamError @@ -247,7 +306,7 @@ def exercise_sync(native: object, api_base: str) -> None: assert_success(route, function(**route_kwargs(route, api_base, "success"))) try: function(**route_kwargs(route, api_base, "429")) - except (HTTPError, RuntimeError, native.RustUpstreamError) as error: + except (WheelRateLimitError, RuntimeError, native.RustUpstreamError) as error: assert_rate_limit(native, route, error) else: raise AssertionError(f"{route} accepted a 429 response") @@ -259,12 +318,34 @@ async def exercise_async(native: object, api_base: str) -> None: assert_success(route, await function(**route_kwargs(route, api_base, "success"))) try: await function(**route_kwargs(route, api_base, "429")) - except (HTTPError, RuntimeError, native.RustUpstreamError) as error: + except (WheelRateLimitError, RuntimeError, native.RustUpstreamError) as error: assert_rate_limit(native, route, error) else: raise AssertionError(f"a{route} accepted a 429 response") +async def exercise_unsupported_ocr(native: object, api_base: str) -> None: + for model, operation in ( + ("azure_ai/doc-intelligence/prebuilt-read", "Document Intelligence OCR polling"), + ("vertex_ai/ocr", "HTTP document URL to data URI conversion"), + ): + arguments: Final = route_kwargs("ocr", api_base, "success")["arguments"] + assert isinstance(arguments, dict) + arguments["model"] = model + arguments["document"] = {"type": "document_url", "document_url": "https://example.invalid/document.pdf"} + for asynchronous in (False, True): + try: + if asynchronous: + await native.aocr(arguments) + else: + native.ocr(arguments) + except NotImplementedError as error: + assert operation in str(error) + else: + raise AssertionError(f"native OCR accepted unsupported operation: {operation}") + assert arguments["litellm_logging_obj"].calls == () + + async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), @@ -278,8 +359,33 @@ def exercise_routes(native_path: Path, api_base: str) -> object: native: Final = load_native(native_path) if hasattr(native, "_trace"): raise AssertionError("release wheel exposed trace-parity diagnostics") - exercise_sync(native, api_base) - asyncio.run(exercise_async(native, api_base)) + transformation: Final = ModuleType("litellm.llms.base_llm.ocr.transformation") + transformation.OCRResponse = WheelOCRResponse + exceptions: Final = ModuleType("litellm.exceptions") + exceptions.RateLimitError = WheelRateLimitError + httpx: Final = ModuleType("httpx") + httpx.Request = WheelHTTPRequest + httpx.Response = WheelHTTPResponse + ocr_bridge: Final = ModuleType("litellm.rust_bridge.ocr") + ocr_bridge.initialize_logging = initialize_ocr_logging + ocr_bridge.invoke_terminal = invoke_ocr_terminal + packages: Final = { + name: ModuleType(name) + for name in ( + "litellm", + "litellm.rust_bridge", + "litellm.llms", + "litellm.llms.base_llm", + "litellm.llms.base_llm.ocr", + ) + } + with patch.dict( + sys.modules, + packages | {module.__name__: module for module in (transformation, exceptions, httpx, ocr_bridge)}, + ): + exercise_sync(native, api_base) + asyncio.run(exercise_async(native, api_base)) + asyncio.run(exercise_unsupported_ocr(native, api_base)) asyncio.run(exercise_async_concurrency(native, api_base)) return native