diff --git a/.gitignore b/.gitignore index deb0acae56e..7da917ce450 100644 --- a/.gitignore +++ b/.gitignore @@ -147,3 +147,6 @@ crash.*.log ui/litellm-dashboard/out/ litellm.log + +.coverage-rust +coverage-rust.xml diff --git a/litellm-rust/ADDING_A_PROVIDER.md b/litellm-rust/ADDING_A_PROVIDER.md deleted file mode 100644 index ae8ae5a6870..00000000000 --- a/litellm-rust/ADDING_A_PROVIDER.md +++ /dev/null @@ -1,29 +0,0 @@ -# Adding a provider / route to litellm-rust - -Everything for a route lives in `crates/core/src//`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint. - -1. **Entrypoint** — `mod.rs`: `pub async fn (request) -> CoreResult`, the Rust equivalent of `litellm.()`, plus a `_stream` variant when the route streams. It is the only thing a host touches. -2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`. -3. **Provider config** — `crates/core/src/providers///transformation.rs`: implement that trait as a `const __CONFIG`, mirroring the Python provider tree. Add parity unit tests. -4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response. - -## Coding standards - -Before writing new logic, look for an existing base to extend. When a change is -“the same behavior for one more provider/endpoint/integration”, the codebase -almost always already has a shared abstraction for it (for example, provider -`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared -helpers in `litellm_core_utils/`, typed request/response models, or factory -functions). Find it first with a search, then add the new variant by inheriting -from or composing that base, overriding only what genuinely differs (model -name, parameter mapping, or auth). - -Never copy an existing implementation and edit it in place, and never hand-roll -a parallel version of logic a base already provides. If you catch yourself -writing a second copy of a pattern that exists twice already, stop and extract a -base instead: put the shared shape in one place and make both call sites thin -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. - -**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run the commands under "Checks" in [CLAUDE.md](CLAUDE.md). diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md deleted file mode 100644 index 17856218e60..00000000000 --- a/litellm-rust/AGENTS.md +++ /dev/null @@ -1,45 +0,0 @@ -# AGENTS.md - -litellm-rust has six crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. - -## 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-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. | -| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. | -| 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-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | -| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | - -Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers, `litellm-token-counter`, and `litellm-python-interop`. The token counter and interop foundations depend on no LiteLLM domain crate. - -## Where a route lives - -A top-level LiteLLM call is a module under `crates/core/src//`, shaped like `messages`: - -``` -core/src/messages/ - mod.rs # pub async fn messages(..) -> CoreResult<..> (+ messages_stream for SSE) - types.rs # request/response types, MessagesRequest - transformation.rs # the provider template trait - prepare.rs # provider resolution, auth headers, URL - handler.rs # the provider call - 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. - -Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. - -Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional. - -## Style - -All Rust in `litellm-rust/` follows the official Rust Style Guide: -https://doc.rust-lang.org/style-guide/ - -`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style. - -Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md deleted file mode 100644 index dfacf37b6cd..00000000000 --- a/litellm-rust/CLAUDE.md +++ /dev/null @@ -1,189 +0,0 @@ -# CLAUDE.md - -This file defines the rules for Rust work in LiteLLM. - -## Provider Coding Standards - -Before writing new logic, look for an existing base to extend. When a change is -“the same behavior for one more provider/endpoint/integration”, the codebase -almost always already has a shared abstraction for it (for example, provider -`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared -helpers in `litellm_core_utils/`, typed request/response models, or factory -functions). Find it first with a search, then add the new variant by inheriting -from or composing that base, overriding only what genuinely differs (model -name, parameter mapping, or auth). - -Never copy an existing implementation and edit it in place, and never hand-roll -a parallel version of logic a base already provides. If you catch yourself -writing a second copy of a pattern that exists twice already, stop and extract a -base instead: put the shared shape in one place and make both call sites thin -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 (see AGENTS.md) - -`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. -`litellm-config` is the config-loading boundary and returns resolved core types. -`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and -`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop` -holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate -is a layer or shared foundation, not a route; add modules, not crates. - -## Core Boundary - -`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()` -is `litellm_core::messages::messages(request).await`: you call it, it does the -provider call, and you get a typed non-streaming response back. - -Route-level Rust structure mirrors LiteLLM's Python responsibilities: -- `core/src//` owns the route end to end: the public entrypoint fn named - after the route in `mod.rs`, the request/response types (`types.rs`), the - provider template trait (`transformation.rs`), the provider/auth/URL - resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that - performs the call (`handler.rs`). `core/src/messages` is the reference. -- `core/src/providers///transformation.rs` owns the - provider-specific transform. For Anthropic Messages, this means - `core/src/providers/anthropic/messages/transformation.rs`. -- Handlers live in `core`, never in a host. `ai-gateway` must not contain a - route handler that talks to a provider; its axum route reads the HTTP request, - picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals - Python objects and calls the same entrypoint. - -Streaming keeps the same shape: the route entrypoint has a `_stream` -variant in `core` that returns the upstream response so a host can splice it to -its own caller; the host still owns no provider logic. - -Call-hook and lifecycle instrumentation, including phase timing, usage -accumulation, and callback payload construction, always lives in `core`. -Hosts feed observed events into core and dispatch the completed payloads through -their I/O logger; hosts must not own callback orchestration. - -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 -- Shared data types and validation errors -- Deterministic token/cost helper logic - -Not allowed in `core`: -- Serving HTTP: axum routes, extractors, and transport concerns stay in the host -- Filesystem access -- Database access -- Config file reading and rollout state -- Logging callbacks, spend writes, or custom callbacks -- Global mutable runtime state - -Env reads in `core` are limited to credential fallback inside a route's -`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when -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. - -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. -A new provider/route may instead be implemented rust-only with no Python -reference; then the Python interface is a thin dispatch that calls Rust with no -fallback, and you state the rust-only choice explicitly in the PR. Either way -the Python side stays minimal (it only marshals inputs and calls the Rust -interface), never add a per-route feature flag, and never push provider -dispatch into `litellm/main.py`; put it in a thin dispatch class under -`litellm/llms///`. - -## 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. - -## Network I/O Rules - -These rules apply to every module that executes network I/O, whether it is a -`core` route handler or a host such as `ai-gateway`: - -- 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. - -## Rust Style Guide - -All Rust in `litellm-rust/` follows the official Rust Style Guide: -https://doc.rust-lang.org/style-guide/ - -`rustfmt` implements the guide's formatting rules by default, so the mechanical -side is enforced for you: run `cargo fmt` before committing and CI gates every -PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add -a `rustfmt.toml` that diverges from the default style; the default style *is* the -guide. - -The guide also covers conventions rustfmt cannot auto-apply; follow these too: -- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for - types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and - statics; acronyms count as one word (`HttpClient`, not `HTTPClient`). -- Ordering and grouping the guide prescribes: imports grouped std / external / - crate-local, derives before other attributes, and consistent item order. -- Idioms the guide recommends over the formatter fighting you (e.g. prefer - restructuring an over-long expression rather than forcing an awkward wrap). - -## Constants - -Magic numbers and fixed strings go in a crate-level `constants.rs`, never -hardcoded inline — the Rust mirror of Python's `litellm/constants.py`. - -- Each crate that needs them has `src/constants.rs` (declared `mod constants;`); - import from it (`use crate::constants::...`). Don't scatter `const` values at - the top of feature modules. -- An env-overridable tunable still lives in `constants.rs` as its `DEFAULT_*` - value; the env read (with fallback to that default) happens at the host/config - resolution layer, not in `core`/`providers`. -- Exception: a value that is purely local to one function and has no meaning - elsewhere may stay inline, but prefer `constants.rs` when in doubt. - -## 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 clippy -p litellm-core --all-targets --features bedrock-auth -- -D warnings -# the ai-gateway binary + server code is behind the `server` feature -cargo clippy -p litellm-ai-gateway --all-targets --all-features -- -D warnings -cargo test --workspace -cargo test -p litellm-core --features bedrock-auth -# the `auth`, `routes`, `state` and `realtime` tests only exist under `server` -cargo test -p litellm-ai-gateway --features server -``` - -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.lock b/litellm-rust/Cargo.lock index 7b0b593b70f..7e3d25e9c5d 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1948,12 +1948,17 @@ dependencies = [ "azure_core", "azure_identity", "base64 0.22.1", + "bytes", "data-url", + "futures-util", "gcp_auth", + "mime_guess", "moka", "rand 0.8.7", "reqwest 0.12.28", "rstest", + "rustls 0.23.42", + "rustls-native-certs", "serde", "serde_json", "serde_path_to_error", @@ -1962,6 +1967,7 @@ dependencies = [ "subtle", "thiserror 2.0.19", "tokio", + "tokio-tungstenite", "tracing", "tracing-subscriber", "url", @@ -1974,12 +1980,12 @@ version = "0.1.0" dependencies = [ "criterion", "futures-util", - "litellm-ai-gateway", "litellm-core", "litellm-python-interop", "litellm-token-counter", "pyo3", "pyo3-async-runtimes", + "rstest", "serde", "serde_json", "tokio", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 5f25e69a1f8..5c72c86d6ef 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -16,6 +16,7 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +bytes = "1" tracing = "0.1" tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } litellm-core = { path = "crates/core" } diff --git a/litellm-rust/README.md b/litellm-rust/README.md deleted file mode 100644 index 650d38753e7..00000000000 --- a/litellm-rust/README.md +++ /dev/null @@ -1,56 +0,0 @@ -# LiteLLM Rust - -This workspace contains the staged Rust implementation for LiteLLM. - -`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call -that makes the LLM call and hands back a typed response, the same shape as -`litellm.messages()` in Python. - -```rust -let response = litellm_core::messages::messages(MessagesRequest { - model: "claude-sonnet-4-5", - body, - api_key: Some(key), - .. -}) -.await?; -``` - -Python continues to own configuration, retries, routing policy, logging, -callbacks, spend tracking, and customer plugins until each Rust path has parity -coverage and production evidence. - -## Crates - -| 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-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. | -| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | -| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | - -Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. - -## Layout - -```text -crates/ - core/ The SDK: route modules + provider transforms. - src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client - src/providers/anthropic/messages/transformation.rs - config/ Config loading and resolved deployments. - ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. - python-interop/ Domain-neutral PyO3 conversion and GIL primitives. - python-bridge/ PyO3 API adapter for Python LiteLLM. -``` - -The folder shape follows the Python provider tree: -`core/src/providers///transformation.rs`. The bridge exposes one -function per top-level route, mirroring the core entrypoints. - -## Checks - -Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust -changes. That list is the single source of truth and matches what GitHub Actions -runs for changes under `litellm-rust/`. diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md deleted file mode 100644 index 952bbc38b43..00000000000 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ /dev/null @@ -1,53 +0,0 @@ -# Provider coding standards (litellm-rust) - -Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response. - -## Provider resolution - -1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string. -2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers. - -## Transforms and the base config - -3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src//transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`). -4. Each provider implements that trait as a `const __CONFIG` in `core/src/providers///transformation.rs`, mirroring the Python provider tree. -5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it. -6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers. - -## Boundaries - -7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request. -8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers///`; a route is a module, never a new crate. -9. Route entry point stays thin: `core::::()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them. -10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`. - -## Types and errors - -11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string. -12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input. -13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating. -14. Early returns over deep nesting; small focused files over god modules. -15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test. - -## Safety and data minimization - -16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary. -17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer. -18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS. - -## Tests and rollout - -19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity. -20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping. -21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR. - -## Python bridge (SDK side) - -22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust. -23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms///` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method. -24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_`. - -## Checks before push - -25. Run, and keep green, the commands under "Checks" in `litellm-rust/CLAUDE.md`. - That list is the single source of truth and matches what GitHub Actions runs. diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md deleted file mode 100644 index b2fd583316b..00000000000 --- a/litellm-rust/crates/ai-gateway/AGENTS.md +++ /dev/null @@ -1,54 +0,0 @@ -# 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. - -``` -src/ - main.rs # entrypoint: build AppState (router + master key), bind, serve - state.rs # AppState — shared Arc + master_key - auth/ # authentication as an axum extractor — added to handler args - mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY) - routes/ # one module per route, all matching the same template - AGENTS.md # ← the route template (read this before adding a route) - mod.rs # app(): merges every module's router() - health.rs # simple route (one file): router() + liveness/readiness - realtime/ # route with logic → axum surface + a no-axum service: - mod.rs # router() + handler + WS<->events adapter (the axum surface) - service.rs # business logic (select deployment, call provider) — no axum, testable -``` - -## Rules - -- **Routes follow one template.** Each route module exposes - `pub fn router() -> Router`; `routes/mod.rs` only merges them. Simple - routes are one file; non-trivial routes are a folder (`handler`/`service`/ - `transport`). See `routes/AGENTS.md`. -- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's - args; it runs during extraction. Never re-implement the check per route. -- **Handlers are thin.** A handler validates and delegates to its `service`. No - business logic, no provider calls, no transforms in handlers. -- **Services call `core`, they don't reimplement it.** A `service` picks the - deployment and calls the `core` route entrypoint. Provider resolution, auth - headers, URL building, and the HTTP call are `core`'s job; a service that - builds a provider request itself is a bug (`routes/messages/service.rs` is - the reference). -- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in - `state.rs`; read env/config only in `main.rs` when building state. - -## Auth (interim) - -A single **master key** (`LITELLM_MASTER_KEY`), enforced by the -`auth::RequireMasterKey` extractor: any caller presenting it as -`Authorization: Bearer ` may invoke the gateway. Fails closed (500) when -unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to -override). Full per-key auth + budgets/rate-limits are delegated to the Python -proxy in a later phase. Health routes don't add the extractor (unauthenticated). - -## Python interop - -Python-backed loading lives in `litellm-config` and is **load-time only**. The -gateway's `python-config` feature forwards to that crate. The realtime data path -never takes the GIL. diff --git a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md deleted file mode 100644 index 6d090cf4c8e..00000000000 --- a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md +++ /dev/null @@ -1,14 +0,0 @@ -# ai-gateway architecture - -The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an -API callback: it POSTs each finished session to the LiteLLM proxy, which records -spend and runs the usual callbacks. - -```mermaid -flowchart LR - C[client] <--> G[Rust ai-gateway
LLM inference] - G <--> O[OpenAI realtime] - G -. spend tracking callback .-> P[litellm proxy] - F[litellm-config
load-time only] --> G - F -. Python backend .-> P -``` diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 74cf66e88a2..dfa61226d4e 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -13,6 +13,11 @@ name = "litellm-ai-gateway" path = "src/main.rs" required-features = ["server"] +[[bin]] +name = "trace-parity-gateway" +path = "src/bin/trace_parity_gateway.rs" +required-features = ["trace-parity"] + [dependencies] tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } diff --git a/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md b/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md deleted file mode 100644 index 84e926af243..00000000000 --- a/litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md +++ /dev/null @@ -1,55 +0,0 @@ -# Realtime gateway benchmark — pool on/off - -Measures what the gateway adds over talking to OpenAI's realtime WebSocket -directly, and what the pre-warmed connection pool removes. See -`../../src/routes/realtime/README.md` for how the pool works. - -## Results - -5000 calls / 500 concurrency, gateway at 10 instances, pool ON -(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice. -Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade, -**session** = upgrade → `session.created` (the phase the pool removes), -**1st-audio** = `response.create` → first audio delta (OpenAI inference), -**total** = full wall-clock. - -| metric | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI | -| ------------------ | ------------- | ----------------- | ------------- | ---------- | -| success rate (%) | 99.8 | 99.8 | — | — | -| dial p50 (ms) | 276 | 158 | −118 | **faster** | -| session p50 (ms) | 7 | 0 | −7 | **faster** | -| 1st-audio p50 (ms) | 440 | 664 | +224 | slower¹ | -| total p50 (ms) | 816 | 1010 | +194 | slower¹ | -| total p95 (ms) | 2152 | 1970 | −182 | **faster** | -| total p99 (ms) | 2692 | 2610 | −82 | **faster** | - -The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the -**session phase sub-millisecond** at the median — ~76% of connects hit the pool, -~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead: -`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran -slower during the gateway legs and drags `total p50` with it. - -**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the -fresh-dial overhead the pool removes. - -## Reproduce - -The load generator lives in a separate repo: -**https://github.com/ishaan-berri/litellm-realtime-bench** - -```bash -git clone https://github.com/ishaan-berri/litellm-realtime-bench -cd litellm-realtime-bench && go build -o wsbench . - -# Direct to OpenAI (baseline) -./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60 - -# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0 -./wsbench -host -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60 -``` - -Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`, -`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At -500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was -used here for 10 instances). The bench repo's README covers running 500-concurrency -legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.** diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 6f48f38c9f6..b17f17de11f 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -277,7 +277,7 @@ fn core_error_kind(error: &Error) -> &'static str { Error::InvalidProvider(_) => "InvalidProvider", Error::InvalidRequest(_) => "InvalidRequest", Error::InvalidType { .. } => "InvalidType", - Error::MissingField(_) => "MissingField", + Error::MissingField(_) | Error::MissingDocumentUrl => "MissingField", Error::Http { .. } => "HttpError", Error::InvalidResponse(_) => "InvalidResponse", Error::Network(_) => "NetworkError", diff --git a/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs new file mode 100644 index 00000000000..9036deb9871 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/bin/trace_parity_gateway.rs @@ -0,0 +1,40 @@ +use std::io::Read; + +use serde::Deserialize; +use serde_json::Value; + +#[derive(Deserialize)] +struct Input { + model_alias: String, + provider_model: String, + api_base: String, + body: Value, +} + +#[tokio::main] +async fn main() { + let mut input = String::new(); + if let Err(error) = std::io::stdin().read_to_string(&mut input) { + fail(error); + } + let input: Input = match serde_json::from_str(&input) { + Ok(input) => input, + Err(error) => fail(error), + }; + let result = litellm_ai_gateway::trace_parity::traced_messages_request( + input.model_alias, + input.provider_model, + input.api_base, + input.body, + ) + .await; + match serde_json::to_string(&result) { + Ok(result) => println!("{result}"), + Err(error) => fail(error), + } +} + +fn fail(error: impl std::fmt::Display) -> ! { + eprintln!("{error}"); + std::process::exit(1) +} diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 7f3b6b0650f..f86dd778424 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; -use std::sync::Arc; use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; @@ -10,106 +8,21 @@ use litellm_core::auth::error::MissingCredential; use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; use litellm_core::responses::types::ResponsesWsEvent; use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; -use tokio::net::TcpStream; -use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; -use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use crate::io::tls::connect_upstream; +use litellm_core::responses::websocket::{ResponsesUpstreamWs, connect_upstream}; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, }; const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -pub type ResponsesUpstreamWs = WebSocketStream>; type UpstreamTx = SplitSink; type UpstreamRx = SplitStream; -#[derive(Clone)] -pub struct ResponsesWebSocketConnection { - socket: Arc>>, -} - -impl ResponsesWebSocketConnection { - pub async fn connect_url( - url: &str, - headers: &HashMap, - timeout: Option, - ) -> Result { - let mut request = url - .into_client_request() - .map_err(|error| Error::Network(error.to_string()))?; - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|error| Error::InvalidRequest(error.to_string()))?; - let header_value = HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string()))?; - request.headers_mut().insert(header_name, header_value); - } - let connect = connect_upstream(request); - let result = match timeout { - Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - Error::Network("Responses WebSocket connection timed out".to_string()) - })?, - None => connect.await, - }; - let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { - status: response.status().as_u16(), - body: String::new(), - }, - other => Error::Network(other.to_string()), - })?; - Ok(Self { - socket: Arc::new(Mutex::new(Some(socket))), - }) - } - - pub async fn send_text(&self, text: String) -> Result<(), Error> { - let mut socket = self.socket.lock().await; - let Some(socket) = socket.as_mut() else { - return Err(Error::Network("Responses WebSocket is closed".to_string())); - }; - socket - .send(Message::Text(text)) - .await - .map_err(|error| Error::Network(error.to_string())) - } - - pub async fn recv_text(&self) -> Result, Error> { - let mut socket_guard = self.socket.lock().await; - let Some(socket) = socket_guard.as_mut() else { - return Ok(None); - }; - match socket.next().await { - Some(Ok(Message::Text(text))) => Ok(Some(text)), - Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) - .map(Some) - .map_err(|error| Error::InvalidResponse(error.to_string())), - Some(Ok(Message::Close(_))) | None => Ok(None), - Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Network(error.to_string())), - } - } - - pub async fn close(&self) -> Result<(), Error> { - let mut socket = self.socket.lock().await; - if let Some(socket) = socket.as_mut() { - socket - .close(None) - .await - .map_err(|error| Error::Network(error.to_string()))?; - } - *socket = None; - Ok(()) - } -} - pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { api_key .map(str::trim) diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 39465e28e84..3334053a0a4 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -118,7 +118,8 @@ impl IntoResponse for MessagesRouteError { | Error::Connect(_) | Error::InvalidResponse(_) | Error::InvalidType { .. } - | Error::MissingField(_) => ( + | Error::MissingField(_) + | Error::MissingDocumentUrl => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/ai-gateway/src/trace_parity.rs index 21123df3f1c..00c9b53e691 100644 --- a/litellm-rust/crates/ai-gateway/src/trace_parity.rs +++ b/litellm-rust/crates/ai-gateway/src/trace_parity.rs @@ -10,6 +10,7 @@ use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; use serde::Serialize; use serde_json::Value; use tower::ServiceExt; +use tracing::instrument::WithSubscriber; use crate::io::realtime_pool::RealtimePool; use crate::routes; @@ -21,6 +22,38 @@ pub struct GatewayResponse { pub body: Value, } +#[derive(Debug, Serialize)] +pub struct TracedGatewayResponse { + pub response: Option, + pub error: Option, + pub trace: Vec, +} + +pub async fn traced_messages_request( + model_alias: String, + provider_model: String, + api_base: String, + body: Value, +) -> TracedGatewayResponse { + let trace = litellm_core::observability::FunctionTrace::default(); + let result = messages_request(model_alias, provider_model, api_base, body) + .with_subscriber(trace.dispatcher()) + .await; + let events = trace.events(); + match result { + Ok(response) => TracedGatewayResponse { + response: Some(response), + error: None, + trace: events, + }, + Err(error) => TracedGatewayResponse { + response: None, + error: Some(error.to_string()), + trace: events, + }, + } +} + pub async fn messages_request( model_alias: String, provider_model: String, diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs index 05f7d9610d5..ac37440d682 100644 --- a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -2,10 +2,10 @@ //! API has to resolve its own crypto provider, in a test binary where nothing //! has installed a process-wide one, and has to leave it uninstalled. -use std::collections::HashMap; use std::time::Duration; -use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; +use futures_util::{sink, stream}; +use litellm_ai_gateway::io::responses_ws::async_responses_websocket; use tokio::net::TcpListener; async fn dead_tls_server() -> u16 { @@ -30,10 +30,15 @@ async fn dead_tls_server() -> u16 { async fn dialing_wss_returns_an_error_instead_of_panicking() { let port = dead_tls_server().await; - let result = ResponsesWebSocketConnection::connect_url( - &format!("wss://127.0.0.1:{port}/"), - &HashMap::new(), + let result = async_responses_websocket( + "gpt-5", + Some("test-key"), + Some(&format!("wss://127.0.0.1:{port}/")), + None, Some(Duration::from_secs(10)), + |_| {}, + stream::empty(), + sink::drain(), ) .await; diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index aee8b4937ef..9ba7bfb5323 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -2,6 +2,6 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve 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`. +Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. 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. diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md deleted file mode 100644 index 5d36305ded5..00000000000 --- a/litellm-rust/crates/core/CLAUDE.md +++ /dev/null @@ -1,66 +0,0 @@ -# CLAUDE.md - -Rules for `litellm-rust/crates/core`. - -## Responsibility - -`core` is the LiteLLM SDK in Rust: it makes the LLM call. Every top-level -LiteLLM call has a public entrypoint here, named after the route -(`messages::messages()` is the Rust equivalent of `litellm.messages()`), and -calling it returns a typed non-streaming response. - -Allowed: -- The public entrypoint for a route, plus its `_stream` variant when the - route supports streaming. -- Provider resolution, auth header construction, URL building, and the provider - HTTP call (shared reused client, connect + request timeouts). -- 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 - `messages::transformation::AnthropicMessagesProviderConfig`. - -Not allowed: -- Serving HTTP: axum routers, extractors, and other transport concerns. -- Filesystem, database, or cache access. -- Config file reading or rollout state; the host resolves those and passes them - in. Env reads are limited to credential fallback in a route's `prepare.rs`. -- Logging callbacks, tracing spans, spend writes, or customer callbacks. -- Provider-specific branching that belongs in `providers`. -- Panics for user/provider-controlled input. - -## Typed Contracts (core rule) - -Trait and function boundaries MUST be strongly typed. No stringly-typed JSON -(`&str` / `String` / `Vec` / bare `serde_json::Value`) as a transform -input or output. Parse wire bytes into typed structs/enums at the host edge; -`core` and `providers` operate only on those types (e.g. `RealtimeEvent`, -`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a -typed field on a struct, not a raw string threaded through the API. - -## Structure - -Use route names directly under `src/`: `messages`, `ocr`, future -`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not -invent broad names like `engine` for route contracts. - -`src/messages` is the reference shape for a route module: - -``` -mod.rs pub async fn messages(..) (+ messages_stream) -types.rs request/response types -transformation.rs the provider template trait -prepare.rs provider resolution, auth headers, URL -handler.rs the provider call -client.rs the shared reqwest client -``` - -## 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 index a2433435e34..09c526f73cf 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -6,25 +6,27 @@ license.workspace = true repository.workspace = true autotests = false -[[test]] -name = "workspace_crate_allowlist" -path = "tests/workspace_crate_allowlist.rs" - [dependencies] +bytes.workspace = true +futures-util.workspace = true base64.workspace = true azure_core.workspace = true azure_identity.workspace = true data-url = "0.3.2" gcp_auth.workspace = true moka.workspace = true +mime_guess = "2.0.5" rand.workspace = true reqwest.workspace = true +rustls.workspace = true +rustls-native-certs.workspace = true serde.workspace = true serde_json.workspace = true serde_path_to_error = "0.1" strum.workspace = true subtle.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["sync"] } +tokio-tungstenite.workspace = true thiserror.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, optional = true } diff --git a/litellm-rust/crates/core/src/auth/credential.rs b/litellm-rust/crates/core/src/auth/credential.rs index b5235b6780c..c64d331b877 100644 --- a/litellm-rust/crates/core/src/auth/credential.rs +++ b/litellm-rust/crates/core/src/auth/credential.rs @@ -9,6 +9,21 @@ use crate::AuthError; use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; +pub fn credential_index(requested: &str, names: &[String]) -> Option { + names.iter().position(|name| name == requested) +} + +pub fn credential_default_fields<'a>( + supplied: &[String], + credential_fields: &'a [String], +) -> Vec<&'a str> { + credential_fields + .iter() + .filter(|name| !supplied.contains(name)) + .map(String::as_str) + .collect() +} + #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { Path(PathBuf), diff --git a/litellm-rust/crates/core/src/auth/mod.rs b/litellm-rust/crates/core/src/auth/mod.rs index 35d9c676f65..2940a983fb9 100644 --- a/litellm-rust/crates/core/src/auth/mod.rs +++ b/litellm-rust/crates/core/src/auth/mod.rs @@ -49,6 +49,7 @@ impl Sourced { pub use credential::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, + credential_default_fields, credential_index, }; pub use http::{CredentialPlacement, RequestAuth}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; diff --git a/litellm-rust/crates/core/src/call_lifecycle/README.md b/litellm-rust/crates/core/src/call_lifecycle/README.md deleted file mode 100644 index 692e249ef27..00000000000 --- a/litellm-rust/crates/core/src/call_lifecycle/README.md +++ /dev/null @@ -1,167 +0,0 @@ -# Call lifecycle - -`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call -types migrated to Rust. It owns lifecycle ordering, phase timing, and trace -observer calls. It must not know about OCR, chat, messages, responses, -completions, provider auth, request transforms, or response normalization. - -Call-type modules own their domain behavior. For example, OCR owns document -payloads, OCR provider transforms, safe document fetch, guardrail payload shape, -callback payload shape, and provider HTTP execution. - -## Runtime order - -Every wrapped call runs in this order: - -1. `async_pre_call_hook` -2. `async_during_call_hook` -3. provider call -4. `async_log_success_event` or `async_log_failure_event` - -`async_pre_call_hook` receives the initial LiteLLM request shape. It is where -pre-call custom guardrails run. - -`async_during_call_hook` converts the initial request into the provider-ready -request. It is where provider config selection, parameter mapping, auth/header -resolution, request transforms, and during-call guardrails belong. - -The provider call receives only the provider-ready request. It should execute -I/O and call the provider response transform. - -Success and failure callbacks receive `CallLifecycleTiming`. Callback failures -must not replace the original provider or guardrail result. - -## Trace contract - -The lifecycle runner records: - -- full call start and end time -- `pre_call` phase timing -- `during_call` phase timing -- `provider_call` phase timing -- `success_callback` phase timing -- `failure_callback` phase timing - -`CallLifecycleObserver` receives phase start and end events. The default -observer is a no-op. Future OTEL support should implement this observer instead -of editing OCR, chat, messages, responses, completions, or provider modules. - -## Required shape - -Each migrated call type should use this folder shape: - -```text -litellm-rust/crates/ai-gateway/src// - mod.rs # thin public entrypoint - types.rs # public request, prepared request, provider request, response types - prepare.rs # model/provider/callback/guardrail setup - hooks.rs # CallLifecycleHooks implementation - handler.rs # provider I/O and response normalization - tests.rs # call-type lifecycle and handler tests -``` - -Provider transforms can live in `litellm-rust/crates/core/src/providers/...`. -Shared call-type helpers can live beside the call type, but generic lifecycle -code stays in this folder. - -## Core API - -The prepared request implements `CallLifecycleRequest`: - -```rust -impl CallLifecycleRequest for PreparedMessagesRequest { - fn lifecycle_context(&self) -> CallLifecycleContext { - CallLifecycleContext::new( - "messages", - self.model.clone(), - self.custom_llm_provider.clone(), - self.litellm_call_id.clone(), - ) - } -} -``` - -The call-type hooks implement `CallLifecycleHooks`: - -```rust -impl CallLifecycleHooks< - PreparedMessagesRequest, - ProviderMessagesRequest, - MessagesResponse, -> for MessagesLifecycleHooks { - fn async_pre_call_hook(...) { - // run pre-call custom guardrails against the LiteLLM request shape - } - - fn async_during_call_hook(...) { - // map params, validate env, transform request, run during-call guardrails - } - - fn async_log_success_event(...) { - // call async_log_success_event on configured custom loggers - } - - fn async_log_failure_event(...) { - // call async_log_failure_event without swallowing the original error - } -} -``` - -The public entrypoint stays thin: - -```rust -pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { - let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?; - - CallLifecycle::default() - .run_request(request, &hooks, execute_messages_provider_call) - .await -} -``` - -Use `run_request` for new call types. Keep `run` available only for specialized -tests or existing code that already has a `CallLifecycleContext`. - -## Adding a new call type - -1. Add `/types.rs` - -Define the public request accepted by the bridge, the prepared request used by -the lifecycle runner, and the provider request consumed by the handler. - -2. Implement `CallLifecycleRequest` - -Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`. -Do not put provider-specific logic here. - -3. Add `/prepare.rs` - -Resolve model/provider once, generate or preserve `litellm_call_id`, construct -callback and guardrail runners, and return `PreparedCall`. - -4. Add `/hooks.rs` - -Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction, -provider config selection, param mapping, request transform, during-call -guardrail payload construction, and callback payload construction here. - -5. Add `/handler.rs` - -Execute the provider request and normalize the provider response. Do not repeat -provider-specific transforms here; call the provider config. - -6. Add tests - -Cover hook order, success callback payload, failure callback payload, pre-call -guardrail blocking before provider I/O, during-call body mutation, and provider -error mapping. - -## Review checklist - -- Core lifecycle has no call-type or provider-specific branches -- Public call-type entrypoint only prepares and calls `run_request` -- Provider behavior lives behind provider config/transformation code -- Hook method names map to the Python custom logger and guardrail concepts -- Phase timing is recorded once in lifecycle, not separately per call type -- Callback failures never hide the original provider or guardrail error -- Tests prove the provider socket is not touched when pre-call guardrails block diff --git a/litellm-rust/crates/core/src/call_lifecycle/host.rs b/litellm-rust/crates/core/src/call_lifecycle/host.rs new file mode 100644 index 00000000000..ac6ddf99b9e --- /dev/null +++ b/litellm-rust/crates/core/src/call_lifecycle/host.rs @@ -0,0 +1,121 @@ +use std::future::Future; +use std::pin::Pin; + +pub enum HostCallStep { + Host(O), + Complete(C), +} + +pub type HostCallFuture<'a, O, C> = + Pin, crate::Error>> + Send + 'a>>; + +pub trait HostCall: Send + Sync { + type Operation: Send + 'static; + type Result: Send + 'static; + type Complete: Send + 'static; + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete>; + + fn interrupt( + &mut self, + failure: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete>; +} + +pub enum HostStep { + Ready(V), + Suspend(S), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum HostPhase { + Setup, + DeploymentPreCall, + Prepare, + Execute, + ConstructResponse, + DeploymentPostCall, + Finalize, + Success, + MapFailure, + DeploymentFailure, + Failure, + AsyncFailure, + Complete, +} + +#[derive(Clone, Debug)] +pub enum HostFailure { + Error(crate::Error), + Cancelled(crate::Error), +} + +pub struct HostLifecycle { + phase: HostPhase, + asynchronous: bool, +} + +impl HostLifecycle { + pub fn new(asynchronous: bool) -> Self { + Self { + phase: HostPhase::Setup, + asynchronous, + } + } + + pub fn phase(&self) -> HostPhase { + self.phase + } + + pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option { + if let Err(failure) = result { + if self.phase == HostPhase::DeploymentFailure { + self.phase = HostPhase::Failure; + return None; + } + let error = match failure { + HostFailure::Cancelled(error) => { + self.phase = HostPhase::Complete; + return Some(error); + } + HostFailure::Error(error) => error, + }; + match self.phase { + HostPhase::Failure | HostPhase::AsyncFailure => { + self.advance(); + return None; + } + HostPhase::Success => self.phase = HostPhase::Complete, + HostPhase::Execute | HostPhase::ConstructResponse => { + self.phase = HostPhase::MapFailure; + } + _ => self.phase = HostPhase::Failure, + } + return Some(error); + } + self.advance(); + None + } + + fn advance(&mut self) { + self.phase = match self.phase { + HostPhase::Setup if self.asynchronous => HostPhase::DeploymentPreCall, + HostPhase::Setup | HostPhase::DeploymentPreCall => HostPhase::Prepare, + HostPhase::Prepare => HostPhase::Execute, + HostPhase::Execute => HostPhase::ConstructResponse, + HostPhase::ConstructResponse if self.asynchronous => HostPhase::DeploymentPostCall, + HostPhase::ConstructResponse | HostPhase::DeploymentPostCall => HostPhase::Finalize, + HostPhase::Finalize => HostPhase::Success, + HostPhase::MapFailure if self.asynchronous => HostPhase::DeploymentFailure, + HostPhase::MapFailure | HostPhase::DeploymentFailure => HostPhase::Failure, + HostPhase::Failure if self.asynchronous => HostPhase::AsyncFailure, + HostPhase::Failure + | HostPhase::AsyncFailure + | HostPhase::Success + | HostPhase::Complete => HostPhase::Complete, + }; + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index 637c156e192..5c752a73899 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -3,6 +3,10 @@ use std::time::{Instant, SystemTime, UNIX_EPOCH}; use crate::Error; +pub mod host; +#[cfg(test)] +#[path = "../../tests/host_lifecycle.rs"] +mod host_tests; pub mod types; pub use types::{ diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index 9469d379462..1babb0078b8 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -46,9 +46,10 @@ pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10; +pub(crate) const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600; pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10; -pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; +pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024; pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10; pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120; @@ -63,3 +64,6 @@ pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY"; pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://"; pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr"; pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1"; + +pub(crate) const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com"; +pub(crate) const COHERE_API_KEY_ENV: &str = "COHERE_API_KEY"; diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index fa4a9d36e03..359ad56c336 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -1,6 +1,6 @@ use thiserror::Error as ThisError; -#[derive(Debug, ThisError, PartialEq, Eq)] +#[derive(Clone, Debug, ThisError, PartialEq, Eq)] pub enum Error { #[error("expected {expected}, got {actual}")] InvalidType { @@ -9,6 +9,8 @@ pub enum Error { }, #[error("missing required field: {0}")] MissingField(&'static str), + #[error("Document URL is required")] + MissingDocumentUrl, #[error("invalid response: {0}")] InvalidResponse(String), #[error("invalid provider: {0}")] @@ -52,6 +54,17 @@ pub enum Error { Unsupported(&'static str), } +impl Error { + pub const fn http_status_code(&self) -> Option { + match self { + Self::InvalidRequest(_) => Some(400), + Self::MissingDocumentUrl => Some(500), + Self::Http { status, .. } => Some(*status), + _ => None, + } + } +} + #[derive(Debug, ThisError)] pub(crate) enum MediaError { #[error("media URL rejected by network policy")] @@ -106,6 +119,7 @@ impl From for Error { fn from(error: crate::ocr::error::OcrRequestError) -> Self { match error { crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field), + crate::ocr::error::OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl, error => Self::InvalidRequest(error.to_string()), } } diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs new file mode 100644 index 00000000000..4c8455a171c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs @@ -0,0 +1,131 @@ +use super::super::OcrAdapter; +use crate::Error; +use crate::ocr::OcrClient; +use crate::ocr::codecs::cohere::{ + CohereParams, CohereResponse, transform_request, transform_response, validate_document, +}; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{credential_env, transform_request_body}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; +use crate::providers::azure_ai::auth::AzureAuthInputs; +use crate::url_utils::ApiUrl; + +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; + +pub(crate) struct AzureCohereAdapter; + +impl OcrAdapter for AzureCohereAdapter { + type ProviderResponse = CohereResponse; + const PROVIDER: OcrProvider = OcrProvider::AzureAi; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let params = super::super::super::wire::decode_request_value::( + serde_json::Value::Object(request.optional_params.clone()), + "optional_params", + )?; + let mut config = AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); + let base = request + .connection + .api_base + .clone() + .or_else(|| credential_env(AZURE_AI_API_BASE_ENV)) + .filter(|base| !base.trim().is_empty()) + .ok_or_else(|| { + Error::Auth( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(), + ) + })?; + let headers = + super::validate_ai_environment(&request.connection, &config, &credential_env).await?; + validate_document(&request.document)?; + let remote = request.document.source().starts_with("http://") + || request.document.source().starts_with("https://"); + let document = inline_remote_document( + client.document_fetcher(), + request.document.clone(), + &request.connection, + ) + .await?; + let body = transform_request(&request.model, document, params)?; + transform_request_body( + client, + request, + &complete_url(&base)?, + &headers, + !remote, + body, + |body| { + validate_document(&body.document)?; + validate_inline_document(&body.document) + }, + ) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + transform_response(&request.model, response) + } +} + +fn complete_url(base: &str) -> Result { + let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + if !matches!(url.scheme(), "http" | "https") { + return Err(invalid_api_base().into()); + } + let path = url.path().trim_end_matches('/').to_string(); + if path.ends_with("/v2/parse") { + url.set_path(&path); + return Ok(url.into()); + } + url.set_path(path.strip_suffix("/models").unwrap_or(&path)); + ApiUrl::parse(url.as_str()) + .and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base().into()) +} + +fn invalid_api_base() -> OcrRequestError { + OcrRequestError::RequestField { + path: "api_base".into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() { + for suffix in [ + "", + "/models", + "/providers/cohere/v2", + "/providers/cohere/v2/parse", + ] { + assert_eq!( + complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(), + "https://example.com/providers/cohere/v2/parse?tenant=a" + ); + } + assert_eq!( + complete_url("https://example.com/v2/parse?tenant=a").unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + assert!(complete_url("relative/path").is_err()); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs index 71ca69ddc58..e90c27ba59d 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs @@ -10,7 +10,6 @@ use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; use crate::ocr::prepare::{credential_env, transform_request_body}; use crate::ocr::registry::OcrProvider; use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat}; -use crate::ocr::wire::DecodedOcrResponse; use crate::providers::azure_ai::auth::AzureAuthInputs; use crate::url_utils::ApiUrl; @@ -32,18 +31,19 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { client: &OcrClient, ) -> Result { let params = map_ocr_params(request)?; - let config = AzureAuthInputs::from_sourced_optional_params( + let mut config = AzureAuthInputs::from_sourced_optional_params( &request.optional_params, &request.input_sources, ) .map_err(Error::from)?; + config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); let headers = validate_environment(&request.connection, &config, &credential_env).await?; let endpoint = nonblank(request.connection.api_base.clone()) .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) .ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?; let url = get_complete_url(&endpoint, &request.model, ¶ms)?; let body = document_intelligence::transform_ocr_request(request.document.clone())?; - transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await + transform_request_body(client, request, &url, &headers, false, body, |_| Ok(())).await } fn transform_ocr_response( @@ -61,7 +61,7 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { url: &str, headers: &[(String, String)], request: &LiteLLMOcrRequest, - ) -> Result, OcrError> { + ) -> Result, OcrError> { polling::read_operation_response( client.polling_http(), response, @@ -69,6 +69,7 @@ impl OcrAdapter for AzureDocumentIntelligenceAdapter { headers, &request.connection, request.response_format()? == OcrResponseFormat::Native, + &request.hooks, ) .await } diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs index 1bddea0da4f..6ed1e4441d4 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Duration; use reqwest::Url; @@ -9,6 +10,7 @@ use crate::ocr::codecs::document_intelligence::{ AzureDocumentIntelligenceOperation, OperationStatus, }; use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError}; +use crate::ocr::hooks::OcrHooks; use crate::ocr::types::OcrConnection; use crate::ocr::wire::DecodedOcrResponse; @@ -19,24 +21,33 @@ pub(super) async fn read_operation_response( headers: &[(String, String)], connection: &OcrConnection, native: bool, + hooks: &Arc, ) -> Result, OcrError> { if response.status() != reqwest::StatusCode::ACCEPTED { - return read_json_response(response, native).await; + let bytes = + crate::ocr::client::read_response_bytes(response, connection.max_response_bytes) + .await?; + crate::ocr::handler::post_call(hooks, &bytes).await?; + return Ok(crate::ocr::wire::decode_response(&bytes, native)?); } let location = response .headers() .get("operation-location") .and_then(|value| value.to_str().ok()) - .ok_or(OcrPollingError::PollLocation)?; + .ok_or(OcrPollingError::PollLocation)? + .to_string(); let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?; - let operation = Url::parse(location).map_err(|_| OcrPollingError::PollOrigin)?; + let operation = Url::parse(&location).map_err(|_| OcrPollingError::PollOrigin)?; if original.origin() != operation.origin() || !operation.username().is_empty() || operation.password().is_some() { return Err(OcrPollingError::PollOrigin.into()); } - poll_operation(http_client, operation, headers, connection, native).await + let bytes = + crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?; + crate::ocr::handler::post_call(hooks, &bytes).await?; + poll_operation(http_client, operation, headers, connection, native, hooks).await } async fn poll_operation( @@ -45,6 +56,7 @@ async fn poll_operation( headers: &[(String, String)], connection: &OcrConnection, native: bool, + hooks: &Arc, ) -> Result, OcrError> { let deadline = Instant::now() .checked_add(connection.poll_timeout) @@ -75,12 +87,19 @@ async fn poll_operation( .max(1); let decoded = tokio::time::timeout_at( deadline, - read_json_response::(response, native), + read_json_response::( + response, + native, + connection.max_response_bytes, + ), ) .await .map_err(|_| OcrPollingError::PollTimeout)??; match &decoded.data.status { - Some(OperationStatus::Succeeded) => return Ok(decoded), + Some(OperationStatus::Succeeded) => { + crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?; + return Ok(decoded); + } Some(OperationStatus::Running | OperationStatus::NotStarted) => { tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry))) .await diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs index 3107494d39e..8639590b05c 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs @@ -33,13 +33,16 @@ impl OcrAdapter for AzureMistralAdapter { known: params, extra_params: _extra_params, } = _prepare_ocr_request::(request)?; - let config = AzureAuthInputs::from_sourced_optional_params( + let mut config = AzureAuthInputs::from_sourced_optional_params( &request.optional_params, &request.input_sources, ) .map_err(Error::from)?; - let headers = validate_environment(&request.connection, &config, &credential_env).await?; + config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; + let headers = validate_environment(&request.connection, &config, &credential_env).await?; + let retains_document = !request.document.source().starts_with("http://") + && !request.document.source().starts_with("https://"); let document = inline_remote_document( client.document_fetcher(), request.document.clone(), @@ -47,9 +50,15 @@ impl OcrAdapter for AzureMistralAdapter { ) .await?; let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; - transform_request_body(client, request, &url, &headers, body, |body| { - validate_inline_document(&body.document) - }) + transform_request_body( + client, + request, + &url, + &headers, + retains_document, + body, + |body| validate_inline_document(&body.document), + ) .await } @@ -83,12 +92,15 @@ fn get_complete_url( }) } -async fn validate_environment( +pub(in crate::ocr::adapters) async fn validate_environment( connection: &OcrConnection, config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, OcrError> { if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + if config.azure_ad_token_provider.is_some() { + super::resolve_entra(config, env_lookup).await?; + } super::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs index 9c02a7471c9..3d30ae6d6bd 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs @@ -1,3 +1,4 @@ +mod cohere; mod document_intelligence; mod mistral; @@ -10,8 +11,10 @@ use crate::ocr::error::OcrError; use crate::ocr::types::OcrConnection; use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService}; +pub(crate) use cohere::AzureCohereAdapter; pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter; pub(crate) use mistral::AzureMistralAdapter; +pub(super) use mistral::validate_environment as validate_ai_environment; async fn resolve_entra( config: &AzureAuthInputs, @@ -22,6 +25,10 @@ async fn resolve_entra( .get_or_init(AzureAuthService::default) .get_azure_ad_token(config, env_lookup) .await + .or_else(|error| match error { + crate::AuthError::EmptyAzureToken => Ok(None), + other => Err(other), + }) .map(|credential| { credential.map(|credential| { let source = credential.source(); diff --git a/litellm-rust/crates/core/src/ocr/adapters/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs new file mode 100644 index 00000000000..933ead7f7f7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs @@ -0,0 +1,123 @@ +use super::OcrAdapter; +use crate::Error; +use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE}; +use crate::ocr::OcrClient; +use crate::ocr::codecs::cohere::{ + CohereParams, CohereResponse, transform_request, transform_response, validate_document, +}; +use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; +use crate::ocr::prepare::{credential_env, transform_request_body}; +use crate::ocr::registry::OcrProvider; +use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; +use crate::url_utils::ApiUrl; + +pub(crate) struct CohereAdapter; + +impl OcrAdapter for CohereAdapter { + type ProviderResponse = CohereResponse; + const PROVIDER: OcrProvider = OcrProvider::Cohere; + + async fn prepare_request( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let params = super::super::wire::decode_request_value::( + serde_json::Value::Object(request.optional_params.clone()), + "optional_params", + )?; + let headers = validate_environment(&request.connection, &credential_env)?; + let url = complete_url( + request + .connection + .api_base + .as_deref() + .unwrap_or(COHERE_PARSE_API_BASE), + )?; + let body = transform_request(&request.model, request.document.clone(), params)?; + transform_request_body(client, request, &url, &headers, true, body, |body| { + validate_document(&body.document) + }) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + transform_response(&request.model, response) + } +} + +fn complete_url(base: &str) -> Result { + let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(invalid_api_base().into()); + } + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base().into()) +} + +fn invalid_api_base() -> OcrRequestError { + OcrRequestError::RequestField { + path: "api_base".into(), + } +} + +fn validate_environment( + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, OcrError> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(COHERE_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| { + Error::Auth("Missing COHERE_API_KEY - set it in the environment or pass api_key".into()) + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() { + for suffix in ["", "/v2", "/v2/parse"] { + assert_eq!( + complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + } + } + + #[test] + fn rejects_invalid_urls_and_blank_keys() { + assert!(complete_url("relative/path").is_err()); + assert!(complete_url("ftp://example.com").is_err()); + assert!(matches!( + validate_environment( + &OcrConnection { + api_key: Some(" ".into()), + ..Default::default() + }, + &|_| None, + ), + Err(OcrError::Public(Error::Auth(_))) + )); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs index ea569ffb34f..cdbc2c3effc 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs @@ -33,7 +33,7 @@ impl OcrAdapter for MistralAdapter { let url = get_complete_url(request.connection.api_base.as_deref())?; let body = mistral::transform_ocr_request(&request.model, request.document.clone(), ¶ms)?; - transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await + transform_request_body(client, request, &url, &headers, true, body, |_| Ok(())).await } fn transform_ocr_response( diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs index 9171d11836c..d473fcad280 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -5,15 +5,16 @@ use serde::de::DeserializeOwned; use super::OcrClient; use super::error::{OcrError, OcrResponseError}; use super::registry::OcrProvider; -use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat}; -use super::wire::DecodedOcrResponse; +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; mod azure; +mod cohere; mod mistral; mod reducto; mod vertex; -pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; +pub(crate) use azure::{AzureCohereAdapter, AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; +pub(crate) use cohere::CohereAdapter; pub(crate) use mistral::MistralAdapter; pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter}; pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter}; @@ -55,18 +56,27 @@ pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static { _url: &str, _headers: &[(String, String)], request: &LiteLLMOcrRequest, - ) -> impl Future, OcrError>> + Send - { - let retain_native = request - .response_format() - .map(|format| format == OcrResponseFormat::Native); - async move { super::client::read_json_response(response, retain_native?).await } + ) -> impl Future< + Output = Result, OcrError>, + > + Send { + async move { + let bytes = + super::client::read_response_bytes(response, request.connection.max_response_bytes) + .await?; + super::handler::post_call(&request.hooks, &bytes).await?; + Ok(super::wire::decode_response( + &bytes, + request.response_format()? == super::types::OcrResponseFormat::Native, + )?) + } } } macro_rules! for_each_ocr_adapter { ($callback:ident) => { $callback! { + Cohere, $crate::ocr::adapters::CohereAdapter, $crate::ocr::adapters::CohereAdapter, Cohere; + AzureCohere, $crate::ocr::adapters::AzureCohereAdapter, $crate::ocr::adapters::AzureCohereAdapter, AzureAi; Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral; AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi; AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs index 062a0071a34..8889bcd1b45 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs @@ -27,7 +27,7 @@ impl OcrAdapter for ReductoLegacyAdapter { } = _prepare_ocr_request::(request)?; let headers = super::validate_environment(&request.connection, &credential_env)?; let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; - let document = guardrail_document(request, &url).await?; + let (document, headers) = guardrail_document(request, &url, &headers).await?; let document = super::prepare_document(client, document, &request.connection, &headers).await?; let body = reducto::transform_legacy_ocr_request(&request.model, document, ¶ms)?; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs index 7621d0d326a..2dafe291674 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs @@ -93,7 +93,7 @@ pub(super) async fn prepare_document( .map_err(crate::error::TransportError::from)?; let uploaded = crate::ocr::client::read_json_response::< crate::ocr::codecs::reducto::ReductoUploadResponse, - >(response, false) + >(response, false, connection.max_response_bytes) .await? .data; let file_id = uploaded diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs index a49f8105e26..c272d31b67e 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs @@ -27,7 +27,7 @@ impl OcrAdapter for ReductoV3Adapter { } = _prepare_ocr_request::(request)?; let headers = super::validate_environment(&request.connection, &credential_env)?; let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; - let document = guardrail_document(request, &url).await?; + let (document, headers) = guardrail_document(request, &url, &headers).await?; let document = super::prepare_document(client, document, &request.connection, &headers).await?; let body = reducto::transform_v3_ocr_request(&request.model, document, ¶ms)?; diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs index ef188f8b9ac..d16b3e7f386 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs @@ -57,9 +57,15 @@ impl OcrAdapter for VertexDeepSeekAdapter { let document = request.document.clone(); let body = deepseek::transform_ocr_request(&provider_model(&request.model), document, ¶ms)?; - transform_request_body(client, request, &url, &authentication.headers, body, |_| { - Ok(()) - }) + transform_request_body( + client, + request, + &url, + &authentication.headers, + false, + body, + |_| Ok(()), + ) .await } diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs index f3335bf497c..88c61725cee 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs @@ -54,6 +54,8 @@ impl OcrAdapter for VertexMistralAdapter { &location, &request.model, )?; + let retains_document = !request.document.source().starts_with("http://") + && !request.document.source().starts_with("https://"); let document = inline_remote_document( client.document_fetcher(), request.document.clone(), @@ -66,6 +68,7 @@ impl OcrAdapter for VertexMistralAdapter { request, &url, &authentication.headers, + retains_document, body, |body| validate_inline_document(&body.document), ) diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index ab2d098d0bb..394ca778d2f 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,10 +1,10 @@ use std::sync::OnceLock; use std::time::Duration; +use bytes::{Bytes, BytesMut}; use serde::de::DeserializeOwned; -use super::error::OcrError; -use super::handler::perform_ocr_request; +use super::error::{OcrError, OcrResponseError}; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; use super::wire::{DecodedOcrResponse, decode_response}; use crate::Error; @@ -32,6 +32,10 @@ impl OcrClient { }) } + pub fn shared() -> Result { + shared_client() + } + #[tracing::instrument( name = "ocr", target = "litellm::function_trace", @@ -39,7 +43,34 @@ impl OcrClient { skip_all )] pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result { - perform_ocr_request(self, request).await + use super::{ + NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost, + OcrHostOperation, OcrHostResult, + }; + + let host = OcrHookHost::new(request.hooks.clone()); + let mut request = Some(request); + let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all()) + else { + return Err(Error::InvalidRequest( + "native OCR host admission declined".into(), + )); + }; + let mut result = None; + loop { + match call.resume(result.take()).await? { + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().ok_or_else(|| { + Error::InvalidRequest("OCR request was already projected".into()) + })?), + false, + )))) + } + OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await), + OcrCallStep::Complete(response) => return Ok(response), + } + } } pub(crate) fn provider_http(&self) -> &reqwest::Client { @@ -77,7 +108,7 @@ fn no_redirect_http() -> Result { .map_err(TransportError::from) } -pub async fn ocr(request: LiteLLMOcrRequest) -> Result { +pub(crate) fn shared_client() -> Result { static CLIENT: OnceLock> = OnceLock::new(); let client = CLIENT .get_or_init(|| { @@ -88,18 +119,50 @@ pub async fn ocr(request: LiteLLMOcrRequest) -> Result Result { + shared_client()?.perform(request).await } pub async fn read_json_response( response: reqwest::Response, native: bool, + max_response_bytes: usize, ) -> Result, OcrError> { + let bytes = read_response_bytes(response, max_response_bytes).await?; + Ok(decode_response(&bytes, native)?) +} + +pub(crate) async fn read_response_bytes( + mut response: reqwest::Response, + max_response_bytes: usize, +) -> Result { let status = response.status(); - let bytes = response - .bytes() - .await - .map_err(crate::error::TransportError::from)?; + let limit = if status.is_success() { + max_response_bytes + } else { + max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)) + }; + if status.is_success() + && response + .content_length() + .is_some_and(|length| length > limit as u64) + { + return Err(OcrResponseError::TooLarge { limit }.into()); + } + let mut bytes = BytesMut::new(); + while let Some(chunk) = response.chunk().await.map_err(transport_error)? { + let remaining = limit.saturating_sub(bytes.len()); + if status.is_success() && chunk.len() > remaining { + return Err(OcrResponseError::TooLarge { limit }.into()); + } + bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); + if !status.is_success() && bytes.len() == limit { + break; + } + } if !status.is_success() { return Err(crate::error::TransportError::Http { status: status.as_u16(), @@ -107,5 +170,41 @@ pub async fn read_json_response( } .into()); } - Ok(decode_response(&bytes, native)?) + Ok(bytes.freeze()) +} + +pub(crate) fn transport_error(error: reqwest::Error) -> Error { + if error.is_timeout() { + return Error::Http { + status: 408, + body: "OCR request timed out".into(), + }; + } + crate::error::TransportError::from(error).into() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn request_timeout_has_an_http_408_status() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let _connection = listener.accept().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + }); + let error = reqwest::Client::new() + .get(format!("http://{address}")) + .timeout(Duration::from_millis(10)) + .send() + .await + .unwrap_err(); + assert!(matches!( + transport_error(error), + Error::Http { status: 408, .. } + )); + server.abort(); + } } diff --git a/litellm-rust/crates/core/src/ocr/codecs/cohere.rs b/litellm-rust/crates/core/src/ocr/codecs/cohere.rs new file mode 100644 index 00000000000..649432f39d3 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/cohere.rs @@ -0,0 +1,254 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; + +use crate::ocr::document::InlineDocument; +use crate::ocr::error::{OcrRequestError, OcrResponseError}; +use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum OutputFormat { + #[default] + Markdown, + Blocks, +} + +#[derive(Deserialize)] +pub(crate) struct CohereParams { + #[serde(default)] + pub output_format: OutputFormat, +} + +#[derive(Deserialize, Serialize)] +pub(crate) struct CohereRequest { + pub model: String, + pub document: OcrDocument, + pub output_format: OutputFormat, +} + +pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), OcrRequestError> { + let OcrDocument::ImageUrl { image_url, .. } = document else { + return Err(OcrRequestError::CohereImageOnly); + }; + if image_url.is_empty() { + return Err(OcrRequestError::CohereImageOnly); + } + if let Some(inline) = InlineDocument::parse(image_url)? { + if !inline.mime_type().type_.eq_ignore_ascii_case("image") { + return Err(OcrRequestError::CohereImageOnly); + } + inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + } + Ok(()) +} + +#[derive(Deserialize)] +pub(crate) struct CohereResponse { + #[serde(default)] + pages: Vec, + meta: Option, +} + +#[derive(Deserialize)] +struct CoherePage { + index: Option, + markdown: Option, + blocks: Option>>, +} + +#[derive(Deserialize)] +struct CohereMarkdown { + #[serde(default)] + content: String, + images: Option>>, +} + +#[derive(Deserialize)] +struct CohereMeta { + billed_units: Option, +} + +#[derive(Deserialize)] +struct CohereBilledUnits { + pages: Option, +} + +pub(crate) fn transform_response( + model: &str, + response: CohereResponse, +) -> Result { + let pages_processed = response + .meta + .and_then(|meta| meta.billed_units) + .and_then(|units| units.pages) + .map(Ok) + .unwrap_or_else(|| { + i64::try_from(response.pages.len()).map_err(|_| OcrResponseError::NumericRange("pages")) + })?; + let pages = response + .pages + .into_iter() + .enumerate() + .map(|(position, page)| { + let index = page.index.map(Ok).unwrap_or_else(|| { + i64::try_from(position).map_err(|_| OcrResponseError::NumericRange("page index")) + })?; + let (content, images) = page + .markdown + .map(|markdown| { + let images = + markdown + .images + .filter(|images| !images.is_empty()) + .map(|images| { + images + .into_iter() + .map(|mut image| { + if let Some(Value::Object(bbox)) = + image.get("bounding_box").cloned() + { + image.insert("bbox".into(), Value::Object(bbox)); + } + Value::Object(image) + }) + .collect::>() + }); + (markdown.content, images) + }) + .unwrap_or_default(); + let mut normalized = json!({"index": index, "markdown": content, "images": images}); + if let Some(blocks) = page.blocks { + normalized["blocks"] = json!(blocks); + } + Ok(normalized) + }) + .collect::, OcrResponseError>>()?; + Ok(LiteLLMOcrResponse { + pages, + model: model.into(), + document_annotation: None, + usage_info: Some(json!({"pages_processed": pages_processed})), + object: "ocr".into(), + extra_fields: Map::new(), + provider_native_response: None, + }) +} + +pub(crate) fn transform_request( + model: &str, + document: OcrDocument, + params: CohereParams, +) -> Result { + validate_document(&document)?; + Ok(CohereRequest { + model: model.into(), + document, + output_format: params.output_format, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn response_normalizes_markdown_images_blocks_and_billed_pages() { + let response = serde_json::from_value(json!({ + "pages": [ + { + "type":"markdown", + "index":4, + "markdown":{ + "content":"receipt", + "images":[{ + "id":"image", + "bounding_box":{"top_left_x":1,"bottom_right_x":48}, + "bounding_box_normalized":{"top_left_x":0.04,"bottom_right_x":0.15}, + "description":"scan", + "category":"logo" + }] + } + }, + {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]} + ], + "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}} + })) + .unwrap(); + let normalized = transform_response("parse-v5.0", response).unwrap(); + assert_eq!(normalized.pages[0]["index"], 4); + assert_eq!(normalized.pages[0]["markdown"], "receipt"); + assert_eq!(normalized.pages[0]["images"][0]["bbox"]["top_left_x"], 1); + assert_eq!( + normalized.pages[0]["images"][0]["bounding_box_normalized"]["bottom_right_x"], + 0.15 + ); + assert_eq!(normalized.pages[0]["images"][0]["description"], "scan"); + assert_eq!(normalized.pages[0]["images"][0]["category"], "logo"); + assert_eq!(normalized.pages[1]["index"], 1); + assert_eq!(normalized.pages[1]["markdown"], ""); + assert_eq!(normalized.pages[1]["blocks"][0]["text"]["content"], "total"); + assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 3); + } + + #[test] + fn response_defaults_and_invalid_fields() { + for value in [ + json!({}), + json!({"meta":null}), + json!({"pages":[],"meta":{"billed_units":null}}), + ] { + let normalized = + transform_response("parse", serde_json::from_value(value).unwrap()).unwrap(); + assert!(normalized.pages.is_empty()); + assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 0); + } + for value in [ + json!({"pages":null}), + json!({"pages":[{"markdown":"text"}]}), + json!({"pages":[{"index":"bad"}]}), + ] { + assert!(serde_json::from_value::(value).is_err()); + } + let normalized = transform_response( + "parse", + serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(), + ) + .unwrap(); + assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 1); + assert!(normalized.pages[0]["images"].is_null()); + } + + #[test] + fn request_requires_image_and_supported_output_format() { + for value in [ + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + json!({"type":"image_url","image_url":""}), + json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}), + ] { + assert_eq!( + validate_document(&serde_json::from_value(value).unwrap()), + Err(OcrRequestError::CohereImageOnly) + ); + } + assert!(serde_json::from_value::(json!({"output_format":"html"})).is_err()); + for format in ["markdown", "blocks"] { + assert!( + serde_json::from_value::(json!({"output_format":format})).is_ok() + ); + } + let request = transform_request( + "parse-v5.0", + serde_json::from_value(json!({ + "type":"image_url", + "image_url":"https://example.com/image.png" + })) + .unwrap(), + serde_json::from_value(json!({})).unwrap(), + ) + .unwrap(); + assert_eq!( + serde_json::to_value(request).unwrap()["output_format"], + "markdown" + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs index 98cfc0db78d..7e8ce63b379 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs @@ -12,13 +12,17 @@ pub(crate) fn transform_ocr_request( params: &DeepSeekOcrParams, ) -> Result { if document.source().is_empty() { - return Err(OcrRequestError::MissingField("document URL")); + return Err(OcrRequestError::MissingDocumentUrl); } + let content = OcrDocument::ImageUrl { + image_url: document.source().to_string(), + extra_fields: serde_json::Map::new(), + }; Ok(DeepSeekOcrRequest { model: provider_model.to_string(), messages: vec![DeepSeekOcrMessage { role: UserRole::User, - content: vec![document], + content: vec![content], }], params: params.clone(), }) diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs index 85d1dafa542..9389f93b8e3 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs @@ -163,6 +163,30 @@ mod tests { ); } + #[rstest] + #[case(json!([0, 1, 2]), Some("1,2,3"))] + #[case(json!([2, 0, 0, 1]), Some("1,2,3"))] + #[case(json!([]), None)] + #[case(json!("3-9"), Some("3-9"))] + #[case(json!("1-3, 5"), Some("1-3,5"))] + #[case(json!(["1", "3-5"]), Some("1,3-5"))] + fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) { + assert_eq!( + map(json!({"pages": input})).unwrap().pages.as_deref(), + expected + ); + } + + #[rstest] + #[case(json!("a,b"))] + #[case(json!([-1]))] + #[case(json!([true, false]))] + #[case(json!([1, "2"]))] + #[case(json!(5))] + fn invalid_page_mapping_matches_python(#[case] input: Value) { + assert!(map(json!({"pages": input})).is_err()); + } + #[rstest] #[case(json!(["keyValuePairs"]), "keyValuePairs")] #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")] diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs index 2b848fcfb7a..f76a7c2b232 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs @@ -13,7 +13,7 @@ pub(crate) fn transform_ocr_request( ) -> Result { let source = document.source(); if source.is_empty() { - return Err(OcrRequestError::MissingField("document URL")); + return Err(OcrRequestError::MissingDocumentUrl); } Ok(if let Some(document) = InlineDocument::parse(source)? { DocumentIntelligenceRequest::Base64Source( @@ -46,10 +46,7 @@ pub(crate) fn transform_ocr_response( let mut extra_fields = Map::new(); extra_fields.insert("content".into(), option_value(result.content)); extra_fields.insert("tables".into(), option_value(result.tables)); - extra_fields.insert( - "key_value_pairs".into(), - option_value(result.key_value_pairs), - ); + extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs)); Ok(LiteLLMOcrResponse { pages, model: model.into(), diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs index 5bd7e555a1e..e60f1f5d3d6 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs @@ -114,6 +114,7 @@ mod tests { #[rstest] #[case("table_format", json!("html"))] #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("block"))] #[case("document_annotation_prompt", json!("extract"))] #[case("include_blocks", json!(true))] #[case("id", json!("req-123"))] @@ -133,6 +134,7 @@ mod tests { #[rstest] #[case("pages", json!([0, 2]))] + #[case("pages", json!("0,2-4"))] #[case("include_image_base64", json!(true))] #[case("image_limit", json!(2))] #[case("image_min_size", json!(100))] @@ -196,8 +198,16 @@ mod tests { #[rstest] fn transform_ocr_response_preserves_blocks_and_confidence_scores() { let response: MistralOcrResponse = serde_json::from_value(json!({ - "pages":[{"index":0,"markdown":"hello","blocks":[{"type":"title"}],"confidence_scores":{"mean":0.99}}], + "pages":[{ + "index":0, + "markdown":"hello", + "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], + "dimensions":{"width":612,"height":792,"dpi":72}, + "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], + "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} + }], "model":"returned-model", + "document_annotation":"{\"language\":\"en\"}", "usage_info":{"pages_processed":1} })) .unwrap(); @@ -205,7 +215,20 @@ mod tests { .unwrap() .into_json(); assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); - assert_eq!(result["pages"][0]["confidence_scores"]["mean"], 0.99); + assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1); + assert_eq!( + result["pages"][0]["blocks"][0]["confidence_scores"]["mean"], + 0.98 + ); + assert_eq!( + result["pages"][0]["confidence_scores"]["average_page_confidence_score"], + 0.99 + ); + assert_eq!(result["pages"][0]["images"][0]["id"], "img-0"); + assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72); + assert_eq!(result["model"], "returned-model"); + assert_eq!(result["document_annotation"], "{\"language\":\"en\"}"); + assert_eq!(result["usage_info"]["pages_processed"], 1); } #[rstest] diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs index 0e601cd8319..e0bc8a267d2 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs @@ -3,10 +3,17 @@ use serde_json::{Map, Value}; use crate::ocr::types::OcrDocument; +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum MistralOcrPages { + Range(String), + Indices(Vec), +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub(crate) struct MistralOcrParams { #[serde(skip_serializing_if = "Option::is_none")] - pub pages: Option>, + pub pages: Option, #[serde(skip_serializing_if = "Option::is_none")] pub include_image_base64: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/litellm-rust/crates/core/src/ocr/codecs/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mod.rs index 7c752749901..639b985b9ae 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mod.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod cohere; pub(crate) mod deepseek; pub(crate) mod document_intelligence; pub(crate) mod mistral; diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index e89b1c5c569..82a32ac1ab5 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -2,13 +2,90 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::mime::Mime; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; use reqwest::Url; +use serde_json::Map; use super::error::{OcrError, OcrRequestError, OcrResponseError}; use super::types::{OcrConnection, OcrDocument}; -use crate::constants::OCR_MAX_FETCH_REDIRECTS; +use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS}; use crate::error::{MediaError, TransportError}; use crate::media::{DownloadPolicy, MediaFetcher}; +pub fn encode_file_document( + bytes: &[u8], + file_name: Option<&str>, + mime_type: Option<&str>, +) -> Result { + if bytes.is_empty() { + return Err(OcrRequestError::EmptyFile); + } + if bytes.len() > OCR_INLINE_MAX_BYTES { + return Err(OcrRequestError::InlineDocumentTooLarge); + } + if let Some(value) = mime_type + && !valid_mime_type(value) + { + return Err(OcrRequestError::InvalidMimeType(value.into())); + } + let mime_type = mime_type + .map(str::to_string) + .or_else(|| file_name.map(|name| mime_type_for_name(name).to_string())) + .unwrap_or_else(|| "application/octet-stream".into()); + let source = format!("data:{mime_type};base64,{}", STANDARD.encode(bytes)); + Ok(if mime_type.starts_with("image/") { + OcrDocument::ImageUrl { + image_url: source, + extra_fields: Map::new(), + } + } else { + OcrDocument::DocumentUrl { + document_url: source, + extra_fields: Map::new(), + } + }) +} + +fn valid_mime_type(value: &str) -> bool { + let Some((kind, subtype)) = value.split_once('/') else { + return false; + }; + !kind.is_empty() + && !subtype.is_empty() + && kind.chars().chain(subtype.chars()).all(|character| { + character.is_alphanumeric() || matches!(character, '.' | '+' | '-' | '_') + }) +} + +pub fn mime_type_for_name(name: &str) -> &'static str { + let extension = std::path::Path::new(name) + .extension() + .and_then(|value| value.to_str()) + .unwrap_or_default(); + match extension.to_ascii_lowercase().as_str() { + "pdf" => "application/pdf", + "png" => "image/png", + "jpg" | "jpeg" => "image/jpeg", + "gif" => "image/gif", + "webp" => "image/webp", + "tiff" | "tif" => "image/tiff", + "bmp" => "image/bmp", + _ => mime_guess::from_path(name) + .first_raw() + .unwrap_or("application/octet-stream"), + } +} + +pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str { + match content_type + .and_then(|value| value.split(';').next()) + .map(str::trim) + { + Some(value) if !value.is_empty() && value != "application/octet-stream" => value, + _ => file_name + .map(mime_type_for_name) + .unwrap_or("application/octet-stream"), + } +} + pub(crate) struct InlineDocument<'a>(DataUrl<'a>); impl<'a> InlineDocument<'a> { @@ -95,9 +172,11 @@ fn map_media_error(error: MediaError) -> OcrError { body: "OCR document download failed".into(), } .into(), - MediaError::Timeout => { - TransportError::Network("OCR document download timed out".into()).into() + MediaError::Timeout => TransportError::Http { + status: 408, + body: "OCR document download timed out".into(), } + .into(), MediaError::Transport(error) => error.into(), } } @@ -114,6 +193,90 @@ mod tests { } } + #[test] + fn file_bytes_are_encoded_with_core_owned_mime_policy() { + assert_eq!( + encode_file_document(b"abc", Some("scan.png"), None).unwrap(), + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,YWJj".into(), + extra_fields: Map::new(), + } + ); + assert_eq!( + encode_file_document(b"abc", None, Some("application/pdf")).unwrap(), + document("data:application/pdf;base64,YWJj") + ); + } + + #[test] + fn file_name_mime_mapping_matches_python() { + for (name, expected) in [ + ("document.pdf", "application/pdf"), + ("image.png", "image/png"), + ("photo.jpg", "image/jpeg"), + ("photo.jpeg", "image/jpeg"), + ("animation.gif", "image/gif"), + ("image.webp", "image/webp"), + ("scan.tiff", "image/tiff"), + ("scan.tif", "image/tiff"), + ("bitmap.bmp", "image/bmp"), + ("DOCUMENT.PDF", "application/pdf"), + ("IMAGE.PNG", "image/png"), + ("file.unknown-extension", "application/octet-stream"), + ] { + assert_eq!(mime_type_for_name(name), expected); + } + } + + #[test] + fn upload_mime_mapping_matches_python() { + assert_eq!( + upload_mime_type(Some("report.pdf"), Some("application/octet-stream")), + "application/pdf" + ); + assert_eq!(upload_mime_type(Some("image.png"), None), "image/png"); + assert_eq!(upload_mime_type(None, None), "application/octet-stream"); + assert_eq!( + upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")), + "application/pdf" + ); + assert_eq!( + upload_mime_type( + Some("img.png"), + Some("image/png; charset=utf-8; boundary=something") + ), + "image/png" + ); + } + + #[test] + fn file_encoding_enforces_decoded_size_limit() { + let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1]; + assert_eq!( + encode_file_document(&bytes, None, None), + Err(OcrRequestError::InlineDocumentTooLarge) + ); + let document = encode_file_document(&bytes[..OCR_INLINE_MAX_BYTES], None, None).unwrap(); + let inline = InlineDocument::parse(document.source()).unwrap().unwrap(); + assert_eq!( + inline.decode(OCR_INLINE_MAX_BYTES).unwrap(), + bytes[..OCR_INLINE_MAX_BYTES] + ); + } + + #[test] + fn file_encoding_rejects_empty_bytes_and_invalid_explicit_mime() { + assert!(encode_file_document(b"", None, None).is_err()); + for mime in [ + "text/plain;bad", + "text/plain/extra", + " text/plain", + "text/plain\n", + ] { + assert!(encode_file_document(b"abc", None, Some(mime)).is_err()); + } + } + #[test] fn decodes_data_urls_and_limits_decoded_size() { for (source, expected) in [ diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 522d059ec48..55ea2cbcdae 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -4,15 +4,27 @@ use crate::error::TransportError; #[derive(Debug, Clone, PartialEq, Eq, Error)] pub enum OcrRequestError { + #[error("File is empty or could not be read")] + EmptyFile, + #[error("Invalid MIME type: {0}")] + InvalidMimeType(String), + #[error( + "Cohere Parse only accepts `image_url` documents; document_url and PDF inputs are not supported" + )] + CohereImageOnly, #[error("Invalid `req_format`. Expected 'native' or 'litellm'.")] RequestFormat, #[error("invalid OCR request field: {path}")] RequestField { path: String }, #[error("missing required field: {0}")] MissingField(&'static str), + #[error("Document URL is required")] + MissingDocumentUrl, #[error("invalid OCR document data URI")] InvalidDataUri, - #[error("Reducto requires a reducto:// id or a data URI")] + #[error( + "Reducto requires a reducto:// id or a data URI; plain HTTP URLs are not supported, upload the file first" + )] ReductoSource, #[error("inline OCR document exceeds the size limit")] InlineDocumentTooLarge, @@ -34,6 +46,8 @@ pub enum OcrRequestError { #[derive(Debug, Clone, PartialEq, Eq, Error)] pub enum OcrResponseError { + #[error("OCR response exceeds the size limit of {limit} bytes")] + TooLarge { limit: usize }, #[error("invalid OCR response field: {path}")] ResponseField { path: String }, #[error("OCR response is missing non-empty content")] diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 0b04319d966..cd1d538aaa8 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,15 +1,17 @@ use super::OcrClient; use super::adapters::OcrAdapter; -use super::hooks::OcrLifecycleHooks; +use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest}; use super::registry::OcrAdapterKind; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; use crate::Error; use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext}; +use std::sync::Arc; pub(crate) async fn perform_ocr_request( client: &OcrClient, request: LiteLLMOcrRequest, ) -> Result { + request.response_format()?; let context = CallLifecycleContext::new( "ocr", request.model.clone(), @@ -23,27 +25,71 @@ pub(crate) async fn perform_ocr_request( hooks: request.hooks.clone(), provider_name: context.custom_llm_provider.clone(), }; - CallLifecycle::default().run(context, request, &hooks, |request| async move { - macro_rules! execute_selected_adapter { + CallLifecycle::default() + .run(context, request, &hooks, |request| async move { + PreparedOcrCall::prepare(client.clone(), request) + .await? + .execute() + .await? + .normalize() + }) + .await +} + +pub(crate) struct PreparedOcrCall { + client: OcrClient, + request: LiteLLMOcrRequest, + http: reqwest::Request, +} + +impl PreparedOcrCall { + pub(crate) async fn prepare( + client: OcrClient, + request: LiteLLMOcrRequest, + ) -> Result { + macro_rules! prepare_adapter { ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { match request.adapter { - $( OcrAdapterKind::$variant => execute_ocr_provider_call(client, &$instance, request).await, )+ + $( OcrAdapterKind::$variant => $instance.prepare_request(&request, &client).await?, )+ } }; } - super::adapters::for_each_ocr_adapter!(execute_selected_adapter) - }).await + let http = super::adapters::for_each_ocr_adapter!(prepare_adapter); + Ok(Self { + client, + request, + http, + }) + } + + pub(crate) async fn execute(self) -> Result { + let url = self.http.url().to_string(); + let headers = request_headers(&self.http)?; + let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts( + self.client.provider_http().clone(), + self.http, + )) + .await + .map_err(super::client::transport_error)?; + macro_rules! read_adapter { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + match self.request.adapter { + $( OcrAdapterKind::$variant => { + let decoded = $instance.read_response(&self.client, response, &url, &headers, &self.request).await?; + Ok(OcrProviderResponse { + request: self.request, + data: OcrProviderData::$variant(decoded), + }) + }, )+ + } + }; + } + super::adapters::for_each_ocr_adapter!(read_adapter) + } } -#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -async fn execute_ocr_provider_call( - client: &OcrClient, - adapter: &A, - request: LiteLLMOcrRequest, -) -> Result { - let provider_request = adapter.prepare_request(&request, client).await?; - let url = provider_request.url().to_string(); - let headers = provider_request +fn request_headers(request: &reqwest::Request) -> Result, Error> { + request .headers() .iter() .map(|(name, value)| { @@ -53,20 +99,41 @@ async fn execute_ocr_provider_call( .map_err(|_| super::error::OcrRequestError::RequestField { path: "headers".into(), }) + .map_err(Error::from) }) - .collect::, _>>()?; - let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts( - client.provider_http().clone(), - provider_request, - )) - .await - .map_err(crate::error::TransportError::from)?; - let decoded = adapter - .read_response(client, response, &url, &headers, &request) - .await?; - let response = adapter.transform_ocr_response(&request, decoded.data)?; - Ok(LiteLLMOcrResponse { - provider_native_response: decoded.native, - ..response - }) + .collect() } + +macro_rules! provider_data { + ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { + enum OcrProviderData { + $( $variant(super::wire::DecodedOcrResponse<<$adapter as OcrAdapter>::ProviderResponse>), )+ + } + + impl OcrProviderResponse { + pub(crate) fn normalize(self) -> Result { + match self.data { + $( OcrProviderData::$variant(decoded) => { + let response = $instance.transform_ocr_response(&self.request, decoded.data)?; + Ok(LiteLLMOcrResponse { provider_native_response: decoded.native, ..response }) + }, )+ + } + } + } + }; +} + +pub(crate) struct OcrProviderResponse { + request: LiteLLMOcrRequest, + data: OcrProviderData, +} + +pub(crate) async fn post_call(hooks: &Arc, bytes: &[u8]) -> Result<(), Error> { + let original_response = serde_json::Value::String(String::from_utf8_lossy(bytes).into_owned()); + hooks + .post_call(OcrPostCallRequest { original_response }) + .await?; + Ok(()) +} + +super::adapters::for_each_ocr_adapter!(provider_data); diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 7dd3c6bf8b2..3e7507e9ed5 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -24,11 +24,19 @@ pub struct OcrDuringCallRequest { pub model: String, pub custom_llm_provider: String, pub url: String, + pub headers: Vec<(String, String)>, pub body: Value, + #[serde(skip)] + pub retained_fields: Vec, +} + +#[derive(Clone, Debug, Serialize)] +pub struct OcrPostCallRequest { + pub original_response: Value, } pub trait OcrHooks: Send + Sync { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { false } fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { @@ -40,6 +48,9 @@ pub trait OcrHooks: Send + Sync { ) -> OcrHookFuture<'_, OcrDuringCallRequest> { Box::pin(async move { Ok(request) }) } + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { Ok(request) }) + } fn success<'a>( &'a self, _context: &'a CallLifecycleContext, @@ -80,7 +91,7 @@ impl CallLifecycleHooks Self::PreCallFuture<'a> { Box::pin(async move { - if !self.hooks.has_guardrails() { + if !self.hooks.intercepts_requests() { return Ok(request); } let changed = self diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs new file mode 100644 index 00000000000..92c9d4b717c --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -0,0 +1,640 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use tokio::sync::{mpsc, oneshot}; + +use super::handler::perform_ocr_request; +use super::hooks::{ + OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, + OcrPreCallRequest, +}; +use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; +use crate::AuthError; +use crate::Error; +use crate::auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use crate::call_lifecycle::host::{ + HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase, +}; +use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; + +pub type NativeResult = Result, Error>; + +#[derive(Debug, PartialEq, Eq)] +pub enum NativeOutcome { + Completed(T), + Declined(OcrDecline), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrDecline { + ProviderWorkflow, + HostOperations, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OcrAdmission { + pub provider_workflow: bool, + pub host_operations: bool, + pub asynchronous: bool, +} + +impl OcrAdmission { + pub const fn all() -> Self { + Self { + provider_workflow: true, + host_operations: true, + asynchronous: false, + } + } +} + +#[derive(Clone, Debug)] +pub enum OcrHostOperation { + ProjectRequest, + Lifecycle(HostPhase), + ConstructResponse(Arc), + MapFailure(Error), + Success { + context: CallLifecycleContext, + response: Arc, + timing: CallLifecycleTiming, + }, + Failure { + context: CallLifecycleContext, + error: Error, + timing: CallLifecycleTiming, + }, + AcquireAzureAdToken, + PreCall(OcrPreCallRequest), + DuringCall(OcrDuringCallRequest), + PostCall(OcrPostCallRequest), +} + +impl OcrHostOperation { + pub const fn phase(&self) -> Option { + match self { + Self::Lifecycle(phase) => Some(*phase), + Self::Success { .. } => Some(HostPhase::Success), + Self::Failure { .. } => Some(HostPhase::Failure), + _ => None, + } + } +} + +pub enum OcrHostResult { + Request(Result<(Box, bool), Error>), + Lifecycle(Result<(), HostFailure>), + AzureAdToken(Result), + PreCall(Result), + DuringCall(Result), + PostCall(Result), +} + +pub type OcrCallStep = HostCallStep; + +pub struct OcrCall { + lifecycle: HostLifecycle, + execution: OcrExecution, + response: Option>, + error: Option, + pending: bool, + completed: bool, + projecting: bool, +} + +impl OcrCall { + pub fn admit(client: OcrClient, admission: OcrAdmission) -> NativeOutcome { + if !admission.provider_workflow { + return NativeOutcome::Declined(OcrDecline::ProviderWorkflow); + } + if !admission.host_operations { + return NativeOutcome::Declined(OcrDecline::HostOperations); + } + NativeOutcome::Completed(Self { + lifecycle: HostLifecycle::new(admission.asynchronous), + execution: OcrExecution::new(client), + response: None, + error: None, + pending: false, + completed: false, + projecting: false, + }) + } + + pub async fn resume(&mut self, result: Option) -> Result { + if self.completed { + return Err(Error::InvalidRequest( + "OCR call cannot be resumed after completion".into(), + )); + } + if self.pending != result.is_some() { + return Err(Error::InvalidRequest( + "OCR host operation result does not match pending state".into(), + )); + } + match &result { + Some(OcrHostResult::Lifecycle(Ok(()))) + if self.lifecycle.phase() == HostPhase::Execute => + { + return Err(Error::InvalidRequest( + "OCR provider operation requires a typed result".into(), + )); + } + Some(result) + if !matches!(result, OcrHostResult::Lifecycle(_)) + && self.lifecycle.phase() != HostPhase::Execute => + { + return Err(Error::InvalidRequest( + "unexpected OCR provider operation result".into(), + )); + } + _ => {} + } + self.pending = false; + let provider_result = match result { + Some(OcrHostResult::Request(result)) if self.projecting => { + self.projecting = false; + match result { + Ok((request, azure_ad_token_provider)) => { + self.execution.request = Some(*request); + self.execution.azure_ad_token_provider = azure_ad_token_provider; + } + Err(error) => self.accept(Err(HostFailure::Error(error))), + } + None + } + Some(OcrHostResult::Request(_)) => { + return Err(Error::InvalidRequest( + "unexpected OCR request projection".into(), + )); + } + Some(OcrHostResult::Lifecycle(result)) => { + self.accept(result); + None + } + result => result, + }; + if self.lifecycle.phase() == HostPhase::Execute { + if self.execution.request.is_none() + && self.execution.execution.is_none() + && !self.execution.completed + { + self.projecting = true; + return Ok(self.host_step(OcrHostOperation::ProjectRequest)); + } + match self.execution.resume(provider_result).await { + Ok(OcrCallStep::Host(operation)) => return Ok(self.host_step(operation)), + Ok(OcrCallStep::Complete(response)) => { + self.response = Some(Arc::new(response)); + self.accept(Ok(())); + } + Err(error) => self.accept(Err(HostFailure::Error(error))), + } + } + if self.error.is_some() { + self.execution.stop().await; + } + let operation = match self.lifecycle.phase() { + HostPhase::Complete => { + self.completed = true; + return match self.error.take() { + Some(error) => Err(error), + None => self + .response + .take() + .map(Arc::unwrap_or_clone) + .map(OcrCallStep::Complete) + .ok_or_else(|| { + Error::InvalidRequest("OCR completed without a response".into()) + }), + }; + } + HostPhase::ConstructResponse => OcrHostOperation::ConstructResponse( + self.response + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))? + .clone(), + ), + HostPhase::MapFailure => OcrHostOperation::MapFailure( + self.error + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))? + .clone(), + ), + HostPhase::Success | HostPhase::Failure => { + let snapshot = self + .execution + .terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone(); + match (self.lifecycle.phase(), snapshot) { + (HostPhase::Success, Some((context, timing))) => OcrHostOperation::Success { + context, + response: self + .response + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR response".into()))? + .clone(), + timing, + }, + (HostPhase::Failure, Some((context, timing))) => OcrHostOperation::Failure { + context, + error: self + .error + .as_ref() + .ok_or_else(|| Error::InvalidRequest("missing OCR failure".into()))? + .clone(), + timing, + }, + (phase, _) => OcrHostOperation::Lifecycle(phase), + } + } + phase => OcrHostOperation::Lifecycle(phase), + }; + Ok(self.host_step(operation)) + } + + fn accept(&mut self, result: Result<(), HostFailure>) { + let cancelled = matches!(&result, Err(HostFailure::Cancelled(_))); + if let Some(error) = self.lifecycle.accept(result) { + if cancelled { + self.error = Some(error); + } else { + self.error.get_or_insert(error); + } + self.execution.cancel(); + } + } + + pub async fn interrupt(&mut self, failure: HostFailure) -> Result { + if self.completed { + return Err(Error::InvalidRequest( + "OCR call cannot be interrupted after completion".into(), + )); + } + self.pending = false; + self.accept(Err(failure)); + self.resume(None).await + } + + fn host_step(&mut self, operation: OcrHostOperation) -> OcrCallStep { + self.pending = true; + OcrCallStep::Host(operation) + } +} + +impl HostCall for OcrCall { + type Operation = OcrHostOperation; + type Result = OcrHostResult; + type Complete = LiteLLMOcrResponse; + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(OcrCall::resume(self, result)) + } + + fn interrupt( + &mut self, + failure: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(OcrCall::interrupt(self, failure)) + } +} + +struct PendingOperation { + operation: OcrHostOperation, + result: oneshot::Sender, +} + +struct OcrExecution { + client: Option, + request: Option, + operations_tx: mpsc::UnboundedSender, + operations_rx: mpsc::UnboundedReceiver, + pending_result: Option>, + execution: Option>>, + completed: bool, + azure_ad_token_provider: bool, + terminal: Arc>>, +} + +impl OcrExecution { + fn new(client: OcrClient) -> Self { + let (operations_tx, operations_rx) = mpsc::unbounded_channel(); + Self { + client: Some(client), + request: None, + operations_tx, + operations_rx, + pending_result: None, + execution: None, + completed: false, + azure_ad_token_provider: false, + terminal: Arc::default(), + } + } + + pub async fn resume(&mut self, result: Option) -> Result { + if self.completed { + return Err(Error::InvalidRequest( + "OCR call cannot be resumed after completion".into(), + )); + } + match (self.pending_result.take(), result) { + (Some(sender), Some(result)) => sender + .send(result) + .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into()))?, + (None, None) if self.execution.is_none() => self.start(), + (Some(sender), None) => { + self.pending_result = Some(sender); + return Err(Error::InvalidRequest( + "OCR host operation result is required".into(), + )); + } + (None, Some(_)) => { + return Err(Error::InvalidRequest( + "unexpected OCR host operation result".into(), + )); + } + (None, None) => {} + } + + let execution = self.execution.as_mut().ok_or_else(|| { + Error::InvalidRequest("OCR call cannot be resumed after completion".into()) + })?; + tokio::select! { + operation = self.operations_rx.recv() => { + let operation = operation.ok_or_else(|| Error::InvalidRequest("OCR operation channel closed".into()))?; + self.pending_result = Some(operation.result); + Ok(OcrCallStep::Host(operation.operation)) + } + result = execution => { + self.execution = None; + self.completed = true; + result + .map_err(|error| Error::Network(format!("OCR execution task failed: {error}")))? + .map(OcrCallStep::Complete) + } + } + } + + fn start(&mut self) { + let client = self.client.take().expect("admitted OCR call has a client"); + let mut request = self + .request + .take() + .expect("admitted OCR call has a request"); + let intercepts_requests = request.hooks.intercepts_requests(); + if self.azure_ad_token_provider { + request.azure_ad_token_provider = Some(TokenProviderHandle::new(Arc::new( + OcrAzureAdTokenProvider { + operations: self.operations_tx.clone(), + }, + ))); + } + request.hooks = Arc::new(ProtocolHooks { + operations: self.operations_tx.clone(), + intercepts_requests, + terminal: self.terminal.clone(), + }); + self.execution = Some(tokio::spawn(async move { + perform_ocr_request(&client, request).await + })); + } + + fn cancel(&mut self) { + self.pending_result = None; + if let Some(execution) = &self.execution { + execution.abort(); + } + } + + async fn stop(&mut self) { + self.cancel(); + if let Some(execution) = self.execution.as_mut() { + let _ = execution.await; + } + self.execution = None; + } +} + +impl Drop for OcrExecution { + fn drop(&mut self) { + if let Some(execution) = &self.execution { + execution.abort(); + } + } +} + +struct ProtocolHooks { + operations: mpsc::UnboundedSender, + intercepts_requests: bool, + terminal: Arc>>, +} + +#[derive(Debug)] +struct OcrAzureAdTokenProvider { + operations: mpsc::UnboundedSender, +} + +impl TokenProvider for OcrAzureAdTokenProvider { + fn acquire(&self) -> TokenFuture<'_> { + Box::pin(async move { + let (result, receiver) = oneshot::channel(); + self.operations + .send(PendingOperation { + operation: OcrHostOperation::AcquireAzureAdToken, + result, + }) + .map_err(|_| { + AuthError::AzureTokenAcquisition("OCR host driver was abandoned".into()) + })?; + match receiver.await.map_err(|_| { + AuthError::AzureTokenAcquisition( + "OCR token provider operation was abandoned".into(), + ) + })? { + OcrHostResult::AzureAdToken(result) => result, + _ => Err(AuthError::AzureTokenAcquisition( + "invalid OCR token provider host result".into(), + )), + } + }) + } +} + +impl ProtocolHooks { + async fn invoke(&self, operation: OcrHostOperation) -> Result { + let (result, receiver) = oneshot::channel(); + self.operations + .send(PendingOperation { operation, result }) + .map_err(|_| Error::InvalidRequest("OCR host driver was abandoned".into()))?; + receiver + .await + .map_err(|_| Error::InvalidRequest("OCR host operation was abandoned".into())) + } +} + +impl OcrHooks for ProtocolHooks { + fn intercepts_requests(&self) -> bool { + self.intercepts_requests + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { + match self.invoke(OcrHostOperation::PreCall(request)).await? { + OcrHostResult::PreCall(result) => result, + _ => Err(Error::InvalidRequest( + "invalid OCR pre-call host result".into(), + )), + } + }) + } + + fn during_call( + &self, + request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + match self.invoke(OcrHostOperation::DuringCall(request)).await? { + OcrHostResult::DuringCall(result) => result, + _ => Err(Error::InvalidRequest( + "invalid OCR during-call host result".into(), + )), + } + }) + } + + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + match self.invoke(OcrHostOperation::PostCall(request)).await? { + OcrHostResult::PostCall(result) => result, + _ => Err(Error::InvalidRequest( + "invalid OCR post-call host result".into(), + )), + } + }) + } + + fn success<'a>( + &'a self, + context: &'a CallLifecycleContext, + _response: &'a LiteLLMOcrResponse, + timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + *self + .terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some((context.clone(), timing.clone())); + }) + } + + fn failure<'a>( + &'a self, + context: &'a CallLifecycleContext, + _error: &'a Error, + timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + *self + .terminal + .lock() + .unwrap_or_else(|error| error.into_inner()) = + Some((context.clone(), timing.clone())); + }) + } +} + +pub type OcrHostFuture<'a> = Pin + Send + 'a>>; + +pub trait OcrHost: Send + Sync { + fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_>; +} + +pub struct NoopOcrHost; + +impl OcrHost for NoopOcrHost { + fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> { + Box::pin(async move { + match operation { + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( + Error::InvalidRequest("OCR host has no request projection".into()), + )), + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => OcrHostResult::Lifecycle(Ok(())), + OcrHostOperation::AcquireAzureAdToken => { + OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition( + "OCR host has no Azure AD token provider".into(), + ))) + } + OcrHostOperation::PreCall(request) => OcrHostResult::PreCall(Ok(request)), + OcrHostOperation::DuringCall(request) => OcrHostResult::DuringCall(Ok(request)), + OcrHostOperation::PostCall(request) => OcrHostResult::PostCall(Ok(request)), + } + }) + } +} + +pub struct OcrHookHost { + hooks: Arc, +} + +impl OcrHookHost { + pub fn new(hooks: Arc) -> Self { + Self { hooks } + } +} + +impl OcrHost for OcrHookHost { + fn invoke(&self, operation: OcrHostOperation) -> OcrHostFuture<'_> { + Box::pin(async move { + match operation { + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( + Error::InvalidRequest("OCR hook host has no request projection".into()), + )), + OcrHostOperation::Success { + context, + response, + timing, + } => { + self.hooks.success(&context, &response, &timing).await; + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Failure { + context, + error, + timing, + } => { + self.hooks.failure(&context, &error, &timing).await; + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) => OcrHostResult::Lifecycle(Ok(())), + OcrHostOperation::AcquireAzureAdToken => { + OcrHostResult::AzureAdToken(Err(AuthError::AzureTokenAcquisition( + "OCR hook host has no Azure AD token provider".into(), + ))) + } + OcrHostOperation::PreCall(request) => { + OcrHostResult::PreCall(self.hooks.pre_call(request).await) + } + OcrHostOperation::DuringCall(request) => { + OcrHostResult::DuringCall(self.hooks.during_call(request).await) + } + OcrHostOperation::PostCall(request) => { + OcrHostResult::PostCall(self.hooks.post_call(request).await) + } + } + }) + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 1e975c3f521..e29fd6ac572 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -5,12 +5,18 @@ mod document; pub mod error; mod handler; pub mod hooks; +mod lifecycle; mod prepare; mod registry; pub mod types; pub mod wire; pub use client::{OcrClient, ocr}; +pub use document::{encode_file_document, mime_type_for_name, upload_mime_type}; +pub use lifecycle::{ + NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, + OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, +}; pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; #[cfg(test)] diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index bf6f924088c..9934a1d9a14 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -62,34 +62,48 @@ pub(crate) async fn transform_request_body( request: &LiteLLMOcrRequest, url: &str, headers: &[(String, String)], + retains_document: bool, body: B, validate: impl FnOnce(&B) -> Result<(), OcrRequestError>, ) -> Result where B: Serialize + DeserializeOwned, { - let body = if request.hooks.has_guardrails() { + let (body, headers) = if request.hooks.intercepts_requests() { + let body = serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { + path: "body".into(), + })?; + let retained_fields = request + .optional_params + .keys() + .filter(|name| body.get(*name).is_some()) + .cloned() + .chain(retains_document.then(|| "document".to_string())) + .collect(); let changed = request .hooks .during_call(OcrDuringCallRequest { model: request.model.clone(), custom_llm_provider: request.adapter.provider().as_str().into(), url: url.into(), - body: serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { - path: "body".into(), - })?, + headers: headers.to_vec(), + body, + retained_fields, }) .await?; let body = OcrWireBody::::decode(changed.body)?; validate(&body.body)?; - body + (body, changed.headers) } else { - OcrWireBody { - body, - extra: Map::new(), - } + ( + OcrWireBody { + body, + extra: Map::new(), + }, + headers.to_vec(), + ) }; - build_http_request(client, request, url, headers, &body) + build_http_request(client, request, url, &headers, &body) } pub(crate) fn build_http_request( @@ -113,9 +127,10 @@ pub(crate) fn build_http_request( pub(crate) async fn guardrail_document( request: &LiteLLMOcrRequest, url: &str, -) -> Result { - if !request.hooks.has_guardrails() { - return Ok(request.document.clone()); + headers: &[(String, String)], +) -> Result<(OcrDocument, Vec<(String, String)>), OcrError> { + if !request.hooks.intercepts_requests() { + return Ok((request.document.clone(), headers.to_vec())); } let changed = request .hooks @@ -123,14 +138,17 @@ pub(crate) async fn guardrail_document( model: request.model.clone(), custom_llm_provider: request.adapter.provider().as_str().into(), url: url.into(), + headers: headers.to_vec(), body: serde_json::to_value(&request.document).map_err(|_| { OcrRequestError::RequestField { path: "document".into(), } })?, + retained_fields: Vec::new(), }) .await?; - super::wire::decode_request_value(changed.body, "guardrail.document").map_err(OcrError::from) + let document = super::wire::decode_request_value(changed.body, "guardrail.document")?; + Ok((document, changed.headers)) } #[derive(Serialize)] diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs index 1b20a91143b..ed7d4fd5cf2 100644 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -23,6 +23,7 @@ super::adapters::for_each_ocr_adapter!(define_adapter_types); #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum OcrProvider { + Cohere, Mistral, AzureAi, Reducto, @@ -32,6 +33,7 @@ pub(crate) enum OcrProvider { impl OcrProvider { pub(crate) const fn as_str(self) -> &'static str { match self { + Self::Cohere => "cohere", Self::Mistral => "mistral", Self::AzureAi => "azure_ai", Self::Reducto => "reducto", @@ -50,6 +52,7 @@ pub(crate) fn resolve_wire_adapter( custom_llm_provider: OcrProvider::Mistral.as_str(), }); let typed_provider = match provider.custom_llm_provider { + "cohere" => OcrProvider::Cohere, "mistral" => OcrProvider::Mistral, "azure_ai" => OcrProvider::AzureAi, "reducto" => OcrProvider::Reducto, @@ -57,10 +60,17 @@ pub(crate) fn resolve_wire_adapter( value => return Err(Error::InvalidProvider(value.to_string())), }; let adapter = match typed_provider { + OcrProvider::Cohere => OcrAdapterKind::Cohere, OcrProvider::Mistral => OcrAdapterKind::Mistral, OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { OcrAdapterKind::AzureDocumentIntelligence } + OcrProvider::AzureAi + if provider.model.to_ascii_lowercase().contains("cohere") + && provider.model.to_ascii_lowercase().contains("parse") => + { + OcrAdapterKind::AzureCohere + } OcrProvider::AzureAi => OcrAdapterKind::AzureMistral, OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => { OcrAdapterKind::ReductoLegacy @@ -68,12 +78,7 @@ pub(crate) fn resolve_wire_adapter( OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-v3") => { OcrAdapterKind::ReductoV3 } - OcrProvider::Reducto => { - return Err(Error::InvalidRequest(format!( - "unsupported Reducto OCR model: {}", - provider.model - ))); - } + OcrProvider::Reducto => OcrAdapterKind::ReductoV3, OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { OcrAdapterKind::VertexDeepSeek } @@ -107,11 +112,10 @@ mod tests { } #[test] - fn unknown_reducto_models_are_rejected() { - assert!(matches!( - resolve_wire_adapter("reducto/future-parse-model", None), - Err(Error::InvalidRequest(_)) - )); + fn unknown_reducto_models_use_the_current_protocol() { + let (model, adapter) = resolve_wire_adapter("reducto/future-parse-model", None).unwrap(); + assert_eq!(model, "future-parse-model"); + assert_eq!(adapter, OcrAdapterKind::ReductoV3); } #[test] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 06519f86c91..76df8b42806 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -8,7 +8,7 @@ use serde_json::{Map, Value}; use super::hooks::{NoopOcrHooks, OcrHooks}; use super::registry::{OcrAdapterKind, resolve_wire_adapter}; use crate::Error; -use crate::auth::InputSource; +use crate::auth::{InputSource, TokenProviderHandle}; use crate::constants::OCR_HTTP_TIMEOUT_SECS; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -68,6 +68,7 @@ pub struct OcrConnection { pub extra_headers_source: InputSource, pub timeout: Duration, pub max_download_bytes: u64, + pub max_response_bytes: usize, pub poll_timeout: Duration, } @@ -82,6 +83,7 @@ impl Default for OcrConnection { extra_headers_source: InputSource::Deployment, timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, + max_response_bytes: crate::constants::OCR_RESPONSE_MAX_BYTES, poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS), } } @@ -95,6 +97,7 @@ pub struct LiteLLMOcrRequest { pub litellm_call_id: Option, pub optional_params: Map, pub input_sources: BTreeMap, + pub azure_ad_token_provider: Option, pub(crate) adapter: OcrAdapterKind, } @@ -115,6 +118,7 @@ impl LiteLLMOcrRequest { litellm_call_id: None, optional_params, input_sources: BTreeMap::new(), + azure_ad_token_provider: None, adapter: adapter_kind, }) } @@ -132,6 +136,10 @@ impl LiteLLMOcrRequest { .map(|format| format.unwrap_or_default()) } + pub fn provider_name(&self) -> &'static str { + self.adapter.provider().as_str() + } + pub fn with_host_hooks( self, hooks: Arc, @@ -169,6 +177,47 @@ mod tests { use super::*; use serde_json::json; + #[test] + fn document_variants_preserve_provider_fields_when_rewriting_sources() { + for (value, original, replacement, expected) in [ + ( + json!({ + "type":"document_url", + "document_url":"https://example.com/input.pdf", + "document_name":"input.pdf" + }), + "https://example.com/input.pdf", + "data:application/pdf;base64,AA==", + json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,AA==", + "document_name":"input.pdf" + }), + ), + ( + json!({ + "type":"image_url", + "image_url":"https://example.com/input.png", + "detail":"high" + }), + "https://example.com/input.png", + "data:image/png;base64,AA==", + json!({ + "type":"image_url", + "image_url":"data:image/png;base64,AA==", + "detail":"high" + }), + ), + ] { + let document: OcrDocument = serde_json::from_value(value).unwrap(); + assert_eq!(document.source(), original); + assert_eq!( + serde_json::to_value(document.with_source(replacement.into())).unwrap(), + expected + ); + } + } + #[test] fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() { let response = LiteLLMOcrResponse { diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 34d0a7d7b86..6dc6b34b73d 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -3,7 +3,6 @@ use crate::ocr::error::OcrResponseError; use std::collections::BTreeMap; use std::time::Duration; -use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest}; use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument}; use crate::Error; use crate::auth::InputSource; @@ -13,10 +12,58 @@ use serde::{ }; use serde_json::{Map, Value}; +const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"]; +const MISTRAL_OPTION_FIELDS: &[&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", + "include_blocks", + "id", +]; +const DEEPSEEK_OPTION_FIELDS: &[&str] = + &["stream", "temperature", "max_tokens", "top_p", "n", "stop"]; +const DOCUMENT_INTELLIGENCE_OPTION_FIELDS: &[&str] = &["pages", "features"]; +const REDUCTO_V3_OPTION_FIELDS: &[&str] = &["formatting", "retrieval", "settings"]; +const REDUCTO_LEGACY_OPTION_FIELDS: &[&str] = &["enhance"]; +const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_scope", + "azure_authority_host", + "azure_credential", + "azure_federated_token_file", + "enable_azure_ad_token_refresh", +]; +const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[ + "vertex_credentials", + "vertex_ai_credentials", + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", +]; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OptionalParamSpec { + pub name: &'static str, + pub secret: bool, +} + #[derive(Debug)] pub struct DecodedOcrResponse { pub data: T, pub native: Option, + pub text: String, } #[derive(Deserialize)] @@ -39,11 +86,65 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok() } +pub fn consumed_optional_param_names( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, Error> { + use super::registry::OcrAdapterKind; + + let (_, adapter) = super::registry::resolve_wire_adapter(model, custom_llm_provider)?; + let provider_fields: &[&str] = match adapter { + OcrAdapterKind::Cohere | OcrAdapterKind::AzureCohere => &["output_format"], + OcrAdapterKind::Mistral | OcrAdapterKind::AzureMistral | OcrAdapterKind::VertexMistral => { + MISTRAL_OPTION_FIELDS + } + OcrAdapterKind::AzureDocumentIntelligence => DOCUMENT_INTELLIGENCE_OPTION_FIELDS, + OcrAdapterKind::ReductoV3 => REDUCTO_V3_OPTION_FIELDS, + OcrAdapterKind::ReductoLegacy => REDUCTO_LEGACY_OPTION_FIELDS, + OcrAdapterKind::VertexDeepSeek => DEEPSEEK_OPTION_FIELDS, + }; + let auth_fields: &[&str] = match adapter { + OcrAdapterKind::AzureMistral + | OcrAdapterKind::AzureDocumentIntelligence + | OcrAdapterKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS, + OcrAdapterKind::VertexMistral | OcrAdapterKind::VertexDeepSeek => VERTEX_AUTH_OPTION_FIELDS, + _ => &[], + }; + Ok(COMMON_OPTION_FIELDS + .iter() + .chain(provider_fields) + .chain(auth_fields) + .copied() + .collect()) +} + +pub fn consumed_optional_params( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, Error> { + consumed_optional_param_names(model, custom_llm_provider).map(|names| { + names + .into_iter() + .map(|name| OptionalParamSpec { + name, + secret: matches!( + name, + "azure_ad_token" + | "client_secret" + | "azure_federated_token_file" + | "vertex_credentials" + | "vertex_ai_credentials" + ), + }) + .collect() + }) +} + pub fn decode_request(wire: OcrWireRequest) -> Result { let api_key_source = source_for(&wire.input_sources, "api_key"); let api_base_source = source_for(&wire.input_sources, "api_base"); let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); - let document = decode_request_value(wire.document, "document")?; + let document = decode_document(wire.document)?; let headers = wire .extra_headers .unwrap_or_default() @@ -66,11 +167,28 @@ pub fn decode_request(wire: OcrWireRequest) -> Result }) .transpose()?; let defaults = OcrConnection::default(); + let max_response_bytes = wire + .optional_params + .get("max_response_bytes") + .map(|value| { + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0 && *value <= defaults.max_response_bytes) + .ok_or_else(|| OcrRequestError::RequestField { + path: "max_response_bytes".into(), + }) + }) + .transpose()? + .unwrap_or(defaults.max_response_bytes); let request = LiteLLMOcrRequest::new( wire.model, document, wire.custom_llm_provider.as_deref(), - wire.optional_params, + wire.optional_params + .into_iter() + .filter(|(name, _)| name != "max_response_bytes") + .collect(), )?; let connection = OcrConnection { api_key: nonblank(wire.api_key), @@ -81,6 +199,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result extra_headers_source, timeout: timeout.unwrap_or(defaults.timeout), max_download_bytes: defaults.max_download_bytes, + max_response_bytes, poll_timeout: defaults.poll_timeout, }; Ok(LiteLLMOcrRequest { @@ -90,6 +209,16 @@ pub fn decode_request(wire: OcrWireRequest) -> Result }) } +fn decode_document(value: Value) -> Result { + let kind = value.get("type").and_then(Value::as_str); + let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none() + || matches!(kind, Some("image_url")) && value.get("image_url").is_none(); + if missing_url { + return Err(OcrRequestError::MissingDocumentUrl); + } + decode_request_value(value, "document") +} + fn source_for(sources: &BTreeMap, name: &str) -> InputSource { sources.get(name).copied().unwrap_or_default() } @@ -134,38 +263,81 @@ pub fn decode_response( } else { None }; - Ok(DecodedOcrResponse { data, native }) -} - -pub fn decode_pre_call_result( - original: OcrPreCallRequest, - value: Value, -) -> Result { - #[derive(Deserialize)] - struct Changed { - document: OcrDocument, - #[serde(default)] - optional_params: Map, - } - let changed: Changed = decode_request_value(value, "guardrail")?; - Ok(OcrPreCallRequest { - document: changed.document, - optional_params: Value::Object(changed.optional_params), - ..original + Ok(DecodedOcrResponse { + data, + native, + text: String::from_utf8_lossy(bytes).into_owned(), }) } -pub fn decode_during_call_result( - original: OcrDuringCallRequest, - value: Value, -) -> Result { - #[derive(Deserialize)] - struct Changed { - body: Value, +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn option_projection_is_provider_specific_and_excludes_opaque_fields() { + let mistral = consumed_optional_param_names("mistral/model", None).unwrap(); + assert!(mistral.contains(&"pages")); + assert!(mistral.contains(&"req_format")); + assert!(!mistral.contains(&"vertex_project")); + assert!(!mistral.contains(&"opaque_extension")); + + let vertex = consumed_optional_param_names("vertex_ai/deepseek-ocr", None).unwrap(); + assert!(vertex.contains(&"temperature")); + assert!(vertex.contains(&"vertex_credentials")); + assert!(!vertex.contains(&"pages")); + } + + #[test] + fn optional_param_metadata_marks_only_credentials_as_secret() { + let azure = consumed_optional_params("model", Some("azure_ai")).unwrap(); + assert!( + azure + .iter() + .any(|spec| spec.name == "client_secret" && spec.secret) + ); + assert!( + azure + .iter() + .any(|spec| spec.name == "tenant_id" && !spec.secret) + ); + let vertex = consumed_optional_params("deepseek-ocr", Some("vertex_ai")).unwrap(); + assert!( + vertex + .iter() + .any(|spec| spec.name == "vertex_credentials" && spec.secret) + ); + assert!( + vertex + .iter() + .any(|spec| spec.name == "vertex_project" && !spec.secret) + ); + } + + #[test] + fn activation_includes_migrated_providers() { + assert!(is_supported_request("model", Some("mistral"))); + assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); + assert!(is_supported_request( + "documentintelligence/prebuilt-read", + Some("azure_ai") + )); + assert!(is_supported_request("parse-v3", Some("reducto"))); + assert!(is_supported_request("parse-legacy", Some("reducto"))); + assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); + assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); + } + + #[test] + fn missing_document_source_has_a_typed_public_error() { + for document in [ + serde_json::json!({"type": "document_url"}), + serde_json::json!({"type": "image_url"}), + ] { + assert_eq!( + decode_document(document), + Err(OcrRequestError::MissingDocumentUrl) + ); + } } - let changed: Changed = decode_request_value(value, "guardrail")?; - Ok(OcrDuringCallRequest { - body: changed.body, - ..original - }) } diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 5d037e9cf1b..34213e5f6c4 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -1,3 +1,21 @@ +use std::collections::HashMap; +use std::io; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use futures_util::{SinkExt, StreamExt}; +use rustls::{ClientConfig, RootCertStore}; +use tokio::net::TcpStream; +use tokio::sync::Mutex; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::TlsError; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue}; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, +}; + use crate::Error; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; @@ -125,6 +143,137 @@ pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool { ) } +pub type ResponsesUpstreamWs = WebSocketStream>; + +static TLS_CONFIG: OnceLock> = OnceLock::new(); + +fn build_tls_config() -> Result> { + let native = rustls_native_certs::load_native_certs(); + let mut store = RootCertStore::empty(); + let (added, _ignored) = store.add_parsable_certificates(native.certs); + if added == 0 { + return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io( + io::Error::other(format!( + "no usable native root certificates: {:?}", + native.errors + )), + ))); + } + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map(|builder| builder.with_root_certificates(store).with_no_client_auth()) + .map_err(|error| { + Box::new(tokio_tungstenite::tungstenite::Error::Tls( + TlsError::Rustls(error), + )) + }) +} + +fn tls_config() -> Result, Box> { + if let Some(config) = TLS_CONFIG.get() { + return Ok(Arc::clone(config)); + } + let built = Arc::new(build_tls_config()?); + Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) +} + +pub async fn connect_upstream( + request: R, +) -> Result<(ResponsesUpstreamWs, Response), Box> +where + R: IntoClientRequest + Unpin, +{ + let request = request.into_client_request().map_err(Box::new)?; + let connector = match request.uri().scheme_str() { + Some("wss") => Some(Connector::Rustls(tls_config()?)), + _ => None, + }; + connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(Box::new) +} + +#[derive(Clone)] +pub struct ResponsesWebSocketConnection { + socket: Arc>>, +} + +impl ResponsesWebSocketConnection { + pub async fn connect_url( + url: &str, + headers: &HashMap, + timeout: Option, + ) -> Result { + let mut request = url + .into_client_request() + .map_err(|error| Error::Network(error.to_string()))?; + for (name, value) in headers { + let header_name = name + .parse::() + .map_err(|error| Error::InvalidRequest(error.to_string()))?; + let header_value = HeaderValue::from_str(value) + .map_err(|error| Error::InvalidRequest(error.to_string()))?; + request.headers_mut().insert(header_name, header_value); + } + let connect = connect_upstream(request); + let result = match timeout { + Some(timeout) => tokio::time::timeout(timeout, connect) + .await + .map_err(|_| Error::Network("Responses WebSocket connection timed out".into()))?, + None => connect.await, + }; + let (socket, _) = result.map_err(|error| match *error { + tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { + status: response.status().as_u16(), + body: String::new(), + }, + other => Error::Network(other.to_string()), + })?; + Ok(Self { + socket: Arc::new(Mutex::new(Some(socket))), + }) + } + + pub async fn send_text(&self, text: String) -> Result<(), Error> { + let mut socket = self.socket.lock().await; + let Some(socket) = socket.as_mut() else { + return Err(Error::Network("Responses WebSocket is closed".into())); + }; + socket + .send(Message::Text(text)) + .await + .map_err(|error| Error::Network(error.to_string())) + } + + pub async fn recv_text(&self) -> Result, Error> { + let mut socket = self.socket.lock().await; + let Some(socket) = socket.as_mut() else { + return Ok(None); + }; + match socket.next().await { + Some(Ok(Message::Text(text))) => Ok(Some(text)), + Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) + .map(Some) + .map_err(|error| Error::InvalidResponse(error.to_string())), + Some(Ok(Message::Close(_))) | None => Ok(None), + Some(Ok(_)) => Ok(None), + Some(Err(error)) => Err(Error::Network(error.to_string())), + } + } + + pub async fn close(&self) -> Result<(), Error> { + let mut socket = self.socket.lock().await; + if let Some(socket) = socket.as_mut() { + socket + .close(None) + .await + .map_err(|error| Error::Network(error.to_string()))?; + } + *socket = None; + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs index d7d532cfef1..b6dc8d90b93 100644 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -70,7 +70,7 @@ async fn facade_acquires_supplied_entra_token_for_final_request() { struct ReplaceBodyDocument; impl OcrHooks for ReplaceBodyDocument { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index e4c81dea5a7..3fca59033cc 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,5 @@ use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use super::wire::{OcrWireRequest, decode_request}; @@ -124,6 +125,14 @@ async fn immediate_response_normalizes_pages_and_preserves_native() { json!({"width":816,"height":1056,"dpi":96}) ); assert_eq!(result.usage_info, Some(json!({"pages_processed":1}))); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); assert_eq!(result.provider_native_response, Some(operation)); } @@ -169,6 +178,55 @@ async fn accepted_response_polls_to_success_with_only_credentials() { } } +struct SubmissionBoundary { + request_count: Arc>>, +} + +impl super::hooks::OcrHooks for SubmissionBoundary { + fn post_call( + &self, + request: super::hooks::OcrPostCallRequest, + ) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> { + Box::pin(async move { + match self.request_count.lock().unwrap().len() { + 1 => assert_eq!(request.original_response, json!(r#"{"submitted":true}"#)), + 2 => assert!( + request + .original_response + .as_str() + .unwrap() + .contains("succeeded") + ), + count => panic!("unexpected callback after {count} requests"), + } + Ok(request) + }) + } +} + +#[tokio::test] +async fn accepted_response_runs_post_call_before_polling() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(SubmissionBoundary { + request_count: seen.clone(), + }), + ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); +} + #[tokio::test] async fn polling_forwards_bearer_credentials() { let (base, seen, server) = mock_server(vec![ @@ -361,7 +419,7 @@ async fn pre_call_guardrail_receives_caller_pages_before_mapping() { struct RewritePages; impl OcrHooks for RewritePages { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs index 875fc9e3dc6..4ba39561dcd 100644 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/deepseek_ocr.rs @@ -34,6 +34,28 @@ fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { assert!(result.get("ignored").is_none()); } +#[rstest] +#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] +#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] +fn request_maps_both_document_types_to_image_content(#[case] document: Value) { + let source = document + .get("image_url") + .or_else(|| document.get("document_url")) + .unwrap() + .clone(); + let request = transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + serde_json::from_value(document).unwrap(), + &DeepSeekOcrParams::default(), + ) + .unwrap(); + let result = serde_json::to_value(request).unwrap(); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":source}) + ); +} + #[rstest] #[case(json!("# hello"), "# hello")] #[case(json!("{broken"), "{broken")] diff --git a/litellm-rust/crates/core/tests/host_lifecycle.rs b/litellm-rust/crates/core/tests/host_lifecycle.rs new file mode 100644 index 00000000000..19fb946afde --- /dev/null +++ b/litellm-rust/crates/core/tests/host_lifecycle.rs @@ -0,0 +1,116 @@ +use crate::Error; +use crate::call_lifecycle::host::{HostFailure, HostLifecycle, HostPhase}; + +fn run(fail_at: Option, asynchronous: bool) -> (Vec, Vec) { + let mut lifecycle = HostLifecycle::new(asynchronous); + let mut events = Vec::new(); + let mut failures = Vec::new(); + while lifecycle.phase() != HostPhase::Complete { + let phase = lifecycle.phase(); + events.push(phase); + let result = if Some(phase) == fail_at { + Err(HostFailure::Error(Error::InvalidRequest( + "selected failure".into(), + ))) + } else { + Ok(()) + }; + if let Some(error) = lifecycle.accept(result) { + failures.push(error); + } + } + (events, failures) +} + +#[test] +fn public_outcome_is_finalized_before_a_single_terminal_dispatch() { + for asynchronous in [false, true] { + let (events, failures) = run(None, asynchronous); + assert!(failures.is_empty()); + assert_eq!( + &events[events.len() - 2..], + &[HostPhase::Finalize, HostPhase::Success] + ); + assert_eq!( + events + .iter() + .filter(|phase| **phase == HostPhase::Execute) + .count(), + 1 + ); + assert_eq!( + events.contains(&HostPhase::DeploymentPostCall), + asynchronous + ); + } +} + +#[test] +fn only_provider_and_response_construction_failures_use_provider_mapping() { + for phase in [ + HostPhase::Setup, + HostPhase::DeploymentPreCall, + HostPhase::Prepare, + HostPhase::Execute, + HostPhase::ConstructResponse, + HostPhase::DeploymentPostCall, + HostPhase::Finalize, + ] { + let (events, failures) = run(Some(phase), true); + assert_eq!(failures.len(), 1); + assert!(!events.contains(&HostPhase::Success)); + let mapped = matches!(phase, HostPhase::Execute | HostPhase::ConstructResponse); + assert_eq!(events.contains(&HostPhase::MapFailure), mapped); + assert_eq!(events.contains(&HostPhase::DeploymentFailure), mapped); + assert_eq!( + &events[events.len() - 2..], + &[HostPhase::Failure, HostPhase::AsyncFailure] + ); + assert!( + events + .iter() + .filter(|phase| **phase == HostPhase::Execute) + .count() + <= 1 + ); + } +} + +#[test] +fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() { + let mut lifecycle = HostLifecycle::new(true); + while lifecycle.phase() != HostPhase::Execute { + lifecycle.accept(Ok(())); + } + let selected = Error::InvalidRequest("provider".into()); + assert_eq!( + lifecycle.accept(Err(HostFailure::Error(selected.clone()))), + Some(selected) + ); + lifecycle.accept(Ok(())); + for phase in [ + HostPhase::DeploymentFailure, + HostPhase::Failure, + HostPhase::AsyncFailure, + ] { + assert_eq!(lifecycle.phase(), phase); + assert_eq!( + lifecycle.accept(Err(HostFailure::Error(Error::InvalidRequest( + "callback".into() + )))), + None + ); + } + assert_eq!(lifecycle.phase(), HostPhase::Complete); +} + +#[test] +fn cancellation_skips_terminal_dispatch() { + let mut lifecycle = HostLifecycle::new(true); + let error = Error::InvalidRequest("cancelled".into()); + assert_eq!( + lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))), + Some(error) + ); + assert_eq!(lifecycle.phase(), HostPhase::Complete); +} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index cecd8869741..55f8713d76e 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -3,9 +3,16 @@ use std::sync::{Arc, Mutex}; use serde_json::{Value, json}; use super::OcrClient; -use super::hooks::{OcrHookFuture, OcrHooks, OcrLogFuture, OcrPreCallRequest}; +use super::hooks::{ + OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, + OcrPreCallRequest, +}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use super::wire::{OcrWireRequest, decode_request}; +use super::{ + NativeOutcome, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHost, + OcrHostOperation, OcrHostResult, +}; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; #[test] @@ -51,7 +58,7 @@ async fn facade_executes_direct_mistral_once() { let result = perform_ocr(wire_request( "mistral/model", &base, - json!({"extract_header":true,"unknown":"ignored"}), + json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), )) .await .unwrap(); @@ -72,6 +79,7 @@ async fn facade_executes_direct_mistral_once() { json!({ "model":"model", "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", "extract_header":true }) ); @@ -124,7 +132,7 @@ struct RecordingHooks { } impl OcrHooks for RecordingHooks { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } @@ -148,6 +156,13 @@ impl OcrHooks for RecordingHooks { }) } + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("post"); + Ok(request) + }) + } + fn success<'a>( &'a self, _context: &'a CallLifecycleContext, @@ -171,6 +186,38 @@ impl OcrHooks for RecordingHooks { } } +struct HeaderEditHooks; + +impl OcrHooks for HeaderEditHooks { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + request + .headers + .push(("x-core-callback".into(), "edited".into())); + Box::pin(async move { Ok(request) }) + } +} + +#[tokio::test] +async fn lifecycle_sends_headers_returned_by_the_typed_during_call_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(HeaderEditHooks), + ..wire_request("mistral/model", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); +} + #[tokio::test] async fn lifecycle_orders_hooks_and_emits_one_success() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; @@ -185,7 +232,10 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { }; perform_ocr(request).await.unwrap(); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["pre", "during", "success"]); + assert_eq!( + *events.lock().unwrap(), + ["pre", "during", "post", "success"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -227,3 +277,562 @@ async fn upstream_failure_emits_one_terminal_failure() { assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]); assert_eq!(seen.lock().unwrap().len(), 1); } + +struct AdmissionSpy { + effects: Arc>, +} + +impl OcrHooks for AdmissionSpy { + fn intercepts_requests(&self) -> bool { + *self.effects.lock().unwrap() += 1; + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + *self.effects.lock().unwrap() += 1; + Box::pin(async move { Ok(request) }) + } +} + +#[test] +fn admission_declines_without_invoking_hooks_or_transport() { + for (admission, expected) in [ + ( + OcrAdmission { + provider_workflow: false, + host_operations: true, + asynchronous: false, + }, + OcrDecline::ProviderWorkflow, + ), + ( + OcrAdmission { + provider_workflow: true, + host_operations: false, + asynchronous: false, + }, + OcrDecline::HostOperations, + ), + ] { + let outcome = OcrCall::admit(super::test_support::ocr_client(), admission); + assert!(matches!(outcome, NativeOutcome::Declined(reason) if reason == expected)); + } +} + +#[tokio::test] +async fn fallible_host_phases_do_not_replay_or_reach_transport() { + for failure_phase in ["pre", "during"] { + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut result = None; + let mut phases = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(operation)) => match operation { + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => { + result = Some(OcrHostResult::Lifecycle(Ok(()))) + } + OcrHostOperation::ProjectRequest => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))) + } + OcrHostOperation::AcquireAzureAdToken => { + panic!("test request has no token provider") + } + OcrHostOperation::PreCall(request) => { + phases.push("pre"); + result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { + Err(crate::Error::InvalidRequest("pre failed".into())) + } else { + Ok(request) + })); + } + OcrHostOperation::DuringCall(request) => { + phases.push("during"); + result = Some(OcrHostResult::DuringCall(if failure_phase == "during" { + Err(crate::Error::InvalidRequest("during failed".into())) + } else { + Ok(request) + })); + } + OcrHostOperation::PostCall(_) => panic!("transport should not be reached"), + }, + Err(error) => break error, + Ok(OcrCallStep::Complete(_)) => panic!("failed call completed"), + } + }; + assert!(matches!(error, crate::Error::InvalidRequest(_))); + assert_eq!( + phases + .iter() + .filter(|phase| **phase == failure_phase) + .count(), + 1 + ); + } +} + +#[tokio::test] +async fn invalid_provider_response_runs_post_call_before_normalization_failure() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; + let mut request = Some(wire_request("mistral/model", &base, json!({}))); + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let host = NoopOcrHost; + let mut result = None; + let mut post_calls = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))); + } + Ok(OcrCallStep::Host(operation)) => { + if let OcrHostOperation::PostCall(request) = &operation { + post_calls.push(request.original_response.clone()); + } + result = Some(host.invoke(operation).await); + } + Err(error) => break error, + Ok(OcrCallStep::Complete(_)) => panic!("invalid provider response completed"), + } + }; + server.await.unwrap(); + assert!(matches!(error, crate::Error::InvalidResponse(_))); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]); +} + +#[tokio::test] +async fn direct_native_host_drives_the_same_state_machine() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"native"}] + }))]) + .await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", &base, json!({})) + }; + let NativeOutcome::Completed(mut call) = OcrCall::admit( + super::test_support::ocr_client(), + OcrAdmission { + asynchronous: true, + ..OcrAdmission::all() + }, + ) else { + panic!("supported call declined") + }; + let mut request = Some(request); + let host = NoopOcrHost; + let mut result = None; + let mut operations = Vec::new(); + let response = loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(operation) => { + operations.push(match &operation { + OcrHostOperation::ProjectRequest => "ProjectRequest".into(), + OcrHostOperation::Lifecycle(phase) => format!("{phase:?}"), + OcrHostOperation::PreCall(_) => "PreCall".into(), + OcrHostOperation::DuringCall(_) => "DuringCall".into(), + OcrHostOperation::PostCall(_) => "PostCall".into(), + OcrHostOperation::ConstructResponse(_) => "ConstructResponse".into(), + OcrHostOperation::Success { response, .. } => { + assert_eq!(response.pages[0]["markdown"], "native"); + "Success".into() + } + _ => panic!("unexpected OCR operation"), + }); + result = Some(match operation { + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } + operation => host.invoke(operation).await, + }); + } + OcrCallStep::Complete(response) => break response, + } + }; + server.await.unwrap(); + assert_eq!(response.pages[0]["markdown"], "native"); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!( + operations, + [ + "Setup", + "DeploymentPreCall", + "Prepare", + "ProjectRequest", + "PreCall", + "DuringCall", + "PostCall", + "ConstructResponse", + "DeploymentPostCall", + "Finalize", + "Success", + ] + ); + assert!(matches!( + call.resume(None).await, + Err(crate::Error::InvalidRequest(_)) + )); +} + +#[tokio::test] +async fn public_finalization_failure_never_dispatches_success_or_replays_provider() { + use crate::call_lifecycle::host::{HostFailure, HostPhase}; + + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = Some(wire_request("mistral/model", &base, json!({}))); + let NativeOutcome::Completed(mut call) = OcrCall::admit( + super::test_support::ocr_client(), + OcrAdmission { + asynchronous: true, + ..OcrAdmission::all() + }, + ) else { + panic!("supported call declined") + }; + let selected = crate::Error::InvalidRequest("public metadata failed".into()); + let host = NoopOcrHost; + let mut result = None; + let mut failures = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(operation)) => { + result = Some(match operation { + OcrHostOperation::Lifecycle(HostPhase::Finalize) => { + OcrHostResult::Lifecycle(Err(HostFailure::Error(selected.clone()))) + } + OcrHostOperation::Failure { error, .. } => { + assert_eq!(error, selected); + failures.push("sync"); + OcrHostResult::Lifecycle(Err(HostFailure::Error( + crate::Error::InvalidRequest("failure callback failed".into()), + ))) + } + OcrHostOperation::Lifecycle(HostPhase::AsyncFailure) => { + failures.push("async"); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Success { .. } + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { + panic!("finalization failure used provider/success dispatch") + } + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } + operation => host.invoke(operation).await, + }); + } + Ok(OcrCallStep::Complete(_)) => panic!("failed call completed successfully"), + Err(error) => break error, + } + }; + server.await.unwrap(); + assert_eq!(error, selected); + assert_eq!(failures, ["sync", "async"]); + assert_eq!(seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption() { + use crate::call_lifecycle::host::HostFailure; + + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let host = NoopOcrHost; + let mut result = None; + loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))) + } + OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await), + OcrCallStep::Complete(_) => panic!("provider executed before pre-call result"), + } + } + let selected = crate::Error::InvalidRequest("cancelled".into()); + assert!(matches!( + call.interrupt(HostFailure::Cancelled(selected.clone())).await, + Err(error) if error == selected + )); + assert!( + call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) + .await + .is_err() + ); +} + +#[tokio::test] +async fn missing_host_result_preserves_pending_operation() { + use crate::call_lifecycle::host::HostPhase; + + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + assert!(matches!( + call.resume(None).await.unwrap(), + OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Setup)) + )); + assert!(call.resume(None).await.is_err()); + assert!(matches!( + call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) + .await + .unwrap(), + OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Prepare)) + )); +} + +async fn read_bounded_response( + response: Vec, + limit: usize, +) -> Result { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::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 = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(&response).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + super::client::read_response_bytes(response, limit), + ) + .await; + server.abort(); + let _ = server.await; + result.expect("bounded reads must finish without waiting for the rest of an oversized body") +} + +#[tokio::test] +async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { + use super::error::{OcrError, OcrResponseError}; + + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", + ] { + assert_eq!( + read_bounded_response(response.as_bytes().to_vec(), 8) + .await + .unwrap(), + "abcdefgh" + ); + } + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", + ] { + assert!(matches!( + read_bounded_response(response.as_bytes().to_vec(), 8).await, + Err(OcrError::Response(OcrResponseError::TooLarge { limit: 8 })) + )); + } +} + +#[tokio::test] +async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining() { + let prefix = "x".repeat(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)); + for headers in ["Content-Length: 1000000", "Transfer-Encoding: chunked"] { + let body = if headers.starts_with("Transfer") { + format!("{:x}\r\n{prefix}\r\n", prefix.len()) + } else { + prefix.clone() + }; + let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); + let error = read_bounded_response(response.into_bytes(), 4096) + .await + .unwrap_err(); + match error { + super::error::OcrError::Transport(crate::error::TransportError::Http { + status, + body, + }) => { + assert_eq!(status, 429); + assert_eq!( + body, + format!( + "{}... (truncated)", + "x".repeat(crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS) + ) + ); + } + error => panic!("unexpected error: {error}"), + } + } +} + +#[test] +fn response_limit_is_validated_and_not_forwarded_to_the_provider() { + let request = wire_request( + "mistral/model", + "http://localhost", + json!({"max_response_bytes": 123}), + ); + assert_eq!(request.connection.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); + for value in [ + json!(0), + json!(-1), + json!(true), + json!("123"), + json!(1.5), + json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1), + Value::Null, + ] { + let wire = serde_json::from_value(json!({ + "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "optional_params": {"max_response_bytes": value} + })).unwrap(); + let Err(error) = decode_request(wire) else { + panic!("invalid response limit accepted") + }; + assert!(error.to_string().contains("max_response_bytes")); + } +} + +#[derive(Debug)] +struct PendingToken { + entered: Arc, + dropped: Arc, +} + +struct TokenFutureDrop(Arc); + +impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } +} + +impl crate::auth::TokenProvider for PendingToken { + fn acquire(&self) -> crate::auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } +} + +#[tokio::test] +async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_is_cancelled() { + use crate::call_lifecycle::host::HostFailure; + use std::future::Future; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::task::Poll; + + for interrupt_acknowledgement in [false, true] { + let entered = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + let request = super::LiteLLMOcrRequest { + connection: super::OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.connection + }, + azure_ad_token_provider: Some(crate::auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))), + ..request + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = entered.notified() => break, + step = call.resume(result.take()) => { + result = Some(match step.unwrap() { + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))), + OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await, + OcrCallStep::Complete(_) => panic!("pending provider completed"), + }); + } + } + } + }).await.unwrap(); + assert!(!dropped.load(Ordering::SeqCst)); + let selected = crate::Error::InvalidRequest("cancelled".into()); + if interrupt_acknowledgement { + let mut acknowledgement = + Box::pin(call.interrupt(HostFailure::Cancelled(selected.clone()))); + std::future::poll_fn(|cx| { + assert!(acknowledgement.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(acknowledgement); + assert!(!dropped.load(Ordering::SeqCst)); + } + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + call.interrupt(HostFailure::Cancelled(selected.clone())), + ) + .await + .unwrap(); + assert!(matches!(result, Err(error) if error == selected)); + assert!( + dropped.load(Ordering::SeqCst), + "cancellation returned while provider captures were still alive" + ); + } +} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 8e86e4713ef..a15e9cae5b5 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use rstest::rstest; use serde_json::{Value, json}; -use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks}; +use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; fn request_body(request: &str) -> Value { @@ -100,6 +100,42 @@ async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { assert!(requests[1].starts_with("POST /parse ")); } +struct ParseBoundary { + request_count: Arc>>, +} + +impl OcrHooks for ParseBoundary { + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + assert_eq!(self.request_count.lock().unwrap().len(), 2); + assert_eq!( + request.original_response, + json!(r#"{"result":{"chunks":[]}}"#) + ); + Ok(request) + }) + } +} + +#[tokio::test] +async fn post_call_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request = super::LiteLLMOcrRequest { + hooks: Arc::new(ParseBoundary { + request_count: seen.clone(), + }), + ..wire_request("reducto/parse-v3", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); +} + #[rstest] #[case(json!({"file_id":""}))] #[case(json!({}))] @@ -148,9 +184,16 @@ async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { fn response_normalization_groups_blocks_and_distinguishes_null_result() { use crate::ocr::codecs::reducto::{ReductoResponse, transform_ocr_response}; - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"chunks":[ - {"blocks":[{"content":"B","bbox":{"page":2},"kind":"table"}]}, - {"blocks":[{"content":"A","bbox":{"page":1},"kind":"text"},{"content":"C","bbox":{"page":1}}]} + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ + {"blocks":[{ + "type":"Table", + "content":"B", + "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, + "confidence":"high", + "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, + "image_url":null + }]}, + {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} ]}}); let response: ReductoResponse = serde_json::from_value(raw).unwrap(); let normalized = transform_ocr_response("parse-v3", response) @@ -158,7 +201,17 @@ fn response_normalization_groups_blocks_and_distinguishes_null_result() { .into_json(); assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["kind"], "table"); + assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["bbox"], + json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) + ); + assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], + 0.95 + ); + assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); assert_eq!(normalized["usage_info"]["pages_processed"], 2); assert_eq!(normalized["usage_info"]["credits"], 3.0); @@ -195,7 +248,7 @@ async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { struct RewriteDocument; impl OcrHooks for RewriteDocument { - fn has_guardrails(&self) -> bool { + fn intercepts_requests(&self) -> bool { true } diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs index 6d3061d8f5d..676799eb2fe 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -49,7 +49,7 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { assert!(body.get("extra_body").is_none()); assert_eq!( body["messages"][0]["content"][0], - json!({"type":"document_url","document_url":"gs://bucket/document.pdf"}) + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) ); } diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs deleted file mode 100644 index e7739fe7312..00000000000 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ /dev/null @@ -1,115 +0,0 @@ -//! Enforcement: the litellm-rust workspace has exactly six crates. -//! -//! `core` (the Rust SDK), `token-counter` (standalone input token counting), -//! `config` (the config-loading boundary), -//! `ai-gateway` (the HTTP/WebSocket host), -//! `python-interop` (domain-neutral PyO3 primitives), 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 -//! artifact — and to keep `litellm-rust/AGENTS.md` in sync. -//! -//! Std-only (no toml crate): we scan the workspace manifest's `members = [...]` -//! block and the `crates/` directory directly. - -use std::collections::BTreeSet; -use std::fs; -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/token-counter", - "crates/config", - "crates/ai-gateway", - "crates/python-interop", - "crates/python-bridge", -]; - -/// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &[ - "core", - "token-counter", - "config", - "ai-gateway", - "python-interop", - "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)."; - -/// Absolute path to the workspace root (`litellm-rust/`). -fn workspace_root() -> PathBuf { - // CARGO_MANIFEST_DIR is `.../litellm-rust/crates/core`; the workspace root is - // two levels up. - Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/../..")) - .canonicalize() - .expect("workspace root should resolve") -} - -/// Parse the `members = [ ... ]` array out of the workspace `[workspace]` table. -/// -/// Minimal hand-rolled scan: find `members`, then collect every double-quoted -/// string up to the closing `]`. Good enough for our fixed manifest shape and -/// keeps this test dependency-free. -fn parse_members(manifest: &str) -> BTreeSet { - let after_members = manifest - .split_once("members") - .map(|(_, rest)| rest) - .expect("workspace manifest should declare members"); - let open = after_members.find('[').expect("members should be an array"); - let close = after_members[open..] - .find(']') - .map(|offset| open + offset) - .expect("members array should be closed"); - let body = &after_members[open + 1..close]; - - let mut members = BTreeSet::new(); - let mut rest = body; - while let Some(start) = rest.find('"') { - let after_quote = &rest[start + 1..]; - let end = after_quote - .find('"') - .expect("opening quote should be matched"); - members.insert(after_quote[..end].to_string()); - rest = &after_quote[end + 1..]; - } - members -} - -/// The crate subdirectory names under `crates/`. -/// -/// A directory counts as a crate only when it holds a `Cargo.toml`; non-crate -/// directories (e.g. docs like `CODING_STANDARDS/`) are ignored so they can live -/// under `crates/` without tripping the crate-proliferation guard. -fn crate_dirs(root: &Path) -> BTreeSet { - fs::read_dir(root.join("crates")) - .expect("crates/ directory should exist") - .filter_map(Result::ok) - .filter(|entry| entry.file_type().map(|ty| ty.is_dir()).unwrap_or(false)) - .filter(|entry| entry.path().join("Cargo.toml").is_file()) - .map(|entry| entry.file_name().to_string_lossy().into_owned()) - .collect() -} - -#[test] -fn workspace_members_match_allowlist() { - let root = workspace_root(); - let manifest = fs::read_to_string(root.join("Cargo.toml")) - .expect("workspace Cargo.toml should be readable"); - - let actual = parse_members(&manifest); - let expected: BTreeSet = EXPECTED_MEMBERS.iter().map(|s| s.to_string()).collect(); - assert_eq!(actual, expected, "{MISMATCH}"); -} - -#[test] -fn crates_directory_matches_allowlist() { - let root = workspace_root(); - - let actual = crate_dirs(&root); - let expected: BTreeSet = EXPECTED_CRATE_DIRS.iter().map(|s| s.to_string()).collect(); - assert_eq!(actual, expected, "{MISMATCH}"); -} diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 42282ca4da4..9262617156b 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,42 @@ -litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop. - -Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. +- Target invariants, not completion claims; these supersede older conflicting bridge guidance +- Keep this crate the product-specific PyO3 consumer of `litellm-python-interop` + - Own registration, input projection, retained Python state, callback invocation, public response/error construction and host scheduling + - Keep value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment in `execution.rs`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` + - Core owns typed native state, admission, lifecycle sequencing, provider preparation/I/O, normalization and terminal-outcome/dispatch decisions + - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers + - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points +- Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work + - Free-threading requires separate runtime/concurrency validation; omitting the attribute does not opt out on PyO3 0.28+ +- Preserve public argument binding and Python object provenance + - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view + - Retain independently captured body/header roots; in-place mutation and logging-envelope field replacement have different effects + - Project only consumed fields at reference read points; no eager whole-graph serialization or equality-based alias reconstruction + - Preserve provider-specific upload/submission/poll observation and encoding boundaries; signed/build-captured bytes must not be silently reserialized +- Only core's typed, effect-free admission may return `Declined`; conversion errors and all post-admission failures are terminal + - Admission cannot invoke hooks, acquire credentials, consume files/iterators, prepare requests or perform I/O + - Disabled/unavailable native execution or an admission decline may select legacy once; callback exceptions never authorize fallback or replay +- Use one ordinary inline `async def` driver in `litellm/rust_bridge/lifecycle.py`, with the native handle in `src/lifecycle.rs` + - Contract: `start`, `resume_value`, `resume_error`, idempotent `close`; explicitly tagged `Await`/`Complete` preserve awaitable final values + - Validate Created/Running/Suspended/Closed protocol states; core alone chooses lifecycle phases and result/error policy + - Defer effectful setup/context reads/timestamps until start; unstarted-handle destruction releases inputs independently of Python `finally` + - Catch only the selected await's errors; start/resume errors propagate, `GeneratorExit` closes without further awaits + - Inline hooks preserve caller task/thread/loop and context writes; `into_future` creates a separate task and cannot satisfy this contract + - Delivery follows the binding, not callable type; keep direct, awaited, worker, background and deferred behavior distinct +- Finalize fallible public response/error construction, replacements and metadata under core control before terminal dispatch + - Success/failure handler entry receives the exact selected public response/exception; logging projections/redaction/snapshots retain their own copy contracts + - Ordinary failure-callback errors cannot suppress later eligible sync/async callbacks or replace the mapped provider error; control-flow exceptions have phase-specific policy + - Dispatch errors never replay provider work/accepted dispatch or trigger the opposite outcome; proxy acceptance/rejection releases core-owned deferred success at most once +- Make ownership safe across suspension, re-entry, cancellation and GC + - Keep native provider state typed in core; do not shuttle it through opaque Python transport/response classes + - Prefer one retained `Py` via `PyErr::into_value(py)`; reconstruct transient `PyErr`s, preserving identity, traceback, cause and context + - Traverse every owned Python edge, including duplicate references; traversal cannot call Python + - Take state out and mark Running under a short borrow, release borrows/locks before Python invocation, publish terminal state before finalizer-capable drops + - Close/GC/deferred release are idempotent and re-entry-safe, including during Rust unwinding; release only owned references, never clear caller containers or mask the selected error + - Cancellation signaling is not termination; retain captures until work actually finishes and use a Rust-selected awaited acknowledgement where required, never synchronous close/GC +- Verify behavior through a fresh, provenance-checked installed extension and positive native execution evidence before replacing the custom coroutine + - Cover admitted provider workflows, binding/read-point/identity behavior, failure continuation, finalization, no replay, deferred gates, re-entry, GC and cancellation termination + - Measure real conversion/copy costs before optimizing; preserve input contracts and capture lifetimes with `PyBackedBytes`, and lookup timing when interning names + - Ship accurate `_native.pyi` declarations and typing markers; distinguish Future-returning bindings from coroutine-returning bindings +- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [exception transfer](https://docs.rs/pyo3/0.29.2/pyo3/struct.PyErr.html#method.into_value), [re-entry](https://pyo3.rs/v0.29.2/class/call.html) + - [GIL policy](https://pyo3.rs/v0.29.2/free-threading.html), [experimental async limits](https://pyo3.rs/v0.29.2/async-await.html), [task conversion](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/fn.into_future_with_locals.html), [native cancellation/delivery](https://docs.rs/pyo3-async-runtimes/0.29.0/pyo3_async_runtimes/tokio/fn.future_into_py.html) + - [performance](https://pyo3.rs/v0.29.2/performance.html), [PyBackedBytes](https://docs.rs/pyo3/0.29.2/pyo3/pybacked/struct.PyBackedBytes.html), [typing](https://pyo3.rs/v0.29.2/python-typing-hints.html) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 337a1e8e5ac..42fad740870 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -17,7 +17,6 @@ panic-test = [] trace-parity = [ "dep:tracing", "litellm-core/observability", - "litellm-ai-gateway/trace-parity", ] [dependencies] @@ -25,7 +24,6 @@ futures-util.workspace = true tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-token-counter.workspace = true -litellm-ai-gateway = { workspace = true, default-features = false } litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true @@ -35,6 +33,7 @@ tokio = { workspace = true, features = ["sync"] } [dev-dependencies] criterion.workspace = true +rstest.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/auth.rs b/litellm-rust/crates/python-bridge/src/auth.rs new file mode 100644 index 00000000000..8dc0b7aabf0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/auth.rs @@ -0,0 +1,194 @@ +use litellm_core::auth::{ResolvedCredential, SecretValue}; +use pyo3::exceptions::{PyException, PyRuntimeError, PyTypeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::PyString; + +#[derive(Clone, Copy)] +pub(crate) struct TokenProviderContract { + callable_error: &'static str, + token_type_error: &'static str, + callback_error: &'static str, +} + +pub(crate) const AZURE_AD_TOKEN_PROVIDER: TokenProviderContract = TokenProviderContract { + callable_error: "Azure AD token provider must be callable", + token_type_error: "Azure AD token must be a string, got {}", + callback_error: "Failed to get Azure AD token: {}", +}; + +pub(crate) struct PythonTokenProvider { + callback: Py, + contract: TokenProviderContract, +} + +impl PythonTokenProvider { + pub(crate) fn select( + provider: Bound<'_, PyAny>, + contract: TokenProviderContract, + ) -> Option { + (provider.is_callable() && provider.is_truthy().unwrap_or(false)).then(|| Self { + callback: provider.unbind(), + contract, + }) + } + + pub(crate) fn acquire(&self, py: Python<'_>) -> PyResult { + let provider = self.callback.bind(py); + if !provider.is_callable() { + return Err(PyTypeError::new_err(self.contract.callable_error)); + } + let token = (|| { + let token = provider.call0()?; + if !token.is_instance_of::() { + let message = PyString::new(py, self.contract.token_type_error) + .call_method1("format", (token.get_type(),))?; + return Err(PyTypeError::new_err(message.unbind())); + } + Ok(token) + })() + .map_err(|error| { + if error.is_instance_of::(py) || !error.is_instance_of::(py) { + return error; + } + match PyString::new(py, self.contract.callback_error) + .call_method1("format", (error.value(py),)) + { + Ok(message) => { + let wrapped = PyRuntimeError::new_err(message.unbind()); + wrapped.set_context(py, Some(error.clone_ref(py))); + wrapped.set_cause(py, Some(error)); + wrapped + } + Err(format_error) => { + format_error.set_context(py, Some(error)); + format_error + } + } + })?; + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new(token.extract::()?), + expires_on: None, + }) + } + + pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.callback) + } +} + +#[cfg(test)] +mod tests { + use pyo3::exceptions::PyRuntimeError; + use pyo3::types::PyDict; + + use super::*; + + #[test] + fn token_callback_preserves_exception_identity_and_explicit_chaining() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +class ProviderError(Exception): + def __format__(self, specification): + return 'unavailable' +ordinary = ProviderError('must use __format__') +type_error = TypeError('signature') +abort = KeyboardInterrupt('cancelled') +def provider(error): + def acquire(): + raise error + return acquire +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + for name in ["ordinary", "type_error", "abort"] { + let original = locals.get_item(name).unwrap().unwrap(); + let callback = locals + .get_item("provider") + .unwrap() + .unwrap() + .call1((&original,)) + .unwrap(); + let provider = + PythonTokenProvider::select(callback, AZURE_AD_TOKEN_PROVIDER).unwrap(); + let error = provider.acquire(py).unwrap_err(); + if name == "ordinary" { + assert!(error.is_instance_of::(py)); + assert!(error.cause(py).unwrap().value(py).is(&original)); + assert!( + error + .value(py) + .getattr("__context__") + .unwrap() + .is(&original) + ); + assert_eq!( + error.value(py).str().unwrap().to_str().unwrap(), + "Failed to get Azure AD token: unavailable" + ); + } else { + assert!(error.value(py).is(&original)); + } + } + }); + } + + #[test] + fn invalid_token_type_formatting_preserves_python_failure_semantics() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +failure = ValueError('formatting failed') +class TokenType(type): + def __format__(cls, specification): + raise failure +class Token(metaclass=TokenType): + pass +def provider(): + return Token() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let provider = PythonTokenProvider::select( + locals.get_item("provider").unwrap().unwrap(), + AZURE_AD_TOKEN_PROVIDER, + ) + .unwrap(); + let error = provider.acquire(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!( + error + .cause(py) + .unwrap() + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn token_string_extraction_errors_are_not_wrapped_as_callback_failures() { + Python::initialize(); + Python::attach(|py| { + let callback = py + .eval(pyo3::ffi::c_str!("lambda: '\\ud800'"), None, None) + .unwrap(); + let provider = PythonTokenProvider::select(callback, AZURE_AD_TOKEN_PROVIDER).unwrap(); + let error = provider.acquire(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index e1f458ea0bc..701c6abb68c 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -22,7 +22,8 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { Error::InvalidProvider(_) | Error::InvalidRequest(_) | Error::InvalidType { .. } - | Error::MissingField(_) => PyValueError::new_err(err.to_string()), + | Error::MissingField(_) + | Error::MissingDocumentUrl => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } @@ -41,6 +42,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidRequest(_) | Error::InvalidType { .. } | Error::MissingField(_) + | Error::MissingDocumentUrl | Error::MissingApiKey { .. } | Error::MissingAzureAiCredentials | Error::MissingAzureDocumentIntelligenceCredentials @@ -49,9 +51,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - Error::Http { status, body } => { - RustUpstreamError::new_err((status, format!("{status}: {body}"))) - } + Error::Http { status, body } => RustUpstreamError::new_err((status, body)), Error::Network(message) | Error::InvalidResponse(message) => { RustUpstreamError::new_err((0u16, message)) } @@ -63,41 +63,3 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } - -pub(crate) fn ocr_error_to_pyerr(err: Error) -> PyErr { - match err { - Error::MissingField("document_url" | "image_url") => { - PyValueError::new_err("Document URL is required") - } - Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - other => core_error_to_pyerr(other), - } -} - -#[cfg(test)] -mod ocr_error_tests { - use super::*; - - #[test] - fn ocr_errors_preserve_python_validation_and_provider_details() { - Python::initialize(); - Python::attach(|py| { - for field in ["document_url", "image_url"] { - let mapped = ocr_error_to_pyerr(Error::MissingField(field)); - assert!(mapped.is_instance_of::(py)); - assert_eq!(mapped.value(py).to_string(), "Document URL is required"); - } - let mapped = ocr_error_to_pyerr(Error::Http { - status: 429, - body: r#"{"message":"rate limited"}"#.to_string(), - }); - assert!(mapped.is_instance_of::(py)); - let args: (u16, String) = mapped - .value(py) - .getattr("args") - .and_then(|args| args.extract()) - .expect("OCR failures retain status and unprefixed provider message"); - assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index b57197b9ddf..d8dda10068d 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -1,5 +1,7 @@ use std::future::Future; use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::task::{Context, Poll, Waker}; use std::time::Duration; use futures_util::FutureExt; @@ -28,6 +30,27 @@ where ) } +pub(crate) fn run_sync_value(py: Python<'_>, future: F) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future) +} + +fn run_sync_value_on(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + if Handle::try_current().is_ok() { + return Err(PyRuntimeError::new_err( + "synchronous native routes cannot run from a Tokio context; use the async route", + )); + } + release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))? +} + fn run_sync_on( py: Python<'_>, runtime: &Runtime, @@ -67,6 +90,32 @@ where }) } +pub(crate) fn run_async_value(py: Python<'_>, future: F) -> PyResult> +where + T: for<'py> IntoPyObject<'py> + Send + 'static, + F: Future> + Send + 'static, +{ + pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? }) +} + +pub(crate) fn poll_async_value(py: Python<'_>, future: Pin<&mut F>) -> PyResult> +where + T: Send, + F: Future> + Send, +{ + let result = release_gil(py, || { + let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter(); + std::panic::catch_unwind(AssertUnwindSafe(|| { + future.poll(&mut Context::from_waker(Waker::noop())) + })) + .map_err(panic_to_pyerr) + })?; + match result { + Poll::Ready(result) => result.map(Poll::Ready), + Poll::Pending => Ok(Poll::Pending), + } +} + fn map_core_result(result: Result, map_error: fn(E) -> PyErr) -> PyResult { match result { Ok(value) => Ok(value), @@ -119,11 +168,30 @@ mod tests { use litellm_core::error::Error; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; + use rstest::{fixture, rstest}; use serde::Serializer; use tokio::runtime::Builder; use super::*; + struct InitializedPython; + + impl InitializedPython { + fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } + } + + #[fixture] + #[once] + fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython + } + fn runtime_error(error: Error) -> PyErr { PyRuntimeError::new_err(error.to_string()) } @@ -194,10 +262,84 @@ mod tests { .expect("result should convert") } - #[test] - fn sync_runner_polls_future_on_the_caller_thread() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn inline_poll_releases_gil_and_enters_runtime( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let (sender, receiver) = mpsc::sync_channel(1); + let worker = thread::spawn(move || Python::attach(|_| sender.send(()).unwrap())); + let mut future = Box::pin(async move { + receiver.recv_timeout(Duration::from_secs(2)).unwrap(); + Ok(Handle::try_current().is_ok()) + }); + assert_eq!( + poll_async_value(py, future.as_mut()).unwrap(), + Poll::Ready(true) + ); + worker.join().unwrap(); + }); + } + + #[rstest] + fn inline_poll_contains_panics_and_preserves_python_errors( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let mut panicking = Box::pin(poll_fn(|_| -> Poll> { + panic!("inline native panic") + })); + let error = poll_async_value(py, panicking.as_mut()).unwrap_err(); + assert!(error.is_instance_of::(py)); + let original = PyRuntimeError::new_err("inline failure"); + let identity = original.value(py).clone().unbind(); + let mut failing = Box::pin(async move { Err::<(), _>(original) }); + let error = poll_async_value(py, failing.as_mut()).unwrap_err(); + assert!(error.value(py).is(identity.bind(py))); + }); + } + + #[pyfunction] + fn pending_after_inline_poll(py: Python<'_>) -> PyResult> { + let starts = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&starts); + let mut future = Box::pin(async move { + starts.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(5)).await; + Ok(starts.load(Ordering::SeqCst)) + }); + assert!(poll_async_value(py, future.as_mut())?.is_pending()); + assert_eq!(observed.load(Ordering::SeqCst), 1); + run_async_value(py, future) + } + + #[rstest] + fn inline_pending_future_resumes_on_tokio_without_restarting( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "pending", + wrap_pyfunction!(pending_after_inline_poll, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + "import asyncio\nasync def exercise():\n assert await asyncio.wait_for(pending(), 2) == 1\nasyncio.run(exercise())" + ), + Some(&locals), + Some(&locals), + ).unwrap(); + }); + } + + #[rstest] + fn sync_runner_polls_future_on_the_caller_thread( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let caller_thread = std::thread::current().id(); let result = run_sync( py, @@ -209,10 +351,11 @@ mod tests { }); } - #[test] - fn sync_runner_releases_gil_while_waiting() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_releases_gil_while_waiting( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let result = run_sync( py, async { @@ -230,16 +373,17 @@ mod tests { }); } - #[test] - fn sync_runner_rejects_calls_from_a_tokio_context() { - Python::initialize(); + #[rstest] + fn sync_runner_rejects_calls_from_a_tokio_context( + #[from(initialized_python)] python: &InitializedPython, + ) { let runtime = Builder::new_current_thread() .enable_all() .build() .expect("runtime should build"); let error = runtime.block_on(async { - Python::attach(|py| { + python.attach(|py| { run_sync::(py, async { Ok(true) }, runtime_error) .expect_err("sync route should reject a nested Tokio runtime") }) @@ -251,14 +395,15 @@ mod tests { ); } - #[test] - fn sync_runner_can_drive_a_current_thread_runtime() { - Python::initialize(); + #[rstest] + fn sync_runner_can_drive_a_current_thread_runtime( + #[from(initialized_python)] python: &InitializedPython, + ) { let runtime = Builder::new_current_thread() .enable_all() .build() .expect("runtime should build"); - Python::attach(|py| { + python.attach(|py| { let result = run_sync_on( py, &runtime, @@ -272,10 +417,9 @@ mod tests { }); } - #[test] - fn sync_runner_maps_a_panicked_future() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_maps_a_panicked_future(#[from(initialized_python)] python: &InitializedPython) { + python.attach(|py| { let error = run_sync::( py, poll_fn(|_| -> Poll> { panic!("route future panicked") }), @@ -288,10 +432,11 @@ mod tests { }); } - #[test] - fn sync_runner_maps_a_panicked_error_mapper() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_maps_a_panicked_error_mapper( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let error = run_sync::( py, async { Err(Error::InvalidRequest("invalid".to_string())) }, @@ -304,10 +449,11 @@ mod tests { }); } - #[test] - fn sync_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_surfaces_serializer_panics( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) .expect_err("serializer panic should become a Python exception"); @@ -316,9 +462,10 @@ mod tests { }); } - #[test] - fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { - Python::initialize(); + #[rstest] + fn sync_runner_supports_concurrent_callers_on_the_shared_runtime( + #[from(initialized_python)] _python: &InitializedPython, + ) { let barrier = Arc::new(tokio::sync::Barrier::new(2)); let callers: Vec<_> = (0..2) .map(|_| { @@ -349,10 +496,11 @@ mod tests { assert_eq!(results, vec![true, true]); } - #[test] - fn async_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn async_runner_surfaces_serializer_panics( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let module = PyModule::new(py, "runtime").expect("module should be created"); module .add_function( @@ -386,11 +534,12 @@ asyncio.run(exercise()) }); } - #[test] - fn async_result_delivery_does_not_stall_tokio_workers() { - Python::initialize(); + #[rstest] + fn async_result_delivery_does_not_stall_tokio_workers( + #[from(initialized_python)] python: &InitializedPython, + ) { ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); - Python::attach(|py| { + python.attach(|py| { let module = PyModule::new(py, "runtime").expect("module should be created"); for function in [ wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"), diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index cf0450a1b30..12bc57a8931 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,14 +1,16 @@ +mod auth; mod constants; mod diagnostics; mod errors; mod execution; #[cfg(feature = "trace-parity")] mod function_trace; +mod lifecycle; mod marshal; mod routes; mod token_counter; -use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; +use litellm_core::responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use pyo3::types::PyAny; use serde_json::Value; @@ -64,7 +66,7 @@ impl ResponsesWebSocketConnection { } } -#[pymodule(gil_used = false)] +#[pymodule(gil_used = true)] mod _native { use pyo3::prelude::*; @@ -152,7 +154,6 @@ mod tests { "amessages", "chat_completions", "achat_completions", - "gateway_messages", ] ); } diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs new file mode 100644 index 00000000000..06b32b67fd5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/bindings.rs @@ -0,0 +1,391 @@ +use pyo3::exceptions::PyBaseException; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +#[derive(FromPyObject)] +pub(crate) struct PythonLogger(Py); + +impl PythonLogger { + pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> { + self.0.bind(py) + } + + pub(crate) fn clone_ref(&self, py: Python<'_>) -> Self { + Self(self.0.clone_ref(py)) + } + + pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + + pub(crate) fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { + if !self + .object(py) + .getattr("_native_callback_fast_path") + .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + { + return Ok(true); + } + py.import("litellm.rust_bridge.lifecycle")? + .getattr("callbacks_needed")? + .call1((self.object(py), phase))? + .extract() + } + + pub(super) fn success_bookkeeping( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + asynchronous: bool, + ) -> PyResult<()> { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("success_bookkeeping")? + .call1((self.object(py), response, start, end, asynchronous))?; + Ok(()) + } + + pub(super) fn defers_async_logging(&self, py: Python<'_>) -> bool { + self.object(py) + .getattr("_defer_async_logging") + .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + } + + pub(super) fn defer_success( + &self, + py: Python<'_>, + pending: Py, + ) -> PyResult<()> { + self.object(py).setattr("_native_pending_logging", pending) + } + + pub(super) fn sync_success_for_async_call( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + if !self.callbacks_needed(py, "sync_success_async")? { + return Ok(()); + } + self.object(py).call_method1( + "handle_sync_success_callbacks_for_async_calls", + (response, start, end), + )?; + Ok(()) + } + + pub(super) fn failure( + &self, + py: Python<'_>, + error: &Py, + start: &Py, + end: &Option>, + asynchronous: bool, + ) -> PyResult>> { + if !self.callbacks_needed( + py, + if asynchronous { + "async_failure" + } else { + "sync_failure" + }, + )? { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("failure_bookkeeping")? + .call1((self.object(py), error, start, end, asynchronous))?; + return Ok(None); + } + let trace = py + .import("traceback")? + .getattr("format_exception")? + .call1((error,))?; + let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?; + let value = self.object(py).call_method1( + if asynchronous { + "async_failure_handler" + } else { + "failure_handler" + }, + (error, trace, start, end), + )?; + Ok(asynchronous.then(|| value.unbind())) + } + + pub(super) fn restore_context(&self, py: Python<'_>) -> PyResult<()> { + py.import("litellm.utils")? + .getattr("_restore_correlation_context_if_supported")? + .call1((self.object(py),))?; + Ok(()) + } + + pub(super) fn submit_success( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + if !self.callbacks_needed(py, "sync_success")? { + return self.success_bookkeeping(py, response, start, end, false); + } + let context = py.import("contextvars")?.call_method0("copy_context")?; + py.import("litellm.litellm_core_utils.litellm_logging")? + .getattr("executor")? + .call_method1( + "submit", + ( + context.getattr("run")?, + self.object(py).getattr("success_handler")?, + response, + start, + end, + ), + )?; + Ok(()) + } + + pub(super) fn enqueue_success( + &self, + py: Python<'_>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> PyResult<()> { + if !self.callbacks_needed(py, "async_success")? { + return self.success_bookkeeping(py, response, start, end, true); + } + let context = py.import("contextvars")?.call_method0("copy_context")?; + let worker = py + .import("litellm.litellm_core_utils.logging_worker")? + .getattr("GLOBAL_LOGGING_WORKER")? + .getattr("ensure_initialized_and_enqueue")?; + let coroutine = self + .object(py) + .call_method1("async_success_handler", (response, start, end))?; + let enqueue = context.call_method1("run", (worker, &coroutine)); + if enqueue.is_err() + && let Err(error) = coroutine.call_method0("close") + { + error.write_unraisable(py, Some(&coroutine)); + } + enqueue.map(|_| ()) + } +} + +pub(super) struct SetupResult<'py>(Bound<'py, PyAny>); + +impl SetupResult<'_> { + pub(super) fn logger(&self) -> PyResult { + self.0.getattr("logger")?.extract() + } + + pub(super) fn kwargs(&self) -> PyResult> { + Ok(self.0.getattr("kwargs")?.extract()?) + } +} + +pub(super) fn setup<'py>( + py: Python<'py>, + call_type: &str, + args: &Py, + kwargs: &Py, + start: &Py, + asynchronous: bool, +) -> PyResult> { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("setup")? + .call1((call_type, args, kwargs, start, asynchronous)) + .map(SetupResult) +} + +pub(super) fn finalize( + py: Python<'_>, + response: &Option>, + logger: &PythonLogger, + kwargs: &Py, + start: &Py, + end: &Option>, +) -> PyResult<()> { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("finalize")? + .call1((response, logger.object(py), kwargs, start, end))?; + Ok(()) +} + +pub(super) fn is_internal_call(py: Python<'_>) -> PyResult { + py.import("litellm._internal_context")? + .getattr("is_internal_call")? + .call_method0("get")? + .extract() +} + +pub(super) struct DeploymentHooks; + +impl DeploymentHooks { + pub(super) fn needed(py: Python<'_>) -> PyResult { + py.import("litellm.rust_bridge.lifecycle")? + .getattr("deployment_callbacks_needed")? + .call0()? + .extract() + } + + pub(super) fn before_call( + py: Python<'_>, + kwargs: &Py, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_pre_call_deployment_hook")? + .call1((kwargs, call_type)) + .map(Bound::unbind) + } + + pub(super) fn after_success( + py: Python<'_>, + kwargs: &Py, + response: &Option>, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_post_call_success_deployment_hook")? + .call1((kwargs, response, call_type)) + .map(Bound::unbind) + } + + pub(super) fn after_failure( + py: Python<'_>, + kwargs: &Py, + error: &Py, + call_type: &str, + ) -> PyResult> { + py.import("litellm.utils")? + .getattr("async_post_call_failure_deployment_hook")? + .call1((kwargs, error, call_type)) + .map(Bound::unbind) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyTypeError; + + #[test] + fn setup_fields_are_checked_in_order_without_eager_logger_method_reads() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +reads = [] +class Logger: + def __getattribute__(self, name): + reads.append(name) + raise AssertionError('logger methods must remain lazy') +logger = Logger() +class Setup: + @property + def logger(self): + reads.append('logger') + return logger + @property + def kwargs(self): + reads.append('kwargs') + return [] +result = Setup() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let result = SetupResult(locals.get_item("result").unwrap().unwrap()); + let logger = result.logger().unwrap(); + assert!( + logger + .object(py) + .is(locals.get_item("logger").unwrap().unwrap()) + ); + assert_eq!( + locals + .get_item("reads") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + ["logger"] + ); + assert!( + result + .kwargs() + .unwrap_err() + .is_instance_of::(py) + ); + assert_eq!( + locals + .get_item("reads") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + ["logger", "kwargs"] + ); + }); + } + + #[test] + fn logger_resolves_each_callback_at_invocation_and_preserves_arguments() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +calls = [] +response, start, end = object(), object(), object() +class Logger: + @property + def handle_sync_success_callbacks_for_async_calls(self): + generation = len(calls) + def callback(*args): + assert args == (response, start, end) + calls.append(generation) + return callback +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let logger: PythonLogger = locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(); + let response = Some(locals.get_item("response").unwrap().unwrap().unbind()); + let start = locals.get_item("start").unwrap().unwrap().unbind(); + let end = Some(locals.get_item("end").unwrap().unwrap().unbind()); + for _ in 0..2 { + logger + .sync_success_for_async_call(py, &response, &start, &end) + .unwrap(); + } + assert_eq!( + locals + .get_item("calls") + .unwrap() + .unwrap() + .extract::>() + .unwrap(), + [0, 1] + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs b/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs new file mode 100644 index 00000000000..17a480a7225 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/handle.rs @@ -0,0 +1,139 @@ +use std::panic::{AssertUnwindSafe, catch_unwind}; + +use litellm_python_interop::panic_to_pyerr; +use pyo3::exceptions::{PyBaseException, PyRuntimeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; + +pub(super) enum ExecutionStep { + Return(Py), + Await(Py), +} + +pub(super) trait ExecutionBody: Send + Sync { + fn resume(&mut self, result: Option>>) -> PyResult; + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; +} + +enum ExecutionState { + Created(Box), + Running, + Suspended(Box), + Closed, +} + +#[pyclass] +pub(super) struct Execution { + state: ExecutionState, +} + +impl Execution { + pub(super) fn new(body: impl ExecutionBody + 'static) -> Self { + Self { + state: ExecutionState::Created(Box::new(body)), + } + } + + fn advance( + slf: &Bound<'_, Self>, + py: Python<'_>, + result: Option>>, + ) -> PyResult> { + let mut body = { + let mut execution = slf.borrow_mut(); + match (&execution.state, result.is_some()) { + (ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {} + (ExecutionState::Running, _) => { + return Err(PyRuntimeError::new_err("execution is already running")); + } + (ExecutionState::Closed, _) => { + return Err(PyRuntimeError::new_err("execution is closed")); + } + _ => { + return Err(PyRuntimeError::new_err( + "execution requires start before resume and can only start once", + )); + } + } + match std::mem::replace(&mut execution.state, ExecutionState::Running) { + ExecutionState::Created(body) | ExecutionState::Suspended(body) => body, + _ => unreachable!(), + } + }; + let outcome = catch_unwind(AssertUnwindSafe(|| { + let step = body.resume(result)?; + let (tag, value, suspended) = match step { + ExecutionStep::Await(value) => ("Await", value, true), + ExecutionStep::Return(value) => ("Complete", value, false), + }; + let step = py + .import("litellm.rust_bridge.lifecycle")? + .getattr(tag)? + .call1((value,))? + .unbind(); + Ok((step, suspended)) + })) + .map_err(panic_to_pyerr) + .and_then(|result| result); + match outcome { + Ok((step, true)) if matches!(slf.borrow().state, ExecutionState::Running) => { + slf.borrow_mut().state = ExecutionState::Suspended(body); + Ok(step) + } + outcome => { + slf.borrow_mut().state = ExecutionState::Closed; + drop(body); + outcome.and_then(|(step, suspended)| { + if suspended { + Err(PyRuntimeError::new_err( + "execution was closed while running", + )) + } else { + Ok(step) + } + }) + } + } + } +} + +#[pymethods] +impl Execution { + fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult> { + Self::advance(slf, py, None) + } + + fn resume_value( + slf: &Bound<'_, Self>, + py: Python<'_>, + value: Py, + ) -> PyResult> { + Self::advance(slf, py, Some(Ok(value))) + } + + fn resume_error( + slf: &Bound<'_, Self>, + py: Python<'_>, + error: Bound<'_, PyBaseException>, + ) -> PyResult> { + Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any())))) + } + + fn close(slf: &Bound<'_, Self>) { + let state = std::mem::replace(&mut slf.borrow_mut().state, ExecutionState::Closed); + drop(state); + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + match &self.state { + ExecutionState::Created(body) | ExecutionState::Suspended(body) => { + body.traverse(&visit) + } + _ => Ok(()), + } + } + + fn __clear__(slf: &Bound<'_, Self>) { + Self::close(slf); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs new file mode 100644 index 00000000000..014564ae89d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/mod.rs @@ -0,0 +1,1175 @@ +use std::sync::Arc; +use std::task::Poll; + +use futures_util::future::{AbortHandle, Abortable}; +#[cfg(test)] +use litellm_core::call_lifecycle::host::HostCallFuture; +use litellm_core::call_lifecycle::host::{ + HostCall as NativeCall, HostCallStep as NativeCallStep, HostFailure, HostPhase, HostStep, +}; +use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; +use pyo3::gc::{PyTraverseError, PyVisit}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; +use tokio::sync::Mutex; + +use crate::execution::{poll_async_value, run_async_value, run_sync_value}; + +mod bindings; +mod handle; +mod preparation; + +use bindings::DeploymentHooks; +pub(crate) use bindings::PythonLogger; +use handle::{Execution, ExecutionBody, ExecutionStep}; + +pub(crate) enum OperationClass { + Phase(HostPhase), + Route, +} + +pub(crate) trait PythonRoute: Send + Sync { + type Call: NativeCall + 'static; + + fn state(&self) -> &PythonCallState; + fn state_mut(&mut self) -> &mut PythonCallState; + fn classify(operation: &::Operation) -> OperationClass; + fn lifecycle_result() -> ::Result; + fn map_error(error: litellm_core::Error) -> PyErr; + fn invoke( + &mut self, + py: Python<'_>, + operation: ::Operation, + ) -> PyResult<::Result>; + fn cleanup(&mut self); + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; +} + +type NativeStep = NativeCallStep<::Operation, ::Complete>; +type NativeResult = Result, litellm_core::Error>; +type HostResumeStep = HostStep::Call>, Py>; + +struct NativeCallState { + call: C, + result: Option>, +} + +enum PendingOperation { + Native, + Host(HostPhase), +} + +struct PythonLifecycle { + route: R, + call: Option>>>, + pending: Option, + native_abort: Option, +} + +pub(crate) fn run_call( + py: Python<'_>, + call: R::Call, + route: R, +) -> PyResult> { + let asynchronous = route.state().asynchronous; + let mut lifecycle = PythonLifecycle { + route, + call: Some(Arc::new(Mutex::new(NativeCallState { call, result: None }))), + pending: None, + native_abort: None, + }; + if asynchronous { + let execution = Py::new(py, Execution::new(lifecycle))?; + return py + .import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) + .map(Bound::unbind); + } + match lifecycle.resume(None)? { + ExecutionStep::Return(value) => Ok(value), + ExecutionStep::Await(_) => Err(pyo3::exceptions::PyRuntimeError::new_err( + "sync call suspended", + )), + } +} + +pub(crate) fn missing_state() -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err("missing native call state") +} + +impl PythonLifecycle { + fn resume_core( + &mut self, + py: Python<'_>, + result: Option::Result, HostFailure>>, + ) -> PyResult> { + let call = Arc::clone(self.call.as_ref().ok_or_else(missing_state)?); + let future = async move { + let mut call = call.lock().await; + let result = match result { + Some(Err(failure)) => call.call.interrupt(failure).await, + Some(Ok(result)) => call.call.resume(Some(result)).await, + None => call.call.resume(None).await, + }; + call.result = Some(result); + Ok(()) + }; + if self.route.state().asynchronous { + let mut future = Box::pin(future); + if let Poll::Ready(()) = poll_async_value(py, future.as_mut())? { + return Ok(HostStep::Ready(self.take_native_result()?)); + } + let (abort, registration) = AbortHandle::new_pair(); + self.native_abort = Some(abort); + self.pending = Some(PendingOperation::Native); + Ok(HostStep::Suspend( + run_async_value(py, async move { + Abortable::new(future, registration) + .await + .map_err(|_| PyRuntimeError::new_err("native execution closed"))? + })? + .unbind(), + )) + } else { + run_sync_value(py, future)?; + Ok(HostStep::Ready(self.take_native_result()?)) + } + } + + fn take_native_result(&self) -> PyResult> { + self.call + .as_ref() + .ok_or_else(missing_state)? + .try_lock() + .map_err(|_| missing_state())? + .result + .take() + .ok_or_else(missing_state)? + .map_err(R::map_error) + } + + fn host_failure( + &mut self, + py: Python<'_>, + error: PyErr, + phase: Option, + ) -> HostFailure { + let native = litellm_core::Error::InvalidRequest(error.to_string()); + let cancelled = !error.is_instance_of::(py); + let failure = if !cancelled { + HostFailure::Error(native) + } else { + HostFailure::Cancelled(native) + }; + let state = self.route.state_mut(); + if state.error.is_none() || (cancelled && phase != Some(HostPhase::DeploymentFailure)) { + state.retain_error(py, error); + } + if state.end.is_none() { + state.end = now(py).ok(); + } + failure + } + + fn drive( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + let mut step = match (self.pending.take(), result) { + (None, None) => self.resume_core(py, None)?, + (Some(PendingOperation::Native), Some(result)) => match result { + Ok(_) => HostStep::Ready(self.take_native_result()?), + Err(error) => { + let failure = self.host_failure(py, error, None); + self.resume_core(py, Some(Err(failure)))? + } + }, + (Some(PendingOperation::Host(phase)), Some(result)) => { + let result = + result.and_then(|value| self.route.state_mut().accept(py, phase, value)); + let result = match result { + Ok(()) => Ok(R::lifecycle_result()), + Err(error) => Err(self.host_failure(py, error, Some(phase))), + }; + self.resume_core(py, Some(result))? + } + _ => return Err(missing_state()), + }; + loop { + let operation = match step { + HostStep::Suspend(awaitable) => return Ok(ExecutionStep::Await(awaitable)), + HostStep::Ready(NativeCallStep::Complete(_)) => { + return self + .route + .state_mut() + .response + .take() + .map(ExecutionStep::Return) + .ok_or_else(missing_state); + } + HostStep::Ready(NativeCallStep::Host(operation)) => operation, + }; + let phase = match R::classify(&operation) { + OperationClass::Phase(phase) => Some(phase), + OperationClass::Route => None, + }; + let result = match phase { + Some(phase) => match self.route.state_mut().invoke(py, phase) { + Ok(HostStep::Suspend(awaitable)) => { + self.pending = Some(PendingOperation::Host(phase)); + return Ok(ExecutionStep::Await(awaitable)); + } + Ok(HostStep::Ready(value)) => self + .route + .state_mut() + .accept(py, phase, value) + .map(|()| R::lifecycle_result()), + Err(error) => Err(error), + }, + None => self.route.invoke(py, operation), + }; + let result = match result { + Ok(result) => Ok(result), + Err(error) => Err(self.host_failure(py, error, phase)), + }; + step = self.resume_core(py, Some(result))?; + } + } +} + +impl ExecutionBody for PythonLifecycle { + fn resume(&mut self, result: Option>>) -> PyResult { + let result = Python::attach(|py| self.drive(py, result)); + match result { + Ok(ExecutionStep::Await(value)) => Ok(ExecutionStep::Await(value)), + result => result.map_err(|error| { + Python::attach(|py| { + self.route + .state_mut() + .error + .take() + .map(|value| PyErr::from_value(value.into_bound(py).into_any())) + .unwrap_or(error) + }) + }), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.route.state().traverse(visit)?; + self.route.traverse(visit) + } +} + +impl PythonLifecycle { + fn clear(&mut self) { + if let Some(abort) = self.native_abort.take() { + abort.abort(); + } + if self.call.take().is_some() { + Python::attach(|py| self.route.state_mut().cleanup(py)); + self.route.cleanup(); + } + } +} + +impl Drop for PythonLifecycle { + fn drop(&mut self) { + self.clear(); + } +} + +pub(crate) struct PythonCallState { + pub args: Py, + pub kwargs: Py, + pub logger: Option, + pub start: Py, + pub end: Option>, + pub response: Option>, + pub error: Option>, + pub asynchronous: bool, + pub internal: bool, + pub call_type: &'static str, +} + +pub(crate) fn now(py: Python<'_>) -> PyResult> { + py.import("datetime")? + .getattr("datetime")? + .call_method0("now") + .map(Bound::unbind) +} + +impl PythonCallState { + fn invoke( + &mut self, + py: Python<'_>, + phase: HostPhase, + ) -> PyResult, Py>> { + match phase { + HostPhase::Setup => self.setup(py)?, + HostPhase::DeploymentPreCall => { + if !DeploymentHooks::needed(py)? { + return Ok(HostStep::Ready(self.kwargs.clone_ref(py).into_any())); + } + return Ok(HostStep::Suspend(DeploymentHooks::before_call( + py, + &self.kwargs, + self.call_type, + )?)); + } + HostPhase::Prepare => self.prepare(py)?, + HostPhase::DeploymentPostCall => { + if !DeploymentHooks::needed(py)? { + return self + .response + .as_ref() + .map(|value| HostStep::Ready(value.clone_ref(py))) + .ok_or_else(missing_state); + } + return Ok(HostStep::Suspend(DeploymentHooks::after_success( + py, + &self.kwargs, + &self.response, + self.call_type, + )?)); + } + HostPhase::Finalize => self.finalize(py)?, + HostPhase::Success => self.dispatch_success(py)?, + HostPhase::DeploymentFailure => { + if let Some(error) = &self.error + && DeploymentHooks::needed(py)? + { + return Ok(HostStep::Suspend(DeploymentHooks::after_failure( + py, + &self.kwargs, + error, + self.call_type, + )?)); + } + } + HostPhase::Failure | HostPhase::AsyncFailure => { + if let Some(awaitable) = + self.dispatch_failure(py, phase == HostPhase::AsyncFailure)? + { + return Ok(HostStep::Suspend(awaitable)); + } + } + HostPhase::Execute + | HostPhase::ConstructResponse + | HostPhase::MapFailure + | HostPhase::Complete => return Err(missing_state()), + } + Ok(HostStep::Ready(py.None())) + } + + fn accept(&mut self, py: Python<'_>, phase: HostPhase, value: Py) -> PyResult<()> { + match phase { + HostPhase::DeploymentPreCall => { + self.kwargs = value.into_bound(py).cast_into::()?.unbind() + } + HostPhase::DeploymentPostCall => self.response = Some(value), + _ => {} + } + Ok(()) + } + + pub fn new( + py: Python<'_>, + args: Py, + kwargs: Py, + asynchronous: bool, + call_type: &'static str, + ) -> PyResult { + Ok(Self { + args, + kwargs, + logger: None, + start: py.None(), + end: None, + response: None, + error: None, + asynchronous, + internal: false, + call_type, + }) + } + + pub fn logger(&self) -> PyResult<&PythonLogger> { + self.logger.as_ref().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err("call logging is not initialized") + }) + } + + pub fn setup(&mut self, py: Python<'_>) -> PyResult<()> { + self.start = now(py)?; + self.internal = bindings::is_internal_call(py)?; + let result = bindings::setup( + py, + self.call_type, + &self.args, + &self.kwargs, + &self.start, + self.asynchronous, + )?; + self.logger = Some(result.logger()?); + self.kwargs = result.kwargs()?; + Ok(()) + } + + pub fn prepare(&mut self, py: Python<'_>) -> PyResult<()> { + self.kwargs = preparation::prepare(py, self.kwargs.bind(py), self.logger()?)?.unbind(); + Ok(()) + } + + pub fn finalize(&self, py: Python<'_>) -> PyResult<()> { + bindings::finalize( + py, + &self.response, + self.logger()?, + &self.kwargs, + &self.start, + &self.end, + ) + } + + pub fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> { + match self.try_dispatch_success(py) { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, self.logger.as_ref().map(|logger| logger.object(py))); + Ok(()) + } + result => result, + } + } + + fn try_dispatch_success(&self, py: Python<'_>) -> PyResult<()> { + let logger = self.logger()?; + let pending = || PendingSuccess { + logger: logger.clone_ref(py), + response: self.response.as_ref().map(|value| value.clone_ref(py)), + start: self.start.clone_ref(py), + end: self.end.as_ref().map(|value| value.clone_ref(py)), + }; + if !self.asynchronous { + if !logger.callbacks_needed(py, "sync_success")? { + return logger.success_bookkeeping( + py, + &self.response, + &self.start, + &self.end, + false, + ); + } + pending().sync(py) + } else { + if !self.internal + && self + .kwargs + .bind(py) + .get_item("fallbacks")? + .is_none_or(|value| value.is_none()) + { + if !logger.callbacks_needed(py, "async_success")? { + logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; + } else if logger.defers_async_logging(py) { + logger.defer_success( + py, + Py::new( + py, + PendingLogging { + pending: Some(pending()), + }, + )?, + )?; + } else { + pending().asynchronous(py)?; + } + } + logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) + } + } + + pub fn dispatch_failure( + &self, + py: Python<'_>, + asynchronous: bool, + ) -> PyResult>> { + if self.logger.is_none() || (self.asynchronous && self.internal) { + return Ok(None); + } + let Some(error) = &self.error else { + return Ok(None); + }; + self.logger()? + .failure(py, error, &self.start, &self.end, asynchronous) + } + + pub fn cleanup(&mut self, py: Python<'_>) { + if let Some(logger) = self.logger.take() + && let Err(error) = logger.restore_context(py) + { + error.write_unraisable(py, None); + } + } + + pub fn retain_error(&mut self, py: Python<'_>, error: PyErr) { + self.error = Some(error.into_value(py)); + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.args)?; + visit.call(&self.kwargs)?; + if let Some(logger) = &self.logger { + logger.traverse(visit)?; + } + visit.call(&self.start)?; + visit.call(&self.end)?; + visit.call(&self.response)?; + visit.call(&self.error) + } +} + +struct PendingSuccess { + logger: PythonLogger, + response: Option>, + start: Py, + end: Option>, +} + +impl PendingSuccess { + fn sync(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .submit_success(py, &self.response, &self.start, &self.end) + } + + fn asynchronous(&self, py: Python<'_>) -> PyResult<()> { + self.logger + .enqueue_success(py, &self.response, &self.start, &self.end) + } +} + +#[pyclass] +struct PendingLogging { + pending: Option, +} + +#[pymethods] +impl PendingLogging { + fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> { + let pending = slf.borrow_mut().pending.take(); + if let Some(pending) = pending + && success + { + match pending.asynchronous(py) { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(pending.logger.object(py))); + } + result => return result, + } + } + Ok(()) + } + + fn __traverse__(&self, visit: pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { + if let Some(pending) = &self.pending { + pending.logger.traverse(&visit)?; + visit.call(&pending.response)?; + visit.call(&pending.start)?; + visit.call(&pending.end)?; + } + Ok(()) + } + + fn __clear__(slf: &Bound<'_, Self>) { + let pending = slf.borrow_mut().pending.take(); + drop(pending); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::types::PyDict; + use std::sync::Mutex; + + static PYTHON_GLOBALS: Mutex<()> = Mutex::new(()); + + fn install_logging_worker(py: Python<'_>, worker: &Bound<'_, PyAny>) -> PyResult<()> { + py.import("litellm.litellm_core_utils.logging_worker")? + .setattr("GLOBAL_LOGGING_WORKER", worker) + } + + struct RetainingHost { + retained: Option>, + } + + impl ExecutionBody for RetainingHost { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| Ok(ExecutionStep::Return(py.None()))) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.retained) + } + } + + #[pyfunction] + fn retaining_coroutine(py: Python<'_>, retained: Py) -> PyResult> { + Py::new( + py, + Execution::new(RetainingHost { + retained: Some(retained), + }), + ) + } + + struct AwaitBody(Option>); + + impl ExecutionBody for AwaitBody { + fn resume(&mut self, result: Option>>) -> PyResult { + match self.0.take() { + Some(awaitable) => Ok(ExecutionStep::Await(awaitable)), + None => result + .expect("selected await completed") + .map(ExecutionStep::Return), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + } + + #[pyfunction] + fn await_execution(awaitable: Py) -> Execution { + Execution::new(AwaitBody(Some(awaitable))) + } + + struct CallingBody(Py); + + impl ExecutionBody for CallingBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + } + + #[pyfunction] + fn calling_execution(callback: Py) -> Execution { + Execution::new(CallingBody(callback)) + } + + struct SyntheticCall(bool); + + impl NativeCall for SyntheticCall { + type Operation = (); + type Result = (); + type Complete = (); + + fn resume( + &mut self, + result: Option, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(async move { + match (self.0, result) { + (false, None) => { + self.0 = true; + Ok(NativeCallStep::Host(())) + } + (true, Some(())) => Ok(NativeCallStep::Complete(())), + _ => Err(litellm_core::Error::InvalidRequest( + "invalid synthetic lifecycle state".into(), + )), + } + }) + } + + fn interrupt( + &mut self, + _: HostFailure, + ) -> HostCallFuture<'_, Self::Operation, Self::Complete> { + Box::pin(async { Ok(NativeCallStep::Complete(())) }) + } + } + + struct SyntheticRoute(PythonCallState); + + impl PythonRoute for SyntheticRoute { + type Call = SyntheticCall; + + fn state(&self) -> &PythonCallState { + &self.0 + } + + fn state_mut(&mut self) -> &mut PythonCallState { + &mut self.0 + } + + fn classify(_: &()) -> OperationClass { + OperationClass::Route + } + + fn lifecycle_result() {} + + fn map_error(error: litellm_core::Error) -> PyErr { + crate::errors::core_error_to_pyerr(error) + } + + fn invoke(&mut self, py: Python<'_>, _: ()) -> PyResult<()> { + self.0.response = Some( + pyo3::types::PyString::new(py, "shared lifecycle") + .into_any() + .unbind(), + ); + Ok(()) + } + + fn cleanup(&mut self) {} + + fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { + Ok(()) + } + } + + #[test] + fn shared_runner_executes_a_non_ocr_adapter() { + Python::initialize(); + Python::attach(|py| { + let route = SyntheticRoute( + PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + false, + "synthetic", + ) + .unwrap(), + ); + let value: String = run_call(py, SyntheticCall(false), route) + .unwrap() + .extract(py) + .unwrap(); + assert_eq!(value, "shared lifecycle"); + }); + } + + #[test] + fn ready_native_lifecycle_completes_without_scheduling() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap(); + let route = SyntheticRoute( + PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + true, + "synthetic", + ) + .unwrap(), + ); + let coroutine = run_call(py, SyntheticCall(false), route).unwrap(); + let completed = coroutine + .call_method1(py, "send", (py.None(),)) + .unwrap_err(); + assert!(completed.is_instance_of::(py)); + assert_eq!( + completed + .value(py) + .getattr("value") + .unwrap() + .extract::() + .unwrap(), + "shared lifecycle", + ); + }); + } + + #[test] + fn python_driver_preserves_inline_await_and_native_ownership() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + py.import("asyncio").unwrap(); + let source = std::ffi::CString::new(include_str!( + "../../../../../litellm/rust_bridge/lifecycle.py" + )) + .unwrap(); + let module = PyModule::from_code( + py, + &source, + pyo3::ffi::c_str!("lifecycle.py"), + pyo3::ffi::c_str!("litellm.rust_bridge.lifecycle"), + ) + .unwrap(); + let locals = PyDict::new(py); + locals + .set_item("drive", module.getattr("drive").unwrap()) + .unwrap(); + locals + .set_item( + "await_execution", + wrap_pyfunction!(await_execution, py).unwrap(), + ) + .unwrap(); + locals + .set_item( + "calling_execution", + wrap_pyfunction!(calling_execution, py).unwrap(), + ) + .unwrap(); + let probe = std::ffi::CString::new(include_str!("../../tests/lifecycle.py")).unwrap(); + py.run(&probe, Some(&locals), Some(&locals)).unwrap(); + }); + } + + struct ErrorBody(PythonCallState); + + impl ExecutionBody for ErrorBody { + fn resume(&mut self, _: Option>>) -> PyResult { + Python::attach(|py| { + Err(PyErr::from_value( + self.0.error.take().unwrap().into_bound(py).into_any(), + )) + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.0.traverse(visit) + } + } + + #[pyfunction] + fn error_execution(py: Python<'_>, error: Bound<'_, PyBaseException>) -> Execution { + let mut state = PythonCallState::new( + py, + PyTuple::empty(py).unbind(), + PyDict::new(py).unbind(), + true, + "test", + ) + .unwrap(); + state.retain_error(py, PyErr::from_value(error.into_any())); + Execution::new(ErrorBody(state)) + } + + #[test] + fn retained_exception_frames_and_duplicate_argument_edges_are_collectable() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "error_execution", + wrap_pyfunction!(error_execution, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained: + pass + +def cycle(): + retained = Retained() + try: + raise ValueError('retained traceback') + except ValueError as error: + retained.owner = error_execution(error) + return weakref.ref(retained) + +reference = cycle() +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + fn state( + py: Python<'_>, + logger: Py, + response: Py, + asynchronous: bool, + ) -> PythonCallState { + PythonCallState { + args: PyTuple::empty(py).unbind(), + kwargs: PyDict::new(py).unbind(), + logger: Some(logger.extract(py).unwrap()), + start: py.None(), + end: Some(py.None()), + response: Some(response), + error: None, + asynchronous, + internal: false, + call_type: "test", + } + } + + #[test] + fn success_dispatch_reports_ordinary_failures_without_replacing_response() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +import sys + +response = object() +failure = ValueError('terminal diagnostic') +diagnostics = [] +old_hook = sys.unraisablehook +sys.unraisablehook = lambda event: diagnostics.append(event.exc_value) + +class Logger: + def handle_sync_success_callbacks_for_async_calls(self, *args): + raise failure + +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let response = locals.get_item("response").unwrap().unwrap().unbind(); + let mut lifecycle_state = state( + py, + locals.get_item("logger").unwrap().unwrap().unbind(), + response.clone_ref(py), + true, + ); + lifecycle_state.internal = true; + lifecycle_state.dispatch_success(py).unwrap(); + assert!(lifecycle_state.response.as_ref().unwrap().is(&response)); + py.run( + pyo3::ffi::c_str!( + r#" +assert diagnostics == [failure] +sys.unraisablehook = old_hook +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn retained_failure_preserves_exception_identity() { + Python::initialize(); + Python::attach(|py| { + let logger = PyDict::new(py).into_any().unbind(); + let response = py.None(); + let failure = pyo3::exceptions::PyValueError::new_err("identity"); + let failure_value = failure.value(py).clone().unbind(); + let mut lifecycle_state = state(py, logger, response, false); + lifecycle_state.retain_error(py, failure); + let retained = lifecycle_state.error.take().unwrap(); + assert!(retained.is(&failure_value)); + }); + } + + #[test] + fn deferred_release_uses_release_context_and_allows_reentry_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +import sys +import types +from contextvars import ContextVar + +litellm = types.ModuleType('litellm') +core_utils = types.ModuleType('litellm.litellm_core_utils') +logging_worker = types.ModuleType('litellm.litellm_core_utils.logging_worker') +litellm.litellm_core_utils = core_utils +core_utils.logging_worker = logging_worker +sys.modules['litellm'] = litellm +sys.modules['litellm.litellm_core_utils'] = core_utils +sys.modules['litellm.litellm_core_utils.logging_worker'] = logging_worker + +marker = ContextVar('marker', default='unset') +observed = [] + +class Coroutine: + def close(self): + observed.append('closed') + +class Worker: + def ensure_initialized_and_enqueue(self, coroutine): + observed.append(marker.get()) + pending.release(True) + coroutine.close() + +class Logger: + def async_success_handler(self, *args): + observed.append('created') + return Coroutine() + +worker = Worker() +logger = Logger() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + install_logging_worker(py, &locals.get_item("worker").unwrap().unwrap()).unwrap(); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(), + response: Some(py.None()), + start: py.None(), + end: Some(py.None()), + }), + }, + ) + .unwrap(); + locals.set_item("pending", &pending).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['created', 'release', 'closed'] +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn deferred_logging_collects_cycles_through_typed_logger() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + pyo3::ffi::c_str!("class Logger: pass\nlogger = Logger()"), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: locals + .get_item("logger") + .unwrap() + .unwrap() + .extract() + .unwrap(), + response: None, + start: py.None(), + end: None, + }), + }, + ) + .unwrap(); + locals.set_item("pending", pending).unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref +logger.pending = pending +reference = weakref.ref(logger) +del logger, pending +gc.collect() +assert reference() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } + + #[test] + fn coroutine_collects_cycles_retained_by_bridge_host() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "retaining_coroutine", + wrap_pyfunction!(retaining_coroutine, py).unwrap(), + ) + .unwrap(); + py.run( + pyo3::ffi::c_str!( + r#" +import gc +import weakref + +class Retained: + pass + +def cycle(): + retained = Retained() + coroutine = retaining_coroutine(retained) + retained.coroutine = coroutine + return weakref.ref(retained) + +retained_ref = cycle() +gc.collect() +assert retained_ref() is None +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs b/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs new file mode 100644 index 00000000000..ba4a8bb3739 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/lifecycle/preparation.rs @@ -0,0 +1,314 @@ +use litellm_core::auth::{credential_default_fields, credential_index}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList}; + +struct CredentialEntry<'py>(Bound<'py, PyAny>); + +impl<'py> CredentialEntry<'py> { + fn name(&self) -> PyResult { + self.0.getattr("credential_name")?.extract() + } + + fn values(&self) -> PyResult> { + Ok(self.0.getattr("credential_values")?.cast_into::()?) + } +} + +pub(super) fn prepare<'py>( + py: Python<'py>, + kwargs: &Bound<'py, PyDict>, + logger: &super::PythonLogger, +) -> PyResult> { + let arguments = kwargs.copy()?; + arguments.set_item("litellm_logging_obj", logger.object(py))?; + let litellm = py.import("litellm")?; + inherit_credentials(py, &litellm, &arguments)?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("check_limits")? + .call1((&arguments,))?; + Ok(arguments) +} + +fn inherit_credentials( + py: Python<'_>, + litellm: &Bound<'_, PyModule>, + arguments: &Bound<'_, PyDict>, +) -> PyResult<()> { + let Some(requested) = arguments + .get_item("litellm_credential_name")? + .filter(|value| !value.is_none()) + else { + return Ok(()); + }; + if !requested.is_truthy()? { + return Ok(()); + } + let requested: String = requested.extract()?; + let credentials = litellm.getattr("credential_list")?.cast_into::()?; + let names = credentials + .iter() + .map(|credential| CredentialEntry(credential).name()) + .collect::>>()?; + let Some(index) = credential_index(&requested, &names) else { + py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1( + "warning", + ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()), + )?; + return Ok(()); + }; + let selected = CredentialEntry(credentials.get_item(index)?); + let values = selected.values()?; + let supplied: Vec = arguments.keys().extract()?; + let fields: Vec = values.keys().extract()?; + for name in credential_default_fields(&supplied, &fields) { + if let Some(value) = values.get_item(name)? { + arguments.set_item(name, value)?; + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> { + let litellm = PyModule::new(py, "credential_host")?; + litellm.setattr( + "credential_list", + locals.get_item("credentials").unwrap().unwrap(), + )?; + inherit_credentials( + py, + &litellm, + &locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::()?, + ) + } + + #[test] + fn duplicate_names_select_the_first_entry_without_reading_other_values() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +accesses = [] +class Credential: + def __init__(self, name, values): + self._name = name + self._values = values + @property + def credential_name(self): + accesses.append(('name', self._name)) + return self._name + @property + def credential_values(self): + accesses.append(('values', self._name)) + return self._values +credentials = [ + Credential('ocr-test', {'api_key': 'first'}), + Credential('other', {'api_key': 'unused'}), + Credential('ocr-test', {'api_key': 'later'}), +] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + inherit(py, &locals).unwrap(); + let arguments = locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert_eq!( + arguments + .get_item("api_key") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "first" + ); + let accesses: Vec<(String, String)> = locals + .get_item("accesses") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!( + accesses, + [ + ("name".into(), "ocr-test".into()), + ("name".into(), "other".into()), + ("name".into(), "ocr-test".into()), + ("values".into(), "ocr-test".into()), + ] + ); + }); + } + + #[test] + fn later_invalid_name_still_fails_after_an_earlier_match() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +failure = LookupError('later name') +class Good: + credential_name = 'ocr-test' + credential_values = {'api_key': 'first'} +class Bad: + @property + def credential_name(self): + raise failure +credentials = [Good(), Bad()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + let error = inherit(py, &locals).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn selected_values_must_be_a_dictionary_and_property_errors_keep_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Listed: + credential_name = 'ocr-test' + credential_values = ['not-a-dict'] +credentials = [Listed()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + assert!( + inherit(py, &locals) + .unwrap_err() + .is_instance_of::(py) + ); + + let locals = eval( + py, + c" +failure = RuntimeError('values failed') +class Broken: + credential_name = 'ocr-test' + @property + def credential_values(self): + raise failure +credentials = [Broken()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + let error = inherit(py, &locals).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn explicit_none_is_not_overwritten_and_inherited_objects_keep_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +opaque = object() +class Credential: + credential_name = 'ocr-test' + credential_values = {'api_key': 'credential-key', 'opaque': opaque} +credentials = [Credential()] +arguments = {'litellm_credential_name': 'ocr-test', 'api_key': None} +", + ); + inherit(py, &locals).unwrap(); + let arguments = locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert!(arguments.get_item("api_key").unwrap().unwrap().is_none()); + assert!( + arguments + .get_item("opaque") + .unwrap() + .unwrap() + .is(locals.get_item("opaque").unwrap().unwrap()) + ); + }); + } + + #[test] + fn selection_rereads_the_list_after_name_properties_run() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class First: + @property + def credential_name(self): + credentials[0] = Second() + return 'ocr-test' + credential_values = {'api_key': 'first'} +class Second: + credential_name = 'ocr-test' + credential_values = {'api_key': 'replaced'} +credentials = [First()] +arguments = {'litellm_credential_name': 'ocr-test'} +", + ); + inherit(py, &locals).unwrap(); + let arguments = locals + .get_item("arguments") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert_eq!( + arguments + .get_item("api_key") + .unwrap() + .unwrap() + .extract::() + .unwrap(), + "replaced" + ); + }); + } + + #[test] + fn falsy_credential_names_return_before_loading_credentials() { + Python::initialize(); + Python::attach(|py| { + let litellm = PyModule::new(py, "credential_host").unwrap(); + for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] { + let arguments = PyDict::new(py); + arguments.set_item("litellm_credential_name", name).unwrap(); + inherit_credentials(py, &litellm, &arguments).unwrap(); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index a14e4b55d82..5f7633a64a0 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -1,10 +1,14 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::time::Duration; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::types::PyDict; use serde_json::{Map, Value}; +use litellm_core::auth::InputSource; +use litellm_python_interop::from_py_preserving_errors as from_py; + pub(crate) struct RouteOptions { pub(crate) model: String, pub(crate) api_key: Option, @@ -36,18 +40,18 @@ impl RouteOptions { } } -pub(crate) fn required_value( - name: &'static str, - value: Value, - expected: fn(&Value) -> bool, - expected_name: &'static str, -) -> PyResult { - if expected(&value) { - return Ok(value); +pub(crate) fn required_array(name: &'static str, value: Value) -> PyResult> { + match value { + Value::Array(values) => Ok(values), + _ => Err(PyValueError::new_err(format!("{name} must be a list"))), + } +} + +pub(crate) fn required_object(name: &'static str, value: Value) -> PyResult> { + match value { + Value::Object(values) => Ok(values), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), } - Err(PyValueError::new_err(format!( - "{name} must be a {expected_name}" - ))) } pub(crate) fn object_or_empty( @@ -55,7 +59,7 @@ pub(crate) fn object_or_empty( value: Option, ) -> PyResult> { match value { - Some(value) => object(name, value), + Some(value) => required_object(name, value), None => Ok(Map::new()), } } @@ -64,14 +68,7 @@ fn optional_object( name: &'static str, value: Option, ) -> PyResult>> { - value.map(|value| object(name, value)).transpose() -} - -fn object(name: &'static str, value: Value) -> PyResult> { - match value { - Value::Object(map) => Ok(map), - _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), - } + value.map(|value| required_object(name, value)).transpose() } pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option { @@ -84,6 +81,72 @@ pub(crate) fn optional_timeout(timeout_seconds: Option) -> Option }) } +pub(crate) fn python_timeout_seconds(py: Python<'_>, timeout: Py) -> PyResult> { + py.import("litellm.rust_bridge.timeouts")? + .getattr("timeout_to_seconds")? + .call1((timeout,))? + .extract() +} + +pub(crate) fn project_optional_fields( + kwargs: &Bound<'_, PyDict>, + names: &[&str], +) -> PyResult> { + names + .iter() + .filter_map(|name| match kwargs.get_item(name) { + Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect() +} + +struct RequestFieldSources<'py> { + body: Option>, + credentials: Option>, +} + +impl<'py> RequestFieldSources<'py> { + fn extract(proxy_request: &Bound<'py, PyAny>) -> PyResult { + let proxy_request = proxy_request.cast::()?; + + let body = proxy_request + .get_item("body_fields")? + .or(proxy_request.get_item("body")?); + + let credentials = proxy_request.get_item("credential_fields")?; + + Ok(Self { body, credentials }) + } + + fn contains(&self, name: &str) -> bool { + self.body + .as_ref() + .is_some_and(|fields| fields.contains(name).unwrap_or(false)) + || self + .credentials + .as_ref() + .is_some_and(|fields| fields.contains(name).unwrap_or(false)) + } +} + +pub(crate) fn request_input_sources<'a>( + kwargs: &Bound<'_, PyDict>, + names: impl Iterator, +) -> PyResult> { + let Some(proxy_request) = kwargs.get_item("proxy_server_request")? else { + return Ok(BTreeMap::new()); + }; + + let sources = RequestFieldSources::extract(&proxy_request)?; + + Ok(names + .filter(|name| sources.contains(name)) + .map(|name| (name.to_string(), InputSource::Request)) + .collect()) +} + pub(crate) fn marshal_headers(headers: Option) -> PyResult> { let value = match headers { Some(headers) => headers, @@ -102,3 +165,199 @@ pub(crate) fn marshal_headers(headers: Option) -> PyResult(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn sources( + py: Python<'_>, + proxy: &Bound<'_, PyAny>, + names: &[&str], + ) -> PyResult> { + let kwargs = PyDict::new(py); + kwargs.set_item("proxy_server_request", proxy)?; + request_input_sources(&kwargs, names.iter().copied()) + } + + #[test] + fn required_shapes_preserve_nested_values_and_existing_errors() { + let nested = json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]); + assert_eq!( + Value::Array(required_array("messages", nested.clone()).unwrap()), + nested + ); + + let body = json!({"model": "claude", "metadata": {"user": "1"}}); + assert_eq!( + Value::Object(required_object("body", body.clone()).unwrap()), + body + ); + + assert_eq!( + required_array("messages", json!({"role": "user"})) + .unwrap_err() + .to_string(), + "ValueError: messages must be a list" + ); + assert_eq!( + required_object("body", json!([])).unwrap_err().to_string(), + "ValueError: body must be a dict" + ); + } + + #[test] + fn optional_parameters_treat_missing_as_empty() { + assert_eq!( + object_or_empty("optional_params", None).unwrap(), + Map::new() + ); + assert_eq!( + object_or_empty("optional_params", Some(json!({"temperature": 0.2}))).unwrap(), + required_object("optional_params", json!({"temperature": 0.2})).unwrap() + ); + } + + #[test] + fn missing_none_and_empty_proxy_metadata_are_distinct() { + Python::initialize(); + Python::attach(|py| { + let kwargs = PyDict::new(py); + assert!( + request_input_sources(&kwargs, ["api_key"].into_iter()) + .unwrap() + .is_empty() + ); + + kwargs.set_item("proxy_server_request", py.None()).unwrap(); + assert!( + request_input_sources(&kwargs, ["api_key"].into_iter()) + .unwrap_err() + .is_instance_of::(py) + ); + + kwargs + .set_item("proxy_server_request", PyDict::new(py)) + .unwrap(); + assert!( + request_input_sources(&kwargs, ["api_key"].into_iter()) + .unwrap() + .is_empty() + ); + }); + } + + #[test] + fn body_fields_win_over_body_and_explicit_none_does_not_fall_back() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +proxy = {'body_fields': ['api_key'], 'body': ['api_base']} +none_fields = {'body_fields': None, 'body': ['api_key']} +body_only = {'body': ['api_base']} +", + ); + let named = sources( + py, + &locals.get_item("proxy").unwrap().unwrap(), + &["api_key", "api_base"], + ) + .unwrap(); + assert_eq!(named.get("api_key").copied(), Some(InputSource::Request)); + assert!(!named.contains_key("api_base")); + + assert!( + sources( + py, + &locals.get_item("none_fields").unwrap().unwrap(), + &["api_key"], + ) + .unwrap() + .is_empty() + ); + + let body_only = sources( + py, + &locals.get_item("body_only").unwrap().unwrap(), + &["api_base"], + ) + .unwrap(); + assert_eq!( + body_only.get("api_base").copied(), + Some(InputSource::Request) + ); + }); + } + + #[test] + fn body_and_credential_membership_can_mark_request_fields() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Raising: + def __contains__(self, item): + raise RuntimeError('credential membership') +proxy = { + 'body_fields': ['api_key'], + 'credential_fields': Raising(), +} +credentials_only = {'credential_fields': ['extra_headers']} +erroring = {'body_fields': Raising()} +extra = {'body_fields': ['api_key', 'unused']} +", + ); + let skipped = sources( + py, + &locals.get_item("proxy").unwrap().unwrap(), + &["api_key"], + ) + .unwrap(); + assert_eq!(skipped.get("api_key").copied(), Some(InputSource::Request)); + + let credentials = sources( + py, + &locals.get_item("credentials_only").unwrap().unwrap(), + &["extra_headers"], + ) + .unwrap(); + assert_eq!( + credentials.get("extra_headers").copied(), + Some(InputSource::Request) + ); + + assert!( + sources( + py, + &locals.get_item("erroring").unwrap().unwrap(), + &["api_key"], + ) + .unwrap() + .is_empty() + ); + + let requested = sources( + py, + &locals.get_item("extra").unwrap().unwrap(), + &["api_key"], + ) + .unwrap(); + assert_eq!(requested.len(), 1); + assert_eq!( + requested.get("api_key").copied(), + Some(InputSource::Request) + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs new file mode 100644 index 00000000000..f2997ee278c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/mod.rs @@ -0,0 +1,12 @@ +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs similarity index 100% rename from litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs rename to litellm-rust/crates/python-bridge/src/routes/audio_transcription/value.rs diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs new file mode 100644 index 00000000000..f2997ee278c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -0,0 +1,12 @@ +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs similarity index 95% rename from litellm-rust/crates/python-bridge/src/routes/chat_completions.rs rename to litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs index 08ab476005c..e67bfa89cc7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/value.rs @@ -9,12 +9,12 @@ use pyo3::prelude::*; use serde_json::Value; use crate::errors::chat_completions_error_to_pyerr; -use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_value}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty, required_array}; fn prepare_chat_completions( inputs: ChatCompletionsInputs, ) -> PyResult> + Send + 'static> { - let messages = required_value("messages", inputs.messages, Value::is_array, "list")?; + let messages = required_array("messages", inputs.messages)?; let optional_params = object_or_empty("optional_params", inputs.optional_params)?; let options = RouteOptions::from_python(RouteOptionsInputs { model: inputs.model, @@ -36,7 +36,7 @@ fn prepare_chat_completions( } = options; run_chat_completions(ChatCompletionsRequest { model: &model, - messages, + messages: Value::Array(messages), optional_params, api_key: api_key.as_deref(), api_base: api_base.as_deref(), diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 97313651011..571042062f5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -389,6 +389,82 @@ mod tests { }); } + #[test] + fn missing_and_explicit_none_optional_params_share_the_next_error() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let messages = PyList::empty(py); + let headers = PyList::empty(py); + let omitted = PyDict::new(py); + omitted + .set_item("extra_headers", &headers) + .expect("kwargs should accept extra_headers"); + let explicit = PyDict::new(py); + explicit + .set_item("optional_params", py.None()) + .expect("kwargs should accept optional_params"); + explicit + .set_item("extra_headers", &headers) + .expect("kwargs should accept extra_headers"); + + let omitted_error = module + .getattr("chat_completions") + .and_then(|function| function.call(("model", &messages), Some(&omitted))) + .expect_err("omitted optional_params should reach header validation"); + let explicit_error = module + .getattr("chat_completions") + .and_then(|function| function.call(("model", &messages), Some(&explicit))) + .expect_err("None optional_params should reach header validation"); + assert_eq!( + omitted_error.to_string(), + "ValueError: extra_headers must be a dict" + ); + assert_eq!(explicit_error.to_string(), omitted_error.to_string()); + }); + } + + #[test] + fn chat_completions_decline_keeps_existing_reasons() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::new(py, "routes").expect("module should be created"); + crate::routes::register(&module).expect("routes should register"); + let decline = module + .getattr("chat_completions_decline") + .expect("decline helper should be registered"); + let empty = PyList::empty(py); + let unreadable = py + .eval(c"'nope'", None, None) + .expect("string messages should convert"); + + let unknown: Option = decline + .call1(("unknown-model", &empty)) + .and_then(|value| value.extract()) + .expect("unknown providers should decline"); + assert_eq!( + unknown.as_deref(), + Some("provider is not on the rust chat completions path") + ); + + let empty_reason: Option = decline + .call1(("anthropic/claude-sonnet-4-5", &empty)) + .and_then(|value| value.extract()) + .expect("empty lists should decline"); + assert_eq!(empty_reason.as_deref(), Some("empty message list")); + + let unreadable_reason: Option = decline + .call1(("anthropic/claude-sonnet-4-5", unreadable)) + .and_then(|value| value.extract()) + .expect("non-list messages should decline"); + assert_eq!( + unreadable_reason.as_deref(), + Some("unreadable message list") + ); + }); + } + #[test] fn generated_routes_execute_sync_and_async_contracts() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs deleted file mode 100644 index 97ff93f299a..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs +++ /dev/null @@ -1,29 +0,0 @@ -use pyo3::prelude::*; -use serde_json::Value; - -use crate::errors::core_error_to_pyerr; - -#[pyfunction] -fn gateway_messages<'py>( - py: Python<'py>, - model_alias: String, - provider_model: String, - api_base: String, - #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value, -) -> PyResult> { - let future = litellm_ai_gateway::trace_parity::messages_request( - model_alias, - provider_model, - api_base, - body, - ); - crate::execution::run_async( - py, - crate::function_trace::capture(future), - core_error_to_pyerr, - ) -} - -pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { - super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs new file mode 100644 index 00000000000..f2997ee278c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -0,0 +1,12 @@ +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages/value.rs similarity index 94% rename from litellm-rust/crates/python-bridge/src/routes/messages.rs rename to litellm-rust/crates/python-bridge/src/routes/messages/value.rs index f69b5e9251d..b741e54f0ca 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/value.rs @@ -6,12 +6,12 @@ use serde_json::Value; use std::future::Future; use crate::errors::core_error_to_pyerr; -use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value}; +use crate::marshal::{RouteOptions, RouteOptionsInputs, required_object}; fn prepare_messages( inputs: MessagesInputs, ) -> PyResult> + Send + 'static> { - let body = required_value("body", inputs.body, Value::is_object, "dict")?; + let body = required_object("body", inputs.body)?; let options = RouteOptions::from_python(RouteOptionsInputs { model: inputs.model, api_key: inputs.api_key, @@ -32,7 +32,7 @@ fn prepare_messages( } = options; run_messages(MessagesRequest { model: &model, - body, + body: Value::Object(body), api_key: api_key.as_deref(), api_base: api_base.as_deref(), custom_llm_provider: custom_llm_provider.as_deref(), diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 7e81f2ffe9b..97c39a5d6b3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -3,9 +3,6 @@ use pyo3::prelude::*; #[macro_use] mod definition; -#[cfg(feature = "trace-parity")] -mod gateway_messages; - mod audio_transcription; mod chat_completions; mod messages; @@ -16,6 +13,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { audio_transcription::register(module)?; messages::register(module)?; chat_completions::register(module)?; + #[cfg(feature = "trace-parity")] { let trace = PyModule::new(module.py(), "_trace")?; @@ -23,7 +21,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { audio_transcription::register_trace(&trace)?; messages::register_trace(&trace)?; chat_completions::register_trace(&trace)?; - gateway_messages::register_trace(&trace)?; module.add_submodule(&trace)?; } Ok(()) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs new file mode 100644 index 00000000000..1cbe8a179e3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -0,0 +1,161 @@ +use pyo3::exceptions::PyBaseException; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::Value; + +use litellm_core::ocr::LiteLLMOcrResponse; +use litellm_core::ocr::hooks::OcrPreCallRequest; +use litellm_python_interop::to_py_preserving_errors as to_py; + +use crate::lifecycle::PythonLogger; + +pub(super) struct OcrLoggingFields { + model: String, + custom_llm_provider: String, + optional_params: Value, +} + +impl From<&OcrPreCallRequest> for OcrLoggingFields { + fn from(request: &OcrPreCallRequest) -> Self { + Self { + model: request.model.clone(), + custom_llm_provider: request.custom_llm_provider.clone(), + optional_params: request.optional_params.clone(), + } + } +} + +impl PythonLogger { + pub(super) fn update_ocr( + &self, + py: Python<'_>, + kwargs: &Py, + pre_call: &OcrLoggingFields, + secret_fields: &[&str], + url: &str, + ) -> PyResult<()> { + let update = PyDict::new(py); + update.set_item("kwargs", redact(py, kwargs.bind(py), secret_fields)?)?; + update.set_item("model", &pre_call.model)?; + update.set_item( + "optional_params", + redact( + py, + &to_py(py, &pre_call.optional_params)? + .into_bound(py) + .cast_into::()?, + secret_fields, + )?, + )?; + let params = PyDict::new(py); + params.set_item( + "litellm_call_id", + kwargs.bind(py).get_item("litellm_call_id")?, + )?; + params.set_item("api_base", url)?; + for name in ["logger_fn", "litellm_request_debug"] { + if let Some(value) = kwargs.bind(py).get_item(name)? { + params.set_item(name, value)?; + } + } + update.set_item("litellm_params", params)?; + update.set_item("custom_llm_provider", &pre_call.custom_llm_provider)?; + self.object(py) + .call_method("update_from_kwargs", (), Some(&update))?; + Ok(()) + } + + pub(crate) fn pre_ocr( + &self, + py: Python<'_>, + api_key: &Option>, + body: &Bound<'_, PyDict>, + headers: &Bound<'_, PyDict>, + url: &str, + ) -> PyResult<()> { + let additional = PyDict::new(py); + additional.set_item("complete_input_dict", body)?; + additional.set_item("headers", headers)?; + additional.set_item("api_base", url)?; + let kwargs = PyDict::new(py); + kwargs.set_item("input", "OCR document processing")?; + kwargs.set_item("api_key", api_key)?; + kwargs.set_item("additional_args", &additional)?; + if self.callbacks_needed(py, "input")? { + self.object(py).call_method("pre_call", (), Some(&kwargs))?; + } else { + self.object(py) + .call_method("_pre_call", (), Some(&kwargs))?; + self.object(py).call_method0("record_api_call_start_time")?; + } + Ok(()) + } + + pub(crate) fn post_ocr( + &self, + py: Python<'_>, + original_response: &Value, + body: Option<&Py>, + headers: Option<&Py>, + ) -> PyResult<()> { + let additional = PyDict::new(py); + additional.set_item("complete_input_dict", body)?; + additional.set_item("headers", headers)?; + if self.callbacks_needed(py, "input")? { + let kwargs = PyDict::new(py); + kwargs.set_item("original_response", to_py(py, original_response)?)?; + kwargs.set_item("additional_args", &additional)?; + self.object(py) + .call_method("post_call", (), Some(&kwargs))?; + } else { + let response = py + .import("json")? + .call_method1("dumps", (to_py(py, original_response)?,))?; + self.object(py).call_method1( + "record_post_call", + (response, py.None(), py.None(), additional), + )?; + } + Ok(()) + } +} + +fn redact( + py: Python<'_>, + params: &Bound<'_, PyDict>, + secret_fields: &[&str], +) -> PyResult> { + let redacted = PyDict::new(py); + for (name, value) in params { + let name = name.extract::()?; + if name == "proxy_server_request" { + continue; + } + if secret_fields.contains(&name.as_str()) { + redacted.set_item(name, "****")?; + } else { + redacted.set_item(name, value)?; + } + } + Ok(redacted.unbind()) +} + +pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult> { + py.import("litellm.rust_bridge.ocr")? + .getattr("_response")? + .call1((to_py(py, response)?,)) + .map(Bound::unbind) +} + +pub(super) fn map_failure( + py: Python<'_>, + error: &Py, + request: &Bound<'_, PyAny>, + provider: &str, +) -> PyResult> { + Ok(py + .import("litellm.rust_bridge.ocr_lifecycle")? + .getattr("map_failure")? + .call1((error, request, provider))? + .extract()?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs new file mode 100644 index 00000000000..d43c2f88775 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -0,0 +1,264 @@ +use std::io::Read; +use std::path::PathBuf; + +use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedBytes; +#[cfg(test)] +use pyo3::types::PyDict; +use pyo3::types::{PyBytes, PyString}; + +use litellm_core::constants::OCR_INLINE_MAX_BYTES; +use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type}; +use litellm_python_interop::to_py_preserving_errors; + +enum FileBytes { + Python(PyBackedBytes), + Native(Vec), +} + +impl AsRef<[u8]> for FileBytes { + fn as_ref(&self) -> &[u8] { + match self { + Self::Python(bytes) => bytes, + Self::Native(bytes) => bytes, + } + } +} + +fn read_file_input( + py: Python<'_>, + file: &Bound<'_, PyAny>, +) -> PyResult<(FileBytes, Option)> { + if file.is_instance_of::() { + return Err(PyValueError::new_err( + "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", + )); + } + if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { + let path: PathBuf = file.extract()?; + let name = path + .file_name() + .map(|value| value.to_string_lossy().into_owned()); + let bytes = py + .detach(|| { + let mut bytes = Vec::new(); + std::fs::File::open(&path)? + .take(OCR_INLINE_MAX_BYTES as u64 + 1) + .read_to_end(&mut bytes)?; + Ok::<_, std::io::Error>(bytes) + }) + .map_err(|error| { + if error.kind() == std::io::ErrorKind::NotFound { + PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) + } else { + error.into() + } + })?; + return Ok((FileBytes::Native(bytes), name)); + } + if file.is_instance_of::() { + return Ok((FileBytes::Python(file.extract()?), None)); + } + let reader = file + .getattr_opt("read")? + .filter(|value| value.is_callable()); + let Some(reader) = reader else { + return Err(PyValueError::new_err(format!( + "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", + file.get_type(), + ))); + }; + let name = file + .getattr_opt("name")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()?; + let value = reader.call0()?; + let bytes = if value.is_instance_of::() { + FileBytes::Native(value.extract::()?.into_bytes()) + } else if value.is_instance_of::() { + FileBytes::Python(value.extract()?) + } else { + return Err(PyTypeError::new_err(format!( + "OCR file read must return bytes or str, got {}", + value.get_type(), + ))); + }; + Ok((bytes, name)) +} + +pub(super) struct FileDocumentInput { + bytes: FileBytes, + name: Option, + mime_type: Option, +} + +impl FromPyObject<'_, '_> for FileDocumentInput { + type Error = PyErr; + + fn extract(document: Borrowed<'_, '_, PyAny>) -> PyResult { + let py = document.py(); + let mime_type = match document.get_item("mime_type") { + Ok(value) => Some(value.extract::()?), + Err(error) if error.is_instance_of::(py) => None, + Err(error) => return Err(error), + }; + let file = document.get_item("file").map_err(|error| { + if error.is_instance_of::(py) { + PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes") + } else { + error + } + })?; + if file.is_none() { + return Err(PyValueError::new_err( + "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", + )); + } + let (bytes, name) = read_file_input(py, &file)?; + Ok(Self { + bytes, + name, + mime_type, + }) + } +} + +pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult { + py.detach(|| { + encode_file_document( + document.bytes.as_ref(), + document.name.as_deref(), + document.mime_type.as_deref(), + ) + }) + .map_err(|error| PyValueError::new_err(error.to_string())) +} + +#[pyfunction] +fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult> { + to_py_preserving_errors(py, &file_document(py, document.extract()?)?) +} + +#[pyfunction] +fn _ocr_mime_type(file_name: &str) -> String { + mime_type_for_name(file_name).into() +} + +#[pyfunction] +#[pyo3(signature = (file_content, file_name=None, content_type=None))] +fn _ocr_upload_document( + py: Python<'_>, + file_content: &Bound<'_, PyBytes>, + file_name: Option<&str>, + content_type: Option<&str>, +) -> PyResult> { + let bytes: PyBackedBytes = file_content.extract()?; + let document = py + .detach(|| { + encode_file_document( + &bytes, + None, + Some(upload_mime_type(file_name, content_type)), + ) + }) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + to_py_preserving_errors(py, &document) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?; + module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?; + module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?; + module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extraction_validates_required_file_and_optional_mime_type() { + Python::initialize(); + Python::attach(|py| { + for expression in [c"{}", c"{'file': None}"] { + let document = py.eval(expression, None, None).unwrap(); + let error = document.extract::().err().unwrap(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("must include a 'file' field")); + } + for expression in [ + c"{'file': b'abc', 'mime_type': None}", + c"{'file': b'abc', 'mime_type': 7}", + ] { + let document = py.eval(expression, None, None).unwrap(); + let error = document.extract::().err().unwrap(); + assert!(error.is_instance_of::(py)); + } + let document = py.eval(c"{'file': b'abc'}", None, None).unwrap(); + let input: FileDocumentInput = document.extract().unwrap(); + assert_eq!(input.bytes.as_ref(), b"abc"); + assert_eq!(input.name, None); + assert_eq!(input.mime_type, None); + }); + } + + #[test] + fn extraction_validates_mime_type_before_consuming_file() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"class Reader: + def __init__(self): + self.reads = 0 + def read(self): + self.reads += 1 + return b'abc' +reader = Reader() +document = {'file': reader, 'mime_type': 7}", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let document = locals.get_item("document").unwrap().unwrap(); + let error = document.extract::().err().unwrap(); + assert!(error.is_instance_of::(py)); + let reads: usize = locals + .get_item("reader") + .unwrap() + .unwrap() + .getattr("reads") + .unwrap() + .extract() + .unwrap(); + assert_eq!(reads, 0); + }); + } + + #[test] + fn extraction_preserves_reader_key_error_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"failure = KeyError('reader failed') +class Reader: + def read(self): + raise failure +document = {'file': Reader()}", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let document = locals.get_item("document").unwrap().unwrap(); + let error = document.extract::().err().unwrap(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs new file mode 100644 index 00000000000..66bdfb7583e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -0,0 +1,72 @@ +use litellm_core::error::Error; +use pyo3::prelude::*; + +use crate::errors::{RustUpstreamError, core_error_to_pyerr}; + +pub(super) fn to_pyerr(error: Error) -> PyErr { + let status = error.http_status_code(); + let mapped = match error { + Error::Http { status, body } => RustUpstreamError::new_err((status, body)), + other => core_error_to_pyerr(other), + }; + attach_status(mapped, status) +} + +fn attach_status(error: PyErr, status: Option) -> PyErr { + if let Some(status) = status { + Python::attach(|py| { + let value = error.value(py); + value.setattr("status_code", status).ok(); + value.setattr("message", value.to_string()).ok(); + }); + } + error +} + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyValueError; + + #[test] + fn preserves_python_validation_and_provider_details() { + Python::initialize(); + Python::attach(|py| { + let mapped = to_pyerr(Error::MissingDocumentUrl); + assert!(mapped.is_instance_of::(py)); + assert_eq!(mapped.value(py).to_string(), "Document URL is required"); + assert_eq!( + mapped + .value(py) + .getattr("status_code") + .unwrap() + .extract::() + .unwrap(), + 500 + ); + let mapped = to_pyerr(Error::Http { + status: 429, + body: r#"{"message":"rate limited"}"#.to_string(), + }); + assert!(mapped.is_instance_of::(py)); + let args: (u16, String) = mapped + .value(py) + .getattr("args") + .and_then(|args| args.extract()) + .expect("OCR failures retain status and unprefixed provider message"); + assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string())); + + let mapped = to_pyerr(Error::InvalidRequest("invalid format".into())); + assert!(mapped.is_instance_of::(py)); + assert_eq!( + mapped + .value(py) + .getattr("status_code") + .unwrap() + .extract::() + .unwrap(), + 400 + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs new file mode 100644 index 00000000000..12d902a3544 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -0,0 +1,311 @@ +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use litellm_core::auth::ResolvedCredential; +use litellm_core::ocr::hooks::{OcrDuringCallRequest, OcrPostCallRequest, OcrPreCallRequest}; +use litellm_core::ocr::{OcrAdmission, OcrCall, OcrClient, OcrHostOperation, OcrHostResult}; +use litellm_python_interop::{ + from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +}; + +use super::callbacks; +use super::errors::to_pyerr as ocr_error_to_pyerr; +use super::project::{ProjectedOcrFields, admitted_call, project_request}; +use crate::lifecycle::{ + OperationClass, PythonCallState, PythonRoute, missing_state, now, run_call, +}; + +struct PythonOcrHost { + state: PythonCallState, + data: OcrHostData, +} + +enum OcrHostData { + Unprojected { request: Py }, + Projected(Box), + Released, +} + +struct ProjectedOcrHost { + fields: ProjectedOcrFields, + pre_call: Option, + retained_fields: Option>, + body: Option>, + headers: Option>, +} + +impl PythonOcrHost { + fn projected(&self) -> PyResult<&ProjectedOcrHost> { + match &self.data { + OcrHostData::Projected(projected) => Ok(projected), + _ => Err(missing_state()), + } + } + + fn projected_mut(&mut self) -> PyResult<&mut ProjectedOcrHost> { + match &mut self.data { + OcrHostData::Projected(projected) => Ok(projected), + _ => Err(missing_state()), + } + } + + fn pre_call( + &mut self, + py: Python<'_>, + request: OcrPreCallRequest, + ) -> PyResult { + let kwargs = self.state.kwargs.bind(py); + let retained_fields = PyDict::new(py); + for name in request + .optional_params + .as_object() + .ok_or_else(missing_state)? + .keys() + { + if let Some(value) = kwargs.get_item(name)? { + retained_fields.set_item(name, value)?; + } + } + retained_fields.set_item("document", &self.projected()?.fields.document)?; + let projected = self.projected_mut()?; + projected.retained_fields = Some(retained_fields.unbind()); + projected.pre_call = Some((&request).into()); + Ok(request) + } + + fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult { + let provider = self + .projected()? + .fields + .azure_ad_token_provider + .as_ref() + .ok_or_else(missing_state)?; + provider.acquire(py) + } + + fn python_pre_call( + &mut self, + py: Python<'_>, + mut request: OcrDuringCallRequest, + ) -> PyResult { + let projected = self.projected()?; + let pre_call = projected.pre_call.as_ref().ok_or_else(missing_state)?; + self.state.logger()?.update_ocr( + py, + &self.state.kwargs, + pre_call, + &projected.fields.secret_fields, + &request.url, + )?; + if !self.state.logger()?.callbacks_needed(py, "payload")? { + self.state + .logger()? + .object(py) + .call_method0("record_api_call_start_time")?; + return Ok(request); + } + if let Some(body) = request.body.as_object_mut() { + for name in &request.retained_fields { + body.remove(name); + } + } + let body = to_py(py, &request.body)? + .into_bound(py) + .cast_into::()?; + if let Some(retained) = &self.projected()?.retained_fields { + for name in &request.retained_fields { + if let Some(value) = retained.bind(py).get_item(name)? { + body.set_item(name, value)?; + } + } + } + let headers = PyDict::new(py); + for (name, value) in &request.headers { + headers.set_item(name, value)?; + } + let api_key = self.projected()?.fields.api_key.clone_ref(py); + let projected = self.projected_mut()?; + projected.body = Some(body.clone().unbind()); + projected.headers = Some(headers.clone().unbind()); + self.state + .logger()? + .pre_ocr(py, &Some(api_key), &body, &headers, &request.url)?; + let headers = headers + .iter() + .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) + .collect::>>()?; + request.body = from_py(&body)?; + request.headers = headers; + Ok(request) + } + + fn python_post_call( + &mut self, + py: Python<'_>, + request: OcrPostCallRequest, + ) -> PyResult { + let logger = self.state.logger()?; + if logger.callbacks_needed(py, "payload")? { + let projected = self.projected()?; + logger.post_ocr( + py, + &request.original_response, + projected.body.as_ref(), + projected.headers.as_ref(), + )?; + } + Ok(request) + } +} + +impl PythonRoute for PythonOcrHost { + type Call = OcrCall; + + fn state(&self) -> &PythonCallState { + &self.state + } + + fn state_mut(&mut self) -> &mut PythonCallState { + &mut self.state + } + + fn classify(operation: &OcrHostOperation) -> OperationClass { + operation + .phase() + .map_or(OperationClass::Route, OperationClass::Phase) + } + + fn lifecycle_result() -> OcrHostResult { + OcrHostResult::Lifecycle(Ok(())) + } + + fn map_error(error: litellm_core::Error) -> PyErr { + ocr_error_to_pyerr(error) + } + + fn invoke(&mut self, py: Python<'_>, operation: OcrHostOperation) -> PyResult { + Ok(match operation { + OcrHostOperation::ProjectRequest => { + let OcrHostData::Unprojected { request } = &self.data else { + return Err(missing_state()); + }; + let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?; + let has_token_provider = projected.fields.azure_ad_token_provider.is_some(); + let request = projected.request; + self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost { + fields: projected.fields, + pre_call: None, + retained_fields: None, + body: None, + headers: None, + })); + OcrHostResult::Request(Ok((Box::new(request), has_token_provider))) + } + OcrHostOperation::AcquireAzureAdToken => { + OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?)) + } + OcrHostOperation::PreCall(request) => { + OcrHostResult::PreCall(Ok(self.pre_call(py, request)?)) + } + OcrHostOperation::DuringCall(request) => { + OcrHostResult::DuringCall(Ok(self.python_pre_call(py, request)?)) + } + OcrHostOperation::PostCall(request) => { + OcrHostResult::PostCall(Ok(self.python_post_call(py, request)?)) + } + OcrHostOperation::ConstructResponse(response) => { + self.state.end = Some(now(py)?); + self.state.response = Some(callbacks::response(py, response.as_ref())?); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::MapFailure(error) => { + if self.state.error.is_none() { + self.state.retain_error(py, ocr_error_to_pyerr(error)); + } + if self.state.end.is_none() { + self.state.end = Some(now(py)?); + } + let error = self.state.error.as_ref().ok_or_else(missing_state)?; + let (request, provider) = match &self.data { + OcrHostData::Unprojected { request } => (request.bind(py), ""), + OcrHostData::Projected(projected) => ( + projected.fields.boundary_request.bind(py), + projected.fields.provider, + ), + OcrHostData::Released => return Err(missing_state()), + }; + let mapped = callbacks::map_failure(py, error, request, provider)?; + self.state + .retain_error(py, PyErr::from_value(mapped.into_bound(py).into_any())); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => return Err(missing_state()), + }) + } + + fn cleanup(&mut self) { + self.data = OcrHostData::Released; + } + fn traverse(&self, visit: &pyo3::gc::PyVisit<'_>) -> Result<(), pyo3::gc::PyTraverseError> { + match &self.data { + OcrHostData::Unprojected { request } => visit.call(request), + OcrHostData::Projected(projected) => { + visit.call(&projected.fields.boundary_request)?; + visit.call(&projected.fields.document)?; + visit.call(&projected.fields.api_key)?; + if let Some(provider) = &projected.fields.azure_ad_token_provider { + provider.traverse(visit)?; + } + visit.call(&projected.retained_fields)?; + visit.call(&projected.body)?; + visit.call(&projected.headers) + } + OcrHostData::Released => Ok(()), + } + } +} + +pub(super) struct BridgeOcrHooks; + +impl litellm_core::ocr::hooks::OcrHooks for BridgeOcrHooks { + fn intercepts_requests(&self) -> bool { + true + } +} + +#[pyfunction] +fn _ocr_lifecycle( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult> { + let client = OcrClient::shared().map_err(ocr_error_to_pyerr)?; + let call = admitted_call(OcrCall::admit( + client, + OcrAdmission { + asynchronous, + ..OcrAdmission::all() + }, + ))?; + let host = PythonOcrHost { + state: PythonCallState::new( + py, + args.unbind(), + kwargs.copy()?.unbind(), + asynchronous, + if asynchronous { "aocr" } else { "ocr" }, + )?, + data: OcrHostData::Unprojected { + request: request.unbind(), + }, + }; + run_call(py, call, host) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(_ocr_lifecycle, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs new file mode 100644 index 00000000000..10fa40b65ea --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -0,0 +1,19 @@ +mod callbacks; +mod document; +mod errors; +mod lifecycle; +mod project; +mod value; + +use pyo3::prelude::*; + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register(module)?; + document::register(module)?; + lifecycle::register(module) +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + value::register_trace(module) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs new file mode 100644 index 00000000000..8b6a1b02e19 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -0,0 +1,579 @@ +use std::sync::Arc; + +use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request}; +use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall}; +use litellm_python_interop::{ + from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::{Map, Value}; + +use super::errors::to_pyerr as ocr_error_to_pyerr; +use super::lifecycle::BridgeOcrHooks; +use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider}; +use crate::errors::RustBridgeDeclined; +use crate::marshal::{project_optional_fields, python_timeout_seconds, request_input_sources}; + +pub(super) struct ProjectedOcrFields { + pub boundary_request: Py, + pub document: Py, + pub api_key: Py, + pub azure_ad_token_provider: Option, + pub provider: &'static str, + pub secret_fields: Vec<&'static str>, +} + +pub(super) struct ProjectedOcrCall { + pub request: LiteLLMOcrRequest, + pub fields: ProjectedOcrFields, +} + +struct OcrArguments<'a, 'py> { + request: &'a Bound<'py, PyAny>, + kwargs: &'a Bound<'py, PyDict>, +} + +impl<'py> OcrArguments<'_, 'py> { + fn lookup(&self, name: &str) -> PyResult> { + match self.kwargs.get_item(name)? { + Some(value) => Ok(value), + None => self.request.getattr(name), + } + } + + fn model(&self) -> PyResult { + self.lookup("model")?.extract() + } + + fn custom_llm_provider(&self) -> PyResult> { + self.lookup("custom_llm_provider")?.extract() + } + + fn document(&self) -> PyResult> { + self.lookup("document") + } + + fn api_key(&self) -> PyResult> { + self.lookup("api_key") + } + + fn api_base(&self) -> PyResult> { + self.lookup("api_base")?.extract() + } + + fn extra_headers(&self) -> PyResult>> { + self.lookup("extra_headers")? + .extract::>>()? + .map(|value| from_py(value.bind(self.request.py()))) + .transpose() + } + + fn timeout_seconds(&self) -> PyResult> { + Ok(self + .lookup("timeout")? + .extract::>>()? + .map(|value| python_timeout_seconds(self.request.py(), value)) + .transpose()? + .flatten()) + } +} + +enum ProjectedDocument { + File { wire: Value, retained: Py }, + Other { wire: Value, retained: Py }, +} + +impl ProjectedDocument { + fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult { + let kind: String = document.get_item("type")?.extract()?; + if kind != "file" { + return Ok(Self::Other { + wire: from_py(document)?, + retained: document.clone().unbind(), + }); + } + let input = document.extract()?; + let encoded = super::document::file_document(py, input)?; + let wire = serde_json::to_value(encoded) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + Ok(Self::File { + retained: to_py(py, &wire)?, + wire, + }) + } + + fn into_parts(self) -> (Value, Py) { + match self { + Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained), + } + } +} + +pub(super) fn project_request( + py: Python<'_>, + request: &Bound<'_, PyAny>, + kwargs: &Bound<'_, PyDict>, +) -> PyResult { + let boundary_request = request.clone().unbind(); + let arguments = OcrArguments { request, kwargs }; + let model = arguments.model()?; + let custom_llm_provider = arguments.custom_llm_provider()?; + let (wire_document, retained_document) = + ProjectedDocument::project(py, &arguments.document()?)?.into_parts(); + let api_key = arguments.api_key()?; + let specs = consumed_optional_params(&model, custom_llm_provider.as_deref()) + .map_err(ocr_error_to_pyerr)?; + let names = specs.iter().map(|spec| spec.name).collect::>(); + let optional_params = project_optional_fields(kwargs, &names)?; + let input_sources = request_input_sources( + kwargs, + names + .iter() + .copied() + .chain(["api_key", "api_base", "extra_headers"]), + )?; + let azure_ad_token_provider = kwargs + .get_item("azure_ad_token_provider")? + .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER)); + let wire = OcrWireRequest { + model, + document: wire_document, + api_key: api_key.extract()?, + api_base: arguments.api_base()?, + custom_llm_provider, + extra_headers: arguments.extra_headers()?, + optional_params, + input_sources, + timeout_seconds: arguments.timeout_seconds()?, + }; + let request = decode_request(wire).map_err(ocr_error_to_pyerr)?; + let provider = request.provider_name(); + Ok(ProjectedOcrCall { + request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None), + fields: ProjectedOcrFields { + boundary_request, + document: retained_document, + api_key: api_key.unbind(), + azure_ad_token_provider, + provider, + secret_fields: specs + .into_iter() + .filter(|spec| spec.secret) + .map(|spec| spec.name) + .collect(), + }, + }) +} + +pub(super) fn admitted_call(outcome: NativeOutcome) -> PyResult { + match outcome { + NativeOutcome::Completed(call) => Ok(call), + NativeOutcome::Declined(reason) => Err(RustBridgeDeclined::new_err(format!( + "native OCR admission declined: {reason:?}" + ))), + } +} + +#[cfg(test)] +mod tests { + use litellm_core::Error; + use litellm_core::ocr::OcrDecline; + use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError}; + + use super::*; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } + + fn arguments<'a, 'py>( + request: &'a Bound<'py, PyAny>, + kwargs: &'a Bound<'py, PyDict>, + ) -> OcrArguments<'a, 'py> { + OcrArguments { request, kwargs } + } + + fn project_document( + py: Python<'_>, + document: &Bound<'_, PyAny>, + ) -> PyResult<(Value, Py)> { + ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts) + } + + fn stub_timeout_conversion(py: Python<'_>) { + eval( + py, + c" +import sys +import types +timeouts = types.ModuleType('litellm.rust_bridge.timeouts') +timeouts.timeout_to_seconds = lambda timeout: None if timeout is None else float(timeout) +sys.modules.setdefault('litellm', types.ModuleType('litellm')) +sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) +sys.modules['litellm.rust_bridge.timeouts'] = timeouts +", + ); + } + + #[test] + fn typed_initial_decline_uses_bridge_decline_contract() { + Python::initialize(); + Python::attach(|py| { + let Err(error) = admitted_call(NativeOutcome::Declined(OcrDecline::HostOperations)) + else { + panic!("unsupported host operations should decline admission"); + }; + assert!(error.is_instance_of::(py)); + }); + } + + #[test] + fn post_admission_error_does_not_use_bridge_decline_contract() { + Python::initialize(); + Python::attach(|py| { + let error = ocr_error_to_pyerr(Error::InvalidRequest("callback result".into())); + assert!(error.is_instance_of::(py)); + assert!(!error.is_instance_of::(py)); + }); + } + + #[test] + fn kwargs_override_request_attributes_including_explicit_none() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Request: + def __init__(self): + self.accesses = [] + def __getattribute__(self, name): + if name != 'accesses': + object.__getattribute__(self, 'accesses').append(name) + return object.__getattribute__(self, name) +request = Request() +request.model = 'from-request' +request.custom_llm_provider = 'mistral' +kwargs = {'model': 'from-kwargs', 'custom_llm_provider': None} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let arguments = arguments(&request, &kwargs); + assert_eq!(arguments.model().unwrap(), "from-kwargs"); + assert_eq!(arguments.custom_llm_provider().unwrap(), None); + let accesses: Vec = request.getattr("accesses").unwrap().extract().unwrap(); + assert_eq!(accesses, Vec::::new()); + }); + } + + #[test] + fn missing_kwargs_read_the_request_property_once() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Request: + def __init__(self): + self.reads = 0 + @property + def model(self): + self.reads += 1 + return 'mistral-ocr-latest' +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + assert_eq!( + arguments(&request, &kwargs).model().unwrap(), + "mistral-ocr-latest" + ); + assert_eq!( + request.getattr("reads").unwrap().extract::().unwrap(), + 1 + ); + }); + } + + #[test] + fn request_property_exceptions_keep_their_identity() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +failure = LookupError('model failed') +class Request: + @property + def model(self): + raise failure +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let error = arguments(&request, &kwargs).model().unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn unused_raising_property_is_never_inspected() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Request: + @property + def unused(self): + raise RuntimeError('unused') + model = 'mistral-ocr-latest' + custom_llm_provider = None +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let arguments = arguments(&request, &kwargs); + assert_eq!(arguments.model().unwrap(), "mistral-ocr-latest"); + assert_eq!(arguments.custom_llm_provider().unwrap(), None); + }); + } + + #[test] + fn document_reader_mutations_are_visible_to_later_field_reads() { + Python::initialize(); + Python::attach(|py| { + stub_timeout_conversion(py); + let locals = eval( + py, + c" +class Request: + api_base = 'original' + timeout = 1 + @property + def document(self): + return document +class Reader: + def read(self): + Request.api_base = 'mutated' + Request.timeout = 9 + return b'abc' +document = {'type': 'file', 'file': Reader()} +request = Request() +kwargs = {} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let arguments = arguments(&request, &kwargs); + let document = arguments.document().unwrap(); + project_document(py, &document).unwrap(); + assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated")); + assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0)); + }); + } + + #[test] + fn captured_api_key_keeps_the_original_python_object() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +key = object() +class Request: + api_key = None +request = Request() +kwargs = {'api_key': key} +", + ); + let request = locals.get_item("request").unwrap().unwrap(); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .unwrap() + .cast_into::() + .unwrap(); + let captured = arguments(&request, &kwargs).api_key().unwrap(); + assert!( + captured + .unbind() + .bind(py) + .is(locals.get_item("key").unwrap().unwrap()) + ); + }); + } + + #[test] + fn file_documents_are_encoded_and_other_documents_keep_the_python_object() { + Python::initialize(); + Python::attach(|py| { + let file = py + .eval( + c"{'type': 'file', 'file': b'%PDF-1.4', 'mime_type': 'application/pdf'}", + None, + None, + ) + .unwrap(); + assert_eq!( + project_document(py, &file).unwrap().0, + serde_json::json!({ + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + }) + ); + + let original = py + .eval( + c"{'type': 'document_url', 'document_url': 'https://example.com/a.pdf'}", + None, + None, + ) + .unwrap(); + let (wire, retained) = project_document(py, &original).unwrap(); + assert_eq!( + wire, + serde_json::json!({ + "type": "document_url", + "document_url": "https://example.com/a.pdf", + }) + ); + assert!(retained.bind(py).is(&original)); + }); + } + + #[test] + fn unknown_document_types_reach_existing_downstream_validation() { + Python::initialize(); + Python::attach(|py| { + let document = py + .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None) + .unwrap(); + let wire_document = project_document(py, &document).unwrap().0; + assert_eq!( + wire_document, + serde_json::json!({"type": "mystery", "mystery": "x"}) + ); + let error = match decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: wire_document, + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Map::new(), + input_sources: Default::default(), + timeout_seconds: None, + }) { + Ok(_) => panic!("unknown discriminators belong to core validation"), + Err(error) => error, + }; + assert!(error.to_string().contains("document")); + }); + } + + #[test] + fn document_discriminator_errors_keep_their_existing_exceptions() { + Python::initialize(); + Python::attach(|py| { + let missing = py.eval(c"{}", None, None).unwrap(); + assert!( + project_document(py, &missing) + .unwrap_err() + .is_instance_of::(py) + ); + + let non_string = py.eval(c"{'type': 1}", None, None).unwrap(); + assert!( + project_document(py, &non_string) + .unwrap_err() + .is_instance_of::(py) + ); + + let locals = eval( + py, + c" +failure = RuntimeError('type lookup failed') +class Document: + def __getitem__(self, key): + raise failure +document = Document() +", + ); + let error = + project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); + } + + #[test] + fn document_classification_happens_once() { + Python::initialize(); + Python::attach(|py| { + let locals = eval( + py, + c" +class Document(dict): + def __init__(self): + super().__init__({'file': b'abc'}) + self.reads = [] + def __getitem__(self, key): + self.reads.append(key) + if key == 'type': + return 'file' if self.reads.count('type') == 1 else 'document_url' + return super().__getitem__(key) +document = Document() +", + ); + let document = locals.get_item("document").unwrap().unwrap(); + let (wire, retained) = project_document(py, &document).unwrap(); + assert_eq!(wire["type"], "document_url"); + assert!(!retained.bind(py).is(&document)); + let reads: Vec = document.getattr("reads").unwrap().extract().unwrap(); + assert_eq!(reads, ["type", "mime_type", "file"]); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs similarity index 52% rename from litellm-rust/crates/python-bridge/src/routes/ocr.rs rename to litellm-rust/crates/python-bridge/src/routes/ocr/value.rs index c5def64c2f1..051ac19d4fb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/value.rs @@ -1,12 +1,11 @@ use litellm_core::Error; use std::future::Future; -use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; -use litellm_core::ocr::wire::{OcrWireRequest, decode_request, is_supported_request}; +use litellm_core::ocr::wire::{OcrWireRequest, decode_request}; use pyo3::prelude::*; use serde_json::Value; -use crate::errors::ocr_error_to_pyerr; +use super::errors::to_pyerr as ocr_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; fn prepare_ocr( @@ -38,37 +37,20 @@ fn prepare_ocr( extra_headers, timeout, } = options; - if is_supported_request(&model, custom_llm_provider.as_deref()) { - let request = decode_request(OcrWireRequest { - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - input_sources, - timeout_seconds: timeout.map(|value| value.as_secs_f64()), - })?; - return litellm_core::ocr::ocr(request) - .await - .map(|response| response.into_json()); - } - run_ocr(OcrRequest { - model: &model, + let request = decode_request(OcrWireRequest { + model, document, - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), + api_key, + api_base, + custom_llm_provider, extra_headers, optional_params, - timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, - }) - .await + input_sources, + timeout_seconds: timeout.map(|value| value.as_secs_f64()), + })?; + litellm_core::ocr::ocr(request) + .await + .map(|response| response.into_json()) }) } @@ -96,22 +78,3 @@ bridge_route! { prepare = prepare_ocr, errors = ocr_error_to_pyerr, } - -#[cfg(test)] -mod tests { - use litellm_core::ocr::wire::is_supported_request; - - #[test] - fn native_activation_includes_migrated_providers() { - assert!(is_supported_request("model", Some("mistral"))); - assert!(is_supported_request("pixtral-12b", Some("azure_ai"))); - assert!(is_supported_request( - "documentintelligence/prebuilt-read", - Some("azure_ai") - )); - assert!(is_supported_request("parse-v3", Some("reducto"))); - assert!(is_supported_request("parse-legacy", Some("reducto"))); - assert!(is_supported_request("mistral-ocr", Some("vertex_ai"))); - assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); - } -} diff --git a/litellm-rust/crates/python-bridge/tests/lifecycle.py b/litellm-rust/crates/python-bridge/tests/lifecycle.py new file mode 100644 index 00000000000..fd6742102a4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/lifecycle.py @@ -0,0 +1,186 @@ +import asyncio +import gc +import threading +import weakref +from contextvars import ContextVar + + +async def exercise(): + caller = asyncio.current_task() + thread = threading.get_ident() + loop = asyncio.get_running_loop() + marker = ContextVar("driver", default="before") + entered = asyncio.Event() + released = asyncio.Event() + result = object() + + class CustomAwaitable: + def __await__(self): + return operation().__await__() + + async def operation(): + assert asyncio.current_task() is caller + assert threading.get_ident() == thread + assert asyncio.get_running_loop() is loop + marker.set("inside") + entered.set() + await released.wait() + assert asyncio.current_task() is caller + assert marker.get() == "inside" + return result + + async def release(): + await entered.wait() + released.set() + + releaser = asyncio.create_task(release()) + execution = await_execution(CustomAwaitable()) + try: + execution.resume_value(None) + except RuntimeError: + pass + else: + raise AssertionError("resumed an unstarted execution") + wrapped = drive(execution) + try: + wrapped.send(1) + except TypeError: + pass + else: + raise AssertionError("accepted initial value") + assert await wrapped is result + assert marker.get() == "inside" + await releaser + execution.close() + execution.close() + try: + await wrapped + except RuntimeError: + pass + else: + raise AssertionError("accepted coroutine reuse") + + final_awaitable = CustomAwaitable() + assert await drive(calling_execution(lambda: final_awaitable)) is final_awaitable + + cause = KeyError("cause") + failure = ValueError("original") + + async def failing(): + await asyncio.sleep(0) + raise failure from cause + + try: + await drive(await_execution(failing())) + except ValueError as error: + assert error is failure + assert error.__cause__ is cause + names = [] + traceback = error.__traceback__ + while traceback: + names.append(traceback.tb_frame.f_code.co_name) + traceback = traceback.tb_next + assert "failing" in names + else: + raise AssertionError("lost original exception") + + for suppress in (False, True): + pending = asyncio.Event() + cleanup_entered = asyncio.Event() + cleanup_release = asyncio.Event() + cleaned = [] + + async def cancel_operation(): + try: + pending.set() + await asyncio.Event().wait() + except asyncio.CancelledError: + if suppress: + return result + raise + finally: + cleanup_entered.set() + try: + await cleanup_release.wait() + except asyncio.CancelledError: + await cleanup_release.wait() + cleaned.append(asyncio.current_task()) + + task = asyncio.create_task(drive(await_execution(cancel_operation()))) + await pending.wait() + task.cancel() + await cleanup_entered.wait() + assert not task.done() + task.cancel() + await asyncio.sleep(0) + cleanup_release.set() + if suppress: + assert await task is result + else: + try: + await task + except asyncio.CancelledError: + pass + else: + raise AssertionError("lost cancellation") + assert cleaned == [task] + + observed = [] + + def reenter(): + try: + active.start() + except RuntimeError as error: + observed.append(str(error)) + return result + + active = calling_execution(reenter) + assert await drive(active) is result + assert observed == ["execution is already running"] + + class Finalizer: + def __call__(self): + return result + + def __del__(self): + self.owner.close() + observed.append("released") + + def cycle(started): + callback = Finalizer() + execution = calling_execution(callback) + callback.owner = execution + if started: + assert execution.start().value is result + return weakref.ref(callback) + + for started in (False, True): + reference = cycle(started) + gc.collect() + assert reference() is None + assert observed[-2:] == ["released", "released"] + + class Awaitable: + def __await__(self): + try: + yield self + finally: + observed.append("unwound") + + def abandoned(started): + awaitable = Awaitable() + coroutine = drive(await_execution(awaitable)) + awaitable.owner = coroutine + if started: + assert coroutine.send(None) is awaitable + coroutine.close() + return weakref.ref(awaitable) + + for started in (False, True): + reference = abandoned(started) + gc.collect() + assert reference() is None + assert observed[-1] == "unwound" + + +asyncio.run(asyncio.wait_for(exercise(), 10)) diff --git a/litellm-rust/crates/python-interop/AGENTS.md b/litellm-rust/crates/python-interop/AGENTS.md index d1d61e5dfa0..63996d3a92b 100644 --- a/litellm-rust/crates/python-interop/AGENTS.md +++ b/litellm-rust/crates/python-interop/AGENTS.md @@ -1 +1,16 @@ -litellm-python-interop is the domain-neutral PyO3 foundation. Keep generic Python/Serde conversion and interpreter primitives here. Do not add LiteLLM domain crates, route types, API registration, or cdylib build features. +- Target invariants; implementation and runtime validation may lag these rules +- Keep this crate a small, domain-neutral foundation: Python/Serde conversion and interpreter-boundary utilities + - No LiteLLM domain dependencies, route types, callback policy, public API registration or cdylib build features + - Generic code alone does not justify extraction: runtime integration stays in `python-bridge/src/execution.rs`, host adaptation in its `lifecycle.rs` +- Use standard PyO3 ownership and conversion APIs + - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads + - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` + - Preserve `PythonizeError`'s standard conversion into `PyErr`; do not stringify original Python exceptions into new `ValueError`s + - Keep serializer-panic containment in `Pythonized`: async output conversion can run in an unjoined blocking task and otherwise strand delivery +- Use `Python::detach` for Rust-only work; Python operations require attachment + - Keep diagnostic counters in the consumer; wrapper invocations do not measure every interpreter release + - Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers; expose retained Python edges to GC without calling Python during traversal +- Keep coroutine driving in the shared Python driver and native adapter + - Driver: `litellm/rust_bridge/lifecycle.py`; handle: `python-bridge/src/lifecycle.rs`; native-backed behavior tests: `python-bridge/tests/lifecycle.py` +- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html) + - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html) diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs index 2e562bdae70..79af79e8c61 100644 --- a/litellm-rust/crates/python-interop/src/lib.rs +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -2,4 +2,6 @@ mod gil; mod marshal; pub use gil::{release_count, release_gil}; -pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py}; +pub use marshal::{ + Pythonized, from_py, from_py_preserving_errors, panic_to_pyerr, to_py, to_py_preserving_errors, +}; diff --git a/litellm-rust/crates/python-interop/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs index a16d1e0ae13..ed4cce862c0 100644 --- a/litellm-rust/crates/python-interop/src/marshal.rs +++ b/litellm-rust/crates/python-interop/src/marshal.rs @@ -14,6 +14,13 @@ where pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) } +pub fn from_py_preserving_errors(value: &Bound<'_, PyAny>) -> PyResult +where + T: DeserializeOwned, +{ + pythonize::depythonize(value).map_err(PyErr::from) +} + pub fn to_py(py: Python<'_>, value: &T) -> PyResult> where T: Serialize + ?Sized, @@ -23,6 +30,15 @@ where .map_err(|error| PyValueError::new_err(error.to_string())) } +pub fn to_py_preserving_errors(py: Python<'_>, value: &T) -> PyResult> +where + T: Serialize + ?Sized, +{ + pythonize::pythonize(py, value) + .map(Bound::unbind) + .map_err(PyErr::from) +} + pub struct Pythonized(pub T); impl<'py, T> IntoPyObject<'py> for Pythonized @@ -89,4 +105,49 @@ mod tests { assert_eq!(error.to_string(), "PanicException: serializer panicked"); }); } + + #[test] + fn depythonize_preserves_python_exception_identity_and_traceback() { + Python::initialize(); + Python::attach(|py| { + let locals = pyo3::types::PyDict::new(py); + py.run( + pyo3::ffi::c_str!( + r#" +failure = LookupError('conversion failed') +cause = ValueError('cause') +class Broken: + def __index__(self): + raise failure from cause +value = Broken() +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let value = locals.get_item("value").unwrap().unwrap(); + let legacy_error = from_py::(&value).unwrap_err(); + assert!(legacy_error.is_instance_of::(py)); + assert!( + !legacy_error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + let error = from_py_preserving_errors::(&value).unwrap_err(); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + assert!( + error + .cause(py) + .unwrap() + .value(py) + .is(locals.get_item("cause").unwrap().unwrap()) + ); + assert!(error.traceback(py).is_some()); + }); + } } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 498d662a906..2d3a99abe81 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -23,7 +23,6 @@ import litellm from litellm import ( _custom_logger_compatible_callbacks_literal, json_logs, - log_raw_request_response, turn_off_message_logging, ) from litellm._logging import ( @@ -563,6 +562,7 @@ class Logging(LiteLLMLoggingBaseClass): self.streaming_chunks: list[Any] = [] # for generating complete stream response self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response + self._native_callback_fast_path: bool = False # Initialize dynamic callbacks self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks @@ -1236,6 +1236,11 @@ class Logging(LiteLLMLoggingBaseClass): additional_args.get("api_base", "") ) + def record_api_call_start_time(self) -> None: + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + if self.model_call_details.get("first_api_call_start_time") is None: + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] + def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API try: @@ -1253,7 +1258,7 @@ class Logging(LiteLLMLoggingBaseClass): additional_args=additional_args, ) # log raw request to provider (like LangFuse) -- if opted in. - if self.log_raw_request_response is True or log_raw_request_response is True: + if self.log_raw_request_response is True or litellm.log_raw_request_response is True: _litellm_params: Final = self.model_call_details.get("litellm_params", {}) _metadata: Final = _litellm_params.get("metadata", {}) or {} try: @@ -1300,15 +1305,7 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - # Set-once first provider-handoff instant. api_call_start_time - # is overwritten on every retry, so it can't measure one-time - # preprocessing; pinning the first attempt excludes retry loops - # + backoff. Logging object only — must NOT go into - # litellm_params["metadata"] (caller request metadata, typed - # Dict[str, str], echoed downstream; a datetime breaks it). - if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] + self.record_api_call_start_time() # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1442,16 +1439,21 @@ class Logging(LiteLLMLoggingBaseClass): """ return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) + def record_post_call( + self, original_response: object, input: object, api_key: object, additional_args: dict[str, object] + ) -> None: + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" + def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" + self.record_post_call(original_response, input, api_key, additional_args) attr: Literal["warning", "debug"] if self.litellm_request_debug: @@ -2116,6 +2118,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, + build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2140,6 +2143,9 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) + if not build_logging_payload: + return + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2201,6 +2207,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, + build_logging_payload: bool = True, ): try: if start_time is None: @@ -2238,6 +2245,7 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, + build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3261,7 +3269,9 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): + def _failure_handler_helper_fn( + self, exception, traceback_exception, start_time=None, end_time=None, build_logging_payload: bool = True + ): if start_time is None: start_time = self.start_time if end_time is None: @@ -3296,6 +3306,9 @@ class Logging(LiteLLMLoggingBaseClass): metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) + if not build_logging_payload: + return start_time, end_time + ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index f5126f81006..3a2af8a5aba 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -92,6 +92,18 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def get_api_key_env_var(self) -> str | None: return AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR + def resolve_connection_params( + self, + *, + api_key: str | None, + api_base: str | None, + dynamic_api_key: str | None, + dynamic_api_base: str | None, + ) -> tuple[str | None, str | None]: + explicit_api_key: Final = None if api_key is None else dynamic_api_key or api_key + explicit_api_base: Final = None if api_base is None else dynamic_api_base or api_base + return explicit_api_key, explicit_api_base + def get_supported_ocr_params(self, model: str) -> list: """ Get supported OCR parameters for Azure Document Intelligence. @@ -618,7 +630,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): except SSRFError as ssrf_err: raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") - poll_headers = {"Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "")} + poll_headers: Final = { + header: raw_response.request.headers[header] + for header in ("Ocp-Apim-Subscription-Key", "Authorization") + if header in raw_response.request.headers + } return operation_url, poll_headers @staticmethod diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 8111f9a194a..bd67dbf1a2a 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -144,9 +144,15 @@ class BaseOCRConfig: """ return None - def supports_rust_bridge(self) -> bool: - """Whether the Rust OCR bridge may serve this config when it is enabled for the provider.""" - return True + def resolve_connection_params( + self, + *, + api_key: str | None, + api_base: str | None, + dynamic_api_key: str | None, + dynamic_api_base: str | None, + ) -> tuple[str | None, str | None]: + return dynamic_api_key or api_key, dynamic_api_base or api_base def get_health_check_document(self) -> DocumentType: return { # mutable-ok: litellm.aocr rejects any document that is not a dict diff --git a/litellm/llms/cohere/ocr/transformation.py b/litellm/llms/cohere/ocr/transformation.py index dd15d5360a6..b55ff4a3cbf 100644 --- a/litellm/llms/cohere/ocr/transformation.py +++ b/litellm/llms/cohere/ocr/transformation.py @@ -144,9 +144,6 @@ class CohereParseConfig(BaseOCRConfig): def get_api_key_env_var(self) -> str | None: return COHERE_API_KEY_ENV_VAR - def supports_rust_bridge(self) -> bool: - return False - def get_health_check_document(self) -> DocumentType: return { # mutable-ok: litellm.aocr rejects any document that is not a dict "type": "image_url", diff --git a/litellm/ocr/input.py b/litellm/ocr/input.py new file mode 100644 index 00000000000..bcb448371c4 --- /dev/null +++ b/litellm/ocr/input.py @@ -0,0 +1,112 @@ +from collections.abc import Mapping +from os import PathLike +from typing import Final, Literal, Protocol, cast # noqa: TID251 # native callables are validated when loaded + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.configuration import rust_ocr_enabled + + +class FileReader(Protocol): + def read(self) -> bytes | str: ... + + +class FileDocument(TypedDict): + type: ReadOnly[Literal["file"]] + file: ReadOnly[bytes | PathLike[str] | FileReader] + mime_type: ReadOnly[NotRequired[str]] + + +class NativeFileDocument(Protocol): + def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ... + + +class NativeUploadDocument(Protocol): + def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ... + + +class NativeMimeType(Protocol): + def __call__(self, file_name: str) -> str: ... + + +_FILE_DOCUMENT: Final = NativeBinding( + "_ocr_file_document", + validate=lambda value: ( + cast( # cast-ok: native export owns the callable signature + NativeFileDocument, value + ) + if callable(value) + else None + ), +) +_UPLOAD_DOCUMENT: Final = NativeBinding( + "_ocr_upload_document", + validate=lambda value: ( + cast( # cast-ok: native export owns the callable signature + NativeUploadDocument, value + ) + if callable(value) + else None + ), +) +_MAX_FILE_BYTES: Final = NativeBinding( + "_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None +) +_MIME_TYPE: Final = NativeBinding( + "_ocr_mime_type", + validate=lambda value: ( + cast( # cast-ok: native export owns the callable signature + NativeMimeType, value + ) + if callable(value) + else None + ), +) +_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024 + + +def get_mime_type(file_path: str) -> str: + native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None + if native is None: + from litellm.ocr import legacy + + return legacy.get_mime_type(file_path) + return native(file_path) + + +def get_max_file_bytes() -> int: + limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None + if limit is None: + return _PYTHON_MAX_FILE_BYTES + return limit + + +def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]: + native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None + if native is None: + from litellm.ocr import legacy + + return legacy.convert_file_document_to_url_document(document) + return native(document) + + +def convert_upload_to_url_document( + file_content: bytes, filename: str | None, content_type: str | None +) -> dict[str, str]: + native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None + if native is None: + from litellm.ocr import legacy + + if len(file_content) > _PYTHON_MAX_FILE_BYTES: + raise ValueError("OCR file exceeds the size limit") + content_mime: Final = content_type.split(";")[0].strip() if content_type else None + mime_type: Final = ( + legacy.get_mime_type(filename) + if filename and (not content_mime or content_mime == "application/octet-stream") + else content_mime or "application/octet-stream" + ) + return legacy.convert_file_document_to_url_document( + {"type": "file", "file": file_content, "mime_type": mime_type} + ) + return native(file_content, filename, content_type) diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py new file mode 100644 index 00000000000..ddf6016dce3 --- /dev/null +++ b/litellm/ocr/legacy.py @@ -0,0 +1,411 @@ +""" +Main OCR function for LiteLLM. +""" + +import asyncio +import base64 +import mimetypes +import os +import re +from collections.abc import Coroutine, Mapping +from dataclasses import dataclass +from io import IOBase +from types import MappingProxyType +from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts + +import httpx + +import litellm +from litellm._logging import verbose_logger +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 ( + OCR_REQUEST_FORMAT_PARAM, + BaseOCRConfig, + OCRResponse, + parse_ocr_request_format, +) +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.ocr.input import FileReader +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager, client + +base_llm_http_handler: Final = BaseLLMHTTPHandler() + + +@dataclass(frozen=True, slots=True) +class _PreparedOCRRequest: + model: str + document: Mapping[str, object] + api_key: str | None + api_base: str | None + custom_llm_provider: str + extra_headers: dict[str, object] | None + provider_config: BaseOCRConfig + optional_params: dict[str, object] + litellm_params: dict[str, object] + effective_timeout: float | httpx.Timeout + litellm_logging_obj: LiteLLMLoggingObj + + +def _prepare_ocr_request( + model: str, + document: Mapping[str, object], + api_key: str | None, + api_base: str | None, + timeout: float | httpx.Timeout | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + kwargs: dict[str, object], +) -> _PreparedOCRRequest: + litellm_logging_obj: Final = cast( # cast-ok: @client supplies the logging object; preserve legacy failure behavior + LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj") + ) + litellm_call_id: Final = cast( # cast-ok: @client supplies the call id without coercion + str | None, kwargs.get("litellm_call_id", None) + ) + + if not isinstance(document, dict): + raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}") + + doc_type = document.get("type") + + if doc_type == "file": + document = convert_file_document_to_url_document(document) + doc_type = document.get("type") + + if doc_type not in ["document_url", "image_url"]: + raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") + + ( + model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + ) + + ocr_provider_config: Final = ProviderConfigManager.get_provider_ocr_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) + + if ocr_provider_config is None: + raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}") + + resolved_api_key, resolved_api_base = ocr_provider_config.resolve_connection_params( + api_key=api_key, + api_base=api_base, + dynamic_api_key=dynamic_api_key, + dynamic_api_base=dynamic_api_base, + ) + + verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider) + + litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) + + supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model) + requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM) + if requested_format is not None: + try: + parsed_format: Final = parse_ocr_request_format(requested_format) + except ValueError as e: + raise litellm.exceptions.UnsupportedParamsError( + message=f"{e}", model=model, llm_provider=custom_llm_provider + ) from e + if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native": + raise litellm.exceptions.UnsupportedParamsError( + message=( + f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, " + f"model: {model}" + ), + model=model, + llm_provider=custom_llm_provider, + ) + + non_default_params: Final = {} + for param in supported_params: + if param in kwargs: + non_default_params[param] = kwargs.pop(param) + + optional_params: Final = ocr_provider_config.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + ) + + verbose_logger.debug("OCR optional_params after mapping: %s", optional_params) + + effective_timeout: Final = timeout or request_timeout + + litellm_logging_obj.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params={ + "litellm_call_id": litellm_call_id, + "api_base": resolved_api_base, + }, + custom_llm_provider=custom_llm_provider, + ) + + return _PreparedOCRRequest( + model=model, + document=document, + api_key=resolved_api_key, + api_base=resolved_api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + provider_config=ocr_provider_config, + optional_params=cast( + dict[str, object], optional_params + ), # cast-ok: provider configs return heterogeneous OCR options + litellm_params=dict(litellm_params), + effective_timeout=effective_timeout, + litellm_logging_obj=litellm_logging_obj, + ) + + +def _error_provider(model: str, custom_llm_provider: str | None) -> str | None: + if custom_llm_provider is not None: + return custom_llm_provider + prefix: Final = model.partition("/")[0] + if prefix in {"mistral", "azure_ai", "vertex_ai"}: + return prefix + return "mistral" if model.startswith("mistral-ocr") else None + + +@client +async def aocr( + model: str, + document: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + **kwargs: object, # kwargs-ok: public OCR accepts provider-specific options +) -> OCRResponse: + completion_kwargs: Final[dict[str, object]] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } + try: + prepared: Final = _prepare_ocr_request( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) + + response = base_llm_http_handler.ocr( + model=prepared.model, + document=cast( # cast-ok: preserve legacy document fields for provider validation + dict[str, str], prepared.document + ), + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=True, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + if asyncio.iscoroutine(response): + response = await response + + if response is None: + raise ValueError(f"Got an unexpected None response from the OCR API: {response}") + + return response + except Exception as e: + error_provider: Final = _error_provider(model, custom_llm_provider) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model + raise litellm.exception_type( + model=error_model, + custom_llm_provider=error_provider, + original_exception=e, + completion_kwargs=completion_kwargs, + extra_kwargs=kwargs, + ) + + +_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$") + +_MIME_TYPE_MAP: Final = MappingProxyType( + { + ".pdf": "application/pdf", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".bmp": "image/bmp", + } +) + + +def get_mime_type(file_path: str) -> str: + ext: Final = os.path.splitext(file_path)[1].lower() + mime: Final = _MIME_TYPE_MAP.get(ext) + if mime: + return mime + guessed, _ = mimetypes.guess_type(file_path) + return guessed or "application/octet-stream" + + +def _read_file(file_input: object) -> tuple[bytes, str, str | None]: + if isinstance(file_input, str): + raise ValueError( + "OCR file input does not accept bare str values. Pass bytes, " + "a pathlib.Path, or a file-like object. To OCR a local file " + "from a path, call open(path, 'rb') yourself." + ) + if isinstance(file_input, os.PathLike): + file_path: Final = str(cast(object, file_input)) # cast-ok: preserve staging's str(PathLike) conversion + if not os.path.isfile(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + mime_type: Final = get_mime_type(file_path) + with open(file_path, "rb") as stream: + return stream.read(), mime_type, os.path.basename(file_path) + if isinstance(file_input, bytes): + return file_input, "application/octet-stream", None + if isinstance(file_input, IOBase) or hasattr(file_input, "read"): + file_name: Final = cast( # cast-ok: retain legacy validation and errors for file-like metadata + str | None, getattr(file_input, "name", None) + ) + inferred_mime: Final = get_mime_type(file_name) if file_name else "application/octet-stream" + reader: Final = cast(FileReader, file_input) # cast-ok: legacy accepts duck-typed file readers + content: Final = reader.read() + return content.encode("utf-8") if isinstance(content, str) else content, inferred_mime, file_name + raise ValueError( + f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object." + ) + + +def convert_file_document_to_url_document(document: Mapping[str, object]) -> dict[str, str]: + file_input: Final = document.get("file") + if file_input is None: + raise ValueError( + "document with type='file' must include a 'file' field containing " + "a pathlib.Path, file-like object, or bytes" + ) + file_bytes, inferred_mime, file_name = _read_file(file_input) + if not file_bytes: + raise ValueError("File is empty or could not be read") + mime_type: Final = cast( # cast-ok: keep staging's MIME validation errors + str, document.get("mime_type", inferred_mime) + ) + if not _MIME_PATTERN.match(mime_type): + raise ValueError(f"Invalid MIME type: {mime_type}") + + base64_data: Final = base64.b64encode(file_bytes).decode("utf-8") + data_uri: Final = f"data:{mime_type};base64,{base64_data}" + + if mime_type.startswith("image/"): + verbose_logger.debug( + "OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)", + mime_type, + len(file_bytes), + file_name, + ) + return {"type": "image_url", "image_url": data_uri} + + verbose_logger.debug( + "OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)", + mime_type, + len(file_bytes), + file_name, + ) + return {"type": "document_url", "document_url": data_uri} + + +@client +def ocr( + model: str, + document: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, object] | None = None, + **kwargs: object, # kwargs-ok: public OCR accepts provider-specific options +) -> OCRResponse | Coroutine[object, object, OCRResponse]: + completion_kwargs: Final[dict[str, object]] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } + try: + _is_async: Final = kwargs.pop("aocr", False) is True + completion_kwargs["aocr"] = _is_async + prepared: Final = _prepare_ocr_request( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + kwargs=kwargs, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout=timeout, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) + + response: Final = base_llm_http_handler.ocr( + model=prepared.model, + document=cast( # cast-ok: preserve legacy document fields for provider validation + dict[str, str], prepared.document + ), + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=_is_async, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + return response + except Exception as e: + error_provider: Final = _error_provider(model, custom_llm_provider) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model + raise litellm.exception_type( + model=error_model, + custom_llm_provider=error_provider, + original_exception=e, + completion_kwargs=completion_kwargs, + extra_kwargs=kwargs, + ) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 56bfd98895d..382c5d6aae4 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -1,460 +1,20 @@ -""" -Main OCR function for LiteLLM. -""" - -import asyncio -import base64 -import mimetypes -import os -import re -from collections.abc import Callable, Coroutine, Mapping -from dataclasses import dataclass -from io import IOBase -from types import MappingProxyType -from typing import Any, Final, cast +from collections.abc import Awaitable, Callable, Coroutine, Mapping +from typing import Final, cast # noqa: TID251 # native binding selects a sync result or an async awaitable import httpx -import litellm -from litellm._logging import verbose_logger -from litellm.constants import request_timeout -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.azure_ai.ocr.common_utils import ( - is_azure_cohere_parse_model, - is_azure_document_intelligence_model, -) -from litellm.llms.base_llm.ocr.transformation import ( - OCR_REQUEST_FORMAT_PARAM, - BaseOCRConfig, - OCRResponse, - parse_ocr_request_format, -) -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.rust_bridge import ocr as rust_ocr_bridge +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import legacy +from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type from litellm.rust_bridge.bindings import native_exception_types -from litellm.rust_bridge.configuration import rust_enabled -from litellm.types.router import GenericLiteLLMParams -from litellm.utils import ProviderConfigManager, client +from litellm.rust_bridge.configuration import rust_ocr_enabled +from litellm.rust_bridge.ocr import LiteLLMOcrRequest +from litellm.rust_bridge.ocr_lifecycle import select -####### ENVIRONMENT VARIABLES ################### -base_llm_http_handler = BaseLLMHTTPHandler() -################################################# +__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr") -@dataclass -class _PreparedOCRRequest: - model: str - document: dict[str, Any] - api_key: str | None - api_base: str | None - custom_llm_provider: str - extra_headers: dict[str, object] | None - provider_config: BaseOCRConfig - optional_params: dict[str, object] - litellm_params: dict[str, object] - effective_timeout: float | httpx.Timeout - litellm_logging_obj: LiteLLMLoggingObj - caller_supplied_api_key: bool = True - caller_supplied_api_base: bool = True - - -_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) -_RUST_OCR_CONFIG_FIELDS: Final = frozenset( - { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_scope", - "azure_authority_host", - "azure_credential", - "azure_federated_token_file", - "vertex_credentials", - "vertex_ai_credentials", - "vertex_project", - "vertex_ai_project", - "vertex_location", - "vertex_ai_location", - } -) -_RUST_OCR_SECRET_FIELDS: Final = frozenset( - {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} -) - - -def _prepare_ocr_request( - model: str, - document: Mapping[str, object], - api_key: str | None, - api_base: str | None, - timeout: float | httpx.Timeout | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - kwargs: dict[str, object], -) -> _PreparedOCRRequest: - litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) - litellm_call_id: Final = cast(str | None, kwargs.get("litellm_call_id", None)) - - if not isinstance(document, dict): - raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}") - - doc_type = document.get("type") - - if doc_type == "file": - document = convert_file_document_to_url_document(document) - doc_type = document.get("type") - - if doc_type not in ["document_url", "image_url"]: - raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") - - caller_supplied_api_key: Final = api_key is not None - caller_supplied_api_base: Final = api_base is not None - - ( - model, - custom_llm_provider, - dynamic_api_key, - dynamic_api_base, - ) = litellm.get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - api_key=api_key, - ) - - suppress_dynamic_api_base: Final = ( - not caller_supplied_api_base - and custom_llm_provider == "azure_ai" - and is_azure_document_intelligence_model(model) - ) - if dynamic_api_key: - api_key = dynamic_api_key - if dynamic_api_base and not suppress_dynamic_api_base: - api_base = dynamic_api_base - - ocr_provider_config: Final = ProviderConfigManager.get_provider_ocr_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) - - if ocr_provider_config is None: - raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}") - - verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider) - - litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs) - - supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model) - requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM) - if requested_format is not None: - try: - parsed_format: Final = parse_ocr_request_format(requested_format) - except ValueError as e: - raise litellm.exceptions.UnsupportedParamsError( - message=f"{e}", model=model, llm_provider=custom_llm_provider - ) from e - if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native": - raise litellm.exceptions.UnsupportedParamsError( - message=( - f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, " - f"model: {model}" - ), - model=model, - llm_provider=custom_llm_provider, - ) - - non_default_params: Final = {} - for param in supported_params: - if param in kwargs: - non_default_params[param] = kwargs.pop(param) - - optional_params: Final = ocr_provider_config.map_ocr_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - ) - - verbose_logger.debug("OCR optional_params after mapping: %s", optional_params) - - effective_timeout: Final = timeout or request_timeout - - litellm_logging_obj.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params={ - "litellm_call_id": litellm_call_id, - "api_base": api_base, - }, - custom_llm_provider=custom_llm_provider, - ) - - return _PreparedOCRRequest( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - provider_config=ocr_provider_config, - optional_params=cast(dict[str, object], optional_params), - litellm_params=dict(litellm_params), - effective_timeout=effective_timeout, - litellm_logging_obj=litellm_logging_obj, - caller_supplied_api_key=caller_supplied_api_key, - caller_supplied_api_base=caller_supplied_api_base, - ) - - -def _rust_ocr_provider(request: rust_ocr_bridge.LiteLLMOcrRequest) -> str | None: - if request.custom_llm_provider is not None: - return request.custom_llm_provider - prefix: Final = request.model.partition("/")[0] - if prefix in _RUST_OCR_PROVIDERS: - return prefix - if request.model.startswith("mistral-ocr"): - return "mistral" - return None - - -def _rust_ocr_supported(request: rust_ocr_bridge.LiteLLMOcrRequest) -> bool: - provider: Final = _rust_ocr_provider(request) - if provider not in _RUST_OCR_PROVIDERS or request.kwargs.get(OCR_REQUEST_FORMAT_PARAM) == "native": - return False - if provider == "azure_ai": - return ( - not is_azure_cohere_parse_model(request.model) - and not callable(request.kwargs.get("azure_ad_token_provider")) - and request.kwargs.get("azure_username") is None - and request.kwargs.get("azure_password") is None - ) - return True - - -def _rust_bridge_optional_params( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], -) -> Mapping[str, object]: - optional_params: Final = MappingProxyType( - { - name: value - for name, value in request.kwargs.items() - if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) - and name not in {"litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request"} - } - ) - provider: Final = _rust_ocr_provider(request) - if provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: - return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) - if provider != "vertex_ai": - return optional_params - project: Final = ( - request.kwargs.get("vertex_project") - or request.kwargs.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - location: Final = ( - request.kwargs.get("vertex_location") - or request.kwargs.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - credentials: Final = ( - request.kwargs.get("vertex_credentials") - or request.kwargs.get("vertex_ai_credentials") - or resolve_secret("VERTEXAI_CREDENTIALS") - ) - vertex_params: Final = MappingProxyType( - { - name: value - for name, value in ( - ("vertex_project", project), - ("vertex_location", location), - ("vertex_credentials", credentials), - ) - if value is not None - } - ) - return MappingProxyType({**optional_params, **vertex_params}) - - -def _rust_bridge_input_sources( - request: rust_ocr_bridge.LiteLLMOcrRequest, - optional_params: Mapping[str, object], -) -> Mapping[str, str]: - proxy_request: Final = request.kwargs.get("proxy_server_request") - if not isinstance(proxy_request, Mapping): - return MappingProxyType({}) - proxy_request_mapping: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types - Mapping[object, object], proxy_request - ) - body_value: Final = proxy_request_mapping.get("body") - if not isinstance(body_value, Mapping): - return MappingProxyType({}) - body: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types - Mapping[object, object], body_value - ) - credential_fields_value: Final = proxy_request_mapping.get("credential_fields", ()) - credential_fields: Final = ( - frozenset(name for name in credential_fields_value if isinstance(name, str)) - if isinstance(credential_fields_value, (list, tuple, set, frozenset)) - else frozenset() - ) - names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) - request_sources: Final = MappingProxyType( - {name: "request" for name in names if name in body or name in credential_fields} - ) - if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: - return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) - return request_sources - - -def _marshal_rust_ocr_request( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], -) -> rust_ocr_bridge.LiteLLMOcrRequest: - if not isinstance(request.document, dict): - raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") - document: Final = ( - convert_file_document_to_url_document(request.document) - if request.document.get("type") == "file" - else request.document - ) - provider: Final = _rust_ocr_provider(request) - api_key: Final = request.api_key or resolve_secret("MISTRAL_API_KEY") if provider == "mistral" else request.api_key - optional_params: Final = _rust_bridge_optional_params(request, resolve_secret) - input_sources: Final = _rust_bridge_input_sources(request, optional_params) - logged_optional_params: Final = MappingProxyType( - {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} - ) - logged_kwargs: Final = MappingProxyType( - { - name: "****" if name in _RUST_OCR_SECRET_FIELDS else value - for name, value in request.kwargs.items() - if name != "proxy_server_request" - } - ) - logging_obj: Final = cast( # cast-ok: bridge kwargs carry the prepared logging object - LiteLLMLoggingObj, request.kwargs["litellm_logging_obj"] - ) - logging_obj.update_from_kwargs( - kwargs=dict(logged_kwargs), # mutable-ok: logging API requires an owned dict - model=request.model, - optional_params=dict(logged_optional_params), # mutable-ok: logging API requires an owned dict - litellm_params={ - "litellm_call_id": request.kwargs.get("litellm_call_id"), - "api_base": request.api_base, - }, # mutable-ok: legacy logging requires a concrete params dict - custom_llm_provider=provider, - ) - logging_obj.pre_call( - input="OCR document processing", - api_key=api_key, - additional_args={ # mutable-ok: pre_call mutates the additional_args dict - "complete_input_dict": { - "model": request.model, - "document": document, - **logged_optional_params, - }, # mutable-ok: callbacks consume a JSON-serializable request dict - "api_base": request.api_base or "", - "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict - }, - ) - return rust_ocr_bridge.LiteLLMOcrRequest( - model=request.model, - document=document, - api_key=api_key, - api_base=request.api_base, - timeout=request.timeout if request.timeout is not None else request_timeout, - custom_llm_provider=request.custom_llm_provider, - extra_headers=request.extra_headers, - kwargs=optional_params, - input_sources=input_sources, - ) - - -def _map_rust_ocr_error( - error: Exception, - request: rust_ocr_bridge.LiteLLMOcrRequest, - exception_types: tuple[type[BaseException], type[BaseException]] | None, -) -> Exception: - if exception_types is None or not isinstance(error, exception_types[1]): - return error - provider: Final = _rust_ocr_provider(request) - if provider is None: - return error - provider_config: Final = ProviderConfigManager.get_provider_ocr_config( - model=request.model.removeprefix(f"{provider}/"), provider=litellm.LlmProviders(provider) - ) - if provider_config is None: - return error - error_args: Final = cast( # cast-ok: Python exceptions expose positional args as a tuple - tuple[object, ...], error.args - ) - status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 - message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) - error_factory: Final = cast( # cast-ok: provider configs expose heterogeneous exception factories - Callable[..., Exception], provider_config.get_error_class - ) - return error_factory( - error_message=message, status_code=status or 500, headers={} - ) # mutable-ok: provider error factories require a concrete headers dict - - -def _run_rust_ocr( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_ocr() is None: - return None - marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) - input_sources: Final = marshalled.input_sources - try: - response: Final = rust_ocr_bridge.ocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_rust_ocr_error(error, request, native_exception_types()) from error - return OCRResponse.model_validate(response) if response is not None else None - - -async def _run_rust_aocr( - request: rust_ocr_bridge.LiteLLMOcrRequest, - resolve_api_key: Callable[[str], str | None], -) -> OCRResponse | None: - if rust_ocr_bridge.load_rust_aocr() is None: - return None - marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) - input_sources: Final = marshalled.input_sources - try: - response: Final = await rust_ocr_bridge.aocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_rust_ocr_error(error, request, native_exception_types()) from error - return OCRResponse.model_validate(response) if response is not None else None - - -@client -async def aocr( +def _bind_request( model: str, document: Mapping[str, object], api_key: str | None = None, @@ -462,77 +22,9 @@ async def aocr( timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, - **kwargs: object, -) -> OCRResponse: - """ - Async OCR function. - - Args: - model: Model name (e.g., "mistral/mistral-ocr-latest") - document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs, - {"type": "image_url", "image_url": "https://..."} for images, or - {"type": "file", "file": } for local files - api_key: Optional API key - api_base: Optional API base URL - timeout: Optional timeout - custom_llm_provider: Optional custom LLM provider - extra_headers: Optional extra headers - **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - - Returns: - OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - - Example: - ```python - import litellm - - # OCR with PDF - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": "https://arxiv.org/pdf/2201.04234" - }, - include_image_base64=True - ) - - # OCR with image - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "image_url", - "image_url": "https://example.com/image.png" - } - ) - - # OCR with base64 encoded PDF - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": f"data:application/pdf;base64,{base64_pdf}" - } - ) - - # OCR with local file - response = await litellm.aocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": "/path/to/document.pdf"} - ) - ``` - """ - completion_kwargs: Final[dict[str, object]] = { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "timeout": timeout, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "kwargs": kwargs, - } - request: Final = rust_ocr_bridge.LiteLLMOcrRequest( + **kwargs: object, # kwargs-ok: public OCR accepts provider-specific options +) -> LiteLLMOcrRequest: + return LiteLLMOcrRequest( model=model, document=document, api_key=api_key, @@ -542,340 +34,50 @@ async def aocr( extra_headers=extra_headers, kwargs=kwargs, ) + + +def _public_request(name: str, args: tuple[object, ...], kwargs: dict[str, object]) -> LiteLLMOcrRequest: try: - if rust_enabled() and _rust_ocr_supported(request): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = await _run_rust_aocr( - request=request, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - - prepared: Final = _prepare_ocr_request( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, - ) - model = prepared.model - custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - - response = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=True, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) - - if asyncio.iscoroutine(response): - response = await response - - if response is None: - raise ValueError(f"Got an unexpected None response from the OCR API: {response}") - - return response - except Exception as e: - error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) - error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model - raise litellm.exception_type( - model=error_model, - custom_llm_provider=error_provider, - original_exception=e, - completion_kwargs=completion_kwargs, - extra_kwargs=kwargs, - ) + return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation + except TypeError as error: + raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None -################################################# -# Public utilities — used by the SDK and the proxy -################################################# - -_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$") - -_MIME_TYPE_MAP: Final = { - ".pdf": "application/pdf", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".webp": "image/webp", - ".tiff": "image/tiff", - ".tif": "image/tiff", - ".bmp": "image/bmp", -} - - -def get_mime_type(file_path: str) -> str: - """ - Determine MIME type from file path extension. - - Falls back to mimetypes.guess_type, then to 'application/octet-stream'. - """ - ext: Final = os.path.splitext(file_path)[1].lower() - mime: Final = _MIME_TYPE_MAP.get(ext) - if mime: - return mime - guessed, _ = mimetypes.guess_type(file_path) - return guessed or "application/octet-stream" - - -def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]: - """ - Convert a file-type document dict to a document_url-type document dict - with an inline base64 data URI. - - Accepts document dicts like: - {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path - {"type": "file", "file": } # file-like object (BinaryIO) - {"type": "file", "file": b"raw bytes"} # raw bytes - - Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or - ``open(path, "rb")`` instead. See the str check below for the rationale. - - Returns: - {"type": "document_url", "document_url": "data:;base64,"} - or {"type": "image_url", "image_url": "data:;base64,"} - """ - file_input: Final = document.get("file") - if file_input is None: - raise ValueError( - "document with type='file' must include a 'file' field containing " - "a pathlib.Path, file-like object, or bytes" - ) - - file_bytes: bytes - mime_type: str = "application/octet-stream" - file_name: str | None = None - - if isinstance(file_input, str): - # Bare strings are rejected here. The OCR ``document`` accepts a - # ``{"type": "file", "file": }`` shape, and when this helper - # runs in a proxy request handler ```` is attacker-controlled. - # Opening it as a path is an arbitrary local file read on the proxy - # host, which is then base64-encoded and forwarded to the OCR - # provider — an exfiltration primitive. - raise ValueError( - "OCR file input does not accept bare str values. Pass bytes, " - "a pathlib.Path, or a file-like object. To OCR a local file " - "from a path, call open(path, 'rb') yourself." - ) - if isinstance(file_input, os.PathLike): - # os.PathLike (pathlib.Path and custom __fspath__ classes) is a - # Python-level type that HTTP form values can't fabricate. - file_path: Final = str(file_input) - if not os.path.isfile(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - mime_type = get_mime_type(file_path) - file_name = os.path.basename(file_path) - with open(file_path, "rb") as f: - file_bytes = f.read() - elif isinstance(file_input, bytes): - file_bytes = file_input - elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): - if hasattr(file_input, "name"): - file_name = getattr(file_input, "name", None) - if file_name: - mime_type = get_mime_type(file_name) - file_bytes = file_input.read() - if isinstance(file_bytes, str): - file_bytes = file_bytes.encode("utf-8") - else: - raise ValueError( - f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object." - ) - - if not file_bytes: - raise ValueError("File is empty or could not be read") - - if "mime_type" in document: - mime_type = document["mime_type"] - - if not _MIME_PATTERN.match(mime_type): - raise ValueError(f"Invalid MIME type: {mime_type}") - - base64_data: Final = base64.b64encode(file_bytes).decode("utf-8") - data_uri: Final = f"data:{mime_type};base64,{base64_data}" - - if mime_type.startswith("image/"): - verbose_logger.debug( - "OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)", - mime_type, - len(file_bytes), - file_name, - ) - return {"type": "image_url", "image_url": data_uri} - - verbose_logger.debug( - "OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)", - mime_type, - len(file_bytes), - file_name, - ) - return {"type": "document_url", "document_url": data_uri} - - -@client def ocr( - model: str, - document: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - timeout: float | httpx.Timeout | None = None, - custom_llm_provider: str | None = None, - extra_headers: dict[str, object] | None = None, - **kwargs: object, + *args: object, + **kwargs: object, # kwargs-ok: preserve the public OCR call shape ) -> OCRResponse | Coroutine[object, object, OCRResponse]: - """ - Synchronous OCR function. - - Args: - model: Model name (e.g., "mistral/mistral-ocr-latest") - document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs, - {"type": "image_url", "image_url": "https://..."} for images, or - {"type": "file", "file": } for local files - api_key: Optional API key - api_base: Optional API base URL - timeout: Optional timeout - custom_llm_provider: Optional custom LLM provider - extra_headers: Optional extra headers - **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - - Returns: - OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - - Example: - ```python - import litellm - - # OCR with PDF - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": "https://arxiv.org/pdf/2201.04234" - }, - include_image_base64=True - ) - - # OCR with image - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "image_url", - "image_url": "https://example.com/image.png" - } - ) - - # OCR with base64 encoded PDF - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": f"data:application/pdf;base64,{base64_pdf}" - } - ) - - # OCR with local file - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": "/path/to/document.pdf"} - ) - - # Access pages - for page in response.pages: - print(f"Page {page.index}: {page.markdown}") - ``` - """ - completion_kwargs: Final[dict[str, object]] = { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "timeout": timeout, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "kwargs": kwargs, - } - request: Final = rust_ocr_bridge.LiteLLMOcrRequest( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - timeout=timeout, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - kwargs=kwargs, - ) - try: - _is_async: Final = kwargs.pop("aocr", False) is True - completion_kwargs["aocr"] = _is_async - if rust_enabled() and _rust_ocr_supported(request): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = _run_rust_ocr( - request=request, - resolve_api_key=get_secret_str, + request: Final = _public_request("ocr", args, kwargs) + native: Final = select(request) if rust_ocr_enabled() else None + if native is not None: + try: + return cast( # cast-ok: False selects the synchronous result + OCRResponse, native(request, args, kwargs, False) ) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response + except _decline_types(): + pass + fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator + Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]], legacy.ocr + ) + return fallback(*args, **kwargs) - prepared: Final = _prepare_ocr_request( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - kwargs=kwargs, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout=timeout, - ) - model = prepared.model - custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - response: Final = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=_is_async, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) +async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape + request: Final = _public_request("aocr", args, kwargs) + native: Final = select(request) if rust_ocr_enabled() else None + if native is not None: + try: + return await cast( # cast-ok: True selects the asynchronous result + Awaitable[OCRResponse], native(request, args, kwargs, True) + ) + except _decline_types(): + pass + fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator + Callable[..., Awaitable[OCRResponse]], legacy.aocr + ) + return await fallback(*args, **kwargs) - return response - except Exception as e: - error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) - error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model - raise litellm.exception_type( - model=error_model, - custom_llm_provider=error_provider, - original_exception=e, - completion_kwargs=completion_kwargs, - extra_kwargs=kwargs, - ) + +def _decline_types() -> tuple[type[BaseException], ...]: + exception_types: Final = native_exception_types() + return (exception_types[0],) if exception_types is not None else () diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9f58aaf24f1..289fe086379 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3181,6 +3181,11 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can exercise the production gating logic directly rather than reimplementing the finally block. """ + if getattr(logging_obj, "call_type", None) in ("ocr", "aocr"): + pending: Final = getattr(logging_obj, "_native_pending_logging", None) + if pending is not None: + logging_obj._native_pending_logging = None # rebind-ok: consume the native OCR release signal once + pending.release(not exception_raised) _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is None: return diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 64a47f4f4ff..ee5cd7c4cb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -315,7 +315,6 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is None: call_type = _infer_call_type(call_type=None, completion_response=response) - # Fallback: resolve call_type from logging_obj for pass-through endpoints if call_type is None: litellm_logging_obj: Final = data.get("litellm_logging_obj") logging_call_type: Final = ( @@ -324,6 +323,8 @@ class UnifiedLLMGuardrails(CustomLogger): if logging_call_type in ( CallTypes.pass_through.value, CallTypes.allm_passthrough_route.value, + CallTypes.ocr.value, + CallTypes.aocr.value, ): call_type = logging_call_type diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index ebf4d988fdd..53ebbe91b54 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -15,7 +15,7 @@ from litellm.llms.base_llm.ocr.transformation import ( OCRResponse, parse_ocr_request_format, ) -from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type +from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -28,24 +28,7 @@ def _build_document_from_upload( filename: str | None, content_type: str | None, ) -> dict[str, str]: - """ - Convert uploaded file bytes into a Mistral-format document dict with base64 data URI. - - Delegates to convert_file_document_to_url_document after resolving MIME type - from the upload's content_type header or filename. - """ - mime_type = content_type.split(";")[0].strip() if content_type else None - if not mime_type or mime_type == "application/octet-stream": - if filename: - mime_type = get_mime_type(filename) - - return convert_file_document_to_url_document( - { - "type": "file", - "file": file_content, - "mime_type": mime_type or "application/octet-stream", - } - ) + return convert_upload_to_url_document(file_content, filename, content_type) def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]: @@ -120,7 +103,7 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: # Seek to start in case the file was already partially read by middleware await uploaded_file.seek(0) - file_content: Final = await uploaded_file.read() + file_content: Final = await uploaded_file.read(get_max_file_bytes() + 1) if not file_content: raise ValueError("Uploaded file is empty") diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 5582027bb5d..ff2e389a6bb 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -42,6 +42,17 @@ def rust_enabled() -> bool: ) +def rust_ocr_enabled() -> bool: + environment: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)) + if environment is False: + return False + return resolve_rust_enabled( + process_override=_CONFIGURATION.override, + environment_override=environment, + release_default=True, + ) + + def reset_rust_configuration() -> None: _CONFIGURATION.override = None diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py new file mode 100644 index 00000000000..f5e0c1b0fc6 --- /dev/null +++ b/litellm/rust_bridge/lifecycle.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import datetime +import os +import uuid +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + + +@dataclass(frozen=True, slots=True) +class Await: + awaitable: Awaitable[object] + + +@dataclass(frozen=True, slots=True) +class Complete: + value: object + + +class Execution(Protocol): + def start(self) -> Await | Complete: ... + + def resume_value(self, value: object) -> Await | Complete: ... + + def resume_error(self, error: BaseException) -> Await | Complete: ... + + def close(self) -> None: ... + + +async def drive(execution: Execution) -> object: + try: + step = execution.start() # rebind-ok: the execution protocol advances after each selected await + while isinstance(step, Await): + try: + value = await step.awaitable # rebind-ok: each selected await produces the next protocol input + except GeneratorExit: + raise + except BaseException as error: + step = execution.resume_error(error) # rebind-ok: advance the execution protocol + else: + step = execution.resume_value(value) # rebind-ok: advance the execution protocol + return step.value + finally: + execution.close() + + +class MetadataUpdater(Protocol): + def __call__( + self, + result: object, + logging_obj: Logging, + model: str | None, + kwargs: dict[str, object], + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: ... + + +@dataclass(frozen=True, slots=True) +class CallSetup: + logger: Logging + kwargs: dict[str, object] + + +def setup( + call_type: str, + args: tuple[object, ...], + kwargs: Mapping[str, object], + start_time: datetime.datetime, + asynchronous: bool, +) -> CallSetup: + from litellm import utils + from litellm.litellm_core_utils.litellm_logging import Logging + + arguments: Final = { # mutable-ok: function_setup consumes an owned kwargs dict + "litellm_call_id": str(uuid.uuid4()), + **kwargs, + } + supplied: Final = arguments.get("litellm_logging_obj") + if isinstance(supplied, Logging): + supplied._native_callback_fast_path = False # pyright: ignore[reportPrivateUsage] # supplied loggers retain all dispatch contracts + return CallSetup(supplied, arguments) + logger, prepared = utils.function_setup( + call_type, utils.Rules(), start_time, *args, is_async_call=asynchronous, **arguments + ) + if type(logger) is Logging and call_type in ("ocr", "aocr"): + logger._native_callback_fast_path = True # pyright: ignore[reportPrivateUsage] # only bridge-created OCR loggers opt into callback elision + return CallSetup(logger, prepared) + + +def check_limits(kwargs: Mapping[str, object]) -> None: + import litellm + + current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor + if litellm.max_budget and current_cost > litellm.max_budget: + raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) + metadata: Final = kwargs.get("metadata") + if isinstance(metadata, Mapping): + typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata + Mapping[str, object], metadata + ) + previous: Final = typed_metadata.get("previous_models") + if ( + isinstance(previous, list) + and litellm.num_retries_per_request is not None + and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history + >= litellm.num_retries_per_request + ): + raise RuntimeError("Max retries per request hit!") + + +def finalize( + response: object, + logger: Logging, + kwargs: dict[str, object], + start_time: datetime.datetime, + end_time: datetime.datetime, +) -> None: + from litellm.litellm_core_utils.llm_response_utils import response_metadata + + model: Final = kwargs.get("model") + update: Final = cast( # cast-ok: legacy metadata function accepts concrete kwargs + MetadataUpdater, response_metadata.update_response_metadata + ) + update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) + + +def deployment_callbacks_needed() -> bool: + import litellm + from litellm.integrations.custom_logger import CustomLogger + + return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) + + +def callbacks_needed(logger: Logging, phase: str) -> bool: + import litellm + from litellm._logging import ( + _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging + ) + + if ( + _is_debugging_on() + or getattr(logger, "litellm_request_debug", False) + or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") + ): + return True + input_needed: Final = bool( + litellm.input_callback + or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + or logger.dynamic_input_callbacks + or callable(getattr(logger, "logger_fn", None)) + or logger.log_raw_request_response + or litellm.log_raw_request_response + ) + match phase: + case "input": + return input_needed + case "sync_success": + return bool(litellm.success_callback or logger.dynamic_success_callbacks) + case "sync_success_async": + return bool( + (litellm.success_callback or logger.dynamic_success_callbacks) + and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks + ) + case "async_success": + return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + case "sync_failure": + return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) + case "async_failure": + return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + case "payload": + return bool( + input_needed + or litellm.success_callback + or litellm.failure_callback + or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor + or logger.dynamic_success_callbacks + or logger.dynamic_async_success_callbacks + or logger.dynamic_failure_callbacks + or logger.dynamic_async_failure_callbacks + ) + case _: + return True + + +def success_bookkeeping( + logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> None: + phase: Final = "async_success" if asynchronous else "sync_success" + if logger.should_run_logging(phase): + logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload + result=response, start_time=start, end_time=end, build_logging_payload=False + ) + logger.has_run_logging(phase) + + +def failure_bookkeeping( + logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> None: + phase: Final = "async_failure" if asynchronous else "sync_failure" + if logger.should_run_logging(phase): + logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload + error, "", start, end, build_logging_payload=False + ) + logger.has_run_logging(phase) diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 89eab71ccba..de8a93dd8b1 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,44 +2,16 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Mapping from dataclasses import dataclass from types import MappingProxyType from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx -import litellm -from litellm.constants import request_timeout -from litellm.llms.azure_ai.ocr.common_utils import is_azure_cohere_parse_model from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse -from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds -from litellm.types.router import GenericLiteLLMParams -from litellm.utils import ProviderConfigManager - -_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) -_RUST_OCR_CONFIG_FIELDS: Final = frozenset( - { - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_scope", - "azure_authority_host", - "azure_credential", - "azure_federated_token_file", - "vertex_credentials", - "vertex_ai_credentials", - "vertex_project", - "vertex_ai_project", - "vertex_location", - "vertex_ai_location", - } -) -_RUST_OCR_SECRET_FIELDS: Final = frozenset( - {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} -) @dataclass(frozen=True, slots=True) @@ -87,26 +59,6 @@ class RustAocr(Protocol): raise NotImplementedError -class _OCRLogging(Protocol): - def update_from_kwargs( - self, - *, - kwargs: dict[str, object], - model: str, - optional_params: dict[str, object], - litellm_params: dict[str, object], - custom_llm_provider: str | None, - ) -> None: ... - - def pre_call( - self, - *, - input: str, - api_key: str | None, - additional_args: dict[str, object], - ) -> None: ... - - def _as_ocr(value: object) -> RustOcr | None: return cast(RustOcr, value) if callable(value) else None @@ -127,204 +79,6 @@ def load_rust_aocr() -> RustAocr | None: return _AOCR.load() -def provider(request: LiteLLMOcrRequest) -> str | None: - if request.custom_llm_provider is not None: - return request.custom_llm_provider - prefix: Final = request.model.partition("/")[0] - if prefix in _RUST_OCR_PROVIDERS: - return prefix - if request.model.startswith("mistral-ocr"): - return "mistral" - return None - - -def supported(request: LiteLLMOcrRequest) -> bool: - request_provider: Final = provider(request) - if request_provider not in _RUST_OCR_PROVIDERS: - return False - if request_provider == "azure_ai": - return ( - not is_azure_cohere_parse_model(request.model) - and not callable(request.kwargs.get("azure_ad_token_provider")) - and request.kwargs.get("azure_username") is None - and request.kwargs.get("azure_password") is None - ) - return True - - -def _optional_params(request: LiteLLMOcrRequest, resolve_secret: Callable[[str], str | None]) -> Mapping[str, object]: - optional_params: Final = MappingProxyType( - { - name: value - for name, value in request.kwargs.items() - if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) - and name not in ("litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request") - } - ) - request_provider: Final = provider(request) - if request_provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: - return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) - if request_provider != "vertex_ai": - return optional_params - project: Final = ( - request.kwargs.get("vertex_project") - or request.kwargs.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - location: Final = ( - request.kwargs.get("vertex_location") - or request.kwargs.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - credentials: Final = ( - request.kwargs.get("vertex_credentials") - or request.kwargs.get("vertex_ai_credentials") - or resolve_secret("VERTEXAI_CREDENTIALS") - ) - vertex_params: Final = MappingProxyType( - { - name: value - for name, value in ( - ("vertex_project", project), - ("vertex_location", location), - ("vertex_credentials", credentials), - ) - if value is not None - } - ) - return MappingProxyType({**optional_params, **vertex_params}) - - -def _input_sources(request: LiteLLMOcrRequest, optional_params: Mapping[str, object]) -> Mapping[str, str]: - proxy_request_value: Final = request.kwargs.get("proxy_server_request") - if not isinstance(proxy_request_value, Mapping): - return MappingProxyType({}) - proxy_request: Final = cast( # cast-ok: runtime Mapping check narrows metadata with unknown key and value types - Mapping[object, object], proxy_request_value - ) - credential_fields_value: Final = proxy_request.get("credential_fields", ()) - credential_fields: Final = ( - frozenset(name for name in credential_fields_value if isinstance(name, str)) - if isinstance(credential_fields_value, (list, tuple, set, frozenset)) - else frozenset() - ) - request_fields_value: Final = proxy_request.get("body_fields") - request_fields: Sequence[object] - if isinstance(request_fields_value, Sequence) and not isinstance(request_fields_value, (str, bytes)): - request_fields = cast( # cast-ok: runtime Sequence check excludes scalar strings and bytes - Sequence[object], request_fields_value - ) - else: - body_value: Final = proxy_request.get("body") - request_fields = ( - tuple(cast(Mapping[object, object], body_value)) # cast-ok: runtime Mapping check establishes iterable keys - if isinstance(body_value, Mapping) - else () - ) - names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) - request_sources: Final = MappingProxyType( - {name: "request" for name in names if name in request_fields or name in credential_fields} - ) - if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: - return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) - return request_sources - - -def _marshal( - request: LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], - convert_file_document: Callable[[dict[str, object]], dict[str, str]], -) -> LiteLLMOcrRequest: - if not isinstance(request.document, dict): - raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") - document: Final = ( - convert_file_document(request.document) if request.document.get("type") == "file" else request.document - ) - request_provider: Final = provider(request) - api_key: Final = ( - request.api_key or resolve_secret("MISTRAL_API_KEY") if request_provider == "mistral" else request.api_key - ) - optional_params: Final = _optional_params(request, resolve_secret) - input_sources: Final = _input_sources(request, optional_params) - logged_optional_params: Final = MappingProxyType( - {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} - ) - logged_kwargs: Final = MappingProxyType( - { - name: "****" if name in _RUST_OCR_SECRET_FIELDS else value - for name, value in request.kwargs.items() - if name != "proxy_server_request" - } - ) - logging_obj: Final = cast( # cast-ok: client decorator injects the logging object through untyped kwargs - _OCRLogging, request.kwargs["litellm_logging_obj"] - ) - logging_obj.update_from_kwargs( - kwargs=dict(logged_kwargs), # mutable-ok: legacy logging mutates its kwargs copy - model=request.model, - optional_params=dict(logged_optional_params), # mutable-ok: legacy logging requires concrete dict params - litellm_params={ # mutable-ok: legacy logging requires a concrete params dict - "litellm_call_id": request.kwargs.get("litellm_call_id"), - "api_base": request.api_base, - }, - custom_llm_provider=request_provider, - ) - logging_obj.pre_call( - input="OCR document processing", - api_key=api_key, - additional_args={ # mutable-ok: pre_call mutates the additional_args dict - "complete_input_dict": { # mutable-ok: callbacks consume a JSON-serializable request dict - "model": request.model, - "document": document, - **logged_optional_params, - }, - "api_base": request.api_base or "", - "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict - }, - ) - return LiteLLMOcrRequest( - model=request.model, - document=document, - api_key=api_key, - api_base=request.api_base, - timeout=request.timeout if request.timeout is not None else request_timeout, - custom_llm_provider=request.custom_llm_provider, - extra_headers=request.extra_headers, - kwargs=optional_params, - input_sources=input_sources, - ) - - -def _map_error(error: Exception, request: LiteLLMOcrRequest) -> Exception: - exception_types: Final = native_exception_types() - if exception_types is None or not isinstance(error, exception_types[1]): - return error - request_provider: Final = provider(request) - if request_provider is None: - return error - provider_config: Final = ProviderConfigManager.get_provider_ocr_config( - model=request.model.removeprefix(f"{request_provider}/"), provider=litellm.LlmProviders(request_provider) - ) - if provider_config is None: - return error - error_args: Final = cast( # cast-ok: BaseException.args exposes Any while native errors carry scalar args - tuple[object, ...], error.args - ) - status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 - message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) - error_factory: Final = cast( # cast-ok: legacy provider error factories have untyped callable parameters - Callable[..., Exception], provider_config.get_error_class - ) - return error_factory( - error_message=message, - status_code=status or 500, - headers={}, # mutable-ok: provider error factories require a concrete headers dict - ) - - def _response(response: Mapping[str, object]) -> OCRResponse: provider_native_response: Final = response.get(PROVIDER_NATIVE_RESPONSE_KEY) normalized: Final = OCRResponse.model_validate( @@ -335,56 +89,6 @@ def _response(response: Mapping[str, object]) -> OCRResponse: return normalized -def run( - request: LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], - convert_file_document: Callable[[dict[str, object]], dict[str, str]], -) -> OCRResponse | None: - if load_rust_ocr() is None: - return None - marshalled: Final = _marshal(request, resolve_secret, convert_file_document) - try: - response: Final = ocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=marshalled.input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_error(error, request) from error - return _response(response) if response is not None else None - - -async def arun( - request: LiteLLMOcrRequest, - resolve_secret: Callable[[str], str | None], - convert_file_document: Callable[[dict[str, object]], dict[str, str]], -) -> OCRResponse | None: - if load_rust_aocr() is None: - return None - marshalled: Final = _marshal(request, resolve_secret, convert_file_document) - try: - response: Final = await aocr( - model=marshalled.model, - document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict - api_key=marshalled.api_key, - api_base=marshalled.api_base, - custom_llm_provider=marshalled.custom_llm_provider, - extra_headers=marshalled.extra_headers, - optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict - input_sources=marshalled.input_sources, - timeout=marshalled.timeout, - ) - except Exception as error: - raise _map_error(error, request) from error - return _response(response) if response is not None else None - - def ocr( *, model: str, diff --git a/litellm/rust_bridge/ocr_lifecycle.py b/litellm/rust_bridge/ocr_lifecycle.py new file mode 100644 index 00000000000..5ca584e1c11 --- /dev/null +++ b/litellm/rust_bridge/ocr_lifecycle.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Mapping, Sequence +from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.bindings import NativeBinding +from litellm.rust_bridge.ocr import LiteLLMOcrRequest + + +class NativeOcrLifecycle(Protocol): + def __call__( + self, + request: LiteLLMOcrRequest, + args: Sequence[object], + kwargs: Mapping[str, object], + asynchronous: bool, + ) -> OCRResponse | Awaitable[OCRResponse]: ... + + +class ExceptionMapper(Protocol): + def __call__( + self, + *, + model: str, + custom_llm_provider: str | None, + original_exception: Exception, + completion_kwargs: dict[str, object], + extra_kwargs: dict[str, object], + ) -> Exception: ... + + +def _binding(value: object) -> NativeOcrLifecycle | None: + if not callable(value): + return None + return cast("NativeOcrLifecycle", value) # cast-ok: callable validated at the native binding boundary + + +NATIVE_OCR_LIFECYCLE: Final = NativeBinding("_ocr_lifecycle", validate=_binding) + + +def select(request: LiteLLMOcrRequest) -> NativeOcrLifecycle | None: + if request.kwargs.get("aocr"): + return None + return NATIVE_OCR_LIFECYCLE.load() + + +def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]: + return request.kwargs + + +def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception: + mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper + ExceptionMapper, litellm.exception_type + ) + try: + return mapper( + model=request.model.removeprefix(f"{request_provider}/"), + custom_llm_provider=request_provider, + original_exception=error, + completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs + extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs + ) + except Exception as public_error: + public_error.__context__ = error + return public_error diff --git a/scripts/benchmark_ocr_callbacks.py b/scripts/benchmark_ocr_callbacks.py new file mode 100644 index 00000000000..5db182c3d0e --- /dev/null +++ b/scripts/benchmark_ocr_callbacks.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +"""Measure serial sync/async OCR latency through a loopback HTTP provider + +Run each callback mode in a fresh process against an installed release wheel: +python -I scripts/benchmark_ocr_callbacks.py --callbacks none --label before \ + --expected-transport rust --iterations 200 --warmup 20 --output before-none.json +Repeat with --callbacks noop and with the candidate wheel in a separate venv +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import importlib.metadata +import json +import statistics +import sys +import threading +import time +from collections.abc import Sequence +from dataclasses import asdict, dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final, cast + +SIZES: Final = ( + 1024, + 4 * 1024, + 16 * 1024, + 64 * 1024, + 256 * 1024, + 1024 * 1024, +) +MODEL: Final = "mistral/mistral-ocr-latest" +EXPECTED_MARKDOWN: Final = "mock remote OCR response" +RESPONSE: Final = json.dumps( + { + "pages": [{"index": 0, "markdown": EXPECTED_MARKDOWN, "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, + }, + separators=(",", ":"), +).encode() + + +class Server(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), Handler) + self.user_agents: set[str] = set() + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + server: Final = cast(Server, self.server) + server.user_agents.add(self.headers.get("User-Agent", "")) + length: Final = int(self.headers["Content-Length"]) + body: Final = self.rfile.read(length) + request: Final = json.loads(body) + if self.path != "/v1/ocr" or request.get("model") != "mistral-ocr-latest": + self.send_error(400) + return + document: Final = request.get("document", {}) + if not isinstance(document, dict) or not str(document.get("document_url", "")).startswith( + "data:application/pdf;base64," + ): + self.send_error(400) + return + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(RESPONSE))) + self.end_headers() + self.wfile.write(RESPONSE) + + def log_message(self, format: str, *args: object) -> None: + return + + +@dataclass(frozen=True, slots=True) +class Result: + label: str + mode: str + size: int + iterations: int + median_ms: float + mean_ms: float + p95_ms: float + requests_per_second: float + + +def document(size: int) -> dict[str, str]: + payload: Final = b"%PDF-1.4\n" + b"x" * max(0, size - 9) + encoded: Final = base64.b64encode(payload[:size]).decode("ascii") + return {"type": "document_url", "document_url": f"data:application/pdf;base64,{encoded}"} + + +def percentile(values: Sequence[float], quantile: float) -> float: + ordered: Final = sorted(values) + index: Final = min(len(ordered) - 1, round((len(ordered) - 1) * quantile)) + return ordered[index] + + +def verify(response: object) -> None: + pages: Final = getattr(response, "pages", ()) + if len(pages) != 1 or getattr(pages[0], "markdown", None) != EXPECTED_MARKDOWN: + raise RuntimeError(f"unexpected OCR response: {response!r}") + + +def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Result: + median: Final = statistics.median(samples) + return Result( + label=label, + mode=mode, + size=size, + iterations=len(samples), + median_ms=median * 1000, + mean_ms=statistics.fmean(samples) * 1000, + p95_ms=percentile(samples, 0.95) * 1000, + requests_per_second=1 / median, + ) + + +def sync_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]: + samples: list[float] = [] + for _ in range(count): + started: Final = time.perf_counter() + response: Final = litellm.ocr( + model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30 + ) + samples.append(time.perf_counter() - started) + verify(response) + return tuple(samples) + + +async def async_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]: + samples: list[float] = [] + for _ in range(count): + started: Final = time.perf_counter() + response: Final = await litellm.aocr( + model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30 + ) + samples.append(time.perf_counter() - started) + verify(response) + return tuple(samples) + + +async def main() -> int: + parser: Final = argparse.ArgumentParser(description="E2E OCR benchmark against a local remote-style HTTP server") + parser.add_argument("--callbacks", choices=("none", "noop"), required=True) + parser.add_argument("--label", required=True) + parser.add_argument("--expected-transport", choices=("python", "rust"), required=True) + parser.add_argument("--iterations", type=int, default=30) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--sizes", type=int, nargs="+", default=SIZES) + parser.add_argument("--output", type=Path, required=True) + args: Final = parser.parse_args() + + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class NoopCallback(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.pre_calls = 0 + self.sync_successes = 0 + self.async_successes = 0 + + def log_pre_api_call(self, model, messages, kwargs): + self.pre_calls += 1 + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.sync_successes += 1 + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.async_successes += 1 + + registry_names: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ) + if any(getattr(litellm, name) for name in registry_names): + raise RuntimeError("benchmark requires initially empty callback registrations") + callback: Final = NoopCallback() + if args.callbacks == "noop": + litellm.callbacks.append(callback) + + rust_toggle: Final = getattr(litellm, "rust", None) + if callable(rust_toggle): + rust_toggle(False) + package: Final = Path(litellm.__file__).resolve() + version: Final = importlib.metadata.version("litellm") + native_path: str | None = None + native_sha256: str | None = None + try: + from litellm.rust_bridge import _native + + native: Final = Path(_native.__file__).resolve() + native_path = str(native) + native_sha256 = hashlib.file_digest(native.open("rb"), "sha256").hexdigest() + except ImportError: + pass + + server: Final = Server() + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + results: list[Result] = [] + try: + for size in args.sizes: + request_document: Final = document(size) + sync_samples(litellm, url, request_document, args.warmup) + sync_result: Final = summarize( + args.label, "sync", size, sync_samples(litellm, url, request_document, args.iterations) + ) + results.append(sync_result) + await async_samples(litellm, url, request_document, args.warmup) + async_result: Final = summarize( + args.label, + "async", + size, + await async_samples(litellm, url, request_document, args.iterations), + ) + results.append(async_result) + sys.stdout.write(json.dumps(asdict(sync_result)) + "\n") + sys.stdout.write(json.dumps(asdict(async_result)) + "\n") + sys.stdout.flush() + finally: + server.shutdown() + server.server_close() + thread.join() + + from litellm.litellm_core_utils.litellm_logging import executor + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.to_thread(executor.shutdown, wait=True) + per_mode: Final = len(args.sizes) * (args.iterations + args.warmup) + if args.callbacks == "noop": + if (callback.pre_calls, callback.sync_successes, callback.async_successes) != ( + 2 * per_mode, + per_mode, + per_mode, + ): + raise RuntimeError(f"callback delivery mismatch: {vars(callback)}") + elif any(getattr(litellm, name) for name in registry_names): + raise RuntimeError("callback registrations appeared in the no-callback case") + await GLOBAL_LOGGING_WORKER.stop() + + user_agents: Final = tuple(sorted(server.user_agents)) + python_transport: Final = any( + value.startswith("python-httpx") or value.startswith("litellm/") for value in user_agents + ) + if (args.expected_transport == "python") != python_transport: + raise RuntimeError(f"unexpected transport for {args.label}: user_agents={user_agents}") + artifact: Final = { + "label": args.label, + "callbacks": args.callbacks, + "python": sys.executable, + "callback_counts": { + "pre": callback.pre_calls, + "sync_success": callback.sync_successes, + "async_success": callback.async_successes, + }, + "version": version, + "package": str(package), + "native": native_path, + "native_sha256": native_sha256, + "user_agents": user_agents, + "results": tuple(asdict(result) for result in results), + } + args.output.write_text(json.dumps(artifact, indent=2) + "\n") + sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package", "user_agents")}) + "\n") + sys.stdout.write(f"results={args.output}\n") + sys.stdout.flush() + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 2bd3a50f39f..860e872dd44 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -1,7 +1,8 @@ from __future__ import annotations -import asyncio -from collections.abc import Awaitable, Callable +import json +import subprocess +from functools import cache from pathlib import Path from typing import Final, Protocol, cast @@ -28,12 +29,12 @@ class _GatewayClient(Protocol): def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: - import litellm from fastapi.testclient import TestClient + import litellm + from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.endpoints import user_api_key_auth - from litellm.proxy import proxy_server provider_model: Final = cast(str, fixture.kwargs["provider_model"]) model_alias: Final = cast(str, fixture.kwargs["model_alias"]) @@ -76,24 +77,24 @@ def _collect_python(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: - from litellm.rust_bridge import get_native_bridge - - bridge: Final[object | None] = get_native_bridge() - trace: Final[object | None] = getattr(bridge, "_trace", None) if bridge is not None else None - gateway_messages: Final[object | None] = getattr(trace, "gateway_messages", None) - if gateway_messages is None or not callable(gateway_messages): - raise RuntimeError("native Rust trace bridge does not expose gateway_messages") - invoke_gateway: Final = cast(Callable[[str, str, str, object], Awaitable[object]], gateway_messages) - - async def invoke() -> object: - return await invoke_gateway( - cast(str, fixture.kwargs["model_alias"]), - cast(str, fixture.kwargs["provider_model"]), - cast(str, fixture.kwargs["api_base"]), - fixture.kwargs["body"], - ) - - result: Final = asyncio.run(invoke()) + payload: Final = json.dumps( + { + "model_alias": fixture.kwargs["model_alias"], + "provider_model": fixture.kwargs["provider_model"], + "api_base": fixture.kwargs["api_base"], + "body": fixture.kwargs["body"], + } + ) + completed: Final = subprocess.run( + (_gateway_trace_binary(),), + input=payload, + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError(f"Rust gateway trace failed: {completed.stderr.strip()}") + result: Final = json.loads(completed.stdout) payload: Final = TraceResponsePayload.model_validate(result) response: Final = _GatewayResponsePayload.model_validate(payload.response) if response.status != 200: @@ -101,6 +102,34 @@ def _collect_rust(fixture: RouteFixture) -> tuple[FunctionTraceEvent, ...]: return native_trace_events(payload) +@cache +def _gateway_trace_binary() -> Path: + repo_root: Final = next(parent for parent in Path(__file__).resolve().parents if (parent / "litellm-rust").is_dir()) + rust_root: Final = repo_root / "litellm-rust" + completed: Final = subprocess.run( + ( + "cargo", + "build", + "--quiet", + "--package", + "litellm-ai-gateway", + "--features", + "trace-parity", + "--bin", + "trace-parity-gateway", + "--target-dir", + rust_root / "target", + ), + cwd=rust_root, + capture_output=True, + text=True, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError(f"Rust gateway trace build failed: {completed.stderr.strip()}") + return rust_root / "target" / "debug" / "trace-parity-gateway" + + def _collect(scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: try: with replay_server() as provider: diff --git a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py b/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py index 3f98e9b6a2d..2f457fcb25b 100644 --- a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py +++ b/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py @@ -1,9 +1,5 @@ -import base64 -import json - import pytest -import litellm from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig from litellm.llms.azure_ai.ocr.common_utils import get_azure_ai_ocr_config from litellm.llms.azure_ai.ocr.document_intelligence.transformation import AzureDocumentIntelligenceOCRConfig @@ -12,27 +8,6 @@ from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig MODEL = "azure_ai/Cohere-parse-v5" API_BASE = "https://resource.services.ai.azure.com" PARSE_URL = f"{API_BASE}/providers/cohere/v2/parse" -IMAGE_URL = "https://example.com/receipt.png" -PNG_BYTES = base64.b64decode( - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" -) -PNG_DATA_URI = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}" - - -def _parse_response() -> dict: - return { - "id": "882bf973-9dfa-4d02-9d30-709247008efd", - "pages": [{"index": 0, "type": "markdown", "markdown": {"content": "# Receipt\n\nTotal Due: $4.00"}}], - "meta": {"api_version": {"version": "2"}, "billed_units": {"pages": 1}}, - } - - -@pytest.fixture() -def disable_aiohttp_transport(monkeypatch): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - litellm.in_memory_llm_clients_cache.flush_cache() - yield - litellm.in_memory_llm_clients_cache.flush_cache() @pytest.mark.parametrize( @@ -95,90 +70,3 @@ def test_validate_environment_requires_api_base(monkeypatch) -> None: with pytest.raises(ValueError, match="AZURE_AI_API_BASE"): AzureAICohereParseConfig().validate_environment(headers={}, model="Cohere-parse-v5", api_key="key") - - -@pytest.mark.asyncio -async def test_aocr_inlines_remote_image_and_posts_to_foundry(disable_aiohttp_transport, respx_mock): - respx_mock.get(IMAGE_URL).respond(content=PNG_BYTES, headers={"Content-Type": "image/png"}) - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - response = await litellm.aocr( - model=MODEL, - document={"type": "image_url", "image_url": IMAGE_URL}, - api_base=API_BASE, - api_key="azure-key", - ) - - request = route.calls.last.request - assert request.headers["Authorization"] == "Bearer azure-key" - assert json.loads(request.content) == { - "model": "Cohere-parse-v5", - "document": {"type": "image_url", "image_url": PNG_DATA_URI}, - "output_format": "markdown", - } - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - assert response.usage_info.pages_processed == 1 - - -@pytest.mark.asyncio -async def test_aocr_passes_data_uri_through_without_fetching(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - await litellm.aocr( - model=MODEL, - document={"type": "image_url", "image_url": PNG_DATA_URI}, - api_base=API_BASE, - api_key="azure-key", - output_format="blocks", - ) - - body = json.loads(route.calls.last.request.content) - assert body["document"]["image_url"] == PNG_DATA_URI - assert body["output_format"] == "blocks" - - -def test_ocr_sync_inlines_remote_image(respx_mock): - respx_mock.get(IMAGE_URL).respond(content=PNG_BYTES, headers={"Content-Type": "image/png"}) - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - response = litellm.ocr( - model=MODEL, - document={"type": "image_url", "image_url": IMAGE_URL}, - api_base=API_BASE, - api_key="azure-key", - ) - - assert json.loads(route.calls.last.request.content)["document"]["image_url"] == PNG_DATA_URI - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - - -@pytest.mark.asyncio -async def test_aocr_rejects_pdf_before_calling_foundry(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - with pytest.raises(litellm.BadRequestError, match="only accepts `image_url` documents") as exc_info: - await litellm.aocr( - model=MODEL, - document={"type": "document_url", "document_url": "https://example.com/doc.pdf"}, - api_base=API_BASE, - api_key="azure-key", - ) - - assert exc_info.value.llm_provider == "azure_ai" - assert not route.called - - -@pytest.mark.asyncio -async def test_ahealth_check_ocr_sends_an_image_to_the_foundry_cohere_parse_deployment( - disable_aiohttp_transport, respx_mock -): - route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) - - result = await litellm.ahealth_check( - model_params={"model": MODEL, "api_base": API_BASE, "api_key": "test-key"}, mode="ocr" - ) - - document = json.loads(route.calls.last.request.content)["document"] - assert document["type"] == "image_url" - assert document["image_url"].startswith("data:image/png;base64,") - assert "error" not in result diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py index ef4c78553f1..be0dfb5724e 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -1,4 +1,5 @@ from unittest.mock import MagicMock +from typing import Final import httpx import pytest @@ -371,3 +372,35 @@ def test_validate_environment_falls_back_to_entra_token(monkeypatch): assert headers["Authorization"] == "Bearer entra-token" assert "Ocp-Apim-Subscription-Key" not in headers + + +@pytest.mark.parametrize( + ("request_headers", "expected_poll_headers"), + ( + ( + {"Ocp-Apim-Subscription-Key": "subscription-key"}, + {"Ocp-Apim-Subscription-Key": "subscription-key"}, + ), + ( + {"Authorization": "Bearer entra-token"}, + {"Authorization": "Bearer entra-token"}, + ), + ), +) +def test_get_polling_target_preserves_request_authentication( + request_headers: dict[str, str], expected_poll_headers: dict[str, str] +) -> None: + response: Final = httpx.Response( + status_code=202, + headers={"Operation-Location": "https://example.cognitiveservices.azure.com/operations/123"}, + request=httpx.Request( + "POST", + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze", + headers=request_headers, + ), + ) + + operation_url, poll_headers = AzureDocumentIntelligenceOCRConfig()._get_polling_target(response) + + assert operation_url == "https://example.cognitiveservices.azure.com/operations/123" + assert poll_headers == expected_poll_headers diff --git a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py index cb9af56f5e0..1f120be6ffa 100644 --- a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py +++ b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py @@ -1,8 +1,11 @@ -import json +from typing import Final +from unittest.mock import Mock +import httpx import pytest import litellm +from litellm.llms.cohere.ocr.transformation import CohereParseConfig PARSE_URL = "https://api.cohere.com/v2/parse" MODEL = "cohere/parse-v5.0" @@ -57,173 +60,38 @@ def _blocks_response() -> dict: } -@pytest.fixture() -def disable_aiohttp_transport(monkeypatch): - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - litellm.in_memory_llm_clients_cache.flush_cache() - yield - litellm.in_memory_llm_clients_cache.flush_cache() - - -@pytest.mark.asyncio -async def test_aocr_sends_markdown_parse_request_and_normalizes_pages(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - request = route.calls.last.request - assert request.headers["Authorization"] == "Bearer test-key" - assert json.loads(request.content) == { - "model": "parse-v5.0", - "document": IMAGE_DOCUMENT, - "output_format": "markdown", - } - assert response.object == "ocr" - assert [page.index for page in response.pages] == [0, 1] - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - assert response.pages[1].markdown == "Page two" - assert response.pages[1].images is None - image = response.pages[0].images[0] - assert image.bbox == BOUNDING_BOX - assert image.model_extra["description"] == "A parking receipt" - assert image.model_extra["bounding_box_normalized"]["bottom_right_x"] == 1 - assert response.usage_info.pages_processed == 2 - assert response.get_provider_native_response() is None - - -@pytest.mark.asyncio -async def test_aocr_usage_prefers_billed_units_over_page_count(disable_aiohttp_transport, respx_mock): - respx_mock.post(PARSE_URL).respond(json=_markdown_response(billed_pages=3)) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - assert response.usage_info.pages_processed == 3 - - -@pytest.mark.asyncio -async def test_aocr_usage_falls_back_to_page_count_without_meta(disable_aiohttp_transport, respx_mock): - respx_mock.post(PARSE_URL).respond(json=_markdown_response(billed_pages=None)) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - assert response.usage_info.pages_processed == 2 - - -@pytest.mark.asyncio -async def test_aocr_blocks_output_format_forwards_param_and_keeps_blocks(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_blocks_response()) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", output_format="blocks") - - assert json.loads(route.calls.last.request.content)["output_format"] == "blocks" - assert response.pages[0].markdown == "" - assert response.pages[0].model_extra["blocks"] == [{"type": "text", "text": "Total Due: $4.00"}] - assert response.usage_info.pages_processed == 1 - - -@pytest.mark.asyncio -async def test_aocr_native_format_carries_provider_payload(disable_aiohttp_transport, respx_mock): - payload = _markdown_response() - route = respx_mock.post(PARSE_URL).respond(json=payload) - - response = await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", req_format="native") - - assert "req_format" not in json.loads(route.calls.last.request.content) - assert response.get_provider_native_response() == payload - assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" - - -@pytest.mark.asyncio -async def test_aocr_rejects_unknown_output_format_before_calling_provider(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - with pytest.raises(litellm.BadRequestError, match="Invalid `output_format`: 'html'") as exc_info: - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", output_format="html") - - assert exc_info.value.status_code == 400 - assert not route.called - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "document", - [ - {"type": "document_url", "document_url": "https://example.com/doc.pdf"}, - {"type": "image_url", "image_url": "data:application/pdf;base64,JVBERi0="}, - {"type": "image_url", "image_url": ""}, - ], -) -async def test_aocr_rejects_non_image_documents_before_calling_provider( - disable_aiohttp_transport, respx_mock, document -): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - with pytest.raises(litellm.BadRequestError, match="only accepts `image_url` documents") as exc_info: - await litellm.aocr(model=MODEL, document=document, api_key="test-key") - - assert exc_info.value.status_code == 400 - assert not route.called - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "api_base, expected_url", - [ - ("https://gateway.example.com", "https://gateway.example.com/v2/parse"), - ("https://gateway.example.com/cohere/", "https://gateway.example.com/cohere/v2/parse"), - ("https://gateway.example.com/v2", "https://gateway.example.com/v2/parse"), - ("https://gateway.example.com/v2/parse", "https://gateway.example.com/v2/parse"), - ], -) -async def test_aocr_posts_to_api_base_variants(disable_aiohttp_transport, respx_mock, api_base, expected_url): - route = respx_mock.post(expected_url).respond(json=_markdown_response()) - - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key", api_base=api_base) - - assert route.called - - -@pytest.mark.asyncio -async def test_aocr_surfaces_provider_error_with_its_status_and_message(disable_aiohttp_transport, respx_mock): - respx_mock.post(PARSE_URL).respond( - status_code=400, json={"id": "83b0d95e", "message": "output_format must be `blocks` or `markdown`"} +@pytest.mark.parametrize("output_format", ["markdown", "blocks"]) +def test_transform_cohere_request_filters_options(output_format: str) -> None: + config: Final = CohereParseConfig() + params: Final = config.map_ocr_params( + {"output_format": output_format, "req_format": "native", "unknown": True}, {}, "parse-v5.0" ) - - with pytest.raises(litellm.BadRequestError, match="output_format must be") as exc_info: - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT, api_key="test-key") - - assert exc_info.value.status_code == 400 + request: Final = config.transform_ocr_request("parse-v5.0", IMAGE_DOCUMENT, params, {}) + assert request.data == {"model": "parse-v5.0", "document": IMAGE_DOCUMENT, "output_format": output_format} -@pytest.mark.asyncio -async def test_aocr_reads_api_key_from_environment(disable_aiohttp_transport, respx_mock, monkeypatch): - monkeypatch.setenv("COHERE_API_KEY", "env-key") - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT) - - assert route.calls.last.request.headers["Authorization"] == "Bearer env-key" +@pytest.mark.parametrize("native", [False, True]) +def test_transform_cohere_response_keeps_images_and_native_payload(native: bool) -> None: + payload: Final = _markdown_response(3) + response: Final = CohereParseConfig().transform_ocr_response( + "parse-v5.0", httpx.Response(200, json=payload), Mock(), {"req_format": "native" if native else "litellm"} + ) + assert response.pages[0].markdown == "# Receipt\n\nTotal Due: $4.00" + assert response.pages[0].images[0].bbox == BOUNDING_BOX + assert response.pages[0].images[0].model_extra["description"] == "A parking receipt" + assert response.pages[1].images is None + assert response.usage_info.pages_processed == 3 + assert response.get_provider_native_response() == (payload if native else None) -@pytest.mark.asyncio -async def test_aocr_without_api_key_names_the_env_var(disable_aiohttp_transport, respx_mock, monkeypatch): - monkeypatch.delenv("COHERE_API_KEY", raising=False) - monkeypatch.setattr(litellm, "cohere_key", None) - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - with pytest.raises(Exception, match="Missing COHERE_API_KEY"): - await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT) - - assert not route.called +def test_transform_cohere_blocks() -> None: + response: Final = CohereParseConfig().transform_ocr_response( + "parse-v5.0", httpx.Response(200, json=_blocks_response()), Mock() + ) + assert response.pages[0].model_extra["blocks"] == [{"type": "text", "text": "Total Due: $4.00"}] + assert response.pages[0].markdown == "" -@pytest.mark.asyncio -async def test_ahealth_check_ocr_sends_an_image_cohere_parse_accepts(disable_aiohttp_transport, respx_mock): - route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) - - result = await litellm.ahealth_check(model_params={"model": MODEL, "api_key": "test-key"}, mode="ocr") - - document = json.loads(route.calls.last.request.content)["document"] - assert document["type"] == "image_url" - assert document["image_url"].startswith("data:image/png;base64,") - assert "error" not in result +def test_transform_cohere_rejects_unsupported_output_format() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="output_format"): + CohereParseConfig().map_ocr_params({"output_format": "html"}, {}, "parse-v5.0") diff --git a/tests/test_litellm/llms/reducto/conftest.py b/tests/test_litellm/llms/reducto/conftest.py new file mode 100644 index 00000000000..4ff3ab43006 --- /dev/null +++ b/tests/test_litellm/llms/reducto/conftest.py @@ -0,0 +1,11 @@ +from collections.abc import Generator + +import pytest + +from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service + + +@pytest.fixture +def reducto_server() -> Generator[RecordingServer]: + with recording_service() as server: + yield server diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/test_litellm/llms/reducto/test_parse_legacy.py index db19460baa3..252369cbd3d 100644 --- a/tests/test_litellm/llms/reducto/test_parse_legacy.py +++ b/tests/test_litellm/llms/reducto/test_parse_legacy.py @@ -1,7 +1,7 @@ -import json +import pytest import litellm -import pytest +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec @pytest.fixture() @@ -17,24 +17,28 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_legacy_wraps_enhance_under_options( - disable_aiohttp_transport, respx_mock -): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://legacy.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json={ - "usage": {"num_pages": 1, "credits": 1}, - "result": { - "chunks": [ - { - "content": "Legacy parse", - "blocks": [{"content": "Legacy parse", "bbox": {"page": 1}}], - } - ] - }, - } +async def test_parse_legacy_wraps_enhance_under_options(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 2 + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://legacy.pdf"})) + reducto_server.enqueue( + ResponseSpec( + body={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Legacy parse", + "blocks": [ + { + "content": "Legacy parse", + "bbox": {"page": 1}, + } + ], + } + ] + }, + } + ) ) response = await litellm.aocr( @@ -45,13 +49,15 @@ async def test_parse_legacy_wraps_enhance_under_options( "mime_type": "application/pdf", }, api_key="legacy-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, enhance={"agentic": [{"type": "table"}]}, ) - assert upload_route.called - assert parse_route.called - request_body = json.loads(parse_route.calls[0].request.read()) + upload_request, parse_request = reducto_server.requests + assert upload_request.path == "/upload" + assert parse_request.path == "/parse" + assert isinstance(parse_request.body, dict) + request_body = parse_request.body assert request_body == { "document_url": "reducto://legacy.pdf", "options": {"enhance": {"agentic": [{"type": "table"}]}}, diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/test_litellm/llms/reducto/test_parse_v3.py index 1d0c826ef8b..0ebc0d926c4 100644 --- a/tests/test_litellm/llms/reducto/test_parse_v3.py +++ b/tests/test_litellm/llms/reducto/test_parse_v3.py @@ -1,8 +1,7 @@ -import json - import pytest import litellm +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec def _reducto_parse_response() -> dict: @@ -69,11 +68,11 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transport, respx_mock): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://uploaded.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) +async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 2 + provider_response = _reducto_parse_response() + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded.pdf"})) + reducto_server.enqueue(ResponseSpec(body=provider_response)) response = await litellm.aocr( model="reducto/parse-v3", @@ -83,25 +82,24 @@ async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transpo "mime_type": "application/pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, formatting={"table_output_format": "html"}, retrieval={"chunk_mode": "section"}, settings={"ocr_system": "standard"}, + req_format="native", ) - assert upload_route.called - assert parse_route.called - assert len(upload_route.calls) == 1 - assert len(parse_route.calls) == 1 - - upload_request = upload_route.calls[0].request + upload_request, parse_request = reducto_server.requests + assert upload_request.path == "/upload" + assert parse_request.path == "/parse" assert upload_request.headers["authorization"] == "Bearer test-key" assert "application/json" not in upload_request.headers["content-type"] - upload_body = upload_request.read() + upload_body = upload_request.raw_body assert b'filename="document"' in upload_body assert b"application/pdf" in upload_body - parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert isinstance(parse_request.body, dict) + parse_request_body = parse_request.body assert parse_request_body["input"] == "reducto://uploaded.pdf" assert parse_request_body["formatting"] == {"table_output_format": "html"} assert parse_request_body["retrieval"] == {"chunk_mode": "section"} @@ -116,15 +114,12 @@ async def test_parse_v3_file_upload_and_response_mapping(disable_aiohttp_transpo assert getattr(response.pages[0], "blocks")[0]["bbox"]["page"] == 1 assert response.pages[1].markdown == "Page 2 block A" assert response.pages[2].markdown == "Page 3 block A" - assert response._hidden_params["reducto_raw"]["usage"]["credits"] == 3 + assert response.get_provider_native_response() == provider_response @pytest.mark.asyncio -async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_transport, respx_mock): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://should-not-upload.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) +async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.enqueue(ResponseSpec(body=_reducto_parse_response())) response = await litellm.aocr( model="reducto/parse-v3", @@ -133,13 +128,15 @@ async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_tran "document_url": "reducto://already-uploaded.pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, retrieval={"chunk_mode": "section"}, ) - assert not upload_route.called - assert parse_route.called - parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert len(reducto_server.requests) == 1 + parse_request = reducto_server.requests[0] + assert parse_request.path == "/parse" + assert isinstance(parse_request.body, dict) + parse_request_body = parse_request.body assert parse_request_body["input"] == "reducto://already-uploaded.pdf" assert parse_request_body["retrieval"]["chunk_mode"] == "section" assert response.pages[0].markdown.startswith("Page 1 block A") @@ -147,11 +144,9 @@ async def test_parse_v3_reducto_id_passthrough_skips_upload(disable_aiohttp_tran @pytest.mark.asyncio async def test_unknown_model_uses_current_protocol_without_local_rejection( - disable_aiohttp_transport, respx_mock + disable_aiohttp_transport, reducto_server: RecordingServer ): - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json=_reducto_parse_response() - ) + reducto_server.enqueue(ResponseSpec(body=_reducto_parse_response())) response = await litellm.aocr( model="reducto/future-parse-model", @@ -160,11 +155,9 @@ async def test_unknown_model_uses_current_protocol_without_local_rejection( "document_url": "reducto://already-uploaded.pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, ) - assert parse_route.called - assert json.loads(parse_route.calls[0].request.read()) == { - "input": "reducto://already-uploaded.pdf" - } + assert reducto_server.requests[0].path == "/parse" + assert reducto_server.requests[0].body == {"input": "reducto://already-uploaded.pdf"} assert response.model == "future-parse-model" diff --git a/tests/test_litellm/llms/reducto/test_upload.py b/tests/test_litellm/llms/reducto/test_upload.py index 4fae90436bb..adfc2663fb0 100644 --- a/tests/test_litellm/llms/reducto/test_upload.py +++ b/tests/test_litellm/llms/reducto/test_upload.py @@ -1,16 +1,16 @@ -import json import os from unittest.mock import AsyncMock, Mock import httpx -import litellm import pytest +import litellm from litellm.llms.reducto.common import ( extract_file_id_or_bytes, upload_bytes_async, upload_bytes_sync, ) +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec @pytest.fixture() @@ -28,7 +28,8 @@ def disable_aiohttp_transport(monkeypatch): @pytest.mark.asyncio -async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport): +async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 0 with pytest.raises(litellm.BadRequestError, match="upload the file first"): await litellm.aocr( model="reducto/parse-v3", @@ -37,29 +38,30 @@ async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport): "document_url": "https://example.com/document.pdf", }, api_key="test-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, ) @pytest.mark.asyncio async def test_parse_v3_image_data_uri_upload_uses_image_mime( - disable_aiohttp_transport, respx_mock + disable_aiohttp_transport, reducto_server: RecordingServer ): - upload_route = respx_mock.post("https://custom.reducto.test/upload").respond( - json={"file_id": "reducto://uploaded-image.png"} - ) - parse_route = respx_mock.post("https://custom.reducto.test/parse").respond( - json={ - "usage": {"num_pages": 1, "credits": 1}, - "result": { - "chunks": [ - { - "content": "Image OCR", - "blocks": [{"content": "Image OCR", "bbox": {"page": 1}}], - } - ] - }, - } + reducto_server.expected_requests = 2 + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded-image.png"})) + reducto_server.enqueue( + ResponseSpec( + body={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Image OCR", + "blocks": [{"content": "Image OCR", "bbox": {"page": 1}}], + } + ] + }, + } + ) ) response = await litellm.aocr( @@ -70,41 +72,43 @@ async def test_parse_v3_image_data_uri_upload_uses_image_mime( "mime_type": "image/png", }, api_key="programmatic-key", - api_base="https://custom.reducto.test/", + api_base=f"{reducto_server.base_url}/", ) - assert upload_route.called - assert parse_route.called - upload_request = upload_route.calls[0].request + upload_request, parse_request = reducto_server.requests + assert upload_request.path == "/upload" + assert parse_request.path == "/parse" assert upload_request.headers["authorization"] == "Bearer programmatic-key" - assert b"image/png" in upload_request.read() + assert b"image/png" in upload_request.raw_body - parse_request_body = json.loads(parse_route.calls[0].request.read()) - assert parse_request_body["input"] == "reducto://uploaded-image.png" + assert isinstance(parse_request.body, dict) + assert parse_request.body["input"] == "reducto://uploaded-image.png" assert response.pages[0].markdown == "Image OCR" @pytest.mark.asyncio -async def test_parse_v3_uses_programmatic_api_key_over_env( - disable_aiohttp_transport, respx_mock -): - upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( - json={"file_id": "reducto://uploaded.pdf"} - ) - parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( - json={ - "usage": {"num_pages": 1, "credits": 1}, - "result": { - "chunks": [ - { - "content": "Programmatic auth", - "blocks": [ - {"content": "Programmatic auth", "bbox": {"page": 1}} - ], - } - ] - }, - } +async def test_parse_v3_uses_programmatic_api_key_over_env(disable_aiohttp_transport, reducto_server: RecordingServer): + reducto_server.expected_requests = 2 + reducto_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded.pdf"})) + reducto_server.enqueue( + ResponseSpec( + body={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Programmatic auth", + "blocks": [ + { + "content": "Programmatic auth", + "bbox": {"page": 1}, + } + ], + } + ] + }, + } + ) ) await litellm.aocr( @@ -115,11 +119,11 @@ async def test_parse_v3_uses_programmatic_api_key_over_env( "mime_type": "application/pdf", }, api_key="passed-key", - api_base="https://platform.reducto.ai", + api_base=reducto_server.base_url, ) - assert upload_route.calls[0].request.headers["authorization"] == "Bearer passed-key" - assert parse_route.calls[0].request.headers["authorization"] == "Bearer passed-key" + assert reducto_server.requests[0].headers["authorization"] == "Bearer passed-key" + assert reducto_server.requests[1].headers["authorization"] == "Bearer passed-key" def test_upload_bytes_sync_uses_shared_client(monkeypatch): diff --git a/tests/test_litellm/ocr/test_legacy.py b/tests/test_litellm/ocr/test_legacy.py new file mode 100644 index 00000000000..a30976f89af --- /dev/null +++ b/tests/test_litellm/ocr/test_legacy.py @@ -0,0 +1,200 @@ +import importlib +from collections.abc import AsyncGenerator +from datetime import datetime +from io import BytesIO +from typing import Final +from unittest.mock import Mock + +import httpx +import orjson +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.llms.custom_httpx import llm_http_handler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.ocr.legacy import _prepare_ocr_request +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE + + +@pytest.fixture +async def provider(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[Mock]: + configuration.reset_rust_configuration() + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust must not load"))) + handler: Final = Mock( + return_value=httpx.Response( + 200, + json={ + "pages": [{"index": 0, "markdown": "parsed document"}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, + }, + ) + ) + transport: Final = httpx.MockTransport(handler) + with httpx.Client(transport=transport) as sync_client: + async with httpx.AsyncClient(transport=transport) as async_client: + sync_handler: Final = HTTPHandler(client=sync_client) + async_handler: Final = AsyncHTTPHandler() + await async_handler.client.aclose() + async_handler.client = async_client + monkeypatch.setattr(llm_http_handler, "_get_httpx_client", lambda: sync_handler) + monkeypatch.setattr(llm_http_handler, "get_async_httpx_client", lambda llm_provider: async_handler) + yield handler + NATIVE_OCR_LIFECYCLE.reset() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["sync", "async", "sync_async"]) +@pytest.mark.parametrize("dispatch", ["disabled", "declined", "unavailable"]) +async def test_python_request_response_and_callbacks( + provider: Mock, monkeypatch: pytest.MonkeyPatch, mode: str, dispatch: str +) -> None: + class Declined(Exception): + pass + + if dispatch != "disabled": + monkeypatch.setenv("LITELLM_RUST", "1") + NATIVE_OCR_LIFECYCLE.override(Mock(side_effect=Declined()) if dispatch == "declined" else None) + main: Final = importlib.import_module("litellm.ocr.main") + monkeypatch.setattr(main, "native_exception_types", lambda: (Declined, RuntimeError)) + logger: Final = Mock(spec=CustomLogger) + monkeypatch.setattr(litellm, "input_callback", [logger]) + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "file", "file": BytesIO(b"pdf"), "mime_type": "application/pdf"}, + "api_key": "test-key", + "api_base": "https://ocr.test/v1", + "timeout": 7.0, + "pages": [0, 2], + "include_image_base64": True, + "extra_headers": {"x-test-header": "preserved"}, + } + + async def call() -> OCRResponse: + if mode == "async": + return await litellm.aocr(**arguments) + if mode == "sync_async": + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = Logging( + model=arguments["model"], + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.now(), + litellm_call_id="test-call", + function_id="test-function", + ) + return await litellm.ocr(**arguments, aocr=True, litellm_logging_obj=logging_obj) + return litellm.ocr(**arguments) + + response: Final = await call() + assert response.pages[0].markdown == "parsed document" + assert response.usage_info.pages_processed == 1 + assert provider.call_count == 1 + request: Final = provider.call_args.args[0] + assert str(request.url) == "https://ocr.test/v1/ocr" + assert request.headers["authorization"] == "Bearer test-key" + assert request.headers["x-test-header"] == "preserved" + assert request.extensions["timeout"] == {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0} + assert orjson.loads(request.content) == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,cGRm"}, + "pages": [0, 2], + "include_image_base64": True, + } + assert logger.log_pre_api_call.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_python_provider_errors_keep_public_exception(provider: Mock, asynchronous: bool) -> None: + provider.return_value = httpx.Response(429, json={"error": "rate limited"}) + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "https://example.com/file.pdf"}, + "api_key": "test-key", + "api_base": "https://ocr.test/v1", + "num_retries": 0, + } + + async def call() -> object: + if asynchronous: + return await litellm.aocr(**arguments) + return litellm.ocr(**arguments) + + with pytest.raises(litellm.RateLimitError) as error: + await call() + assert error.value.status_code == 429 + assert error.value.model == "mistral-ocr-latest" + assert error.value.llm_provider == "mistral" + assert provider.call_count == 1 + + +def test_document_intelligence_environment_key_is_not_replaced_by_generic_azure_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key") + monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", "document-key") + monkeypatch.setenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", "https://document.example.com") + prepared: Final = _prepare_ocr_request( + model="azure_ai/doc-intelligence/prebuilt-layout", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key=None, + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": Mock()}, + ) + + assert prepared.api_key is None + headers: Final = prepared.provider_config.validate_environment( + headers={}, + model=prepared.model, + api_key=prepared.api_key, + api_base=prepared.api_base, + litellm_params=prepared.litellm_params, + ) + assert headers["Ocp-Apim-Subscription-Key"] == "document-key" + + +def test_document_intelligence_explicit_connection_is_preserved(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://generic.example.com") + prepared: Final = _prepare_ocr_request( + model="azure_ai/doc-intelligence/prebuilt-layout", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key="explicit-key", + api_base="https://document.example.com", + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": Mock()}, + ) + + assert prepared.api_key == "explicit-key" + assert prepared.api_base == "https://document.example.com" + + +def test_generic_azure_connection_still_applies_to_foundry_ocr(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("AZURE_AI_API_KEY", "generic-key") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://generic.example.com") + prepared: Final = _prepare_ocr_request( + model="azure_ai/mistral-document-ai-2505", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key=None, + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": Mock()}, + ) + + assert prepared.api_key == "generic-key" + assert prepared.api_base == "https://generic.example.com" diff --git a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py deleted file mode 100644 index 460aff3e8d1..00000000000 --- a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py +++ /dev/null @@ -1,73 +0,0 @@ -""" -Regression tests for Azure Document Intelligence api_base ownership in OCR. - -`azure_ai` exposes two OCR services on one provider; the `doc-intelligence` -sub-route must defer environment resolution to Rust, not accept the generic -`AZURE_AI_API_BASE` fallback that `get_llm_provider` injects. An explicitly -supplied api_base is still always honoured. -""" - -from litellm.llms.azure_ai.ocr.common_utils import ( - is_azure_document_intelligence_model, -) -from litellm.ocr.main import _prepare_ocr_request - -_DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -_AZURE_AI_API_BASE = "https://generic-azure-ai.example.com" - - -class _FakeLogging: - def update_from_kwargs(self, **kwargs: object) -> None: - return None - - -def _prepare(model: str, api_base: str | None): - return _prepare_ocr_request( - model=model, - document=dict(_DOC), - api_key="test-key", - api_base=api_base, - timeout=None, - custom_llm_provider=None, - extra_headers=None, - kwargs={"litellm_logging_obj": _FakeLogging()}, - ) - - -class TestIsAzureDocumentIntelligenceModel: - def test_matches_doc_intelligence_route(self): - assert is_azure_document_intelligence_model("doc-intelligence/prebuilt-layout") - - def test_matches_documentintelligence_and_is_case_insensitive(self): - assert is_azure_document_intelligence_model("azure_ai/DocumentIntelligence/x") - - def test_does_not_match_mistral_route(self): - assert not is_azure_document_intelligence_model("mistral-document-ai-2505") - - -class TestDocIntelligenceApiBaseResolution: - def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch): - """The generic Azure base must not overwrite Rust-owned DI resolution.""" - monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False) - - prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None) - - assert prepared.api_base is None - - def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch): - """A caller-supplied api_base must always win, even for doc-intelligence.""" - monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - - custom = "https://my-di.cognitiveservices.azure.com" - prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom) - - assert prepared.api_base == custom - - def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch): - """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE.""" - monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) - - prepared = _prepare("azure_ai/mistral-document-ai-2505", None) - - assert prepared.api_base == _AZURE_AI_API_BASE diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index feb98d14c03..3526d8c00d6 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -12,15 +12,32 @@ Tests that: import base64 import os import tempfile +from collections.abc import Generator from io import BytesIO from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from typing import Final +from unittest.mock import AsyncMock, MagicMock, Mock import orjson import pytest from starlette.datastructures import FormData -from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type +from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type + + +@pytest.fixture(autouse=True, params=["native", "disabled", "unavailable"]) +def document_runtime(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + from litellm.rust_bridge import bindings, configuration + + configuration.reset_rust_configuration() + monkeypatch.delenv("LITELLM_RUST", raising=False) + if request.param == "disabled": + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust is disabled"))) + elif request.param == "unavailable": + monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) + yield + configuration.reset_rust_configuration() class TestGetMimeType: @@ -480,3 +497,37 @@ class TestProxySecurityGuard: "data:application/pdf;base64," ) assert result["model"] == "mistral/mistral-ocr-latest" + + +@pytest.mark.asyncio +async def test_proxy_upload_stops_reading_at_size_limit() -> None: + from starlette.datastructures import UploadFile + + from litellm.ocr.input import get_max_file_bytes + from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form + + limit: Final = get_max_file_bytes() + with tempfile.TemporaryFile() as stream: + stream.truncate(limit * 2) + upload: Final = UploadFile(file=stream, filename="large.pdf") + request: Final = MagicMock(form=AsyncMock(return_value=FormData({"file": upload}))) + with pytest.raises(ValueError, match="exceeds the size limit"): + await _parse_multipart_form(request) + assert stream.tell() == limit + 1 + + +@pytest.mark.asyncio +async def test_proxy_upload_filename_is_only_metadata(tmp_path: Path) -> None: + from starlette.datastructures import UploadFile + + from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form + + secret: Final = tmp_path / "secret.pdf" + secret.write_bytes(b"server secret") + upload: Final = UploadFile(file=BytesIO(b"uploaded bytes"), filename=str(secret)) + request: Final = MagicMock(form=AsyncMock(return_value=FormData({"file": upload}))) + result: Final = await _parse_multipart_form(request) + assert result["document"] == { + "type": "document_url", + "document_url": "data:application/pdf;base64,dXBsb2FkZWQgYnl0ZXM=", + } diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 46e9a4d3729..4ad556f6941 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -2,37 +2,7 @@ Tests for the OCR `req_format` option in the SDK request path. """ -import pytest - -import litellm from litellm.rust_bridge import ocr as rust_ocr_bridge -from litellm.rust_bridge.ocr import LiteLLMOcrRequest - -DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} - - -def _request( - optional_params: dict[str, object], model: str = "azure_ai/doc-intelligence/prebuilt-layout" -) -> LiteLLMOcrRequest: - return LiteLLMOcrRequest( - model=model, - document=DOCUMENT, - api_key="fake-key", - api_base=None, - custom_llm_provider=None, - extra_headers=None, - timeout=60.0, - kwargs=optional_params, - ) - - -@pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) -def test_rust_ocr_serves_default_format(optional_params): - assert rust_ocr_bridge.supported(_request(optional_params)) is True - - -def test_rust_ocr_serves_native_format_for_document_intelligence(): - assert rust_ocr_bridge.supported(_request({"req_format": "native"})) is True def test_rust_ocr_response_retains_provider_native_response(): @@ -50,34 +20,3 @@ def test_rust_ocr_response_retains_provider_native_response(): assert response.get_provider_native_response() == provider_response assert response.model_dump().get("provider_native_response") is None - - -@pytest.mark.parametrize("model", ["cohere/cohere-parse", "azure_ai/cohere-parse"]) -def test_rust_ocr_skipped_for_unsupported_models(model): - assert rust_ocr_bridge.supported(_request({}, model)) is False - - -@pytest.mark.asyncio -async def test_native_format_rejected_for_provider_without_support_as_bad_request(): - with pytest.raises(litellm.BadRequestError, match="not supported for provider") as exc_info: - await litellm.aocr( - model="mistral/mistral-ocr-latest", - document=DOCUMENT, - api_key="fake-key", - req_format="native", - ) - - assert exc_info.value.status_code == 400 - - -@pytest.mark.asyncio -async def test_unknown_format_rejected_for_provider_without_support_as_bad_request(): - with pytest.raises(litellm.BadRequestError, match="Invalid `req_format`") as exc_info: - await litellm.aocr( - model="mistral/mistral-ocr-latest", - document=DOCUMENT, - api_key="fake-key", - req_format="raw", - ) - - assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py deleted file mode 100644 index dbb4f822d0b..00000000000 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ /dev/null @@ -1,1161 +0,0 @@ -"""Tests for the optional Rust-backed OCR path.""" - -import builtins -import importlib -import types - -import httpx -import pytest - -import litellm -from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import configuration - -# `litellm/__init__.py` does `from .ocr.main import *`, which binds the `ocr` -# function onto `litellm.ocr` and shadows the submodule, so import the modules -# explicitly via importlib rather than attribute traversal. -ocr_main = importlib.import_module("litellm.ocr.main") -rust_bridge = importlib.import_module("litellm.rust_bridge.ocr") -rust_bridge_bindings = importlib.import_module("litellm.rust_bridge.bindings") -rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") - -MODEL = "mistral/mistral-ocr-latest" -DOCUMENT: dict[str, object] = { - "type": "document_url", - "document_url": "https://example.com/doc.pdf", -} - -FAKE_OCR_RESPONSE: dict[str, object] = { - "pages": [{"index": 0, "markdown": "hello world"}], - "model": "mistral-ocr-2505-completion", - "document_annotation": None, - "usage_info": {"pages_processed": 1}, - "object": "ocr", -} - - -class CapturedException(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - -class RecordingBridge: - """A fake ``RustOcr`` callable that records the args it was handed.""" - - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - self.calls.append( - { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "optional_params": optional_params, - "input_sources": input_sources, - "timeout_seconds": timeout_seconds, - } - ) - return dict(FAKE_OCR_RESPONSE) - - -class RecordingAsyncBridge: - """A fake async ``RustAocr`` callable that records the args it was handed.""" - - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - self.calls.append( - { - "model": model, - "document": document, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": custom_llm_provider, - "extra_headers": extra_headers, - "optional_params": optional_params, - "input_sources": input_sources, - "timeout_seconds": timeout_seconds, - } - ) - return dict(FAKE_OCR_RESPONSE) - - -class RaisingBridge: - def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise RuntimeError("bridge failed") - - -class RaisingAsyncBridge: - async def __call__( - self, - model: str, - document: dict[str, object], - api_key: str | None, - api_base: str | None, - custom_llm_provider: str | None, - extra_headers: dict[str, object] | None, - optional_params: dict[str, object], - input_sources: dict[str, str], - timeout_seconds: float | None, - ) -> dict[str, object]: - raise RuntimeError("bridge failed") - - -class RecordingLogging: - """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" - - def __init__(self) -> None: - self.pre_call_kwargs: dict[str, object] | None = None - - def update_from_kwargs(self, **kwargs: object) -> None: - self.update_kwargs = kwargs - - def pre_call( - self, - *, - input: str, - api_key: str | None, - additional_args: dict[str, object], - ) -> None: - self.pre_call_kwargs = { - "input": input, - "api_key": api_key, - "additional_args": additional_args, - } - - -def build_request( - *, - logging_obj: RecordingLogging | None = None, - model: str = "mistral-ocr-latest", - document: dict[str, object] = DOCUMENT, - api_key: str | None = "sk-test", - api_base: str | None = None, - custom_llm_provider: str | None = "mistral", - extra_headers: dict[str, object] | None = None, - optional_params: dict[str, object] | None = None, - litellm_params: dict[str, object] | None = None, - timeout: float | httpx.Timeout | None = 12.5, -) -> rust_bridge.LiteLLMOcrRequest: - return rust_bridge.LiteLLMOcrRequest( - model=model, - document=document, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, - extra_headers=extra_headers, - timeout=timeout, - kwargs={ - **(optional_params or {}), - **(litellm_params or {}), - "litellm_logging_obj": logging_obj or RecordingLogging(), - }, - ) - - -@pytest.fixture(autouse=True) -def _reset_rust_flag(): - """Keep the global toggle isolated between tests.""" - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() - configuration.reset_rust_configuration() - rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - yield - rust_bridge._OCR.reset() - rust_bridge._AOCR.reset() - configuration.reset_rust_configuration() - rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL - - -@pytest.fixture -def fake_bridge(): - """Enable the Rust path with an injected recording bridge (no native wheel).""" - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - return bridge - - -@pytest.fixture -def fake_async_bridge(): - """Enable the async Rust path with an injected recording bridge.""" - bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._AOCR.override(bridge) - return bridge - - -def test_load_rust_ocr_returns_injected_impl(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - assert rust_bridge.load_rust_ocr() is bridge - - -def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch): - real_import = builtins.__import__ - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - if name == "litellm.rust_bridge" and "_native" in fromlist: - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - - -def test_native_bridge_loader_caches_absent_extension(monkeypatch): - real_import = builtins.__import__ - attempts = 0 - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal attempts - if name == "litellm.rust_bridge" and "_native" in fromlist: - attempts += 1 - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - assert rust_bridge_loader.get_native_bridge() is None - assert attempts == 1 - - -def test_native_bridge_loader_reset_forces_relookup(monkeypatch): - real_import = builtins.__import__ - attempts = 0 - - def fake_import(name, globals=None, locals=None, fromlist=(), level=0): - nonlocal attempts - if name == "litellm.rust_bridge" and "_native" in fromlist: - attempts += 1 - raise ImportError - return real_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - assert rust_bridge_loader.get_native_bridge() is None - rust_bridge_loader.reset_native_bridge_cache() - assert rust_bridge_loader.get_native_bridge() is None - assert attempts == 2 - - -def test_native_bridge_available_reflects_loader(monkeypatch): - fake_module = types.ModuleType("litellm.rust_bridge._native") - monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module) - - assert rust_bridge_loader.native_bridge_available() is True - - -def test_load_rust_aocr_returns_injected_impl(): - bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._AOCR.override(bridge) - assert rust_bridge.load_rust_aocr() is bridge - - -def test_toggle_without_ocr_arg_preserves_injected_impl(): - """The public flag must not clobber an internal test binding.""" - bridge = RecordingBridge() - async_bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) - - litellm.rust(False) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge - litellm.rust(True) - assert rust_bridge.load_rust_ocr() is bridge - assert rust_bridge.load_rust_aocr() is async_bridge - - -def test_explicit_ocr_none_clears_injected_impl(monkeypatch): - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: None, - ) - bridge = RecordingBridge() - async_bridge = RecordingAsyncBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - rust_bridge._AOCR.override(async_bridge) - - rust_bridge._OCR.override(None) - rust_bridge._AOCR.override(None) - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None - - -def test_load_rust_ocr_none_when_extension_absent(monkeypatch): - """With no injected impl and no compiled wheel, the loader returns None so the - caller degrades to the Python path instead of raising ImportError.""" - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: None, - ) - litellm.rust(True) # no impl injected; extension isn't built in CI - assert rust_bridge.load_rust_ocr() is None - assert rust_bridge.load_rust_aocr() is None - - -def test_load_rust_ocr_uses_compiled_extension(monkeypatch): - """With no injected impl but a packaged ``litellm.rust_bridge._native`` importable, - the loader returns the extension's ``ocr`` callable. The native wheel isn't - built in CI, so stand in a fake module via the bridge loader.""" - fake_module = types.ModuleType("litellm.rust_bridge._native") - fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - monkeypatch.setattr( - rust_bridge_bindings, - "get_native_bridge", - lambda: fake_module, - ) - - litellm.rust(True) # enabled, no impl injected -> import the extension - assert rust_bridge.load_rust_ocr() is fake_module.ocr - assert rust_bridge.load_rust_aocr() is fake_module.aocr - - -def test_timeout_to_seconds_handles_float_timeout_and_none(): - assert rust_bridge._timeout_to_seconds(12.5) == 12.5 - assert rust_bridge._timeout_to_seconds(None) is None - assert rust_bridge._timeout_to_seconds(httpx.Timeout(30.0, read=42.0)) == 42.0 - - -def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): - bridge = RecordingBridge() - - litellm.rust(True) - - rust_bridge._OCR.override(bridge) - response = rust_bridge.ocr( - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - custom_llm_provider="mistral", - extra_headers={"Authorization": "Bearer sk-test", "x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True, "pages": [0]}, - timeout=12.5, - ) - - assert response == FAKE_OCR_RESPONSE - call = bridge.calls[0] - assert call == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": { - "Authorization": "Bearer sk-test", - "x-trace-id": "trace-1", - }, - "optional_params": {"include_image_base64": True, "pages": [0]}, - "input_sources": {}, - "timeout_seconds": 12.5, - } - - -@pytest.mark.asyncio -async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): - bridge = RecordingAsyncBridge() - - litellm.rust(True) - - rust_bridge._AOCR.override(bridge) - response = await rust_bridge.aocr( - model="mistral-ocr-maas", - document=DOCUMENT, - api_key=None, - api_base=None, - custom_llm_provider="vertex_ai", - extra_headers=None, - optional_params={"vertex_project": "project-1"}, - timeout=httpx.Timeout(30.0, read=42.0), - ) - - assert response == FAKE_OCR_RESPONSE - assert bridge.calls[0] == { - "model": "mistral-ocr-maas", - "document": DOCUMENT, - "api_key": None, - "api_base": None, - "custom_llm_provider": "vertex_ai", - "extra_headers": None, - "optional_params": {"vertex_project": "project-1"}, - "input_sources": {}, - "timeout_seconds": 42.0, - } - - -def test_run_rust_ocr_prepares_request_and_wraps_response(): - bridge = RecordingBridge() - logging_obj = RecordingLogging() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - response = ocr_main._run_rust_ocr( - request=build_request( - logging_obj=logging_obj, - api_base="https://proxy.internal", - extra_headers={"x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True}, - timeout=12.5, - ), - resolve_api_key=lambda _name: None, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert bridge.calls[0] == { - "model": "mistral-ocr-latest", - "document": DOCUMENT, - "api_key": "sk-test", - "api_base": "https://proxy.internal", - "custom_llm_provider": "mistral", - "extra_headers": { - "x-trace-id": "trace-1", - }, - "optional_params": {"include_image_base64": True}, - "input_sources": {}, - "timeout_seconds": 12.5, - } - - -def test_rust_upstream_error_uses_ocr_provider_error_mapping(): - error = RustUpstreamError(400, '{"message":"invalid model"}') - - mapped = ocr_main._map_rust_ocr_error( - error, - build_request(), - (RuntimeError, RustUpstreamError), - ) - - assert isinstance(mapped, BaseLLMException) - assert mapped.status_code == 400 - assert mapped.message == '{"message":"invalid model"}' - - -def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request(api_key=None, timeout=None), - resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, - ) - - assert bridge.calls[0]["api_key"] == "sk-from-vault" - - -def test_run_rust_ocr_prefers_explicit_key_over_resolver(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name: str) -> str | None: - raise AssertionError(f"resolver should not be called for {name}") - - ocr_main._run_rust_ocr( - request=build_request( - api_key="sk-explicit", - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["api_key"] == "sk-explicit" - - -def test_run_rust_ocr_uses_mistral_secret_manager_without_provider_config(): - bridge = RecordingBridge() - resolver_calls = [] - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name): - resolver_calls.append(name) - return "sk-provider-env" - - ocr_main._run_rust_ocr( - request=build_request( - model="mistral-ocr-latest", - api_key=None, - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert resolver_calls == ["MISTRAL_API_KEY"] - assert bridge.calls[0]["api_key"] == "sk-provider-env" - - -def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="vertex_ai", - model="mistral-ocr-maas", - litellm_params={ - "vertex_project": "project-1", - "vertex_location": "us-central1", - "vertex_credentials": "redacted", - }, - optional_params={"include_image_base64": True}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["optional_params"] == { - "include_image_base64": True, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "vertex_credentials": "redacted", - } - - -def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - def _resolver(name: str) -> str | None: - return { - "VERTEXAI_PROJECT": "project-from-secret", - "VERTEXAI_LOCATION": "us-east5", - "VERTEXAI_CREDENTIALS": "credentials-from-secret", - }.get(name) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="vertex_ai", - model="mistral-ocr-maas", - timeout=None, - ), - resolve_api_key=_resolver, - ) - - assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" - assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" - assert bridge.calls[0]["optional_params"]["vertex_credentials"] == "credentials-from-secret" - - -def test_prepare_rust_ocr_call_defers_azure_environment_resolution_to_rust(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - api_base=None, - timeout=None, - ), - resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), - ) - - assert bridge.calls[0]["api_base"] is None - assert bridge.calls[0]["api_key"] is None - assert bridge.calls[0]["extra_headers"] is None - - -def test_prepare_rust_ocr_call_defers_document_intelligence_environment_to_rust(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="doc-intelligence/prebuilt-layout", - api_base=None, - timeout=None, - ), - resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), - ) - - assert bridge.calls[0]["api_base"] is None - - -def test_prepare_rust_ocr_call_forwards_raw_azure_auth_inputs(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - api_base="https://azure.example.com", - extra_headers={"x-trace-id": "trace-1"}, - litellm_params={ - "azure_ad_token": "entra-token", - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_scope": "scope", - "azure_authority_host": "https://login.example.com", - "azure_credential": "ClientSecretCredential", - "azure_federated_token_file": "/token", - }, - timeout=None, - ), - resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), - ) - - call = bridge.calls[0] - assert call["api_key"] is None - assert call["api_base"] == "https://azure.example.com" - assert call["extra_headers"] == {"x-trace-id": "trace-1"} - assert call["optional_params"] == { - "azure_ad_token": "entra-token", - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_scope": "scope", - "azure_authority_host": "https://login.example.com", - "azure_credential": "ClientSecretCredential", - "azure_federated_token_file": "/token", - } - assert call["input_sources"] == {} - - -def test_prepare_rust_ocr_call_preserves_proxy_input_sources(): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - request_values = { - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_authority_host": "https://login.example.com", - "api_base": "https://azure.example.com", - } - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key="request-key", - api_base="https://azure.example.com", - litellm_params={ - "tenant_id": "tenant", - "client_id": "client", - "client_secret": "secret", - "azure_authority_host": "https://login.example.com", - "proxy_server_request": {"body": request_values, "credential_fields": ("api_key",)}, - }, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["input_sources"] == { - **{name: "request" for name in request_values}, - "api_key": "request", - } - - marshaled = rust_bridge._marshal( - build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key="request-key", - api_base="https://azure.example.com", - litellm_params={ - "proxy_server_request": { - "body": {"api_base": "https://azure.example.com"}, - "credential_fields": ("api_key",), - } - }, - ), - lambda _name: None, - lambda document: document, - ) - assert marshaled.input_sources == {"api_base": "request", "api_key": "request"} - - -def test_rust_ocr_logging_redacts_azure_credentials(): - bridge = RecordingBridge() - logging_obj = RecordingLogging() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - logging_obj=logging_obj, - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - litellm_params={"azure_ad_token": "token", "client_secret": "secret"}, - ), - resolve_api_key=lambda _name: None, - ) - - assert logging_obj.update_kwargs["optional_params"] == { - "azure_ad_token": "****", - "client_secret": "****", - } - assert logging_obj.pre_call_kwargs is not None - additional_args = logging_obj.pre_call_kwargs["additional_args"] - assert isinstance(additional_args, dict) - complete_input = additional_args["complete_input_dict"] - assert isinstance(complete_input, dict) - assert complete_input["azure_ad_token"] == "****" - assert complete_input["client_secret"] == "****" - - -def test_rust_eligibility_rejects_python_only_azure_auth_modes(): - for params in ( - {"azure_ad_token_provider": lambda: "token"}, - {"azure_username": "user"}, - {"azure_password": "password"}, - ): - assert not ocr_main._rust_ocr_supported( - build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - litellm_params=params, - ) - ) - - -def test_prepare_rust_ocr_call_forwards_global_azure_refresh(monkeypatch: pytest.MonkeyPatch): - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) - - ocr_main._run_rust_ocr( - request=build_request( - custom_llm_provider="azure_ai", - model="pixtral-12b-2409", - api_key=None, - api_base="https://azure.example.com", - litellm_params={"proxy_server_request": {"body": {"enable_azure_ad_token_refresh": True}}}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert bridge.calls[0]["optional_params"] == {"enable_azure_ad_token_refresh": True} - assert bridge.calls[0]["input_sources"] == {"enable_azure_ad_token_refresh": "deployment"} - - -def test_run_rust_ocr_runs_pre_call_logging(): - logging_obj = RecordingLogging() - bridge = RecordingBridge() - litellm.rust(True) - rust_bridge._OCR.override(bridge) - - ocr_main._run_rust_ocr( - request=build_request( - logging_obj=logging_obj, - api_base="https://api.mistral.ai/v1", - extra_headers={"x-trace-id": "trace-1"}, - optional_params={"include_image_base64": True}, - timeout=None, - ), - resolve_api_key=lambda _name: None, - ) - - assert logging_obj.pre_call_kwargs is not None - assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" - additional_args = logging_obj.pre_call_kwargs["additional_args"] - complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT - assert complete_input["include_image_base64"] is True - assert additional_args["api_base"] == "https://api.mistral.ai/v1" - assert additional_args["headers"] == { - "x-trace-id": "trace-1", - } - - -def test_ocr_routes_to_rust_when_enabled(fake_bridge): - response = litellm.ocr( - model=MODEL, - document=DOCUMENT, - api_key="sk-test", - extra_headers={"x-trace-id": "trace-1"}, - include_image_base64=True, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert len(fake_bridge.calls) == 1 - call = fake_bridge.calls[0] - assert call["model"] == MODEL - assert call["document"] == DOCUMENT - assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] is None - assert call["extra_headers"] == { - "x-trace-id": "trace-1", - } - assert call["optional_params"].get("include_image_base64") is True - - -def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): - response = litellm.ocr( - model="azure_ai/pixtral-12b-2409", - document=DOCUMENT, - api_key="sk-test", - api_base="https://example.services.ai.azure.com", - ) - - assert isinstance(response, OCRResponse) - assert len(fake_bridge.calls) == 1 - assert fake_bridge.calls[0]["model"] == "azure_ai/pixtral-12b-2409" - assert fake_bridge.calls[0]["custom_llm_provider"] is None - assert fake_bridge.calls[0]["extra_headers"] is None - - -def test_ocr_routes_azure_entra_inputs_to_rust_without_python_auth(fake_bridge): - response = litellm.ocr( - model="azure_ai/pixtral-12b-2409", - document=DOCUMENT, - api_base="https://example.services.ai.azure.com", - azure_ad_token="entra-token", - tenant_id="tenant", - client_id="client", - ) - - assert isinstance(response, OCRResponse) - assert fake_bridge.calls[0]["api_key"] is None - assert fake_bridge.calls[0]["extra_headers"] is None - assert fake_bridge.calls[0]["optional_params"] == { - "azure_ad_token": "entra-token", - "tenant_id": "tenant", - "client_id": "client", - } - - -def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): - response = litellm.ocr( - model=MODEL, - document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, - api_key="sk-test", - ) - - assert isinstance(response, OCRResponse) - document = fake_bridge.calls[0]["document"] - assert document["type"] == "document_url" - assert document["document_url"].startswith("data:application/pdf;base64,") - - -def test_ocr_exception_type_uses_resolved_provider_context( - monkeypatch: pytest.MonkeyPatch, -): - captured: dict[str, object] = {} - - def fake_exception_type(**kwargs: object) -> CapturedException: - captured.update(kwargs) - return CapturedException("wrapped") - - monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - litellm.rust(True) - rust_bridge._OCR.override(RaisingBridge()) - - with pytest.raises(CapturedException): - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured["model"] == "mistral-ocr-latest" - assert captured["custom_llm_provider"] == "mistral" - - -@pytest.mark.asyncio -async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): - response = await litellm.aocr( - model=MODEL, - document=DOCUMENT, - api_key="sk-test", - extra_headers={"x-trace-id": "trace-1"}, - include_image_base64=True, - ) - - assert isinstance(response, OCRResponse) - assert response.pages[0].markdown == "hello world" - assert len(fake_async_bridge.calls) == 1 - call = fake_async_bridge.calls[0] - assert call["model"] == MODEL - assert call["document"] == DOCUMENT - assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] is None - assert call["extra_headers"] == { - "x-trace-id": "trace-1", - } - assert call["optional_params"].get("include_image_base64") is True - - -@pytest.mark.asyncio -async def test_aocr_exception_type_uses_resolved_provider_context( - monkeypatch: pytest.MonkeyPatch, -): - captured: dict[str, object] = {} - - def fake_exception_type(**kwargs: object) -> CapturedException: - captured.update(kwargs) - return CapturedException("wrapped") - - monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) - litellm.rust(True) - rust_bridge._AOCR.override(RaisingAsyncBridge()) - - with pytest.raises(CapturedException): - await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured["model"] == "mistral-ocr-latest" - assert captured["custom_llm_provider"] == "mistral" - - -def test_ocr_forwards_timeout_to_rust(fake_bridge): - """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s - client ceiling doesn't silently override shorter deadlines.""" - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test", timeout=12.5) - - assert fake_bridge.calls[0]["timeout_seconds"] == 12.5 - - -def test_ocr_passes_default_request_timeout_to_rust(fake_bridge): - litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - from litellm.constants import request_timeout - - assert fake_bridge.calls[0]["timeout_seconds"] == float(request_timeout) - - -def test_ocr_does_not_route_to_rust_when_disabled(): - """With the flag off, the bridge must not be consulted even if an impl exists.""" - bridge = RecordingBridge() - litellm.rust(False) - rust_bridge._OCR.override(bridge) - # The impl stays available for injection, but the disabled flag gates usage, - # so ocr() never reaches the Rust path (asserted via the enabled-path test). - assert bridge.calls == [] - - -def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): - """Rust enabled but no bridge available (no injected impl, no compiled wheel): - ocr() must degrade to the Python HTTP handler instead of raising.""" - monkeypatch.setattr(rust_bridge, "load_rust_ocr", lambda: None) - litellm.rust(True) # enabled, but load_rust_ocr() returns None in CI - - captured = {} - - def fake_handler_ocr(**kwargs): - captured["called"] = True - return OCRResponse(pages=[], model="mistral-ocr-latest", object="ocr") - - monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) - - response = litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") - - assert captured.get("called") is True # Python path was used - assert isinstance(response, OCRResponse) - - -def test_ocr_provider_configs_expose_api_key_env_vars(): - from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( - AzureDocumentIntelligenceOCRConfig, - ) - from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig - from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig - from litellm.llms.mistral.ocr.transformation import MistralOCRConfig - from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( - VertexAIDeepSeekOCRConfig, - ) - from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig - - assert BaseOCRConfig().get_api_key_env_var() is None - assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" - assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" - assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" - assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" - assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.asyncio -async def test_rust_receives_unmapped_azure_options(asynchronous, fake_bridge, fake_async_bridge): - from typing import Final - - arguments: Final = { - "model": "azure_ai/doc-intelligence/prebuilt-layout", - "document": DOCUMENT, - "api_key": "test-key", - "pages": [0, 2], - "features": ["languages", "style"], - "provider_extension": {"enabled": True}, - } - if asynchronous: - await litellm.aocr(**arguments) - else: - litellm.ocr(**arguments) - call: Final = (fake_async_bridge if asynchronous else fake_bridge).calls[0] - assert call["model"] == arguments["model"] - assert call["custom_llm_provider"] is None - assert call["extra_headers"] is None - assert call["optional_params"] == { - "pages": [0, 2], - "features": ["languages", "style"], - "provider_extension": {"enabled": True}, - } - - -@pytest.mark.parametrize("enabled", [False, True]) -@pytest.mark.asyncio -async def test_python_fallback_maps_original_options_once(enabled, monkeypatch): - from io import BytesIO - from typing import Final - - class PythonHandler: - def __init__(self): - self.calls = [] - - def ocr(self, **kwargs): - self.calls.append(kwargs) - return OCRResponse(pages=[], model=kwargs["model"]) - - handler: Final = PythonHandler() - monkeypatch.setattr(ocr_main, "base_llm_http_handler", handler) - litellm.rust(enabled) - rust_bridge._OCR.override(None) - rust_bridge._AOCR.override(None) - for asynchronous in (False, True): - file: Final = BytesIO(b"test document") - arguments: Final = { - "model": "azure_ai/doc-intelligence/prebuilt-layout", - "document": {"type": "file", "file": file}, - "api_key": "test-key", - "pages": [0, 2], - } - if asynchronous: - await litellm.aocr(**arguments) - else: - litellm.ocr(**arguments) - assert handler.calls[-1]["optional_params"]["pages"] == "1,3" - assert handler.calls[-1]["document"]["document_url"].endswith("dGVzdCBkb2N1bWVudA==") - assert len(handler.calls) == 2 - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) -@pytest.mark.asyncio -async def test_native_public_ocr_matches_python(model, asynchronous): - import json - from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer - from threading import Thread - from typing import Final - from urllib.parse import parse_qsl, urlsplit - - native: Final = rust_bridge_loader.get_native_bridge() - if native is None: - pytest.skip("requires the compiled Rust extension") - calls: Final = [] - - class Handler(BaseHTTPRequestHandler): - def do_POST(self): - body: Final = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) - target: Final = urlsplit(self.path) - calls.append( - ( - target.path, - parse_qsl(target.query), - self.headers.get("Authorization"), - self.headers.get("Ocp-Apim-Subscription-Key"), - body, - ) - ) - payload: Final = ( - {"status": "succeeded", "analyzeResult": {"pages": []}} - if "doc-intelligence" in model - else {"pages": [{"index": 0, "markdown": "hello"}]} - ) - encoded: Final = json.dumps(payload).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(encoded))) - self.end_headers() - self.wfile.write(encoded) - - def log_message(self, *_args): - pass - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread: Final = Thread(target=server.serve_forever, daemon=True) - thread.start() - responses: Final = [] - try: - for enabled in (False, True): - litellm.rust(enabled) - arguments: Final = { - "model": model, - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "pages": [0, 2], - "timeout": 3.0, - } - response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - responses.append(response.model_dump()) - assert len(calls) == 2 - assert calls[0] == calls[1] - for key in ("model", "pages", "object"): - assert responses[0][key] == responses[1][key] - finally: - server.shutdown() - server.server_close() - thread.join(timeout=3) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3a7ae7aba61..2932373c77e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1,6 +1,8 @@ """Tests for unified guardrail.""" import logging +from types import SimpleNamespace +from typing import Final import pytest @@ -19,14 +21,14 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( openai_messages_without_system, openai_messages_without_tool, ) +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse +from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) -from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse -from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) @@ -644,6 +646,64 @@ class TestUnifiedLLMGuardrails: class TestOCRGuardrailE2E: """End-to-end tests: UnifiedLLMGuardrails -> OCRHandler.""" + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", [CallTypes.ocr, CallTypes.aocr, CallTypes.aresponses]) + async def test_post_call_logging_fallback_is_limited_to_ocr(self, call_type: CallTypes) -> None: + guardrail: Final = RecordingGuardrail() + response: Final = ( + TestUnifiedLLMGuardrails.TestResponsesRouteAliases._responses_api_response() + if call_type == CallTypes.aresponses + else OCRResponse(model="mistral-ocr-latest", pages=[OCRPage(index=0, markdown="Scan this page")]) + ) + + result: Final = await UnifiedLLMGuardrails().async_post_call_success_hook( + data={ + "guardrail_to_apply": guardrail, + "litellm_logging_obj": SimpleNamespace(call_type=call_type.value), + }, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert result is response + if call_type in (CallTypes.ocr, CallTypes.aocr): + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Scan this page"] + else: + assert guardrail.apply_calls == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"]) + async def test_ocr_logging_fallback_preserves_route_and_response_precedence( + self, request_route: str | None, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.types.utils import ModelResponse + + _patch_translation_mappings( + monkeypatch, + { + CallTypes.completion: OpenAIChatCompletionsHandler, + CallTypes.acompletion: OpenAIChatCompletionsHandler, + CallTypes.aocr: OCRHandler, + }, + ) + guardrail: Final = RecordingGuardrail() + response: Final = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Chat output"}}]) + + result: Final = await guardrail.async_post_call_success_deployment_hook( + request_data={ + "guardrails": [guardrail.guardrail_name], + "user_api_key_request_route": request_route, + "litellm_logging_obj": SimpleNamespace(call_type=CallTypes.aocr.value), + }, + response=response, + call_type=CallTypes.aocr, + ) + + assert result is response + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Chat output"] + @pytest.mark.asyncio async def test_pre_call_hook_invokes_ocr_handler_for_input(self): """ diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 8fde4cc9d5e..11c3d2f8b20 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -15,7 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. """ import asyncio -from typing import Any +from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -297,6 +297,38 @@ async def test_no_flag_fires_create_task_normally(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("call_type", ["ocr", "aocr", "completion", "acompletion", "embedding", "responses"]) +@pytest.mark.parametrize("exception_raised", [False, True]) +def test_native_pending_logging_is_released_only_for_ocr(call_type: str, exception_raised: bool) -> None: + pending: Final = MagicMock() + enqueue: Final = MagicMock() + logger: Final = MagicMock( + call_type=call_type, + _native_pending_logging=pending, + _enqueue_deferred_logging=enqueue, + ) + + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( + logging_obj=logger, + exception_raised=exception_raised, + ) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( + logging_obj=logger, + exception_raised=exception_raised, + ) + + if call_type in ("ocr", "aocr"): + pending.release.assert_called_once_with(not exception_raised) + assert logger._native_pending_logging is None + else: + pending.release.assert_not_called() + assert logger._native_pending_logging is pending + if exception_raised: + enqueue.assert_not_called() + else: + enqueue.assert_called_once_with() + + def test_flush_deferred_async_logging_fires_on_success(): """ Happy path: with no exception, the production flush helper invokes the diff --git a/tests/test_litellm/rust_bridge/test_configuration.py b/tests/test_litellm/rust_bridge/test_configuration.py index aff9d5acac1..08fa3bfc053 100644 --- a/tests/test_litellm/rust_bridge/test_configuration.py +++ b/tests/test_litellm/rust_bridge/test_configuration.py @@ -52,6 +52,18 @@ def test_resolution_precedence( def test_release_default_remains_disabled() -> None: assert configuration.DEFAULT_RUST_ENABLED is False assert configuration.rust_enabled() is False + assert configuration.rust_ocr_enabled() is True + + +@pytest.mark.parametrize("process", [None, False, True]) +@pytest.mark.parametrize("environment", [None, "0", "1", "off"]) +def test_ocr_configuration(monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None) -> None: + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + if process is not None: + configuration.rust(process) + + assert configuration.rust_ocr_enabled() is (environment not in {"0", "off"} and process is not False) def test_process_override_wins_over_environment(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py b/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py new file mode 100644 index 00000000000..501a4e986c0 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py @@ -0,0 +1,230 @@ +from collections.abc import Generator, Mapping +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.ocr import legacy +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge.ocr import LiteLLMOcrRequest +from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE + + +@pytest.fixture(autouse=True) +def isolated_ocr_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + monkeypatch.delenv("LITELLM_RUST", raising=False) + configuration.reset_rust_configuration() + yield + NATIVE_OCR_LIFECYCLE.reset() + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_unavailable_native_uses_legacy(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None: + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + NATIVE_OCR_LIFECYCLE.override(None) + document: Final = {"type": "document_url", "document_url": "https://example.com"} + + result: Final = ( + await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[0]) + if asynchronous + else litellm.ocr("mistral/mistral-ocr-latest", document, pages=[0]) + ) + + assert result is response + fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[0]) + + +def test_admitted_failure_is_returned_without_replay() -> None: + failure: Final = RuntimeError("admitted") + native: Final = Mock(side_effect=failure) + litellm.rust(True) + NATIVE_OCR_LIFECYCLE.override(native) + try: + with pytest.raises(RuntimeError) as caught: + litellm.ocr("mistral/mistral-ocr-latest", {"type": "document_url", "document_url": "https://example.com"}) + assert caught.value is failure + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + assert native.call_count == 1 + + +def test_public_binding_keeps_positional_fields_and_defaults_out_of_native_hook_kwargs() -> None: + document: Final = {"type": "document_url", "document_url": "https://example.com"} + captured: Final = [] + + def native( + request: LiteLLMOcrRequest, + args: tuple[object, ...], + kwargs: Mapping[str, object], + asynchronous: bool, + ) -> OCRResponse: + captured.append((request, args, kwargs, asynchronous)) + return OCRResponse(pages=[], model=request.model) + + litellm.rust(True) + NATIVE_OCR_LIFECYCLE.override(native) + try: + response: Final = litellm.ocr("mistral/mistral-ocr-latest", document) + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + + request, call_args, hook_kwargs, asynchronous = captured[0] + assert response.model == "mistral/mistral-ocr-latest" + assert request.model == "mistral/mistral-ocr-latest" + assert request.document is document + assert call_args == ("mistral/mistral-ocr-latest", document) + assert hook_kwargs == {} + assert asynchronous is False + + +def test_public_binding_keeps_keyword_model_and_document_in_native_hook_kwargs() -> None: + document: Final = {"type": "document_url", "document_url": "https://example.com"} + captured: Final = [] + + def native( + request: LiteLLMOcrRequest, + args: tuple[object, ...], + kwargs: Mapping[str, object], + asynchronous: bool, + ) -> OCRResponse: + assert args == () + captured.append(kwargs) + return OCRResponse(pages=[], model=request.model) + + litellm.rust(True) + NATIVE_OCR_LIFECYCLE.override(native) + try: + litellm.ocr(model="mistral/mistral-ocr-latest", document=document) + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + + assert captured[0]["model"] == "mistral/mistral-ocr-latest" + assert captured[0]["document"] is document + assert "timeout" not in captured[0] + + +@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"]) +def test_public_duplicate_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None: + native: Final = Mock(side_effect=AssertionError("binding errors precede admission")) + document: Final = {"type": "document_url", "document_url": "https://example.com"} + litellm.rust(enabled) + NATIVE_OCR_LIFECYCLE.override(native) + try: + with pytest.raises(TypeError, match=r"ocr\(\) got multiple values for argument 'model'"): + litellm.ocr("mistral/mistral-ocr-latest", document, model="duplicate") + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + assert native.call_count == 0 + + +@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"]) +def test_public_missing_required_argument_error_does_not_depend_on_native_selection(enabled: bool) -> None: + native: Final = Mock(side_effect=AssertionError("binding errors precede admission")) + litellm.rust(enabled) + NATIVE_OCR_LIFECYCLE.override(native) + try: + with pytest.raises(TypeError, match=r"ocr\(\) missing 1 required positional argument: 'document'"): + litellm.ocr("mistral/mistral-ocr-latest") + finally: + NATIVE_OCR_LIFECYCLE.reset() + litellm.rust(None) + assert native.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("enabled", [False, True, None]) +async def test_environment_opt_out_never_loads_native( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, enabled: bool | None +) -> None: + monkeypatch.setenv("LITELLM_RUST", "0") + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + load: Final = Mock(side_effect=AssertionError("native must not be loaded")) + monkeypatch.setattr(bindings, "get_native_bridge", load) + litellm.rust(enabled) + document: Final = {"type": "file", "file": b"pdf"} + + result: Final = ( + await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[1]) + if asynchronous + else litellm.ocr("mistral/mistral-ocr-latest", document, pages=[1]) + ) + + assert result is response + fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[1]) + load.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("environment", [None, "1"]) +async def test_native_is_enabled_by_default( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, environment: str | None +) -> None: + if environment is not None: + monkeypatch.setenv("LITELLM_RUST", environment) + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + native: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + NATIVE_OCR_LIFECYCLE.override(native) + fallback: Final = Mock(side_effect=AssertionError("legacy must not run")) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + + result: Final = ( + await litellm.aocr("mistral/mistral-ocr-latest", {}) + if asynchronous + else litellm.ocr("mistral/mistral-ocr-latest", {}) + ) + + assert result is response + assert native.call_count == 1 + fallback.assert_not_called() + + +class Declined(Exception): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("declined", [False, True]) +async def test_only_native_declines_replay_on_legacy( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, declined: bool +) -> None: + failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called") + native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure) + NATIVE_OCR_LIFECYCLE.override(native) + import importlib + + main: Final = importlib.import_module("litellm.ocr.main") + monkeypatch.setattr(main, "native_exception_types", lambda: (Declined, RuntimeError)) + response: Final = OCRResponse(pages=[], model="mistral-ocr-latest") + fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response) + monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback) + document: Final = {"type": "file", "file": b"pdf"} + + async def call() -> object: + if asynchronous: + return await litellm.aocr("mistral/mistral-ocr-latest", document, pages=[0]) + return litellm.ocr("mistral/mistral-ocr-latest", document, pages=[0]) + + if declined: + assert await call() is response + fallback.assert_called_once_with("mistral/mistral-ocr-latest", document, pages=[0]) + else: + with pytest.raises(RuntimeError) as caught: + await call() + assert caught.value is failure + fallback.assert_not_called() + assert native.call_count == 1 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index c7e46829aba..19ed31c7b22 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import contextvars import json import logging import os @@ -7,6 +8,7 @@ import queue import threading from datetime import datetime, timedelta, timezone from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -60,6 +62,36 @@ from litellm.utils import ( # Adds the parent directory to the system path +def test_non_ocr_wrapper_preserves_logging_executor_and_context(monkeypatch: pytest.MonkeyPatch) -> None: + marker: Final = contextvars.ContextVar("non-ocr-logging-context", default="missing") + token: Final = marker.set("caller-context") + caller_thread: Final = threading.get_ident() + response: Final = object() + logger: Final = MagicMock() + observed: Final = queue.Queue[tuple[object, str, int]]() + + def record_success(result: object, start_time: datetime, end_time: datetime) -> None: + observed.put((result, marker.get(), threading.get_ident())) + + def embedding(**kwargs: object) -> object: + return response + + logger.success_handler.side_effect = record_success + monkeypatch.setattr("litellm.utils.function_setup", MagicMock(return_value=(logger, {}))) + try: + with ThreadPoolExecutor(max_workers=1) as executor: + monkeypatch.setattr("litellm.utils.executor", executor) + result: Final = client(embedding)() + logged_response, context, worker_thread = observed.get_nowait() + assert result is response + assert logged_response is response + assert context == "caller-context" + assert worker_thread != caller_thread + assert observed.empty() + finally: + marker.reset(token) + + def test_cloudflare_model_info_includes_rpm(local_model_cost_map: None) -> None: assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md deleted file mode 100644 index 4c117fb846b..00000000000 --- a/tests/test_litellm_rust/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# Rust OCR bridge tests - -This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR behavior tests live under `ocr/`; reusable OCR request, callback, and recording-server fixtures live under `support/` - -A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions - -`ocr/test_requests.py` covers provider payloads, file preparation, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. These contract modules call the Rust bridge directly. `ocr/test_dispatch.py` has the single public API dispatch test, covering enabled native dispatch and disabled Python dispatch. `test_ocr.py` is a strict smoke test of the compiled Rust OCR transport - -Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports - -Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture isolates callback and configuration state but does not select a backend. Native contract tests call `litellm.rust_bridge.ocr` directly, while the strict dispatch test explicitly enables and disables Rust and records which OCR entrypoint runs - -The OCR contract modules are non-strict expected failures until the retained callback implementation from #40070 lands. The public dispatch test remains strict. Passing contract cases appear as XPASS so staging coverage stays visible diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index b0c75d9d2f5..4387ea2e2fd 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -11,7 +11,7 @@ import pytest_asyncio import litellm from litellm import utils -from litellm.litellm_core_utils import litellm_logging +from litellm.litellm_core_utils import litellm_logging, thread_pool_executor from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # preserve raw configuration state in test isolation _CONFIGURATION, @@ -29,11 +29,6 @@ CALLBACK_ATTRIBUTES: Final = ( "_async_success_callback", "_async_failure_callback", ) -EXPECTED_FAILURE_REASONS: Final = { - "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070", - "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070", - "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070", -} def _list_attribute(container: ModuleType, attribute: str) -> list[object]: @@ -76,7 +71,9 @@ async def isolate_ocr_test_state() -> AsyncIterator[None]: stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache stack.enter_context(_rebound(_CONFIGURATION, "override", None)) executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") + stack.enter_context(_rebound(litellm_logging, "executor", executor)) stack.enter_context(_rebound(utils, "executor", executor)) + stack.enter_context(_rebound(thread_pool_executor, "executor", executor)) try: yield finally: @@ -94,14 +91,6 @@ def recording_server() -> Generator[RecordingServer]: def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - for item in items: - if "test_litellm_rust" not in item.path.parts: - continue - relative_path: Final = "/".join(item.path.parts[item.path.parts.index("test_litellm_rust") + 1 :]) - reason: Final = EXPECTED_FAILURE_REASONS.get(relative_path) - if reason is not None: - item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) - if not _parse_env_bool(os.environ.get("LITELLM_RUST")): skip: Final = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") for item in items: diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index b08446412c0..1cfd04b1bff 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -41,7 +41,7 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_ observations: Final = [] class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observations.append((model, copy.deepcopy(kwargs["additional_args"]))) call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0]) @@ -64,13 +64,13 @@ def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider( observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["include_image_base64"] = True if raise_after_edit: raise RuntimeError("pre-call callback failed") class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(copy.deepcopy(request_body(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False) @@ -83,11 +83,11 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ observed: Final = [] class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_headers(kwargs)["x-audit-tag"] = "reviewed" class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(dict(request_headers(kwargs))) call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()]) @@ -96,6 +96,29 @@ def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_ assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" +def test_native_ocr_pre_call_header_rebinding_does_not_replace_execution_root(ocr_server: RecordingServer) -> None: + retained: Final = [] + observed: Final = [] + + class RetainMutateAndRebind(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + headers = request_headers(kwargs) + retained.append(headers) + kwargs["additional_args"]["headers"] = {"x-rebound": "not-sent"} + headers["x-retained"] = "sent" + + class ObserveRebinding(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(dict(request_headers(kwargs))) + + call_native_ocr_with_callbacks(ocr_server, [RetainMutateAndRebind(), ObserveRebinding()]) + + assert observed == [{"x-rebound": "not-sent"}] + assert retained[0]["x-retained"] == "sent" + assert ocr_server.requests[0].headers["x-retained"] == "sent" + assert "x-rebound" not in ocr_server.requests[0].headers + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references( @@ -107,12 +130,12 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ aliases: Final = [] class Retain(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): aliases.append(request_body(kwargs)["document"] is original) retained.append(request_body(kwargs)["document"]) class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): original["document_url"] = replacement_url arguments: Final = { @@ -143,7 +166,7 @@ def test_native_ocr_pre_call_document_replacement_does_not_mutate_original_docum retained: Final = [] class RetainAndReplace(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): body = request_body(kwargs) retained.append(body["document"]) body["document"] = replacement @@ -165,11 +188,11 @@ def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_prov observed: Final = [] class Rebind(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} class Observe(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): observed.append(request_body(kwargs)) call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()]) @@ -182,11 +205,11 @@ def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_ queued: Final = [] class QueuePayload(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): queued.append(request_body(kwargs)) class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): request_body(kwargs)["queued-edit"] = True call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()]) @@ -200,7 +223,7 @@ def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(o finished: Final = threading.Event() class Stash(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["test-token"] = token def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -280,7 +303,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal observed: Final = [] class TrackInFlightRequest(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): kwargs["request-token"] = token def log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -364,7 +387,7 @@ async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context return "caller-token" class Edit(CustomLogger): - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): assert request_headers(kwargs)["Authorization"] == "Bearer caller-token" observations.append("pre_call") request_headers(kwargs)["Authorization"] = "Bearer edited" diff --git a/tests/test_litellm_rust/ocr/test_cohere.py b/tests/test_litellm_rust/ocr/test_cohere.py new file mode 100644 index 00000000000..2a35dc62bd1 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_cohere.py @@ -0,0 +1,141 @@ +from typing import Final + +import pytest + +import litellm +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension +MODELS: Final = ("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0") +IMAGE: Final = {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} +BOX: Final = {"top_left_x": 0, "top_left_y": 0, "bottom_right_x": 32, "bottom_right_y": 32} +PAYLOAD: Final = { + "pages": [ + { + "index": 4, + "markdown": {"content": "receipt", "images": [{"id": "image", "bounding_box": BOX, "description": "scan"}]}, + }, + {"markdown": {"content": "page two"}}, + ], + "meta": {"billed_units": {"pages": 3}}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_public_cohere_request_and_normalization( + recording_server: RecordingServer, model: str, asynchronous: bool +) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + args: Final = { + "model": model, + "document": IMAGE, + "api_base": recording_server.base_url, + "api_key": "test-key", + "req_format": "native", + "unrecognized": True, + } + response: Final = await litellm.aocr(**args) if asynchronous else litellm.ocr(**args) + request: Final = recording_server.requests[0] + assert request.path == ("/providers/cohere/v2/parse" if model.startswith("azure_ai/") else "/v2/parse") + assert request.headers["authorization"] == "Bearer test-key" + assert request.body == {"model": model.split("/", 1)[1], "document": IMAGE, "output_format": "markdown"} + assert [page.index for page in response.pages] == [4, 1] + assert response.pages[0].markdown == "receipt" + assert response.pages[0].images[0].bbox == BOX + assert response.pages[0].images[0].model_extra["description"] == "scan" + assert response.pages[1].images is None + assert response.usage_info.pages_processed == 3 + assert response.get_provider_native_response() == PAYLOAD + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_blocks_and_usage_fallback(recording_server: RecordingServer, model: str) -> None: + blocks: Final = [{"type": "text", "text": "total"}] + recording_server.enqueue(ResponseSpec(body={"pages": [{"blocks": blocks}]})) + response: Final = await litellm.aocr( + model=model, document=IMAGE, api_base=recording_server.base_url, api_key="test-key", output_format="blocks" + ) + assert recording_server.requests[0].body["output_format"] == "blocks" + assert response.pages[0].model_extra["blocks"] == blocks + assert response.pages[0].markdown == "" + assert response.usage_info.pages_processed == 1 + assert response.get_provider_native_response() is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +@pytest.mark.parametrize( + "document", + [ + {"type": "document_url", "document_url": "https://example.com/file.pdf"}, + {"type": "image_url", "image_url": "data:application/pdf;base64,YQ=="}, + {"type": "image_url", "image_url": ""}, + ], +) +async def test_public_cohere_rejects_non_images_before_network( + recording_server: RecordingServer, model: str, document: dict[str, str] +) -> None: + recording_server.expected_requests = 0 + with pytest.raises(litellm.BadRequestError, match="only accepts `image_url`"): + await litellm.aocr(model=model, document=document, api_base=recording_server.base_url, api_key="test-key") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_rejects_unknown_format(recording_server: RecordingServer, model: str) -> None: + recording_server.expected_requests = 0 + with pytest.raises(litellm.BadRequestError, match="output_format"): + await litellm.aocr( + model=model, document=IMAGE, api_base=recording_server.base_url, api_key="test-key", output_format="html" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_provider_failure(recording_server: RecordingServer, model: str) -> None: + recording_server.enqueue(ResponseSpec(status=400, body={"message": "output_format must be blocks or markdown"})) + with pytest.raises(litellm.BadRequestError, match="output_format must be") as caught: + await litellm.aocr(model=model, document=IMAGE, api_base=recording_server.base_url, api_key="test-key") + assert caught.value.status_code == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", MODELS) +async def test_public_cohere_health_check(recording_server: RecordingServer, model: str) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + response: Final = await litellm.ahealth_check( + model_params={"model": model, "api_key": "test-key", "api_base": recording_server.base_url}, mode="ocr" + ) + assert "error" not in response + assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("suffix", ["", "/cohere/", "/v2", "/v2/parse"]) +async def test_public_cohere_url_variants(recording_server: RecordingServer, suffix: str) -> None: + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + await litellm.aocr(model=MODELS[0], document=IMAGE, api_base=recording_server.base_url + suffix, api_key="test-key") + assert recording_server.requests[0].path == ("/cohere/v2/parse" if suffix == "/cohere/" else "/v2/parse") + + +@pytest.mark.asyncio +async def test_public_cohere_environment_key_and_remote_url( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("COHERE_API_KEY", "env-key") + recording_server.enqueue(ResponseSpec(body=PAYLOAD)) + document: Final = {"type": "image_url", "image_url": "https://example.com/receipt.png"} + await litellm.aocr(model=MODELS[0], document=document, api_base=recording_server.base_url) + assert recording_server.requests[0].headers["authorization"] == "Bearer env-key" + assert recording_server.requests[0].body["document"] == document + + +@pytest.mark.asyncio +async def test_public_cohere_missing_key(recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("COHERE_API_KEY", raising=False) + recording_server.expected_requests = 0 + with pytest.raises(Exception, match="Missing COHERE_API_KEY"): + await litellm.aocr(model=MODELS[0], document=IMAGE, api_base=recording_server.base_url) diff --git a/tests/test_litellm_rust/ocr/test_dispatch.py b/tests/test_litellm_rust/ocr/test_dispatch.py index a6c76bc5d0e..7b4b9fab579 100644 --- a/tests/test_litellm_rust/ocr/test_dispatch.py +++ b/tests/test_litellm_rust/ocr/test_dispatch.py @@ -1,11 +1,9 @@ from typing import Final -from unittest.mock import Mock import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.ocr import main as ocr_main from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE @@ -18,18 +16,9 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer: return recording_server -@pytest.mark.parametrize("rust_enabled", [True, False], ids=["enabled", "disabled"]) -def test_public_ocr_dispatches_according_to_rust_setting( - ocr_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - rust_enabled: bool, -) -> None: - rust_call: Final = Mock(wraps=ocr_main.rust_ocr_bridge.ocr) - python_call: Final = Mock(wraps=ocr_main.base_llm_http_handler.ocr) - monkeypatch.setattr(ocr_main.rust_ocr_bridge, "ocr", rust_call) - monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", python_call) - litellm.rust(rust_enabled) - +@pytest.mark.parametrize("enabled", [False, True, None]) +def test_public_ocr_uses_native_route_independently_of_flag(ocr_server: RecordingServer, enabled: bool | None) -> None: + litellm.rust(enabled) response: Final = litellm.ocr( model=OCR_MODEL, document=OCR_DOCUMENT, @@ -39,6 +28,26 @@ def test_public_ocr_dispatches_according_to_rust_setting( assert isinstance(response, OCRResponse) assert response.pages[0].markdown == "native OCR response" - assert rust_call.call_count == int(rust_enabled) - assert python_call.call_count == int(not rust_enabled) + assert len(ocr_server.requests) == 1 + assert not ocr_server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("caching", [None, False, True]) +async def test_ocr_does_not_depend_on_chat_cache( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, caching: bool | None +) -> None: + from litellm.caching.caching import Cache + + monkeypatch.setattr(litellm, "cache", Cache(type="local", supported_call_types=["completion", "acompletion"])) + arguments: Final = { + "model": OCR_MODEL, + "document": OCR_DOCUMENT, + "api_key": "test-key", + "api_base": ocr_server.base_url, + "caching": caching, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert response.pages[0].markdown == "native OCR response" assert len(ocr_server.requests) == 1 diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py new file mode 100644 index 00000000000..1acad5527d8 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -0,0 +1,996 @@ +import asyncio +import datetime +import gc +import json +import sys +import threading +import weakref +from collections.abc import Coroutine +from contextvars import ContextVar +from typing import Final + +import pytest + +import litellm +from litellm._logging import trace_id_var +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_aocr, call_ocr + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["deployment", "failure"]) +async def test_cancellation_during_failure_obeys_phase_policy(ocr_server: RecordingServer, phase: str) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider failure"}, status=500)) + entered: Final = asyncio.Event() + observed: Final = [] + + class Observer(CustomLogger): + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, **kwargs): + if phase == "deployment": + entered.set() + await asyncio.Event().wait() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(kwargs["exception"]) + if phase == "failure": + entered.set() + await asyncio.Event().wait() + + observer: Final = Observer() + litellm.callbacks.append(observer) + task: Final = asyncio.create_task(call_aocr(ocr_server, callbacks=[observer])) + await asyncio.wait_for(entered.wait(), 5) + task.cancel() + if phase == "deployment": + with pytest.raises(litellm.InternalServerError) as caught: + await task + assert observed == [caught.value] + else: + with pytest.raises(asyncio.CancelledError): + await task + assert len(observed) == 1 + assert isinstance(observed[0], litellm.InternalServerError) + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +@pytest.mark.asyncio +async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) -> None: + from litellm.proxy._types import UserAPIKeyAuth + + recorder: Final = RecordingLogger() + auth: Final = UserAPIKeyAuth(user_id="ocr-user") + response: Final = await call_aocr( + ocr_server, callbacks=[recorder], metadata={"user_api_key_auth": auth}, shared_session=object() + ) + events: Final = await recorder.wait_for_async("async_log_success_event") + assert response.pages[0].markdown == "native OCR response" + assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user" + assert "metadata" not in ocr_server.requests[0].body + + +@pytest.mark.asyncio +async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None: + caller: Final = asyncio.current_task() + context: Final = ContextVar("lifecycle-test", default="before") + observations: Final = [] + recorder: Final = RecordingLogger() + + class Replace(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + context.set("pre") + observations.append(("pre", asyncio.current_task(), context.get())) + return {**kwargs, "pages": [2]} + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + observations.append(("post", asyncio.current_task(), context.get())) + return response.model_copy(update={"model": "replaced"}) + + litellm.callbacks.append(Replace()) + response: Final = await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="native-final") + events: Final = await recorder.wait_for_async("async_log_success_event") + assert observations == [("pre", caller, "pre"), ("post", caller, "pre")] + assert context.get() == "pre" + assert ocr_server.requests[0].body["pages"] == [2] + assert response.model == "replaced" + assert events[0].response is response + assert response._hidden_params["litellm_call_id"] == "native-final" + assert "response_cost" in response._hidden_params + + +@pytest.mark.asyncio +async def test_deployment_hook_replaces_complete_routing_request(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.05)) + original: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + observed: Final = [] + + class Replace(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + return { + **kwargs, + "model": "azure_ai/mistral-ocr-latest", + "custom_llm_provider": "azure_ai", + "document": replacement, + "api_key": "replacement-key", + "api_base": ocr_server.base_url, + "extra_headers": {"x-deployment": "replacement"}, + "timeout": 2, + "pages": [2], + } + + class Observe(Logging): + def pre_call(self, input, api_key, additional_args): + observed.append((additional_args["complete_input_dict"]["document"], api_key)) + + litellm.callbacks.append(Replace()) + logger: Final = Observe( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deployment-routing", + function_id="deployment-routing", + ) + response: Final = await call_aocr( + ocr_server, + document=original, + timeout=0.001, + litellm_logging_obj=logger, + ) + + assert response.pages[0].markdown == "native OCR response" + assert observed == [(replacement, "replacement-key")] + assert observed[0][0] is replacement + assert replacement == original + assert replacement is not original + assert original == {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + assert ocr_server.requests[0].path == "/providers/mistral/azure/ocr" + assert ocr_server.requests[0].headers["authorization"] == "Bearer replacement-key" + assert ocr_server.requests[0].headers["x-deployment"] == "replacement" + assert ocr_server.requests[0].body["document"] == replacement + assert ocr_server.requests[0].body["pages"] == [2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_metadata_failure_dispatches_only_failure_and_releases_logger( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + failure: Final = RuntimeError("metadata failed") + seen: Final = [] + + class FailingMetadata(Logging): + def _response_cost_calculator(self, *args, **kwargs): + raise failure + + def success_handler(self, *args, **kwargs): + seen.append("success") + + def failure_handler(self, exception, *args, **kwargs): + seen.append(("sync", exception)) + + async def async_failure_handler(self, exception, *args, **kwargs): + seen.append(("async", exception)) + + async def invoke(): + logger: Final = FailingMetadata( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr" if asynchronous else "ocr", + start_time=datetime.datetime.now(), + litellm_call_id="metadata", + function_id="metadata", + ) + reference: Final = weakref.ref(logger) + with pytest.raises(RuntimeError) as caught: + await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr( + ocr_server, litellm_logging_obj=logger + ) + assert caught.value is failure + failure.__traceback__ = None + return reference + + reference: Final = await invoke() + await drain_logging() + gc.collect() + assert seen == ([("sync", failure), ("async", failure)] if asynchronous else [("sync", failure)]) + assert reference() is None + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_mapped_failure_identity_and_deployment_snapshot(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "unavailable"}, status=500)) + recorder: Final = RecordingLogger() + snapshots: Final = [] + + class Observe(CustomLogger): + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, **kwargs): + snapshots.append(exception) + exception.status_code = 418 + + litellm.callbacks.append(Observe()) + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, callbacks=[recorder]) + failures: Final = tuple(event for event in recorder.events if "failure" in event.name) + assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"] + assert all(event.kwargs["exception"] is caught.value for event in failures) + assert caught.value.status_code == 500 + assert snapshots[0] is not caught.value + assert snapshots[0].status_code == 418 + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["pre", "http", "post"]) +async def test_cancellation_cleans_up_in_caller_task_without_terminal_dispatch( + ocr_server: RecordingServer, phase: str +) -> None: + entered: Final = asyncio.Event() + recorder: Final = RecordingLogger() + + class Pause(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + if phase == "pre": + entered.set() + await asyncio.Event().wait() + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + if phase == "post": + entered.set() + await asyncio.Event().wait() + + litellm.callbacks.append(Pause()) + if phase == "http": + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) + if phase == "pre": + ocr_server.expected_requests = 0 + restored: Final = [] + + async def invoke(): + trace_id_var.set("parent") + try: + await call_aocr(ocr_server, callbacks=[recorder], litellm_trace_id="native-call") + finally: + restored.append(trace_id_var.get()) + + task: Final = asyncio.create_task(invoke()) + if phase == "http": + await ocr_server.wait_for_requests(1) + else: + await asyncio.wait_for(entered.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await drain_logging() + assert restored == ["parent"] + assert not any("success" in name or "failure" in name for name in recorder.names) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True]) +async def test_deferred_logging_requires_release_and_runs_at_most_once( + ocr_server: RecordingServer, blocked: bool +) -> None: + recorder: Final = RecordingLogger() + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deferred", + function_id="deferred", + dynamic_async_success_callbacks=[recorder], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger) + await drain_logging() + assert "async_log_success_event" not in recorder.names + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + await drain_logging() + events: Final = tuple(event for event in recorder.events if event.name == "async_log_success_event") + assert len(events) == int(not blocked) + if events: + assert events[0].response is response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [RuntimeError("native enqueue failed"), asyncio.CancelledError("cancelled")]) +async def test_deferred_release_handles_enqueue_failure_once_without_replay( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: BaseException +) -> None: + import inspect + + from litellm.litellm_core_utils import logging_worker + + attempts: Final[list[Coroutine[object, object, object]]] = [] + diagnostics: Final = [] + + class FailingWorker: + def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: + attempts.append(coroutine) + raise failure + + recorder: Final = RecordingLogger() + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="release-failure", + function_id="release-failure", + dynamic_async_success_callbacks=[recorder], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(ocr_server, litellm_logging_obj=logger) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", FailingWorker()) + monkeypatch.setattr(sys, "unraisablehook", lambda event: diagnostics.append(event.exc_value)) + + if isinstance(failure, asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError, match="cancelled") as caught: + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + assert caught.value is failure + assert diagnostics == [] + else: + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + assert diagnostics == [failure] + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + + assert len(attempts) == 1 + assert inspect.getcoroutinestate(attempts[0]) == inspect.CORO_CLOSED + assert response.pages[0].markdown == "native OCR response" + assert len(ocr_server.requests) == 1 + assert not any("success" in name or "failure" in name for name in recorder.names) + + +@pytest.mark.asyncio +async def test_abandoned_deferred_logging_is_collectable(ocr_server: RecordingServer) -> None: + async def invoke(): + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="abandoned", + function_id="abandoned", + ) + logger._defer_async_logging = True + await call_aocr(ocr_server, litellm_logging_obj=logger) + return weakref.ref(logger) + + reference: Final = await invoke() + await drain_logging() + gc.collect() + assert reference() is None + + +def test_sync_success_uses_executor_and_copied_caller_context(ocr_server: RecordingServer) -> None: + context: Final = ContextVar("sync-lifecycle", default="missing") + context.set("caller") + thread: Final = threading.current_thread() + finished: Final = threading.Event() + observations: Final = [] + + class Observe(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append((threading.current_thread(), context.get(), response_obj)) + finished.set() + + response: Final = call_ocr(ocr_server, callbacks=[Observe()]) + assert finished.wait(5) + assert observations[0][0] is not thread + assert observations[0][1] == "caller" + assert observations[0][2] is response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_invalid_response_runs_post_call_before_failure(ocr_server: RecordingServer, asynchronous: bool) -> None: + ocr_server.enqueue(ResponseSpec(body={"pages": "invalid"})) + events: Final = [] + + class Observe(Logging): + def pre_call(self, *args, **kwargs): + events.append("pre") + return super().pre_call(*args, **kwargs) + + def post_call(self, *args, **kwargs): + events.append(("post", kwargs["original_response"])) + return super().post_call(*args, **kwargs) + + def success_handler(self, *args, **kwargs): + events.append("success") + + def failure_handler(self, exception, *args, **kwargs): + events.append(("failure", exception)) + + async def async_failure_handler(self, exception, *args, **kwargs): + events.append(("async_failure", exception)) + + logger: Final = Observe( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr" if asynchronous else "ocr", + start_time=datetime.datetime.now(), + litellm_call_id="invalid", + function_id="invalid", + ) + with pytest.raises(litellm.APIConnectionError) as caught: + await call_aocr(ocr_server, litellm_logging_obj=logger) if asynchronous else call_ocr( + ocr_server, litellm_logging_obj=logger + ) + assert events[0] == "pre" + assert events[1] == ("post", '{"pages": "invalid"}') + assert events[2] == ("failure", caught.value) + if asynchronous: + assert events[3] == ("async_failure", caught.value) + assert "success" not in events + + +@pytest.mark.asyncio +async def test_failing_terminal_handler_preserves_public_failure_and_runs_async_handler( + ocr_server: RecordingServer, +) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider failure"}, status=500)) + failures: Final = [] + + class BrokenHandler(Logging): + def failure_handler(self, exception, *args, **kwargs): + failures.append(exception) + raise RuntimeError("handler failed") + + async def async_failure_handler(self, exception, *args, **kwargs): + failures.append(exception) + + logger: Final = BrokenHandler( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="broken", + function_id="broken", + ) + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, litellm_logging_obj=logger) + assert failures == [caught.value, caught.value] + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_nested_native_calls_preserve_context_and_dispatch_each_outcome(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 2 + recorder: Final = RecordingLogger() + outcomes: Final = [] + + class Nested(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + if kwargs.get("litellm_call_id") == "outer": + outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="inner")) + + litellm.callbacks.append(Nested()) + outcomes.append(await call_aocr(ocr_server, callbacks=[recorder], litellm_call_id="outer")) + events: Final = await recorder.wait_for_async("async_log_success_event", count=2) + assert [event.kwargs["litellm_call_id"] for event in events] == ["inner", "outer"] + assert events[0].response is outcomes[0] + assert events[1].response is outcomes[1] + assert len(ocr_server.requests) == 2 + + +def test_sync_pre_call_can_make_nested_native_request(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 2 + observed: Final = [] + + class Nested(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + if kwargs["litellm_call_id"] == "outer-sync": + observed.append(call_ocr(ocr_server, litellm_call_id="inner-sync")) + + response: Final = call_ocr(ocr_server, callbacks=[Nested()], litellm_call_id="outer-sync") + assert observed[0].pages[0].markdown == response.pages[0].markdown + assert len(ocr_server.requests) == 2 + + +@pytest.mark.asyncio +async def test_retained_argument_aliases_and_body_roots_survive_envelope_replacement( + ocr_server: RecordingServer, +) -> None: + pages: Final = [0] + document: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} + opaque: Final = object() + observed: Final = [] + + class Observe(Logging): + def pre_call(self, input, api_key, additional_args): + body: Final = additional_args["complete_input_dict"] + headers: Final = additional_args["headers"] + observed.append((body["document"] is document, body["pages"] is pages)) + pages.append(2) + headers["x-retained"] = "yes" + additional_args["complete_input_dict"] = {"discarded": True} + additional_args["headers"] = {} + observed.append((body, headers)) + + def post_call(self, original_response, additional_args): + observed.append( + (additional_args["complete_input_dict"] is observed[2][0], additional_args["headers"] is observed[2][1]) + ) + + class Deployment(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + observed.append(("model" in kwargs, "document" in kwargs, kwargs["opaque"] is opaque)) + + litellm.callbacks.append(Deployment()) + logger: Final = Observe( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="roots", + function_id="roots", + ) + response: Final = await litellm.aocr( + "mistral/mistral-ocr-latest", + document, + api_key="test-key", + api_base=ocr_server.base_url, + pages=pages, + opaque=opaque, + litellm_logging_obj=logger, + ) + assert response.pages[0].markdown == "native OCR response" + assert observed[0] == (False, False, True) + assert observed[1] == (True, True) + assert observed[3] == (True, True) + assert ocr_server.requests[0].body["pages"] == [0, 2] + assert ocr_server.requests[0].headers["x-retained"] == "yes" + + +def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None: + from litellm.ocr.main import _public_request + from litellm.rust_bridge import _native + + ocr_server.expected_requests = 0 + effects: Final = [] + + class File: + def read(self): + effects.append("read") + return b"abc" + + def create(): + file: Final = File() + kwargs: Final = {"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": file}} + coroutine: Final = _native._ocr_lifecycle(_public_request("aocr", (), kwargs), (), kwargs, True) + file.owner = coroutine + coroutine.close() + return weakref.ref(file) + + reference: Final = create() + gc.collect() + assert reference() is None + assert effects == [] + + +@pytest.mark.asyncio +async def test_file_read_happens_after_deployment_hook_in_caller_task(ocr_server: RecordingServer) -> None: + effects: Final = [] + caller: Final = asyncio.current_task() + + class File: + def read(self): + effects.append(("read", asyncio.current_task())) + return b"abc" + + class Deployment(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + await asyncio.sleep(0) + effects.append(("hook", asyncio.current_task())) + + litellm.callbacks.append(Deployment()) + await call_aocr(ocr_server, document={"type": "file", "file": File()}) + assert effects == [("hook", caller), ("read", caller)] + + +@pytest.mark.asyncio +async def test_failure_callbacks_continue_within_both_families(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "failed"}, status=500)) + observed: Final = [] + + class Broken(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("broken-sync", kwargs["exception"])) + raise RuntimeError("sync observer") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("broken-async", kwargs["exception"])) + raise RuntimeError("async observer") + + class Following(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("following-sync", kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("following-async", kwargs["exception"])) + + with pytest.raises(litellm.InternalServerError) as caught: + await call_aocr(ocr_server, callbacks=[Broken(), Following()]) + assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"] + assert all(error is caught.value for _, error in observed) + + +@pytest.mark.asyncio +async def test_cancelling_native_transport_closes_connection_before_return() -> None: + received: Final = asyncio.Event() + disconnected: Final = asyncio.Event() + + async def provider(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + headers: Final = await reader.readuntil(b"\r\n\r\n") + length: Final = next( + int(line.split(b":", 1)[1]) + for line in headers.split(b"\r\n") + if line.lower().startswith(b"content-length:") + ) + await reader.readexactly(length) + received.set() + assert await reader.read() == b"" + disconnected.set() + writer.close() + await writer.wait_closed() + + server: Final = await asyncio.start_server(provider, "127.0.0.1", 0) + async with server: + port: Final = server.sockets[0].getsockname()[1] + task: Final = asyncio.create_task( + litellm.aocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + api_key="test-key", + api_base=f"http://127.0.0.1:{port}", + ) + ) + await asyncio.wait_for(received.wait(), 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(disconnected.wait(), 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["reducto/parse-v3", "reducto/parse-legacy"]) +async def test_reducto_lifecycle_retains_upload_parse_and_post_call_boundaries( + ocr_server: RecordingServer, model: str +) -> None: + ocr_server.expected_requests = 2 + ocr_server.enqueue(ResponseSpec(body={"file_id": "reducto://uploaded.pdf"})) + ocr_server.enqueue(ResponseSpec(body={"result": {"chunks": [{"content": "parsed"}]}})) + boundaries: Final = [] + recorder: Final = RecordingLogger() + + class Observe(Logging): + def post_call(self, *args, **kwargs): + boundaries.append(tuple(request.path for request in ocr_server.requests)) + return super().post_call(*args, **kwargs) + + logger: Final = Observe( + model=model, + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="upload", + function_id="upload", + dynamic_async_success_callbacks=[recorder], + ) + response: Final = await call_aocr(ocr_server, model=model, litellm_logging_obj=logger) + events: Final = await recorder.wait_for_async("async_log_success_event") + assert boundaries == [("/upload", "/parse")] + assert b"abc" in ocr_server.requests[0].raw_body + assert "multipart/form-data" in ocr_server.requests[0].headers["content-type"] + assert ocr_server.requests[1].body["input" if model.endswith("v3") else "document_url"] == "reducto://uploaded.pdf" + assert response.pages[0].markdown == "parsed" + assert events[0].response is response + + +@pytest.mark.asyncio +async def test_document_intelligence_post_call_observes_submission_and_final_result( + ocr_server: RecordingServer, +) -> None: + ocr_server.expected_requests = 2 + ocr_server.enqueue( + ResponseSpec( + body={"status": "running"}, + status=202, + headers={"Operation-Location": f"{ocr_server.base_url}/operations/1", "Retry-After": "0"}, + ) + ) + ocr_server.enqueue(ResponseSpec(body={"status": "succeeded", "analyzeResult": {"pages": []}})) + boundaries: Final = [] + + class Observe(Logging): + def post_call(self, *args, **kwargs): + boundaries.append((tuple(request.method for request in ocr_server.requests), kwargs["original_response"])) + return super().post_call(*args, **kwargs) + + logger: Final = Observe( + model="azure_ai/doc-intelligence/prebuilt-read", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="poll", + function_id="poll", + ) + response: Final = await call_aocr( + ocr_server, model="azure_ai/doc-intelligence/prebuilt-read", litellm_logging_obj=logger + ) + assert [methods for methods, _ in boundaries] == [("POST",), ("POST", "GET")] + assert json.loads(boundaries[0][1])["status"] == "running" + assert json.loads(boundaries[1][1])["status"] == "succeeded" + assert [request.method for request in ocr_server.requests] == ["POST", "GET"] + assert ocr_server.requests[1].path == "/operations/1" + assert response.pages == [] + + +@pytest.mark.asyncio +async def test_vertex_deepseek_public_lifecycle_normalizes_before_success(ocr_server: RecordingServer) -> None: + ocr_server.enqueue( + ResponseSpec(body={"choices": [{"message": {"content": "recognized"}}], "usage": {"prompt_tokens": 1}}) + ) + recorder: Final = RecordingLogger() + response: Final = await call_aocr( + ocr_server, + model="vertex_ai/deepseek-ocr-maas", + document={"type": "document_url", "document_url": "gs://bucket/document.pdf"}, + vertex_project="project-1", + vertex_location="europe-west4", + callbacks=[recorder], + ) + events: Final = await recorder.wait_for_async("async_log_success_event") + assert response.pages[0].markdown == "recognized" + assert events[0].response is response + assert ( + ocr_server.requests[0].path + == "/v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("limit", ["budget", "retries"]) +async def test_shared_call_limits_still_reject_before_reading_ocr_file( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool, limit: str +) -> None: + ocr_server.expected_requests = 0 + reads: Final = [] + + class File: + def read(self): + reads.append("read") + return b"abc" + + monkeypatch.setattr(litellm, "max_budget", 1 if limit == "budget" else None) + monkeypatch.setattr(litellm, "_current_cost", 2) + monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None) + expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"previous_models": ["earlier"]}} + with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"): + await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) + assert reads == [] + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("extra_bytes", [0, 1]) +async def test_response_limit_is_enforced_at_the_public_boundary( + ocr_server: RecordingServer, asynchronous: bool, extra_bytes: int +) -> None: + limit: Final = len(json.dumps(OCR_RESPONSE).encode()) - extra_bytes + if extra_bytes: + with pytest.raises(litellm.APIConnectionError, match="OCR response exceeds the size limit"): + await call_aocr(ocr_server, max_response_bytes=limit) if asynchronous else call_ocr( + ocr_server, max_response_bytes=limit + ) + else: + response: Final = ( + await call_aocr(ocr_server, max_response_bytes=limit) + if asynchronous + else call_ocr(ocr_server, max_response_bytes=limit) + ) + assert response.pages[0].markdown == "native OCR response" + assert len(ocr_server.requests) == 1 + body: Final = ocr_server.requests[0].body + assert isinstance(body, dict) + assert "max_response_bytes" not in body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("failure", [False, True]) +async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + failure: bool, + created_loggers: list[Logging], +) -> None: + from litellm import utils + from litellm.litellm_core_utils import litellm_logging, logging_worker + + class DispatchProbe: + deployments = 0 + submissions = 0 + enqueues = 0 + + def deployment(self, *args: object, **kwargs: object) -> None: + self.deployments += 1 + + def submit(self, *args: object, **kwargs: object) -> None: + self.submissions += 1 + + def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: + self.enqueues += 1 + coroutine.close() + + probe: Final = DispatchProbe() + for name in ( + "async_pre_call_deployment_hook", + "async_post_call_success_deployment_hook", + "async_post_call_failure_deployment_hook", + ): + monkeypatch.setattr(utils, name, probe.deployment) + monkeypatch.setattr(litellm_logging, "executor", probe) + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) + if failure: + ocr_server.enqueue(ResponseSpec(body={"message": "provider failed"}, status=500)) + trace_id_var.set("callback-free-parent") + arguments: Final = {"litellm_trace_id": "callback-free-call", "litellm_call_id": "callback-free-id"} + if failure: + with pytest.raises(litellm.InternalServerError): + await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) + else: + response: Final = ( + await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert response._hidden_params["litellm_call_id"] == "callback-free-id" + assert response._hidden_params["response_cost"] is not None + assert response._hidden_params["_response_ms"] > 0 + assert trace_id_var.get() == "callback-free-parent" + assert probe.deployments == probe.submissions == probe.enqueues == 0 + assert len(created_loggers) == 1 + logger: Final = created_loggers[0] + assert not hasattr(logger, "_native_pending_logging") + assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] + assert "standard_logging_object" not in logger.model_call_details + assert ( + "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None + ) + assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) + assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "registration", ["success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"] +) +async def test_terminal_registration_added_during_http_is_observed( + ocr_server: RecordingServer, registration: str +) -> None: + failure: Final = "failure" in registration + observer: Final = RecordingLogger() + ocr_server.enqueue( + ResponseSpec( + body={"message": "provider failed"} if failure else OCR_RESPONSE, status=500 if failure else 200, delay=0.1 + ) + ) + task: Final = asyncio.create_task( + asyncio.to_thread(call_ocr, ocr_server) if registration == "success_callback" else call_aocr(ocr_server) + ) + await ocr_server.wait_for_requests(1) + getattr(litellm, registration).append(observer) + if failure: + with pytest.raises(litellm.InternalServerError): + await task + else: + await task + event: Final = ("async_" if registration.startswith("_async") else "") + ( + "log_failure_event" if failure else "log_success_event" + ) + await observer.wait_for_async(event) + assert event in observer.names + + +@pytest.fixture +def created_loggers(monkeypatch: pytest.MonkeyPatch) -> list[Logging]: + from litellm import utils + + original_setup: Final = utils.function_setup + loggers: Final[list[Logging]] = [] + + def setup( + call_type: str, + rules: utils.Rules, + start: datetime.datetime, + *args: object, + is_async_call: bool = True, + **kwargs: object, + ) -> tuple[Logging, dict[str, object]]: + logger, prepared = original_setup(call_type, rules, start, *args, is_async_call=is_async_call, **kwargs) + assert isinstance(logger, Logging) + setattr(logger, "_defer_async_logging", True) + loggers.append(logger) + return logger, prepared + + monkeypatch.setattr(utils, "function_setup", setup) + return loggers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("consumer", ["logger_fn", "raw_global", "request_debug"]) +async def test_explicit_logging_consumers_keep_request_and_response_payloads( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging], consumer: str +) -> None: + snapshots: Final[list[dict[str, object]]] = [] + if consumer == "raw_global": + monkeypatch.setattr(litellm, "log_raw_request_response", True) + arguments: Final = { + "logger_fn": {"logger_fn": lambda details: snapshots.append(dict(details))}, + "raw_global": {}, + "request_debug": {"litellm_request_debug": True}, + }[consumer] + response: Final = await call_aocr(ocr_server, **arguments) + details: Final = created_loggers[0].model_call_details + assert details["additional_args"]["complete_input_dict"]["model"] == "mistral-ocr-latest" + assert json.loads(details["original_response"])["pages"][0]["markdown"] == response.pages[0].markdown + if consumer.startswith("raw_"): + assert details["raw_request_typed_dict"]["raw_request_body"]["model"] == "mistral-ocr-latest" + if consumer == "logger_fn": + assert [item["log_event_type"] for item in snapshots] == ["pre_api_call", "post_api_call"] + + +@pytest.mark.asyncio +async def test_registration_removed_before_deferred_release_skips_queue( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, created_loggers: list[Logging] +) -> None: + from litellm.litellm_core_utils import logging_worker + + class QueueProbe: + enqueues = 0 + + def ensure_initialized_and_enqueue(self, coroutine: Coroutine[object, object, object]) -> None: + self.enqueues += 1 + coroutine.close() + + observer: Final = RecordingLogger() + litellm._async_success_callback.append(observer) + await call_aocr(ocr_server) + logger: Final = created_loggers[0] + assert hasattr(logger, "_native_pending_logging") + litellm._async_success_callback.clear() + probe: Final = QueueProbe() + monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, False) + assert probe.enqueues == 0 + assert not observer.names + assert logger.model_call_details["response_cost"] is not None diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index d241fe08fc8..4f4b39fa6c6 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -1,3 +1,4 @@ +from pathlib import Path from typing import Final import pytest @@ -5,13 +6,13 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, call_native_aocr, call_native_ocr, ) -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -79,6 +80,22 @@ def test_native_ocr_prepares_file_document_like_python(ocr_server: RecordingServ } +def test_native_ocr_reads_sdk_path_input(ocr_server: RecordingServer, tmp_path: Path) -> None: + document_path: Final = tmp_path / "document.pdf" + document_path.write_bytes(b"%PDF-1.4") + + response: Final = call_native_ocr( + ocr_server, + document={"type": "file", "file": document_path}, + ) + + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].body["document"] == { + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", + } + + def test_native_ocr_sends_pages_and_image_options(ocr_server: RecordingServer) -> None: call_native_ocr(ocr_server, pages=[0, 2], include_image_base64=True) @@ -149,7 +166,7 @@ def test_native_ocr_normalizes_provider_response_model_and_usage(ocr_server: Rec assert response.usage_info.pages_processed == 1 -def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: RecordingServer) -> None: +def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: RecordingServer) -> None: ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) with pytest.raises(litellm.BadRequestError) as caught: @@ -158,13 +175,23 @@ def test_native_ocr_maps_provider_400_without_exposing_response_body(ocr_server: assert caught.value.status_code == 400 assert caught.value.model == "mistral-ocr-latest" assert caught.value.llm_provider == "mistral" - assert "invalid OCR request" not in str(caught.value) + assert "invalid OCR request" in str(caught.value) -def test_native_ocr_raises_transport_error_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: +def test_native_ocr_rejects_unknown_response_format_before_provider_request(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 0 + + with pytest.raises(litellm.BadRequestError, match="Invalid `req_format`"): + call_native_ocr(ocr_server, req_format="raw") + + assert ocr_server.requests == [] + + +def test_ocr_raises_public_timeout_when_request_exceeds_timeout(ocr_server: RecordingServer) -> None: + litellm.rust(True) ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) - with pytest.raises(RuntimeError, match="OCR transport failed"): + with pytest.raises(litellm.Timeout): call_native_ocr(ocr_server, timeout=0.01) assert len(ocr_server.requests) == 1 @@ -301,13 +328,10 @@ async def test_native_azure_ocr_token_provider_failure_prevents_pre_call_callbac @pytest.mark.parametrize( "configuration", - [ - {"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}, - {"model": "azure_ai/doc-intelligence/prebuilt-read"}, - ], - ids=["oidc-assertion", "document-intelligence-model"], + [{"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}], + ids=["invalid-oidc-assertion"], ) -def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_callbacks( +def test_public_azure_ocr_maps_invalid_oidc_configuration_before_token_or_request( ocr_server: RecordingServer, isolated_azure_auth: None, configuration: dict[str, object], @@ -327,10 +351,10 @@ def test_native_azure_ocr_rejects_unsupported_configuration_before_token_or_call "callbacks": [recorder], **configuration, } - with pytest.raises(NotImplementedError): + with pytest.raises(litellm.APIConnectionError): call_native_ocr(ocr_server, **arguments) assert calls == [] - assert recorder.events == () + assert "log_pre_api_call" not in recorder.names assert ocr_server.requests == [] @@ -432,3 +456,170 @@ async def test_native_azure_ocr_rejects_coroutine_returned_by_sync_token_provide coroutine.close() assert calls == [] assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "override, expected_key", + [ + ({}, "credential-key"), + ({"api_key": "explicit-key"}, "explicit-key"), + ({"api_key": None}, "environment-key"), + ], + ids=["inherit", "explicit", "explicit-none"], +) +async def test_native_ocr_inherits_named_credentials_without_overwriting_arguments( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + override: dict[str, object], + expected_key: str, +) -> None: + from litellm.models.credentials import CredentialItem + + pages: Final = [0] + opaque: Final = object() + monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem(credential_name="other", credential_info={}, credential_values={"api_key": "wrong-key"}), + CredentialItem( + credential_name="ocr-test", + credential_info={}, + credential_values={ + "api_key": "credential-key", + "api_base": ocr_server.base_url, + "pages": pages, + "opaque": opaque, + }, + ), + CredentialItem(credential_name="ocr-test", credential_info={}, credential_values={"api_key": "later-key"}), + ], + ) + + class Observer(RecordingLogger): + def log_pre_api_call(self, model, messages, kwargs): + super().log_pre_api_call(model, messages, kwargs) + pages.append(2) + + arguments: Final = { + "model": "mistral/mistral-ocr-latest", + "document": OCR_DOCUMENT, + "litellm_credential_name": "ocr-test", + "callbacks": [Observer()], + **override, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].headers["authorization"] == f"Bearer {expected_key}" + assert ocr_server.requests[0].body["pages"] == [0, 2] + + +@pytest.mark.parametrize("source", ["sdk", "proxy"]) +@pytest.mark.parametrize( + "filename,mime", [("scan.PNG", "image/png"), ("document.pdf", "application/pdf"), ("note.txt", "text/plain")] +) +def test_ocr_file_helpers_use_native_document_preparation(source: str, filename: str, mime: str) -> None: + from io import BytesIO + + from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type + from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload + + file: Final = BytesIO(b"abc") + file.name = filename + document: Final = ( + convert_file_document_to_url_document({"type": "file", "file": file}) + if source == "sdk" + else _build_document_from_upload(b"abc", filename, "application/octet-stream; charset=utf-8") + ) + field: Final = "image_url" if mime.startswith("image/") else "document_url" + assert get_mime_type(filename) == mime + assert document == {"type": field, field: f"data:{mime};base64,YWJj"} + + +@pytest.mark.parametrize("attribute", ["read", "name"]) +def test_native_file_preparation_preserves_property_errors(attribute: str) -> None: + from litellm.ocr.input import convert_file_document_to_url_document + + failure: Final = LookupError("file property failed") + + class File: + def __getattribute__(self, name: str): + if name == attribute: + raise failure + return super().__getattribute__(name) + + def read(self): + return b"abc" + + with pytest.raises(LookupError) as caught: + convert_file_document_to_url_document({"type": "file", "file": File()}) + assert caught.value is failure + + +@pytest.mark.parametrize("kind", ["bytes", "path", "reader"]) +def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Path) -> None: + from litellm.ocr.input import FileDocument, convert_file_document_to_url_document, get_max_file_bytes + + limit: Final = get_max_file_bytes() + path: Final = tmp_path / "large.pdf" + with path.open("wb") as stream: + stream.truncate(limit + 1) + + class Reader: + def read(self) -> bytes: + return b"a" * (limit + 1) + + document: Final[FileDocument] = { + "type": "file", + "file": path if kind == "path" else Reader() if kind == "reader" else b"a" * (limit + 1), + } + with pytest.raises(ValueError, match="exceeds the size limit"): + convert_file_document_to_url_document(document) + + +@pytest.mark.parametrize("kind", ["str", "path", "reader"]) +def test_native_upload_binding_rejects_filesystem_inputs(kind: str, tmp_path: Path) -> None: + from io import BytesIO + from typing import cast # noqa: TID251 # deliberately invalid inputs exercise the native runtime boundary + + from litellm.ocr.input import convert_upload_to_url_document + + path: Final = tmp_path / "secret.pdf" + path.write_bytes(b"server secret") + source: Final = str(path) if kind == "str" else path if kind == "path" else BytesIO(b"abc") + with pytest.raises(TypeError): + convert_upload_to_url_document(cast(bytes, source), "document.pdf", None) + + +@pytest.mark.parametrize("extra_bytes", [0, 1]) +def test_native_upload_enforces_file_size_limit(extra_bytes: int) -> None: + import base64 + + from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes + + content: Final = b"a" * (get_max_file_bytes() + extra_bytes) + if extra_bytes: + with pytest.raises(ValueError, match="exceeds the size limit"): + convert_upload_to_url_document(content, "scan.pdf", None) + return + document: Final = convert_upload_to_url_document(content, "scan.pdf", None) + assert document["type"] == "document_url" + assert base64.b64decode(document["document_url"].split(",", 1)[1]) == content + + +def test_native_file_preparation_preserves_reader_exception() -> None: + from litellm.ocr.input import convert_file_document_to_url_document + + failure: Final = RuntimeError("reader failed") + + class Reader: + def read(self) -> bytes: + raise failure + + with pytest.raises(RuntimeError) as caught: + convert_file_document_to_url_document({"type": "file", "file": Reader()}) + assert caught.value is failure diff --git a/tests/test_litellm_rust/support/callback_recorder.py b/tests/test_litellm_rust/support/callback_recorder.py index 6de011b1414..d3749ccc095 100644 --- a/tests/test_litellm_rust/support/callback_recorder.py +++ b/tests/test_litellm_rust/support/callback_recorder.py @@ -81,7 +81,7 @@ class RecordingLogger(CustomLogger): await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=timeout) return tuple(event for event in self.events if event.name == name) - def log_pre_api_call(self, model, _messages, kwargs): + def log_pre_api_call(self, model, messages, kwargs): self._record("log_pre_api_call", kwargs) def log_success_event(self, kwargs, response_obj, start_time, end_time): diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 5a9b9497c6e..228ed2cc454 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -58,7 +58,9 @@ def recording_service() -> Iterator[RecordingServer]: def _handle(self) -> None: content_length: Final = int(self.headers.get("Content-Length", "0")) raw_body: Final = self.rfile.read(content_length) if content_length else b"" - body: Final = json.loads(raw_body) if raw_body else None + body: Final = ( + json.loads(raw_body) if raw_body and self.headers.get_content_type() == "application/json" else None + ) requests.append( RecordedRequest( method=self.command, @@ -84,6 +86,7 @@ def recording_service() -> Iterator[RecordingServer]: pass do_POST = _handle + do_GET = _handle def log_message(self, format: str, *args: object) -> None: pass diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index d681d752ffd..7114e42a59e 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -2,7 +2,6 @@ from typing import Final import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse -from litellm.rust_bridge import ocr as native_ocr from tests.test_litellm_rust.support.recording_server import RecordingServer OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} @@ -36,11 +35,11 @@ async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return native_ocr.ocr(ocr_arguments(server, **kwargs)) + return call_ocr(server, **kwargs) async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse: - return await native_ocr.aocr(ocr_arguments(server, **kwargs)) + return await call_aocr(server, **kwargs) def request_body(kwargs: dict[str, object]) -> dict[str, object]: diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index ad1c8c652bb..e0e06d685b8 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -2,6 +2,7 @@ import json import threading from collections.abc import Generator from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from io import BytesIO from typing import Final import pytest @@ -99,6 +100,38 @@ def test_native_ocr_with_compiled_rust_extension( } +@pytest.mark.parametrize( + "file_input,mime_type,expected_type,expected_field,expected_uri", + [ + (b"abc", "application/pdf", "document_url", "document_url", "data:application/pdf;base64,YWJj"), + (BytesIO(b"abc"), "image/png", "image_url", "image_url", "data:image/png;base64,YWJj"), + ], +) +def test_native_lifecycle_core_encodes_python_file_input( + ocr_server, + file_input, + mime_type, + expected_type, + expected_field, + expected_uri, +): + server, requests = ocr_server + litellm.rust(True) + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "file", "file": file_input, "mime_type": mime_type}, + api_key="test-key", + api_base=f"http://127.0.0.1:{server.server_port}", + opaque_extension=object(), + ) + assert response.pages[0].markdown == "native OCR response" + assert requests[0]["body"]["document"] == { + "type": expected_type, + expected_field: expected_uri, + } + assert "opaque_extension" not in requests[0]["body"] + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) @pytest.mark.asyncio @@ -145,24 +178,20 @@ async def test_native_public_ocr_matches_python(model, asynchronous): server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) thread: Final = Thread(target=server.serve_forever, daemon=True) thread.start() - responses: Final = [] try: - for enabled in (False, True): - litellm.rust(enabled) - arguments: Final = { - "model": model, - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "pages": [0, 2], - "timeout": 3.0, - } - response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - responses.append(response.model_dump()) - assert len(calls) == 2 - assert calls[0] == calls[1] - for key in ("model", "pages", "object"): - assert responses[0][key] == responses[1][key] + litellm.rust(True) + arguments: Final = { + "model": model, + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "pages": [0, 2], + "timeout": 3.0, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + response_data: Final = response.model_dump() + assert len(calls) == 1 + assert response_data["object"] == "ocr" finally: server.shutdown() server.server_close() @@ -195,7 +224,7 @@ def test_native_ocr_rejects_invalid_input_before_network(ocr_server, custom_prov from litellm.rust_bridge import _native server, requests = ocr_server - with pytest.raises(ValueError, match=r"invalid (OCR request field|provider)|invalid request"): + with pytest.raises(ValueError, match="Document URL is required"): _native.ocr( model="mistral-ocr-latest", custom_llm_provider=custom_provider, @@ -224,7 +253,7 @@ async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, "num_retries": 0, } started = time.monotonic() - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.Timeout): await asyncio.wait_for( litellm.aocr(**arguments) if asynchronous else asyncio.to_thread(litellm.ocr, **arguments), timeout=3,