scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_govcloud_profiles_lit6421

Browse source
This commit is contained in:
mateo-berri 2026-09-04 09:28:33 -07:00
commit 46974fe46e
298 changed files with 13741 additions and 6630 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
.github/actions/setup-uv-with-retries/action.yml vendored
View file

@ -1,11 +1,7 @@
name: "Set up uv with retries"
description: >-
Install uv via astral-sh/setup-uv, retrying on transient failures. Even with
an exact pinned version, the action resolves the artifact URL by fetching
https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a
single request with no retry, timeout, or fallback, so one connection-level
network error ("fetch failed") fails the whole job before any test runs.
Retrying the full step covers the manifest fetch and the binary download.
Install uv via astral-sh/setup-uv, retrying the full setup step so manifest
resolution and binary downloads get fresh attempts after transient failures.
inputs:
version:
@ -18,7 +14,7 @@ runs:
- name: Set up uv (attempt 1)
id: attempt-1
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
@ -31,7 +27,7 @@ runs:
id: attempt-2
if: steps.attempt-1.outcome == 'failure'
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
@ -42,6 +38,6 @@ runs:
- name: Set up uv (attempt 3)
if: steps.attempt-2.outcome == 'failure'
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}

29
.github/workflows/_test-unit-base.yml vendored
View file

@ -57,9 +57,15 @@ permissions:
jobs:
run:
name: Run tests
name: ${{ matrix.python-version == '3.12' && 'Run tests' || format('Run tests (Python {0})', matrix.python-version) }}
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.job-timeout-minutes }}
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
env:
UV_PYTHON: ${{ matrix.python-version }}
permissions:
contents: read
pull-requests: read
@ -82,7 +88,7 @@ jobs:
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
python-version: ${{ matrix.python-version }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
@ -96,12 +102,10 @@ jobs:
timeout-minutes: 5
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
path: ${{ env.UV_CACHE_DIR }}
key: ${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
@ -113,6 +117,7 @@ jobs:
timeout-minutes: 8
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -134,13 +139,7 @@ jobs:
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
# coverage.py's sys.monitoring backend (PEP 669), the cheapest core it has.
# It is only the default from Python 3.14, and these shards run 3.12, so it
# has to be asked for. Coverage refuses it when branch measurement is on
# (`branch_right_left` needs > 3.14.0a5) and falls back to the slow core with
# a `no-sysmon` warning, so turning on `branch = true` here means giving this
# back until the runners move to 3.14.
COVERAGE_CORE: sysmon
COVERAGE_CORE: ${{ contains(fromJSON('["3.10", "3.11"]'), matrix.python-version) && 'ctrace' || 'sysmon' }}
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
@ -167,7 +166,7 @@ jobs:
fi
- name: Save coverage report
if: always() && steps.changes.outputs.decision != 'skip'
if: always() && matrix.python-version == '3.12' && steps.changes.outputs.decision != 'skip'
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}

1
.github/workflows/test-unit.yml vendored
View file

@ -96,7 +96,6 @@ jobs:
- shard: misc
artifact-name: misc
test-path: >-
tests/sdk_function_trace
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol

10
basedpyright-code-budget.json
View file

@ -18,7 +18,7 @@
"limit": 40
},
"reportDeprecated": {
"limit": 211
"limit": 209
},
"reportDuplicateImport": {
"limit": 19
@ -45,7 +45,7 @@
"limit": 24
},
"reportInvalidTypeForm": {
"limit": 34
"limit": 30
},
"reportInvalidTypeVarUse": {
"limit": 2
@ -105,13 +105,13 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38311
"limit": 38309
},
"reportUnknownParameterType": {
"limit": 19624
"limit": 19622
},
"reportUnknownVariableType": {
"limit": 29847
"limit": 29846
},
"reportUnnecessaryCast": {
"limit": 111

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -5,6 +5,7 @@ datasource client {
generator client {
provider = "prisma-client-py"
recursive_type_depth = -1
binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"]
}

5
litellm-rust/AGENTS.md
View file

@ -1,17 +1,18 @@
# AGENTS.md
litellm-rust has four 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.
litellm-rust has five 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-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-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate.
## Where a route lives

1
litellm-rust/CLAUDE.md
View file

@ -24,6 +24,7 @@ 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

33
litellm-rust/Cargo.lock generated
View file

@ -1412,8 +1412,8 @@ dependencies = [
"base64",
"futures-channel",
"futures-util",
"litellm-config",
"litellm-core",
"pyo3",
"reqwest",
"serde",
"serde_json",
@ -1425,6 +1425,16 @@ dependencies = [
"tracing",
]
[[package]]
name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-core",
"pyo3",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
@ -1435,14 +1445,17 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"base64",
"rand 0.8.7",
"reqwest",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1461,7 +1474,6 @@ dependencies = [
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1511,6 +1523,16 @@ version = "0.3.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]]
name = "mime_guess"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
dependencies = [
"mime",
"unicase",
]
[[package]]
name = "mio"
version = "1.2.2"
@ -1969,6 +1991,7 @@ dependencies = [
"hyper-util",
"js-sys",
"log",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -2736,6 +2759,12 @@ version = "1.20.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
[[package]]
name = "unicase"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-ident"
version = "1.0.24"

4
litellm-rust/Cargo.toml
View file

@ -1,6 +1,7 @@
[workspace]
members = [
"crates/core",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
"crates/python-bridge",
@ -17,6 +18,7 @@ repository = "https://github.com/BerriAI/litellm"
tracing = "0.1"
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
litellm-core = { path = "crates/core" }
litellm-config = { path = "crates/config" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }
axum = "0.7"
@ -24,7 +26,7 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }

4
litellm-rust/README.md
View file

@ -25,11 +25,12 @@ coverage and production evidence.
| Crate | Role |
|-------|------|
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
| litellm-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: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
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
@ -38,6 +39,7 @@ 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.

2
litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md
View file

@ -45,7 +45,7 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages`
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/<provider>/<route>/` 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 (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
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_<ROUTE>`.
## Checks before push

9
litellm-rust/crates/ai-gateway/AGENTS.md
View file

@ -9,19 +9,15 @@ 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<Router> + master_key
gil.rs # GIL-activity tracker (records Python acquisitions)
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
gil.rs # simple route (one file): router() + GET /health/gil
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
python/ # Python interop (feature: python-config) — load-time only
mod.rs, config.rs, AGENTS.md
```
## Rules
@ -53,5 +49,6 @@ proxy in a later phase. Health routes don't add the extractor (unauthenticated).
## Python interop
Anything that calls into Python lives in `python/` and is **load-time only** — see
`python/AGENTS.md`. The realtime data path never takes the GIL.
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.

2
litellm-rust/crates/ai-gateway/ARCHITECTURE.md
View file

@ -9,4 +9,6 @@ flowchart LR
C[client] <--> G[Rust ai-gateway<br/>LLM inference]
G <--> O[OpenAI realtime]
G -. spend tracking callback .-> P[litellm proxy]
F[litellm-config<br/>load-time only] --> G
F -. Python backend .-> P
```

6
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -16,6 +16,7 @@ required-features = ["server"]
[dependencies]
tracing.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-config.workspace = true
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
@ -31,14 +32,15 @@ subtle = { workspace = true, optional = true }
# sha2 hashes the master key into user_api_key_hash (matches the proxy's
# SHA-256 hash_token) so the plaintext credential never enters a log payload.
sha2 = { workspace = true, optional = true }
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
tower = { version = "0.5.3", features = ["util"], optional = true }
[features]
default = []
server = ["dep:axum", "dep:subtle", "dep:sha2"]
# Build the gateway's config from the proxy YAML via an embedded Python
# interpreter (links libpython; requires `litellm` importable at runtime).
python-config = ["dep:pyo3"]
python-config = ["litellm-config/python"]
trace-parity = ["server", "dep:tower", "litellm-core/observability"]
[dev-dependencies]
futures-channel = "0.3"

22
litellm-rust/crates/ai-gateway/README.md
View file

@ -6,25 +6,30 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
## Crates
`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route:
`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route:
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
| litellm-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. |
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
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.
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil`
- **Health:** `GET /health/readiness`, `GET /health/liveness`
- **Request logs:** POSTed to a LiteLLM proxy at `/v1/rust_control_plane/logs` (see [Request logging](#request-logging))
> **Realtime serving is pure Rust.** Python is used at **load time only** — to
> read the config once at boot. The realtime hot path never touches Python.
The former `/health/gil` route and its acquisition counter were removed. They
only observed the single startup config load and did not prove that every GIL
acquisition was instrumented
## Configuration (config.yaml)
The gateway loads its `model_list` from a **config.yaml**, the same as the
@ -43,9 +48,10 @@ model_list:
LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway
```
At boot the gateway calls into `litellm.proxy.read_model_list`, which reuses the
**real proxy config reader** (`ProxyConfig.get_config`). That means everything
the proxy supports in config.yaml works here too:
At boot `litellm-config` calls into `litellm.proxy.read_model_list` and returns
resolved deployments to the gateway, which constructs the router. The Python
backend still reuses the **real proxy config reader** (`ProxyConfig.get_config`),
so everything the proxy supports in config.yaml works here too:
- `include:` to merge in other config files,
- `os.environ/VAR` secret references (resolved via the secret manager, never
@ -82,8 +88,8 @@ stand-in built from the environment:
|---|---|---|
| `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). |
This mode links no libpython and needs no config file, but it only supports one
hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
The default workspace build links no libpython and needs no config file. This
fallback mode only supports one hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
stand-in only for the leanest possible build.
## Request logging

6
litellm-rust/crates/ai-gateway/config.yaml
View file

@ -1,8 +1,8 @@
# Sample realtime config for the LiteLLM Rust AI Gateway.
#
# The gateway loads this model_list at boot via the embedded python config
# reader (litellm.proxy.read_model_list), which reuses the proxy's own reader —
# so include:, os.environ/ secrets, and DB-stored models all work here too.
# litellm-config resolves this model_list at boot through the Python config
# reader (litellm.proxy.read_model_list), then the gateway builds its router.
# Includes, environment secrets, and database-stored models still work.
#
# Secrets are referenced (never inlined) via os.environ/. A real deploy can
# override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH).

58
litellm-rust/crates/ai-gateway/src/gil.rs
View file

@ -1,58 +0,0 @@
//! GIL-activity tracking.
//!
//! Every acquisition of the Python GIL is recorded here so the `/health/gil`
//! endpoint can report whether Python was touched recently. The design goal is
//! that the GIL is acquired **only at load time** (config read) and never on the
//! realtime hot path — polling this endpoint during traffic should show the
//! count holding steady and `acquired_last_30s` falling to `false`.
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
/// Window (seconds) for the "recently acquired" signal.
pub const RECENT_WINDOW_SECS: u64 = 30;
static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0);
/// Unix seconds of the last acquisition; `0` means "never".
static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0);
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
/// Record that the GIL was just acquired. Call immediately before taking the GIL.
///
/// Only invoked under the `python-config` feature; without it the gateway never
/// touches Python, so the recorder is unused (and the endpoint reports zero).
#[cfg_attr(not(feature = "python-config"), allow(dead_code))]
pub fn record_acquisition() {
GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed);
LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed);
}
/// Point-in-time view of GIL activity.
pub struct GilSnapshot {
pub total_acquisitions: u64,
pub seconds_since_last: Option<u64>,
pub acquired_last_30s: bool,
}
/// Read the current GIL-activity snapshot.
pub fn snapshot() -> GilSnapshot {
let total = GIL_ACQUISITIONS.load(Ordering::Relaxed);
let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed);
let seconds_since_last = if last == 0 {
None
} else {
Some(now_unix_secs().saturating_sub(last))
};
let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS);
GilSnapshot {
total_acquisitions: total,
seconds_since_last,
acquired_last_30s,
}
}

12
litellm-rust/crates/ai-gateway/src/lib.rs
View file

@ -10,29 +10,23 @@
//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
//! for the load-time config reader.
//! binary turns on.
pub mod audio_transcription;
mod client;
pub mod io;
pub mod ocr;
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
/// the `python-config` reader, so it is available without either feature.
pub mod gil;
#[cfg(feature = "server")]
pub mod auth;
#[cfg(feature = "server")]
pub mod routes;
#[cfg(feature = "server")]
pub mod state;
#[cfg(feature = "trace-parity")]
pub mod trace_parity;
mod constants;
pub mod integrations;
#[cfg(feature = "server")]
mod realtime;
#[cfg(feature = "python-config")]
pub mod python;

10
litellm-rust/crates/ai-gateway/src/main.rs
View file

@ -14,12 +14,12 @@ use std::sync::Arc;
use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key};
use litellm_ai_gateway::routes;
use litellm_ai_gateway::state::AppState;
#[cfg(feature = "python-config")]
use litellm_config::load_model_list;
use litellm_core::router::{Deployment, LiteLLMParams, Router};
use litellm_ai_gateway::integrations::custom_logger::CustomLogger;
use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger;
#[cfg(feature = "python-config")]
use litellm_ai_gateway::python;
/// Bind to localhost by default so the gateway is not a public, unauthenticated
/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`).
@ -124,10 +124,10 @@ fn resolve_port() -> u16 {
fn build_router() -> Router {
#[cfg(feature = "python-config")]
if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") {
match python::config::load_router_from_config(&config_path) {
Ok(router) => {
match load_model_list(std::path::Path::new(&config_path)) {
Ok(deployments) => {
eprintln!("loaded model_list from {config_path} via python config reader");
return router;
return Router::new(deployments);
}
Err(err) => {
eprintln!("config load failed ({err}); falling back to env deployment");

90
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
View file

@ -12,6 +12,7 @@ use litellm_core::providers::azure_ai::ocr::transformation::{
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
};
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::reducto::ocr::transformation as reducto;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
@ -39,6 +40,7 @@ pub(super) fn ocr_provider_config(
) -> Option<&'static dyn OcrProviderConfig> {
match provider {
"mistral" => Some(&MISTRAL_OCR_CONFIG),
"reducto" => reducto::config_for_model(model),
"azure_ai" if is_azure_document_intelligence_model(model) => {
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
}
@ -334,6 +336,7 @@ fn operation_status(response_json: &Value) -> Result<&str, Error> {
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn poll_document_intelligence(
operation_url: &str,
original_url: &str,
@ -392,9 +395,11 @@ pub(super) async fn poll_document_intelligence(
#[cfg(test)]
mod tests {
use super::*;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::json;
use super::*;
#[test]
fn blocks_private_and_metadata_ips() {
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
@ -438,4 +443,87 @@ mod tests {
assert_eq!(transformed, document);
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
let body = "Unauthorized";
assert_eq!(truncate_error_body(body), "Unauthorized");
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}
}

12
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
View file

@ -50,7 +50,11 @@ pub(crate) async fn execute_ocr_provider_call(
.await?;
return Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json());
}
@ -71,6 +75,10 @@ pub(crate) async fn execute_ocr_provider_call(
Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

90
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -1,11 +1,15 @@
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use litellm_core::providers::reducto::ocr::transformation::{
build_upload_request, extract_document_source, extract_upload_file_id,
};
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::common_utils::{convert_document_url_to_data_uri, string_headers};
use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::client::http_client;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
@ -89,22 +93,39 @@ impl OcrLifecycleHooks {
)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let document = if config.requires_data_uri_document() {
let is_reducto = custom_llm_provider == "reducto";
let document = if is_reducto {
let guarded_document = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
.await?;
upload_reducto_document(
&guarded_document,
request.api_base.as_deref(),
request.timeout,
&upstream_headers,
)
.await?
} else if config.requires_data_uri_document() {
convert_document_url_to_data_uri(request.document).await?
} else {
request.document
};
let optional_params = request.optional_params;
let body = config
.transform_ocr_request(&request.model, document, request.optional_params)?
.transform_ocr_request(&request.model, document, optional_params.clone())?
.data;
let body = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?;
let body = if is_reducto {
body
} else {
self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?
};
Ok(ProviderOcrRequest {
model,
config,
url,
body,
optional_params,
upstream_headers,
timeout: request.timeout,
})
@ -165,6 +186,63 @@ impl OcrLifecycleHooks {
}
}
async fn upload_reducto_document(
document: &Value,
api_base: Option<&str>,
timeout: Option<std::time::Duration>,
upstream_headers: &[(String, String)],
) -> Result<Value, Error> {
let source = extract_document_source(document)?;
let Some(authorization) = upstream_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
else {
return Err(Error::Auth(
"Reducto upload requires an Authorization header".to_string(),
));
};
let Some(upload) = build_upload_request(source, authorization, api_base) else {
return Ok(document.clone());
};
let part = reqwest::multipart::Part::bytes(upload.bytes)
.file_name(upload.file_name)
.mime_str(&upload.mime_type)
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let form = reqwest::multipart::Form::new().part("file", part);
let mut request_builder = http_client().post(upload.url).multipart(form);
for (name, value) in upstream_headers {
if !name.eq_ignore_ascii_case("content-type")
&& !name.eq_ignore_ascii_case("content-length")
{
request_builder = request_builder.header(name, value);
}
}
if let Some(timeout) = timeout {
request_builder = request_builder.timeout(timeout);
}
let response = request_builder
.send()
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let response_json: Value = serde_json::from_str(&body).map_err(|error| {
Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
})?;
let file_id = extract_upload_file_id(&response_json)?;
Ok(json!({"type": "document_url", "document_url": file_id}))
}
impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;

149
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
View file

@ -24,4 +24,151 @@ pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
}
#[cfg(test)]
mod tests;
mod tests {
use serde_json::{Map, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::{OcrRequest, ocr};
use crate::integrations::types::RequestMetadata;
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
#[tokio::test]
async fn reducto_file_upload_then_parse_maps_response() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request");
let upload_request = read_http_request(&mut upload_socket).await;
let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#;
let upload_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
upload_body.len(),
upload_body
);
upload_socket
.write_all(upload_response.as_bytes())
.await
.expect("writes upload response");
let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request");
let parse_request = read_http_request(&mut parse_socket).await;
let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#;
let parse_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
parse_body.len(),
parse_body
);
parse_socket
.write_all(parse_response.as_bytes())
.await
.expect("writes parse response");
(upload_request, parse_request)
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.api_key = None;
request.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer test-key")),
("x-trace-id".to_string(), json!("trace-1")),
]));
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.optional_params = Map::from_iter([
(
"formatting".to_string(),
json!({"table_output_format": "html"}),
),
("retrieval".to_string(), json!({"chunk_mode": "section"})),
("settings".to_string(), json!({"ocr_system": "standard"})),
]);
let response = ocr(request).await.expect("Reducto OCR succeeds");
assert_eq!(response["pages"].as_array().map(Vec::len), Some(3));
assert_eq!(
response["pages"][0]["markdown"],
"Page 1 block A\n\nPage 1 block B"
);
assert_eq!(response["pages"][1]["markdown"], "Page 2 block A");
assert_eq!(response["pages"][2]["markdown"], "Page 3 block A");
assert_eq!(response["usage_info"]["pages_processed"], 3);
assert_eq!(response["usage_info"]["credits"], 3);
assert_eq!(response["provider_native_response"]["job_id"], "job_123");
let (upload_request, parse_request) = server.await.expect("server task completes");
assert!(
upload_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(upload_request.contains("application/pdf"));
assert!(upload_request.contains("%PDF-1.4"));
assert!(upload_request.contains("x-trace-id: trace-1"));
assert!(
parse_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#));
assert!(parse_request.contains(r#""table_output_format":"html""#));
assert!(parse_request.contains(r#""chunk_mode":"section""#));
assert!(parse_request.contains(r#""ocr_system":"standard""#));
}
}

96
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
View file

@ -2,6 +2,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use serde_json::{Map, Value};
use super::common_utils::ocr_provider_config;
use super::hooks::OcrLifecycleHooks;
@ -28,17 +29,33 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let model = provider_info.model.to_string();
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
let config = ocr_provider_config(&custom_llm_provider, &model)
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()));
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
.and_then(|config| {
validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
Ok(config)
});
let optional_params = match &config {
Ok(config) => {
let supported = config.supported_ocr_params();
config.map_ocr_params(
let mut mapped = config.map_ocr_params(
&request
.optional_params
.into_iter()
.iter()
.filter(|(name, _)| supported.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
)
);
for name in [
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
] {
if let Some(value) = request.optional_params.get(name) {
mapped.insert(name.to_string(), value.clone());
}
}
mapped
}
Err(_) => request.optional_params,
};
@ -64,6 +81,26 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
}
}
fn validate_request_format(
config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
optional_params: &Map<String, Value>,
provider: &str,
) -> Result<(), litellm_core::Error> {
let Some(format) = optional_params.get("req_format") else {
return Ok(());
};
match format.as_str() {
Some("litellm") => Ok(()),
Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
"`req_format=native` is not supported for provider {provider}"
))),
_ => Err(litellm_core::Error::InvalidRequest(format!(
"Invalid `req_format`: {format}. Expected `litellm` or `native`"
))),
}
}
fn new_ocr_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
@ -73,3 +110,54 @@ fn new_ocr_call_id() -> String {
.unwrap_or(0);
format!("ocr-{timestamp}-{sequence}")
}
#[cfg(test)]
mod tests {
use litellm_core::error::Error;
use serde_json::{Map, json};
use super::{OcrRequest, prepare_ocr_call};
use crate::integrations::types::RequestMetadata;
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
fn request_with_format(format: &str) -> OcrRequest<'_> {
let mut request = base_ocr_request("mistral/mistral-ocr-latest");
request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]);
request
}
#[test]
fn native_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("native"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider"))
);
}
#[test]
fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("raw"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`"))
);
}
}

1
litellm-rust/crates/ai-gateway/src/ocr/types.rs
View file

@ -53,6 +53,7 @@ pub(crate) struct ProviderOcrRequest {
pub(crate) config: &'static dyn OcrProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) optional_params: Map<String, Value>,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
View file

@ -1,27 +0,0 @@
# ai-gateway/src/python — Python interop (load-time only)
Functions here embed the Python interpreter (pyo3) and take the GIL to call into
`litellm` (e.g. read the proxy `model_list`). Compiled only under the
`python-config` feature.
## Hard rule: non-hot-path functions only
Everything in this folder MUST run **at most once per process lifetime — at
startup / load time** (config read, warm-up). NEVER call into Python on the
request path:
- No GIL acquisition per request, per connection, or per realtime event.
- No Python call inside a route handler, the router's hot path, or any loop that
scales with traffic.
**Why:** the GIL serializes execution and would cap throughput; the realtime data
path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll
`GET /health/gil`, and `total_acquisitions` MUST stay flat under load.
## How to add one
Resolve whatever Python-derived data you need **once at boot** and hand the rest
of the gateway an owned, plain-Rust value (e.g. build a `Router` from the
resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()`
immediately before taking the GIL. If a function would need to run per request,
it does not belong here — move the work to Rust, or pre-resolve it at startup.

37
litellm-rust/crates/ai-gateway/src/python/config.rs
View file

@ -1,37 +0,0 @@
//! Build the router by calling the Python proxy config reader (load time only).
//!
//! Embeds the interpreter via pyo3 and calls
//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's
//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot**
//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python.
//!
//! Compiled only under the `python-config` feature.
use litellm_core::error::Error;
use litellm_core::router::{Deployment, Router};
use pyo3::prelude::*;
use crate::gil;
/// Load the router's `model_list` from `config_path` via the Python reader.
pub fn load_router_from_config(config_path: &str) -> Result<Router, Error> {
gil::record_acquisition();
Python::attach(|py| {
let model_list = py
.import("litellm.proxy.read_model_list")
.and_then(|module| module.getattr("read_model_list"))
.and_then(|reader| reader.call1((config_path,)))
.map_err(|err| Error::Routing(format!("read_model_list failed: {err}")))?;
let model_list_json: String = py
.import("json")
.and_then(|json| json.getattr("dumps"))
.and_then(|dumps| dumps.call1((model_list,)))
.and_then(|encoded| encoded.extract())
.map_err(|err| Error::Routing(format!("serializing model_list failed: {err}")))?;
let deployments: Vec<Deployment> = serde_json::from_str(&model_list_json)
.map_err(|err| Error::Routing(format!("parsing model_list failed: {err}")))?;
Ok(Router::new(deployments))
})
}

4
litellm-rust/crates/ai-gateway/src/python/mod.rs
View file

@ -1,4 +0,0 @@
//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path
//! only.** Compiled only under the `python-config` feature.
pub mod config;

2
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
View file

@ -13,7 +13,7 @@ private). This is the norm — don't split until it hurts.
pub fn router() -> Router<AppState> { Router::new().route(PATH, get(handle)) }
async fn handle(...) -> impl IntoResponse { ... }
```
`health.rs` and `gil.rs` are examples.
`health.rs` is the example.
## Split out `service` when there's real logic
When a route has business logic worth testing without axum, put it in a sibling

30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
View file

@ -1,30 +0,0 @@
//! `GET /health/gil` — poll to confirm Python is only touched at load time.
//! Simple-route template: a `router()` plus its handler, in one file.
use axum::routing::get;
use axum::{Json, Router};
use serde::Serialize;
use crate::gil;
use crate::state::AppState;
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new().route("/health/gil", get(status))
}
#[derive(Debug, Serialize)]
struct GilStatusResponse {
gil_acquired_last_30s: bool,
total_acquisitions: u64,
seconds_since_last: Option<u64>,
}
async fn status() -> Json<GilStatusResponse> {
let snapshot = gil::snapshot();
Json(GilStatusResponse {
gil_acquired_last_30s: snapshot.acquired_last_30s,
total_acquisitions: snapshot.total_acquisitions,
seconds_since_last: snapshot.seconds_since_last,
})
}

6
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -21,6 +21,12 @@ pub fn router() -> Router<AppState> {
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
}
#[tracing::instrument(
name = "messages_gateway_route",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
async fn handle(
_auth: RequireMasterKey,
State(state): State<AppState>,

6
litellm-rust/crates/ai-gateway/src/routes/messages/service.rs
View file

@ -12,6 +12,12 @@ pub(crate) enum MessagesResponse {
Stream(reqwest::Response),
}
#[tracing::instrument(
name = "messages_gateway_service",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub async fn run(
router: &Arc<Router>,
body: Value,

4
litellm-rust/crates/ai-gateway/src/routes/mod.rs
View file

@ -2,10 +2,9 @@
//!
//! **Template:** every route module exposes `pub fn router() -> Router<AppState>`
//! that mounts its own paths; [`app`] merges them. A trivial route is a single
//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with
//! file (`health.rs`); a non-trivial one is a folder (`realtime/`) with
//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md.
pub mod gil;
pub mod health;
pub mod messages;
pub mod realtime;
@ -19,7 +18,6 @@ use crate::state::AppState;
pub fn app(state: AppState) -> Router {
Router::new()
.merge(health::router())
.merge(gil::router())
.merge(messages::router())
.merge(realtime::router())
.merge(responses::router())

65
litellm-rust/crates/ai-gateway/src/trace_parity.rs Normal file
View file

@ -0,0 +1,65 @@
//! Harness-only in-process adapters. Never mounted as production routes.
use std::sync::Arc;
use axum::body::{Body, to_bytes};
use axum::http::header::{AUTHORIZATION, CONTENT_TYPE};
use axum::http::{Request, StatusCode};
use litellm_core::Error;
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde::Serialize;
use serde_json::Value;
use tower::ServiceExt;
use crate::io::realtime_pool::RealtimePool;
use crate::routes;
use crate::state::AppState;
#[derive(Debug, Serialize)]
pub struct GatewayResponse {
pub status: u16,
pub body: Value,
}
pub async fn messages_request(
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> Result<GatewayResponse, Error> {
let state = AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias,
litellm_params: LiteLLMParams {
model: provider_model,
api_key: Some("trace-provider-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: Some(Arc::from("trace-master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
};
let request = Request::builder()
.method("POST")
.uri("/v1/messages")
.header(AUTHORIZATION, "Bearer trace-master-key")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(body.to_string()))
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let response = routes::app(state)
.oneshot(request)
.await
.map_err(|error| match error {})?;
let status: StatusCode = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
let body = serde_json::from_slice(&bytes).map_err(|error| {
Error::InvalidResponse(format!("gateway returned invalid JSON: {error}"))
})?;
Ok(GatewayResponse {
status: status.as_u16(),
body,
})
}

196
litellm-rust/crates/ai-gateway/src/ocr/tests.rs → litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -1,23 +1,19 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use litellm_core::error::Error;
use litellm_core::http_utils::has_header;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body};
use super::{OcrRequest, ocr};
use crate::integrations::custom_guardrail::{
use litellm_ai_gateway::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
};
use crate::integrations::custom_logger::{
use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use crate::integrations::types::RequestMetadata;
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -136,6 +132,7 @@ struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
}
impl RecordingOcrGuardrail {
@ -144,6 +141,7 @@ impl RecordingOcrGuardrail {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
}
}
@ -152,6 +150,16 @@ impl RecordingOcrGuardrail {
hooks: vec![GuardrailEventHook::PreCall],
events: Mutex::new(Vec::new()),
block_pre_call: true,
block_during_call: false,
}
}
fn blocking_during_call() -> Self {
Self {
hooks: vec![GuardrailEventHook::DuringCall],
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: true,
}
}
@ -193,91 +201,95 @@ impl CustomGuardrail for RecordingOcrGuardrail {
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_moderation_hook");
if self.block_during_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["body"]["guarded_during"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
let body = "Unauthorized";
assert_eq!(truncate_error_body(body), "Unauthorized");
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
#[tokio::test]
async fn reducto_during_call_guardrail_blocks_before_upload() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let api_base = format!("http://{address}");
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call());
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.guardrails = vec![guardrail.clone()];
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
let error = ocr(request).await.expect_err("guardrail blocks upload");
assert!(matches!(error, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_moderation_hook"]);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "upload socket should not be touched");
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[tokio::test]
async fn reducto_upload_error_body_is_truncated() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts upload request");
let _request = read_http_request(&mut socket).await;
let body = "x".repeat(300);
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes upload response");
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
let error = ocr(request).await.expect_err("upload should fail");
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)"))
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn auth_header_detection_is_case_insensitive() {
let headers = vec![
("x-trace-id".to_string(), "trace-1".to_string()),
("authorization".to_string(), "Bearer sk-test".to_string()),
];
assert!(has_header(&headers, "authorization"));
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
assert!(has_header(&headers, "authorization"));
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
assert!(!has_header(&headers, "authorization"));
server.await.expect("server task completes");
}
#[tokio::test]
@ -595,21 +607,3 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() {
"{poll_request}"
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}

16
litellm-rust/crates/config/Cargo.toml Normal file
View file

@ -0,0 +1,16 @@
[package]
name = "litellm-config"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-core.workspace = true
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
serde_json.workspace = true
thiserror.workspace = true
[features]
default = []
python = ["dep:pyo3"]

11
litellm-rust/crates/config/src/error.rs Normal file
View file

@ -0,0 +1,11 @@
use thiserror::Error as ThisError;
#[derive(Debug, ThisError)]
pub enum Error {
#[error("read_model_list failed: {0}")]
PythonLoading(String),
#[error("serializing model_list failed: {0}")]
Serialization(String),
#[error("parsing model_list failed: {0}")]
ModelListParsing(#[source] serde_json::Error),
}

7
litellm-rust/crates/config/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
mod error;
#[cfg(feature = "python")]
mod python;
pub use error::Error;
#[cfg(feature = "python")]
pub use python::load_model_list;

76
litellm-rust/crates/config/src/python.rs Normal file
View file

@ -0,0 +1,76 @@
use std::path::Path;
use litellm_core::router::Deployment;
use pyo3::prelude::*;
use crate::Error;
pub fn load_model_list(config_path: &Path) -> Result<Vec<Deployment>, Error> {
Python::attach(|python| {
let model_list = python
.import("litellm.proxy.read_model_list")
.and_then(|module| module.getattr("read_model_list"))
.and_then(|reader| reader.call1((config_path.to_string_lossy().as_ref(),)))
.map_err(|error| Error::PythonLoading(error.to_string()))?;
let model_list_json = python
.import("json")
.and_then(|json| json.getattr("dumps"))
.and_then(|dumps| dumps.call1((model_list,)))
.and_then(|encoded| encoded.extract::<String>())
.map_err(|error| Error::Serialization(error.to_string()))?;
parse_model_list(&model_list_json)
})
}
fn parse_model_list(model_list_json: &str) -> Result<Vec<Deployment>, Error> {
serde_json::from_str(model_list_json).map_err(Error::ModelListParsing)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_resolved_model_list() {
let deployments = parse_model_list(
r#"[
{
"model_name": "realtime",
"litellm_params": {
"model": "openai/gpt-realtime",
"api_key": "resolved-secret",
"api_base": "https://api.example.test/v1"
}
},
{
"model_name": "without-optional-values",
"litellm_params": {"model": "openai/gpt-4.1"}
}
]"#,
)
.expect("resolved model list should parse");
assert_eq!(deployments.len(), 2);
assert_eq!(deployments[0].model_name, "realtime");
assert_eq!(
deployments[0].litellm_params.api_key.as_deref(),
Some("resolved-secret")
);
assert_eq!(
deployments[0].litellm_params.api_base.as_deref(),
Some("https://api.example.test/v1")
);
assert_eq!(deployments[1].litellm_params.api_key, None);
assert_eq!(deployments[1].litellm_params.api_base, None);
}
#[test]
fn malformed_model_list_returns_parsing_error() {
let error = parse_model_list(r#"[{"model_name":"missing-params"}]"#)
.expect_err("missing litellm_params should fail");
assert!(matches!(error, Error::ModelListParsing(_)));
}
}

5
litellm-rust/crates/core/Cargo.toml
View file

@ -6,12 +6,14 @@ license.workspace = true
repository.workspace = true
[dependencies]
base64.workspace = true
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
@ -30,6 +32,9 @@ bedrock-auth = [
"dep:aws-types",
"dep:aws-smithy-runtime-api",
]
observability = ["dep:tracing-subscriber"]
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber.workspace = true

2
litellm-rust/crates/core/src/constants.rs
View file

@ -41,3 +41,5 @@ pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

17
litellm-rust/crates/core/src/http_utils.rs
View file

@ -101,6 +101,23 @@ mod tests {
assert!(!has_header(&headers, "authorization"));
}
#[test]
fn auth_header_detection_is_case_insensitive() {
let headers = vec![
("x-trace-id".to_string(), "trace-1".to_string()),
("authorization".to_string(), "Bearer sk-test".to_string()),
];
assert!(has_header(&headers, "authorization"));
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
assert!(has_header(&headers, "authorization"));
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
assert!(!has_header(&headers, "authorization"));
}
#[test]
fn bearer_detection_requires_a_non_empty_token() {
assert!(has_bearer_auth(&[(

2
litellm-rust/crates/core/src/lib.rs
View file

@ -6,6 +6,8 @@ pub mod constants;
pub mod error;
pub mod http_utils;
pub mod messages;
#[cfg(any(feature = "observability", test))]
pub mod observability;
pub mod ocr;
pub mod providers;
pub mod realtime;

215
litellm-rust/crates/core/src/observability/function_trace.rs Normal file
View file

@ -0,0 +1,215 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde::Serialize;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Subscriber};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
use super::function_trace_filter;
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionTraceEvent {
pub id: usize,
pub parent_id: Option<usize>,
pub function: &'static str,
pub module_path: Option<&'static str>,
pub file: Option<&'static str>,
pub line: Option<u32>,
}
#[derive(Clone, Default)]
pub struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
span_events: Arc<Mutex<HashMap<Id, usize>>>,
}
impl FunctionTrace {
pub fn dispatcher(&self) -> Dispatch {
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(function_trace_filter()),
),
)
}
pub fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let parent_id = context.span(id).and_then(|span| {
let span_events = self
.trace
.span_events
.lock()
.unwrap_or_else(|error| error.into_inner());
span.scope()
.skip(1)
.find_map(|ancestor| span_events.get(&ancestor.id()).copied())
});
let mut events = self
.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner());
let event_id = events.len();
events.push(FunctionTraceEvent {
id: event_id,
parent_id,
function: attributes.metadata().name(),
module_path: attributes.metadata().module_path(),
file: attributes.metadata().file(),
line: attributes.metadata().line(),
});
self.trace
.span_events
.lock()
.unwrap_or_else(|error| error.into_inner())
.insert(id.clone(), event_id);
}
}
#[cfg(test)]
mod tests {
use crate::constants::FUNCTION_TRACE_TARGET;
use super::*;
fn event(
id: usize,
parent_id: Option<usize>,
function: &'static str,
) -> (usize, Option<usize>, &'static str) {
(id, parent_id, function)
}
fn structural_events(
events: &[FunctionTraceEvent],
) -> Vec<(usize, Option<usize>, &'static str)> {
events
.iter()
.map(|event| (event.id, event.parent_id, event.function))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn outer() {
tokio::task::yield_now().await;
inner().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn inner() {
tokio::task::yield_now().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn concurrent_parent() {
tokio::join!(inner(), inner());
}
#[tokio::test]
async fn concurrent_futures_keep_separate_traces_across_yields() {
use tracing::instrument::WithSubscriber;
let first = FunctionTrace::default();
let second = FunctionTrace::default();
let outside = FunctionTrace::default();
async {
tokio::join!(
outer().with_subscriber(first.dispatcher()),
inner().with_subscriber(second.dispatcher()),
);
inner().await;
}
.with_subscriber(outside.dispatcher())
.await;
assert_eq!(
structural_events(&first.events()),
vec![event(0, None, "outer"), event(1, Some(0), "inner")],
);
assert_eq!(
structural_events(&second.events()),
vec![event(0, None, "inner")],
);
assert_eq!(
structural_events(&outside.events()),
vec![event(0, None, "inner")],
);
}
#[tokio::test]
async fn concurrent_siblings_keep_the_same_parent() {
use tracing::instrument::WithSubscriber;
let trace = FunctionTrace::default();
concurrent_parent()
.with_subscriber(trace.dispatcher())
.await;
assert_eq!(
structural_events(&trace.events()),
vec![
event(0, None, "concurrent_parent"),
event(1, Some(0), "inner"),
event(2, Some(0), "inner"),
]
);
}
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
});
assert_eq!(
structural_events(&trace.events()),
vec![event(0, None, "same_name"), event(1, None, "same_name")]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
});
assert_eq!(
structural_events(&trace.events()),
vec![event(0, None, "outer"), event(1, Some(0), "inner")]
);
}
}

59
litellm-rust/crates/core/src/observability/mod.rs Normal file
View file

@ -0,0 +1,59 @@
use tracing::span::Id;
use tracing::{Level, Metadata, Subscriber};
use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::registry::LookupSpan;
use crate::constants::FUNCTION_TRACE_TARGET;
pub mod function_trace;
pub use function_trace::{FunctionTrace, FunctionTraceEvent};
pub fn function_trace_filter() -> FilterFn<impl Fn(&Metadata<'_>) -> bool> {
filter_fn(|metadata| {
metadata.is_span()
&& metadata.target() == FUNCTION_TRACE_TARGET
&& *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE)
}
pub fn span_depth<S>(context: &Context<'_, S>, id: &Id) -> usize
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use tracing::instrument::WithSubscriber;
use super::*;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn instrumented_with_literal_target() {}
#[tokio::test]
async fn literal_instrument_target_matches_filter_constant() {
assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace");
let trace = FunctionTrace::default();
instrumented_with_literal_target()
.with_subscriber(trace.dispatcher())
.await;
let events = trace.events();
assert_eq!(events.len(), 1);
assert_eq!(events[0].id, 0);
assert_eq!(events[0].parent_id, None);
assert_eq!(events[0].function, "instrumented_with_literal_target");
assert_eq!(events[0].module_path, Some(module_path!()));
assert_eq!(events[0].file, Some(file!()));
assert!(events[0].line.is_some());
}
}

9
litellm-rust/crates/core/src/ocr/transformation.rs
View file

@ -51,6 +51,15 @@ pub trait OcrProviderConfig: Sync {
response_json: Value,
) -> Result<OcrResponseData, Error>;
fn transform_ocr_response_with_params(
&self,
model: &str,
response_json: Value,
_optional_params: &Map<String, Value>,
) -> Result<OcrResponseData, Error> {
self.transform_ocr_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,

15
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -1,5 +1,5 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrRequestData {
@ -14,16 +14,25 @@ pub struct OcrResponseData {
pub document_annotation: Option<Value>,
pub usage_info: Option<Value>,
pub object: String,
pub extra_fields: Map<String, Value>,
pub provider_native_response: Option<Value>,
}
impl OcrResponseData {
pub fn into_json(self) -> Value {
serde_json::json!({
let mut response = serde_json::json!({
"pages": self.pages,
"model": self.model,
"document_annotation": self.document_annotation,
"usage_info": self.usage_info,
"object": self.object,
})
});
if let Value::Object(object) = &mut response {
object.extend(self.extra_fields);
if let Some(native_response) = self.provider_native_response {
object.insert("provider_native_response".to_string(), native_response);
}
}
response
}
}

887
litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs
View file

File diff suppressed because it is too large Load diff

2
litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs
View file

@ -134,6 +134,8 @@ impl OcrProviderConfig for MistralOcrConfig {
document_annotation,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}

1
litellm-rust/crates/core/src/providers/mod.rs
View file

@ -4,4 +4,5 @@ pub mod azure_ai;
pub mod bedrock;
pub mod mistral;
pub mod openai;
pub mod reducto;
pub mod vertex_ai;

1
litellm-rust/crates/core/src/providers/reducto/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod ocr;

4
litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub mod transformation;
#[cfg(test)]
mod tests;

202
litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs Normal file
View file

@ -0,0 +1,202 @@
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::transformation::*;
use crate::ocr::transformation::OcrProviderConfig;
#[fixture]
fn parse_response() -> Value {
json!({
"job_id": "job_123",
"usage": {"num_pages": 3, "credits": 3},
"result": {
"chunks": [
{
"content": "Page 1 block A",
"blocks": [{
"content": "Page 1 block A",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 2 block A",
"blocks": [{
"content": "Page 2 block A",
"bbox": {"page": 2},
"kind": "table",
}],
},
{
"content": "Page 1 block B",
"blocks": [{
"content": "Page 1 block B",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 3 block A",
"blocks": [{
"content": "Page 3 block A",
"bbox": {"page": 3},
"kind": "figure",
}],
},
],
},
})
}
#[rstest]
fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) {
let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=")
.expect("PDF data URI should be valid");
let upload = build_upload_request(
source,
"Bearer test-key",
Some("https://platform.reducto.ai"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://platform.reducto.ai/upload");
assert_eq!(upload.authorization, "Bearer test-key");
assert_eq!(upload.file_name, "document");
assert_eq!(upload.mime_type, "application/pdf");
assert_eq!(upload.bytes, b"%PDF-1.4");
let optional_params = json!({
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
.as_object()
.expect("params should be an object")
.clone();
let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params);
assert_eq!(
request.data,
json!({
"input": "reducto://uploaded.pdf",
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
);
let transformed = transform_reducto_response("parse-v3", parse_response.clone())
.expect("response should transform");
assert_eq!(
transformed.usage_info,
Some(json!({"pages_processed": 3, "credits": 3}))
);
assert_eq!(transformed.pages.len(), 3);
assert_eq!(
transformed.pages[0],
json!({
"index": 0,
"markdown": "Page 1 block A\n\nPage 1 block B",
"blocks": [
{"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"},
{"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"},
],
})
);
assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A");
assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A");
assert_eq!(transformed.provider_native_response, Some(parse_response));
}
#[rstest]
fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) {
let document = json!({
"type": "document_url",
"document_url": "reducto://already-uploaded.pdf",
});
let source = extract_document_source(&document).expect("Reducto ID should be valid");
assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none());
assert_eq!(
source,
ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string())
);
let request = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_request(
"parse-v3",
document,
json!({"retrieval": {"chunk_mode": "section"}})
.as_object()
.expect("params should be object")
.clone(),
)
.expect("direct ID should transform");
assert_eq!(request.data["input"], "reducto://already-uploaded.pdf");
assert_eq!(request.data["retrieval"]["chunk_mode"], "section");
let response = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_response("parse-v3", parse_response)
.expect("response should transform");
assert!(
response.pages[0]["markdown"]
.as_str()
.expect("markdown should be string")
.starts_with("Page 1 block A")
);
}
#[rstest]
fn test_parse_legacy_wraps_enhance_under_options() {
let request = build_parse_legacy_request(
"reducto://legacy.pdf",
json!({"enhance": {"agentic": [{"type": "table"}]}})
.as_object()
.expect("params should be object"),
);
assert_eq!(
request.data,
json!({
"document_url": "reducto://legacy.pdf",
"options": {"enhance": {"agentic": [{"type": "table"}]}},
})
);
}
#[rstest]
fn test_parse_v3_image_data_uri_upload_uses_image_mime() {
let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=")
.expect("PNG data URI should be valid");
let upload = build_upload_request(
source,
"Bearer programmatic-key",
Some("https://custom.reducto.test/"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://custom.reducto.test/upload");
assert_eq!(upload.authorization, "Bearer programmatic-key");
assert_eq!(upload.mime_type, "image/png");
assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n");
}
#[rstest]
#[case::http("http://example.com/document.pdf")]
#[case::https("https://example.com/document.pdf")]
fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) {
let error = classify_document_source(source).expect_err("plain URL should be rejected");
assert!(error.to_string().contains("upload the file first"));
}
#[rstest]
fn test_parse_v3_uses_programmatic_api_key_over_env() {
let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string()))
.expect("explicit key should resolve");
assert_eq!(key, "passed-key");
let headers = REDUCTO_PARSE_V3_CONFIG
.validate_environment(Vec::new(), Some("passed-key"), &|_| {
Some("env-reducto-key".to_string())
})
.expect("headers should validate");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer passed-key".to_string())]
);
}

407
litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs Normal file
View file

@ -0,0 +1,407 @@
use std::collections::BTreeMap;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use serde_json::{Map, Value, json};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub const REDUCTO_ID_PREFIX: &str = "reducto://";
const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"];
const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"];
const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()";
const DATA_URI_UPLOAD_REQUIRED: &str =
"Reducto data URI upload must complete before OCR request transformation";
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReductoDocumentSource {
FileId(String),
Upload { bytes: Vec<u8>, mime_type: String },
}
#[derive(Clone, PartialEq, Eq)]
pub struct ReductoUploadRequest {
pub url: String,
pub authorization: String,
pub file_name: &'static str,
pub bytes: Vec<u8>,
pub mime_type: String,
}
pub struct ReductoParseV3Config;
pub struct ReductoParseLegacyConfig;
pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config;
pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig;
pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> {
match model {
"parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG),
"parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG),
_ => None,
}
}
pub fn normalize_api_base(api_base: Option<&str>) -> String {
api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE)
.trim_end_matches('/')
.to_string()
}
pub fn parse_url(api_base: Option<&str>) -> String {
format!("{}/parse", normalize_api_base(api_base))
}
pub fn upload_url(api_base: Option<&str>) -> String {
format!("{}/upload", normalize_api_base(api_base))
}
pub fn resolve_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
pub fn extract_document_source(document: &Value) -> Result<ReductoDocumentSource, Error> {
let document = document.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let source = document
.get("document_url")
.and_then(Value::as_str)
.filter(|source| !source.is_empty())
.or_else(|| document.get("image_url").and_then(Value::as_str))
.ok_or_else(|| {
Error::InvalidRequest(
"Reducto expected OCR preprocessing to produce document_url or image_url"
.to_string(),
)
})?;
classify_document_source(source)
}
pub fn classify_document_source(source: &str) -> Result<ReductoDocumentSource, Error> {
if source.starts_with(REDUCTO_ID_PREFIX) {
return Ok(ReductoDocumentSource::FileId(source.to_string()));
}
if source.starts_with("http://") || source.starts_with("https://") {
return Err(Error::InvalidRequest(
"Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first."
.to_string(),
));
}
if !source.starts_with("data:") {
return Err(Error::InvalidRequest(
"Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing."
.to_string(),
));
}
let (header, encoded) = source
.split_once(',')
.ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?;
if !header.split(';').any(|part| part == "base64") {
return Err(Error::InvalidRequest(
"Reducto only supports base64-encoded data URIs.".to_string(),
));
}
let mime_type = header
.strip_prefix("data:")
.and_then(|header| header.split(';').next())
.filter(|mime| !mime.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| {
Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string())
})?;
Ok(ReductoDocumentSource::Upload { bytes, mime_type })
}
pub fn build_upload_request(
source: ReductoDocumentSource,
authorization: &str,
api_base: Option<&str>,
) -> Option<ReductoUploadRequest> {
let ReductoDocumentSource::Upload { bytes, mime_type } = source else {
return None;
};
Some(ReductoUploadRequest {
url: upload_url(api_base),
authorization: authorization.to_string(),
file_name: "document",
bytes,
mime_type,
})
}
pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> {
response_json
.as_object()
.and_then(|response| response.get("file_id"))
.and_then(Value::as_str)
.filter(|file_id| !file_id.is_empty())
.ok_or_else(|| {
Error::InvalidResponse(format!(
"Reducto /upload returned 200 without a file_id; got payload={response_json}"
))
})
}
pub fn build_parse_v3_request(
file_id: &str,
optional_params: Map<String, Value>,
) -> OcrRequestData {
let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string())))
.chain(optional_params)
.collect();
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
pub fn build_parse_legacy_request(
file_id: &str,
optional_params: &Map<String, Value>,
) -> OcrRequestData {
let options = optional_params
.get("enhance")
.filter(|enhance| !enhance.is_null())
.map(|enhance| json!({"options": {"enhance": enhance}}));
let data = match options {
Some(Value::Object(options)) => std::iter::once((
"document_url".to_string(),
Value::String(file_id.to_string()),
))
.chain(options)
.collect(),
_ => Map::from_iter([(
"document_url".to_string(),
Value::String(file_id.to_string()),
)]),
};
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
fn source_file_id(document: &Value) -> Result<String, Error> {
match extract_document_source(document)? {
ReductoDocumentSource::FileId(file_id) => Ok(file_id),
ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)),
}
}
fn page_number(block: &Map<String, Value>) -> Option<i64> {
let page = block.get("bbox")?.as_object()?.get("page")?;
page.as_i64()
.or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok()))
.or_else(|| page.as_str().and_then(|page| page.parse().ok()))
}
fn chunks(result: &Map<String, Value>) -> &[Value] {
result
.get("chunks")
.and_then(Value::as_array)
.map(Vec::as_slice)
.unwrap_or_default()
}
fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
let blocks_by_page = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array))
.flatten()
.filter_map(|block| block.as_object().map(|object| (block, object)))
.filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone())))
.fold(
BTreeMap::<i64, Vec<Value>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
return if markdown.is_empty() {
Vec::new()
} else {
vec![json!({"index": 0, "markdown": markdown})]
};
}
blocks_by_page
.into_iter()
.map(|(page, blocks)| {
let markdown = blocks
.iter()
.filter_map(Value::as_object)
.filter_map(|block| block.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
json!({
"index": page.saturating_sub(1).max(0),
"markdown": markdown,
"blocks": blocks,
})
})
.collect()
}
pub fn transform_reducto_response(
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let empty_result = Map::new();
let result = match response.get("result") {
Some(Value::Object(result)) => result,
Some(Value::Null) => &empty_result,
Some(_) => {
return Err(Error::InvalidResponse(
"Reducto result must be an object".to_string(),
));
}
None => response,
};
let usage = response
.get("usage")
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let usage_info = Some(json!({
"pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null),
"credits": usage.get("credits").cloned().unwrap_or(Value::Null),
}));
Ok(OcrResponseData {
pages: build_pages(result),
model: model.to_string(),
document_annotation: None,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: Some(response_json),
})
}
impl OcrProviderConfig for ReductoParseV3Config {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_V3_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_v3_request(&file_id, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}
impl OcrProviderConfig for ReductoParseLegacyConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_LEGACY_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_legacy_request(&file_id, &optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}

34
litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs
View file

@ -212,6 +212,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -229,6 +230,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
@ -253,10 +255,21 @@ impl OcrProviderConfig for VertexAiOcrConfig {
}
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
DEEPSEEK_SUPPORTED_OCR_PARAMS
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
non_default_params
.iter()
.filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -283,6 +296,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_response(
&self,
model: &str,
@ -335,9 +349,12 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
document_annotation: object.get("document_annotation").cloned(),
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
@ -360,6 +377,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[test]
fn vertex_mistral_url_uses_project_location_and_model() {
@ -411,6 +429,22 @@ mod tests {
);
}
#[rstest]
#[case::bare_model("deepseek-ocr-maas")]
#[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")]
fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) {
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_request(
model,
json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
}
#[test]
fn vertex_deepseek_response_wraps_markdown_content() {
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG

14
litellm-rust/crates/core/tests/workspace_crate_allowlist.rs
View file

@ -1,6 +1,7 @@
//! Enforcement: the litellm-rust workspace has exactly four crates.
//! Enforcement: the litellm-rust workspace has exactly five crates.
//!
//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host),
//! `core` (the Rust SDK), `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
@ -19,13 +20,20 @@ use std::path::{Path, PathBuf};
/// workspace legitimately gains or loses a crate.
const EXPECTED_MEMBERS: &[&str] = &[
"crates/core",
"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", "ai-gateway", "python-interop", "python-bridge"];
const EXPECTED_CRATE_DIRS: &[&str] = &[
"core",
"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).";

9
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -14,11 +14,15 @@ default = ["abi3"]
abi3 = ["pyo3/abi3-py310"]
extension-module = ["pyo3/extension-module"]
panic-test = []
trace-parity = [
"dep:tracing",
"litellm-core/observability",
"litellm-ai-gateway/trace-parity",
]
[dependencies]
futures-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
tracing = { workspace = true, optional = true }
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
litellm-python-interop.workspace = true
@ -31,6 +35,7 @@ tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
tokio-tungstenite.workspace = true
tracing.workspace = true
[[bench]]
name = "serialization"

1
litellm-rust/crates/python-bridge/src/constants.rs
View file

@ -1 +0,0 @@
pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

208
litellm-rust/crates/python-bridge/src/function_trace.rs
View file

@ -1,216 +1,22 @@
use std::future::Future;
use std::sync::{Arc, Mutex};
use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
use serde::Serialize;
use tracing::instrument::WithSubscriber;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Level, Subscriber};
use tracing_subscriber::filter::{LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
use crate::constants::FUNCTION_TRACE_TARGET;
#[derive(Serialize)]
#[serde(untagged)]
pub(crate) enum TraceResponse<T> {
Plain(T),
Traced {
response: T,
trace: Vec<FunctionTraceEvent>,
},
pub(crate) struct TracedResponse<T> {
response: T,
trace: Vec<FunctionTraceEvent>,
}
pub(crate) async fn trace_call<T, E>(
pub(crate) async fn capture<T, E>(
future: impl Future<Output = Result<T, E>>,
enabled: bool,
) -> Result<TraceResponse<T>, E> {
if !enabled {
return future.await.map(TraceResponse::Plain);
}
) -> Result<TracedResponse<T>, E> {
let trace = FunctionTrace::default();
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(TraceResponse::Traced {
Ok(TracedResponse {
response,
trace: trace.events(),
})
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionTraceEvent {
pub function: &'static str,
pub depth: usize,
}
#[derive(Clone, Default)]
pub struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
}
impl FunctionTrace {
pub fn dispatcher(&self) -> Dispatch {
let filter = filter_fn(|metadata| {
metadata.is_span()
&& metadata.target() == FUNCTION_TRACE_TARGET
&& *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE);
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(filter),
),
)
}
pub fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let depth = context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default();
self.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.push(FunctionTraceEvent {
function: attributes.metadata().name(),
depth,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn outer() {
tokio::task::yield_now().await;
inner().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn inner() {
tokio::task::yield_now().await;
}
#[tokio::test]
async fn concurrent_futures_keep_separate_traces_across_yields() {
use tracing::instrument::WithSubscriber;
let first = FunctionTrace::default();
let second = FunctionTrace::default();
let outside = FunctionTrace::default();
async {
tokio::join!(
outer().with_subscriber(first.dispatcher()),
inner().with_subscriber(second.dispatcher()),
);
inner().await;
}
.with_subscriber(outside.dispatcher())
.await;
assert_eq!(
first.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0
},
FunctionTraceEvent {
function: "inner",
depth: 1
},
],
);
assert_eq!(
second.events(),
vec![FunctionTraceEvent {
function: "inner",
depth: 0
}],
);
assert_eq!(
outside.events(),
vec![FunctionTraceEvent {
function: "inner",
depth: 0
}],
);
}
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0,
},
FunctionTraceEvent {
function: "inner",
depth: 1,
},
]
);
}
}

40
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,8 +1,8 @@
mod constants;
mod diagnostics;
mod errors;
mod execution;
pub mod function_trace;
#[cfg(feature = "trace-parity")]
mod function_trace;
mod marshal;
mod routes;
@ -115,9 +115,43 @@ mod tests {
.extract::<Vec<String>>()
.expect("module names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.filter(|name| !name.starts_with('_'))
.collect();
assert_eq!(public_names, expected);
#[cfg(not(feature = "trace-parity"))]
assert!(!module.hasattr("_trace").expect("module lookup should work"));
#[cfg(feature = "trace-parity")]
{
let trace = module
.getattr("_trace")
.expect("trace build should expose its diagnostic namespace");
let trace_names: Vec<String> = trace
.cast::<PyModule>()
.expect("trace namespace should be a module")
.dict()
.keys()
.extract::<Vec<String>>()
.expect("trace names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.collect();
assert_eq!(
trace_names,
[
"ocr",
"aocr",
"transcription",
"atranscription",
"messages",
"amessages",
"chat_completions",
"achat_completions",
"gateway_messages",
]
);
}
});
}

6
litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs
View file

@ -54,16 +54,16 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
audio: Value,
audio: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
optional_params: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_transcription,

6
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -73,16 +73,16 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
messages: Value,
messages: serde_json::Value,
},
optional = {
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
optional_params: Option<serde_json::Value>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_chat_completions,

149
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -20,43 +20,33 @@ macro_rules! bridge_route {
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
trace: bool,
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_sync(
py,
$crate::function_trace::trace_call(future, trace),
$map_error,
)
$crate::execution::run_sync(py, future, $map_error)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None,)* trace=false))]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
trace: bool,
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_async(
py,
$crate::function_trace::trace_call(future, trace),
$map_error,
)
$crate::execution::run_async(py, future, $map_error)
}
pub(super) fn register(
@ -67,6 +57,71 @@ macro_rules! bridge_route {
$crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?;
Ok(())
}
#[cfg(feature = "trace-parity")]
mod trace {
use pyo3::prelude::*;
use super::{$inputs, $map_error, $prepare};
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $sync_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
) -> pyo3::PyResult<pyo3::Py<pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_sync(
py,
$crate::function_trace::capture(future),
$map_error,
)
}
#[pyfunction]
#[pyo3(signature = ($($required_name),*, $($optional_name=None),*))]
#[allow(clippy::too_many_arguments)]
fn $async_name(
py: pyo3::Python<'_>,
$($(#[$required_attr])* $required_name: $required_type,)*
$($(#[$optional_attr])* $optional_name: $optional_type,)*
) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
let future = $prepare($inputs {
$($required_name,)*
$($optional_name),*
})?;
$crate::execution::run_async(
py,
$crate::function_trace::capture(future),
$map_error,
)
}
pub(super) fn register(
module: &pyo3::Bound<'_, pyo3::types::PyModule>,
) -> pyo3::PyResult<()> {
$crate::routes::definition::add_function(
module,
pyo3::wrap_pyfunction!($sync_name, module)?,
)?;
$crate::routes::definition::add_function(
module,
pyo3::wrap_pyfunction!($async_name, module)?,
)?;
Ok(())
}
}
#[cfg(feature = "trace-parity")]
pub(super) fn register_trace(
module: &pyo3::Bound<'_, pyo3::types::PyModule>,
) -> pyo3::PyResult<()> {
trace::register(module)
}
};
}
@ -130,20 +185,26 @@ mod tests {
) -> PyResult<impl Future<Output = Result<String, Error>> + Send + 'static> {
FUTURE_DROPPED.store(false, Ordering::SeqCst);
let drop_guard = (inputs.value == "pending").then_some(DropGuard);
Ok(async move {
let _drop_guard = drop_guard;
tokio::task::yield_now().await;
match inputs.value.as_str() {
"error" => Err(Error::InvalidRequest("synthetic error".to_string())),
"map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
"panic" => panic!("synthetic panic"),
"pending" => {
pending::<()>().await;
unreachable!()
}
_ => Ok(inputs.value),
Ok(execute_echo(inputs, drop_guard))
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn execute_echo(
inputs: EchoInputs,
drop_guard: Option<DropGuard>,
) -> Result<String, Error> {
let _drop_guard = drop_guard;
tokio::task::yield_now().await;
match inputs.value.as_str() {
"error" => Err(Error::InvalidRequest("synthetic error".to_string())),
"map_panic" => Err(Error::InvalidRequest("panic in mapper".to_string())),
"panic" => panic!("synthetic panic"),
"pending" => {
pending::<()>().await;
unreachable!()
}
})
_ => Ok(inputs.value),
}
}
fn map_error(error: Error) -> PyErr {
@ -164,22 +225,22 @@ mod tests {
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"transcription",
"atranscription",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None, trace=False)",
"(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
),
(
"messages",
"amessages",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
"(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
(
"chat_completions",
"achat_completions",
"(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, trace=False)",
"(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)",
),
];
@ -411,6 +472,32 @@ asyncio.run(exercise())
});
}
#[cfg(feature = "trace-parity")]
#[test]
fn diagnostic_route_returns_the_response_and_filtered_trace() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "synthetic").expect("module should be created");
synthetic::register_trace(&module).expect("trace routes should register");
let locals = PyDict::new(py);
locals
.set_item("routes", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
result = routes.echo("traced")
assert result == {
"response": "traced",
"trace": [{"function": "execute_echo", "depth": 0}],
}
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("diagnostic route should return its response and trace");
});
}
#[test]
fn route_registration_rejects_duplicate_python_names() {
Python::initialize();

29
litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs Normal file
View file

@ -0,0 +1,29 @@
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<Bound<'py, PyAny>> {
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)?)
}

4
litellm-rust/crates/python-bridge/src/routes/messages.rs
View file

@ -50,14 +50,14 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
body: Value,
body: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_messages,

16
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -3,6 +3,9 @@ use pyo3::prelude::*;
#[macro_use]
mod definition;
#[cfg(feature = "trace-parity")]
mod gateway_messages;
mod audio_transcription;
mod chat_completions;
mod messages;
@ -12,5 +15,16 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
ocr::register(module)?;
audio_transcription::register(module)?;
messages::register(module)?;
chat_completions::register(module)
chat_completions::register(module)?;
#[cfg(feature = "trace-parity")]
{
let trace = PyModule::new(module.py(), "_trace")?;
ocr::register_trace(&trace)?;
audio_transcription::register_trace(&trace)?;
messages::register_trace(&trace)?;
chat_completions::register_trace(&trace)?;
gateway_messages::register_trace(&trace)?;
module.add_submodule(&trace)?;
}
Ok(())
}

6
litellm-rust/crates/python-bridge/src/routes/ocr.rs
View file

@ -56,16 +56,16 @@ bridge_route! {
required = {
model: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
document: Value,
document: serde_json::Value,
},
optional = {
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<Value>,
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<Value>,
optional_params: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,

423
litellm-rust/crates/python-bridge/src/routes/runtime.rs
View file

@ -1,423 +0,0 @@
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::time::Duration;
use futures_util::FutureExt;
use litellm_core::error::Error;
use litellm_python_interop::{Pythonized, panic_to_pyerr, release_gil};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pub(super) fn run_sync<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
}
fn run_sync_on<T, F>(
py: Python<'_>,
runtime: &Runtime,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Py<PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + 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",
));
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;
let result = map_core_result(result, map_error)?;
Pythonized(result).into_pyobject(py).map(Bound::unbind)
}
pub(super) fn run_async<T, F>(
py: Python<'_>,
future: F,
map_error: fn(Error) -> PyErr,
) -> PyResult<Bound<'_, PyAny>>
where
T: Serialize + Send + 'static,
F: Future<Output = Result<T, Error>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let result = catch_route_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
})
}
fn map_core_result<T>(result: Result<T, Error>, map_error: fn(Error) -> PyErr) -> PyResult<T> {
match result {
Ok(value) => Ok(value),
Err(error) => Err(
std::panic::catch_unwind(AssertUnwindSafe(|| map_error(error)))
.map_err(panic_to_pyerr)?,
),
}
}
async fn catch_route_panic<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
AssertUnwindSafe(future)
.catch_unwind()
.await
.map_err(panic_to_pyerr)
}
async fn wait_for_sync_result<T, F>(future: F) -> PyResult<Result<T, Error>>
where
F: Future<Output = Result<T, Error>>,
{
let future = catch_route_panic(future);
tokio::pin!(future);
let signal_interval = Duration::from_millis(50);
let mut signal_checks =
time::interval_at(time::Instant::now() + signal_interval, signal_interval);
signal_checks.set_missed_tick_behavior(MissedTickBehavior::Delay);
loop {
tokio::select! {
result = &mut future => return result,
_ = signal_checks.tick() => Python::attach(|py| py.check_signals())?,
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::CString;
use std::future::poll_fn;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, mpsc};
use std::task::Poll;
use std::thread;
use std::time::Instant;
use pyo3::panic::PanicException;
use pyo3::types::{PyDict, PyModule};
use serde::Serializer;
use tokio::runtime::Builder;
use super::*;
fn runtime_error(error: Error) -> PyErr {
PyRuntimeError::new_err(error.to_string())
}
fn panicking_error_mapper(_error: Error) -> PyErr {
panic!("error mapper panicked")
}
struct PanickingOutput;
static ASYNC_PROBE_COMPLETED: AtomicUsize = AtomicUsize::new(0);
impl Serialize for PanickingOutput {
fn serialize<S>(&self, _serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
panic!("serializer panicked")
}
}
#[pyfunction]
fn async_serialization_panic(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(py, async { Ok(PanickingOutput) }, runtime_error)
}
#[pyfunction]
fn async_runtime_probe(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
run_async(
py,
async {
ASYNC_PROBE_COMPLETED.fetch_add(1, Ordering::SeqCst);
Ok(true)
},
runtime_error,
)
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
result
.expect("route should complete")
.bind(py)
.extract()
.expect("result should convert")
}
#[test]
fn sync_runner_polls_future_on_the_caller_thread() {
Python::initialize();
Python::attach(|py| {
let caller_thread = std::thread::current().id();
let result = run_sync(
py,
async move { Ok(std::thread::current().id() == caller_thread) },
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_releases_gil_while_waiting() {
Python::initialize();
Python::attach(|py| {
let result = run_sync(
py,
async {
let gil_acquired = tokio::time::timeout(
Duration::from_secs(2),
tokio::task::spawn_blocking(|| Python::attach(|_| true)),
)
.await;
Ok(matches!(gil_acquired, Ok(Ok(true))))
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_rejects_calls_from_a_tokio_context() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
let error = runtime.block_on(async {
Python::attach(|py| {
run_sync::<bool, _>(py, async { Ok(true) }, runtime_error)
.expect_err("sync route should reject a nested Tokio runtime")
})
});
assert_eq!(
error.to_string(),
"RuntimeError: synchronous native routes cannot run from a Tokio context; use the async route"
);
}
#[test]
fn sync_runner_can_drive_a_current_thread_runtime() {
Python::initialize();
let runtime = Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime should build");
Python::attach(|py| {
let result = run_sync_on(
py,
&runtime,
async {
tokio::task::yield_now().await;
Ok(true)
},
runtime_error,
);
assert!(extract_bool(py, result));
});
}
#[test]
fn sync_runner_maps_a_panicked_future() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
poll_fn(|_| -> Poll<Result<bool, Error>> { panic!("route future panicked") }),
runtime_error,
)
.expect_err("panicked route should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: route future panicked");
});
}
#[test]
fn sync_runner_maps_a_panicked_error_mapper() {
Python::initialize();
Python::attach(|py| {
let error = run_sync::<bool, _>(
py,
async { Err(Error::InvalidRequest("invalid".to_string())) },
panicking_error_mapper,
)
.expect_err("panicked mapper should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: error mapper panicked");
});
}
#[test]
fn sync_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error)
.expect_err("serializer panic should become a Python exception");
assert!(error.is_instance_of::<PanicException>(py));
assert_eq!(error.to_string(), "PanicException: serializer panicked");
});
}
#[test]
fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() {
Python::initialize();
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let callers: Vec<_> = (0..2)
.map(|_| {
let barrier = Arc::clone(&barrier);
thread::spawn(move || {
Python::attach(|py| {
extract_bool(
py,
run_sync(
py,
async move {
Ok(tokio::time::timeout(Duration::from_secs(2), barrier.wait())
.await
.is_ok())
},
runtime_error,
),
)
})
})
})
.collect();
let results: Vec<_> = callers
.into_iter()
.map(|caller| caller.join().expect("caller should not panic"))
.collect();
assert_eq!(results, vec![true, true]);
}
#[test]
fn async_runner_surfaces_serializer_panics() {
Python::initialize();
Python::attach(|py| {
let module = PyModule::new(py, "runtime").expect("module should be created");
module
.add_function(
wrap_pyfunction!(async_serialization_panic, &module)
.expect("function should wrap"),
)
.expect("function should register");
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
try:
await runtime.async_serialization_panic()
except BaseException as error:
assert type(error).__name__ == "PanicException"
assert str(error) == "serializer panicked"
else:
raise AssertionError("serializer panic was not raised")
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("serializer panic should reach the Python awaiter");
});
}
#[test]
fn async_result_delivery_does_not_stall_tokio_workers() {
Python::initialize();
ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst);
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"),
wrap_pyfunction!(runtime_worker_count, &module).expect("function should wrap"),
wrap_pyfunction!(runtime_is_responsive, &module).expect("function should wrap"),
] {
module
.add_function(function)
.expect("function should register");
}
let locals = PyDict::new(py);
locals
.set_item("runtime", &module)
.expect("module should enter Python locals");
let code = CString::new(
r#"
import asyncio
async def exercise():
worker_count = runtime.runtime_worker_count()
awaitables = [runtime.async_runtime_probe() for _ in range(worker_count)]
assert runtime.runtime_is_responsive(worker_count)
assert await asyncio.gather(*awaitables) == [True] * worker_count
asyncio.run(exercise())
"#,
)
.expect("Python source should not contain null bytes");
py.run(&code, Some(&locals), Some(&locals))
.expect("result delivery should leave Tokio workers responsive");
});
}
}

2
litellm/__init__.py
View file

@ -1421,7 +1421,7 @@ from .skills.main import (
)
from .containers.main import *
from .ocr.main import *
from .rust_bridge import use_litellm_rust
from .rust_bridge import rust
from .rag.main import *
from .sandbox.main import *
from .search.main import *

30
litellm/caching/caching.py
View file

@ -100,6 +100,7 @@ class Cache:
qdrant_semantic_cache_vector_size: int | None = None,
semantic_cache_embedding_max_input_tokens: int | None = None,
semantic_cache_embedding_timeout: float | None = None,
semantic_cache_scope: str = SemanticCacheScope.KEY.value,
# GCP IAM authentication parameters
gcp_service_account: str | None = None,
gcp_ssl_ca_certs: str | None = None,
@ -127,6 +128,7 @@ class Cache:
similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS.
semantic_cache_scope (str, optional): "key" isolates semantic-cache buckets per key/team/org. "end_user" additionally isolates per end user (falls back to the key scope when the request carries no end-user id). Defaults to "key".
# Disk Cache Args
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
@ -274,6 +276,7 @@ class Cache:
self.redis_flush_size = redis_flush_size
self.ttl = ttl
self.mode: CacheMode = mode or CacheMode.default_on
self.semantic_cache_scope: str = SemanticCacheScope(semantic_cache_scope).value
if self.type == LiteLLMCacheType.LOCAL and default_in_memory_ttl is not None:
self.ttl = default_in_memory_ttl
@ -301,6 +304,7 @@ class Cache:
"user_api_key_team_id",
"user_api_key_org_id",
)
_SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id"
def _is_semantic_cache(self) -> bool:
return self.type in (
@ -309,19 +313,21 @@ class Cache:
LiteLLMCacheType.VALKEY_SEMANTIC,
)
def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str:
metadata: Final[dict] = kwargs.get("metadata") or {}
litellm_params: Final[dict] = kwargs.get("litellm_params") or {}
metadata_in_litellm_params: Final[dict] = litellm_params.get("metadata") or {}
def _semantic_cache_scope_fields(self) -> tuple[str, ...]:
if self.semantic_cache_scope == SemanticCacheScope.END_USER:
return (*self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS, self._SEMANTIC_CACHE_END_USER_SCOPE_FIELD)
return self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS
scope = ""
for field in self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS:
value = metadata.get(field)
if value is None:
value = metadata_in_litellm_params.get(field)
if value is not None:
scope += f"{field}: {value}"
return scope
def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str:
litellm_params: Final[dict] = kwargs.get("litellm_params") or {}
metadata_sources: Final[tuple[dict, ...]] = tuple(
source.get(key) or {} for source in (kwargs, litellm_params) for key in ("metadata", "litellm_metadata")
)
scope_values: Final = (
(field, next((source[field] for source in metadata_sources if source.get(field) is not None), None))
for field in self._semantic_cache_scope_fields()
)
return "".join(f"{field}: {value}" for field, value in scope_values if value is not None)
def get_cache_key(self, **kwargs) -> str:
"""

9
litellm/llms/base_llm/ocr/transformation.py
View file

@ -2,6 +2,7 @@
Base OCR transformation configuration.
"""
import builtins
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
@ -93,8 +94,8 @@ class OCRResponse(LiteLLMPydanticObjectBase):
document_annotation: Any | None = None
usage_info: OCRUsageInfo | None = None
content: str | None = None
tables: list[dict[str, object]] | None = None
keyValuePairs: list[dict[str, object]] | None = None
tables: list[dict[str, builtins.object]] | None = None
keyValuePairs: list[dict[str, builtins.object]] | None = None
object: str = "ocr"
model_config = {"extra": "allow"}
@ -102,11 +103,11 @@ class OCRResponse(LiteLLMPydanticObjectBase):
# Define private attributes using PrivateAttr
_hidden_params: dict = PrivateAttr(default_factory=dict)
def set_provider_native_response(self, native_response: Mapping[str, object]) -> None:
def set_provider_native_response(self, native_response: Mapping[str, builtins.object]) -> None:
"""Keep the provider's own response payload alongside the normalized one."""
self._hidden_params[PROVIDER_NATIVE_RESPONSE_KEY] = native_response
def get_provider_native_response(self) -> Mapping[str, object] | None:
def get_provider_native_response(self) -> Mapping[str, builtins.object] | None:
"""The provider's own response payload, when `req_format=native` was requested."""
native_response: Final = self._hidden_params.get(PROVIDER_NATIVE_RESPONSE_KEY)
return native_response if isinstance(native_response, dict) else None

3
litellm/proxy/_types.py
View file

@ -737,6 +737,9 @@ class LiteLLMRoutes(enum.Enum):
"/.well-known/litellm-ui-config",
"/public/model_hub",
"/public/v1/model_hub",
"/public/v1/model_hub/providers",
"/public/v1/model_hub/modes",
"/public/v1/model_hub/features",
"/public/model_hub/info",
"/public/agent_hub",
"/public/mcp_hub",

2
litellm/proxy/client/models.py
View file

@ -32,7 +32,7 @@ class ModelsManagementClient:
headers["Authorization"] = f"Bearer {self._api_key}"
return headers
def list(self, return_request: bool = False) -> list[dict[str, Any]] | requests.Request:
def list(self, return_request: bool = False) -> builtins.list[dict[str, Any]] | requests.Request:
"""
Get the list of models supported by the server.

2
litellm/proxy/client/teams.py
View file

@ -40,7 +40,7 @@ class TeamsManagementClient:
self,
user_id: str | None = None,
organization_id: str | None = None,
) -> list[dict[str, Any]]:
) -> builtins.list[dict[str, Any]]:
"""
List teams that the user belongs to.

12
litellm/proxy/list_api/in_memory.py
View file

@ -141,3 +141,15 @@ class InMemoryListExecutor(Generic[TRow]):
async def find_many(self, plan: QueryPlan) -> Sequence[TRow]:
page: Final = _ordered(self._matching(plan.where), plan.order)[plan.skip : plan.skip + plan.take]
return await self.enrich_page(tuple(row for _, row in page))
async def distinct(self, field: str, where: tuple[Predicate, ...]) -> Sequence[str]:
"""A repeated field contributes each of its elements, so a facet over `providers`
lists providers rather than the tuples rows happen to carry."""
cells: Final = (cells.get(field) for cells, _ in self._matching(where))
values: Final = (
value
for cell in cells
for value in (cell if isinstance(cell, tuple) else (cell,))
if isinstance(value, str) and value
)
return tuple(sorted(frozenset(values)))

