scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_all_proxy_mcps_for_teams

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-06-25 16:20:35 -07:00
commit 707ab320fd
59 changed files with 4241 additions and 645 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm-rust/.cargo/config.toml → .cargo/config.toml
View file

@ -1,3 +1,12 @@
[http]
# CI has seen transient crates.io failures from libcurl's HTTP/2 multiplexing
# during `maturin` metadata resolution. Disable multiplexing and retry more
# aggressively so editable `uv sync` builds are not failed by one flaky frame.
multiplexing = false
[net]
retry = 5
# PyO3 cdylib (`litellm-python-bridge`) links against the host interpreter's
# symbols, which are not present at link time when building an extension module.
# On macOS, tell the linker to resolve undefined `_Py*` symbols dynamically at

42
.circleci/config.yml
View file

@ -190,6 +190,8 @@ jobs:
working_directory: ~/project
environment:
UV_PYTHON: "3.11"
CARGO_HTTP_MULTIPLEXING: "false"
CARGO_NET_RETRY: "5"
steps:
- checkout
- run:
@ -205,6 +207,24 @@ jobs:
environment:
UV_HTTP_TIMEOUT: "300"
command: |
$rustupInit = Join-Path $env:TEMP "rustup-init.exe"
$rustupVersion = "1.28.2"
$rustupUrl = "https://static.rust-lang.org/rustup/archive/$rustupVersion/x86_64-pc-windows-msvc/rustup-init.exe"
Invoke-WebRequest -Uri $rustupUrl -OutFile $rustupInit
$rustupExpected = "88d8258dcf6ae4f7a80c7d1088e1f36fa7025a1cfd1343731b4ee6f385121fc0"
$rustupActual = (Get-FileHash -Path $rustupInit -Algorithm SHA256).Hash.ToLower()
if ($rustupActual -ne $rustupExpected) {
throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual"
}
& $rustupInit -y --profile minimal --default-toolchain stable
if ($LASTEXITCODE -ne 0) {
exit $LASTEXITCODE
}
Remove-Item $rustupInit
$cargoBin = Join-Path $HOME ".cargo\bin"
$env:Path = "$cargoBin;$env:Path"
rustc --version
cargo --version
$installer = Join-Path $env:TEMP "uv-install.ps1"
Invoke-WebRequest -Uri https://astral.sh/uv/0.10.9/install.ps1 -OutFile $installer
$expected = "d43ffff8d28e7d1e7d1831a212465f12b24a43c7f87f386e95e2a5915aee5d7d"
@ -222,7 +242,20 @@ jobs:
if (-not (Select-String -Path $PROFILE -SimpleMatch $uvBin -Quiet)) {
Add-Content -Path $PROFILE -Value "`$env:Path = `"$uvBin;`$env:Path`""
}
uv sync --frozen --group dev --python 3.11
if (-not (Select-String -Path $PROFILE -SimpleMatch $cargoBin -Quiet)) {
Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`""
}
for ($attempt = 1; $attempt -le 5; $attempt++) {
Write-Host "uv sync attempt $attempt/5"
uv sync --frozen --group dev --python 3.11
if ($LASTEXITCODE -eq 0) {
break
}
if ($attempt -eq 5) {
exit $LASTEXITCODE
}
Start-Sleep -Seconds 15
}
- run:
name: Run Windows-specific test
command: |
@ -232,6 +265,9 @@ jobs:
environment:
UV_HTTP_TIMEOUT: "300"
command: |
$env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
cargo --version
Get-ChildItem -Path "litellm\rust_bridge" -Filter "_native*" -File -ErrorAction SilentlyContinue | Remove-Item -Force
uv build --wheel --out-dir dist
uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
@ -1020,7 +1056,9 @@ jobs:
name: Run tests
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
TEST_FILES=$(printf "%s\n%s\n" \
"$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \
"tests/test_litellm/ocr/test_rust_bridge.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \

2
.dockerignore
View file

@ -49,6 +49,8 @@ build/
*.egg-info/
.DS_Store
**/node_modules
litellm-rust/target/
litellm/rust_bridge/_native*.so
*.log
.env
.env.local

31
.github/scripts/uv_sync_with_retries.sh vendored Executable file
View file

@ -0,0 +1,31 @@
#!/usr/bin/env bash
set -euo pipefail
max_attempts="${UV_SYNC_MAX_ATTEMPTS:-5}"
delay_seconds="${UV_SYNC_RETRY_DELAY_SECONDS:-15}"
export CARGO_HTTP_MULTIPLEXING="${CARGO_HTTP_MULTIPLEXING:-false}"
export CARGO_NET_RETRY="${CARGO_NET_RETRY:-5}"
if [[ "$#" -eq 0 ]]; then
echo "usage: $0 <uv sync args...>" >&2
exit 2
fi
for attempt in $(seq 1 "${max_attempts}"); do
echo "uv sync attempt ${attempt}/${max_attempts}"
status=0
if uv sync "$@"; then
exit 0
else
status=$?
fi
if [[ "${attempt}" -eq "${max_attempts}" ]]; then
echo "uv sync failed after ${max_attempts} attempts" >&2
exit "${status}"
fi
echo "uv sync failed; retrying in ${delay_seconds}s..."
sleep "${delay_seconds}"
done

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

@ -73,7 +73,7 @@ jobs:
- name: Install dependencies
run: |
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:

2
.github/workflows/check-ui-api-types.yml vendored
View file

@ -46,7 +46,7 @@ jobs:
${{ runner.os }}-uv-
- name: Install backend dependencies
run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:

2
.github/workflows/mutation-test.yml vendored
View file

@ -55,7 +55,7 @@ jobs:
- name: Install dependencies
run: |
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:

11
.github/workflows/test-linting.yml vendored
View file

@ -51,10 +51,15 @@ jobs:
uv sync --frozen
- name: Check ruff format
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
cd litellm
uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' .
cd ..
git diff --name-only "$BASE_SHA"...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
echo "No changed litellm Python files to check with ruff format."
exit 0
fi
xargs uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt"
- name: Debug - Check file state
run: |

2
.github/workflows/test-mcp.yml vendored
View file

@ -39,7 +39,7 @@ jobs:
- name: Install dependencies
run: |
uv lock --check
uv sync --frozen --group proxy-dev --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
- name: Run MCP tests
run: |

2
.github/workflows/test-unit-documentation.yml vendored
View file

@ -54,7 +54,7 @@ jobs:
- name: Install dependencies
run: |
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:

2
.github/workflows/test-unit-proxy-legacy.yml vendored
View file

@ -71,7 +71,7 @@ jobs:
- name: Install dependencies
run: |
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:

8
.gitignore vendored
View file

@ -9,12 +9,17 @@ litellm/proxy/myenv/*
litellm_uuid.txt
__pycache__/
*.pyc
# Rust bridge build artifacts (compiled, platform-specific; regenerated by maturin/cargo)
litellm/rust_bridge/_native*.so
litellm/rust_bridge/_native*.pyd
litellm-rust/target/
bun.lockb
**/.DS_Store
.aider*
litellm_results.jsonl
secrets.toml
.gitignore
litellm/proxy/litellm_secrets.toml
litellm/proxy/api_log.json
.idea/
@ -36,7 +41,6 @@ litellm/tests/dynamo*.log
.vscode/settings.json
litellm/proxy/log.txt
proxy_server_config_@.yaml
.gitignore
proxy_server_config_2.yaml
litellm/proxy/secret_managers/credentials.json
hosted_config.yaml

1
Dockerfile
View file

@ -21,6 +21,7 @@ RUN apk add --no-cache \
gcc \
python3 \
python3-dev \
rust \
openssl \
openssl-dev \
nodejs \

3
codecov.yaml
View file

@ -3,6 +3,9 @@ codecov:
notify:
wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI
ignore:
- "litellm-rust/**"
# Uploads are flagged per workflow/shard (GHA) or "circleci". carryforward makes
# a re-upload of a flag replace its prior session instead of accumulating a
# conflicting one, and lets a commit reuse a flag from its parent when that flag

1
docker/Dockerfile.non_root
View file

@ -19,6 +19,7 @@ RUN for i in 1 2 3; do \
python3 \
python3-dev \
gcc \
rust \
bash \
coreutils \
curl \

162
litellm-rust/Cargo.lock generated
View file

@ -206,12 +206,24 @@ dependencies = [
"syn",
]
[[package]]
name = "equivalent"
version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
[[package]]
name = "find-msvc-tools"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582"
[[package]]
name = "fnv"
version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
[[package]]
name = "form_urlencoded"
version = "1.2.2"
@ -221,6 +233,21 @@ dependencies = [
"percent-encoding",
]
[[package]]
name = "futures"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
dependencies = [
"futures-channel",
"futures-core",
"futures-executor",
"futures-io",
"futures-sink",
"futures-task",
"futures-util",
]
[[package]]
name = "futures-channel"
version = "0.3.32"
@ -237,12 +264,34 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
[[package]]
name = "futures-executor"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d"
dependencies = [
"futures-core",
"futures-task",
"futures-util",
]
[[package]]
name = "futures-io"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
[[package]]
name = "futures-macro"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "futures-sink"
version = "0.3.32"
@ -261,8 +310,10 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
dependencies = [
"futures-channel",
"futures-core",
"futures-io",
"futures-macro",
"futures-sink",
"futures-task",
"memchr",
@ -307,6 +358,31 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "h2"
version = "0.4.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
dependencies = [
"atomic-waker",
"bytes",
"fnv",
"futures-core",
"futures-sink",
"http",
"indexmap",
"slab",
"tokio",
"tokio-util",
"tracing",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
[[package]]
name = "heck"
version = "0.5.0"
@ -368,6 +444,7 @@ dependencies = [
"bytes",
"futures-channel",
"futures-core",
"h2",
"http",
"http-body",
"httparse",
@ -521,6 +598,16 @@ dependencies = [
"icu_properties",
]
[[package]]
name = "indexmap"
version = "2.14.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
dependencies = [
"equivalent",
"hashbrown",
]
[[package]]
name = "indoc"
version = "2.0.7"
@ -544,9 +631,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
[[package]]
name = "js-sys"
version = "0.3.102"
version = "0.3.103"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31"
checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102"
dependencies = [
"cfg-if",
"futures-util",
@ -564,6 +651,7 @@ name = "litellm-ai-gateway"
version = "0.1.0"
dependencies = [
"axum",
"base64",
"futures-channel",
"futures-util",
"litellm-core",
@ -594,7 +682,9 @@ dependencies = [
"litellm-ai-gateway",
"litellm-core",
"pyo3",
"pyo3-async-runtimes",
"serde_json",
"tokio",
]
[[package]]
@ -728,6 +818,19 @@ dependencies = [
"unindent",
]
[[package]]
name = "pyo3-async-runtimes"
version = "0.23.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e"
dependencies = [
"futures",
"once_cell",
"pin-project-lite",
"pyo3",
"tokio",
]
[[package]]
name = "pyo3-build-config"
version = "0.23.5"
@ -913,6 +1016,7 @@ dependencies = [
"futures-channel",
"futures-core",
"futures-util",
"h2",
"http",
"http-body",
"http-body-util",
@ -932,12 +1036,14 @@ dependencies = [
"sync_wrapper",
"tokio",
"tokio-rustls",
"tokio-util",
"tower",
"tower-http",
"tower-service",
"url",
"wasm-bindgen",
"wasm-bindgen-futures",
"wasm-streams",
"web-sys",
"webpki-roots",
]
@ -1335,6 +1441,19 @@ dependencies = [
"tungstenite",
]
[[package]]
name = "tokio-util"
version = "0.7.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
dependencies = [
"bytes",
"futures-core",
"futures-sink",
"pin-project-lite",
"tokio",
]
[[package]]
name = "tower"
version = "0.5.3"
@ -1507,9 +1626,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen"
version = "0.2.125"
version = "0.2.126"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a"
checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4"
dependencies = [
"cfg-if",
"once_cell",
@ -1520,9 +1639,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-futures"
version = "0.4.75"
version = "0.4.76"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280"
checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d"
dependencies = [
"js-sys",
"wasm-bindgen",
@ -1530,9 +1649,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-macro"
version = "0.2.125"
version = "0.2.126"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d"
checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1"
dependencies = [
"quote",
"wasm-bindgen-macro-support",
@ -1540,9 +1659,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-macro-support"
version = "0.2.125"
version = "0.2.126"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd"
checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e"
dependencies = [
"bumpalo",
"proc-macro2",
@ -1553,18 +1672,31 @@ dependencies = [
[[package]]
name = "wasm-bindgen-shared"
version = "0.2.125"
version = "0.2.126"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f"
checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24"
dependencies = [
"unicode-ident",
]
[[package]]
name = "web-sys"
version = "0.3.102"
name = "wasm-streams"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d"
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
dependencies = [
"futures-util",
"js-sys",
"wasm-bindgen",
"wasm-bindgen-futures",
"web-sys",
]
[[package]]
name = "web-sys"
version = "0.3.103"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141"
dependencies = [
"js-sys",
"wasm-bindgen",

6
litellm-rust/Cargo.toml
View file

@ -16,13 +16,15 @@ litellm-core = { path = "crates/core" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
axum = "0.7"
pyo3 = "0.23.5"
pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] }
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"

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

@ -23,6 +23,7 @@ tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "tim
tokio-tungstenite.workspace = true
futures-util.workspace = true
serde_json.workspace = true
base64.workspace = true
axum = { workspace = true, features = ["ws"], optional = true }
serde = { workspace = true, optional = true }
subtle = { workspace = true, optional = true }

359
litellm-rust/crates/ai-gateway/src/io/ocr.rs
View file

@ -1,6 +1,6 @@
//! End-to-end OCR orchestration.
//!
//! Owns the whole Mistral OCR call so the Python side stays a thin bridge:
//! Owns supported OCR provider calls so the Python side stays a thin bridge:
//! resolve the API key, build the URL + body via the pure transforms, POST it,
//! and normalize the response. The HTTP client is built once and reused.
@ -8,75 +8,135 @@ use std::sync::OnceLock;
use std::time::Duration;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrProviderConfig;
use litellm_core::ocr::transformation::{OcrAuthStrategy, OcrResponseHandling};
use litellm_core::CoreResult;
use serde_json::{Map, Value};
use litellm_core::providers::mistral::ocr::transformation as mistral;
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
mod common_utils;
use common_utils::{
convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence,
string_headers, truncate_error_body,
};
/// OCR over large documents can take a while; bound it generously rather than
/// hanging forever on an unresponsive upstream. The client-level limit is the
/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``.
const OCR_TIMEOUT_SECS: u64 = 600;
/// Maximum upstream body characters retained in error messages. OCR responses
/// can echo document contents and prompts; keep enough for debugging without
/// forwarding sensitive payloads across the host boundary.
const ERROR_BODY_MAX_CHARS: usize = 256;
/// Process-wide blocking HTTP client (connection pool + TLS reused across calls).
fn http_client() -> &'static reqwest::blocking::Client {
static CLIENT: OnceLock<reqwest::blocking::Client> = OnceLock::new();
/// Process-wide async HTTP client (connection pool + TLS reused across calls).
///
/// The Python fallback path uses LiteLLM's standard `BaseLLMHTTPHandler`. This
/// Rust path is opt-in and owns end-to-end OCR I/O, so it cannot call the
/// Python handler directly; keep this route-scoped until litellm-rust has a
/// shared HTTP abstraction.
fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::blocking::Client::builder()
reqwest::Client::builder()
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
.build()
.expect("failed to build reqwest client")
})
}
fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
fn upstream_headers(
headers: &[(String, String)],
auth_strategy: OcrAuthStrategy,
api_key: Option<&str>,
) -> Vec<(String, String)> {
let auth_header = api_key.map(|api_key| match auth_strategy {
OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
});
auth_header
.into_iter()
.chain(headers.iter().cloned())
.collect()
}
/// Perform a Mistral OCR call end to end and return the normalized response as
pub struct OcrRequest<'a> {
pub model: &'a str,
pub document: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: &'a str,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}
/// Perform an OCR call end to end and return the normalized response as
/// JSON (the shape the Python `OCRResponse` model expects).
///
/// Blocking: intended to be called with the GIL released from the Python bridge.
pub fn run_ocr(
model: &str,
document: Value,
api_key: Option<&str>,
api_base: Option<&str>,
optional_params: Map<String, Value>,
timeout: Option<Duration>,
) -> CoreResult<Value> {
let config = &MISTRAL_OCR_CONFIG;
/// Async: intended to be awaited directly by the Python bridge's async entrypoint.
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
let model = request.model;
let config = ocr_provider_config(request.custom_llm_provider, model)
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.to_string()))?;
let env_lookup = |key: &str| std::env::var(key).ok();
let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?;
let url = mistral::complete_url(api_base);
let filtered_params = config.map_ocr_params(&optional_params);
let headers = string_headers(request.extra_headers)?;
let auth_strategy = config.auth_strategy();
let api_key = (!has_header(&headers, auth_strategy.header_name()))
.then(|| config.resolve_api_key(request.api_key, &env_lookup))
.transpose()?;
let url = config.complete_url(
request.api_base,
model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_ocr_params(&request.optional_params);
let document = if config.requires_data_uri_document() {
convert_document_url_to_data_uri(request.document).await?
} else {
request.document
};
let body = config
.transform_ocr_request(model, document, filtered_params)?
.data;
let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body);
if let Some(duration) = timeout {
request = request.timeout(duration);
let mut request_builder = http_client().post(&url).json(&body);
for (key, value) in &upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request
let response = request_builder
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
let status = response.status();
if config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
&& status.as_u16() == 202
{
let operation_url = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.ok_or_else(|| {
CoreError::InvalidResponse(
"Azure Document Intelligence returned 202 but no Operation-Location header found"
.to_string(),
)
})?;
let response_json =
poll_document_intelligence(&operation_url, &url, &upstream_headers, request.timeout)
.await?;
return Ok(config
.transform_ocr_response(model, response_json)?
.into_json());
}
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
if !status.is_success() {
@ -97,6 +157,25 @@ pub fn run_ocr(
#[cfg(test)]
mod tests {
use super::*;
use serde_json::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();
let mut buffer = [0_u8; 1024];
loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
String::from_utf8(request).expect("request is utf8")
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
@ -106,7 +185,7 @@ mod tests {
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50);
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
@ -115,13 +194,213 @@ mod tests {
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS);
assert_eq!(prefix_chars, 256);
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10);
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document());
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature"));
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn auth_header_detection_is_case_insensitive() {
let headers = vec![
("x-trace-id".to_string(), "trace-1".to_string()),
("authorization".to_string(), "Bearer sk-test".to_string()),
];
assert!(has_header(&headers, "authorization"));
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
assert!(has_header(&headers, "authorization"));
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
assert!(!has_header(&headers, "authorization"));
}
#[tokio::test]
async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_headers(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer sk-from-python".to_string()),
);
headers.insert(
"x-trace-id".to_string(),
Value::String("trace-1".to_string()),
);
let response = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-for-rust-fallback"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: "mistral",
extra_headers: Some(headers),
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let request = server.await.expect("server task completes");
let authorization_count = request
.lines()
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
.count();
assert_eq!(authorization_count, 1, "{request}");
assert!(
request.contains("authorization: Bearer sk-from-python")
|| request.contains("Authorization: Bearer sk-from-python"),
"{request}"
);
}
#[tokio::test]
async fn document_intelligence_poll_uses_resolved_subscription_key() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let operation_url = format!("http://{addr}/operations/1");
let server = tokio::spawn(async move {
let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
let post_request = read_http_headers(&mut post_socket).await;
let post_response = format!(
"HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
);
post_socket
.write_all(post_response.as_bytes())
.await
.expect("writes post response");
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
let poll_request = read_http_headers(&mut poll_socket).await;
let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
let poll_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
poll_socket
.write_all(poll_response.as_bytes())
.await
.expect("writes poll response");
(post_request, poll_request)
});
let response = ocr(OcrRequest {
model: "prebuilt-read",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("di-key"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: "azure_ai/doc-intelligence",
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("document intelligence request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let (post_request, poll_request) = server.await.expect("server task completes");
assert!(
post_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{post_request}"
);
assert!(
poll_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{poll_request}"
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
CoreError::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}
}

448
litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs Normal file
View file

@ -0,0 +1,448 @@
use std::net::IpAddr;
use std::time::{Duration, Instant};
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrProviderConfig;
use litellm_core::CoreResult;
use reqwest::Url;
use serde_json::{Map, Value};
use litellm_core::providers::azure_ai::ocr::transformation::{
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
};
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
};
use super::http_client;
const ERROR_BODY_MAX_CHARS: usize = 256;
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
pub(super) fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
}
pub(super) fn ocr_provider_config(
provider: &str,
model: &str,
) -> Option<&'static dyn OcrProviderConfig> {
match provider {
"mistral" => Some(&MISTRAL_OCR_CONFIG),
"azure_ai" if is_azure_document_intelligence_model(model) => {
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
}
"azure_ai/doc-intelligence" => Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG),
"azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
"vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
"vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),
_ => None,
}
}
fn is_azure_document_intelligence_model(model: &str) -> bool {
let model = model.to_ascii_lowercase();
model.contains("doc-intelligence") || model.contains("documentintelligence")
}
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"OCR extra_headers.{key} must be a string, got {}",
litellm_core::error::json_type_name(&value)
))
})
})
.collect()
}
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
headers
.iter()
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}
fn document_url_field(document: &Value) -> CoreResult<Option<(&str, &str)>> {
let Some(object) = document.as_object() else {
return Ok(None);
};
let Some(doc_type) = object.get("type").and_then(Value::as_str) else {
return Ok(None);
};
let field = match doc_type {
"document_url" => "document_url",
"image_url" => "image_url",
_ => return Ok(None),
};
let Some(url) = object.get(field).and_then(Value::as_str) else {
return Ok(None);
};
Ok(Some((field, url)))
}
fn is_url_requiring_fetch(url: &str) -> bool {
!url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://"))
}
fn max_document_download_bytes() -> u64 {
let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
.ok()
.and_then(|value| value.parse::<f64>().ok())
.unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB);
(max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64
}
fn is_blocked_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(ip) => {
ip.is_private()
|| ip.is_loopback()
|| ip.is_link_local()
|| ip.is_broadcast()
|| ip.is_multicast()
|| ip.is_unspecified()
}
IpAddr::V6(ip) => {
let first_segment = ip.segments()[0];
let is_unique_local = (first_segment & 0xfe00) == 0xfc00;
let is_link_local = (first_segment & 0xffc0) == 0xfe80;
ip.is_loopback()
|| ip.is_unspecified()
|| ip.is_multicast()
|| is_unique_local
|| is_link_local
|| ip
.to_ipv4_mapped()
.or_else(|| ip.to_ipv4())
.map(|v4| is_blocked_ip(IpAddr::V4(v4)))
.unwrap_or(false)
}
}
}
fn blocked_url_error(url: &Url) -> CoreError {
CoreError::InvalidRequest(format!(
"OCR document URL rejected by SSRF protection: {url}"
))
}
async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
if !matches!(url.scheme(), "http" | "https") {
return Err(blocked_url_error(url));
}
let host = url.host_str().ok_or_else(|| blocked_url_error(url))?;
if let Ok(ip) = host.parse::<IpAddr>() {
if is_blocked_ip(ip) {
return Err(blocked_url_error(url));
}
return Ok(());
}
let port = url
.port_or_known_default()
.ok_or_else(|| blocked_url_error(url))?;
let addresses = tokio::net::lookup_host((host, port))
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
let mut saw_address = false;
for address in addresses {
saw_address = true;
if is_blocked_ip(address.ip()) {
return Err(blocked_url_error(url));
}
}
if !saw_address {
return Err(blocked_url_error(url));
}
Ok(())
}
fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult<Url> {
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or_else(|| {
CoreError::InvalidResponse("OCR document redirect missing Location header".to_string())
})?;
url.join(location)
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}")))
}
async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> {
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|err| CoreError::Network(err.to_string()))?;
let mut current_url = Url::parse(url)
.map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
validate_safe_fetch_url(&current_url).await?;
let response = client
.get(current_url.clone())
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
if !response.status().is_redirection() {
return Ok((current_url, response));
}
current_url = redirect_location(&response, &current_url)?;
}
Err(CoreError::InvalidRequest(
"Too many redirects while fetching OCR document URL".to_string(),
))
}
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> {
if max_bytes == 0 {
return Err(CoreError::InvalidRequest(format!(
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
)));
}
if content_length > max_bytes {
let size_mb = content_length as f64 / (1024.0 * 1024.0);
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
return Err(CoreError::InvalidRequest(format!(
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
)));
}
Ok(())
}
async fn read_response_with_limit(
mut response: reqwest::Response,
url: &Url,
) -> CoreResult<Vec<u8>> {
let max_bytes = max_document_download_bytes();
if let Some(content_length) = response.content_length() {
enforce_download_size(content_length, max_bytes, url)?;
} else {
enforce_download_size(0, max_bytes, url)?;
}
let mut bytes = Vec::new();
let mut bytes_downloaded: u64 = 0;
while let Some(chunk) = response
.chunk()
.await
.map_err(|err| CoreError::Network(err.to_string()))?
{
bytes_downloaded += chunk.len() as u64;
enforce_download_size(bytes_downloaded, max_bytes, url)?;
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult<Value> {
let Some((field, url)) = document_url_field(&document)? else {
return Ok(document);
};
if !is_url_requiring_fetch(url) {
return Ok(document);
}
let (final_url, response) = safe_get_document_url(url).await?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = read_response_with_limit(response, &final_url).await?;
let data_uri = format!(
"data:{content_type};base64,{}",
BASE64_STANDARD.encode(bytes)
);
let mut transformed = document
.as_object()
.cloned()
.ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?;
transformed.insert(field.to_string(), Value::String(data_uri));
Ok(Value::Object(transformed))
}
fn same_origin(left: &str, right: &str) -> bool {
let Ok(left) = reqwest::Url::parse(left) else {
return false;
};
let Ok(right) = reqwest::Url::parse(right) else {
return false;
};
left.scheme() == right.scheme()
&& left.host_str() == right.host_str()
&& left.port_or_known_default() == right.port_or_known_default()
}
fn retry_after_secs(response: &reqwest::Response) -> u64 {
response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(2)
}
fn operation_status(response_json: &Value) -> CoreResult<&str> {
let status = response_json
.get("status")
.and_then(Value::as_str)
.ok_or(CoreError::MissingField("status"))?;
match status {
"succeeded" => Ok("succeeded"),
"running" | "notStarted" => Ok("running"),
"failed" => {
let message = response_json
.get("error")
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.unwrap_or("Unknown error");
Err(CoreError::InvalidResponse(format!(
"Azure Document Intelligence analysis failed: {message}"
)))
}
other => Err(CoreError::InvalidResponse(format!(
"Unknown operation status: {other}"
))),
}
}
pub(super) async fn poll_document_intelligence(
operation_url: &str,
original_url: &str,
headers: &[(String, String)],
timeout: Option<Duration>,
) -> CoreResult<Value> {
if !same_origin(operation_url, original_url) {
return Err(CoreError::InvalidResponse(
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
));
}
let start = Instant::now();
let timeout = timeout.unwrap_or(Duration::from_secs(
AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
));
loop {
if start.elapsed() > timeout {
return Err(CoreError::Network(format!(
"Azure Document Intelligence operation polling timed out after {} seconds",
timeout.as_secs()
)));
}
let mut request_builder = http_client().get(operation_url);
for (key, value) in headers {
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
request_builder = request_builder.header(key, value);
}
}
let response = request_builder
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
let retry_after = retry_after_secs(&response);
let status = response.status();
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
})?;
if operation_status(&response_json)? == "succeeded" {
return Ok(response_json);
}
tokio::time::sleep(Duration::from_secs(retry_after)).await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn blocks_private_and_metadata_ips() {
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
assert!(is_blocked_ip("10.0.0.1".parse().unwrap()));
assert!(is_blocked_ip("169.254.169.254".parse().unwrap()));
assert!(is_blocked_ip("::1".parse().unwrap()));
assert!(is_blocked_ip("fd00::1".parse().unwrap()));
assert!(is_blocked_ip("fe80::1".parse().unwrap()));
assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap()));
assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap()));
assert!(!is_blocked_ip("8.8.8.8".parse().unwrap()));
assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap()));
}
#[tokio::test]
async fn convert_document_url_rejects_loopback_fetch() {
let error = convert_document_url_to_data_uri(json!({
"type": "image_url",
"image_url": "http://127.0.0.1/image.png"
}))
.await
.unwrap_err();
assert!(matches!(
error,
CoreError::InvalidRequest(message)
if message.contains("SSRF protection")
));
}
#[tokio::test]
async fn convert_document_url_leaves_data_uri_untouched() {
let document = json!({
"type": "image_url",
"image_url": "data:image/png;base64,abcd"
});
let transformed = convert_document_url_to_data_uri(document.clone())
.await
.unwrap();
assert_eq!(transformed, document);
}
}

4
litellm-rust/crates/core/src/error.rs
View file

@ -13,6 +13,10 @@ pub enum CoreError {
MissingField(&'static str),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("{0}")]
Auth(String),
#[error("OCR request failed with status {status}: {body}")]

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

@ -4,7 +4,28 @@ use crate::CoreResult;
use super::types::{OcrRequestData, OcrResponseData};
pub trait OcrProviderConfig {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrAuthStrategy {
Bearer,
Header(&'static str),
}
impl OcrAuthStrategy {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "authorization",
Self::Header(header_name) => header_name,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrResponseHandling {
Json,
AzureDocumentIntelligencePoll,
}
pub trait OcrProviderConfig: Sync {
fn supported_ocr_params(&self) -> &'static [&'static str];
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
@ -29,4 +50,30 @@ pub trait OcrProviderConfig {
model: &str,
response_json: Value,
) -> CoreResult<OcrResponseData>;
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;
fn auth_strategy(&self) -> OcrAuthStrategy {
OcrAuthStrategy::Bearer
}
fn requires_data_uri_document(&self) -> bool {
false
}
fn response_handling(&self) -> OcrResponseHandling {
OcrResponseHandling::Json
}
}

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

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

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

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

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

@ -0,0 +1,520 @@
use std::collections::BTreeSet;
use crate::error::{json_type_name, CoreError, CoreResult};
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{json, Map, Value};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30";
const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96;
const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages"];
pub struct AzureAiOcrConfig;
pub struct AzureDocumentIntelligenceOcrConfig;
pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig;
pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig =
AzureDocumentIntelligenceOcrConfig;
fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
fn resolve_value(
explicit: Option<&str>,
env_name: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
missing_message: &str,
) -> CoreResult<String> {
non_empty(explicit)
.map(str::to_string)
.or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| CoreError::Auth(missing_message.to_string()))
}
pub fn resolve_azure_ai_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_value(
api_key,
AZURE_AI_API_KEY_ENV,
env_lookup,
"Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params",
)
}
pub fn resolve_azure_ai_api_base(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_value(
api_base,
AZURE_AI_API_BASE_ENV,
env_lookup,
"Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter",
)
}
pub fn complete_azure_ai_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let base = resolve_azure_ai_api_base(api_base, env_lookup)?;
Ok(format!(
"{}/providers/mistral/azure/ocr",
base.trim_end_matches('/')
))
}
pub fn resolve_document_intelligence_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_value(
api_key,
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV,
env_lookup,
"Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter",
)
}
pub fn resolve_document_intelligence_endpoint(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_value(
api_base,
AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV,
env_lookup,
"Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter",
)
}
fn encode_model_id(model: &str) -> String {
let model_id = model.rsplit('/').next().unwrap_or(model);
model_id
.bytes()
.flat_map(|byte| match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
vec![byte as char]
}
_ => format!("%{byte:02X}").chars().collect(),
})
.collect()
}
fn pages_token_is_valid(token: &str) -> bool {
let mut parts = token.split('-');
let Some(start) = parts.next() else {
return false;
};
if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) {
return false;
}
match parts.next() {
None => true,
Some(end) => {
!end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none()
}
}
}
fn normalize_pages_param(pages: &Value) -> CoreResult<Option<String>> {
match pages {
Value::String(value) => {
let normalized = value
.split(',')
.map(str::trim)
.collect::<Vec<_>>()
.join(",");
if normalized.split(',').all(pages_token_is_valid) {
Ok(Some(normalized))
} else {
Err(CoreError::InvalidRequest(format!(
"Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'."
)))
}
}
Value::Array(values) => {
if values.is_empty() {
return Ok(None);
}
if values.iter().all(Value::is_i64) {
let mut pages = BTreeSet::new();
for value in values {
let page = value.as_i64().expect("checked is_i64");
if page < 0 {
return Err(CoreError::InvalidRequest(
"`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(),
));
}
pages.insert(page + 1);
}
return Ok(Some(
pages
.into_iter()
.map(|page| page.to_string())
.collect::<Vec<_>>()
.join(","),
));
}
if values.iter().all(Value::is_string) {
let normalized = values
.iter()
.filter_map(Value::as_str)
.map(str::trim)
.collect::<Vec<_>>()
.join(",");
if normalized.split(',').all(pages_token_is_valid) {
return Ok(Some(normalized));
}
return Err(CoreError::InvalidRequest(format!(
"Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'."
)));
}
Err(CoreError::InvalidRequest(
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
.to_string(),
))
}
_ => Err(CoreError::InvalidRequest(
"`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'."
.to_string(),
)),
}
}
pub fn complete_document_intelligence_url(
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?;
let mut url = format!(
"{}/documentintelligence/documentModels/{}:analyze?api-version={}",
endpoint.trim_end_matches('/'),
encode_model_id(model),
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION
);
if let Some(pages) = optional_params.get("pages") {
if let Some(normalized) = normalize_pages_param(pages)? {
url.push_str("&pages=");
url.push_str(&normalized);
}
}
Ok(url)
}
fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let doc_type = object
.get("type")
.and_then(Value::as_str)
.ok_or(CoreError::MissingField("document.type"))?;
let field_name = match doc_type {
"document_url" => "document_url",
"image_url" => "image_url",
other => {
return Err(CoreError::InvalidRequest(format!(
"Invalid document type: {other}. Must be 'document_url' or 'image_url'"
)))
}
};
object
.get(field_name)
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.ok_or(CoreError::MissingField(field_name))
}
fn extract_base64_from_data_uri(data_uri: &str) -> &str {
data_uri
.split_once(',')
.map(|(_, data)| data)
.unwrap_or(data_uri)
}
fn page_markdown(page: &Map<String, Value>) -> String {
page.get("lines")
.and_then(Value::as_array)
.map(|lines| {
lines
.iter()
.filter_map(|line| line.get("content").and_then(Value::as_str))
.collect::<Vec<_>>()
.join("\n")
})
.unwrap_or_default()
}
fn page_dimensions(page: &Map<String, Value>) -> Value {
let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5);
let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0);
let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch");
let (width, height) = if unit == "inch" {
(
(width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64,
(height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64,
)
} else {
(width as i64, height as i64)
};
json!({
"width": width,
"height": height,
"dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI,
})
}
impl OcrProviderConfig for AzureAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
fn transform_ocr_request(
&self,
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> CoreResult<OcrRequestData> {
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> CoreResult<OcrResponseData> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
complete_azure_ai_url(api_base, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_azure_ai_api_key(api_key, env_lookup)
}
fn requires_data_uri_document(&self) -> bool {
true
}
}
impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
_optional_params: Map<String, Value>,
) -> CoreResult<OcrRequestData> {
let document_url = document_url_from_mistral_document(&document)?;
let mut data = Map::new();
if document_url.starts_with("data:") {
data.insert(
"base64Source".to_string(),
Value::String(extract_base64_from_data_uri(document_url).to_string()),
);
} else {
data.insert(
"urlSource".to_string(),
Value::String(document_url.to_string()),
);
}
Ok(OcrRequestData {
data: Value::Object(data),
files: None,
})
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> CoreResult<OcrResponseData> {
let response = response_json
.as_object()
.ok_or_else(|| CoreError::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let status = response
.get("status")
.and_then(Value::as_str)
.ok_or(CoreError::MissingField("status"))?;
if status != "succeeded" {
return Err(CoreError::InvalidResponse(format!(
"Azure Document Intelligence analysis failed with status: {status}"
)));
}
let azure_pages = response
.get("analyzeResult")
.and_then(|result| result.get("pages"))
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let pages = azure_pages
.iter()
.filter_map(Value::as_object)
.map(|page| {
let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1);
json!({
"index": page_number - 1,
"markdown": page_markdown(page),
"dimensions": page_dimensions(page),
})
})
.collect::<Vec<_>>();
Ok(OcrResponseData {
usage_info: Some(json!({
"pages_processed": pages.len(),
"doc_size_bytes": null,
})),
pages,
model: model.to_string(),
document_annotation: None,
object: "ocr".to_string(),
})
}
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
complete_document_intelligence_url(api_base, model, optional_params, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_document_intelligence_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> OcrAuthStrategy {
OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key")
}
fn response_handling(&self) -> OcrResponseHandling {
OcrResponseHandling::AzureDocumentIntelligencePoll
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn azure_ai_reuses_mistral_body_transform() {
let body = AZURE_AI_OCR_CONFIG
.transform_ocr_request(
"pixtral-12b-2409",
json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}),
serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "pixtral-12b-2409");
assert_eq!(body["include_image_base64"], true);
assert_eq!(
body["document"]["document_url"],
"data:application/pdf;base64,abc"
);
}
#[test]
fn document_intelligence_url_normalizes_zero_based_pages() {
let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]);
let url = complete_document_intelligence_url(
Some("https://example.cognitiveservices.azure.com/"),
"azure_ai/doc-intelligence/prebuilt-layout",
&params,
&|_| None,
)
.expect("url builds");
assert_eq!(
url,
"https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3"
);
}
#[test]
fn document_intelligence_request_uses_base64_source_for_data_uri() {
let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
.transform_ocr_request(
"prebuilt-read",
json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body, json!({"base64Source": "abc123"}));
}
#[test]
fn document_intelligence_response_normalizes_pages() {
let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
.transform_ocr_response(
"prebuilt-layout",
json!({
"status": "succeeded",
"analyzeResult": {
"pages": [{
"pageNumber": 2,
"width": 8.5,
"height": 11,
"unit": "inch",
"lines": [{"content": "hello"}, {"content": "world"}]
}]
}
}),
)
.expect("response transforms");
assert_eq!(response.pages[0]["index"], 1);
assert_eq!(response.pages[0]["markdown"], "hello\nworld");
assert_eq!(response.pages[0]["dimensions"]["width"], 816);
assert_eq!(
response.usage_info,
Some(json!({"pages_processed": 1, "doc_size_bytes": null}))
);
}
}

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

