mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
wip
This commit is contained in:
Yujong Lee
parent
ff958cc30b
commit
b1f74b3603
19 changed files with 3516 additions and 1492 deletions
212
.opencode/plugins/litellm.ts
Normal file
212
.opencode/plugins/litellm.ts
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
import type { Plugin } from "@opencode-ai/plugin"
|
||||
|
||||
type JsonObject = Record<string, unknown>
|
||||
|
||||
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<string, { reasoningEffort: string }>
|
||||
}
|
||||
|
||||
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<string, LiteLLMModelGroup> {
|
||||
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<unknown> {
|
||||
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 {}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Map<String, Value>>,
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
|
|
|
|||
|
|
@ -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<OcrResponseData, Error> {
|
||||
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)
|
||||
}
|
||||
|
|
|
|||
115
litellm-rust/crates/core/src/ocr/prepare.rs
Normal file
115
litellm-rust/crates/core/src/ocr/prepare.rs
Normal file
|
|
@ -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<PreparedOcr, Error> {
|
||||
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")),
|
||||
}
|
||||
}
|
||||
|
|
@ -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<String>,
|
||||
) -> Result<Vec<(String, String)>, 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)]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,75 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub struct OcrRequest {
|
||||
pub model: String,
|
||||
pub custom_llm_provider: Option<String>,
|
||||
pub api_key: Option<String>,
|
||||
pub api_base: Option<String>,
|
||||
pub extra_headers: Vec<(String, String)>,
|
||||
pub timeout_seconds: f64,
|
||||
pub request_format: Option<String>,
|
||||
pub document: OcrDocument,
|
||||
pub azure_ad_token: Option<String>,
|
||||
pub vertex_project: Option<String>,
|
||||
pub vertex_location: Option<String>,
|
||||
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<String, Value>,
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
) -> Result<Vec<(String, String)>, 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<String>,
|
||||
) -> Result<Vec<(String, String)>, 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
|
||||
|
|
|
|||
|
|
@ -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<String, Value>, 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<Value>, 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
|
||||
|
|
|
|||
632
litellm-rust/crates/core/tests/ocr.rs
Normal file
632
litellm-rust/crates/core/tests/ocr.rs
Normal file
|
|
@ -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<_>>(),
|
||||
vec!["content-type: application/vnd.ocr+json"]
|
||||
);
|
||||
assert_eq!(
|
||||
headers
|
||||
.lines()
|
||||
.filter(|line| line.starts_with("accept-encoding:"))
|
||||
.collect::<Vec<_>>(),
|
||||
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"
|
||||
})
|
||||
);
|
||||
}
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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<Py<PyDict>>,
|
||||
body: Option<Py<PyDict>>,
|
||||
headers: Option<Py<PyDict>>,
|
||||
logging: Option<Py<PyAny>>,
|
||||
prepared: Option<PreparedOcr>,
|
||||
}
|
||||
|
||||
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<Py<PyAny>> {
|
||||
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<Py<PyAny>> {
|
||||
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<Request> {
|
||||
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<Self::Output> {
|
||||
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<PyErr> {
|
||||
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<Option<String>> {
|
||||
arguments
|
||||
.get_item(name)?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| value.extract::<String>())
|
||||
.transpose()
|
||||
.map(|value| value.filter(|value| !value.trim().is_empty()))
|
||||
}
|
||||
|
||||
fn header_pairs(headers: &Bound<'_, PyDict>) -> PyResult<Vec<(String, String)>> {
|
||||
headers
|
||||
.iter()
|
||||
.map(|(name, value)| Ok((name.extract()?, value.extract()?)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn send<'a>(
|
||||
boundary: &'a Bound<'a, PyAny>,
|
||||
roots: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
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<PyDict>, asynchronous: bool) -> PyResult<Py<OcrState>> {
|
||||
let bag = arguments.bind(py);
|
||||
let document = bag
|
||||
.get_item("document")?
|
||||
.ok_or_else(|| PyValueError::new_err("OCR requires document"))?
|
||||
.cast_into::<PyDict>()?;
|
||||
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::<f64>()?,
|
||||
Some(timeout) => match timeout.extract::<f64>() {
|
||||
Ok(seconds) => seconds,
|
||||
Err(_) => timeout.getattr("read")?.extract::<f64>()?,
|
||||
},
|
||||
};
|
||||
let extra_headers = match bag
|
||||
.get_item("extra_headers")?
|
||||
.filter(|value| !value.is_none())
|
||||
{
|
||||
Some(headers) => header_pairs(headers.cast::<PyDict>()?)?,
|
||||
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::<bool>())
|
||||
.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::<PyDict>()?;
|
||||
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<OcrState>) -> PyResult<OcrWireRequest> {
|
||||
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<OcrState>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<Py<PyAny>> {
|
||||
let step = if asynchronous {
|
||||
FINISH_ASYNC
|
||||
} else {
|
||||
FINISH_SYNC
|
||||
};
|
||||
invoke(boundary, step, PyTuple::new(boundary.py(), [wire])?)
|
||||
fn finish(py: Python<'_>, response: Py<PyDict>) -> PyResult<Py<PyAny>> {
|
||||
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<Py<PyAny>> {
|
||||
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<OcrState>) -> PyResult<Py<PyAny>> {
|
||||
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::<PyDict>()?;
|
||||
let response = finish(py, fields.unbind());
|
||||
drop(state);
|
||||
response
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn aocr<'a>(boundary: &'a Bound<'a, PyAny>) -> PyResult<Bound<'a, PyAny>> {
|
||||
driver(boundary.py())?.getattr("drive")?.call1((boundary,))
|
||||
fn ocr(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<PyDict>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<Py<PyModule>> = 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::<u16>()
|
||||
.unwrap(),
|
||||
status
|
||||
);
|
||||
assert_eq!(
|
||||
value.getattr("model").unwrap().extract::<String>().unwrap(),
|
||||
"mistral-ocr-latest"
|
||||
);
|
||||
assert_eq!(
|
||||
value
|
||||
.getattr("llm_provider")
|
||||
.unwrap()
|
||||
.extract::<String>()
|
||||
.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::<PyRuntimeError>(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<OcrState>) -> PyResult<Py<PyAny>> {
|
||||
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();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<PyDict>,
|
||||
#[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")]
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue