diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml new file mode 100644 index 00000000000..c085e0a29b3 --- /dev/null +++ b/.github/workflows/test-rust.yml @@ -0,0 +1,54 @@ +name: LiteLLM Rust + +on: + push: + paths: + - "litellm-rust/**" + - ".github/workflows/test-rust.yml" + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_branch + - "litellm_**" + paths: + - "litellm-rust/**" + - ".github/workflows/test-rust.yml" + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + rust-checks: + name: rustfmt, clippy, test + runs-on: ubuntu-latest + timeout-minutes: 10 + defaults: + run: + working-directory: litellm-rust + env: + CARGO_TERM_COLOR: always + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Rust + run: | + rustup toolchain install stable --profile minimal --component clippy,rustfmt + rustup default stable + + - name: Check Rust formatting + run: cargo fmt --check + + - name: Run Clippy + run: cargo clippy --workspace --all-targets -- -D warnings + + - name: Run Rust tests + run: cargo test --workspace diff --git a/litellm-rust/.gitignore b/litellm-rust/.gitignore new file mode 100644 index 00000000000..ca98cd96efd --- /dev/null +++ b/litellm-rust/.gitignore @@ -0,0 +1,2 @@ +/target/ +Cargo.lock diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md new file mode 100644 index 00000000000..1d2987e0a1a --- /dev/null +++ b/litellm-rust/CLAUDE.md @@ -0,0 +1,88 @@ +# CLAUDE.md + +This file defines the rules for Rust work in LiteLLM. + +## Core Boundary + +The `core` and `providers` crates describe work; hosts execute work. + +Route-level Rust structure mirrors LiteLLM's Python responsibilities: +- `core/src//` owns the route contract, shared types, and provider + template traits. For OCR, this means `core/src/ocr`. +- `providers/src///transformation.rs` owns the + provider-specific transform. For Mistral OCR, this means + `providers/src/mistral/ocr/transformation.rs`. +- Future network execution belongs in a host/transport layer such as + `llm_http_handler`, not inside `core` or `providers`. + +Allowed in `core` and `providers`: +- Pure request transforms +- Pure response transforms +- Pure stream chunk normalization +- Shared data types and validation errors +- Deterministic token/cost helper logic + +Not allowed in `core` or `providers`: +- Network calls +- Environment variable or secret reads +- Filesystem access +- Database or cache access +- Provider SDK signing or auth flows +- Logging callbacks, spend writes, or custom callbacks +- Global mutable runtime state + +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. + +## Production Bar + +Rust code in this workspace is held to a strict parity and robustness bar from +the first PR: + +- Correctness parity is proven with tests. Do not rely on README claims or + manual inspection for a port that mirrors Python behavior. +- Every provider transform must have unit tests for supported-parameter + filtering, request body shape, response normalization, missing/null fields, + and bad-input errors. +- When Rust is exposed through Python, add Python tests that prove disabled, + enabled, and unavailable-bridge fallback behavior. +- Avoid panics on user/provider input. Return typed errors and let the host map + them to Python exceptions or HTTP responses. +- OCR handles documents that often contain personal data. Do not log document + contents, base64 payloads, provider response bodies, or secrets. +- Error messages must be useful but data-minimized. Truncate or sanitize any + upstream body before it crosses a host boundary. +- Treat empty or whitespace-only credentials, URLs, and config values as absent + at the host/config resolution layer. +- Preserve Python output shape intentionally. If a field is always serialized as + `null` for Python parity, leave a short comment explaining that parity choice. + +## Host I/O Rules + +These rules apply when adding future crates or modules that execute network I/O, +such as `ai-gateway`, router hosts, or standalone servers: + +- Set connect and full-request timeouts. No unbounded waits. +- Reuse HTTP clients; do not construct clients per request. +- Prefer rustls TLS for portable Python wheels and Linux images unless there is + a documented reason not to. +- Add request IDs and structured tracing at the host layer, without logging OCR + document contents or secrets. +- Do not echo raw upstream response bodies to callers. Sanitize and bound them. +- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is + impossible by construction and documented. + +## Checks + +Run these before pushing Rust changes. The same checks run in GitHub Actions +for changes under `litellm-rust/`. + +```bash +cd litellm-rust +cargo fmt --check +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace +``` + +When a Rust path is exposed through Python, add Python parity tests that compare +the existing Python output with the Rust-backed output. diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml new file mode 100644 index 00000000000..4117bd5f6ac --- /dev/null +++ b/litellm-rust/Cargo.toml @@ -0,0 +1,20 @@ +[workspace] +members = [ + "crates/core", + "crates/providers", + "crates/python-bridge", +] +resolver = "2" + +[workspace.package] +edition = "2021" +license = "MIT" +repository = "https://github.com/BerriAI/litellm" + +[workspace.dependencies] +litellm-core = { path = "crates/core" } +litellm-providers = { path = "crates/providers" } +pyo3 = "0.23.5" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +thiserror = "2.0" diff --git a/litellm-rust/README.md b/litellm-rust/README.md new file mode 100644 index 00000000000..15ad1855420 --- /dev/null +++ b/litellm-rust/README.md @@ -0,0 +1,34 @@ +# LiteLLM Rust + +This workspace contains the staged Rust implementation for LiteLLM. + +Rust starts as a pure transform core used by the existing Python host. Python +continues to own auth, configuration, network I/O, retries, routing, logging, +callbacks, spend tracking, and customer plugins until each Rust path has parity +coverage and production evidence. + +## Layout + +```text +crates/ + core/ Route contracts, shared pure types, errors, and templates. + src/ocr/ + providers/ Provider-specific pure transforms. + src/mistral/ocr/transformation.rs + python-bridge/ PyO3 bridge for Python LiteLLM. +``` + +The folder shape should follow the Python provider tree: +`providers/src///transformation.rs`. The bridge should expose +one function per top-level route, starting with `ocr(payload)`. + +## Checks + +Run these before pushing Rust changes. GitHub Actions runs the same checks for +changes under `litellm-rust/`. + +```bash +cargo fmt --check +cargo clippy --workspace --all-targets -- -D warnings +cargo test --workspace +``` diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md new file mode 100644 index 00000000000..51e0d215e6a --- /dev/null +++ b/litellm-rust/crates/core/CLAUDE.md @@ -0,0 +1,38 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/core`. + +## Responsibility + +`core` owns shared data types, typed errors, and deterministic helper contracts. +It must stay pure and host-independent. + +Allowed: +- Shared request/response structs. +- Typed errors with stable, non-sensitive messages. +- Deterministic validation helpers. +- Serialization helpers that intentionally mirror Python output shape. +- Route templates that match Python base config responsibilities, such as + `ocr::transformation::OcrProviderConfig`. + +Not allowed: +- Network, filesystem, database, cache, or environment access. +- Secret reads or auth/header construction. +- Logging callbacks, tracing spans, spend writes, or customer callbacks. +- Provider-specific branching that belongs in `providers`. +- Panics for user/provider-controlled input. + +## Structure + +Use route names directly under `src/`: `ocr`, future `messages`, +`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not +invent broad names like `engine` for route contracts. + +## Parity Rules + +- Every shared type used by a provider transform needs unit tests for + serialization shape. +- If Python parity requires always emitting a `null` field instead of omitting + it, document that in code and pin it with a test. +- Error enums should preserve enough detail for Python/HTTP hosts to map errors + consistently without exposing document contents or upstream bodies. diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml new file mode 100644 index 00000000000..e54002fe5e8 --- /dev/null +++ b/litellm-rust/crates/core/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "litellm-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs new file mode 100644 index 00000000000..08a7a2b5ed9 --- /dev/null +++ b/litellm-rust/crates/core/src/error.rs @@ -0,0 +1,27 @@ +use thiserror::Error; + +pub type CoreResult = Result; + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum CoreError { + #[error("expected {expected}, got {actual}")] + InvalidType { + expected: &'static str, + actual: &'static str, + }, + #[error("missing required field: {0}")] + MissingField(&'static str), + #[error("invalid response: {0}")] + InvalidResponse(String), +} + +pub fn json_type_name(value: &serde_json::Value) -> &'static str { + match value { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "bool", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs new file mode 100644 index 00000000000..a88204867c4 --- /dev/null +++ b/litellm-rust/crates/core/src/lib.rs @@ -0,0 +1,4 @@ +pub mod error; +pub mod ocr; + +pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs new file mode 100644 index 00000000000..ec2fbb969a6 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -0,0 +1,2 @@ +pub mod transformation; +pub mod types; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs new file mode 100644 index 00000000000..7353d9d22c4 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -0,0 +1,32 @@ +use serde_json::{Map, Value}; + +use crate::CoreResult; + +use super::types::{OcrRequestData, OcrResponseData}; + +pub trait OcrProviderConfig { + fn supported_ocr_params(&self) -> &'static [&'static str]; + + fn map_ocr_params(&self, non_default_params: &Map) -> Map { + let mut mapped_params = Map::new(); + for (param, value) in non_default_params { + if self.supported_ocr_params().contains(¶m.as_str()) { + mapped_params.insert(param.clone(), value.clone()); + } + } + mapped_params + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult; + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult; +} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs new file mode 100644 index 00000000000..1a72b8f1d66 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -0,0 +1,29 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrRequestData { + pub data: Value, + pub files: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct OcrResponseData { + pub pages: Vec, + pub model: String, + pub document_annotation: Option, + pub usage_info: Option, + pub object: String, +} + +impl OcrResponseData { + pub fn into_json(self) -> Value { + serde_json::json!({ + "pages": self.pages, + "model": self.model, + "document_annotation": self.document_annotation, + "usage_info": self.usage_info, + "object": self.object, + }) + } +} diff --git a/litellm-rust/crates/providers/CLAUDE.md b/litellm-rust/crates/providers/CLAUDE.md new file mode 100644 index 00000000000..0f7fdcda2aa --- /dev/null +++ b/litellm-rust/crates/providers/CLAUDE.md @@ -0,0 +1,53 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/providers`. + +## Responsibility + +`providers` owns provider-specific pure transforms. It mirrors the existing +Python provider modules closely enough that parity review is mechanical. + +Provider files should map to the Python provider tree: + +```text +providers/src///transformation.rs +``` + +For example, Mistral OCR lives at +`providers/src/mistral/ocr/transformation.rs`, matching +`litellm/llms/mistral/ocr/transformation.py`. + +Allowed: +- Provider request transforms. +- Provider response normalization. +- Supported-parameter filtering. +- Provider-specific validation that does not require I/O or secrets. + +Not allowed: +- HTTP clients or provider SDK calls. +- Environment variable reads. +- API key resolution or auth header construction. +- Logging, callbacks, spend tracking, retries, routing, cooldowns, or fallbacks. +- Panics on bad user/provider input. + +## Required Tests + +Every provider transform must include focused unit tests for: +- Supported params matching the Python provider config. +- Unknown params being dropped or transformed the same way as Python. +- Request body shape matching Python output. +- Response normalization with complete, missing, null, and extra fields. +- Bad input returning typed errors. + +For OCR specifically, assume documents can contain personal data. Tests should +prove transforms do not copy document contents into error messages. + +## Implementation Rules + +- Prefer static supported-parameter lists over allocating strings on every call. +- Keep transforms deterministic and allocation-conscious, but choose clarity over + premature micro-optimization for tiny parameter lists. +- Use typed errors from `core`; avoid stringly-typed error plumbing. +- Add comments only when they explain Python-parity decisions or provider quirks. +- Put route-level provider dispatch in a route file such as `providers/src/ocr.rs`. + Do not move provider-specific transform logic into the Python bridge. diff --git a/litellm-rust/crates/providers/Cargo.toml b/litellm-rust/crates/providers/Cargo.toml new file mode 100644 index 00000000000..68483af6266 --- /dev/null +++ b/litellm-rust/crates/providers/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "litellm-providers" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-core.workspace = true +serde_json.workspace = true + +[dev-dependencies] +serde_json.workspace = true diff --git a/litellm-rust/crates/providers/src/lib.rs b/litellm-rust/crates/providers/src/lib.rs new file mode 100644 index 00000000000..1a0ca5b7e42 --- /dev/null +++ b/litellm-rust/crates/providers/src/lib.rs @@ -0,0 +1,2 @@ +pub mod mistral; +pub mod ocr; diff --git a/litellm-rust/crates/providers/src/mistral/mod.rs b/litellm-rust/crates/providers/src/mistral/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/providers/src/mistral/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/providers/src/mistral/ocr/mod.rs b/litellm-rust/crates/providers/src/mistral/ocr/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/providers/src/mistral/ocr/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs b/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs new file mode 100644 index 00000000000..ca9073a14a0 --- /dev/null +++ b/litellm-rust/crates/providers/src/mistral/ocr/transformation.rs @@ -0,0 +1,212 @@ +use litellm_core::error::{json_type_name, CoreError, CoreResult}; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{Map, Value}; + +const SUPPORTED_OCR_PARAMS: &[&str] = &[ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "id", +]; + +pub struct MistralOcrConfig; + +pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig; + +impl OcrProviderConfig for MistralOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + if !document.is_object() { + return Err(CoreError::InvalidType { + expected: "object", + actual: json_type_name(&document), + }); + } + + let mut data = Map::new(); + data.insert("model".to_string(), Value::String(model.to_string())); + data.insert("document".to_string(), document); + for (param, value) in optional_params { + data.insert(param, value); + } + + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response_object = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + + let pages = response_object + .get("pages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let model = response_object + .get("model") + .and_then(Value::as_str) + .unwrap_or(model) + .to_string(); + let document_annotation = response_object.get("document_annotation").cloned(); + let usage_info = response_object.get("usage_info").cloned(); + + Ok(OcrResponseData { + pages, + model, + document_annotation, + usage_info, + object: "ocr".to_string(), + }) + } +} + +pub fn supported_ocr_params() -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() +} + +pub fn map_ocr_params(non_default_params: &Map) -> Map { + MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params) +} + +pub fn transform_ocr_request( + model: &str, + document: Value, + optional_params: Map, +) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) +} + +pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn supported_params_match_python_mistral_ocr_config() { + assert_eq!( + supported_ocr_params(), + &[ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "id", + ] + ); + } + + #[test] + fn map_ocr_params_drops_unknown_params() { + let params = json!({ + "extract_header": true, + "unsupported_param": "value", + "pages": [0, 1] + }); + let mapped = map_ocr_params(params.as_object().unwrap()); + + assert_eq!(mapped.get("extract_header"), Some(&json!(true))); + assert_eq!(mapped.get("pages"), Some(&json!([0, 1]))); + assert!(!mapped.contains_key("unsupported_param")); + } + + #[test] + fn transform_ocr_request_builds_mistral_body() { + let document = json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }); + let optional_params = json!({ + "include_image_base64": true, + "table_format": "html" + }) + .as_object() + .unwrap() + .clone(); + + let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params) + .expect("request should transform"); + + assert_eq!( + result.data, + json!({ + "model": "mistral-ocr-latest", + "document": document, + "include_image_base64": true, + "table_format": "html" + }) + ); + assert_eq!(result.files, None); + } + + #[test] + fn transform_ocr_request_rejects_non_object_document() { + let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new()) + .expect_err("string document should be rejected"); + + assert_eq!( + err, + CoreError::InvalidType { + expected: "object", + actual: "string", + } + ); + } + + #[test] + fn transform_ocr_response_normalizes_mistral_json() { + let response = json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": null, + "usage_info": {"pages_processed": 1} + }); + + let result = transform_ocr_response("mistral-ocr-latest", response) + .expect("response should transform"); + + assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]); + assert_eq!(result.model, "mistral-ocr-2505-completion"); + assert_eq!(result.document_annotation, Some(Value::Null)); + assert_eq!(result.usage_info, Some(json!({"pages_processed": 1}))); + assert_eq!(result.object, "ocr"); + } +} diff --git a/litellm-rust/crates/providers/src/ocr.rs b/litellm-rust/crates/providers/src/ocr.rs new file mode 100644 index 00000000000..48e3bbe8894 --- /dev/null +++ b/litellm-rust/crates/providers/src/ocr.rs @@ -0,0 +1,156 @@ +use litellm_core::error::{json_type_name, CoreError}; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::CoreResult; +use serde_json::{Map, Value}; + +use crate::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +pub fn transform(payload: Value) -> CoreResult { + let payload = payload_object(&payload)?; + let provider = required_string(payload, "provider")?; + let operation = required_string(payload, "operation")?; + + match provider { + "mistral" => transform_with_provider(&MISTRAL_OCR_CONFIG, operation, payload), + _ => Err(CoreError::InvalidResponse(format!( + "unsupported OCR provider: {provider}" + ))), + } +} + +fn transform_with_provider( + provider_config: &impl OcrProviderConfig, + operation: &str, + payload: &Map, +) -> CoreResult { + match operation { + "map_params" => { + let params = required_object(payload, "non_default_params")?; + Ok(Value::Object(provider_config.map_ocr_params(¶ms))) + } + "transform_request" => { + let model = required_string(payload, "model")?; + let document = required_value(payload, "document")?; + let optional_params = required_object(payload, "optional_params")?; + let transformed = + provider_config.transform_ocr_request(model, document, optional_params)?; + Ok(serde_json::json!({ + "data": transformed.data, + "files": transformed.files, + })) + } + "transform_response" => { + let model = required_string(payload, "model")?; + let response_json = required_value(payload, "response_json")?; + let transformed = provider_config.transform_ocr_response(model, response_json)?; + Ok(transformed.into_json()) + } + _ => Err(CoreError::InvalidResponse(format!( + "unsupported OCR operation: {operation}" + ))), + } +} + +fn payload_object(payload: &Value) -> CoreResult<&Map> { + payload.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(payload), + }) +} + +fn required_string<'a>(payload: &'a Map, key: &'static str) -> CoreResult<&'a str> { + let value = payload.get(key).ok_or(CoreError::MissingField(key))?; + value.as_str().ok_or_else(|| CoreError::InvalidType { + expected: "string", + actual: json_type_name(value), + }) +} + +fn required_object( + payload: &Map, + key: &'static str, +) -> CoreResult> { + let value = payload.get(key).ok_or(CoreError::MissingField(key))?; + value + .as_object() + .cloned() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(value), + }) +} + +fn required_value(payload: &Map, key: &'static str) -> CoreResult { + payload + .get(key) + .cloned() + .ok_or(CoreError::MissingField(key)) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn transform_dispatches_mistral_map_params() { + let result = transform(json!({ + "provider": "mistral", + "operation": "map_params", + "non_default_params": { + "extract_header": true, + "unsupported_param": "value" + } + })) + .expect("payload should transform"); + + assert_eq!(result, json!({"extract_header": true})); + } + + #[test] + fn transform_dispatches_mistral_request() { + let document = json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }); + + let result = transform(json!({ + "provider": "mistral", + "operation": "transform_request", + "model": "mistral-ocr-latest", + "document": document, + "optional_params": {"include_image_base64": true} + })) + .expect("payload should transform"); + + assert_eq!( + result, + json!({ + "data": { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + "include_image_base64": true + }, + "files": null + }) + ); + } + + #[test] + fn transform_rejects_unknown_provider() { + let err = transform(json!({ + "provider": "azure_ai", + "operation": "map_params", + "non_default_params": {} + })) + .expect_err("unsupported provider should fail"); + + assert_eq!( + err, + CoreError::InvalidResponse("unsupported OCR provider: azure_ai".to_string()) + ); + } +} diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md new file mode 100644 index 00000000000..efa1a554c9c --- /dev/null +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -0,0 +1,36 @@ +# CLAUDE.md + +Rules for `litellm-rust/crates/python-bridge`. + +## Responsibility + +`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. +Keep this crate thin. It adapts Python objects to Rust payloads and returns +Python-compatible dictionaries. + +## Bridge Shape + +- Prefer one stable method per top-level LiteLLM route, for example + `ocr(payload)`. +- Do not add one exported PyO3 function per provider helper unless there is a + measured reason. +- Provider dispatch belongs in Rust route modules such as + `litellm_providers::ocr`, not in this PyO3 crate. +- Python owns rollout state and fallback. Rust should return errors; Python + decides whether to raise or fall back. + +## Data Handling + +- OCR payloads can contain personal data and large base64 images. Do not log + payloads or provider responses. +- Avoid copying large payloads more than needed. The current JSON round-trip is + acceptable for the first scaffold, but future performance work should evaluate + direct PyO3 conversion before expanding Rust coverage to image-heavy paths. +- Do not expose raw Rust errors that include document contents or upstream + bodies. + +## Tests + +- `cargo test --workspace` must compile this crate. +- Python tests must cover bridge disabled, bridge enabled, and module-missing + fallback behavior for every exposed route. diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml new file mode 100644 index 00000000000..40275879af5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-python-bridge" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[lib] +name = "litellm_python_bridge" +crate-type = ["cdylib"] + +[dependencies] +litellm-providers.workspace = true +pyo3 = { workspace = true, features = ["extension-module"] } +serde_json.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs new file mode 100644 index 00000000000..a723daeb65a --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -0,0 +1,31 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyAny; +use serde_json::Value; + +fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { + let json = py.import("json")?; + let encoded: String = json.call_method1("dumps", (value,))?.extract()?; + serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) +} + +fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { + let json = py.import("json")?; + let encoded = + serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; + Ok(json.call_method1("loads", (encoded,))?.unbind()) +} + +#[pyfunction] +fn ocr(py: Python<'_>, payload: Py) -> PyResult> { + let payload = py_to_json(py, payload.bind(py))?; + let transformed = litellm_providers::ocr::transform(payload) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + json_to_py(py, transformed) +} + +#[pymodule] +fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(ocr, module)?)?; + Ok(()) +} diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b27082c361a..fda2cf1115f 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -20,6 +20,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.rust_bridge.ocr import get_rust_ocr_provider_config from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -278,6 +279,11 @@ def ocr( f"OCR call - model: {model}, provider: {custom_llm_provider}" ) + ocr_provider_config = get_rust_ocr_provider_config( + custom_llm_provider=custom_llm_provider, + fallback_config=ocr_provider_config, + ) + # Get litellm params using GenericLiteLLMParams (same as responses API) litellm_params = GenericLiteLLMParams(**kwargs) diff --git a/litellm/rust_bridge/CLAUDE.md b/litellm/rust_bridge/CLAUDE.md new file mode 100644 index 00000000000..a538defc557 --- /dev/null +++ b/litellm/rust_bridge/CLAUDE.md @@ -0,0 +1,39 @@ +# CLAUDE.md + +Rules for `litellm/rust_bridge`. + +## Responsibility + +This package is the Python-side bridge to optional Rust transforms. It should +route to Rust when explicitly enabled and safely return the existing Python path +when Rust is disabled, unavailable, or unsupported for a provider. + +## Naming And Shape + +- Keep this package named `rust_bridge`; do not reintroduce a vague `_rust` + package. +- Organize by LiteLLM route (`ocr/`, future `rerank/`, etc.). +- Keep route entrypoints such as `litellm/ocr/main.py` small. They should only + ask this package for a Rust-backed config or callable. +- Keep provider rollout explicit with enums or small provider registries. +- Keep rollout controlled by Python bridge APIs such as + `set_rust_core_enabled(...)`; do not add new environment variables here + unless the matching docs-repo update lands in the same rollout. +- For each route, expose a single Python-to-Rust call that passes one payload to + the PyO3 module, such as `ocr(payload)`. Do not split provider transform + operations into multiple PyO3 bridge functions. + +## Fallback Rules + +- Rust paths are off by default. +- Missing PyO3 modules must fall back unless strict mode is enabled. +- Unknown providers must return the original Python config unchanged. +- Tests must cover disabled, enabled, module-missing, and unknown-provider paths. + +## Data Handling + +- OCR inputs frequently contain personal data. Do not log documents, base64 + payloads, provider response bodies, or secrets. +- Bridge errors should be bounded and sanitized. Do not surface raw upstream + OCR bodies through Python exceptions. +- Treat blank configuration values as absent at host/config resolution time. diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py new file mode 100644 index 00000000000..e1ed9e40e0a --- /dev/null +++ b/litellm/rust_bridge/__init__.py @@ -0,0 +1,19 @@ +from litellm.rust_bridge.loader import ( + rust_core_available, + set_rust_core_enabled, + set_rust_core_strict, +) +from litellm.rust_bridge.ocr import ( + RUST_OCR_PROVIDERS, + RustOcrProvider, + get_rust_ocr_provider_config, +) + +__all__ = [ + "RUST_OCR_PROVIDERS", + "RustOcrProvider", + "get_rust_ocr_provider_config", + "rust_core_available", + "set_rust_core_enabled", + "set_rust_core_strict", +] diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py new file mode 100644 index 00000000000..51342685e8e --- /dev/null +++ b/litellm/rust_bridge/loader.py @@ -0,0 +1,68 @@ +import importlib +from types import ModuleType +from typing import Any, Iterable, Optional, Union + +_rust_module: Optional[ModuleType] = None +_rust_module_load_attempted = False +_enabled_rust_core_scopes: set[str] = set() +_rust_core_strict = False + + +def _load_rust_module() -> Optional[ModuleType]: + global _rust_module, _rust_module_load_attempted + + if _rust_module_load_attempted: + return _rust_module + + _rust_module_load_attempted = True + try: + _rust_module = importlib.import_module("litellm_python_bridge") + except Exception: + _rust_module = None + return _rust_module + + +def rust_core_available() -> bool: + return _load_rust_module() is not None + + +def set_rust_core_enabled(scopes: Union[bool, str, Iterable[str]]) -> None: + global _enabled_rust_core_scopes + + if scopes is True: + _enabled_rust_core_scopes = {"all"} + return + if scopes is False: + _enabled_rust_core_scopes = set() + return + if isinstance(scopes, str): + _enabled_rust_core_scopes = { + scope.strip() + for scope in scopes.replace(";", ",").split(",") + if scope.strip() + } + return + + _enabled_rust_core_scopes = {scope for scope in scopes if scope} + + +def rust_core_enabled(scope: str) -> bool: + return "all" in _enabled_rust_core_scopes or scope in _enabled_rust_core_scopes + + +def set_rust_core_strict(enabled: bool) -> None: + global _rust_core_strict + _rust_core_strict = enabled + + +def call_rust_function(function_name: str, *args: Any) -> Optional[Any]: + module = _load_rust_module() + if module is None: + return None + + try: + return getattr(module, function_name)(*args) + except Exception: + if _rust_core_strict: + raise + return None diff --git a/litellm/rust_bridge/ocr/__init__.py b/litellm/rust_bridge/ocr/__init__.py new file mode 100644 index 00000000000..273179030fb --- /dev/null +++ b/litellm/rust_bridge/ocr/__init__.py @@ -0,0 +1,8 @@ +from litellm.rust_bridge.ocr.config import get_rust_ocr_provider_config +from litellm.rust_bridge.ocr.providers import RUST_OCR_PROVIDERS, RustOcrProvider + +__all__ = [ + "RUST_OCR_PROVIDERS", + "RustOcrProvider", + "get_rust_ocr_provider_config", +] diff --git a/litellm/rust_bridge/ocr/config.py b/litellm/rust_bridge/ocr/config.py new file mode 100644 index 00000000000..b302a62bdd4 --- /dev/null +++ b/litellm/rust_bridge/ocr/config.py @@ -0,0 +1,160 @@ +from typing import Any, Optional, cast + +import httpx + +from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig, + DocumentType, + OCRRequestData, + OCRResponse, +) +from litellm.rust_bridge.ocr.providers import RustOcrProvider, call_ocr + + +def get_rust_ocr_provider_config( + custom_llm_provider: Optional[str], + fallback_config: BaseOCRConfig, +) -> BaseOCRConfig: + if custom_llm_provider is None: + return fallback_config + + provider_value = getattr(custom_llm_provider, "value", custom_llm_provider) + try: + rust_ocr_provider = RustOcrProvider(str(provider_value)) + except ValueError: + return fallback_config + + return cast( + BaseOCRConfig, + _RustOCRProviderConfig( + rust_ocr_provider=rust_ocr_provider, + fallback_config=fallback_config, + ), + ) + + +class _RustOCRProviderConfig: + def __init__( + self, + rust_ocr_provider: RustOcrProvider, + fallback_config: BaseOCRConfig, + ) -> None: + self.rust_ocr_provider = rust_ocr_provider + self.fallback_config = fallback_config + + def __getattr__(self, name: str) -> Any: + return getattr(self.fallback_config, name) + + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + mapped_params = call_ocr( + { + "provider": self.rust_ocr_provider.value, + "operation": "map_params", + "non_default_params": non_default_params, + } + ) + if mapped_params is not None: + return mapped_params + return self.fallback_config.map_ocr_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + ) + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + if isinstance(document, dict): + transformed_request = call_ocr( + { + "provider": self.rust_ocr_provider.value, + "operation": "transform_request", + "model": model, + "document": document, + "optional_params": optional_params, + }, + ) + if transformed_request is not None: + request_data = transformed_request.get("data") + if not isinstance(request_data, dict): + raise ValueError( + f"Rust OCR provider {self.rust_ocr_provider.value} " + "returned invalid request data" + ) + return OCRRequestData( + data=request_data, + files=transformed_request.get("files"), + ) + + return self.fallback_config.transform_ocr_request( + model=model, + document=document, + optional_params=optional_params, + headers=headers, + **kwargs, + ) + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + return self.transform_ocr_request( + model=model, + document=document, + optional_params=optional_params, + headers=headers, + **kwargs, + ) + + def transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: Any, + **kwargs, + ) -> OCRResponse: + transformed_response = call_ocr( + { + "provider": self.rust_ocr_provider.value, + "operation": "transform_response", + "model": model, + "response_json": raw_response.json(), + }, + ) + if transformed_response is not None: + return OCRResponse(**transformed_response) + + return self.fallback_config.transform_ocr_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + **kwargs, + ) + + async def async_transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: Any, + **kwargs, + ) -> OCRResponse: + return self.transform_ocr_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + **kwargs, + ) diff --git a/litellm/rust_bridge/ocr/providers.py b/litellm/rust_bridge/ocr/providers.py new file mode 100644 index 00000000000..c5f806e3117 --- /dev/null +++ b/litellm/rust_bridge/ocr/providers.py @@ -0,0 +1,35 @@ +from enum import Enum +from typing import Any, Optional + +from litellm.rust_bridge.loader import call_rust_function, rust_core_enabled + + +class RustOcrProvider(str, Enum): + MISTRAL = "mistral" + + +RUST_OCR_PROVIDERS = frozenset({RustOcrProvider.MISTRAL.value}) + + +def call_ocr(payload: dict[str, Any]) -> Optional[dict[str, Any]]: + provider = payload.get("provider") + if not isinstance(provider, str) or provider not in RUST_OCR_PROVIDERS: + return None + + if not _rust_ocr_provider_enabled(provider): + return None + + result = call_rust_function("ocr", payload) + if result is None: + return None + if not isinstance(result, dict): + raise ValueError("Rust OCR bridge returned invalid response") + return result + + +def _rust_ocr_provider_enabled(provider: str) -> bool: + return ( + rust_core_enabled("ocr") + or rust_core_enabled(f"ocr:{provider}") + or rust_core_enabled(f"{provider}_ocr") + ) diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index aededaaca77..509ded8c704 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -165,7 +165,6 @@ class TestResponseCompliance: "status", "created", "updated", - "role", "steps", "usage", ] diff --git a/tests/test_litellm/test_mistral_ocr_rust_bridge.py b/tests/test_litellm/test_mistral_ocr_rust_bridge.py new file mode 100644 index 00000000000..42ea7758637 --- /dev/null +++ b/tests/test_litellm/test_mistral_ocr_rust_bridge.py @@ -0,0 +1,183 @@ +import httpx +import pytest + +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig +from litellm.rust_bridge import loader +from litellm.rust_bridge import ocr as rust_ocr +from litellm.rust_bridge.ocr import providers +from litellm.rust_bridge.ocr import RustOcrProvider + +MODEL = "mistral-ocr-latest" +DOCUMENT = { + "type": "document_url", + "document_url": "https://example.com/doc.pdf", +} + + +@pytest.fixture(autouse=True) +def reset_rust_bridge_state(): + loader.set_rust_core_enabled(False) + loader.set_rust_core_strict(False) + yield + loader.set_rust_core_enabled(False) + loader.set_rust_core_strict(False) + + +class _FakeRustModule: + @staticmethod + def ocr(payload): + provider = payload["provider"] + operation = payload["operation"] + + assert provider == "mistral" + if operation == "map_params": + return { + key: value + for key, value in payload["non_default_params"].items() + if key != "unsupported_param" + } + if operation == "transform_request": + return { + "data": { + "model": payload["model"], + "document": payload["document"], + **payload["optional_params"], + }, + "files": None, + } + if operation == "transform_response": + response_json = payload["response_json"] + return { + "pages": response_json.get("pages", []), + "model": response_json.get("model", payload["model"]), + "document_annotation": response_json.get("document_annotation"), + "usage_info": response_json.get("usage_info"), + "object": "ocr", + } + raise AssertionError(f"Unexpected operation: {operation}") + + +class _FakeLoggingObj: + pass + + +def test_rust_ocr_provider_enum_is_explicit(): + assert providers.RUST_OCR_PROVIDERS == {RustOcrProvider.MISTRAL.value} + + +def test_unknown_ocr_provider_uses_python_fallback(): + fallback_config = MistralOCRConfig() + + config = rust_ocr.get_rust_ocr_provider_config("azure_ai", fallback_config) + + assert config is fallback_config + + +def test_rust_ocr_provider_returns_none_when_scope_disabled(monkeypatch): + monkeypatch.setattr(loader, "_load_rust_module", lambda: _FakeRustModule) + + assert ( + providers.call_ocr( + { + "provider": RustOcrProvider.MISTRAL.value, + "operation": "map_params", + "non_default_params": {"extract_header": True}, + } + ) + is None + ) + + +def test_mistral_ocr_map_params_uses_provider_gated_rust(monkeypatch): + loader.set_rust_core_enabled("ocr:mistral") + monkeypatch.setattr(loader, "_load_rust_module", lambda: _FakeRustModule) + + result = providers.call_ocr( + { + "provider": RustOcrProvider.MISTRAL.value, + "operation": "map_params", + "non_default_params": { + "extract_header": True, + "unsupported_param": "value", + }, + }, + ) + + assert result == {"extract_header": True} + + +def test_mistral_ocr_provider_wrapper_uses_rust_when_enabled(monkeypatch): + loader.set_rust_core_enabled("ocr:mistral") + monkeypatch.setattr(loader, "_load_rust_module", lambda: _FakeRustModule) + + config = rust_ocr.get_rust_ocr_provider_config("mistral", MistralOCRConfig()) + + request = config.transform_ocr_request( + model=MODEL, + document=DOCUMENT, + optional_params={"include_image_base64": True}, + headers={}, + ) + assert request.data == { + "model": MODEL, + "document": DOCUMENT, + "include_image_base64": True, + } + assert request.files is None + + response = config.transform_ocr_response( + model=MODEL, + raw_response=httpx.Response( + 200, + json={ + "pages": [{"index": 0, "markdown": "hello"}], + "model": "mistral-ocr-2505-completion", + "document_annotation": None, + "usage_info": {"pages_processed": 1}, + }, + ), + logging_obj=_FakeLoggingObj(), + ) + + assert response.pages[0].index == 0 + assert response.model == "mistral-ocr-2505-completion" + assert response.usage_info.pages_processed == 1 + + +def test_mistral_ocr_provider_wrapper_falls_back_when_rust_module_missing( + monkeypatch, +): + loader.set_rust_core_enabled("ocr:mistral") + monkeypatch.setattr(loader, "_load_rust_module", lambda: None) + + config = rust_ocr.get_rust_ocr_provider_config("mistral", MistralOCRConfig()) + + result = config.transform_ocr_request( + model=MODEL, + document=DOCUMENT, + optional_params={"include_image_base64": True}, + headers={}, + ) + + assert result.data == { + "model": MODEL, + "document": DOCUMENT, + "include_image_base64": True, + } + + +def test_mistral_ocr_config_stays_python_fallback(monkeypatch): + loader.set_rust_core_enabled("ocr:mistral") + monkeypatch.setattr(loader, "_load_rust_module", lambda: _FakeRustModule) + + config = MistralOCRConfig() + result = config.map_ocr_params( + non_default_params={ + "extract_header": True, + "unsupported_param": "value", + }, + optional_params={}, + model=MODEL, + ) + + assert result == {"extract_header": True}