82
litellm/proxy/list_api/list_framework.py
View file

@ -28,12 +28,15 @@ from litellm.proxy.list_api.common import (
PROBLEM_TYPE_BASE,
ManagementProblem,
build_list_links,
build_page_links,
escape_like,
unknown_query_param_problem,
)
from litellm.types.proxy.management_endpoints.management_v1 import (
FacetListResponse,
ListMeta,
ListResponse,
PageMeta,
ProblemDetail,
)
@ -186,6 +189,13 @@ class ListExecutor(Protocol[TRow_co]):
async def find_many(self, plan: QueryPlan) -> Sequence[TRow_co]: ...
class FacetExecutor(Protocol):
"""The half of a facet that knows the rows. Separate from `ListExecutor` so a SQL
executor is not forced to implement `distinct` to keep serving entity lists."""
async def distinct(self, field: str, where: tuple[Predicate, ...]) -> Sequence[str]: ...
def order_by_sql(order: tuple[SortKey, ...]) -> str:
"""`ORDER BY` body for a plan, NULLS LAST in both directions.
@ -515,6 +525,78 @@ def build_query_plan(
)
def _facet_allowed_params(spec: ListSpec[TRow, TOut]) -> tuple[str, ...]:
"""A facet's values are always ascending, so `sort` is not one of its parameters."""
return tuple(name for name in _allowed_params(spec) if name != SORT_PARAM)
def _facet_where(
spec: ListSpec[TRow, TOut],
params: Mapping[str, str],
caller: UserAPIKeyAuth,
) -> tuple[Predicate, ...] | ProblemDetail:
scope_predicates: Final = _scope_predicates(spec.scope(caller))
if isinstance(scope_predicates, ProblemDetail):
return scope_predicates
filters: Final = _parse_filters(spec, params)
if isinstance(filters, ProblemDetail):
return filters
search: Final = _search_predicate(spec, params)
return scope_predicates + filters + ((search,) if search is not None else ())
async def handle_facet(
spec: ListSpec[TRow, TOut],
executor: FacetExecutor,
request: Request,
caller: UserAPIKeyAuth,
field: str,
) -> FacetListResponse:
"""The distinct values one column takes over a filtered query on a resource.
Carries the parent's parameters so a filter dropdown offers exactly the values the
table can show, and `has_more` rather than a total, which would cost a COUNT(*) over
the whole match set on every keystroke.
"""
params: Final = request.query_params
unknown: Final = tuple(sorted(name for name in params if name == SORT_PARAM or not _is_known_param(spec, name)))
if unknown:
raise ManagementProblem(unknown_query_param_problem(unknown=unknown, allowed=_facet_allowed_params(spec)))
duplicates: Final = _duplicate_params(request)
if duplicates:
raise ManagementProblem(
_problem(
"duplicate-query-parameter",
"Duplicate query parameter",
400,
f"Repeated query parameter(s): {', '.join(duplicates)}. Each may appear once; "
f"use a comma-separated list for multiple filter values.",
)
)
page: Final = _parse_page(params)
if isinstance(page, ProblemDetail):
raise ManagementProblem(page)
page_size: Final = _parse_page_size(spec, params)
if isinstance(page_size, ProblemDetail):
raise ManagementProblem(page_size)
where: Final = _facet_where(spec, params, caller)
if isinstance(where, ProblemDetail):
raise ManagementProblem(where)
values: Final = await executor.distinct(field, where)
skip: Final = (page - 1) * page_size
window: Final = values[skip : skip + page_size + 1]
has_more: Final = len(window) > page_size
return FacetListResponse(
data=tuple(window[:page_size]),
meta=PageMeta(page=page, page_size=page_size, has_more=has_more),
links=build_page_links(request=request, page=page, has_more=has_more),
)
def _duplicate_params(request: Request) -> tuple[str, ...]:
names: Final = tuple(name for name, _ in request.query_params.multi_items())
return tuple(sorted(frozenset(name for name in names if names.count(name) > 1)))

141
litellm/proxy/public_endpoints/public_v1/model_hub.py
View file

