mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
refactor(rust): introduce OCR runtime layer
This commit is contained in:
Yujong Lee
parent
4ba8517134
commit
0d16842e67
39 changed files with 1166 additions and 975 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
13
litellm-rust/Cargo.lock
generated
13
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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/<provider>/<route>/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
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
|
||||
3. Every route defines a base transformation trait in `core/src/<route>/transformation.rs`. OCR uses `OcrProviderTransformation`; its runtime configuration is separate.
|
||||
4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/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/<provider>/<route>/`; a route is a module, never a new crate.
|
||||
9. Route entry point stays thin: `core::<route>::<route>()` -> `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`.
|
||||
|
|
|
|||
|
|
@ -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/
|
||||
|
|
|
|||
|
|
@ -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<br/>LLM inference]
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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://<host>/v1/realtime?model=<model>` (WebSocket)
|
||||
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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<ProviderOcrRequest> {
|
||||
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<PreparedOcrRequest, ProviderOcrRequest, Value> 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,
|
||||
|
|
|
|||
|
|
@ -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<Value> {
|
||||
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)]
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
}
|
||||
|
|
@ -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()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub(crate) api_base: Option<String>,
|
||||
pub(crate) extra_headers: Option<Map<String, Value>>,
|
||||
pub(crate) optional_params: Map<String, Value>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
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<Duration>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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/<route>/` 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`.
|
||||
|
|
|
|||
|
|
@ -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<String, Value>) -> Map<String, Value> {
|
||||
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<String, Value>,
|
||||
) -> CoreResult<OcrRequestData>;
|
||||
|
||||
fn transform_ocr_response(
|
||||
fn transform_ocr_response_data(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData>;
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn auth_strategy(&self) -> OcrAuthStrategy {
|
||||
OcrAuthStrategy::Bearer
|
||||
}
|
||||
|
||||
fn requires_data_uri_document(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn response_handling(&self) -> OcrResponseHandling {
|
||||
OcrResponseHandling::Json
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
missing_message: &str,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<Option<String>> {
|
||||
match pages {
|
||||
Value::String(value) => {
|
||||
let normalized = value
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.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::<Vec<_>>()
|
||||
.join(","),
|
||||
));
|
||||
}
|
||||
if values.iter().all(Value::is_string) {
|
||||
let normalized = values
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>) -> 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<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
complete_azure_ai_url(api_base, env_lookup)
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
Ok(complete_url(api_base))
|
||||
}
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>) -> Map<String, Value> {
|
||||
|
|
@ -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<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
pub fn transform_ocr_response_data(
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
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()));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String, Value>, 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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> 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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<Value> {
|
||||
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<Value>, 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<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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"}}],
|
||||
|
|
|
|||
|
|
@ -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).";
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
5
litellm-rust/crates/runtime/AGENTS.md
Normal file
5
litellm-rust/crates/runtime/AGENTS.md
Normal file
|
|
@ -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.
|
||||
13
litellm-rust/crates/runtime/Cargo.toml
Normal file
13
litellm-rust/crates/runtime/Cargo.toml
Normal file
|
|
@ -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
|
||||
22
litellm-rust/crates/runtime/src/constants.rs
Normal file
22
litellm-rust/crates/runtime/src/constants.rs
Normal file
|
|
@ -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";
|
||||
2
litellm-rust/crates/runtime/src/lib.rs
Normal file
2
litellm-rust/crates/runtime/src/lib.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
mod constants;
|
||||
pub mod ocr;
|
||||
15
litellm-rust/crates/runtime/src/ocr/client.rs
Normal file
15
litellm-rust/crates/runtime/src/ocr/client.rs
Normal file
|
|
@ -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<reqwest::Client> = 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")
|
||||
})
|
||||
}
|
||||
|
|
@ -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<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
|
|
@ -188,9 +157,14 @@ fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult<Url>
|
|||
.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<Duration>,
|
||||
) -> 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<Value> {
|
||||
pub(super) async fn convert_document_url_to_data_uri(
|
||||
document: Value,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
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();
|
||||
|
||||
|
|
@ -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<Value> {
|
||||
pub(super) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
|
||||
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())
|
||||
}
|
||||
129
litellm-rust/crates/runtime/src/ocr/mod.rs
Normal file
129
litellm-rust/crates/runtime/src/ocr/mod.rs
Normal file
|
|
@ -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<Value> {
|
||||
ocr_with_hooks(prepare_ocr_request(request), &NoopOcrLifecycleHooks).await
|
||||
}
|
||||
|
||||
pub async fn ocr_with_hooks<Hooks>(request: PreparedOcrRequest, hooks: &Hooks) -> CoreResult<Value>
|
||||
where
|
||||
Hooks: CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value>,
|
||||
{
|
||||
CallLifecycle::default()
|
||||
.run_request(request, hooks, execute_ocr_provider_call)
|
||||
.await
|
||||
}
|
||||
|
||||
pub struct NoopOcrLifecycleHooks;
|
||||
|
||||
impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for NoopOcrLifecycleHooks {
|
||||
type PreCallFuture<'a> = std::future::Ready<CoreResult<PreparedOcrRequest>>;
|
||||
type DuringCallFuture<'a> =
|
||||
Pin<Box<dyn Future<Output = CoreResult<ProviderOcrRequest>> + 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");
|
||||
}
|
||||
}
|
||||
104
litellm-rust/crates/runtime/src/ocr/prepare.rs
Normal file
104
litellm-rust/crates/runtime/src/ocr/prepare.rs
Normal file
|
|
@ -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<ProviderOcrRequest> {
|
||||
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}")
|
||||
}
|
||||
500
litellm-rust/crates/runtime/src/ocr/provider.rs
Normal file
500
litellm-rust/crates/runtime/src/ocr/provider.rs
Normal file
|
|
@ -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<String>,
|
||||
missing_message: &str,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>, 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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> 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<Option<String>> {
|
||||
match pages {
|
||||
Value::String(value) => {
|
||||
let normalized = value
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.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::<CoreResult<std::collections::BTreeSet<_>>>()?;
|
||||
Ok(Some(
|
||||
pages
|
||||
.into_iter()
|
||||
.map(|page| page.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(","),
|
||||
))
|
||||
}
|
||||
Value::Array(values) if values.iter().all(Value::is_string) => {
|
||||
let normalized = values
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::trim)
|
||||
.collect::<Vec<_>>()
|
||||
.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<String, Value>,
|
||||
_env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
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<String>,
|
||||
) -> CoreResult<String> {
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
100
litellm-rust/crates/runtime/src/ocr/types.rs
Normal file
100
litellm-rust/crates/runtime/src/ocr/types.rs
Normal file
|
|
@ -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<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> litellm_core::CoreResult<String>;
|
||||
|
||||
fn resolve_api_key(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> litellm_core::CoreResult<String>;
|
||||
|
||||
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<Map<String, Value>>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
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<String>,
|
||||
pub api_base: Option<String>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
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<Duration>,
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue