diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 36a5ad5a8f4..71262b32cd6 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,16 +1,17 @@ # AGENTS.md -litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers. +litellm-rust has exactly FOUR crates. A crate is a LAYER, not a route. Routes and providers are modules inside the layers. `litellm-runtime` is an intentional reusable execution layer, introduced OCR-first. ## Crates | Crate | Role | |-------|------| -| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | -| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-core | Provider transformations and shared types. OCR core code is deterministic transformation only. Existing non-OCR route entrypoints remain in core. | +| litellm-runtime | Reusable OCR execution: provider/model and config selection, environment/auth, URL/headers, document materialization, HTTP, polling, and lifecycle ordering. | +| litellm-ai-gateway | The axum server and WebSocket hosts. Adapts host logger, guardrail, and metadata types to runtime interfaces. | +| litellm-python-bridge | PyO3 cdylib exposing Rust to the Python SDK. OCR calls runtime directly; unrelated legacy paths may still call gateway or core. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction (acyclic): ai-gateway/python-bridge -> runtime -> core. Python bridge retains an ai-gateway dependency for unrelated audio and WebSocket code. ## Where a route lives @@ -26,7 +27,7 @@ core/src/messages/ client.rs # the shared reqwest client ``` -Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched. +OCR provider handlers live in `runtime`, not `ai-gateway`. This boundary is OCR-first; audio, realtime, Messages, Chat, Responses, and generic lifecycle code do not move with it. Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index fe6ceedbb86..5995bf50526 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -21,12 +21,12 @@ variants of it. The test for a good abstraction is that adding the next provider is a few declarative lines, not a new file of duplicated flow. Only diverge from the base when behavior is genuinely different, and say so explicitly in the PR. -## Crates (exactly three — see AGENTS.md) +## Crates (exactly four, see AGENTS.md) -`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. -`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and -`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not -a route — add modules, not crates. +`litellm-core` owns provider transformations and shared types. `litellm-runtime` +is the reusable OCR execution layer. `litellm-ai-gateway` is an HTTP/WebSocket +host, and `litellm-python-bridge` exposes Rust to the Python SDK. The runtime +boundary is OCR-first; do not migrate other routes wholesale. ## Core Boundary @@ -61,8 +61,8 @@ Allowed in `core`: - The public entrypoint for a top-level LiteLLM call - Request/response transforms and stream chunk normalization - Provider resolution, auth header construction, and URL building -- The provider HTTP call itself, through a shared reused client with connect and - request timeouts +- The provider HTTP call itself for existing non-OCR core entrypoints, through a + shared reused client with connect and request timeouts - Shared data types and validation errors - Deterministic token/cost helper logic @@ -79,9 +79,14 @@ Env reads in `core` are limited to credential fallback inside a route's no key is passed. Everything else config-shaped is resolved by the host and passed in. -Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`) -predate this rule and are being moved into `core` route modules; do not add new -ones there, and prefer moving one when you touch it. +Audio transcription and realtime still have legacy gateway-owned call paths. +Do not broaden the OCR runtime boundary to those routes as part of OCR work. + +OCR is the exception to the older end-to-end core rule. Core owns supported +parameter metadata, parameter mapping, request transformation, and decoded +provider response normalization. Runtime owns provider/model resolution, +configuration, environment/auth, URL/headers, document materialization, HTTP, +Azure polling, and lifecycle ordering. Gateway OCR only adapts host hooks. Python owns rollout state and fallback while Rust is being introduced. Rust paths must be off by default until parity tests prove equivalence with Python. diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4388e561026..5fba60d08b4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1395,6 +1395,7 @@ dependencies = [ "futures-channel", "futures-util", "litellm-core", + "litellm-runtime", "pyo3", "reqwest", "serde", @@ -1432,6 +1433,7 @@ dependencies = [ "criterion", "litellm-ai-gateway", "litellm-core", + "litellm-runtime", "pyo3", "pyo3-async-runtimes", "pythonize", @@ -1440,6 +1442,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-runtime" +version = "0.1.0" +dependencies = [ + "base64", + "litellm-core", + "reqwest", + "serde_json", + "tokio", +] + [[package]] name = "litemap" version = "0.8.2" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 481ea3f8f66..3df581e8d3e 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "crates/core", + "crates/runtime", "crates/ai-gateway", "crates/python-bridge", ] @@ -14,6 +15,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } +litellm-runtime = { path = "crates/runtime" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" pyo3 = "0.29.0" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index bcccf93300b..dfacd9c31e6 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -24,26 +24,28 @@ coverage and production evidence. | Crate | Role | |-------|------| -| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | -| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-core | Provider transformations and shared types. OCR is deterministic transformation only. | +| litellm-runtime | Reusable OCR resolution, auth, preparation, lifecycle, HTTP, and polling. | +| litellm-ai-gateway | Axum and WebSocket host plus host-specific OCR logger/guardrail adapters. | +| litellm-python-bridge | PyO3 cdylib. Calls runtime directly for OCR and retains gateway for unrelated legacy paths. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction (acyclic): ai-gateway/python-bridge -> runtime -> core. ## Layout ```text crates/ - core/ The SDK: route modules + provider transforms. - src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client + core/ Shared types and provider transformations. + src/messages/ Legacy whole-call route implementation. src/providers/anthropic/messages/transformation.rs - ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. + runtime/ Reusable OCR execution. + ai-gateway/ Axum server + WebSocket hosts; adapts OCR runtime hooks. python-bridge/ PyO3 bridge for Python LiteLLM. ``` The folder shape follows the Python provider tree: `core/src/providers///transformation.rs`. The bridge exposes one -function per top-level route, mirroring the core entrypoints. +function per top-level route, calling the matching core or runtime entrypoint. ## Checks diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index a1860d8a9c9..0f4ba6bc620 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -9,14 +9,14 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` ## Transforms and the base config -3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src//transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`). +3. Every route defines a base transformation trait in `core/src//transformation.rs`. OCR uses `OcrProviderTransformation`; its runtime configuration is separate. 4. Each provider implements that trait as a `const __CONFIG` in `core/src/providers///transformation.rs`, mirroring the Python provider tree. 5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it. 6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers. ## Boundaries -7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request. +7. Layers never cross: for OCR, `core` owns provider transformations and `runtime` owns resolution, auth, provider HTTP, polling, and lifecycle ordering; `ai-gateway` owns HTTP/WS hosting and host-hook adapters; `python-bridge` is a thin PyO3 adapter. Existing non-OCR core entrypoints keep their current ownership until deliberately migrated. 8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers///`; a route is a module, never a new crate. 9. Route entry point stays thin: `core::::()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them. 10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`. diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md index 92567091cd3..e19c270faf3 100644 --- a/litellm-rust/crates/ai-gateway/AGENTS.md +++ b/litellm-rust/crates/ai-gateway/AGENTS.md @@ -1,9 +1,11 @@ # ai-gateway — folder architecture -The Axum server that fronts the Rust gateway. It owns transport + config + auth -only; deployment selection lives in `core::router`, and the LLM call itself -(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint -such as `litellm_core::messages::messages`. No provider handler lives here. +The Axum server that fronts the Rust gateway. It owns host transport + config + auth +only; deployment selection lives in `core::router`. Existing non-OCR calls live +behind core route entrypoints such as `litellm_core::messages::messages`, while +OCR executes through `litellm-runtime`. No provider handler lives here. + +OCR is executed by `litellm-runtime`. Gateway OCR is a compatibility wrapper that adapts custom logger, guardrail, and `RequestMetadata` types to runtime lifecycle hooks. Runtime must never depend on this crate. ``` src/ diff --git a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md index 733953bbdb3..f777c45b1e6 100644 --- a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md +++ b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md @@ -4,6 +4,10 @@ The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is a API callback: it POSTs each finished session to the LiteLLM proxy, which records spend and runs the usual callbacks. +OCR provider execution lives in `litellm-runtime`. The gateway OCR module is a +compatibility host adapter for custom logger, guardrail, and request metadata +types; runtime has no dependency on the gateway. + ```mermaid flowchart LR C[client] <--> G[Rust ai-gateway
LLM inference] diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 541beabe170..6687e8f1793 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -15,8 +15,8 @@ required-features = ["server"] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } -# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the -# Python proxy callbacks API. +litellm-runtime.workspace = true +# reqwest ships realtime logs to the Python proxy callbacks API. reqwest.workspace = true # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 7a6c620ee84..ac4918d5380 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,15 +6,16 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` is exactly three crates (a crate is a **layer**, not a route): +`litellm-rust` is exactly four crates (a crate is a **layer**, not a route): | Crate | Role | |-------|------| -| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | -| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-core | Provider transformations and shared types. OCR is transformation-only. | +| litellm-runtime | Reusable OCR preparation, lifecycle ordering, HTTP, and polling. | +| litellm-ai-gateway | Axum and WebSocket host plus OCR host-hook adapters. | +| litellm-python-bridge | PyO3 adapter calling runtime for OCR and retaining gateway for unrelated legacy paths. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction (acyclic): ai-gateway/python-bridge -> runtime -> core. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 057db6457c4..9f446072dfc 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -3,10 +3,8 @@ //! Two layers, split by feature so the Python `cdylib` can depend on the I/O //! without pulling in the HTTP server: //! -//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks, -//! and provider I/O. Always available — no feature required. These predate the -//! rule that a route's entrypoint and handler live in `litellm-core` (see -//! `litellm_core::messages`) and move there as they are touched. +//! - Call-type modules such as [`ocr`]: host-specific lifecycle adapters and +//! compatibility entrypoints. OCR provider execution lives in `litellm-runtime`. //! - [`io`]: compatibility exports and realtime WebSocket splice helpers. //! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling //! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway` diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 95df566dc53..a7dff87dca5 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -4,13 +4,9 @@ use std::pin::Pin; use litellm_core::CoreResult; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::CoreError; -use litellm_core::ocr::transformation::OcrAuthStrategy; +use litellm_runtime::ocr::{PreparedOcrRequest, ProviderOcrRequest}; use serde_json::{Map, Value, json}; -use super::common_utils::{ - convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, -}; -use super::types::{PreparedOcrRequest, ProviderOcrRequest}; use crate::integrations::custom_guardrail::{ CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, }; @@ -75,43 +71,19 @@ impl OcrLifecycleHooks { &self, request: PreparedOcrRequest, ) -> CoreResult { - let config = ocr_provider_config(&request.custom_llm_provider, &request.model) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; - let env_lookup = |key: &str| std::env::var(key).ok(); - let headers = string_headers(request.extra_headers)?; - let auth_strategy = config.auth_strategy(); - let api_key = (!has_header(&headers, auth_strategy.header_name())) - .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) - .transpose()?; - let url = config.complete_url( - request.api_base.as_deref(), - &request.model, - &request.optional_params, - &env_lookup, - )?; - let filtered_params = config.map_ocr_params(&request.optional_params); let model = request.model.clone(); let custom_llm_provider = request.custom_llm_provider.clone(); - let document = if config.requires_data_uri_document() { - convert_document_url_to_data_uri(request.document).await? - } else { - request.document - }; - let body = config - .transform_ocr_request(&request.model, document, filtered_params)? - .data; - let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); + let mut provider_request = litellm_runtime::ocr::prepare_provider_request(request).await?; let body = self - .run_during_call_guardrails(&model, &custom_llm_provider, &url, body) + .run_during_call_guardrails( + &model, + &custom_llm_provider, + &provider_request.url, + provider_request.body, + ) .await?; - Ok(ProviderOcrRequest { - model, - config, - url, - body, - upstream_headers, - timeout: request.timeout, - }) + provider_request.body = body; + Ok(provider_request) } async fn run_during_call_guardrails( @@ -249,21 +221,6 @@ impl CallLifecycleHooks for OcrLi } } -fn upstream_headers( - headers: &[(String, String)], - auth_strategy: OcrAuthStrategy, - api_key: Option<&str>, -) -> Vec<(String, String)> { - api_key - .map(|api_key| match auth_strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), - }) - .into_iter() - .chain(headers.iter().cloned()) - .collect() -} - fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { GuardrailContext { call_type: CallType::Ocr, diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index c4c13e2300c..3c826ca912d 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,23 +1,31 @@ use litellm_core::CoreResult; -use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; -mod common_utils; -mod handler; mod hooks; -mod prepare; mod types; pub use types::OcrRequest; -use handler::execute_ocr_provider_call; -use prepare::{PreparedOcrCall, prepare_ocr_call}; +use hooks::OcrLifecycleHooks; pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { - let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); - CallLifecycle::default() - .run_request(request, &hooks, execute_ocr_provider_call) - .await + let hooks = OcrLifecycleHooks::new( + crate::integrations::custom_logger::CustomLoggerRunner::new(request.callbacks), + crate::integrations::custom_guardrail::CustomGuardrailRunner::new(request.guardrails), + request.request_metadata, + ); + let request = litellm_runtime::ocr::prepare_ocr_request(litellm_runtime::ocr::OcrRequest { + model: request.model, + document: request.document, + api_key: request.api_key, + api_base: request.api_base, + custom_llm_provider: request.custom_llm_provider, + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + litellm_call_id: request.litellm_call_id, + }); + litellm_runtime::ocr::ocr_with_hooks(request, &hooks).await } #[cfg(test)] diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs deleted file mode 100644 index 6231393c889..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ /dev/null @@ -1,57 +0,0 @@ -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::{SystemTime, UNIX_EPOCH}; - -use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; - -use super::hooks::OcrLifecycleHooks; -use super::types::{OcrRequest, PreparedOcrRequest}; -use crate::integrations::custom_guardrail::CustomGuardrailRunner; -use crate::integrations::custom_logger::CustomLoggerRunner; - -pub(crate) struct PreparedOcrCall { - pub(crate) request: PreparedOcrRequest, - pub(crate) hooks: OcrLifecycleHooks, -} - -pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { - let call_id = request - .litellm_call_id - .map(str::to_string) - .unwrap_or_else(new_ocr_call_id); - let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) - .unwrap_or(CustomLlmProvider { - model: request.model, - custom_llm_provider: "mistral", - }); - let model = provider_info.model.to_string(); - let custom_llm_provider = provider_info.custom_llm_provider.to_string(); - - PreparedOcrCall { - request: PreparedOcrRequest { - model, - custom_llm_provider, - litellm_call_id: call_id, - document: request.document, - api_key: request.api_key.map(str::to_string), - api_base: request.api_base.map(str::to_string), - extra_headers: request.extra_headers, - optional_params: request.optional_params, - timeout: request.timeout, - }, - hooks: OcrLifecycleHooks::new( - CustomLoggerRunner::new(request.callbacks), - CustomGuardrailRunner::new(request.guardrails), - request.request_metadata, - ), - } -} - -fn new_ocr_call_id() -> String { - static COUNTER: AtomicU64 = AtomicU64::new(1); - let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_nanos()) - .unwrap_or(0); - format!("ocr-{timestamp}-{sequence}") -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index bb2a6b06501..667f4642269 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -2,12 +2,10 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use litellm_core::error::CoreError; -use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; use super::{OcrRequest, ocr}; use crate::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, @@ -198,87 +196,6 @@ impl CustomGuardrail for RecordingOcrGuardrail { } } -#[test] -fn truncate_error_body_passes_short_strings_through() { - let body = "Unauthorized"; - assert_eq!(truncate_error_body(body), "Unauthorized"); -} - -#[test] -fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(306); - let truncated = truncate_error_body(&body); - - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); -} - -#[test] -fn truncate_error_body_does_not_split_multibyte_chars() { - let body = "é".repeat(266); - let truncated = truncate_error_body(&body); - assert!(truncated.is_char_boundary(truncated.len())); -} - -#[test] -fn ocr_dispatch_supports_migrated_providers() { - assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); - assert!( - ocr_provider_config("azure_ai", "pixtral-12b-2409") - .expect("azure ai config resolves") - .requires_data_uri_document() - ); - assert_eq!( - ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") - .expect("document intelligence config resolves") - .response_handling(), - OcrResponseHandling::AzureDocumentIntelligencePoll - ); - assert!( - ocr_provider_config("vertex_ai", "deepseek-ocr-maas") - .expect("vertex deepseek config resolves") - .supported_ocr_params() - .contains(&"temperature") - ); - assert!(ocr_provider_config("openai", "gpt-4o").is_none()); -} - -#[test] -fn string_headers_accepts_string_values() { - let headers = json!({ - "x-trace-id": "trace-1" - }) - .as_object() - .unwrap() - .clone(); - - assert_eq!( - string_headers(Some(headers)).expect("string headers accepted"), - vec![("x-trace-id".to_string(), "trace-1".to_string())] - ); -} - -#[test] -fn auth_header_detection_is_case_insensitive() { - let headers = vec![ - ("x-trace-id".to_string(), "trace-1".to_string()), - ("authorization".to_string(), "Bearer sk-test".to_string()), - ]; - - assert!(has_header(&headers, "authorization")); - - let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; - assert!(has_header(&headers, "authorization")); - - let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; - assert!(!has_header(&headers, "authorization")); -} - #[tokio::test] async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { let listener = TcpListener::bind("127.0.0.1:0") @@ -594,21 +511,3 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() { "{poll_request}" ); } - -#[test] -fn string_headers_rejects_non_string_values() { - let headers = json!({ - "x-retry-count": 3 - }) - .as_object() - .unwrap() - .clone(); - - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - CoreError::InvalidRequest( - "OCR extra_headers.x-retry-count must be a string, got number".to_string() - ) - ); -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index bde734a4dd1..e96d2df1adb 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -1,8 +1,6 @@ use std::sync::Arc; use std::time::Duration; -use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; -use litellm_core::ocr::transformation::OcrProviderConfig; use serde_json::{Map, Value}; use crate::integrations::custom_guardrail::CustomGuardrail; @@ -23,35 +21,3 @@ pub struct OcrRequest<'a> { pub request_metadata: RequestMetadata, pub litellm_call_id: Option<&'a str>, } - -pub(crate) struct PreparedOcrRequest { - pub(crate) model: String, - pub(crate) custom_llm_provider: String, - pub(crate) litellm_call_id: String, - pub(crate) document: Value, - pub(crate) api_key: Option, - pub(crate) api_base: Option, - pub(crate) extra_headers: Option>, - pub(crate) optional_params: Map, - pub(crate) timeout: Option, -} - -impl CallLifecycleRequest for PreparedOcrRequest { - fn lifecycle_context(&self) -> CallLifecycleContext { - CallLifecycleContext::new( - "ocr", - self.model.clone(), - self.custom_llm_provider.clone(), - self.litellm_call_id.clone(), - ) - } -} - -pub(crate) struct ProviderOcrRequest { - pub(crate) model: String, - pub(crate) config: &'static dyn OcrProviderConfig, - pub(crate) url: String, - pub(crate) body: Value, - pub(crate) upstream_headers: Vec<(String, String)>, - pub(crate) timeout: Option, -} diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index aee8b4937ef..c7a3c419c4f 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -1,7 +1,9 @@ -litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src//` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back. +litellm-core owns provider transformations and shared types. Existing non-OCR top-level calls may still own their full call path. A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate. Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback dispatch. Env reads are limited to credential fallback in a route's `prepare.rs`. Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates. + +OCR is transformation-only in core: supported parameter metadata, parameter mapping, request body/data transformation, and decoded provider response normalization. OCR resolution, auth, environment, URL/headers, document materialization, HTTP, polling, and lifecycle ordering belong in `litellm-runtime`. diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index cb3e735e533..5e071186a36 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -4,34 +4,13 @@ use crate::CoreResult; use super::types::{OcrRequestData, OcrResponseData}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum OcrAuthStrategy { - Bearer, - Header(&'static str), -} - -impl OcrAuthStrategy { - pub fn header_name(self) -> &'static str { - match self { - Self::Bearer => "authorization", - Self::Header(header_name) => header_name, - } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum OcrResponseHandling { - Json, - AzureDocumentIntelligencePoll, -} - -pub trait OcrProviderConfig: Sync { - fn supported_ocr_params(&self) -> &'static [&'static str]; +pub trait OcrProviderTransformation: Sync { + fn get_supported_ocr_params(&self) -> &'static [&'static str]; fn map_ocr_params(&self, non_default_params: &Map) -> Map { let mut mapped_params = Map::new(); for (param, value) in non_default_params { - if self.supported_ocr_params().contains(¶m.as_str()) { + if self.get_supported_ocr_params().contains(¶m.as_str()) { mapped_params.insert(param.clone(), value.clone()); } } @@ -45,35 +24,9 @@ pub trait OcrProviderConfig: Sync { optional_params: Map, ) -> CoreResult; - fn transform_ocr_response( + fn transform_ocr_response_data( &self, model: &str, response_json: Value, ) -> CoreResult; - - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; - - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Bearer - } - - fn requires_data_uri_document(&self) -> bool { - false - } - - fn response_handling(&self) -> OcrResponseHandling { - OcrResponseHandling::Json - } } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index eabd15677cc..f51f36096c9 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -1,17 +1,10 @@ -use std::collections::BTreeSet; - use crate::error::{CoreError, CoreResult, json_type_name}; -use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; +use crate::ocr::transformation::OcrProviderTransformation; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; -const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; -const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; -const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; -const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; -const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96; const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages"]; @@ -23,199 +16,6 @@ pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig; pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig = AzureDocumentIntelligenceOcrConfig; -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -fn resolve_value( - explicit: Option<&str>, - env_name: &str, - env_lookup: &dyn Fn(&str) -> Option, - missing_message: &str, -) -> CoreResult { - non_empty(explicit) - .map(str::to_string) - .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| CoreError::Auth(missing_message.to_string())) -} - -pub fn resolve_azure_ai_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - resolve_value( - api_key, - AZURE_AI_API_KEY_ENV, - env_lookup, - "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params", - ) -} - -pub fn resolve_azure_ai_api_base( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - resolve_value( - api_base, - AZURE_AI_API_BASE_ENV, - env_lookup, - "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", - ) -} - -pub fn complete_azure_ai_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - let base = resolve_azure_ai_api_base(api_base, env_lookup)?; - Ok(format!( - "{}/providers/mistral/azure/ocr", - base.trim_end_matches('/') - )) -} - -pub fn resolve_document_intelligence_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - resolve_value( - api_key, - AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, - env_lookup, - "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter", - ) -} - -pub fn resolve_document_intelligence_endpoint( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - resolve_value( - api_base, - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, - env_lookup, - "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter", - ) -} - -fn encode_model_id(model: &str) -> String { - let model_id = model.rsplit('/').next().unwrap_or(model); - model_id - .bytes() - .flat_map(|byte| match byte { - b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { - vec![byte as char] - } - _ => format!("%{byte:02X}").chars().collect(), - }) - .collect() -} - -fn pages_token_is_valid(token: &str) -> bool { - let mut parts = token.split('-'); - let Some(start) = parts.next() else { - return false; - }; - if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) { - return false; - } - match parts.next() { - None => true, - Some(end) => { - !end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none() - } - } -} - -fn normalize_pages_param(pages: &Value) -> CoreResult> { - match pages { - Value::String(value) => { - let normalized = value - .split(',') - .map(str::trim) - .collect::>() - .join(","); - if normalized.split(',').all(pages_token_is_valid) { - Ok(Some(normalized)) - } else { - Err(CoreError::InvalidRequest(format!( - "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." - ))) - } - } - Value::Array(values) => { - if values.is_empty() { - return Ok(None); - } - if values.iter().all(Value::is_i64) { - let mut pages = BTreeSet::new(); - for value in values { - let page = value.as_i64().expect("checked is_i64"); - if page < 0 { - return Err(CoreError::InvalidRequest( - "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), - )); - } - pages.insert(page + 1); - } - return Ok(Some( - pages - .into_iter() - .map(|page| page.to_string()) - .collect::>() - .join(","), - )); - } - if values.iter().all(Value::is_string) { - let normalized = values - .iter() - .filter_map(Value::as_str) - .map(str::trim) - .collect::>() - .join(","); - if normalized.split(',').all(pages_token_is_valid) { - return Ok(Some(normalized)); - } - return Err(CoreError::InvalidRequest(format!( - "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." - ))); - } - Err(CoreError::InvalidRequest( - "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." - .to_string(), - )) - } - _ => Err(CoreError::InvalidRequest( - "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." - .to_string(), - )), - } -} - -pub fn complete_document_intelligence_url( - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; - let mut url = format!( - "{}/documentintelligence/documentModels/{}:analyze?api-version={}", - endpoint.trim_end_matches('/'), - encode_model_id(model), - AZURE_DOCUMENT_INTELLIGENCE_API_VERSION - ); - - if let Some(pages) = optional_params.get("pages") - && let Some(normalized) = normalize_pages_param(pages)? - { - url.push_str("&pages="); - url.push_str(&normalized); - } - - Ok(url) -} - fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { let object = document.as_object().ok_or_else(|| CoreError::InvalidType { expected: "object", @@ -280,9 +80,9 @@ fn page_dimensions(page: &Map) -> Value { }) } -impl OcrProviderConfig for AzureAiOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() +impl OcrProviderTransformation for AzureAiOcrConfig { + fn get_supported_ocr_params(&self) -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.get_supported_ocr_params() } fn transform_ocr_request( @@ -294,39 +94,17 @@ impl OcrProviderConfig for AzureAiOcrConfig { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } - fn transform_ocr_response( + fn transform_ocr_response_data( &self, model: &str, response_json: Value, ) -> CoreResult { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) - } - - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - _optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - complete_azure_ai_url(api_base, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_azure_ai_api_key(api_key, env_lookup) - } - - fn requires_data_uri_document(&self) -> bool { - true + MISTRAL_OCR_CONFIG.transform_ocr_response_data(model, response_json) } } -impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { +impl OcrProviderTransformation for AzureDocumentIntelligenceOcrConfig { + fn get_supported_ocr_params(&self) -> &'static [&'static str] { AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS } @@ -355,7 +133,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { }) } - fn transform_ocr_response( + fn transform_ocr_response_data( &self, model: &str, response_json: Value, @@ -407,32 +185,6 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { object: "ocr".to_string(), }) } - - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - complete_document_intelligence_url(api_base, model, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_document_intelligence_api_key(api_key, env_lookup) - } - - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") - } - - fn response_handling(&self) -> OcrResponseHandling { - OcrResponseHandling::AzureDocumentIntelligencePoll - } } #[cfg(test)] @@ -458,23 +210,6 @@ mod tests { ); } - #[test] - fn document_intelligence_url_normalizes_zero_based_pages() { - let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com/"), - "azure_ai/doc-intelligence/prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3" - ); - } - #[test] fn document_intelligence_request_uses_base64_source_for_data_uri() { let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG @@ -492,7 +227,7 @@ mod tests { #[test] fn document_intelligence_response_normalizes_pages() { let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( + .transform_ocr_response_data( "prebuilt-layout", json!({ "status": "succeeded", diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index dc720cc4244..370b1df5a84 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -1,5 +1,5 @@ use crate::error::{CoreError, CoreResult, json_type_name}; -use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::transformation::OcrProviderTransformation; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value}; @@ -19,58 +19,12 @@ const SUPPORTED_OCR_PARAMS: &[&str] = &[ "id", ]; -/// Default Mistral API base, used when the caller does not override `api_base`. -pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1"; - -/// Environment variable holding the Mistral API key. -pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; - -/// Error message raised when no Mistral API key can be resolved. -pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params"; - -/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`. -/// -/// Blank/whitespace `api_base` is treated as absent (guard at resolution time). -pub fn complete_url(api_base: Option<&str>) -> String { - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(MISTRAL_DEFAULT_API_BASE) - .trim_end_matches('/'); - - if base.ends_with("/v1") { - format!("{base}/ocr") - } else { - format!("{base}/v1/ocr") - } -} - -/// Resolve the Mistral API key from the explicit param or the environment. -/// -/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth` -/// when no usable key is available. -/// -/// Note: the env fallback only reads the process environment. Secret-manager -/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in -/// via `api_key`; this fallback is a last resort for direct/standalone use. -pub fn resolve_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - api_key - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) -} - pub struct MistralOcrConfig; pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig; -impl OcrProviderConfig for MistralOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { +impl OcrProviderTransformation for MistralOcrConfig { + fn get_supported_ocr_params(&self) -> &'static [&'static str] { SUPPORTED_OCR_PARAMS } @@ -100,7 +54,7 @@ impl OcrProviderConfig for MistralOcrConfig { }) } - fn transform_ocr_response( + fn transform_ocr_response_data( &self, model: &str, response_json: Value, @@ -133,28 +87,10 @@ impl OcrProviderConfig for MistralOcrConfig { object: "ocr".to_string(), }) } - - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - _optional_params: &Map, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - Ok(complete_url(api_base)) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_api_key(api_key, env_lookup) - } } -pub fn supported_ocr_params() -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() +pub fn get_supported_ocr_params() -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.get_supported_ocr_params() } pub fn map_ocr_params(non_default_params: &Map) -> Map { @@ -169,8 +105,11 @@ pub fn transform_ocr_request( MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } -pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) +pub fn transform_ocr_response_data( + model: &str, + response_json: Value, +) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response_data(model, response_json) } #[cfg(test)] @@ -181,7 +120,7 @@ mod tests { #[test] fn supported_params_match_python_mistral_ocr_config() { assert_eq!( - supported_ocr_params(), + get_supported_ocr_params(), &[ "pages", "include_image_base64", @@ -266,7 +205,7 @@ mod tests { "usage_info": {"pages_processed": 1} }); - let result = transform_ocr_response("mistral-ocr-latest", response) + let result = transform_ocr_response_data("mistral-ocr-latest", response) .expect("response should transform"); assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]); @@ -275,38 +214,4 @@ mod tests { assert_eq!(result.usage_info, Some(json!({"pages_processed": 1}))); assert_eq!(result.object, "ocr"); } - - #[test] - fn complete_url_defaults_and_dedupes_v1() { - assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr"); - assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr"); - assert_eq!( - complete_url(Some("https://proxy.internal")), - "https://proxy.internal/v1/ocr" - ); - assert_eq!( - complete_url(Some("https://proxy.internal/v1/")), - "https://proxy.internal/v1/ocr" - ); - } - - #[test] - fn resolve_api_key_prefers_param_then_env() { - let no_env = |_: &str| None; - assert_eq!( - resolve_api_key(Some("sk-param"), &no_env).unwrap(), - "sk-param" - ); - - let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string()); - assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env"); - // Blank param falls through to the environment. - assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env"); - } - - #[test] - fn resolve_api_key_errors_when_absent() { - let err = resolve_api_key(None, &|_| None).expect_err("missing key should error"); - assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string())); - } } diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index 6300149c237..f3e6a9f21a6 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,18 +1,10 @@ use crate::error::{CoreError, CoreResult, json_type_name}; -use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::transformation::OcrProviderTransformation; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; -const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; -const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; -const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; -const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; -const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; -const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; -const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; - #[rustfmt::skip] const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[ "stream", @@ -29,102 +21,6 @@ pub struct VertexAiDeepSeekOcrConfig; pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig; pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; -fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { - keys.iter() - .find_map(|key| params.get(*key).and_then(Value::as_str)) - .map(str::trim) - .filter(|value| !value.is_empty()) -} - -pub fn is_deepseek_model(model: &str) -> bool { - model.to_ascii_lowercase().contains("deepseek") -} - -pub fn resolve_vertex_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - api_key - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| { - CoreError::Auth( - "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" - .to_string(), - ) - }) -} - -fn vertex_project( - params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - string_param(params, &["vertex_project", "vertex_ai_project"]) - .map(str::to_string) - .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - CoreError::InvalidRequest( - "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" - .to_string(), - ) - }) -} - -fn vertex_location( - params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> String { - string_param(params, &["vertex_location", "vertex_ai_location"]) - .map(str::to_string) - .or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty())) - .or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty())) - .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) -} - -fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String { - api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")) - .trim_end_matches('/') - .to_string() -} - -pub fn complete_vertex_mistral_url( - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - let project = vertex_project(optional_params, env_lookup)?; - let location = vertex_location(optional_params, env_lookup); - let base = vertex_mistral_api_base(api_base, &location); - Ok(format!( - "{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict" - )) -} - -pub fn complete_vertex_deepseek_url( - api_base: Option<&str>, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - let project = vertex_project(optional_params, env_lookup)?; - let location = vertex_location(optional_params, env_lookup); - let base = api_base - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE) - .trim_end_matches('/'); - Ok(format!( - "{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions" - )) -} - fn document_content_item(document: &Value) -> CoreResult { let object = document.as_object().ok_or_else(|| CoreError::InvalidType { expected: "object", @@ -209,9 +105,9 @@ fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> V } } -impl OcrProviderConfig for VertexAiOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { - MISTRAL_OCR_CONFIG.supported_ocr_params() +impl OcrProviderTransformation for VertexAiOcrConfig { + fn get_supported_ocr_params(&self) -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.get_supported_ocr_params() } fn transform_ocr_request( @@ -223,39 +119,17 @@ impl OcrProviderConfig for VertexAiOcrConfig { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } - fn transform_ocr_response( + fn transform_ocr_response_data( &self, model: &str, response_json: Value, ) -> CoreResult { - MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) - } - - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_vertex_api_key(api_key, env_lookup) - } - - fn requires_data_uri_document(&self) -> bool { - true + MISTRAL_OCR_CONFIG.transform_ocr_response_data(model, response_json) } } -impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { - fn supported_ocr_params(&self) -> &'static [&'static str] { +impl OcrProviderTransformation for VertexAiDeepSeekOcrConfig { + fn get_supported_ocr_params(&self) -> &'static [&'static str] { DEEPSEEK_SUPPORTED_OCR_PARAMS } @@ -285,7 +159,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { }) } - fn transform_ocr_response( + fn transform_ocr_response_data( &self, model: &str, response_json: Value, @@ -339,46 +213,12 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { object: "ocr".to_string(), }) } - - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - complete_vertex_deepseek_url(api_base, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_vertex_api_key(api_key, env_lookup) - } } #[cfg(test)] mod tests { use super::*; - #[test] - fn vertex_mistral_url_uses_project_location_and_model() { - let params = Map::from_iter([ - ("vertex_project".to_string(), json!("proj-1")), - ("vertex_location".to_string(), json!("europe-west4")), - ]); - - let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) - .expect("url builds"); - - assert_eq!( - url, - "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - } - #[test] fn vertex_mistral_reuses_mistral_body_transform() { let body = VERTEX_AI_OCR_CONFIG @@ -416,7 +256,7 @@ mod tests { #[test] fn vertex_deepseek_response_wraps_markdown_content() { let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG - .transform_ocr_response( + .transform_ocr_response_data( "deepseek-ocr-maas", json!({ "choices": [{"message": {"content": "# OCR text"}}], diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index 656ba033b62..af029fda0e7 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,7 +1,7 @@ -//! Enforcement: the litellm-rust workspace has exactly three crates. +//! Enforcement: the litellm-rust workspace has exactly four crates. //! -//! `core` (pure translation), `ai-gateway` (routes + all network I/O), and -//! `python-bridge` (the PyO3 cdylib). Adding or removing a crate must be a +//! `core` (provider transformation), `runtime` (reusable call execution), +//! `ai-gateway` (host adapters), and `python-bridge` (the PyO3 cdylib). Adding or removing a crate must be a //! deliberate act: this test fails until the allowlist here is updated, forcing //! whoever changes the crate set to justify the new crate per the rule that a //! crate is a layer needing independent compilation / its own deps / a separate @@ -16,10 +16,15 @@ use std::path::{Path, PathBuf}; /// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the /// workspace legitimately gains or loses a crate. -const EXPECTED_MEMBERS: &[&str] = &["crates/core", "crates/ai-gateway", "crates/python-bridge"]; +const EXPECTED_MEMBERS: &[&str] = &[ + "crates/core", + "crates/runtime", + "crates/ai-gateway", + "crates/python-bridge", +]; /// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-bridge"]; +const EXPECTED_CRATE_DIRS: &[&str] = &["core", "runtime", "ai-gateway", "python-bridge"]; const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact)."; diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index ad3cddfa5fd..3c56e8a5c7d 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,3 @@ -litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over the litellm-core route entrypoints (e.g. `litellm_core::messages::messages`). +litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK. It calls `litellm-runtime` directly for OCR and core entrypoints for routes such as Messages. The ai-gateway dependency remains for unrelated audio and WebSocket paths. -Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. +Keep it thin: no business logic, no transforms, no I/O orchestration. Marshal in/out and call the runtime or core entrypoint. diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 0c4a753f762..d3743398a66 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -15,6 +15,7 @@ extension-module = ["pyo3/extension-module"] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } +litellm-runtime.workspace = true litellm-ai-gateway = { workspace = true, default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f9e75f45f75..a1753df5aa7 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -4,7 +4,6 @@ use std::time::Duration; use litellm_ai_gateway::io::audio_transcription::{ AudioTranscriptionRequest, audio_transcription as run_audio_transcription, }; -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; use litellm_core::chat_completions::{ @@ -13,6 +12,7 @@ use litellm_core::chat_completions::{ use litellm_core::error::CoreError; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_runtime::ocr::{OcrRequest, ocr as run_ocr}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; @@ -240,9 +240,6 @@ fn ocr( extra_headers, optional_params, timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), litellm_call_id: None, })) }); @@ -285,9 +282,6 @@ fn aocr( extra_headers, optional_params, timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), litellm_call_id: None, }) .await diff --git a/litellm-rust/crates/runtime/AGENTS.md b/litellm-rust/crates/runtime/AGENTS.md new file mode 100644 index 00000000000..52438f72550 --- /dev/null +++ b/litellm-rust/crates/runtime/AGENTS.md @@ -0,0 +1,5 @@ +`litellm-runtime` is an OCR-first reusable execution layer. It depends on core and must not depend on ai-gateway or python-bridge. + +Runtime OCR owns provider/model resolution, runtime config selection, environment and credential resolution, URL and headers, document materialization, HTTP, Azure polling, and lifecycle ordering. Provider-specific request and decoded-response transformations remain in core. Public request types and lifecycle seams must stay host-neutral so ai-gateway can adapt its logger, guardrail, and metadata types without creating a reverse dependency. + +Do not migrate audio, realtime, Messages, Chat, Responses, or generic lifecycle code into this crate as part of OCR work. diff --git a/litellm-rust/crates/runtime/Cargo.toml b/litellm-rust/crates/runtime/Cargo.toml new file mode 100644 index 00000000000..8dcd3e8a842 --- /dev/null +++ b/litellm-rust/crates/runtime/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "litellm-runtime" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-core.workspace = true +base64.workspace = true +reqwest.workspace = true +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/runtime/src/constants.rs b/litellm-rust/crates/runtime/src/constants.rs new file mode 100644 index 00000000000..aa195ad8279 --- /dev/null +++ b/litellm-rust/crates/runtime/src/constants.rs @@ -0,0 +1,22 @@ +pub(crate) const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600; +pub(crate) const HTTP_CONNECT_TIMEOUT_SECS: u64 = 10; +pub(crate) const ERROR_BODY_MAX_CHARS: usize = 256; +pub(crate) const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; +pub(crate) const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0; +pub(crate) const MAX_SAFE_FETCH_REDIRECTS: usize = 10; +pub(crate) const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1"; +pub(crate) const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; +pub(crate) const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; +pub(crate) const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; +pub(crate) const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = + "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; +pub(crate) const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = + "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; +pub(crate) const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; +pub(crate) const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; +pub(crate) const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; +pub(crate) const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; +pub(crate) const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; +pub(crate) const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; +pub(crate) const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; +pub(crate) const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; diff --git a/litellm-rust/crates/runtime/src/lib.rs b/litellm-rust/crates/runtime/src/lib.rs new file mode 100644 index 00000000000..288a983c73e --- /dev/null +++ b/litellm-rust/crates/runtime/src/lib.rs @@ -0,0 +1,2 @@ +mod constants; +pub mod ocr; diff --git a/litellm-rust/crates/runtime/src/ocr/client.rs b/litellm-rust/crates/runtime/src/ocr/client.rs new file mode 100644 index 00000000000..650d14c1b97 --- /dev/null +++ b/litellm-rust/crates/runtime/src/ocr/client.rs @@ -0,0 +1,15 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use crate::constants::{HTTP_CLIENT_TIMEOUT_SECS, HTTP_CONNECT_TIMEOUT_SECS}; + +pub(super) fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(HTTP_CONNECT_TIMEOUT_SECS)) + .timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS)) + .build() + .expect("failed to build reqwest client") + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/runtime/src/ocr/common_utils.rs similarity index 85% rename from litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs rename to litellm-rust/crates/runtime/src/ocr/common_utils.rs index 9bc2818b6e7..93886b06d4a 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/runtime/src/ocr/common_utils.rs @@ -5,25 +5,15 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use litellm_core::CoreResult; use litellm_core::error::CoreError; -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 super::client::http_client; +use crate::constants::{ + AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB, + ERROR_BODY_MAX_CHARS, HTTP_CLIENT_TIMEOUT_SECS, HTTP_CONNECT_TIMEOUT_SECS, + MAX_SAFE_FETCH_REDIRECTS, }; -use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; -use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; -use litellm_core::providers::vertex_ai::ocr::transformation::{ - VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, -}; - -use crate::client::http_client; - -const ERROR_BODY_MAX_CHARS: usize = 256; -const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; -const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0; -const MAX_SAFE_FETCH_REDIRECTS: usize = 10; pub(super) fn truncate_error_body(body: &str) -> String { if body.chars().count() <= ERROR_BODY_MAX_CHARS { @@ -33,27 +23,6 @@ pub(super) fn truncate_error_body(body: &str) -> String { format!("{truncated}... (truncated)") } -pub(super) fn ocr_provider_config( - provider: &str, - model: &str, -) -> Option<&'static dyn OcrProviderConfig> { - match provider { - "mistral" => Some(&MISTRAL_OCR_CONFIG), - "azure_ai" if is_azure_document_intelligence_model(model) => { - Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) - } - "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), - "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), - "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), - _ => None, - } -} - -fn is_azure_document_intelligence_model(model: &str) -> bool { - let model = model.to_ascii_lowercase(); - model.contains("doc-intelligence") || model.contains("documentintelligence") -} - pub(super) fn string_headers( extra_headers: Option>, ) -> CoreResult> { @@ -188,9 +157,14 @@ fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}"))) } -async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> { +async fn safe_get_document_url( + url: &str, + timeout: Option, +) -> CoreResult<(Url, reqwest::Response)> { let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(Duration::from_secs(HTTP_CONNECT_TIMEOUT_SECS)) + .timeout(timeout.unwrap_or(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS))) .build() .map_err(|err| CoreError::Network(err.to_string()))?; let mut current_url = Url::parse(url) @@ -255,7 +229,10 @@ async fn read_response_with_limit( Ok(bytes) } -pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult { +pub(super) async fn convert_document_url_to_data_uri( + document: Value, + timeout: Option, +) -> CoreResult { let Some((field, url)) = document_url_field(&document)? else { return Ok(document); }; @@ -263,7 +240,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes return Ok(document); } - let (final_url, response) = safe_get_document_url(url).await?; + let (final_url, response) = safe_get_document_url(url, timeout).await?; let status = response.status(); if !status.is_success() { let body = response.text().await.unwrap_or_default(); @@ -401,6 +378,48 @@ mod tests { use super::*; use serde_json::json; + #[test] + fn truncate_error_body_is_bounded_and_utf8_safe() { + assert_eq!(truncate_error_body("Unauthorized"), "Unauthorized"); + let truncated = truncate_error_body(&"é".repeat(306)); + assert!(truncated.ends_with("... (truncated)")); + assert_eq!( + truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(), + 256 + ); + assert!(truncated.is_char_boundary(truncated.len())); + } + + #[test] + fn string_headers_validate_values_and_match_case_insensitively() { + let headers = string_headers(Some( + json!({"Authorization": "Bearer sk-test"}) + .as_object() + .expect("object") + .clone(), + )) + .expect("string headers accepted"); + assert!(has_header(&headers, "authorization")); + + let error = string_headers(Some( + json!({"x-retry-count": 3}) + .as_object() + .expect("object") + .clone(), + )) + .expect_err("non-string header rejected"); + assert_eq!( + error, + CoreError::InvalidRequest( + "OCR extra_headers.x-retry-count must be a string, got number".to_string() + ) + ); + } + #[test] fn blocks_private_and_metadata_ips() { assert!(is_blocked_ip("127.0.0.1".parse().unwrap())); @@ -417,10 +436,13 @@ mod tests { #[tokio::test] async fn convert_document_url_rejects_loopback_fetch() { - let error = convert_document_url_to_data_uri(json!({ - "type": "image_url", - "image_url": "http://127.0.0.1/image.png" - })) + let error = convert_document_url_to_data_uri( + json!({ + "type": "image_url", + "image_url": "http://127.0.0.1/image.png" + }), + None, + ) .await .unwrap_err(); @@ -438,7 +460,7 @@ mod tests { "image_url": "data:image/png;base64,abcd" }); - let transformed = convert_document_url_to_data_uri(document.clone()) + let transformed = convert_document_url_to_data_uri(document.clone(), None) .await .unwrap(); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/runtime/src/ocr/handler.rs similarity index 85% rename from litellm-rust/crates/ai-gateway/src/ocr/handler.rs rename to litellm-rust/crates/runtime/src/ocr/handler.rs index 1de34eb400e..152ad6a5fa9 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/runtime/src/ocr/handler.rs @@ -1,13 +1,12 @@ use litellm_core::CoreResult; use litellm_core::error::CoreError; -use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; +use super::client::http_client; use super::common_utils::{poll_document_intelligence, truncate_error_body}; -use super::types::ProviderOcrRequest; -use crate::client::http_client; +use super::types::{OcrResponseHandling, ProviderOcrRequest}; -pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { +pub(super) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -45,7 +44,8 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co .await?; return Ok(request .config - .transform_ocr_response(&request.model, response_json)? + .transformation() + .transform_ocr_response_data(&request.model, response_json)? .into_json()); } @@ -66,6 +66,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co Ok(request .config - .transform_ocr_response(&request.model, response_json)? + .transformation() + .transform_ocr_response_data(&request.model, response_json)? .into_json()) } diff --git a/litellm-rust/crates/runtime/src/ocr/mod.rs b/litellm-rust/crates/runtime/src/ocr/mod.rs new file mode 100644 index 00000000000..b3b11d824e6 --- /dev/null +++ b/litellm-rust/crates/runtime/src/ocr/mod.rs @@ -0,0 +1,129 @@ +use std::future::{Future, ready}; +use std::pin::Pin; + +use litellm_core::call_lifecycle::{ + CallLifecycle, CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming, +}; +use litellm_core::{CoreError, CoreResult}; +use serde_json::Value; + +mod client; +mod common_utils; +mod handler; +mod prepare; +mod provider; +mod types; + +pub use prepare::{prepare_ocr_request, prepare_provider_request}; +pub use types::{OcrRequest, PreparedOcrRequest, ProviderOcrRequest}; + +use handler::execute_ocr_provider_call; + +pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { + ocr_with_hooks(prepare_ocr_request(request), &NoopOcrLifecycleHooks).await +} + +pub async fn ocr_with_hooks(request: PreparedOcrRequest, hooks: &Hooks) -> CoreResult +where + Hooks: CallLifecycleHooks, +{ + CallLifecycle::default() + .run_request(request, hooks, execute_ocr_provider_call) + .await +} + +pub struct NoopOcrLifecycleHooks; + +impl CallLifecycleHooks for NoopOcrLifecycleHooks { + type PreCallFuture<'a> = std::future::Ready>; + type DuringCallFuture<'a> = + Pin> + Send + 'a>>; + type SuccessFuture<'a> = std::future::Ready<()>; + type FailureFuture<'a> = std::future::Ready<()>; + + fn async_pre_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::PreCallFuture<'a> { + ready(Ok(request)) + } + + fn async_during_call_hook<'a>( + &'a self, + _context: &'a CallLifecycleContext, + request: PreparedOcrRequest, + ) -> Self::DuringCallFuture<'a> { + Box::pin(prepare_provider_request(request)) + } + + fn async_log_success_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a Value, + _timing: &'a CallLifecycleTiming, + ) -> Self::SuccessFuture<'a> { + ready(()) + } + + fn async_log_failure_event<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a CoreError, + _timing: &'a CallLifecycleTiming, + ) -> Self::FailureFuture<'a> { + ready(()) + } +} + +#[cfg(test)] +mod tests { + use serde_json::{Map, json}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + use super::*; + + #[tokio::test] + async fn direct_runtime_ocr_uses_noop_lifecycle() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = vec![0_u8; 4096]; + let _ = socket.read(&mut request).await.unwrap(); + let body = json!({ + "pages": [{"index": 0, "markdown": "runtime"}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1} + }) + .to_string(); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + socket.write_all(response.as_bytes()).await.unwrap(); + }); + + let api_base = format!("http://{address}"); + let response = ocr(OcrRequest { + model: "mistral/mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/document.pdf" + }), + api_key: Some("test-key"), + api_base: Some(&api_base), + custom_llm_provider: None, + extra_headers: None, + optional_params: Map::new(), + timeout: None, + litellm_call_id: Some("runtime-test"), + }) + .await + .unwrap(); + + server.await.unwrap(); + assert_eq!(response["pages"][0]["markdown"], "runtime"); + } +} diff --git a/litellm-rust/crates/runtime/src/ocr/prepare.rs b/litellm-rust/crates/runtime/src/ocr/prepare.rs new file mode 100644 index 00000000000..c7ef3246538 --- /dev/null +++ b/litellm-rust/crates/runtime/src/ocr/prepare.rs @@ -0,0 +1,104 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_core::CoreResult; +use litellm_core::error::CoreError; +use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + +use super::common_utils::{convert_document_url_to_data_uri, has_header, string_headers}; +use super::provider::ocr_provider_config; +use super::types::{OcrAuthStrategy, OcrRequest, PreparedOcrRequest, ProviderOcrRequest}; + +pub fn prepare_ocr_request(request: OcrRequest<'_>) -> PreparedOcrRequest { + let litellm_call_id = request + .litellm_call_id + .map(str::to_string) + .unwrap_or_else(new_ocr_call_id); + let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) + .unwrap_or(CustomLlmProvider { + model: request.model, + custom_llm_provider: "mistral", + }); + + PreparedOcrRequest { + model: provider_info.model.to_string(), + custom_llm_provider: provider_info.custom_llm_provider.to_string(), + litellm_call_id, + document: request.document, + api_key: request.api_key.map(str::to_string), + api_base: request.api_base.map(str::to_string), + extra_headers: request.extra_headers, + optional_params: request.optional_params, + timeout: request.timeout, + } +} + +pub async fn prepare_provider_request( + request: PreparedOcrRequest, +) -> CoreResult { + let config = ocr_provider_config(&request.custom_llm_provider, &request.model) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); + let headers = string_headers(request.extra_headers)?; + let auth_strategy = config.auth_strategy(); + let api_key = (!has_header(&headers, auth_strategy.header_name())) + .then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup)) + .transpose()?; + let url = config.complete_url( + request.api_base.as_deref(), + &request.model, + &request.optional_params, + &env_lookup, + )?; + let transformation = config.transformation(); + let supported_params = transformation.get_supported_ocr_params(); + let non_default_params = request + .optional_params + .iter() + .filter(|(param, _)| supported_params.contains(¶m.as_str())) + .map(|(param, value)| (param.clone(), value.clone())) + .collect(); + let filtered_params = transformation.map_ocr_params(&non_default_params); + let document = if config.requires_data_uri_document() { + convert_document_url_to_data_uri(request.document, request.timeout).await? + } else { + request.document + }; + let body = transformation + .transform_ocr_request(&request.model, document, filtered_params)? + .data; + + Ok(ProviderOcrRequest { + model: request.model, + config, + url, + body, + upstream_headers: upstream_headers(&headers, auth_strategy, api_key.as_deref()), + timeout: request.timeout, + }) +} + +fn upstream_headers( + headers: &[(String, String)], + auth_strategy: OcrAuthStrategy, + api_key: Option<&str>, +) -> Vec<(String, String)> { + api_key + .map(|api_key| match auth_strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), + }) + .into_iter() + .chain(headers.iter().cloned()) + .collect() +} + +fn new_ocr_call_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(1); + let sequence = COUNTER.fetch_add(1, Ordering::Relaxed); + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + format!("ocr-{timestamp}-{sequence}") +} diff --git a/litellm-rust/crates/runtime/src/ocr/provider.rs b/litellm-rust/crates/runtime/src/ocr/provider.rs new file mode 100644 index 00000000000..f736b8a01c7 --- /dev/null +++ b/litellm-rust/crates/runtime/src/ocr/provider.rs @@ -0,0 +1,500 @@ +use litellm_core::CoreResult; +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrProviderTransformation; +use litellm_core::providers::azure_ai::ocr::transformation::{ + AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG, +}; +use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use litellm_core::providers::vertex_ai::ocr::transformation::{ + VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, +}; +use serde_json::{Map, Value}; + +use super::types::{OcrAuthStrategy, OcrResponseHandling, OcrRuntimeConfig}; +use crate::constants::{ + AZURE_AI_API_BASE_ENV, AZURE_AI_API_KEY_ENV, AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, + AZURE_DOCUMENT_INTELLIGENCE_API_VERSION, AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, + MISTRAL_API_KEY_ENV, MISTRAL_DEFAULT_API_BASE, VERTEX_AI_API_KEY_ENV, + VERTEX_DEFAULT_DEEPSEEK_API_BASE, VERTEX_DEFAULT_LOCATION, VERTEX_LOCATION_ENV, + VERTEXAI_API_KEY_ENV, VERTEXAI_LOCATION_ENV, VERTEXAI_PROJECT_ENV, +}; + +struct MistralRuntimeConfig; +struct AzureAiRuntimeConfig; +struct AzureDocumentIntelligenceRuntimeConfig; +struct VertexAiRuntimeConfig; +struct VertexAiDeepSeekRuntimeConfig; + +static MISTRAL_RUNTIME_CONFIG: MistralRuntimeConfig = MistralRuntimeConfig; +static AZURE_AI_RUNTIME_CONFIG: AzureAiRuntimeConfig = AzureAiRuntimeConfig; +static AZURE_DOCUMENT_INTELLIGENCE_RUNTIME_CONFIG: AzureDocumentIntelligenceRuntimeConfig = + AzureDocumentIntelligenceRuntimeConfig; +static VERTEX_AI_RUNTIME_CONFIG: VertexAiRuntimeConfig = VertexAiRuntimeConfig; +static VERTEX_AI_DEEPSEEK_RUNTIME_CONFIG: VertexAiDeepSeekRuntimeConfig = + VertexAiDeepSeekRuntimeConfig; + +pub(super) fn ocr_provider_config( + provider: &str, + model: &str, +) -> Option<&'static dyn OcrRuntimeConfig> { + match provider { + "mistral" => Some(&MISTRAL_RUNTIME_CONFIG), + "azure_ai" if is_azure_document_intelligence_model(model) => { + Some(&AZURE_DOCUMENT_INTELLIGENCE_RUNTIME_CONFIG) + } + "azure_ai" => Some(&AZURE_AI_RUNTIME_CONFIG), + "vertex_ai" if model.to_ascii_lowercase().contains("deepseek") => { + Some(&VERTEX_AI_DEEPSEEK_RUNTIME_CONFIG) + } + "vertex_ai" => Some(&VERTEX_AI_RUNTIME_CONFIG), + _ => None, + } +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +fn resolve_value( + explicit: Option<&str>, + env_name: &str, + env_lookup: &dyn Fn(&str) -> Option, + missing_message: &str, +) -> CoreResult { + non_empty(explicit) + .map(str::to_string) + .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| CoreError::Auth(missing_message.to_string())) +} + +fn resolve_vertex_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + non_empty(api_key) + .map(str::to_string) + .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| { + CoreError::Auth( + "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" + .to_string(), + ) + }) +} + +fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { + keys.iter() + .find_map(|key| params.get(*key).and_then(Value::as_str)) + .and_then(|value| non_empty(Some(value))) +} + +fn vertex_project( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + string_param(params, &["vertex_project", "vertex_ai_project"]) + .map(str::to_string) + .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + CoreError::InvalidRequest( + "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" + .to_string(), + ) + }) +} + +fn vertex_location( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + string_param(params, &["vertex_location", "vertex_ai_location"]) + .map(str::to_string) + .or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty())) + .or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty())) + .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) +} + +fn encode_model_id(model: &str) -> String { + model + .rsplit('/') + .next() + .unwrap_or(model) + .bytes() + .flat_map(|byte| match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { + vec![byte as char] + } + _ => format!("%{byte:02X}").chars().collect(), + }) + .collect() +} + +fn pages_token_is_valid(token: &str) -> bool { + let mut parts = token.split('-'); + let Some(start) = parts.next() else { + return false; + }; + if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) { + return false; + } + match parts.next() { + None => true, + Some(end) => { + !end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none() + } + } +} + +fn normalize_pages_param(pages: &Value) -> CoreResult> { + match pages { + Value::String(value) => { + let normalized = value + .split(',') + .map(str::trim) + .collect::>() + .join(","); + if normalized.split(',').all(pages_token_is_valid) { + Ok(Some(normalized)) + } else { + Err(CoreError::InvalidRequest(format!( + "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." + ))) + } + } + Value::Array(values) if values.is_empty() => Ok(None), + Value::Array(values) if values.iter().all(Value::is_i64) => { + let pages = values + .iter() + .map(|value| value.as_i64().expect("checked is_i64")) + .map(|page| { + (page >= 0).then_some(page + 1).ok_or_else(|| { + CoreError::InvalidRequest( + "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), + ) + }) + }) + .collect::>>()?; + Ok(Some( + pages + .into_iter() + .map(|page| page.to_string()) + .collect::>() + .join(","), + )) + } + Value::Array(values) if values.iter().all(Value::is_string) => { + let normalized = values + .iter() + .filter_map(Value::as_str) + .map(str::trim) + .collect::>() + .join(","); + if normalized.split(',').all(pages_token_is_valid) { + Ok(Some(normalized)) + } else { + Err(CoreError::InvalidRequest(format!( + "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." + ))) + } + } + _ => Err(CoreError::InvalidRequest( + "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." + .to_string(), + )), + } +} + +fn is_azure_document_intelligence_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.contains("doc-intelligence") || model.contains("documentintelligence") +} + +impl OcrRuntimeConfig for MistralRuntimeConfig { + fn transformation(&self) -> &'static dyn OcrProviderTransformation { + &MISTRAL_OCR_CONFIG + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let base = non_empty(api_base) + .unwrap_or(MISTRAL_DEFAULT_API_BASE) + .trim_end_matches('/'); + Ok(if base.ends_with("/v1") { + format!("{base}/ocr") + } else { + format!("{base}/v1/ocr") + }) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + non_empty(api_key) + .map(str::to_string) + .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| CoreError::Auth("Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params".to_string())) + } +} + +impl OcrRuntimeConfig for AzureAiRuntimeConfig { + fn transformation(&self) -> &'static dyn OcrProviderTransformation { + &AZURE_AI_OCR_CONFIG + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let base = resolve_value( + api_base, + AZURE_AI_API_BASE_ENV, + env_lookup, + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", + )?; + Ok(format!( + "{}/providers/mistral/azure/ocr", + base.trim_end_matches('/') + )) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_value( + api_key, + AZURE_AI_API_KEY_ENV, + env_lookup, + "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params", + ) + } + + fn requires_data_uri_document(&self) -> bool { + true + } +} + +impl OcrRuntimeConfig for AzureDocumentIntelligenceRuntimeConfig { + fn transformation(&self) -> &'static dyn OcrProviderTransformation { + &AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let endpoint = resolve_value( + api_base, + AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, + env_lookup, + "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter", + )?; + let mut url = format!( + "{}/documentintelligence/documentModels/{}:analyze?api-version={}", + endpoint.trim_end_matches('/'), + encode_model_id(model), + AZURE_DOCUMENT_INTELLIGENCE_API_VERSION + ); + if let Some(pages) = optional_params.get("pages") + && let Some(normalized) = normalize_pages_param(pages)? + { + url.push_str("&pages="); + url.push_str(&normalized); + } + Ok(url) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_value( + api_key, + AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, + env_lookup, + "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter", + ) + } + + fn auth_strategy(&self) -> OcrAuthStrategy { + OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") + } + + fn response_handling(&self) -> OcrResponseHandling { + OcrResponseHandling::AzureDocumentIntelligencePoll + } +} + +impl OcrRuntimeConfig for VertexAiRuntimeConfig { + fn transformation(&self) -> &'static dyn OcrProviderTransformation { + &VERTEX_AI_OCR_CONFIG + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let project = vertex_project(optional_params, env_lookup)?; + let location = vertex_location(optional_params, env_lookup); + let base = non_empty(api_base) + .map(str::to_string) + .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")); + Ok(format!( + "{}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict", + base.trim_end_matches('/') + )) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_vertex_api_key(api_key, env_lookup) + } + + fn requires_data_uri_document(&self) -> bool { + true + } +} + +impl OcrRuntimeConfig for VertexAiDeepSeekRuntimeConfig { + fn transformation(&self) -> &'static dyn OcrProviderTransformation { + &VERTEX_AI_DEEPSEEK_OCR_CONFIG + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + let project = vertex_project(optional_params, env_lookup)?; + let location = vertex_location(optional_params, env_lookup); + let base = non_empty(api_base).unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE); + Ok(format!( + "{}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions", + base.trim_end_matches('/') + )) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_vertex_api_key(api_key, env_lookup) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dispatch_supports_ocr_providers() { + assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); + assert!( + ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document() + ); + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") + .expect("document intelligence config resolves") + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert!( + ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .transformation() + .get_supported_ocr_params() + .contains(&"temperature") + ); + assert!(ocr_provider_config("openai", "gpt-4o").is_none()); + } + + #[test] + fn mistral_url_and_credentials_resolve_in_runtime() { + assert_eq!( + MISTRAL_RUNTIME_CONFIG + .complete_url(None, "mistral-ocr-latest", &Map::new(), &|_| None) + .expect("url builds"), + "https://api.mistral.ai/v1/ocr" + ); + assert_eq!( + MISTRAL_RUNTIME_CONFIG + .complete_url( + Some("https://proxy.internal/v1/"), + "mistral-ocr-latest", + &Map::new(), + &|_| None, + ) + .expect("url builds"), + "https://proxy.internal/v1/ocr" + ); + let env_lookup = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string()); + assert_eq!( + MISTRAL_RUNTIME_CONFIG + .resolve_api_key(Some(" "), &env_lookup) + .expect("env key resolves"), + "sk-env" + ); + } + + #[test] + fn document_intelligence_url_normalizes_zero_based_pages() { + let params = Map::from_iter([("pages".to_string(), serde_json::json!([2, 0, 2]))]); + let url = AZURE_DOCUMENT_INTELLIGENCE_RUNTIME_CONFIG + .complete_url( + Some("https://example.cognitiveservices.azure.com/"), + "azure_ai/doc-intelligence/prebuilt-layout", + ¶ms, + &|_| None, + ) + .expect("url builds"); + assert_eq!( + url, + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3" + ); + } + + #[test] + fn vertex_url_uses_project_location_and_model() { + let params = Map::from_iter([ + ("vertex_project".to_string(), serde_json::json!("proj-1")), + ( + "vertex_location".to_string(), + serde_json::json!("europe-west4"), + ), + ]); + let url = VERTEX_AI_RUNTIME_CONFIG + .complete_url(None, "mistral-ocr-maas", ¶ms, &|_| None) + .expect("url builds"); + assert_eq!( + url, + "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + } +} diff --git a/litellm-rust/crates/runtime/src/ocr/types.rs b/litellm-rust/crates/runtime/src/ocr/types.rs new file mode 100644 index 00000000000..46821c66fad --- /dev/null +++ b/litellm-rust/crates/runtime/src/ocr/types.rs @@ -0,0 +1,100 @@ +use std::time::Duration; + +use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use litellm_core::ocr::transformation::OcrProviderTransformation; +use serde_json::{Map, Value}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrAuthStrategy { + Bearer, + Header(&'static str), +} + +impl OcrAuthStrategy { + pub fn header_name(self) -> &'static str { + match self { + Self::Bearer => "authorization", + Self::Header(header_name) => header_name, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrResponseHandling { + Json, + AzureDocumentIntelligencePoll, +} + +pub(crate) trait OcrRuntimeConfig: Sync { + fn transformation(&self) -> &'static dyn OcrProviderTransformation; + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> litellm_core::CoreResult; + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> litellm_core::CoreResult; + + fn auth_strategy(&self) -> OcrAuthStrategy { + OcrAuthStrategy::Bearer + } + + fn requires_data_uri_document(&self) -> bool { + false + } + + fn response_handling(&self) -> OcrResponseHandling { + OcrResponseHandling::Json + } +} + +pub struct OcrRequest<'a> { + pub model: &'a str, + pub document: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: Option<&'a str>, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, + pub litellm_call_id: Option<&'a str>, +} + +pub struct PreparedOcrRequest { + pub model: String, + pub custom_llm_provider: String, + pub litellm_call_id: String, + pub document: Value, + pub api_key: Option, + pub api_base: Option, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, +} + +impl CallLifecycleRequest for PreparedOcrRequest { + fn lifecycle_context(&self) -> CallLifecycleContext { + CallLifecycleContext::new( + "ocr", + self.model.clone(), + self.custom_llm_provider.clone(), + self.litellm_call_id.clone(), + ) + } +} + +pub struct ProviderOcrRequest { + pub model: String, + pub(crate) config: &'static dyn OcrRuntimeConfig, + pub url: String, + pub body: Value, + pub upstream_headers: Vec<(String, String)>, + pub timeout: Option, +} diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index d1c77186ea8..acf023319cc 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -251,6 +251,13 @@ class BaseOCRConfig: """ raise NotImplementedError("transform_ocr_response must be implemented by provider") + def transform_ocr_response_data( + self, + model: str, + response_data: Mapping[str, object], + ) -> OCRResponse: + raise NotImplementedError("transform_ocr_response_data must be implemented by provider") + async def async_transform_ocr_response( self, model: str, diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index 78c8dd11171..7a835bffffc 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -2,9 +2,11 @@ Mistral OCR transformation implementation. """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.llms.base_llm.ocr.transformation import ( @@ -19,6 +21,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj MISTRAL_OCR_API_KEY_ENV_VAR: Final = "MISTRAL_API_KEY" +_OCR_RESPONSE_DATA: Final = TypeAdapter(dict[str, object]) class MistralOCRConfig(BaseOCRConfig): @@ -227,18 +230,26 @@ class MistralOCRConfig(BaseOCRConfig): } """ try: - response_json: Final = raw_response.json() + response_data: Final = _OCR_RESPONSE_DATA.validate_python(raw_response.json()) - verbose_logger.debug("Mistral OCR response keys: %s", response_json.keys()) + verbose_logger.debug("Mistral OCR response keys: %s", response_data.keys()) - # Return native Mistral format - no transformation - return OCRResponse( - pages=response_json.get("pages", []), - model=response_json.get("model", model), - document_annotation=response_json.get("document_annotation"), - usage_info=response_json.get("usage_info"), - object="ocr", - ) + return self.transform_ocr_response_data(model=model, response_data=response_data) except Exception as e: verbose_logger.error("Error parsing Mistral OCR response: %s", e) raise e + + def transform_ocr_response_data( + self, + model: str, + response_data: Mapping[str, object], + ) -> OCRResponse: + return OCRResponse.model_validate( + { + "pages": response_data.get("pages", []), + "model": response_data.get("model", model), + "document_annotation": response_data.get("document_annotation"), + "usage_info": response_data.get("usage_info"), + "object": "ocr", + } + ) diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py index 8ce8b777f71..ec4f3a63a50 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py @@ -234,3 +234,27 @@ class TestTransformOcrResponseOcr4Fields: dumped_page = result.model_dump()["pages"][0] for field in ("tables", "hyperlinks", "header", "footer"): assert dumped_page[field] == page[field] + + def test_decoded_response_data_uses_the_same_normalization(self, config: MistralOCRConfig) -> None: + page = { + "index": 0, + "markdown": "decoded response", + "blocks": [{"type": "text", "content": "decoded response"}], + } + + response_data = { + "pages": [page], + "model": "mistral-ocr-4-0", + "usage_info": {"pages_processed": 1}, + } + result = config.transform_ocr_response_data( + model="mistral-ocr-4-0", + response_data=response_data, + ) + wrapped_result = config.transform_ocr_response( + model="mistral-ocr-4-0", + raw_response=httpx.Response(200, json=response_data), + logging_obj=None, + ) + + assert result == wrapped_result