@ -132,6 +132,24 @@ impl OcrProviderConfig for MistralOcrConfig {
object: "ocr".to_string(),
})
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
Ok(complete_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_api_key(api_key, env_lookup)
}
}
pub fn supported_ocr_params() -> &'static [&'static str] {

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

@ -1,2 +1,4 @@
pub mod azure_ai;
pub mod mistral;
pub mod openai;
pub mod vertex_ai;

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

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

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

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

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

@ -0,0 +1,435 @@
use crate::error::{json_type_name, CoreError, CoreResult};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{json, Map, Value};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
const VERTEX_DEFAULT_LOCATION: &str = "us-central1";
const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com";
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[rustfmt::skip]
const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[
"stream",
"temperature",
"max_tokens",
"top_p",
"n",
"stop",
];
pub struct VertexAiOcrConfig;
pub struct VertexAiDeepSeekOcrConfig;
pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig;
pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig;
fn string_param<'a>(params: &'a Map<String, Value>, keys: &[&str]) -> Option<&'a str> {
keys.iter()
.find_map(|key| params.get(*key).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn is_deepseek_model(model: &str) -> bool {
model.to_ascii_lowercase().contains("deepseek")
}
pub fn resolve_vertex_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or_else(|| {
CoreError::Auth(
"Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers"
.to_string(),
)
})
}
fn vertex_project(
params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
string_param(params, &["vertex_project", "vertex_ai_project"])
.map(str::to_string)
.or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| {
CoreError::InvalidRequest(
"Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter"
.to_string(),
)
})
}
fn vertex_location(
params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
string_param(params, &["vertex_location", "vertex_ai_location"])
.map(str::to_string)
.or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
.or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
.unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string())
}
fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String {
api_base
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com"))
.trim_end_matches('/')
.to_string()
}
pub fn complete_vertex_mistral_url(
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let project = vertex_project(optional_params, env_lookup)?;
let location = vertex_location(optional_params, env_lookup);
let base = vertex_mistral_api_base(api_base, &location);
Ok(format!(
"{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict"
))
}
pub fn complete_vertex_deepseek_url(
api_base: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let project = vertex_project(optional_params, env_lookup)?;
let location = vertex_location(optional_params, env_lookup);
let base = api_base
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE)
.trim_end_matches('/');
Ok(format!(
"{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions"
))
}
fn document_content_item(document: &Value) -> CoreResult<Value> {
let object = document.as_object().ok_or_else(|| CoreError::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let doc_type = object
.get("type")
.and_then(Value::as_str)
.ok_or(CoreError::MissingField("document.type"))?;
let url_field = match doc_type {
"image_url" => "image_url",
"document_url" => "document_url",
other => {
return Err(CoreError::InvalidRequest(format!(
"Unsupported document type: {other}. Expected 'image_url' or 'document_url'"
)))
}
};
let url = object
.get(url_field)
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.ok_or(CoreError::MissingField(url_field))?;
Ok(json!({
"type": "image_url",
"image_url": url,
}))
}
fn deepseek_model_name(model: &str) -> String {
if model.starts_with("deepseek-ai/") {
model.to_string()
} else {
format!("deepseek-ai/{model}")
}
}
fn first_choice_content(response: &Value) -> CoreResult<Value> {
response
.get("choices")
.and_then(Value::as_array)
.and_then(|choices| choices.first())
.and_then(|choice| choice.get("message"))
.and_then(|message| message.get("content"))
.cloned()
.filter(|content| match content {
Value::String(value) => !value.is_empty(),
Value::Object(_) => true,
_ => false,
})
.ok_or_else(|| {
CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string())
})
}
fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> Value {
match content {
Value::String(content) => {
if content.trim_start().starts_with('{') {
serde_json::from_str(&content).unwrap_or_else(|_| {
json!({
"pages": [{"index": 0, "markdown": content}],
"model": model,
"usage_info": usage.unwrap_or_else(|| json!({})),
})
})
} else {
json!({
"pages": [{"index": 0, "markdown": content}],
"model": model,
"usage_info": usage.unwrap_or_else(|| json!({})),
})
}
}
Value::Object(_) => content,
other => json!({
"pages": [{"index": 0, "markdown": other.to_string()}],
"model": model,
"usage_info": usage.unwrap_or_else(|| json!({})),
}),
}
}
impl OcrProviderConfig for VertexAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
fn transform_ocr_request(
&self,
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> CoreResult<OcrRequestData> {
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> CoreResult<OcrResponseData> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
complete_vertex_mistral_url(api_base, model, optional_params, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_vertex_api_key(api_key, env_lookup)
}
fn requires_data_uri_document(&self) -> bool {
true
}
}
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
DEEPSEEK_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> CoreResult<OcrRequestData> {
let mut data = Map::new();
data.insert(
"model".to_string(),
Value::String(deepseek_model_name(model)),
);
data.insert(
"messages".to_string(),
json!([{"role": "user", "content": [document_content_item(&document)?]}]),
);
for (key, value) in optional_params {
if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) {
data.insert(key, value);
}
}
Ok(OcrRequestData {
data: Value::Object(data),
files: None,
})
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> CoreResult<OcrResponseData> {
let response = response_json
.as_object()
.ok_or_else(|| CoreError::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let usage = response.get("usage").cloned();
let content = first_choice_content(&response_json)?;
let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model);
if !ocr_data.get("pages").is_some_and(Value::is_array) {
ocr_data = json!({
"pages": [{
"index": 0,
"markdown": match content {
Value::String(value) => value,
other => other.to_string(),
}
}],
"model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model),
"usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})),
});
}
let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType {
expected: "object",
actual: json_type_name(&ocr_data),
})?;
let pages = object
.get("pages")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let usage_info = object
.get("usage_info")
.cloned()
.or_else(|| response.get("usage").cloned());
Ok(OcrResponseData {
pages,
model: object
.get("model")
.and_then(Value::as_str)
.unwrap_or(model)
.to_string(),
document_annotation: object.get("document_annotation").cloned(),
usage_info,
object: "ocr".to_string(),
})
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
complete_vertex_deepseek_url(api_base, optional_params, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
resolve_vertex_api_key(api_key, env_lookup)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn vertex_mistral_url_uses_project_location_and_model() {
let params = Map::from_iter([
("vertex_project".to_string(), json!("proj-1")),
("vertex_location".to_string(), json!("europe-west4")),
]);
let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", &params, &|_| None)
.expect("url builds");
assert_eq!(
url,
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[test]
fn vertex_mistral_reuses_mistral_body_transform() {
let body = VERTEX_AI_OCR_CONFIG
.transform_ocr_request(
"mistral-ocr-maas",
json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "mistral-ocr-maas");
assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc");
}
#[test]
fn vertex_deepseek_request_uses_ocr_endpoint_shape() {
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_request(
"deepseek-ocr-maas",
json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}),
Map::from_iter([("temperature".to_string(), json!(0.1))]),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(body["temperature"], 0.1);
assert_eq!(
body["messages"][0]["content"][0],
json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"})
);
}
#[test]
fn vertex_deepseek_response_wraps_markdown_content() {
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_response(
"deepseek-ocr-maas",
json!({
"choices": [{"message": {"content": "# OCR text"}}],
"usage": {"prompt_tokens": 1}
}),
)
.expect("response transforms");
assert_eq!(
response.pages,
vec![json!({"index": 0, "markdown": "# OCR text"})]
);
assert_eq!(response.model, "deepseek-ocr-maas");
assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1})));
}
}

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

@ -6,11 +6,13 @@ license.workspace = true
repository.workspace = true
[lib]
name = "litellm_python_bridge"
name = "_native"
crate-type = ["cdylib"]
[dependencies]
litellm-core.workspace = true
litellm-ai-gateway = { workspace = true, default-features = false }
pyo3 = { workspace = true, features = ["extension-module"] }
pyo3-async-runtimes.workspace = true
serde_json.workspace = true
tokio.workspace = true

6
litellm-rust/crates/python-bridge/build.rs Normal file
View file

@ -0,0 +1,6 @@
fn main() {
if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") {
println!("cargo:rustc-cdylib-link-arg=-undefined");
println!("cargo:rustc-cdylib-link-arg=dynamic_lookup");
}
}

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

@ -1,6 +1,6 @@
use std::time::Duration;
use litellm_ai_gateway::io::ocr::run_ocr;
use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest};
use litellm_core::error::CoreError;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
@ -9,6 +9,13 @@ use serde_json::{Map, Value};
mod gil;
type MarshaledOcrInputs = (
Value,
Option<Map<String, Value>>,
Map<String, Value>,
Option<Duration>,
);
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
let json = py.import("json")?;
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
@ -22,59 +29,93 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
Ok(json.call_method1("loads", (encoded,))?.unbind())
}
/// Map a core error to the closest Python exception. Caller-input problems
/// (auth, bad types, missing fields) -> `ValueError`; everything else
/// (network, upstream status, parse failures) -> `RuntimeError`.
fn core_error_to_pyerr(err: CoreError) -> PyErr {
match err {
CoreError::Auth(message) => PyValueError::new_err(message),
CoreError::InvalidType { .. } | CoreError::MissingField(_) => {
PyValueError::new_err(err.to_string())
}
CoreError::InvalidProvider(_)
| CoreError::InvalidRequest(_)
| CoreError::InvalidType { .. }
| CoreError::MissingField(_) => PyValueError::new_err(err.to_string()),
other => PyRuntimeError::new_err(other.to_string()),
}
}
/// Perform a Mistral OCR call end to end and return the response as a dict.
fn optional_object_to_map(
py: Python<'_>,
name: &'static str,
value: Option<Py<PyAny>>,
) -> PyResult<Map<String, Value>> {
match value {
Some(value) => match py_to_json(py, value.bind(py))? {
Value::Object(map) => Ok(map),
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
},
None => Ok(Map::new()),
}
}
fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
timeout_seconds.and_then(|secs| {
if secs.is_finite() && secs > 0.0 {
Some(Duration::from_secs_f64(secs))
} else {
None
}
})
}
fn marshal_inputs(
py: Python<'_>,
document: Py<PyAny>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledOcrInputs> {
let document = py_to_json(py, document.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
};
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let timeout = optional_timeout(timeout_seconds);
Ok((document, extra_headers, optional_params, timeout))
}
#[pyfunction]
#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))]
#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
fn ocr(
py: Python<'_>,
model: String,
document: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let document = py_to_json(py, document.bind(py))?;
let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string());
let (document, extra_headers, optional_params, timeout) = marshal_inputs(
py,
document,
extra_headers,
optional_params,
timeout_seconds,
)?;
let optional_params = match optional_params {
Some(params) => match py_to_json(py, params.bind(py))? {
Value::Object(map) => map,
_ => return Err(PyValueError::new_err("optional_params must be a dict")),
},
None => Map::new(),
};
let timeout = timeout_seconds.and_then(|secs| {
if secs.is_finite() && secs > 0.0 {
Some(Duration::from_secs_f64(secs))
} else {
None
}
});
// Release the GIL during the blocking HTTP call (counted for observability).
let result = gil::release_gil(py, || {
run_ocr(
&model,
pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest {
model: &model,
document,
api_key.as_deref(),
api_base.as_deref(),
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: &custom_llm_provider,
extra_headers,
optional_params,
timeout,
)
}))
});
match result {
@ -83,8 +124,47 @@ fn ocr(
}
}
/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe
/// how often the bridge has dropped the GIL for blocking work.
#[pyfunction]
#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
fn aocr(
py: Python<'_>,
model: String,
document: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string());
let (document, extra_headers, optional_params, timeout) = marshal_inputs(
py,
document,
extra_headers,
optional_params,
timeout_seconds,
)?;
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let value = run_ocr(OcrRequest {
model: &model,
document,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: &custom_llm_provider,
extra_headers,
optional_params,
timeout,
})
.await
.map_err(core_error_to_pyerr)?;
Python::with_gil(|py| json_to_py(py, value))
})
}
#[pyfunction]
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
let stats = PyDict::new(py);
@ -93,8 +173,9 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
}
#[pymodule]
fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> {
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(ocr, module)?)?;
module.add_function(wrap_pyfunction!(aocr, module)?)?;
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
Ok(())
}

33
litellm/caching/redis_cache.py
View file