@ -3,7 +3,7 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Annotated, Final, Protocol
from typing import Annotated, Final, Literal, Protocol
from fastapi import APIRouter, Depends, Request
from typing_extensions import ReadOnly, TypedDict
@ -20,10 +20,12 @@ from litellm.proxy.list_api.list_framework import (
Scope,
ScopeAll,
SortKey,
handle_facet,
handle_list,
)
from litellm.proxy.utils import PrismaClient
from litellm.types.proxy.management_endpoints.management_v1 import (
FacetListResponse,
ListResponse,
ProblemDetail,
)
@ -95,16 +97,37 @@ class HealthEnricher:
return tuple(_with_health(row, health.get(row.model_group)) for row in rows)
FEATURE_PREFIX: Final = "supports_"
def _features(row: ModelGroupInfoProxy) -> tuple[str, ...]:
"""A row's capabilities as one repeated field, so selecting two of them matches either.
The hub's feature control has always been a multi-select over the `supports_*` flags.
One boolean filter per flag would AND them, which is the opposite of what it does.
"""
return tuple(
sorted(
name.removeprefix(FEATURE_PREFIX)
for name, value in row.model_dump().items()
if name.startswith(FEATURE_PREFIX) and value is True
)
)
def _cells(row: ModelGroupInfoProxy) -> Cells:
return MappingProxyType(
{
"model_group": row.model_group,
"mode": row.mode,
"providers": tuple(row.providers),
"features": _features(row),
"max_input_tokens": row.max_input_tokens,
"max_output_tokens": row.max_output_tokens,
"input_cost_per_token": row.input_cost_per_token,
"output_cost_per_token": row.output_cost_per_token,
"rpm": row.rpm,
"tpm": row.tpm,
}
)
@ -126,20 +149,28 @@ def _scope(_caller: UserAPIKeyAuth) -> Scope:
MODEL_HUB_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType(
{
"mode": FilterSpec(type=str, ops=frozenset(("eq", "in"))),
"providers": FilterSpec(type=str, ops=frozenset(("contains",))),
"providers": FilterSpec(type=str, ops=frozenset(("contains", "in"))),
"features": FilterSpec(type=str, ops=frozenset(("in",))),
}
)
MODEL_HUB_FACETS: Final[Mapping[str, str]] = MappingProxyType(
{"providers": "providers", "modes": "mode", "features": "features"}
)
MODEL_HUB_LIST_SPEC: Final[ListSpec[ModelGroupInfoProxy, ModelGroupInfoProxy]] = ListSpec(
resource="model groups",
sortable=frozenset(
(
"model_group",
"mode",
"providers",
"max_input_tokens",
"max_output_tokens",
"input_cost_per_token",
"output_cost_per_token",
"rpm",
"tpm",
)
),
searchable=frozenset(("model_group",)),
@ -153,6 +184,32 @@ MODEL_HUB_LIST_SPEC: Final[ListSpec[ModelGroupInfoProxy, ModelGroupInfoProxy]] =
)
def _published_rows() -> Sequence[ModelGroupInfoProxy]:
from litellm.proxy.proxy_server import (
_get_model_group_info, # pyright: ignore[reportPrivateUsage] # /public/model_hub imports it the same way
llm_router,
)
if llm_router is None:
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}no-llm-router",
title="No models configured",
status=400,
detail=CommonProxyErrors.no_llm_router.value,
)
)
if litellm.public_model_groups is None:
return ()
return tuple(
_get_model_group_info(
llm_router=llm_router,
all_models_str=litellm.public_model_groups,
model_group=None,
)
)
def _executor(
rows: Sequence[ModelGroupInfoProxy],
prisma_client: PrismaClient | None,
@ -191,37 +248,11 @@ async def public_model_hub_list(
```
"""
try:
from litellm.proxy.proxy_server import (
_get_model_group_info, # pyright: ignore[reportPrivateUsage] # /public/model_hub imports it the same way
llm_router,
prisma_client,
)
if llm_router is None:
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}no-llm-router",
title="No models configured",
status=400,
detail=CommonProxyErrors.no_llm_router.value,
)
)
rows: Final[Sequence[ModelGroupInfoProxy]] = (
()
if litellm.public_model_groups is None
else tuple(
_get_model_group_info(
llm_router=llm_router,
all_models_str=litellm.public_model_groups,
model_group=None,
)
)
)
from litellm.proxy.proxy_server import prisma_client
return await handle_list(
spec=MODEL_HUB_LIST_SPEC,
executor=_executor(rows, prisma_client),
executor=_executor(_published_rows(), prisma_client),
request=request,
caller=user_api_key_dict,
)
@ -240,3 +271,53 @@ async def public_model_hub_list(
detail="Failed to list public model groups.",
)
)
@router.get(
"/model_hub/{facet}",
tags=["public", "model management"], # mutable-ok: fastapi types tags as list[str | Enum]
dependencies=(Depends(user_api_key_auth),),
response_model=FacetListResponse,
)
async def public_model_hub_facet(
request: Request,
facet: Literal["providers", "modes", "features"],
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> FacetListResponse:
"""
The distinct providers, modes or features across the published model groups, for the
Model Hub's filter dropdowns. No authentication.
Carries the same filters and search as the list route, so a dropdown offers exactly
the values the table can show: asking for providers under `filter[mode][in]=chat`
lists only the providers that serve a chat model.
Example curl:
```
curl --location --globoff \
'http://0.0.0.0:4000/public/v1/model_hub/providers?filter[mode][in]=chat&page_size=50'
```
"""
try:
return await handle_facet(
spec=MODEL_HUB_LIST_SPEC,
executor=InMemoryListExecutor(rows=_published_rows(), cells=_cells),
request=request,
caller=user_api_key_dict,
field=MODEL_HUB_FACETS[facet],
)
except ManagementProblem:
raise
except Exception as e: # noqa: BLE001 # a router error answers as a problem document, not the OpenAI error shape
verbose_proxy_logger.exception(
"litellm.proxy.public_endpoints.public_v1.model_hub.public_model_hub_facet(): Exception occured - %s", e
)
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}internal-server-error",
title="Internal server error",
status=500,
detail="Failed to list public model group values.",
)
)

1
litellm/proxy/schema.prisma
View file

@ -5,6 +5,7 @@ datasource client {
generator client {
provider = "prisma-client-py"
recursive_type_depth = -1
binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"]
}

5
litellm/rust_bridge/__init__.py
View file

@ -1,9 +1,10 @@
"""LiteLLM Rust bridge package."""
from litellm.rust_bridge.configuration import use_litellm_rust
from litellm.rust_bridge.configuration import rust
from litellm.rust_bridge.loader import (
get_native_bridge,
native_bridge_available,
reset_native_bridge_cache,
)
__all__ = ["get_native_bridge", "native_bridge_available", "use_litellm_rust"]
__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "rust"]

87
litellm/rust_bridge/configuration.py
View file

@ -2,13 +2,7 @@ from __future__ import annotations
import os
import warnings
from typing import TYPE_CHECKING, Final
if TYPE_CHECKING:
from litellm.rust_bridge.messages import RustAmessages, RustMessages
from litellm.rust_bridge.ocr import RustAocr, RustOcr
from litellm.rust_bridge.responses_websocket import RustResponsesWebSocketConnection
from litellm.rust_bridge.transcription import RustAtranscription, RustTranscription
from typing import Final
DEFAULT_RUST_ENABLED: Final = False
_TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"})
@ -16,13 +10,6 @@ _GLOBAL_ENV_NAME: Final = "LITELLM_RUST"
_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR"
class _Unset:
pass
_UNSET: Final = _Unset()
class _RustConfiguration:
def __init__(self) -> None:
self.override: bool | None = None
@ -42,7 +29,7 @@ def resolve_rust_enabled(
request_override: bool | None,
process_override: bool | None,
environment_override: bool | None,
legacy_ocr_override: bool | None = None,
legacy_environment_override: bool | None = None,
release_default: bool = DEFAULT_RUST_ENABLED,
) -> bool:
if request_override is not None:
@ -51,25 +38,12 @@ def resolve_rust_enabled(
return process_override
if environment_override is not None:
return environment_override
if legacy_ocr_override is not None:
return legacy_ocr_override
if legacy_environment_override is not None:
return legacy_environment_override
return release_default
def rust_enabled(*, request_override: bool | None = None) -> bool:
if request_override is not None:
return request_override
process_override: Final = _CONFIGURATION.override
if process_override is not None:
return process_override
return resolve_rust_enabled(
request_override=None,
process_override=None,
environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)),
)
def rust_ocr_enabled(*, request_override: bool | None = None) -> bool:
if request_override is not None:
return request_override
process_override: Final = _CONFIGURATION.override
@ -87,62 +61,21 @@ def rust_ocr_enabled(*, request_override: bool | None = None) -> bool:
request_override=None,
process_override=None,
environment_override=global_override,
legacy_ocr_override=legacy_override,
legacy_environment_override=legacy_override,
)
def rust_ocr_enabled(*, request_override: bool | None = None) -> bool:
return rust_enabled(request_override=request_override)
def reset_rust_configuration() -> None:
_CONFIGURATION.override = None
def use_litellm_rust(
enabled: bool = True,
*,
ocr: RustOcr | None | _Unset = _UNSET,
aocr: RustAocr | None | _Unset = _UNSET,
messages: RustMessages | None | _Unset = _UNSET,
amessages: RustAmessages | None | _Unset = _UNSET,
responses_websocket: type[RustResponsesWebSocketConnection] | None | _Unset = _UNSET,
transcription: RustTranscription | None | _Unset = _UNSET,
atranscription: RustAtranscription | None | _Unset = _UNSET,
) -> None:
def rust(enabled: bool) -> None:
"""Set the process override for optional Rust paths.
Rust-only paths, including Bedrock transcription, are not controlled by this switch.
"""
_CONFIGURATION.override = enabled
bindings: Final = (ocr, aocr, messages, amessages, responses_websocket, transcription, atranscription)
if all(isinstance(binding, _Unset) for binding in bindings):
return
warnings.warn(
"Injecting Rust bridge implementations through use_litellm_rust() is deprecated; "
"use the internal bridge setters in tests",
DeprecationWarning,
stacklevel=2,
)
if not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset):
from litellm.rust_bridge.ocr import set_rust_ocr
if not isinstance(ocr, _Unset):
set_rust_ocr(ocr=ocr)
if not isinstance(aocr, _Unset):
set_rust_ocr(aocr=aocr)
if not isinstance(messages, _Unset) or not isinstance(amessages, _Unset):
from litellm.rust_bridge.messages import set_rust_messages
if not isinstance(messages, _Unset):
set_rust_messages(messages=messages)
if not isinstance(amessages, _Unset):
set_rust_messages(amessages=amessages)
if not isinstance(responses_websocket, _Unset):
from litellm.rust_bridge.responses_websocket import set_rust_responses_websocket
set_rust_responses_websocket(connection=responses_websocket)
if not isinstance(transcription, _Unset) or not isinstance(atranscription, _Unset):
from litellm.rust_bridge.transcription import configure_rust_transcription
if not isinstance(transcription, _Unset):
configure_rust_transcription(transcription=transcription)
if not isinstance(atranscription, _Unset):
configure_rust_transcription(atranscription=atranscription)

6
litellm/rust_bridge/loader.py
View file

@ -24,6 +24,12 @@ def get_native_bridge() -> ModuleType | None:
return _native
def reset_native_bridge_cache() -> None:
"""Forget the cached extension so the next lookup reimports it from disk."""
global _cached_bridge
_cached_bridge = _BRIDGE_SENTINEL
def native_bridge_available() -> bool:
"""Whether the packaged Rust extension is importable."""
return get_native_bridge() is not None

2
litellm/rust_bridge/ocr.py
View file

@ -11,7 +11,7 @@ from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
rust_ocr_enabled = _configuration.rust_ocr_enabled
use_litellm_rust = _configuration.use_litellm_rust
rust = _configuration.rust
class RustOcr(Protocol):

5
litellm/types/caching.py
View file

@ -18,6 +18,11 @@ class LiteLLMCacheType(str, Enum):
GCS = "gcs"
class SemanticCacheScope(str, Enum):
KEY = "key"
END_USER = "end_user"
CachingSupportedCallTypes = Literal[
"completion",
"acompletion",

2
litellm/types/llms/openai.py
View file

@ -632,7 +632,7 @@ class ChatCompletionReasoningItem(TypedDict, total=False):
type: Required[Literal["reasoning"]]
id: str
encrypted_content: str | None
summary: list["ChatCompletionReasoningSummaryTextBlock"]
summary: ReadOnly[list[ChatCompletionReasoningSummaryTextBlock]]
class WebSearchOptionsUserLocationApproximate(TypedDict, total=False):

13
litellm/types/management_endpoints/cache_settings_endpoints.py
View file

@ -187,6 +187,19 @@ CACHE_SETTINGS_FIELDS: Final[list[CacheSettingsField]] = [
ui_field_name="Embedding Model",
redis_type="semantic",
),
CacheSettingsField(
field_name="semantic_cache_scope",
field_type="String",
field_value=None,
field_description=(
"Isolation granularity for semantic cache hits. 'key' shares hits between all end users of a key/team/org."
" 'end_user' also isolates per end user; requests without an end user fall back to the key scope."
),
field_default="key",
options=["key", "end_user"],
ui_field_name="Semantic Cache Scope",
redis_type="semantic",
),
# GCP IAM authentication fields
CacheSettingsField(
field_name="gcp_service_account",

3
litellm/types/proxy/management_endpoints/management_v1.py
View file

@ -1,5 +1,6 @@
"""Shared response shapes for the `/management/v1` control-plane surface."""
from collections.abc import Sequence
from typing import Generic, TypeVar
from pydantic import BaseModel, ConfigDict, Field
@ -38,7 +39,7 @@ class PageMeta(BaseModel):
class FacetListResponse(BaseModel):
"""The distinct values one column takes over a filtered query. `data` holds bare values, not entity rows."""
data: list[str]
data: Sequence[str]
meta: PageMeta
links: PageLinks

10
litellm/utils.py
View file

@ -5121,14 +5121,8 @@ def get_response_string(response_obj: ModelResponse | ModelResponseStream) -> st
return "".join(response_parts)
def get_utc_datetime():
import datetime as dt
from datetime import datetime
if hasattr(dt, "UTC"):
return datetime.now(dt.UTC)
else:
return datetime.utcnow()
def get_utc_datetime() -> datetime.datetime:
return datetime.datetime.now(datetime.timezone.utc)
def get_max_tokens(model: str) -> int | None:

10
litellm/vector_stores/main.py
View file

@ -54,7 +54,7 @@ def _direct_vector_store_embedding_executor(
def mock_vector_store_search_response(
mock_results: list[VectorStoreSearchResult] | None = None,
mock_results: builtins.list[VectorStoreSearchResult] | None = None,
):
"""Mock response for vector store search"""
if mock_results is None:
@ -108,7 +108,7 @@ def mock_vector_store_create_response(
@client
async def acreate(
name: str | None = None,
file_ids: list[str] | None = None,
file_ids: builtins.list[str] | None = None,
expires_after: dict | None = None,
chunking_strategy: dict | None = None,
metadata: dict[str, str] | None = None,
@ -172,7 +172,7 @@ async def acreate(
@client
def create(
name: str | None = None,
file_ids: list[str] | None = None,
file_ids: builtins.list[str] | None = None,
expires_after: dict | None = None,
chunking_strategy: dict | None = None,
metadata: dict[str, str] | None = None,
@ -285,7 +285,7 @@ def create(
@client
async def asearch(
vector_store_id: str,
query: str | list[str],
query: str | builtins.list[str],
filters: dict | None = None,
max_num_results: int | None = None,
ranking_options: dict | None = None,
@ -360,7 +360,7 @@ async def asearch(
@client
def search(
vector_store_id: str,
query: str | list[str],
query: str | builtins.list[str],
filters: dict | None = None,
max_num_results: int | None = None,
ranking_options: dict | None = None,

3
pyproject.toml
View file

@ -46,7 +46,7 @@ proxy = [
"gunicorn>=23.0.0,<24.0",
"uvicorn>=0.33.0,<1.0",
"granian>=2.7.4,<3.0",
"uvloop>=0.21.0,<1.0; sys_platform != 'win32'",
"uvloop>=0.22.1,<1.0; sys_platform != 'win32'",
"fastapi>=0.136.3,<1.0",
"starlette>=1.0.1,<2.0",
"backoff>=2.2.1,<3.0",
@ -179,6 +179,7 @@ dev = [
"basedpyright==1.39.7",
"keyring==25.7.0",
"pytest==9.0.3",
"tomli==2.4.1; python_version < '3.11'",
"pytest-mock==3.15.1",
"pytest-asyncio==1.3.0",
"pytest-postgresql==7.0.2",

4
ruff-strict-budget.json
View file

@ -9,7 +9,7 @@
"limit": 809
},
"ANN201": {
"limit": 2000
"limit": 1999
},
"ANN202": {
"limit": 835
@ -87,7 +87,7 @@
"limit": 2
},
"DTZ003": {
"limit": 26
"limit": 24
},
"DTZ005": {
"limit": 233

1
schema.prisma
View file

@ -5,6 +5,7 @@ datasource client {
generator client {
provider = "prisma-client-py"
recursive_type_depth = -1
binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"]
}

7
scripts/budget_ratchet_check.py
View file

@ -24,7 +24,6 @@ seen the red and accepted it.
Usage:
python scripts/budget_ratchet_check.py [--base REF] [budget.json ...]
Stdlib only.
"""
from __future__ import annotations
@ -33,11 +32,15 @@ import argparse
import json
import subprocess
import sys
import tomllib
from pathlib import Path
from types import MappingProxyType
from typing import NamedTuple
if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib
REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_BASE = "origin/litellm_internal_staging"
DEFAULT_BUDGETS: tuple[str, ...] = (

6
scripts/mutation_report.py
View file

@ -18,13 +18,17 @@ import json
import re
import subprocess
import sys
import tomllib
from collections import defaultdict
from difflib import SequenceMatcher
from pathlib import Path
from typing import Final, NamedTuple
from textwrap import dedent
if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib
ROOT = Path(__file__).resolve().parent.parent
MUTMUT_INVOCATION = ["uv", "run", "--no-sync", "--with", "mutmut==3.5.0", "mutmut"]

2
test-quality-budget.json
View file

@ -21,6 +21,6 @@
"limit": 117
},
"TQ008": {
"limit": 11135
"limit": 11003
}
}

6
tests/code_coverage_tests/check_licenses.py
View file

@ -6,12 +6,16 @@ from pathlib import Path
import re
import sys
import time
import tomllib
from typing import Callable, Dict, Final, List, Optional, Protocol, Set, Tuple
from packaging.requirements import Requirement
import requests
if sys.version_info >= (3, 11):
import tomllib
else:
import tomli as tomllib
DEFAULT_TRANSITIVE_PIN_PACKAGES = (
"aiofiles",
"anyio",

4
tests/proxy_unit_tests/test_reducto_ocr_route.py
View file

@ -100,6 +100,8 @@ def test_proxy_reducto_ocr_json_passthrough_data_uri(client_no_auth):
pages=[OCRPage(index=0, markdown="Proxy OCR")],
model="parse-v3",
usage_info=OCRUsageInfo(pages_processed=1, credits=1),
tables=[{"cells": [["Total", 42]], "page": 1}],
keyValuePairs=[{"key": "approved", "value": True, "confidence": 0.9}],
)
data_uri = "data:application/pdf;base64,JVBERi0xLjQK"
@ -135,3 +137,5 @@ def test_proxy_reducto_ocr_json_passthrough_data_uri(client_no_auth):
assert response_body["object"] == "ocr"
assert response_body["usage_info"]["credits"] == 1
assert response_body["pages"][0]["markdown"] == "Proxy OCR"
assert response_body["tables"] == [{"cells": [["Total", 42]], "page": 1}]
assert response_body["keyValuePairs"] == [{"key": "approved", "value": True, "confidence": 0.9}]

79
tests/rust-python-harness/AGENTS.md
View file

@ -3,42 +3,75 @@
```text
tests/rust-python-harness/
├── __main__.py
├── cli/
│ ├── __init__.py
│ ├── catalog.py
│ └── commands.py
│
├── strategies/
│ ├── e2e_parity/
│ │ ├── runner.py
│ │ ├── __init__.py
│ │ ├── reporting.py
│ │ ├── sdk/
│ │ │ ├── ocr/
│ │ │ ├── messages/
│ │ │ ├── chat_completions/
│ │ │ └── responses/
│ │ └── gateway/
│ │ │ └── ocr/
│ │
│ ├── trace_parity/
│ │ ├── runner.py
│ │ ├── sdk/
│ │ └── gateway/
│ │ ├── __init__.py
│ │ ├── models.py
│ │ ├── reporting.py
│ │ └── sdk/
│ │ ├── chat_completions/
│ │ ├── messages/
│ │ ├── ocr/
│ │ └── transcription/
│ │
│ └── unit_tests/
│ ├── runner.py
│ ├── mapping_validator.py
│ ├── python_runner.py
│ └── rust_runner.py
│ ├── unit_tests_mapping/
│ │ ├── __init__.py
│ │ ├── contracts.py
│ │ ├── cases/
│ │ │ └── ocr.py
│ │ ├── mapping_report.py
│ │ ├── mappings.py
│ │ ├── mapping_validator.py
│ │ ├── reporting.py
│ │ └── runner.py
│ │
│ ├── unit_tests_parity/
│ │ ├── __init__.py
│ │ ├── reporting.py
│ │ └── runner.py
│ │
│ └── unit_tests_rust/
│ ├── __init__.py
│ ├── reporting.py
│ └── runner.py
│
└── shared/
├── parity/
├── tracing/
└── reporting/
├── reporting/
│ └── strategy.py
└── unit_runners/
└── suite_runner.py
```
- A strategy is a folder under `strategies/` with a one-line `AGENTS.md` and an `__init__.py` exporting exactly one `STRATEGY: StrategyDefinition`; its id must equal the folder name
- `shared/reporting/strategy.py` is the contract: runnable module/suite specs, not-implemented/skipped specs, the runner protocol, and `StrategyDefinition`
- Every `STRATEGY` explicitly classifies every SDK function; surface-aware strategies declare their surfaces and classify the complete surface-by-function matrix
- Run locally only; no CI integration
- `__main__.py` selects strategies and combines their reports; each strategy also runs independently
- `python -m tests.rust-python-harness run <strategy>|all` runs the selected strategy; `--function` is common, while each strategy exposes only its supported options
- Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr`
- `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases
- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
- `trace_parity/` compares mapped operations, call counts, and required execution ordering
- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders
- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs
- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts
- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results
- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation
- `shared/` contains reusable parity, tracing, and reporting machinery
- `trace_parity/` compares mapped operations, call counts, and required execution ordering; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`)
- E2E and trace strategies load their registered module cases and run surface-specific execution from their folders
- `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest
- `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report
- `unit_tests_parity/runner.py` runs each contract's `unit_parity_scope` with `LITELLM_RUST=0` and `LITELLM_RUST=1` in separate processes and requires matching outcomes, including failures; exclusions require a reason in the contract
- `unit_tests_rust/runner.py` runs each contract's focused Cargo test suite; native Rust unit tests stay beside their implementation
- `shared/unit_runners/suite_runner.py` runs typed suites registered in code with nodeids of the form `suite:<strategy_id>:<function>:<suite>`
- Every strategy declares its report sections and presentation in its own `reporting.py`; shared reporting code only provides reusable models and cell-formatting primitives
- `shared/` contains reusable parity, tracing, reporting primitives, and unit-runner machinery
- Keep fixtures with their owning API and existing Python tests in their current locations
- Each strategy folder carries an `AGENTS.md` one-liner stating what it should be doing
- Run the harness's own checks with `uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/cli tests/rust-python-harness/strategies/unit_tests_mapping tests/rust-python-harness/strategies/unit_tests_parity tests/rust-python-harness/strategies/unit_tests_rust tests/test_rust_python_harness.py -q`

105
tests/rust-python-harness/README.md
View file

@ -1,105 +0,0 @@
# Rust/Python migration harness
This local harness follows [the agreed structure](AGENTS.md). The root command selects strategies and combines their reports. Each strategy has an independent entry point
```text
strategies/
e2e_parity/runner.py
sdk/ocr/fixtures/
sdk/messages/
sdk/chat_completions/
sdk/responses/
gateway/
existing_e2e_test_sdk/runner.py
trace_parity/runner.py
sdk/
gateway/
unit_tests/
runner.py
mapping_validator.py
python_runner.py
rust_runner.py
shared/
parity/
tracing/
reporting/
```
## Run locally
```bash
uv run python -m tests.rust-python-harness --list
uv run python -m tests.rust-python-harness --function ocr --plain
uv run python -m tests.rust-python-harness --strategy e2e_parity --surface sdk --function ocr --plain
uv run python -m tests.rust-python-harness.strategies.e2e_parity.runner --function ocr --plain
uv run python -m tests.rust-python-harness.strategies.trace_parity.runner --plain
uv run python -m tests.rust-python-harness.strategies.unit_tests.runner --plain
uv run python -m tests.rust-python-harness.strategies.existing_e2e_test_sdk.runner --function transcription --plain
```
Use `--interactive` for strategy and function selection, `--pytest-arg=-x` to stop pytest on its first failure, and `--coverage` to write Python coverage under `target/rust-python-harness/`. The harness enables pytest namespace-package discovery only for its own invocations
This harness has no CI execution. A configured test that fails or disappears makes the command fail. An unconfigured strategy cell remains planned and contributes no passing evidence. Interruptions and collection errors stop execution; ordinary test failures remain in the combined report while later strategies run
## Strategy responsibilities
E2E parity compares SDK objects, exceptions, callbacks, streams, and provider requests. Gateway tests compare HTTP responses. Both surfaces use the same strategy runner and keep execution details and fixtures in their own folders. OCR has recorded sync/async SDK coverage; the existing Messages and Responses bridge checks remain partial
Trace parity compares operation names through an explicit Python/Rust mapping, call counts, and required completion-before-start ordering with `shared/tracing/compare.py`. Surface tests supply captured operation intervals. No production trace instrumentation or trace case is configured yet
Unit testing combines test mapping validation, separate Python processes with Rust disabled and enabled, backend verification, result comparison, and native Cargo tests. Native tests stay beside their Rust implementation. Existing Python tests stay at their original paths. No complete Python/native unit mapping is configured yet, so these cells remain planned
The existing E2E SDK strategy retains the live provider tests configured upstream. It runs OCR, Chat Completions, and Transcription checks from their existing paths and reports them separately from parity tests. These tests require provider credentials
## Configure cases
Each strategy has a `strategy.json`. Its `functions` object defines SDK cases for OCR, Messages, Responses, Count Tokens, Chat Completions, and Transcription. E2E and trace manifests also accept a `gateway` object keyed by API name. A case has `coverage`, `selectors`, and an optional `note`
```json
{
"coverage": "partial",
"selectors": ["tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py"]
}
```
Selectors use pytest file or node syntax. A selector ending in `/` includes tests recursively from that directory
Use `planned` with no selectors until an executable contract exists, `partial` for incomplete coverage, `complete` for the full contract, and `not_applicable` when a strategy does not apply. The dashboard shows passing evidence separately from coverage completeness and LOC coverage
Unit cases use `unit_suite` instead of `selectors`, pointing to a repository-relative JSON file with this shape:
```json
{
"python_selectors": ["tests/test_api.py::test_decode"],
"cargo_manifest": "litellm-rust/Cargo.toml",
"cargo_package": "litellm-core",
"cargo_filter": "ocr::",
"backend": {
"environment_variable": "LITELLM_USE_RUST_OCR",
"probe": "tests.rust-python-harness.strategies.unit_tests.python_runner:ocr_backend"
},
"mappings": [{"python": "tests/test_api.py::test_decode", "rust": "ocr::test_decode"}]
}
```
Names match automatically when the collected Python and Rust test names agree. Explicit `mappings` handle different names, class names, and parametrized cases. Missing or ambiguous counterparts fail validation in either direction. The Cargo filter must select the same behavior as the Python selectors
The backend probe returns `python` or `rust` and runs at startup and before every test call, after fixtures have run. The OCR probe verifies the dispatch flag and native extension availability. Surface tests must also assert that calls reach their intended implementation to catch per-call fallback. Python outcomes must agree, and failed runs remain failures even if both backends fail identically
## OCR fixtures
Fixtures, provider configuration, input strategies, and recording commands live in [the OCR package](strategies/e2e_parity/sdk/ocr/fixtures/README.md). Record with provider credentials:
```bash
uv run python -m tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.fixtures.record --examples 1000
```
`LITELLM_OCR_FIXTURE_DIR` and `--fixture-dir` override the default directory. Shared recording, replay, comparison, streaming, and cassette persistence live in `shared/parity/`
Run the harness's own checks locally:
```bash
uv run pytest -o consider_namespace_packages=true tests/rust-python-harness/shared tests/rust-python-harness/strategies/unit_tests tests/test_rust_python_harness.py -q
```
Existing OCR parity gaps remain visible: invalid-model provider errors differ, Reducto lacks a native contract, and the expanded Azure corpus exposes duplicate Content-Type headers. Moving the harness does not change provider responses or weaken assertions

7
tests/rust-python-harness/__init__.py
View file

@ -1,5 +1,4 @@
"""Interactive Rust/Python SDK parity test harness."""
from .cli import main
from .cli.catalog import load_catalog
from .catalog import load_catalog
__all__ = ["load_catalog"]
__all__ = ["load_catalog", "main"]

75
tests/rust-python-harness/catalog.py
View file

@ -1,75 +0,0 @@
from __future__ import annotations
import json
from pathlib import Path
from typing import Final
from pydantic import BaseModel, ConfigDict, ValidationError
from .shared.reporting.models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy
STRATEGIES_ROOT: Final = Path(__file__).parent / "strategies"
class CaseSpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
coverage: Coverage
selectors: tuple[str, ...] = ()
note: str = ""
unit_suite: str | None = None
class StrategySpec(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
order: int
id: str
label: str
description: str
functions: dict[str, CaseSpec]
gateway: dict[str, CaseSpec] = {}
def _load_strategy(source: Path) -> Strategy:
data: Final = StrategySpec.model_validate_json(source.read_text(encoding="utf-8"))
if set(data.functions) != set(SDK_FUNCTIONS):
raise ValueError(f"{source}: functions must exactly match {SDK_FUNCTIONS}")
cases: Final = tuple(
HarnessCase(
strategy_id=data.id,
strategy_label=data.label,
sdk_function=name,
coverage=case.coverage,
selectors=case.selectors,
note=case.note,
surface=surface,
unit_suite=case.unit_suite,
)
for surface, functions in (("sdk", data.functions), ("gateway", data.gateway))
for name in (SDK_FUNCTIONS if surface == "sdk" else functions)
for case in (functions[name],)
)
for case in cases:
if case.coverage in {Coverage.PLANNED, Coverage.NOT_APPLICABLE} and (case.selectors or case.unit_suite):
raise ValueError(f"{source}: {case.coverage.value} case {case.key} cannot configure tests")
if any(not selector.strip() for selector in case.selectors):
raise ValueError(f"{source}: empty selector in {case.key}")
if data.id == "unit_tests" and case.selectors:
raise ValueError(f"{source}: unit_tests must configure unit_suite instead of pytest selectors")
if data.id != "unit_tests" and case.unit_suite:
raise ValueError(f"{source}: unit_suite is only valid for unit_tests")
return Strategy(data.order, data.id, data.label, data.description, source.parent, cases)
def load_catalog(root: Path = STRATEGIES_ROOT) -> tuple[Strategy, ...]:
sources: Final = tuple(sorted(root.glob("*/strategy.json")))
if not sources:
raise ValueError(f"No strategy manifests found below {root}")
try:
strategies: Final = tuple(sorted((_load_strategy(source) for source in sources), key=lambda item: item.order))
except (ValidationError, json.JSONDecodeError) as error:
raise ValueError(str(error)) from error
if len({strategy.id for strategy in strategies}) != len(strategies):
raise ValueError(f"Duplicate strategy id in {root}")
return strategies

Some files were not shown because too many files have changed in this diff Show more