@ -435,6 +435,7 @@ class RedisCache(BaseCache):
_redis_client = self.redis_client
start_time = time.time()
set_ttl = self.get_ttl(ttl=ttl)
key = self.check_and_fix_namespace(key=key)
try:
start_time = time.time()
result: int = _redis_client.incr(name=key, amount=value) # type: ignore
@ -498,6 +499,7 @@ class RedisCache(BaseCache):
)
return []
pattern = self.check_and_fix_namespace(key=pattern)
async for key in _redis_client.scan_iter(match=pattern + "*", count=count): # type: ignore
keys.append(key)
if len(keys) >= count:
@ -538,6 +540,11 @@ class RedisCache(BaseCache):
Register a Lua script with Redis asynchronously.
Works with both standalone Redis and Redis Cluster.
The returned callable namespaces every key it is invoked with, so Lua
scripts hit the same prefixed keys as get/set/increment. Without this,
scripts would operate on raw keys while the rest of the cache uses the
namespace, leaving rate-limit and lock keys outside the configured prefix.
Args:
script (str): The Lua script to register
@ -548,7 +555,15 @@ class RedisCache(BaseCache):
_redis_client = self.init_async_client()
# For standalone Redis
if hasattr(_redis_client, "register_script"):
return _redis_client.register_script(script) # type: ignore
registered_script = _redis_client.register_script(script) # type: ignore
async def namespaced_script(
keys: list[str], args: list[Any], client: Any = None
) -> Any:
keys = [self.check_and_fix_namespace(key=key) for key in keys]
return await registered_script(keys=keys, args=args, client=client)
return namespaced_script
# For Redis Cluster
elif hasattr(_redis_client, "script_load"):
# Load the script and get its SHA
@ -556,6 +571,7 @@ class RedisCache(BaseCache):
# Return a callable that uses evalsha
async def script_callable(keys: List[str], args: List[Any]) -> Any:
keys = [self.check_and_fix_namespace(key=key) for key in keys]
return _redis_client.evalsha(script_sha, len(keys), *keys, *args) # type: ignore
return script_callable
@ -1257,6 +1273,7 @@ class RedisCache(BaseCache):
async def delete_cache_keys(self, keys):
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete`
_redis_client: Any = self.init_async_client()
keys = [self.check_and_fix_namespace(key=key) for key in keys]
# keys is a list, unpack it so it gets passed as individual elements to delete
await _redis_client.delete(*keys)
@ -1322,10 +1339,12 @@ class RedisCache(BaseCache):
async def async_delete_cache(self, key: str):
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete`
_redis_client: Any = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
# keys is str
return await _redis_client.delete(key)
def delete_cache(self, key):
key = self.check_and_fix_namespace(key=key)
self.redis_client.delete(key)
async def _pipeline_increment_helper(
@ -1432,6 +1451,7 @@ class RedisCache(BaseCache):
try:
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl`
_redis_client: Any = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
ttl = await _redis_client.ttl(key)
if ttl <= -1: # -1 means the key does not exist, -2 key does not exist
return None
@ -1460,6 +1480,7 @@ class RedisCache(BaseCache):
int: The length of the list after the push operation
"""
_redis_client: Any = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
start_time = time.time()
try:
response = await _redis_client.rpush(key, *values)
@ -1499,7 +1520,8 @@ class RedisCache(BaseCache):
) -> List[int]:
"""Helper function for pipeline rpush operations"""
for rpush_op in rpush_list:
pipe.rpush(rpush_op["key"], *rpush_op["values"])
key = self.check_and_fix_namespace(key=rpush_op["key"])
pipe.rpush(key, *rpush_op["values"])
results = await pipe.execute()
# Preserve positional correspondence — raise on per-command errors
for r in results:
@ -1586,6 +1608,7 @@ class RedisCache(BaseCache):
**kwargs,
) -> Union[Any, List[Any]]:
_redis_client: Any = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
start_time = time.time()
print_verbose(f"LPOP from Redis list: key: {key}, count: {count}")
try:
@ -1658,17 +1681,19 @@ class RedisCache(BaseCache):
if major_version >= 7:
for lpop_op in lpop_list:
pipe.lpop(lpop_op["key"], lpop_op["count"])
key = self.check_and_fix_namespace(key=lpop_op["key"])
pipe.lpop(key, lpop_op["count"])
raw_results = await pipe.execute()
else:
# For Redis < 7, LPOP doesn't support count param.
# Issue `count` individual LPOP commands per key, all in one pipeline.
counts: List[int] = []
for lpop_op in lpop_list:
key = self.check_and_fix_namespace(key=lpop_op["key"])
count = lpop_op["count"] or 1
counts.append(count)
for _ in range(count):
pipe.lpop(lpop_op["key"])
pipe.lpop(key)
flat_results = await pipe.execute()
# Re-group the flat results back into per-key lists

47
litellm/integrations/opentelemetry.py
View file

@ -189,6 +189,37 @@ def _normalize_team_metadata_keys(value: Any) -> List[str]:
return [str(item).strip() for item in value if str(item).strip()]
_FREEZE_MAX_DEPTH = 16
HashableScope = Union[
str,
int,
float,
bool,
bytes,
None,
tuple["HashableScope", ...],
frozenset["HashableScope"],
]
def _freeze_for_dedupe(value: object, _depth: int = 0) -> HashableScope:
if _depth >= _FREEZE_MAX_DEPTH:
return repr(value)
if isinstance(value, (list, tuple)):
return tuple(_freeze_for_dedupe(item, _depth + 1) for item in value)
if isinstance(value, set):
return frozenset(_freeze_for_dedupe(item, _depth + 1) for item in value)
if isinstance(value, dict):
return frozenset(
(_freeze_for_dedupe(key, _depth + 1), _freeze_for_dedupe(item, _depth + 1))
for key, item in value.items()
)
if isinstance(value, (str, int, float, bytes)) or value is None:
return value
return repr(value)
@dataclass
class OpenTelemetryConfig:
exporter: Union[str, SpanExporter] = "console"
@ -1073,10 +1104,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
can be re-read with mutated entries between calls, so dedupe
must be at entry granularity. Scope: the entry's stable identity.
``scope`` parts can be any hashable identity. The marker is stored
in ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it
is request-local (kwargs is shared across the sync/async callbacks
and lifecycle hooks for one request).
``scope`` parts may include unhashable containers (list, dict, set);
they are normalized into a hashable shape via ``_freeze_for_dedupe``
before keying the marker dict. The marker is stored in
``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it is
request-local (kwargs is shared across the sync/async callbacks and
lifecycle hooks for one request).
"""
litellm_params = kwargs.get("litellm_params")
if not isinstance(litellm_params, dict):
@ -1098,7 +1131,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
spans_logged = {}
_otel_internal["spans_logged"] = spans_logged
dedupe_key = (self.__class__.__name__, id(self), *scope)
dedupe_key = (
self.__class__.__name__,
id(self),
*(_freeze_for_dedupe(part) for part in scope),
)
if spans_logged.get(dedupe_key) is True:
return False

19
litellm/llms/azure_ai/ocr/document_intelligence/transformation.py
View file

@ -11,7 +11,7 @@ The operation location must be polled until the analysis completes.
import asyncio
import re
import time
from typing import Any, Dict, Optional
from typing import Any, Dict
from urllib.parse import quote
import httpx
@ -35,6 +35,8 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.secret_managers.main import get_secret_str
AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"
class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
"""
@ -54,6 +56,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
def __init__(self) -> None:
super().__init__()
def get_api_key_env_var(self) -> str | None:
return AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR
def get_supported_ocr_params(self, model: str) -> list:
"""
Get supported OCR parameters for Azure Document Intelligence.
@ -144,9 +149,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
self,
headers: Dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_key: str | None = None,
api_base: str | None = None,
litellm_params: dict | None = None,
**kwargs,
) -> Dict:
"""
@ -156,7 +161,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
"""
# Get API key from environment if not provided
if api_key is None:
api_key = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_API_KEY")
api_key = get_secret_str(AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR)
if api_key is None:
raise ValueError(
@ -182,10 +187,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
def get_complete_url(
self,
api_base: Optional[str],
api_base: str | None,
model: str,
optional_params: dict,
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
**kwargs,
) -> str:
"""

19
litellm/llms/azure_ai/ocr/transformation.py
View file

@ -2,7 +2,7 @@
Azure AI OCR transformation implementation.
"""
from typing import Dict, Optional
from typing import Dict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.image_handling import (
@ -13,6 +13,8 @@ from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestDat
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
from litellm.secret_managers.main import get_secret_str
AZURE_AI_OCR_API_KEY_ENV_VAR = "AZURE_AI_API_KEY"
class AzureAIOCRConfig(MistralOCRConfig):
"""
@ -30,13 +32,16 @@ class AzureAIOCRConfig(MistralOCRConfig):
def __init__(self) -> None:
super().__init__()
def get_api_key_env_var(self) -> str | None:
return AZURE_AI_OCR_API_KEY_ENV_VAR
def validate_environment(
self,
headers: Dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_key: str | None = None,
api_base: str | None = None,
litellm_params: dict | None = None,
**kwargs,
) -> Dict:
"""
@ -46,7 +51,7 @@ class AzureAIOCRConfig(MistralOCRConfig):
"""
# Get API key from environment if not provided
if api_key is None:
api_key = get_secret_str("AZURE_AI_API_KEY")
api_key = get_secret_str(AZURE_AI_OCR_API_KEY_ENV_VAR)
if api_key is None:
raise ValueError(
@ -72,10 +77,10 @@ class AzureAIOCRConfig(MistralOCRConfig):
def get_complete_url(
self,
api_base: Optional[str],
api_base: str | None,
model: str,
optional_params: dict,
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
**kwargs,
) -> str:
"""

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

@ -2,7 +2,7 @@
Base OCR transformation configuration.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, List, Union
import httpx
from pydantic import PrivateAttr
@ -25,16 +25,16 @@ DocumentType = Dict[str, str]
class OCRPageDimensions(LiteLLMPydanticObjectBase):
"""Page dimensions from OCR response."""
dpi: Optional[int] = None
height: Optional[int] = None
width: Optional[int] = None
dpi: int | None = None
height: int | None = None
width: int | None = None
class OCRPageImage(LiteLLMPydanticObjectBase):
"""Image extracted from OCR page."""
image_base64: Optional[str] = None
bbox: Optional[Dict[str, Any]] = None
image_base64: str | None = None
bbox: Dict[str, Any] | None = None
model_config = {"extra": "allow"}
@ -44,8 +44,8 @@ class OCRPage(LiteLLMPydanticObjectBase):
index: int
markdown: str
images: Optional[List[OCRPageImage]] = None
dimensions: Optional[OCRPageDimensions] = None
images: List[OCRPageImage] | None = None
dimensions: OCRPageDimensions | None = None
model_config = {"extra": "allow"}
@ -53,9 +53,9 @@ class OCRPage(LiteLLMPydanticObjectBase):
class OCRUsageInfo(LiteLLMPydanticObjectBase):
"""Usage information from OCR response."""
pages_processed: Optional[int] = None
credits: Optional[float] = None
doc_size_bytes: Optional[int] = None
pages_processed: int | None = None
credits: float | None = None
doc_size_bytes: int | None = None
model_config = {"extra": "allow"}
@ -68,8 +68,8 @@ class OCRResponse(LiteLLMPydanticObjectBase):
pages: List[OCRPage]
model: str
document_annotation: Optional[Any] = None
usage_info: Optional[OCRUsageInfo] = None
document_annotation: Any | None = None
usage_info: OCRUsageInfo | None = None
object: str = "ocr"
model_config = {"extra": "allow"}
@ -81,8 +81,8 @@ class OCRResponse(LiteLLMPydanticObjectBase):
class OCRRequestData(LiteLLMPydanticObjectBase):
"""OCR request data structure."""
data: Optional[Union[Dict, bytes]] = None
files: Optional[Dict[str, Any]] = None
data: Union[Dict, bytes] | None = None
files: Dict[str, Any] | None = None
class BaseOCRConfig:
@ -101,6 +101,12 @@ class BaseOCRConfig:
"""
return []
def get_api_key_env_var(self) -> str | None:
"""
Return the provider-specific API key environment variable name, if any.
"""
return None
def map_ocr_params(
self,
non_default_params: dict,
@ -114,9 +120,9 @@ class BaseOCRConfig:
self,
headers: Dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_key: str | None = None,
api_base: str | None = None,
litellm_params: dict | None = None,
**kwargs,
) -> Dict:
"""
@ -127,10 +133,10 @@ class BaseOCRConfig:
def get_complete_url(
self,
api_base: Optional[str],
api_base: str | None,
model: str,
optional_params: dict,
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
**kwargs,
) -> str:
"""

19
litellm/llms/mistral/ocr/transformation.py
View file

@ -2,7 +2,7 @@
Mistral OCR transformation implementation.
"""
from typing import Any, Dict, Optional
from typing import Any, Dict
import httpx
@ -15,6 +15,8 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.secret_managers.main import get_secret_str
MISTRAL_OCR_API_KEY_ENV_VAR = "MISTRAL_API_KEY"
class MistralOCRConfig(BaseOCRConfig):
"""
@ -59,6 +61,9 @@ class MistralOCRConfig(BaseOCRConfig):
"id",
]
def get_api_key_env_var(self) -> str | None:
return MISTRAL_OCR_API_KEY_ENV_VAR
def map_ocr_params(
self,
non_default_params: dict,
@ -85,9 +90,9 @@ class MistralOCRConfig(BaseOCRConfig):
self,
headers: Dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_key: str | None = None,
api_base: str | None = None,
litellm_params: dict | None = None,
**kwargs,
) -> Dict:
"""
@ -95,7 +100,7 @@ class MistralOCRConfig(BaseOCRConfig):
"""
# Get API key from environment if not provided
if api_key is None:
api_key = get_secret_str("MISTRAL_API_KEY")
api_key = get_secret_str(MISTRAL_OCR_API_KEY_ENV_VAR)
if api_key is None:
raise ValueError(
@ -113,10 +118,10 @@ class MistralOCRConfig(BaseOCRConfig):
def get_complete_url(
self,
api_base: Optional[str],
api_base: str | None,
model: str,
optional_params: dict,
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
**kwargs,
) -> str:
"""

72
litellm/llms/vertex_ai/ocr/deepseek_transformation.py
View file

@ -3,7 +3,7 @@ Vertex AI DeepSeek OCR transformation implementation.
"""
import json
from typing import TYPE_CHECKING, Any, Dict, Optional
from typing import TYPE_CHECKING, Any, Dict
import httpx
@ -18,6 +18,8 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY"
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
else:
@ -28,21 +30,24 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
"""
Vertex AI DeepSeek OCR transformation configuration.
Vertex AI DeepSeek OCR uses the chat completion API format through the openapi endpoint.
This transformation converts OCR requests to chat completion format and vice versa.
This transformation converts standard LiteLLM OCR requests to the
Vertex AI DeepSeek OCR OpenAPI endpoint shape and normalizes the response.
"""
def __init__(self) -> None:
super().__init__()
self.vertex_base = VertexBase()
def get_api_key_env_var(self) -> str | None:
return VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR
def validate_environment(
self,
headers: Dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_key: str | None = None,
api_base: str | None = None,
litellm_params: dict | None = None,
**kwargs,
) -> Dict:
"""
@ -50,6 +55,13 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
Vertex AI uses Bearer token authentication with access token from credentials.
"""
if api_key is not None:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
**headers,
}
# Extract Vertex AI parameters using safe helpers from VertexBase
# Use safe_get_* methods that don't mutate litellm_params dict
litellm_params = litellm_params or {}
@ -77,18 +89,15 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
def get_complete_url(
self,
api_base: Optional[str],
api_base: str | None,
model: str,
optional_params: dict,
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
**kwargs,
) -> str:
"""
Get complete URL for Vertex AI DeepSeek OCR endpoint.
Vertex AI endpoint format:
https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions
Args:
api_base: Vertex AI API base URL (optional)
model: Model name (e.g., "deepseek-ai/deepseek-ocr-maas")
@ -123,8 +132,6 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
# Ensure no trailing slash
api_base = api_base.rstrip("/")
# Vertex AI DeepSeek OCR endpoint format
# Format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/endpoints/openapi/chat/completions
return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions"
def transform_ocr_request(
@ -136,9 +143,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
**kwargs,
) -> OCRRequestData:
"""
Transform OCR request to chat completion format for Vertex AI DeepSeek OCR.
Transform OCR request for Vertex AI DeepSeek OCR.
Converts OCR document format to chat completion messages format:
Converts OCR document format to the Vertex AI DeepSeek OCR payload:
- Input: {"type": "image_url", "image_url": "gs://..."}
- Output: {"model": "deepseek-ai/deepseek-ocr-maas", "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "gs://..."}]}]}
@ -150,7 +157,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
**kwargs: Additional arguments
Returns:
OCRRequestData with JSON data in chat completion format
OCRRequestData with JSON data for the DeepSeek OCR endpoint
"""
verbose_logger.debug(
"Vertex AI DeepSeek OCR transform_ocr_request (sync) called"
@ -173,7 +180,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'"
)
# Build chat completion message content
# Build DeepSeek OCR message content
content_item = {}
if image_url:
content_item = {"type": "image_url", "image_url": image_url}
@ -181,25 +188,21 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
# For document URLs, we use image_url type as well (Vertex AI supports both)
content_item = {"type": "image_url", "image_url": document_url}
# Build chat completion request
# Build DeepSeek OCR request
data = {
"model": "deepseek-ai/" + model,
"messages": [{"role": "user", "content": [content_item]}],
}
# Add optional parameters (stream, temperature, etc.)
# Filter out OCR-specific params that don't apply to chat completion
chat_completion_params = {}
deepseek_ocr_params = {}
for key, value in optional_params.items():
# Include common chat completion params
if key in ["stream", "temperature", "max_tokens", "top_p", "n", "stop"]:
chat_completion_params[key] = value
deepseek_ocr_params[key] = value
data.update(chat_completion_params)
data.update(deepseek_ocr_params)
verbose_logger.debug(
"Vertex AI DeepSeek OCR: Transformed request to chat completion format"
)
verbose_logger.debug("Vertex AI DeepSeek OCR: Transformed request")
return OCRRequestData(data=data, files=None)
@ -212,7 +215,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
**kwargs,
) -> OCRRequestData:
"""
Transform OCR request to chat completion format for Vertex AI DeepSeek OCR (async).
Transform OCR request for Vertex AI DeepSeek OCR (async).
Same as sync version - no async-specific logic needed.
@ -224,7 +227,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
**kwargs: Additional arguments
Returns:
OCRRequestData with JSON data in chat completion format
OCRRequestData with JSON data for the DeepSeek OCR endpoint
"""
return self.transform_ocr_request(
model=model,
@ -242,12 +245,11 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
**kwargs,
) -> OCRResponse:
"""
Transform chat completion response to OCR format.
Transform Vertex AI DeepSeek OCR response to OCR format.
Vertex AI DeepSeek OCR returns chat completion format:
Vertex AI DeepSeek OCR returns an OpenAPI response:
{
"id": "...",
"object": "chat.completion",
"choices": [{
"message": {
"role": "assistant",
@ -274,16 +276,16 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
try:
response_json = raw_response.json()
# Extract content from chat completion response
# Extract OCR content from provider response
choices = response_json.get("choices", [])
if not choices:
raise ValueError("No choices in chat completion response")
raise ValueError("No choices in DeepSeek OCR response")
message = choices[0].get("message", {})
content = message.get("content", "")
if not content:
raise ValueError("No content in chat completion response")
raise ValueError("No content in DeepSeek OCR response")
# Try to parse content as JSON (OCR result might be JSON string)
ocr_data = None
@ -376,7 +378,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig):
**kwargs,
) -> OCRResponse:
"""
Async transform chat completion response to OCR format.
Async transform Vertex AI DeepSeek OCR response to OCR format.
Same as sync version - no async-specific logic needed.

24
litellm/llms/vertex_ai/ocr/transformation.py
View file

@ -2,7 +2,7 @@
Vertex AI Mistral OCR transformation implementation.
"""
from typing import Dict, Optional
from typing import Dict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.image_handling import (
@ -14,6 +14,8 @@ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
VERTEX_AI_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY"
class VertexAIOCRConfig(MistralOCRConfig):
"""
@ -32,13 +34,16 @@ class VertexAIOCRConfig(MistralOCRConfig):
super().__init__()
self.vertex_base = VertexBase()
def get_api_key_env_var(self) -> str | None:
return VERTEX_AI_OCR_API_KEY_ENV_VAR
def validate_environment(
self,
headers: Dict,
model: str,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_key: str | None = None,
api_base: str | None = None,
litellm_params: dict | None = None,
**kwargs,
) -> Dict:
"""
@ -46,6 +51,13 @@ class VertexAIOCRConfig(MistralOCRConfig):
Vertex AI uses Bearer token authentication with access token from credentials.
"""
if api_key is not None:
return {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
**headers,
}
# Extract Vertex AI parameters using safe helpers from VertexBase
# Use safe_get_* methods that don't mutate litellm_params dict
litellm_params = litellm_params or {}
@ -73,10 +85,10 @@ class VertexAIOCRConfig(MistralOCRConfig):
def get_complete_url(
self,
api_base: Optional[str],
api_base: str | None,
model: str,
optional_params: dict,
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
**kwargs,
) -> str:
"""

806
litellm/ocr/main.py
View file

@ -4,13 +4,12 @@ Main OCR function for LiteLLM.
import asyncio
import base64
import contextvars
import mimetypes
import os
import re
from functools import partial
from dataclasses import dataclass
from io import IOBase
from typing import Any, Callable, Coroutine, Dict, Optional, Union, cast
from typing import Any, Callable, Coroutine, Union, cast
import httpx
@ -20,7 +19,13 @@ from litellm.constants import request_timeout
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.ocr.rust_bridge import RustOcr, load_rust_ocr, rust_ocr_enabled
from litellm.ocr.rust_bridge import (
RustAocr,
RustOcr,
load_rust_aocr,
load_rust_ocr,
rust_ocr_enabled,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -29,9 +34,40 @@ base_llm_http_handler = BaseLLMHTTPHandler()
#################################################
@dataclass
class _PreparedOCRRequest:
model: str
document: dict[str, Any]
api_key: str | None
api_base: str | None
custom_llm_provider: str
extra_headers: dict[str, object] | None
provider_config: BaseOCRConfig
optional_params: dict[str, object]
litellm_params: dict[str, object]
effective_timeout: Union[float, httpx.Timeout]
litellm_logging_obj: LiteLLMLoggingObj
@dataclass
class _PreparedRustOCRCall:
api_key: str | None
api_base: str | None
headers: dict[str, object]
optional_params: dict[str, object]
_RUST_OCR_PROVIDERS = {
"mistral",
"azure_ai",
"azure_ai/doc-intelligence",
"vertex_ai",
}
def _timeout_to_seconds(
timeout: Optional[Union[float, httpx.Timeout]],
) -> Optional[float]:
timeout: Union[float, httpx.Timeout] | None,
) -> float | None:
"""Convert the Python OCR timeout to a single seconds value for the Rust bridge.
The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate
@ -45,18 +81,206 @@ def _timeout_to_seconds(
return float(timeout)
def _prepare_ocr_request(
model: str,
document: dict[str, Any],
api_key: str | None,
api_base: str | None,
timeout: Union[float, httpx.Timeout] | None,
custom_llm_provider: str | None,
extra_headers: dict[str, Any] | None,
kwargs: dict[str, Any],
) -> _PreparedOCRRequest:
litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj"))
litellm_call_id = cast(str | None, kwargs.get("litellm_call_id", None))
if not isinstance(document, dict):
raise ValueError(
f"document must be a dict with 'type' and URL/file field, got {type(document)}"
)
doc_type = document.get("type")
if doc_type == "file":
document = convert_file_document_to_url_document(document)
doc_type = document.get("type")
if doc_type not in ["document_url", "image_url"]:
raise ValueError(
f"Invalid document type: {doc_type}. "
"Must be 'document_url', 'image_url', or 'file'"
)
(
model,
custom_llm_provider,
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
if dynamic_api_key:
api_key = dynamic_api_key
if dynamic_api_base:
api_base = dynamic_api_base
ocr_provider_config = ProviderConfigManager.get_provider_ocr_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if ocr_provider_config is None:
raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}")
verbose_logger.debug(f"OCR call - model: {model}, provider: {custom_llm_provider}")
litellm_params = GenericLiteLLMParams(**kwargs)
supported_params = ocr_provider_config.get_supported_ocr_params(model=model)
non_default_params = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
optional_params = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}")
effective_timeout = timeout or request_timeout
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
optional_params=optional_params,
litellm_params={
"litellm_call_id": litellm_call_id,
"api_base": api_base,
},
custom_llm_provider=custom_llm_provider,
)
return _PreparedOCRRequest(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=cast(dict[str, object] | None, extra_headers),
provider_config=ocr_provider_config,
optional_params=cast(dict[str, object], optional_params),
litellm_params=dict(litellm_params),
effective_timeout=effective_timeout,
litellm_logging_obj=litellm_logging_obj,
)
def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> dict[str, object]:
optional_params = dict(prepared_request.optional_params)
if prepared_request.custom_llm_provider == "vertex_ai":
vertex_project = (
prepared_request.litellm_params.get("vertex_project")
or prepared_request.litellm_params.get("vertex_ai_project")
or litellm.vertex_project
or resolve_secret("VERTEXAI_PROJECT")
)
vertex_location = (
prepared_request.litellm_params.get("vertex_location")
or prepared_request.litellm_params.get("vertex_ai_location")
or litellm.vertex_location
or resolve_secret("VERTEXAI_LOCATION")
or resolve_secret("VERTEX_LOCATION")
)
if vertex_project is not None:
optional_params["vertex_project"] = vertex_project
if vertex_location is not None:
optional_params["vertex_location"] = vertex_location
return optional_params
def _rust_bridge_api_base(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
) -> str | None:
if prepared_request.api_base is not None:
return prepared_request.api_base
if prepared_request.custom_llm_provider == "azure_ai/doc-intelligence":
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
if prepared_request.custom_llm_provider == "azure_ai":
if (
"doc-intelligence" in prepared_request.model
or "documentintelligence" in prepared_request.model
):
return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
return resolve_secret("AZURE_AI_API_BASE")
return None
def _prepare_rust_ocr_call(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> _PreparedRustOCRCall:
provider_config = prepared_request.provider_config
api_key_env_var = provider_config.get_api_key_env_var()
resolved_api_key = prepared_request.api_key or (
resolve_api_key(api_key_env_var) if api_key_env_var is not None else None
)
resolved_headers = provider_config.validate_environment(
headers=prepared_request.extra_headers or {},
model=prepared_request.model,
api_key=resolved_api_key,
api_base=prepared_request.api_base,
litellm_params=prepared_request.litellm_params,
)
resolved_complete_url = provider_config.get_complete_url(
api_base=prepared_request.api_base,
model=prepared_request.model,
optional_params=prepared_request.optional_params,
litellm_params=prepared_request.litellm_params,
)
rust_api_base = _rust_bridge_api_base(prepared_request, resolve_api_key)
rust_optional_params = _rust_bridge_optional_params(
prepared_request, resolve_api_key
)
prepared_request.litellm_logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": prepared_request.model,
"document": prepared_request.document,
**rust_optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
)
return _PreparedRustOCRCall(
api_key=resolved_api_key,
api_base=rust_api_base,
headers=cast(dict[str, object], resolved_headers),
optional_params=rust_optional_params,
)
def _run_rust_ocr(
rust_ocr: RustOcr,
logging_obj: LiteLLMLoggingObj,
provider_config: BaseOCRConfig,
resolve_api_key: Callable[[str], Optional[str]],
model: str,
document: dict[str, object],
api_key: Optional[str],
api_base: Optional[str],
optional_params: dict[str, object],
litellm_params: dict[str, object],
timeout_seconds: Optional[float],
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> OCRResponse:
"""Run the Mistral OCR call through the Rust bridge and wrap the result.
@ -66,41 +290,43 @@ def _run_rust_ocr(
headers) is mirrored into pre_call so logs match the wire. Dependencies are
injected so this stays unit-testable without patching module globals.
"""
resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY")
resolved_headers = provider_config.validate_environment(
headers={},
model=model,
api_key=resolved_api_key,
api_base=api_base,
litellm_params=litellm_params,
)
resolved_complete_url = provider_config.get_complete_url(
api_base=api_base,
model=model,
optional_params=optional_params,
litellm_params=litellm_params,
)
logging_obj.pre_call(
input="OCR document processing",
api_key=resolved_api_key,
additional_args={
"complete_input_dict": {
"model": model,
"document": document,
**optional_params,
},
"api_base": resolved_complete_url,
"headers": resolved_headers,
},
prepared = _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
return OCRResponse.model_validate(
rust_ocr(
model=model,
document=document,
api_key=resolved_api_key,
api_base=api_base,
optional_params=optional_params,
timeout_seconds=timeout_seconds,
model=prepared_request.model,
document=cast(dict[str, object], prepared_request.document),
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
)
)
async def _run_rust_aocr(
rust_aocr: RustAocr,
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
) -> OCRResponse:
prepared = _prepare_rust_ocr_call(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
return OCRResponse.model_validate(
await rust_aocr(
model=prepared_request.model,
document=cast(dict[str, object], prepared_request.document),
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout),
)
)
@ -108,12 +334,12 @@ def _run_rust_ocr(
@client
async def aocr(
model: str,
document: Dict[str, Any],
api_key: Optional[str] = None,
api_base: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: Optional[Dict[str, Any]] = None,
document: dict[str, Any],
api_key: str | None = None,
api_base: str | None = None,
timeout: Union[float, httpx.Timeout] | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, Any] | None = None,
**kwargs,
) -> OCRResponse:
"""
@ -174,19 +400,18 @@ async def aocr(
)
```
"""
local_vars = locals()
completion_kwargs: dict[str, object] = {
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
try:
loop = asyncio.get_event_loop()
kwargs["aocr"] = True
# Get custom llm provider
if custom_llm_provider is None:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=model, api_base=api_base
)
func = partial(
ocr,
prepared = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
@ -194,17 +419,47 @@ async def aocr(
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
**kwargs,
kwargs=kwargs,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update(
{"model": model, "custom_llm_provider": custom_llm_provider}
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if _rust_ocr_supported(prepared) and rust_ocr_enabled():
rust_aocr = load_rust_aocr()
if rust_aocr is None:
verbose_logger.debug(
"Async Rust OCR bridge unavailable; falling back to Python path"
)
else:
from litellm.secret_managers.main import get_secret_str
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
response = await _run_rust_aocr(
rust_aocr=rust_aocr,
prepared_request=prepared,
resolve_api_key=get_secret_str,
)
return response
response = base_llm_http_handler.ocr(
model=prepared.model,
document=prepared.document,
optional_params=prepared.optional_params,
timeout=prepared.effective_timeout,
logging_obj=prepared.litellm_logging_obj,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared.custom_llm_provider,
aocr=True,
headers=prepared.extra_headers,
provider_config=prepared.provider_config,
litellm_params=prepared.litellm_params,
)
if asyncio.iscoroutine(response):
response = await response
if response is None:
raise ValueError(
@ -217,20 +472,145 @@ async def aocr(
model=model,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)
#################################################
# Public utilities — used by the SDK and the proxy
#################################################
_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP = {
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".tiff": "image/tiff",
".tif": "image/tiff",
".bmp": "image/bmp",
}
def get_mime_type(file_path: str) -> str:
"""
Determine MIME type from file path extension.
Falls back to mimetypes.guess_type, then to 'application/octet-stream'.
"""
ext = os.path.splitext(file_path)[1].lower()
mime = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]:
"""
Convert a file-type document dict to a document_url-type document dict
with an inline base64 data URI.
Accepts document dicts like:
{"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path
{"type": "file", "file": <binary file-like object>} # file-like object (BinaryIO)
{"type": "file", "file": b"raw bytes"} # raw bytes
Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or
``open(path, "rb")`` instead. See the str check below for the rationale.
Returns:
{"type": "document_url", "document_url": "data:<mime>;base64,<data>"}
or {"type": "image_url", "image_url": "data:<mime>;base64,<data>"}
"""
file_input = document.get("file")
if file_input is None:
raise ValueError(
"document with type='file' must include a 'file' field containing "
"a pathlib.Path, file-like object, or bytes"
)
file_bytes: bytes
mime_type: str = "application/octet-stream"
file_name: str | None = None
if isinstance(file_input, str):
# Bare strings are rejected here. The OCR ``document`` accepts a
# ``{"type": "file", "file": <value>}`` shape, and when this helper
# runs in a proxy request handler ``<value>`` is attacker-controlled.
# Opening it as a path is an arbitrary local file read on the proxy
# host, which is then base64-encoded and forwarded to the OCR
# provider — an exfiltration primitive.
raise ValueError(
"OCR file input does not accept bare str values. Pass bytes, "
"a pathlib.Path, or a file-like object. To OCR a local file "
"from a path, call open(path, 'rb') yourself."
)
if isinstance(file_input, os.PathLike):
# os.PathLike (pathlib.Path and custom __fspath__ classes) is a
# Python-level type that HTTP form values can't fabricate.
file_path = str(file_input)
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type = get_mime_type(file_path)
file_name = os.path.basename(file_path)
with open(file_path, "rb") as f:
file_bytes = f.read()
elif isinstance(file_input, bytes):
file_bytes = file_input
elif isinstance(file_input, IOBase) or hasattr(file_input, "read"):
if hasattr(file_input, "name"):
file_name = getattr(file_input, "name", None)
if file_name:
mime_type = get_mime_type(file_name)
file_bytes = file_input.read()
if isinstance(file_bytes, str):
file_bytes = file_bytes.encode("utf-8")
else:
raise ValueError(
f"Unsupported file input type: {type(file_input)}. "
"Expected pathlib.Path, bytes, or a file-like object."
)
if not file_bytes:
raise ValueError("File is empty or could not be read")
if "mime_type" in document:
mime_type = document["mime_type"]
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data = base64.b64encode(file_bytes).decode("utf-8")
data_uri = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
f"OCR file input: Converted file to image_url data URI "
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
)
return {"type": "image_url", "image_url": data_uri}
verbose_logger.debug(
f"OCR file input: Converted file to document_url data URI "
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
)
return {"type": "document_url", "document_url": data_uri}
@client
def ocr(
model: str,
document: Dict[str, Any],
api_key: Optional[str] = None,
api_base: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
custom_llm_provider: Optional[str] = None,
extra_headers: Optional[Dict[str, Any]] = None,
document: dict[str, Any],
api_key: str | None = None,
api_base: str | None = None,
timeout: Union[float, httpx.Timeout] | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, Any] | None = None,
**kwargs,
) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]:
"""
@ -295,96 +675,37 @@ def ocr(
print(f"Page {page.index}: {page.markdown}")
```
"""
local_vars = locals()
completion_kwargs: dict[str, object] = {
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"kwargs": kwargs,
}
try:
litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj"))
litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None)
_is_async = kwargs.pop("aocr", False) is True
# Validate document parameter format
if not isinstance(document, dict):
raise ValueError(
f"document must be a dict with 'type' and URL/file field, got {type(document)}"
)
doc_type = document.get("type")
# Handle file type: convert to document_url/image_url with base64 data URI
if doc_type == "file":
document = convert_file_document_to_url_document(document)
doc_type = document.get("type")
if doc_type not in ["document_url", "image_url"]:
raise ValueError(
f"Invalid document type: {doc_type}. "
"Must be 'document_url', 'image_url', or 'file'"
)
(
model,
custom_llm_provider,
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
completion_kwargs["aocr"] = _is_async
prepared = _prepare_ocr_request(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
document=document,
api_key=api_key,
)
# Update with dynamic values if available
if dynamic_api_key:
api_key = dynamic_api_key
if dynamic_api_base:
api_base = dynamic_api_base
ocr_provider_config: Optional[BaseOCRConfig] = (
ProviderConfigManager.get_provider_ocr_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
)
if ocr_provider_config is None:
raise ValueError(
f"OCR is not supported for provider: {custom_llm_provider}"
)
verbose_logger.debug(
f"OCR call - model: {model}, provider: {custom_llm_provider}"
)
litellm_params = GenericLiteLLMParams(**kwargs)
supported_params = ocr_provider_config.get_supported_ocr_params(model=model)
non_default_params = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
optional_params = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}")
effective_timeout = timeout or request_timeout
litellm_logging_obj.update_from_kwargs(
api_base=api_base,
kwargs=kwargs,
model=model,
optional_params=optional_params,
litellm_params={
"litellm_call_id": litellm_call_id,
"api_base": api_base,
},
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout=timeout,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update(
{"model": model, "custom_llm_provider": custom_llm_provider}
)
# Optional Rust path: hand the whole Mistral OCR call to the Rust bridge.
if custom_llm_provider == "mistral" and rust_ocr_enabled():
# Optional Rust path: hand supported OCR calls to the Rust bridge.
if _rust_ocr_supported(prepared) and rust_ocr_enabled():
rust_ocr = load_rust_ocr()
if rust_ocr is None:
verbose_logger.debug(
@ -395,31 +716,23 @@ def ocr(
return _run_rust_ocr(
rust_ocr=rust_ocr,
logging_obj=litellm_logging_obj,
provider_config=ocr_provider_config,
prepared_request=prepared,
resolve_api_key=get_secret_str,
model=model,
document=document,
api_key=api_key,
api_base=api_base,
optional_params=optional_params,
litellm_params=dict(litellm_params),
timeout_seconds=_timeout_to_seconds(effective_timeout),
)
response = base_llm_http_handler.ocr(
model=model,
document=document,
optional_params=optional_params,
timeout=effective_timeout,
logging_obj=litellm_logging_obj,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
model=prepared.model,
document=prepared.document,
optional_params=prepared.optional_params,
timeout=prepared.effective_timeout,
logging_obj=prepared.litellm_logging_obj,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared.custom_llm_provider,
aocr=_is_async,
headers=extra_headers,
provider_config=ocr_provider_config,
litellm_params=dict(litellm_params),
headers=prepared.extra_headers,
provider_config=prepared.provider_config,
litellm_params=prepared.litellm_params,
)
return response
@ -428,131 +741,6 @@ def ocr(
model=model,
custom_llm_provider=custom_llm_provider,
original_exception=e,
completion_kwargs=local_vars,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)
#################################################
# Public utilities — used by the SDK and the proxy
#################################################
_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP = {
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".tiff": "image/tiff",
".tif": "image/tiff",
".bmp": "image/bmp",
}
def get_mime_type(file_path: str) -> str:
"""
Determine MIME type from file path extension.
Falls back to mimetypes.guess_type, then to 'application/octet-stream'.
"""
ext = os.path.splitext(file_path)[1].lower()
mime = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str, str]:
"""
Convert a file-type document dict to a document_url-type document dict
with an inline base64 data URI.
Accepts document dicts like:
{"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path
{"type": "file", "file": <binary file-like object>} # file-like object (BinaryIO)
{"type": "file", "file": b"raw bytes"} # raw bytes
Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or
``open(path, "rb")`` instead. See the str check below for the rationale.
Returns:
{"type": "document_url", "document_url": "data:<mime>;base64,<data>"}
or {"type": "image_url", "image_url": "data:<mime>;base64,<data>"}
"""
file_input = document.get("file")
if file_input is None:
raise ValueError(
"document with type='file' must include a 'file' field containing "
"a pathlib.Path, file-like object, or bytes"
)
file_bytes: bytes
mime_type: str = "application/octet-stream"
file_name: Optional[str] = None
if isinstance(file_input, str):
# Bare strings are rejected here. The OCR ``document`` accepts a
# ``{"type": "file", "file": <value>}`` shape, and when this helper
# runs in a proxy request handler ``<value>`` is attacker-controlled.
# Opening it as a path is an arbitrary local file read on the proxy
# host, which is then base64-encoded and forwarded to the OCR
# provider — an exfiltration primitive.
raise ValueError(
"OCR file input does not accept bare str values. Pass bytes, "
"a pathlib.Path, or a file-like object. To OCR a local file "
"from a path, call open(path, 'rb') yourself."
)
if isinstance(file_input, os.PathLike):
# os.PathLike (pathlib.Path and custom __fspath__ classes) is a
# Python-level type that HTTP form values can't fabricate.
file_path = str(file_input)
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type = get_mime_type(file_path)
file_name = os.path.basename(file_path)
with open(file_path, "rb") as f:
file_bytes = f.read()
elif isinstance(file_input, bytes):
file_bytes = file_input
elif isinstance(file_input, IOBase) or hasattr(file_input, "read"):
if hasattr(file_input, "name"):
file_name = getattr(file_input, "name", None)
if file_name:
mime_type = get_mime_type(file_name)
file_bytes = file_input.read()
if isinstance(file_bytes, str):
file_bytes = file_bytes.encode("utf-8")
else:
raise ValueError(
f"Unsupported file input type: {type(file_input)}. "
"Expected pathlib.Path, bytes, or a file-like object."
)
if not file_bytes:
raise ValueError("File is empty or could not be read")
if "mime_type" in document:
mime_type = document["mime_type"]
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data = base64.b64encode(file_bytes).decode("utf-8")
data_uri = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
f"OCR file input: Converted file to image_url data URI "
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
)
return {"type": "image_url", "image_url": data_uri}
else:
verbose_logger.debug(
f"OCR file input: Converted file to document_url data URI "
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
)
return {"type": "document_url", "document_url": data_uri}

82
litellm/ocr/rust_bridge.py
View file

@ -2,7 +2,7 @@
Optional Rust-backed OCR path.
Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint
then routes supported Mistral calls through the compiled ``litellm_python_bridge``
then routes supported Mistral calls through the compiled ``litellm.rust_bridge._native``
extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust.
No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py``
@ -11,11 +11,12 @@ can import it statically without forming an import cycle.
from __future__ import annotations
from typing import Final, Protocol, cast
import os
from typing import Awaitable, Final, Protocol, cast
class RustOcr(Protocol):
"""Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint."""
"""Signature of the compiled Rust OCR entrypoint."""
def __call__(
self,
@ -23,9 +24,29 @@ class RustOcr(Protocol):
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]: ...
) -> dict[str, object]:
raise NotImplementedError
class RustAocr(Protocol):
"""Signature of the compiled ``litellm_python_bridge.aocr`` entrypoint."""
def __call__(
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
class _Unset:
@ -34,23 +55,39 @@ class _Unset:
_UNSET: Final[_Unset] = _Unset()
_rust_ocr_enabled = False
def _env_enables_rust_ocr() -> bool:
return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in {
"1",
"true",
"yes",
"on",
}
_rust_ocr_enabled = _env_enables_rust_ocr()
_rust_ocr_impl: RustOcr | None = None
_rust_aocr_impl: RustAocr | None = None
def use_litellm_rust(
enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET
enabled: bool = True,
*,
ocr: RustOcr | None | _Unset = _UNSET,
aocr: RustAocr | None | _Unset = _UNSET,
) -> None:
"""Route supported OCR calls through the Rust ``litellm_python_bridge`` extension.
"""Route supported OCR calls through the packaged Rust extension.
``ocr`` injects the bridge callable; when omitted the compiled extension is
loaded on demand and any previously injected bridge is preserved. Pass
``ocr=None`` explicitly to clear a prior injection.
``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled
extension is loaded on demand and any previously injected bridge is
preserved. Pass ``None`` explicitly to clear a prior injection.
"""
global _rust_ocr_enabled, _rust_ocr_impl
global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl
_rust_ocr_enabled = enabled
if not isinstance(ocr, _Unset):
_rust_ocr_impl = ocr
if not isinstance(aocr, _Unset):
_rust_aocr_impl = aocr
def rust_ocr_enabled() -> bool:
@ -62,13 +99,26 @@ def load_rust_ocr() -> RustOcr | None:
"""Return the Rust OCR callable, or ``None`` when no bridge is available.
Prefers an injected implementation, otherwise loads the compiled
``litellm_python_bridge`` extension; a missing extension yields ``None`` so
``litellm.rust_bridge._native`` extension; a missing extension yields ``None`` so
the caller can fall back to the Python path instead of hard-failing.
"""
if _rust_ocr_impl is not None:
return _rust_ocr_impl
try:
import litellm_python_bridge
except ImportError:
from litellm.rust_bridge import get_native_bridge
native_bridge = get_native_bridge()
if native_bridge is None:
return None
return cast(RustOcr, litellm_python_bridge.ocr)
return cast(RustOcr, native_bridge.ocr)
def load_rust_aocr() -> RustAocr | None:
"""Return the async Rust OCR callable, or ``None`` when unavailable."""
if _rust_aocr_impl is not None:
return _rust_aocr_impl
from litellm.rust_bridge import get_native_bridge
native_bridge = get_native_bridge()
if native_bridge is None:
return None
return cast(RustAocr, getattr(native_bridge, "aocr", None))

8
litellm/rust_bridge/__init__.py Normal file
View file

@ -0,0 +1,8 @@
"""LiteLLM Rust bridge package."""
from litellm.rust_bridge.loader import (
get_native_bridge,
native_bridge_available,
)
__all__ = ["get_native_bridge", "native_bridge_available"]

28
litellm/rust_bridge/loader.py Normal file
View file

@ -0,0 +1,28 @@
"""Loader for the packaged LiteLLM Rust extension."""
from __future__ import annotations
from types import ModuleType
_BRIDGE_SENTINEL = object()
_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL
def get_native_bridge() -> ModuleType | None:
"""Return the packaged Rust extension, or ``None`` when unavailable."""
global _cached_bridge
if _cached_bridge is not _BRIDGE_SENTINEL:
return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None
try:
from litellm.rust_bridge import _native
except ImportError:
_cached_bridge = None
return None
_cached_bridge = _native
return _native
def native_bridge_available() -> bool:
"""Whether the packaged Rust extension is importable."""
return get_native_bridge() is not None

192
model_prices_and_context_window.json
View file

@ -16153,6 +16153,47 @@
"supports_service_tier": true,
"supports_image_size": false
},
"gemini-3-pro-image": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_batches": 6e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query",
"supports_service_tier": true
},
"gemini-3-pro-image-preview": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@ -16194,6 +16235,44 @@
"web_search_billing_unit": "per_query",
"supports_service_tier": true
},
"gemini-3.1-flash-image": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query"
},
"gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
@ -17765,6 +17844,49 @@
"supports_service_tier": true,
"supports_image_size": false
},
"gemini/gemini-3-pro-image": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "gemini",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"rpm": 1000,
"tpm": 4000000,
"output_cost_per_token_batches": 6e-06,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query",
"supports_service_tier": true
},
"gemini/gemini-3-pro-image-preview": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@ -17808,6 +17930,48 @@
"web_search_billing_unit": "per_query",
"supports_service_tier": true
},
"gemini/gemini-3.1-flash-image": {
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_batches": 1.25e-07,
"litellm_provider": "gemini",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.045,
"output_cost_per_image_token": 6e-05,
"output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 1.5e-06,
"output_cost_per_token_batches": 7.5e-07,
"rpm": 1000,
"tpm": 4000000,
"source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
"search_context_size_high": 0.014
},
"web_search_billing_unit": "per_query"
},
"gemini/gemini-3.1-flash-image-preview": {
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_batches": 1.25e-07,
@ -35573,6 +35737,21 @@
"tpm": 8000000,
"supports_image_size": false
},
"vertex_ai/gemini-3-pro-image": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
"input_cost_per_token_batches": 1e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_batches": 6e-06,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3-pro-image-preview": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@ -35588,6 +35767,19 @@
"output_cost_per_token_batches": 6e-06,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3.1-flash-image": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,

33
pyproject.toml
View file

@ -234,8 +234,24 @@ healthcheck = [
]
[build-system]
requires = ["uv_build==0.11.8"]
build-backend = "uv_build"
requires = ["maturin==1.9.4"]
build-backend = "maturin"
[tool.maturin]
manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml"
module-name = "litellm.rust_bridge._native"
python-source = "."
bindings = "pyo3"
exclude = [
"litellm/proxy/enterprise",
"litellm/proxy/enterprise/**",
"**/__pycache__",
"**/__pycache__/**",
"**/.pytest_cache",
"**/.pytest_cache/**",
"**/.ruff_cache",
"**/.ruff_cache/**",
]
[tool.uv]
constraint-dependencies = [
@ -253,18 +269,6 @@ litellm-enterprise = { workspace = true }
[tool.uv.workspace]
members = ["enterprise", "litellm-proxy-extras"]
[tool.uv.build-backend]
module-root = ""
source-exclude = [
"litellm/proxy/enterprise",
"**/__pycache__",
"**/__pycache__/**",
"**/.pytest_cache",
"**/.pytest_cache/**",
"**/.ruff_cache",
"**/.ruff_cache/**",
]
[tool.isort]
profile = "black"
@ -328,4 +332,3 @@ pytest_add_cli_args = [
[tool.coverage.run]
source = ["litellm"]
relative_files = true

16
ruff-strict-budget.json
View file

@ -1,7 +1,7 @@
{
"ANN001": {
"baseline": 2865,
"slack": 50
"slack": 287
},
"ANN002": {
"baseline": 64,
@ -9,19 +9,19 @@
},
"ANN003": {
"baseline": 759,
"slack": 30
"slack": 76
},
"ANN201": {
"baseline": 1944,
"slack": 50
"slack": 194
},
"ANN202": {
"baseline": 858,
"slack": 30
"slack": 86
},
"ANN204": {
"baseline": 658,
"slack": 20
"slack": 66
},
"ANN205": {
"baseline": 117,
@ -33,7 +33,7 @@
},
"ANN401": {
"baseline": 1886,
"slack": 50
"slack": 189
},
"ASYNC230": {
"baseline": 11,
@ -231,10 +231,6 @@
"baseline": 6,
"slack": 3
},
"PLR0913": {
"baseline": 1813,
"slack": 50
},
"PLR1704": {
"baseline": 3,
"slack": 3

2
ruff-strict.toml
View file

@ -2,7 +2,7 @@ extend = "ruff.toml"
[lint]
preview = true
select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR0913", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"]
select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"]
extend-select = []
[lint.mccabe]

97
scripts/budget_ratchet_check.py
View file

@ -1,16 +1,19 @@
#!/usr/bin/env python3
"""Non-gating ratchet guard: budget ceilings may only fall, never rise.
"""Non-gating ratchet guard: budget baselines and ceilings may only fall, never rise.
Every `*-budget.json` file (ruff-strict, type-discipline, basedpyright-code) is a
one-way ratchet: each rule's ceiling is `baseline + slack`, and the whole point is
to drive that number DOWN over time. This check compares every budget file against
its own content at the merge-base with the target branch and fails (exits 1, red) if:
one-way ratchet: each rule's ceiling is `baseline + slack`, and both the recorded
`baseline` (the live violation count) and that ceiling are meant to be driven DOWN
over time. This check compares every budget file against its own content at the
merge-base with the target branch and fails (exits 1, red) if:
* a rule's ceiling went up,
* a rule's ceiling (`baseline + slack`) went up,
* a rule's `baseline` went up, even if `slack` was lowered to keep the ceiling
flat (a higher baseline bakes in more accepted debt and must be acknowledged),
* a rule was dropped from a budget (its ceiling effectively became infinite), or
* an entire budget file was deleted.
New rules and lowered/equal ceilings are fine.
New rules and lowered/equal baselines and ceilings are fine.
This is deliberately NOT a gating check. It should turn the run red so that a
loosening is impossible to miss in review, but it must stay OUT of the
@ -66,7 +69,12 @@ def _load_head(rel: str) -> dict | None:
def _ref_is_commit(ref: str) -> bool:
return _run(["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"]).returncode == 0
return (
_run(
["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"]
).returncode
== 0
)
def _load_base(rel: str, ref: str) -> dict | None:
@ -81,13 +89,55 @@ def _load_base(rel: str, ref: str) -> dict | None:
return json.loads(proc.stdout)
def _baselines(budget: dict) -> dict[str, int]:
"""Map each rule to its recorded baseline; skip malformed specs."""
return {
rule: int(spec.get("baseline", 0))
for rule, spec in budget.items()
if isinstance(spec, dict)
}
def _caps(budget: dict) -> dict[str, int]:
"""Map each rule to its ceiling (baseline + slack); skip malformed specs."""
caps: dict[str, int] = {}
for rule, spec in budget.items():
if isinstance(spec, dict):
caps[rule] = int(spec.get("baseline", 0)) + int(spec.get("slack", 0))
return caps
return {
rule: int(spec.get("baseline", 0)) + int(spec.get("slack", 0))
for rule, spec in budget.items()
if isinstance(spec, dict)
}
def _regression_detail(
rule: str,
base_caps: dict[str, int],
head_caps: dict[str, int],
base_baselines: dict[str, int],
head_baselines: dict[str, int],
) -> str | None:
"""Why `rule` regressed vs base, or None when it held flat or fell.
A dropped rule is terminal; otherwise a raised ceiling and a raised baseline are
independent loosenings (the latter catches a baseline bump masked by a slack cut),
so both reasons are reported when both apply.
"""
base_cap = base_caps[rule]
if rule not in head_caps:
return f"rule dropped (ceiling {base_cap} -> removed)"
reasons = tuple(
message
for raised, message in (
(
head_caps[rule] > base_cap,
f"ceiling raised {base_cap} -> {head_caps[rule]}",
),
(
head_baselines[rule] > base_baselines[rule],
f"baseline raised {base_baselines[rule]} -> {head_baselines[rule]}",
),
)
if raised
)
return "; ".join(reasons) or None
def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]:
@ -96,18 +146,17 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr
if head is None:
return [Regression(rel, "*", "budget file was deleted (every ceiling removed)")]
base_caps = _caps(base)
head_caps = _caps(head)
base_caps, head_caps = _caps(base), _caps(head)
base_baselines, head_baselines = _baselines(base), _baselines(head)
return [
Regression(
rel,
rule,
f"rule dropped (ceiling {base_cap} -> removed)"
if rule not in head_caps
else f"ceiling raised {base_cap} -> {head_caps[rule]}",
Regression(rel, rule, detail)
for rule in sorted(base_caps)
if (
detail := _regression_detail(
rule, base_caps, head_caps, base_baselines, head_baselines
)
)
for rule, base_cap in sorted(base_caps.items())
if rule not in head_caps or head_caps[rule] > base_cap
is not None
]
@ -141,7 +190,9 @@ def main() -> int:
regressions.extend(regressions_for(rel, base, head))
if regressions:
print(f"FAIL: budget ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):")
print(
f"FAIL: budget baseline(s)/ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):"
)
for reg in regressions:
print(f" {reg.budget} {reg.rule}: {reg.detail}")
print(

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -51,6 +51,7 @@ IGNORE_FUNCTIONS = [
"_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard.
"resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input).
"sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth.
"_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap.
]

7
tests/documentation_tests/test_env_keys.py
View file

@ -24,6 +24,12 @@ EXCLUDED_GUARD_ONLY_VARS = {
"MAVVRIK_FOCUS_FREQUENCY",
}
# Temporary/internal rollout flags are intentionally not added to the public
# environment settings docs until the feature is ready for broad use.
EXCLUDED_ROLLOUT_FLAGS = {
"LITELLM_USE_RUST_OCR",
}
EXCLUDED_TERMINAL_VARS = {
"TERM",
"TERM_PROGRAM",
@ -71,6 +77,7 @@ for root, dirs, files in os.walk(repo_base):
for match in getenv_matches
if match not in EXCLUDED_TERMINAL_VARS
and match not in EXCLUDED_GUARD_ONLY_VARS
and match not in EXCLUDED_ROLLOUT_FLAGS
) # Extract only the key part, excluding terminal vars
# Find all keys using litellm.get_secret()

30
tests/e2e/gateway/litellm-config.yml
View file

@ -140,6 +140,35 @@ model_list:
model_info:
mode: realtime
- model_name: rust-ocr-mistral
litellm_params:
model: mistral/mistral-ocr-latest
api_key: os.environ/MISTRAL_API_KEY
- model_name: rust-ocr-azure-ai
litellm_params:
model: azure_ai/mistral-document-ai-2505
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
- model_name: rust-ocr-azure-document-intelligence
litellm_params:
model: azure_ai/doc-intelligence/prebuilt-layout
api_base: os.environ/AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT
api_key: os.environ/AZURE_DOCUMENT_INTELLIGENCE_API_KEY
- model_name: rust-ocr-vertex-mistral
litellm_params:
model: vertex_ai/mistral-ocr-2505
vertex_project: os.environ/VERTEXAI_PROJECT
vertex_location: us-central1
- model_name: rust-ocr-vertex-deepseek
litellm_params:
model: vertex_ai/deepseek-ocr-maas
vertex_project: os.environ/VERTEXAI_PROJECT
vertex_location: us-central1
mcp_servers:
deepwiki_mcp:
@ -167,4 +196,3 @@ guardrails:
US_SSN: BLOCK
PHONE_NUMBER: BLOCK

148
tests/e2e/gateway/test_ocr_rust_e2e.py Normal file
View file

@ -0,0 +1,148 @@
"""
Gateway E2E smoke for Rust-backed OCR.
Start the proxy with:
LITELLM_USE_RUST_OCR=1 litellm --config tests/e2e/gateway/litellm-config.yml --port 4000
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import httpx
import pytest
import yaml
TEST_PDF_URL = (
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
"/tests/llm_translation/fixtures/dummy.pdf"
)
TEST_IMAGE_URL = (
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
"/tests/image_gen_tests/test_image.png"
)
RUST_OCR_GATEWAY_CASES = [
pytest.param(
"rust-ocr-mistral",
{"type": "document_url", "document_url": TEST_PDF_URL},
id="mistral",
),
pytest.param(
"rust-ocr-azure-ai",
{"type": "document_url", "document_url": TEST_PDF_URL},
id="azure_ai",
),
pytest.param(
"rust-ocr-azure-document-intelligence",
{"type": "document_url", "document_url": TEST_PDF_URL},
id="azure_document_intelligence",
),
pytest.param(
"rust-ocr-vertex-mistral",
{"type": "document_url", "document_url": TEST_PDF_URL},
id="vertex_mistral",
),
pytest.param(
"rust-ocr-vertex-deepseek",
{
"type": "image_url",
"image_url": os.getenv("RUST_OCR_IMAGE_URL", TEST_IMAGE_URL),
},
id="vertex_deepseek",
),
]
CONFIG_PATH = Path(__file__).with_name("litellm-config.yml")
@dataclass(frozen=True)
class OcrGateway:
base_url: str
master_key: str
def model_names(self) -> set[str]:
with httpx.Client(
timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120"))
) as client:
response = client.get(
f"{self.base_url.rstrip('/')}/model/info",
headers={"Authorization": f"Bearer {self.master_key}"},
)
assert response.status_code == 200, response.text
return {
model["model_name"]
for model in response.json().get("data", [])
if "model_name" in model
}
def ocr(self, model: str, document: dict[str, str]) -> httpx.Response:
with httpx.Client(
timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120"))
) as client:
return client.post(
f"{self.base_url.rstrip('/')}/v1/ocr",
headers={"Authorization": f"Bearer {self.master_key}"},
json={"model": model, "document": document},
)
@dataclass(frozen=True)
class OcrResources:
gateway: OcrGateway
@pytest.fixture
def resources() -> OcrResources:
proxy_url = os.getenv("LITELLM_PROXY_URL")
if not proxy_url:
pytest.skip(
"Start a Rust OCR proxy and set LITELLM_PROXY_URL, e.g. http://localhost:4000"
)
return OcrResources(
gateway=OcrGateway(
base_url=proxy_url,
master_key=os.getenv("LITELLM_MASTER_KEY", "sk-1234"),
)
)
def _assert_ocr_response_shape(response_json: dict[str, Any]) -> None:
assert response_json["object"] == "ocr"
assert response_json["model"]
assert isinstance(response_json["pages"], list)
assert len(response_json["pages"]) > 0
assert "index" in response_json["pages"][0]
assert "markdown" in response_json["pages"][0]
class TestRustOcrGateway:
def test_rust_ocr_models_are_on_gateway_config(self) -> None:
config = yaml.safe_load(CONFIG_PATH.read_text())
configured_models = {
model_config["model_name"] for model_config in config["model_list"]
}
expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES}
assert expected_models.issubset(configured_models)
def test_running_gateway_loaded_rust_ocr_models(
self, resources: OcrResources
) -> None:
expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES}
assert expected_models.issubset(resources.gateway.model_names())
@pytest.mark.parametrize(("model", "document"), RUST_OCR_GATEWAY_CASES)
def test_rust_ocr_model_gateway_response(
self, resources: OcrResources, model: str, document: dict[str, str]
) -> None:
response = resources.gateway.ocr(model, document)
assert response.status_code == 200, response.text
_assert_ocr_response_shape(response.json())

174
tests/test_litellm/caching/test_redis_cache.py
View file

@ -517,3 +517,177 @@ async def test_async_lpop_with_float_redis_version(
# Verify the method completed without error
assert result is not None
# LIT-3374: the namespace must be applied uniformly across every key-taking
# Redis operation, not just get/set/increment. Before the fix these paths wrote
# or read raw keys, so with a namespace configured the prefixed keys other
# operations created were silently missed.
@pytest.mark.parametrize(
"namespace, raw_keys, expected_keys",
[
(None, ["{k:v}:tokens", "{k:v}:requests"], ["{k:v}:tokens", "{k:v}:requests"]),
(
"litellm_sandbox",
["{k:v}:tokens", "{k:v}:requests"],
["litellm_sandbox:{k:v}:tokens", "litellm_sandbox:{k:v}:requests"],
),
],
)
@pytest.mark.asyncio
async def test_async_register_script_namespaces_keys(
namespace, raw_keys, expected_keys, monkeypatch, redis_no_ping
):
"""The callable returned by async_register_script (used by the rate limiter
Lua scripts, pod-lock release, and budget limiters) must namespace every key
it is invoked with. The hash tag is preserved so cluster slotting is intact."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
registered_script = AsyncMock(return_value="ok")
mock_redis_instance = MagicMock()
mock_redis_instance.register_script = MagicMock(return_value=registered_script)
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
script = redis_cache.async_register_script("return 1")
result = await script(keys=raw_keys, args=[60])
assert result == "ok"
registered_script.assert_awaited_once_with(
keys=expected_keys, args=[60], client=None
)
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
@pytest.mark.asyncio
async def test_async_delete_cache_namespaces_key(
namespace, expected, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_redis_instance = AsyncMock()
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_delete_cache("k")
mock_redis_instance.delete.assert_awaited_once_with(expected)
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
@pytest.mark.asyncio
async def test_delete_cache_keys_namespaces_keys(
namespace, expected, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_redis_instance = AsyncMock()
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.delete_cache_keys(["k"])
mock_redis_instance.delete.assert_awaited_once_with(expected)
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
@pytest.mark.asyncio
async def test_async_get_ttl_namespaces_key(
namespace, expected, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_redis_instance = AsyncMock()
mock_redis_instance.ttl = AsyncMock(return_value=42)
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
ttl = await redis_cache.async_get_ttl("k")
assert ttl == 42
mock_redis_instance.ttl.assert_awaited_once_with(expected)
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
@pytest.mark.asyncio
async def test_async_lpop_namespaces_key(
namespace, expected, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_redis_instance = AsyncMock()
mock_redis_instance.lpop = AsyncMock(return_value=b"value")
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_lpop(key="k")
mock_redis_instance.lpop.assert_awaited_once_with(expected, None)
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
@pytest.mark.asyncio
async def test_async_rpush_namespaces_key(
namespace, expected, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_redis_instance = AsyncMock()
mock_redis_instance.rpush = AsyncMock(return_value=1)
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_rpush("k", ["v"])
mock_redis_instance.rpush.assert_awaited_once_with(expected, "v")
@pytest.mark.parametrize("namespace, expected_match", [(None, "k*"), ("ns", "ns:k*")])
@pytest.mark.asyncio
async def test_async_scan_iter_namespaces_pattern(
namespace, expected_match, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
captured = {}
def scan_iter(match, count):
captured["match"] = match
async def gen():
for _ in ():
yield _
return gen()
mock_redis_instance = MagicMock()
mock_redis_instance.scan_iter = scan_iter
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_scan_iter(pattern="k")
assert captured["match"] == expected_match
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
def test_increment_cache_namespaces_key(
namespace, expected, monkeypatch, redis_no_ping
):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_client = MagicMock()
mock_client.incr.return_value = 5
mock_client.ttl.return_value = 100
redis_cache.redis_client = mock_client
redis_cache.increment_cache(key="k", value=1)
mock_client.incr.assert_called_once_with(name=expected, amount=1)
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=namespace)
mock_client = MagicMock()
redis_cache.redis_client = mock_client
redis_cache.delete_cache(key="k")
mock_client.delete.assert_called_once_with(expected)

103
tests/test_litellm/integrations/test_opentelemetry.py
View file

@ -4690,6 +4690,109 @@ class TestOpenTelemetrySpanDedupe(unittest.TestCase):
self.assertTrue(otel._emit_once(kwargs, "success"))
self.assertFalse(otel._emit_once(kwargs, "success"))
def test_emit_once_accepts_list_valued_scope_part(self):
"""Regression for LIT-3428 / LIT-3764: a list-valued ``guardrail_mode``
(the shape Presidio expands to with ``output_parse_pii: true``) must
not raise ``TypeError: unhashable type: 'list'`` when building the
dedupe key. Pre-fix, this call crashed inside ``dict.get``."""
otel = OpenTelemetry()
kwargs = self._build_kwargs()
self.assertTrue(
otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"])
)
self.assertFalse(
otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]),
"Same list scope must dedupe to False on the second call",
)
def test_emit_once_distinct_list_scopes_dont_collide(self):
"""Two different list-valued scopes on the same handler/kwargs must
each emit exactly once. Catches a regression where every list collapses
to the same key (e.g. ``str(list)`` collisions on near-identical input)."""
otel = OpenTelemetry()
kwargs = self._build_kwargs()
self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call"]))
self.assertTrue(
otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]),
"Distinct list scopes must produce distinct dedupe keys",
)
self.assertFalse(otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call"]))
self.assertFalse(
otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"])
)
def test_emit_once_accepts_dict_and_set_scope_parts(self):
"""``guardrail_mode`` can also arrive as a ``GuardrailMode`` TypedDict
(i.e. a plain dict at runtime). Sets are not produced today but flow
through the same normalization. Both must hash without raising."""
otel = OpenTelemetry()
kwargs = self._build_kwargs()
self.assertTrue(
otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"tags": ["pre", "post"]})
)
self.assertFalse(
otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"tags": ["pre", "post"]})
)
self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"a", "b"}))
def test_emit_once_handles_self_referential_scope_without_recursion_error(self):
"""``_freeze_for_dedupe`` caps recursion at ``_FREEZE_MAX_DEPTH`` and
falls back to ``repr`` past the cap, so a self-referential container
in scope must not crash ``_emit_once``. ``guardrail_mode`` cannot
construct such input today, but the cap is the bound that justifies
recursion on the logging hot path."""
otel = OpenTelemetry()
kwargs = self._build_kwargs()
cyclic: list = []
cyclic.append(cyclic)
self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, cyclic))
self.assertFalse(otel._emit_once(kwargs, "guardrail", "pii", 1.0, cyclic))
def test_create_guardrail_span_does_not_raise_on_list_mode(self):
"""End-to-end regression for LIT-3428: ``_create_guardrail_span``
must produce exactly one span (not raise ``TypeError``) when the
guardrail entry's ``guardrail_mode`` is a list."""
span_exporter = InMemorySpanExporter()
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))
otel = OpenTelemetry(tracer_provider=tracer_provider)
otel.tracer = tracer_provider.get_tracer(__name__)
kwargs = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
"litellm_params": {"custom_llm_provider": "openai", "metadata": {}},
"standard_logging_object": {
"id": "test-id",
"call_type": "completion",
"metadata": {},
"hidden_params": {},
"guardrail_information": [
{
"guardrail_name": "presidio-pii",
"guardrail_mode": ["pre_call", "post_call"],
"guardrail_response": "ok",
"start_time": 1.0,
"end_time": 2.0,
}
],
},
}
otel._create_guardrail_span(kwargs=kwargs, context=None)
otel._create_guardrail_span(kwargs=kwargs, context=None)
guardrail_spans = [
s for s in span_exporter.get_finished_spans() if s.name == "guardrail"
]
self.assertEqual(
len(guardrail_spans),
1,
"List-valued guardrail_mode must emit exactly one guardrail span "
"across repeated lifecycle entrypoints",
)
def test_handle_success_emits_single_litellm_request_span_on_double_call(self):
"""Sync + async callback paths firing for the same kwargs must
result in exactly one litellm_request span."""

550
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -1,8 +1,9 @@
"""Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``)."""
import importlib
import sys
import builtins
import types
from typing import Any
import httpx
import pytest
@ -15,11 +16,15 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse
# explicitly via importlib rather than attribute traversal.
ocr_main = importlib.import_module("litellm.ocr.main")
rust_bridge = importlib.import_module("litellm.ocr.rust_bridge")
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
MODEL = "mistral/mistral-ocr-latest"
DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"}
DOCUMENT: dict[str, object] = {
"type": "document_url",
"document_url": "https://example.com/doc.pdf",
}
FAKE_OCR_RESPONSE = {
FAKE_OCR_RESPONSE: dict[str, object] = {
"pages": [{"index": 0, "markdown": "hello world"}],
"model": "mistral-ocr-2505-completion",
"document_annotation": None,
@ -28,21 +33,35 @@ FAKE_OCR_RESPONSE = {
}
class CapturedException(Exception):
pass
class RecordingBridge:
"""A fake ``RustOcr`` callable that records the args it was handed."""
def __init__(self):
self.calls = []
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
def __call__(
self, model, document, api_key, api_base, optional_params, timeout_seconds
):
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append(
{
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": timeout_seconds,
}
@ -50,13 +69,81 @@ class RecordingBridge:
return dict(FAKE_OCR_RESPONSE)
class RecordingAsyncBridge:
"""A fake async ``RustAocr`` callable that records the args it was handed."""
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
async def __call__(
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append(
{
"model": model,
"document": document,
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"optional_params": optional_params,
"timeout_seconds": timeout_seconds,
}
)
return dict(FAKE_OCR_RESPONSE)
class RaisingBridge:
def __call__(
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
raise RuntimeError("bridge failed")
class RaisingAsyncBridge:
async def __call__(
self,
model: str,
document: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str,
extra_headers: dict[str, object] | None,
optional_params: dict[str, object],
timeout_seconds: float | None,
) -> dict[str, object]:
raise RuntimeError("bridge failed")
class RecordingLogging:
"""A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``."""
def __init__(self):
self.pre_call_kwargs = None
def __init__(self) -> None:
self.pre_call_kwargs: dict[str, object] | None = None
def pre_call(self, *, input, api_key, additional_args):
def pre_call(
self,
*,
input: str,
api_key: str | None,
additional_args: dict[str, object],
) -> None:
self.pre_call_kwargs = {
"input": input,
"api_key": api_key,
@ -67,21 +154,71 @@ class RecordingLogging:
class FakeOCRConfig:
"""A stand-in ``BaseOCRConfig`` that echoes the request it would build."""
def validate_environment(
self, *, headers, model, api_key, api_base, litellm_params
):
return {"authorization": f"Bearer {api_key}"}
def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None:
self.api_key_env_var = api_key_env_var
def get_complete_url(self, *, api_base, model, optional_params, litellm_params):
def get_api_key_env_var(self) -> str:
return self.api_key_env_var
def validate_environment(
self,
*,
headers: dict[str, object],
model: str,
api_key: str | None,
api_base: str | None,
litellm_params: dict[str, object],
) -> dict[str, object]:
return {"Authorization": f"Bearer {api_key}", **headers}
def get_complete_url(
self,
*,
api_base: str | None,
model: str,
optional_params: dict[str, object],
litellm_params: dict[str, object],
) -> str:
return f"{api_base or 'https://api.mistral.ai/v1'}/ocr"
def build_prepared_request(
*,
logging_obj: RecordingLogging | None = None,
provider_config: FakeOCRConfig | None = None,
model: str = "mistral-ocr-latest",
document: dict[str, object] = DOCUMENT,
api_key: str | None = "sk-test",
api_base: str | None = None,
custom_llm_provider: str = "mistral",
extra_headers: dict[str, object] | None = None,
optional_params: dict[str, object] | None = None,
litellm_params: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = 12.5,
) -> Any:
return ocr_main._PreparedOCRRequest(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
provider_config=provider_config or FakeOCRConfig(),
optional_params=optional_params or {},
litellm_params=litellm_params or {},
effective_timeout=timeout,
litellm_logging_obj=logging_obj or RecordingLogging(),
)
@pytest.fixture(autouse=True)
def _reset_rust_flag():
"""Keep the global toggle isolated between tests."""
rust_bridge.use_litellm_rust(False, ocr=None)
rust_bridge.use_litellm_rust(False, ocr=None, aocr=None)
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
yield
rust_bridge.use_litellm_rust(False, ocr=None)
rust_bridge.use_litellm_rust(False, ocr=None, aocr=None)
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
@pytest.fixture
@ -92,6 +229,14 @@ def fake_bridge():
return bridge
@pytest.fixture
def fake_async_bridge():
"""Enable the async Rust path with an injected recording bridge."""
bridge = RecordingAsyncBridge()
litellm.use_litellm_rust(True, aocr=bridge)
return bridge
def test_use_litellm_rust_toggles_flag():
assert rust_bridge.rust_ocr_enabled() is False
litellm.use_litellm_rust()
@ -100,12 +245,61 @@ def test_use_litellm_rust_toggles_flag():
assert rust_bridge.rust_ocr_enabled() is False
def test_env_var_enables_rust_ocr(monkeypatch):
monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1")
assert rust_bridge._env_enables_rust_ocr() is True
def test_load_rust_ocr_returns_injected_impl():
bridge = RecordingBridge()
litellm.use_litellm_rust(True, ocr=bridge)
assert rust_bridge.load_rust_ocr() is bridge
def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch):
real_import = builtins.__import__
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
if name == "litellm.rust_bridge" and "_native" in fromlist:
raise ImportError
return real_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(builtins, "__import__", fake_import)
assert rust_bridge_loader.get_native_bridge() is None
def test_native_bridge_loader_caches_absent_extension(monkeypatch):
real_import = builtins.__import__
attempts = 0
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
nonlocal attempts
if name == "litellm.rust_bridge" and "_native" in fromlist:
attempts += 1
raise ImportError
return real_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(builtins, "__import__", fake_import)
assert rust_bridge_loader.get_native_bridge() is None
assert rust_bridge_loader.get_native_bridge() is None
assert attempts == 1
def test_native_bridge_available_reflects_loader(monkeypatch):
fake_module = types.ModuleType("litellm.rust_bridge._native")
monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module)
assert rust_bridge_loader.native_bridge_available() is True
def test_load_rust_aocr_returns_injected_impl():
bridge = RecordingAsyncBridge()
litellm.use_litellm_rust(True, aocr=bridge)
assert rust_bridge.load_rust_aocr() is bridge
def test_toggle_without_ocr_arg_preserves_injected_impl():
"""Regression: routine enable/disable calls must not clobber a prior injection.
@ -114,39 +308,61 @@ def test_toggle_without_ocr_arg_preserves_injected_impl():
a caller toggled the flag without re-passing ``ocr=``.
"""
bridge = RecordingBridge()
litellm.use_litellm_rust(True, ocr=bridge)
async_bridge = RecordingAsyncBridge()
litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge)
litellm.use_litellm_rust(False)
assert rust_bridge.load_rust_ocr() is bridge
assert rust_bridge.load_rust_aocr() is async_bridge
litellm.use_litellm_rust(True)
assert rust_bridge.load_rust_ocr() is bridge
assert rust_bridge.load_rust_aocr() is async_bridge
def test_explicit_ocr_none_clears_injected_impl():
def test_explicit_ocr_none_clears_injected_impl(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
"get_native_bridge",
lambda: None,
)
bridge = RecordingBridge()
litellm.use_litellm_rust(True, ocr=bridge)
async_bridge = RecordingAsyncBridge()
litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge)
litellm.use_litellm_rust(True, ocr=None)
litellm.use_litellm_rust(True, ocr=None, aocr=None)
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
def test_load_rust_ocr_none_when_extension_absent():
def test_load_rust_ocr_none_when_extension_absent(monkeypatch):
"""With no injected impl and no compiled wheel, the loader returns None so the
caller degrades to the Python path instead of raising ImportError."""
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
"get_native_bridge",
lambda: None,
)
litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI
assert rust_bridge.load_rust_ocr() is None
assert rust_bridge.load_rust_aocr() is None
def test_load_rust_ocr_uses_compiled_extension(monkeypatch):
"""With no injected impl but a compiled ``litellm_python_bridge`` importable,
"""With no injected impl but a packaged ``litellm.rust_bridge._native`` importable,
the loader returns the extension's ``ocr`` callable. The native wheel isn't
built in CI, so stand in a fake module via ``sys.modules``."""
fake_module = types.ModuleType("litellm_python_bridge")
built in CI, so stand in a fake module via the bridge loader."""
fake_module = types.ModuleType("litellm.rust_bridge._native")
fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module)
fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined]
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
"get_native_bridge",
lambda: fake_module,
)
litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension
assert rust_bridge.load_rust_ocr() is fake_module.ocr
assert rust_bridge.load_rust_aocr() is fake_module.aocr
def test_timeout_to_seconds_handles_float_timeout_and_none():
@ -161,16 +377,14 @@ def test_run_rust_ocr_forwards_args_and_wraps_response():
response = ocr_main._run_rust_ocr(
rust_ocr=bridge,
logging_obj=logging_obj,
provider_config=FakeOCRConfig(),
prepared_request=build_prepared_request(
logging_obj=logging_obj,
api_base="https://proxy.internal",
extra_headers={"x-trace-id": "trace-1"},
optional_params={"include_image_base64": True},
timeout=12.5,
),
resolve_api_key=lambda _name: None,
model="mistral-ocr-latest",
document=DOCUMENT,
api_key="sk-test",
api_base="https://proxy.internal",
optional_params={"include_image_base64": True},
litellm_params={},
timeout_seconds=12.5,
)
assert isinstance(response, OCRResponse)
@ -181,6 +395,11 @@ def test_run_rust_ocr_forwards_args_and_wraps_response():
"document": DOCUMENT,
"api_key": "sk-test",
"api_base": "https://proxy.internal",
"custom_llm_provider": "mistral",
"extra_headers": {
"Authorization": "Bearer sk-test",
"x-trace-id": "trace-1",
},
"optional_params": {"include_image_base64": True},
"timeout_seconds": 12.5,
}
@ -193,23 +412,127 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing():
ocr_main._run_rust_ocr(
rust_ocr=bridge,
logging_obj=RecordingLogging(),
provider_config=FakeOCRConfig(),
prepared_request=build_prepared_request(api_key=None, timeout=None),
resolve_api_key=lambda name: (
"sk-from-vault" if name == "MISTRAL_API_KEY" else None
),
model="mistral-ocr-latest",
document=DOCUMENT,
api_key=None,
api_base=None,
optional_params={},
litellm_params={},
timeout_seconds=None,
)
assert bridge.calls[0]["api_key"] == "sk-from-vault"
def test_run_rust_ocr_uses_provider_api_key_env_var():
bridge = RecordingBridge()
resolver_calls = []
def _resolver(name):
resolver_calls.append(name)
return "sk-provider-env"
ocr_main._run_rust_ocr(
rust_ocr=bridge,
prepared_request=build_prepared_request(
provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"),
model="provider-ocr-model",
api_key=None,
timeout=None,
),
resolve_api_key=_resolver,
)
assert resolver_calls == ["PROVIDER_OCR_API_KEY"]
assert bridge.calls[0]["api_key"] == "sk-provider-env"
def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata():
bridge = RecordingBridge()
ocr_main._run_rust_ocr(
rust_ocr=bridge,
prepared_request=build_prepared_request(
custom_llm_provider="vertex_ai",
model="mistral-ocr-maas",
litellm_params={
"vertex_project": "project-1",
"vertex_location": "us-central1",
"vertex_credentials": "redacted",
},
optional_params={"include_image_base64": True},
timeout=None,
),
resolve_api_key=lambda _name: None,
)
assert bridge.calls[0]["optional_params"] == {
"include_image_base64": True,
"vertex_project": "project-1",
"vertex_location": "us-central1",
}
def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager():
bridge = RecordingBridge()
def _resolver(name: str) -> str | None:
return {
"VERTEXAI_PROJECT": "project-from-secret",
"VERTEXAI_LOCATION": "us-east5",
}.get(name)
ocr_main._run_rust_ocr(
rust_ocr=bridge,
prepared_request=build_prepared_request(
custom_llm_provider="vertex_ai",
model="mistral-ocr-maas",
timeout=None,
),
resolve_api_key=_resolver,
)
assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret"
assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5"
def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager():
bridge = RecordingBridge()
ocr_main._run_rust_ocr(
rust_ocr=bridge,
prepared_request=build_prepared_request(
custom_llm_provider="azure_ai",
model="pixtral-12b-2409",
api_base=None,
timeout=None,
),
resolve_api_key=lambda name: (
"https://azure.example.com" if name == "AZURE_AI_API_BASE" else None
),
)
assert bridge.calls[0]["api_base"] == "https://azure.example.com"
def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint():
bridge = RecordingBridge()
ocr_main._run_rust_ocr(
rust_ocr=bridge,
prepared_request=build_prepared_request(
custom_llm_provider="azure_ai/doc-intelligence",
model="prebuilt-layout",
api_base=None,
timeout=None,
),
resolve_api_key=lambda name: (
"https://document-intelligence.example.com"
if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"
else None
),
)
assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com"
def test_run_rust_ocr_prefers_explicit_key_over_resolver():
bridge = RecordingBridge()
resolver_calls = []
@ -220,16 +543,8 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver():
ocr_main._run_rust_ocr(
rust_ocr=bridge,
logging_obj=RecordingLogging(),
provider_config=FakeOCRConfig(),
prepared_request=build_prepared_request(api_key="sk-explicit", timeout=None),
resolve_api_key=_resolver,
model="mistral-ocr-latest",
document=DOCUMENT,
api_key="sk-explicit",
api_base=None,
optional_params={},
litellm_params={},
timeout_seconds=None,
)
assert bridge.calls[0]["api_key"] == "sk-explicit"
@ -242,16 +557,14 @@ def test_run_rust_ocr_runs_pre_call_logging():
ocr_main._run_rust_ocr(
rust_ocr=RecordingBridge(),
logging_obj=logging_obj,
provider_config=FakeOCRConfig(),
prepared_request=build_prepared_request(
logging_obj=logging_obj,
api_base="https://api.mistral.ai/v1",
extra_headers={"x-trace-id": "trace-1"},
optional_params={"include_image_base64": True},
timeout=None,
),
resolve_api_key=lambda _name: None,
model="mistral-ocr-latest",
document=DOCUMENT,
api_key="sk-test",
api_base="https://api.mistral.ai/v1",
optional_params={"include_image_base64": True},
litellm_params={},
timeout_seconds=None,
)
assert logging_obj.pre_call_kwargs is not None
@ -262,7 +575,10 @@ def test_run_rust_ocr_runs_pre_call_logging():
assert complete_input["include_image_base64"] is True
# The logged request mirrors what Rust sends: resolved URL + headers.
assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr"
assert additional_args["headers"] == {"authorization": "Bearer sk-test"}
assert additional_args["headers"] == {
"Authorization": "Bearer sk-test",
"x-trace-id": "trace-1",
}
def test_ocr_routes_to_rust_when_enabled(fake_bridge):
@ -270,6 +586,7 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge):
model=MODEL,
document=DOCUMENT,
api_key="sk-test",
extra_headers={"x-trace-id": "trace-1"},
include_image_base64=True,
)
@ -281,10 +598,93 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge):
assert call["model"] == "mistral-ocr-latest"
assert call["document"] == DOCUMENT
assert call["api_key"] == "sk-test"
assert call["custom_llm_provider"] == "mistral"
assert call["extra_headers"] == {
"Authorization": "Bearer sk-test",
"x-trace-id": "trace-1",
}
# Raw OCR params ride along in optional_params; Rust filters to supported keys.
assert call["optional_params"].get("include_image_base64") is True
def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge):
response = litellm.ocr(
model="azure_ai/pixtral-12b-2409",
document=DOCUMENT,
api_key="sk-test",
api_base="https://example.services.ai.azure.com",
)
assert isinstance(response, OCRResponse)
assert len(fake_bridge.calls) == 1
assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409"
assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai"
def test_ocr_exception_type_uses_resolved_provider_context(
monkeypatch: pytest.MonkeyPatch,
):
captured: dict[str, object] = {}
def fake_exception_type(**kwargs: object) -> CapturedException:
captured.update(kwargs)
return CapturedException("wrapped")
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
litellm.use_litellm_rust(True, ocr=RaisingBridge())
with pytest.raises(CapturedException):
litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
assert captured["model"] == "mistral-ocr-latest"
assert captured["custom_llm_provider"] == "mistral"
@pytest.mark.asyncio
async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge):
response = await litellm.aocr(
model=MODEL,
document=DOCUMENT,
api_key="sk-test",
extra_headers={"x-trace-id": "trace-1"},
include_image_base64=True,
)
assert isinstance(response, OCRResponse)
assert response.pages[0].markdown == "hello world"
assert len(fake_async_bridge.calls) == 1
call = fake_async_bridge.calls[0]
assert call["model"] == "mistral-ocr-latest"
assert call["document"] == DOCUMENT
assert call["api_key"] == "sk-test"
assert call["custom_llm_provider"] == "mistral"
assert call["extra_headers"] == {
"Authorization": "Bearer sk-test",
"x-trace-id": "trace-1",
}
assert call["optional_params"].get("include_image_base64") is True
@pytest.mark.asyncio
async def test_aocr_exception_type_uses_resolved_provider_context(
monkeypatch: pytest.MonkeyPatch,
):
captured: dict[str, object] = {}
def fake_exception_type(**kwargs: object) -> CapturedException:
captured.update(kwargs)
return CapturedException("wrapped")
monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type)
litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge())
with pytest.raises(CapturedException):
await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test")
assert captured["model"] == "mistral-ocr-latest"
assert captured["custom_llm_provider"] == "mistral"
def test_ocr_forwards_timeout_to_rust(fake_bridge):
"""Caller-supplied timeout must flow into the Rust bridge so the fixed 600s
client ceiling doesn't silently override shorter deadlines."""
@ -317,6 +717,7 @@ def test_ocr_does_not_route_to_rust_when_disabled():
def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch):
"""Rust enabled but no bridge available (no injected impl, no compiled wheel):
ocr() must degrade to the Python HTTP handler instead of raising."""
monkeypatch.setattr(ocr_main, "load_rust_ocr", lambda: None)
litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI
captured = {}
@ -331,3 +732,26 @@ def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch):
assert captured.get("called") is True # Python path was used
assert isinstance(response, OCRResponse)
def test_ocr_provider_configs_expose_api_key_env_vars():
from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
AzureDocumentIntelligenceOCRConfig,
)
from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
from litellm.llms.vertex_ai.ocr.deepseek_transformation import (
VertexAIDeepSeekOCRConfig,
)
from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig
assert BaseOCRConfig().get_api_key_env_var() is None
assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY"
assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY"
assert (
AzureDocumentIntelligenceOCRConfig().get_api_key_env_var()
== "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"
)
assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY"
assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY"

33
tests/test_litellm/test_budget_ratchet_check.py
View file

@ -1,7 +1,8 @@
"""Tests for scripts/budget_ratchet_check.py.
The guard's whole contract is "ceilings may only fall": a raised ceiling, a dropped
rule, or a deleted file is a regression, while a lowered/equal ceiling, a brand-new
The guard's contract is "baselines and ceilings may only fall": a raised ceiling, a
raised baseline (even when slack is cut to keep the ceiling flat), a dropped rule, or
a deleted file is a regression, while a lowered/equal baseline and ceiling, a brand-new
rule, or a brand-new budget file is fine. Each branch is pinned here.
"""
@ -10,7 +11,9 @@ import subprocess
import sys
from pathlib import Path
_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py"
_MODULE_PATH = (
Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py"
)
_spec = importlib.util.spec_from_file_location("budget_ratchet_check", _MODULE_PATH)
ratchet = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(ratchet)
@ -35,10 +38,30 @@ def test_raised_ceiling_is_a_regression():
def test_lowered_or_equal_ceiling_is_clean():
base = {"LIT006": _spec_of(1013, 10)}
# baseline drops, slack flat -> ceiling falls
assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1000, 10)}) == []
# nothing changes
assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 10)}) == []
# slack traded for baseline at the same ceiling is fine
assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)}) == []
# slack cut while baseline holds -> ceiling falls, baseline flat
assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 0)}) == []
def test_raised_baseline_is_a_regression_even_when_ceiling_held_flat():
# baseline 1013 -> 1023 with slack cut 10 -> 0 keeps the ceiling at 1023, but a
# higher baseline bakes in more accepted debt and must still surface as a regression
base = {"LIT006": _spec_of(1013, 10)}
regs = ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)})
assert [r.rule for r in regs] == ["LIT006"]
assert "baseline raised 1013 -> 1023" in regs[0].detail
assert "ceiling raised" not in regs[0].detail
def test_raised_baseline_and_ceiling_report_both_reasons():
base = {"LIT006": _spec_of(1013, 10)}
regs = ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1100, 10)})
assert [r.rule for r in regs] == ["LIT006"]
assert "ceiling raised 1023 -> 1110" in regs[0].detail
assert "baseline raised 1013 -> 1100" in regs[0].detail
def test_dropped_rule_is_a_regression():

8
uv.lock generated
View file

@ -9,7 +9,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-06-22T14:51:56.5801Z"
exclude-newer = "2026-06-22T19:05:55.080417Z"
exclude-newer-span = "P3D"
[manifest]
@ -3121,15 +3121,15 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0"
version = "4.1.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "langchain-core" },
{ name = "ormsgpack" },
]
sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" }
sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" },
{ url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" },
]
[[package]]