scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 156

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

Browse source
This commit is contained in:
mateo-berri 2026-09-04 20:32:45 -07:00
commit d5bd788314
570 changed files with 45032 additions and 9492 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
.circleci/config.yml
View file

@ -575,7 +575,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"langfuse\""
@ -630,7 +630,7 @@ jobs:
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 \
-v -x \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -737,7 +737,7 @@ jobs:
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 \
-v -x \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -782,7 +782,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"assistants\""
@ -909,7 +909,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
@ -999,7 +999,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1054,7 +1054,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 8 \
@ -1090,7 +1090,7 @@ jobs:
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 \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1134,7 +1134,7 @@ jobs:
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 \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1178,7 +1178,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1222,7 +1222,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1267,7 +1267,7 @@ jobs:
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 \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1312,7 +1312,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -1391,7 +1391,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1444,7 +1444,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 -n 2 \
@ -1705,7 +1705,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -1794,7 +1794,7 @@ jobs:
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 \
-s -v -x \
-s -v \
--junitxml=test-results/junit.xml \
-n 4 \
--durations=5"
@ -2012,7 +2012,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -2092,7 +2092,7 @@ jobs:
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 \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2195,7 +2195,7 @@ jobs:
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 \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2266,7 +2266,7 @@ jobs:
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 \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2350,7 +2350,7 @@ jobs:
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 \
-vv -x \
-vv \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -2446,7 +2446,7 @@ jobs:
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 \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2516,7 +2516,7 @@ jobs:
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 \
-vv -x -s \
-vv -s \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m

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

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

5
.github/ci-coverage-allowlist.yml vendored
View file

@ -79,6 +79,11 @@ test_paths:
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
locally rather than in pull-request jobs
paths:
- tests/load_tests/test_granian_admission_saturation.py
- reason: >-
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a

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

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

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

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

18
basedpyright-code-budget.json
View file

@ -3,7 +3,7 @@
"limit": 14074
},
"reportArgumentType": {
"limit": 2214
"limit": 2206
},
"reportAssignmentType": {
"limit": 319
@ -18,7 +18,7 @@
"limit": 40
},
"reportDeprecated": {
"limit": 211
"limit": 209
},
"reportDuplicateImport": {
"limit": 19
@ -45,7 +45,7 @@
"limit": 24
},
"reportInvalidTypeForm": {
"limit": 34
"limit": 30
},
"reportInvalidTypeVarUse": {
"limit": 2
@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15287
"limit": 15284
},
"reportMissingTypeStubs": {
"limit": 40
@ -93,25 +93,25 @@
"limit": 181
},
"reportTypedDictNotRequiredAccess": {
"limit": 24
"limit": 22
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44362
"limit": 44358
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38323
"limit": 38309
},
"reportUnknownParameterType": {
"limit": 19624
"limit": 19621
},
"reportUnknownVariableType": {
"limit": 29861
"limit": 29844
},
"reportUnnecessaryCast": {
"limit": 111

1
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -161,6 +161,7 @@ class CheckBatchCost:
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql Normal file
View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "cost" DOUBLE PRECISION;
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "untracked_units" BIGINT NOT NULL DEFAULT 0;

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

@ -5,6 +5,7 @@ datasource client {
generator client {
provider = "prisma-client-py"
recursive_type_depth = -1
binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"]
}
@ -1123,6 +1124,8 @@ model LiteLLM_DailyGuardrailUsageUnits {
api_key String // hashed virtual key; empty string when unknown
usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits
units BigInt @default(0)
cost Float? // USD for the priced share of units; null only on rows written before this column existed
untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out
created_at DateTime @default(now())
updated_at DateTime @updatedAt

5
litellm-rust/AGENTS.md
View file

@ -1,17 +1,18 @@
# AGENTS.md
litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
## Crates
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate.
## Where a route lives

1
litellm-rust/CLAUDE.md
View file

@ -24,6 +24,7 @@ the base when behavior is genuinely different, and say so explicitly in the PR.
## Crates (see AGENTS.md)
`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
`litellm-config` is the config-loading boundary and returns resolved core types.
`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop`
holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate

33
litellm-rust/Cargo.lock generated
View file

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

4
litellm-rust/Cargo.toml
View file

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

4
litellm-rust/README.md
View file

@ -25,11 +25,12 @@ coverage and production evidence.
| Crate | Role |
|-------|------|
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
## Layout
@ -38,6 +39,7 @@ crates/
core/ The SDK: route modules + provider transforms.
src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client
src/providers/anthropic/messages/transformation.rs
config/ Config loading and resolved deployments.
ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints.
python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
python-bridge/ PyO3 API adapter for Python LiteLLM.

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

@ -45,7 +45,7 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages`
22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
## Checks before push

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

@ -9,19 +9,15 @@ such as `litellm_core::messages::messages`. No provider handler lives here.
src/
main.rs # entrypoint: build AppState (router + master key), bind, serve
state.rs # AppState — shared Arc<Router> + master_key
gil.rs # GIL-activity tracker (records Python acquisitions)
auth/ # authentication as an axum extractor — added to handler args
mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
routes/ # one module per route, all matching the same template
AGENTS.md # ← the route template (read this before adding a route)
mod.rs # app(): merges every module's router()
health.rs # simple route (one file): router() + liveness/readiness
gil.rs # simple route (one file): router() + GET /health/gil
realtime/ # route with logic → axum surface + a no-axum service:
mod.rs # router() + handler + WS<->events adapter (the axum surface)
service.rs # business logic (select deployment, call provider) — no axum, testable
python/ # Python interop (feature: python-config) — load-time only
mod.rs, config.rs, AGENTS.md
```
## Rules
@ -53,5 +49,6 @@ proxy in a later phase. Health routes don't add the extractor (unauthenticated).
## Python interop
Anything that calls into Python lives in `python/` and is **load-time only** — see
`python/AGENTS.md`. The realtime data path never takes the GIL.
Python-backed loading lives in `litellm-config` and is **load-time only**. The
gateway's `python-config` feature forwards to that crate. The realtime data path
never takes the GIL.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

File diff suppressed because it is too large Load diff

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@ -1,6 +1,7 @@
//! Enforcement: the litellm-rust workspace has exactly four crates.
//! Enforcement: the litellm-rust workspace has exactly five crates.
//!
//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host),
//! `core` (the Rust SDK), `config` (the config-loading boundary),
//! `ai-gateway` (the HTTP/WebSocket host),
//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the
//! PyO3 cdylib). Adding or removing a crate must be a
//! deliberate act: this test fails until the allowlist here is updated, forcing
@ -19,13 +20,20 @@ use std::path::{Path, PathBuf};
/// workspace legitimately gains or loses a crate.
const EXPECTED_MEMBERS: &[&str] = &[
"crates/core",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
"crates/python-bridge",
];
/// The crate subdirectory names that must exist under `crates/`.
const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"];
const EXPECTED_CRATE_DIRS: &[&str] = &[
"core",
"config",
"ai-gateway",
"python-interop",
"python-bridge",
];
const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact).";

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@ -0,0 +1,29 @@
use pyo3::prelude::*;
use serde_json::Value;
use crate::errors::core_error_to_pyerr;
#[pyfunction]
fn gateway_messages<'py>(
py: Python<'py>,
model_alias: String,
provider_model: String,
api_base: String,
#[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value,
) -> PyResult<Bound<'py, PyAny>> {
let future = litellm_ai_gateway::trace_parity::messages_request(
model_alias,
provider_model,
api_base,
body,
);
crate::execution::run_async(
py,
crate::function_trace::capture(future),
core_error_to_pyerr,
)
}
pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> {
super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?)
}

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

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

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

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

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

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

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

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

2
litellm/__init__.py
View file

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

4
litellm/anthropic_interface/exceptions/exceptions.py
View file

@ -1,8 +1,9 @@
"""Anthropic error format type definitions."""
from collections.abc import Mapping
from typing import Literal
from typing_extensions import Required, TypedDict
from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
# Known Anthropic error types
# Source: https://docs.anthropic.com/en/api/errors
@ -23,6 +24,7 @@ class AnthropicErrorDetail(TypedDict):
type: AnthropicErrorType
message: str
provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]]
class AnthropicErrorResponse(TypedDict, total=False):

30
litellm/caching/caching.py
View file

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

93
litellm/caching/redis_cluster_node_isolation.py
View file

@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov
retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
topology changed.
redis-py 8.x fixed this upstream with gentler machinery than this override's
``node.disconnect()`` (which also kills connections other coroutines are mid-operation
on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
killed connection): it marks in-use connections for reconnect only after their current
operation completes, disconnects only the idle pooled ones, and defers reinitialization
to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
recovery API, the factory returns the base ``RedisCluster`` unmodified.
redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream
still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent
caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full
teardown. For those versions the factory returns a thin wrapper around upstream's
``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError,
a third consecutive timeout on the same node, or a concurrent request from any other command
or ``aclose()`` still reinits).
"""
import asyncio
@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol):
mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
own logic fully typed without a banned ``typing.cast``."""
name: str
async def execute_command(
self,
*args: object,
@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol):
#: this override can't see (Python won't error -- it'll just run our now-stale copy), so
#: construction logs a loud warning rather than silently trusting an unverified copy.
_VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3
def get_litellm_async_redis_cluster_class(
def get_litellm_async_redis_cluster_class( # noqa: C901 # supports redis-py version-specific cluster implementations
cluster_node_class: type | None = None,
base_cluster_class: type | None = None,
) -> type["_AsyncRedisClusterType"]:
"""Returns the base ``RedisCluster`` when the installed redis-py already recovers a
node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
subclass with the per-node isolation fix for older versions whose upstream branch
tears down the whole cluster client.
"""Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already
recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster``
subclass with the per-node isolation fix for older versions whose upstream branch tears
down the whole cluster client.
``cluster_node_class`` exists for dependency injection in tests; production callers
leave it unset and the installed ``ClusterNode`` is used.
``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests;
production callers leave them unset and the installed redis-py classes are used.
Imported lazily because this module is reachable from a base ``import litellm`` while
redis is not a base dependency. Cheap to call repeatedly: the underlying redis
@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class(
from redis.exceptions import TimeoutError as _RedisTimeoutError
node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster
if hasattr(node_class, "update_active_connections_for_reconnect"):
verbose_logger.debug(
"redis-py %s recovers a node-level connection error per-connection upstream; "
"using the base RedisCluster without litellm's node-isolation override.",
"redis-py %s recovers node connections per-connection upstream; using "
"LiteLLM's timeout-tolerant RedisCluster wrapper.",
redis.__version__,
)
return _BaseAsyncRedisCluster
class LiteLLMAsyncRedisClusterTimeoutTolerant(
base_class # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass] # the injected base class is selected at runtime
):
def __init__(
self,
*args: object,
**kwargs: object, # kwargs-ok: passes redis-py's constructor kwargs through untouched
) -> None:
self._litellm_initialize = False
self._litellm_reinit_requests = 0
self._litellm_tolerated_timeouts = 0
super().__init__(*args, **kwargs)
self._litellm_consecutive_timeouts: dict[ # mutable-ok: per-node counter updated on the command hot path
str, int
] = {}
@property
def _initialize(self) -> bool:
return self._litellm_initialize
@_initialize.setter
def _initialize(self, value: bool) -> None:
if value:
self._litellm_reinit_requests += 1
self._litellm_initialize = value
async def _execute_command(
self,
target_node: _ClusterNodeAttrs,
*args: object,
**kwargs: object, # kwargs-ok: matches redis-py's own command dispatch signature
) -> object:
outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts
pending_before: Final = self._litellm_initialize
try:
result: Final = await super()._execute_command(target_node, *args, **kwargs)
except _RedisTimeoutError:
timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1
if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT:
self._litellm_consecutive_timeouts.pop(target_node.name, None)
raise
self._litellm_consecutive_timeouts[target_node.name] = timeouts
self._litellm_tolerated_timeouts += 1
if (
not pending_before
and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before
):
self._initialize = False
raise
if self._litellm_consecutive_timeouts:
self._litellm_consecutive_timeouts.pop(target_node.name, None)
return result
return LiteLLMAsyncRedisClusterTimeoutTolerant
if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
verbose_logger.warning(

1
litellm/constants.py
View file

@ -151,6 +151,7 @@ DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200"))
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600"))
MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60"))
# Default npm cache directory for STDIO MCP servers.
# npm/npx needs a writable cache dir; in containers the default (~/.npm)

2
litellm/integrations/azure_storage/azure_storage.py
View file

@ -36,7 +36,7 @@ AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
def _cached_credential_chain_token_provider() -> Callable[[], str]:
return get_azure_ad_token_provider(
azure_scope=AZURE_STORAGE_TOKEN_SCOPE,
azure_credential=AzureCredentialType.DefaultAzureCredential,
azure_credential=AzureCredentialType.DeploymentIdentityCredential,
)

16
litellm/integrations/cloudzero/database.py
View file

@ -94,8 +94,18 @@ class LiteLLMDatabase:
try:
db_response: Final = await client.db.query_raw(query, *params)
# Convert the response to polars DataFrame with full schema inference
# This prevents schema mismatch errors when data types vary across rows
return pl.DataFrame(db_response, infer_schema_length=None)
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
)
usage_rows: Final = (
db_response.to_dicts()
if isinstance(db_response, pl.DataFrame)
else db_response
if isinstance(db_response, list)
else []
)
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
except Exception as e:
raise Exception(f"Error retrieving usage data: {e}")

64
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -19,7 +19,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.constants import REDACTED_BY_LITELLM
from litellm.constants import REDACTED_BY_LITELLM, REDACTED_BY_LITELM_STRING
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_handler import (
get_datadog_base_url_from_env,
@ -46,9 +46,10 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
from litellm.types.integrations.datadog_llm_obs import *
from litellm.types.utils import (
AUDIT_GUARDRAIL_FIELDS,
PROMPT_CARRYING_GUARDRAIL_FIELDS,
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
CallTypes,
StandardLoggingGuardrailInformation,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -60,6 +61,8 @@ _SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset(
{"agent", "assistant", "developer", "function", "model", "system", "tool", "user"}
)
_CLASSIFIED_GUARDRAIL_FIELDS: Final = AUDIT_GUARDRAIL_FIELDS | PROMPT_CARRYING_GUARDRAIL_FIELDS
_PROMPT_CARRYING_METADATA_FIELDS: Final = frozenset(
{
"routing_decision",
@ -108,6 +111,49 @@ def _router_span_fields(
)
def _guardrail_entries(guardrail_information: object) -> tuple[Mapping[str, object], ...]:
"""The guardrail records as a sequence, whatever shape the payload carries.
`guardrail_information` is typed as a list, but a guardrail that writes the metadata key itself
can leave a single record there; Prometheus normalizes the same shape at
`_guardrail_overhead_seconds`.
"""
if isinstance(guardrail_information, Mapping):
return (guardrail_information,)
if isinstance(guardrail_information, (list, tuple)):
return tuple(entry for entry in guardrail_information if isinstance(entry, Mapping))
return ()
def _guardrail_entry_without_prompt_carriers(entry: Mapping[str, object]) -> Mapping[str, object]:
"""One guardrail record kept as its audit fields, with the prompt-quoting ones marked redacted.
Built as an allow-list rather than a deny-list: a key neither set classifies is dropped, so a
guardrail that records its own extra detail cannot put the caller's prompt on a redacted span.
"""
return { # mutable-ok: a fresh record built per entry, handed straight to the span serializer
field: REDACTED_BY_LITELM_STRING if field in PROMPT_CARRYING_GUARDRAIL_FIELDS else value
for field, value in entry.items()
if field in _CLASSIFIED_GUARDRAIL_FIELDS
}
def _guardrail_information_without_prompt_carriers(
guardrail_information: object,
) -> tuple[Mapping[str, object], ...] | None:
"""The guardrail records reduced to what a redacted span may carry.
Redaction removes the prompt, not the record that a guardrail ran: the name, mode, status,
timings and masked-entity counts are what an operator reads to answer whether a guardrail
caught anything on a request, and none of them reproduce the prompt. Field-level rather than
dropping the list, which is what `_sanitize_guardrail_information_for_spend_logs` already does
for spend logs.
"""
if guardrail_information is None:
return None
return tuple(_guardrail_entry_without_prompt_carriers(entry) for entry in _guardrail_entries(guardrail_information))
def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]:
"""The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text."""
return MappingProxyType(
@ -872,7 +918,9 @@ class DataDogLLMObsLogger(CustomBatchLogger):
"cache_key": standard_logging_payload.get("cache_key", "unknown"),
"saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0),
"guardrail_information": (
None if redact_prompt_text else standard_logging_payload.get("guardrail_information", None)
_guardrail_information_without_prompt_carriers(standard_logging_payload.get("guardrail_information"))
if redact_prompt_text
else standard_logging_payload.get("guardrail_information", None)
),
"is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload),
"latency_metrics": dict(self._get_latency_metrics(standard_logging_payload)),
@ -904,14 +952,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms
# Guardrail overhead latency
guardrail_info: Final[list[StandardLoggingGuardrailInformation] | None] = standard_logging_payload.get(
"guardrail_information"
)
if guardrail_info is not None:
guardrail_info: Final = _guardrail_entries(standard_logging_payload.get("guardrail_information"))
if guardrail_info:
total_duration = 0.0
for info in guardrail_info:
_guardrail_duration_seconds: float | None = info.get("duration")
if _guardrail_duration_seconds is not None:
_guardrail_duration_seconds = info.get("duration")
if isinstance(_guardrail_duration_seconds, (int, float, str)):
total_duration += float(_guardrail_duration_seconds)
if total_duration > 0:

14
litellm/integrations/focus/database.py
View file

@ -96,7 +96,19 @@ class FocusLiteLLMDatabase:
try:
db_response: Final = await client.db.query_raw(query, *query_params)
return pl.DataFrame(db_response, infer_schema_length=None)
from litellm.proxy.spend_tracking.key_metadata_recovery import (
fill_missing_api_key_aliases,
)
usage_rows: Final = (
db_response.to_dicts()
if isinstance(db_response, pl.DataFrame)
else db_response
if isinstance(db_response, list)
else []
)
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
except Exception as exc:
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc

163
litellm/integrations/shadow_eval_logger.py
View file

@ -165,28 +165,91 @@ def _chat_request_from_responses(
)
def _chat_final_text(response_obj: object) -> str:
"""The assistant's text, or empty when the turn carries tool calls: only text-final
turns produce a judgeable A/B comparison."""
def _chat_choice(response_obj: object) -> object | None:
"""The response's first choice, from a payload mapping or a duck-typed ModelResponse."""
try:
message: Final = (
response_obj["choices"][0]["message"]
if isinstance(response_obj, Mapping)
else response_obj.choices[0].message # pyright: ignore[reportAttributeAccessIssue] # duck-typed ModelResponse
)
if isinstance(response_obj, Mapping):
return response_obj["choices"][0]
return response_obj.choices[0] # pyright: ignore[reportAttributeAccessIssue] # duck-typed ModelResponse
except (AttributeError, KeyError, IndexError, TypeError):
return None
def _field_reader(obj: object) -> Callable[[str], object]:
return obj.get if isinstance(obj, Mapping) else lambda key: getattr(obj, key, None)
def _chat_message_reader(response_obj: object) -> Callable[[str], object] | None:
"""Field access over the assistant message of a chat response, or None for a payload
with no readable message."""
choice: Final = _chat_choice(response_obj)
if choice is None:
return None
message: Final = _field_reader(choice)("message")
return _field_reader(message) if message is not None else None
def _chat_final_text(response_obj: object) -> str:
"""The turn's judgeable text: prose, or every tool call serialized alongside it as
`[tool call] name(arguments)` when the assistant chose to act instead of, or as well
as, answering directly. A tool call is a real turn, not a gap, so this is what both
the real arm's sampling decision and the shadow arm's reply compare against."""
read: Final = _chat_message_reader(response_obj)
if read is None:
return ""
read: Final = message.get if isinstance(message, Mapping) else lambda key: getattr(message, key, None)
if read("tool_calls") or read("function_call"):
return ""
return extract_text_from_content(read("content"))
prose: Final = extract_text_from_content(read("content"))
if not (read("tool_calls") or read("function_call")):
return prose
serialized: Final = _serialize_tool_calls(read)
return f"{prose} {serialized}".strip() if prose else serialized
def _chat_finish_reason(response_obj: object) -> str:
choice: Final = _chat_choice(response_obj)
raw: Final = _field_reader(choice)("finish_reason") if choice is not None else None
return str(raw) if raw else "unknown"
_RESPONSES_TOOL_CALL_TYPES: Final = frozenset(("function_call", "custom_tool_call"))
def _tool_calls_list(read: Callable[[str], object]) -> tuple[object, ...]:
calls: Final = read("tool_calls")
listed: Final = tuple(calls) if isinstance(calls, Sequence) and not isinstance(calls, str) else ()
single: Final = read("function_call")
return listed if listed else ((single,) if single is not None else ())
def _tool_call_invocation(call: object) -> str:
"""One tool call as `name(arguments)`. Custom tool calls name themselves and carry their
arguments under `custom` rather than `function`."""
read_call: Final = _field_reader(call)
payload: Final = read_call("function") or read_call("custom") or call
read_payload: Final = _field_reader(payload)
name: Final = read_payload("name")
arguments: Final = read_payload("arguments") or read_payload("input") or ""
return f"{name or 'unnamed'}({arguments})"
def _serialize_tool_calls(read: Callable[[str], object]) -> str:
"""Every tool call in a reply as text a judge built for prose can still read."""
return ", ".join(f"[tool call] {_tool_call_invocation(call)}" for call in _tool_calls_list(read))
def _shadow_empty_reply_error(response_obj: object, routed_model: str) -> str:
"""Why a shadow reply yielded no judgeable text at all: no prose, and no tool call to
serialize either. The stable sentence comes first and every varying part after the
semicolon, so grouping rows by error still yields one row per cause."""
detail: Final = f"finish_reason={_chat_finish_reason(response_obj)}, model={routed_model or 'unknown'}"
return f"shadow router returned an empty response; {detail}"
def _responses_final_text(response_obj: object) -> str:
"""The turn's aggregated output text, or empty when the turn carries tool calls. A
dict-shaped payload is validated into the owner type first, because ``output_text``
is a derived property rather than a serialized field, so it never exists on a dict;
a dict the owner type rejects is unjudgeable and skipped."""
"""The turn's judgeable text: the aggregated output plus any tool call serialized
alongside it, the same way the chat surface renders one. A dict-shaped payload is
validated into the owner type first, because ``output_text`` is a derived property
rather than a serialized field, so it never exists on a dict; a dict the owner type
rejects is unjudgeable and skipped."""
from litellm.types.llms.openai import ResponsesAPIResponse
try:
@ -199,11 +262,16 @@ def _responses_final_text(response_obj: object) -> str:
if not isinstance(output, Sequence):
return ""
items: Final = tuple(item.model_dump() if isinstance(item, BaseModel) else item for item in output)
if any(
not isinstance(item, Mapping) or item.get("type") in ("function_call", "custom_tool_call") for item in items
):
if any(not isinstance(item, Mapping) for item in items):
return ""
return str(getattr(response, "output_text", "") or "")
calls: Final = tuple(
item for item in items if isinstance(item, Mapping) and item.get("type") in _RESPONSES_TOOL_CALL_TYPES
)
prose: Final = str(getattr(response, "output_text", "") or "")
if not calls:
return prose
serialized: Final = ", ".join(f"[tool call] {_tool_call_invocation(call)}" for call in calls)
return f"{prose} {serialized}".strip() if prose else serialized
class _SurfaceOps:
@ -273,8 +341,8 @@ def _judgeable_sample(
response_obj: object,
) -> tuple[tuple[Mapping[str, object], ...], Mapping[str, object], str] | None:
"""The normalized chat conversation, the forwardable generation params, and the
judgeable final text; None when this request's shapes cannot be sampled (tool-final
turn, empty text, or a shape the owner transformations reject)."""
judgeable final text; None when this request's shapes cannot be sampled (no text and no
tool call to serialize, or a shape the owner transformations reject)."""
try:
request: Final = ops.chat_request(kwargs, model_parameters)
items: Final = _MESSAGE_ITEMS_ADAPTER.validate_python(request.get("messages"))
@ -307,6 +375,11 @@ PAIRWISE_JUDGE_SYSTEM_PROMPT: Final = """You are an impartial quality judge comp
The responses are labeled A and B in random order. You do not know which system produced which.
A response may be prose, or a tool call shown as `[tool call] name(arguments)` if the
assistant chose to act instead of answering directly. A tool call is not a defect: judge
whether calling that tool was the right response to the conversation, the same as you
would judge prose.
Criteria: correctness, completeness, clarity, conciseness.
Return ONLY valid JSON in this exact format, no other text:
@ -376,14 +449,37 @@ def _unmask_preference(raw_preference: str, real_is_a: bool) -> str:
return "tie"
def _judge_user_prompt(conversation: str, response_a: str, response_b: str) -> str:
_MAX_JUDGE_TOOL_DEFS_CHARS: Final = 2_000
def _tool_definitions_text(tools: object) -> str:
"""The tools available to both arms, name and description only: enough for the judge
to tell whether the chosen tool, and not some other one, was the right call, without
forwarding parameter schemas it does not need to score that."""
if not isinstance(tools, Sequence) or isinstance(tools, str):
return ""
entries: Final = tuple(
_field_reader(t)("function") or _field_reader(t)("custom") or t for t in tools if not isinstance(t, str)
)
lines: Final = tuple(
f"- {_field_reader(e)('name') or 'unnamed'}: {_field_reader(e)('description') or 'no description'}"
for e in entries
)
if not lines:
return ""
return ("Tools available to both responses:\n" + "\n".join(lines))[:_MAX_JUDGE_TOOL_DEFS_CHARS]
def _judge_user_prompt(conversation: str, response_a: str, response_b: str, tool_definitions: str = "") -> str:
"""The judge prompt under one total character budget: each response is capped, and
the conversation tail gets whatever budget the responses left over."""
the conversation tail gets whatever budget the responses and tool definitions left
over."""
a: Final = response_a[:_MAX_JUDGE_RESPONSE_CHARS]
b: Final = response_b[:_MAX_JUDGE_RESPONSE_CHARS]
conversation_budget: Final = _MAX_JUDGE_PROMPT_CHARS - len(a) - len(b)
prefix: Final = f"{tool_definitions}\n\n" if tool_definitions else ""
conversation_budget: Final = _MAX_JUDGE_PROMPT_CHARS - len(a) - len(b) - len(prefix)
return (
f"Conversation:\n{conversation[-conversation_budget:]}\n\n"
f"{prefix}Conversation:\n{conversation[-conversation_budget:]}\n\n"
f"Response A:\n{a}\n\n"
f"Response B:\n{b}\n\n"
"Which response is better?"
@ -942,6 +1038,7 @@ class ShadowEvalLogger(CustomLogger):
messages=messages,
real_text=real_text,
shadow_text=shadow.text,
tools=shadow_params.get("tools"),
parent_metadata=parent_metadata,
)
if isinstance(verdict, _CallFailure):
@ -1080,15 +1177,18 @@ class ShadowEvalLogger(CustomLogger):
classifier_cost=_decision_classifier_cost(shadow_metadata),
)
text: Final = _chat_final_text(response)
routed_model: Final = str(
getattr(response, "model", None) or _routing_decision(shadow_metadata).get("routed_model") or ""
)
if not text:
return _CallFailure(
"shadow router returned an empty response",
_shadow_empty_reply_error(response, routed_model),
cost=_call_cost(response),
classifier_cost=_decision_classifier_cost(shadow_metadata),
)
return _ShadowResponse(
text=text,
model=str(getattr(response, "model", None) or _routing_decision(shadow_metadata).get("routed_model") or ""),
model=routed_model,
tier=_routed_tier(shadow_metadata),
cost=_call_cost(response),
classifier_cost=_decision_classifier_cost(shadow_metadata),
@ -1100,9 +1200,12 @@ class ShadowEvalLogger(CustomLogger):
messages: Sequence[Mapping[str, object]],
real_text: str,
shadow_text: str,
tools: object,
parent_metadata: Mapping[str, object],
) -> "_JudgeVerdict | _CallFailure":
"""Blind pairwise judge with A/B labels randomized to cancel position bias."""
"""Blind pairwise judge with A/B labels randomized to cancel position bias. Both
arms were offered the same tools, so the judge is shown their definitions too: a
tool call is only assessable against what else was available to call instead."""
real_is_a: Final = random.random() < 0.5
response_a: Final = real_text if real_is_a else shadow_text
response_b: Final = shadow_text if real_is_a else real_text
@ -1117,7 +1220,7 @@ class ShadowEvalLogger(CustomLogger):
{"role": "system", "content": PAIRWISE_JUDGE_SYSTEM_PROMPT}, # mutable-ok: SDK message
{
"role": "user",
"content": _judge_user_prompt(conversation, response_a, response_b),
"content": _judge_user_prompt(conversation, response_a, response_b, _tool_definitions_text(tools)),
}, # mutable-ok: SDK message
]
try:

38
litellm/litellm_core_utils/litellm_logging.py
View file

@ -10,7 +10,7 @@ import subprocess
import sys
import time
import traceback
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime as dt_object
from functools import lru_cache
from types import MappingProxyType, TracebackType
@ -576,6 +576,7 @@ class Logging(LiteLLMLoggingBaseClass):
# enqueue closure here instead of firing it immediately.
self._defer_async_logging: bool = False
self._enqueue_deferred_logging: Callable[[], None] | None = None
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
@ -1894,6 +1895,11 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
"""Stash what an interrupted stream already consumed so the failure log bills it instead of zero."""
self.model_call_details["combined_usage_object"] = usage
self.model_call_details["response_cost"] = response_cost
async def dispatch_failure_handlers(
self,
exception: Exception,
@ -5878,14 +5884,28 @@ def _get_status_fields(
#########################################################
# Map - guardrail_information.guardrail_status to guardrail_status
#########################################################
guardrail_status: GuardrailStatus = "not_run"
if guardrail_information and isinstance(guardrail_information, list):
for information in guardrail_information:
if isinstance(information, dict):
raw_status = information.get("guardrail_status", "not_run")
if raw_status != "not_run":
guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
break
# Severity order, least severe first. The status aggregates across ALL
# guardrail entries rather than taking the first non-"not_run" one: a
# pre_call guardrail that passed (e.g. a mask) records its entry before a
# later guardrail's block, and first-wins would report a blocked request
# as "success".
GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = (
"not_run",
"success",
"guardrail_failed_to_respond",
"guardrail_intervened",
)
entries: Final[Sequence[object]] = guardrail_information if isinstance(guardrail_information, list) else ()
raw_statuses: Final[Iterator[object]] = (
entry.get("guardrail_status", "not_run") for entry in entries if isinstance(entry, dict)
)
# A guardrail is free to write any value here, and an unhashable one would
# raise TypeError on the mapping lookup and drop the whole payload.
guardrail_status: Final[GuardrailStatus] = max(
(GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") for raw_status in raw_statuses if isinstance(raw_status, str)),
key=GUARDRAIL_STATUS_SEVERITY.index,
default="not_run",
)
return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status)

56
litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py
View file

@ -1,8 +1,8 @@
import math
from collections.abc import Mapping
from typing import Final
from typing import Annotated, Final
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
@ -30,6 +30,31 @@ class GuardrailCostEntry(BaseModel):
_GUARDRAIL_COST_ENTRY_ADAPTER: Final[TypeAdapter[GuardrailCostEntry]] = TypeAdapter(GuardrailCostEntry)
class GuardrailCostByUnitEntry(BaseModel):
"""The rollup-side view of a ``guardrail_information`` entry, validated apart from
``GuardrailCostEntry`` so a forged per-counter map can never zero the spend path."""
model_config = ConfigDict(extra="ignore", frozen=True)
guardrail_cost_by_unit: Mapping[str, Annotated[float, Field(ge=0, allow_inf_nan=False)] | None] | None = None
guardrail_cost_in_spend: bool | None = True
_GUARDRAIL_COST_BY_UNIT_ADAPTER: Final[TypeAdapter[GuardrailCostByUnitEntry]] = TypeAdapter(GuardrailCostByUnitEntry)
def billed_guardrail_cost_by_unit(raw: object) -> Mapping[str, float | None] | None:
"""Per-counter USD the daily rollup may record for one raw ``guardrail_information``
entry; None when the entry is unpriced, report-only, or malformed, and None per
counter the hook had no price for."""
try:
entry: Final = _GUARDRAIL_COST_BY_UNIT_ADAPTER.validate_python(raw)
except ValidationError as e:
verbose_logger.warning("Ignoring malformed guardrail_information entry for guardrail cost rollup: %s", e)
return None
return None if entry.guardrail_cost_in_spend is False else entry.guardrail_cost_by_unit
def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing | None:
regional_key: Final = f"bedrock/{aws_region_name}/guardrails" if aws_region_name else None
for key in (regional_key, BEDROCK_GUARDRAIL_PRICING_KEY):
@ -42,11 +67,32 @@ def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing
return None
def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float:
def _priced_units(units: int, price_per_unit: float | None) -> float | None:
return None if price_per_unit is None else units * price_per_unit
def bedrock_guardrail_cost_by_unit(
usage_units: Mapping[str, int], aws_region_name: str | None
) -> Mapping[str, float | None] | None:
"""USD per counter, keyed like ``usage_units``; None when no pricing entry exists,
and None for a counter the entry has no price for, since only an explicit 0.0 means free."""
pricing: Final = _bedrock_guardrail_pricing(aws_region_name)
if pricing is None:
return 0.0
return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items())
return None
return { # mutable-ok: stamped into guardrail_information, which safe_dumps only serializes as a plain dict
counter: _priced_units(units, pricing.guardrail_cost_per_unit.get(counter))
for counter, units in usage_units.items()
}
def guardrail_cost_total(cost_by_unit: Mapping[str, float | None] | None) -> float:
"""The scalar the spend path bills: unknown-priced counters count as 0 here, the
rollup keeps them unknown."""
return sum(cost for cost in cost_by_unit.values() if cost is not None) if cost_by_unit is not None else 0.0
def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float:
return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name))
AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records"

39
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -229,6 +229,45 @@ def _content_parts_contain_image(parts: Sequence[object]) -> bool:
return False
def anthropic_image_source_to_openai_url(image_source: Mapping[str, object]) -> str | None:
"""Data or remote URL for an Anthropic ``source`` block, in the form chat completions expects."""
source_type: Final = image_source.get("type")
if source_type == "base64":
media_type: Final = image_source.get("media_type") or "image/jpeg"
image_data: Final = image_source.get("data") or ""
return f"data:{media_type};base64,{image_data}" if image_data else None
if source_type == "url":
url: Final = image_source.get("url")
return url if isinstance(url, str) else ""
return None
def _image_part_url(part: Mapping[str, object]) -> str | None:
"""The image URL carried by one content part, whichever of the three dialects wrote it."""
part_type: Final = part.get("type")
if part_type == "image_url":
image_url: Final = part.get("image_url")
if isinstance(image_url, str):
return image_url
return image_url.get("url") if isinstance(image_url, Mapping) else None
if part_type == "input_image":
responses_url: Final = part.get("image_url")
return responses_url if isinstance(responses_url, str) else None
if part_type == "image":
source: Final = part.get("source")
return anthropic_image_source_to_openai_url(source) if isinstance(source, Mapping) else None
return None
def as_openai_image_part(part: Mapping[str, object]) -> ChatCompletionImageObject | None:
"""One image content part rewritten into chat-completions dialect, or None when it is not one.
Rebuilt rather than forwarded so no caller-controlled key beyond the URL rides along.
"""
url: Final = _image_part_url(part)
return {"type": "image_url", "image_url": {"url": url}} if url else None
def request_contains_image_content(messages: Sequence[Mapping[str, object]]) -> bool:
"""Whether any message carries an image content part, across the dialects that reach
pre-routing hooks untranslated: chat-completions ``image_url``, Responses ``input_image``,

48
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -1,4 +1,5 @@
from collections.abc import Mapping, Sequence
from collections.abc import Set as AbstractSet
from typing import Any, Final
from pydantic import BaseModel
@ -6,38 +7,45 @@ from pydantic import BaseModel
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
from litellm.litellm_core_utils.secret_redaction import REDACTED
_DEFAULT_SENSITIVE_PATTERNS: Final = frozenset(
(
"password",
"secret",
"key",
"token",
"auth",
"authorization",
"credential",
# Plural form: Vertex uses ``vertex_credentials``; segment-exact
# matching otherwise misses it because "credential" != "credentials".
"credentials",
"access",
"private",
"certificate",
"fingerprint",
"tenancy",
)
)
class SensitiveDataMasker:
def __init__(
self,
sensitive_patterns: set[str] | None = None,
non_sensitive_overrides: set[str] | None = None,
sensitive_patterns: AbstractSet[str] | None = None,
non_sensitive_overrides: AbstractSet[str] | None = None,
visible_prefix: int = 4,
visible_suffix: int = 4,
mask_char: str = "*",
mask_short_values: bool = True,
extra_sensitive_patterns: AbstractSet[str] | None = None,
):
self.sensitive_patterns = sensitive_patterns or {
"password",
"secret",
"key",
"token",
"auth",
"authorization",
"credential",
# Plural form: Vertex uses ``vertex_credentials``; segment-exact
# matching otherwise misses it because "credential" != "credentials".
"credentials",
"access",
"private",
"certificate",
"fingerprint",
"tenancy",
}
self.sensitive_patterns = (sensitive_patterns or _DEFAULT_SENSITIVE_PATTERNS) | (
extra_sensitive_patterns or frozenset()
)
# If any key segment matches one of these, the key is not considered sensitive
# even if it also matches a sensitive pattern. For example, "input_cost_per_token"
# contains "token" but "cost" overrides that — it's a pricing field, not a secret.
self.non_sensitive_overrides = non_sensitive_overrides or {"cost"}
self.non_sensitive_overrides = non_sensitive_overrides or frozenset(("cost",))
self.visible_prefix = visible_prefix
self.visible_suffix = visible_suffix

22
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
View file

@ -99,6 +99,7 @@ def create_tool_name_mapping(
from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingChoice
from litellm.litellm_core_utils.prompt_templates.common_utils import (
anthropic_image_source_to_openai_url,
parse_tool_call_arguments,
reasoning_content_from_thinking_blocks,
with_prompt_cache_breakpoint,
@ -524,18 +525,20 @@ class LiteLLMAnthropicMessagesAdapter:
self._add_cache_control_if_applicable(content, tool_call, model)
tool_calls.append(tool_call)
elif content.get("type") == "thinking":
# Anthropic's schema has no cache_control on thinking or
# redacted_thinking blocks, and anthropic_messages_pt replays
# these verbatim at content[0], so carrying one here (or
# inventing an empty one) is a guaranteed 400 on the way back.
thinking_block = ChatCompletionThinkingBlock(
type="thinking",
thinking=content.get("thinking") or "",
signature=content.get("signature") or "",
cache_control=content.get("cache_control", {}),
)
thinking_blocks.append(thinking_block)
elif content.get("type") == "redacted_thinking":
redacted_thinking_block = ChatCompletionRedactedThinkingBlock(
type="redacted_thinking",
data=content.get("data") or "",
cache_control=content.get("cache_control", {}),
)
thinking_blocks.append(redacted_thinking_block)
@ -1223,20 +1226,7 @@ class LiteLLMAnthropicMessagesAdapter:
"""
if not isinstance(image_source, dict):
return None
source_type: Final = image_source.get("type")
if source_type == "base64":
# Base64 image format
media_type: Final = image_source.get("media_type", "image/jpeg")
image_data: Final = image_source.get("data", "")
if image_data:
return f"data:{media_type};base64,{image_data}"
elif source_type == "url":
# URL-referenced image format
return image_source.get("url", "")
return None
return anthropic_image_source_to_openai_url(image_source)
def _tool_result_content(self, raw_content: object) -> ToolResultContent:
if isinstance(raw_content, str):

54
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import AsyncIterator, Mapping
from collections.abc import AsyncIterator, Mapping, Sequence
from datetime import datetime
from typing import Any, Final, Protocol, runtime_checkable
@ -177,12 +177,6 @@ def _try_claim_detached_drain_slot() -> bool:
def _exception_left_unconsumed(queue: "asyncio.Queue[bytes | None | BaseException]", exc: BaseException) -> bool:
"""After client detach the relay never reads the queue again, so drain it here.
The forwarded exception still sitting in the queue means the relay tore
down before re-raising it, so the proxy's failure handling never ran and
the caller must salvage spend itself.
"""
remaining: Final = tuple(queue.get_nowait() for _ in range(queue.qsize()))
return any(item is exc for item in remaining)
@ -671,27 +665,41 @@ class BaseAnthropicMessagesStreamingIterator:
self,
queue: "asyncio.Queue[bytes | None | BaseException]",
client_detached: "asyncio.Event",
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks
exc: BaseException,
collected_chunks: Sequence[bytes],
exc: Exception,
) -> None:
"""Forward a provider error to a still-connected client, else salvage partial spend.
"""Log the request as failed with its partial usage, then make sure the proxy's failure hook runs once.
Handing the original exception to the client-facing generator lets it
re-raise so the proxy's failure handling keeps the provider status and
owns logging (no success-bill). If the client already went away, or
disconnects before ever consuming the queued exception, no failure hook
runs, so bill the partial instead of dropping the request.
A still-connected client gets the original exception through the queue,
the relay re-raises it, and the proxy's own failure handling records the
failed spend. When the client already left, or leaves before consuming
the queued exception, that handling never runs, so the detached-failure
hook the proxy armed on the logging object fires here instead.
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=self.litellm_logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
request_body=self.request_body,
raw_bytes=collected_chunks,
exception=exc,
)
if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc):
await client_detached.wait()
if not _exception_left_unconsumed(queue, exc):
return
verbose_proxy_logger.warning(
"async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)",
len(collected_chunks),
type(exc).__name__,
exc,
)
await self._bill_collected_chunks(collected_chunks, stream_teardown=True)
await self._fire_detached_failure_hook(exc)
async def _fire_detached_failure_hook(self, exc: Exception) -> None:
from litellm._logging import verbose_proxy_logger
on_detached_failure: Final = getattr(self.litellm_logging_obj, "_on_detached_stream_failure", None)
if on_detached_failure is None:
return
try:
await on_detached_failure(exc)
except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not crash the detached pump
verbose_proxy_logger.warning(
"async_sse_wrapper detached failure hook raised: %s(%s)", type(hook_failure).__name__, hook_failure
)

11
litellm/llms/anthropic/experimental_pass_through/messages/utils.py
View file

@ -6,6 +6,7 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.llms.openai import ChatCompletionSystemMessage
if TYPE_CHECKING:
from litellm.exceptions import ContentPolicyViolationError
@ -36,6 +37,16 @@ def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "
)
def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMessage | None:
"""
Return the Anthropic Messages top-level ``system`` (a string or a list of text
blocks) as an OpenAI-style system message, or None when the request has none.
"""
if not isinstance(system, (str, list)) or not system:
return None
return ChatCompletionSystemMessage(role="system", content=system)
@lru_cache(maxsize=1)
def _anthropic_messages_optional_param_keys() -> frozenset[str]:
"""

9
litellm/llms/azure/responses/transformation.py
View file

@ -6,6 +6,7 @@ from openai.types.responses import ResponseReasoningItem
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.types.llms.openai import *
@ -29,6 +30,14 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE
@staticmethod
def _supports_reasoning_effort_none(model: str) -> bool:
return AzureOpenAIGPT5Config._supports_reasoning_effort_level(model, "none")
@staticmethod
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
return AzureOpenAIGPT5Config.effort_resolves_to_none(model, effort)
def get_supported_openai_params(self, model: str) -> list:
"""
Azure Responses API does not support context_management (compaction).

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

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

27
litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
View file

@ -89,11 +89,24 @@ def _extract_converse_texts(
top-level ``text`` blocks this scans the arbitrary-JSON fields a caller can
hide prompt content in -- ``toolUse.input`` and
``toolResult.content[].json`` (alongside ``toolResult.content[].text``) --
as well as the request-level fields still forwarded to Bedrock that a caller
can route blocked content through: ``toolConfig.tools`` (tool names,
descriptions and input schemas) and ``additionalModelRequestFields``. Tool
message blocks are skipped when tool messages are excluded, but tool
definitions are always scanned to match the chat-completions guardrail path.
as well as ``additionalModelRequestFields``, a free-form model-parameter bag
with no schema that a caller can route blocked content through.
``toolConfig.tools`` is deliberately NOT scanned. Tool definitions are
app-authored config, so their names, descriptions and JSON-schema strings
("object", property names, titles, type names, enum values) would each reach
the guardrail as a separate INPUT item, producing false positives and
inflating guardrail usage for a request whose only prompt is one user
message. No other guardrail translation handler puts tool definitions in
``texts``; the chat and messages handlers carry them in the structured
``tools`` input instead, which this handler does not populate because a
Bedrock ``toolSpec`` is not the OpenAI tool shape those consumers expect.
``additionalModelRequestFields`` is treated differently on purpose. Bedrock
gives ``toolConfig.tools`` a fixed schema whose contents are tool metadata by
contract, while ``additionalModelRequestFields`` is free-form and defined by
the target model, so what it carries cannot be classified without knowing
that model. Scanning it stays the fail-closed default.
"""
holders: Final[list[_StringHolder]] = []
@ -121,10 +134,6 @@ def _extract_converse_texts(
_collect_block_text(inner, holders)
_collect_strings(inner.get("json"), holders)
tool_config: Final = body.get("toolConfig")
if isinstance(tool_config, dict):
_collect_strings(tool_config.get("tools"), holders)
_collect_strings(body.get("additionalModelRequestFields"), holders)
texts: Final = [container[key] for container, key in holders]

2
litellm/llms/bedrock_mantle/common_utils.py
View file

@ -29,7 +29,7 @@ from litellm.secret_managers.main import get_secret_str
BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
# Standard Mantle host: https://bedrock-mantle.<region>.api.aws (group 1 = region).
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE)
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws(?=/|$)", re.IGNORECASE)
def resolve_mantle_bearer_token(api_key: str | None) -> str | None:

8
litellm/llms/fireworks_ai/chat/transformation.py
View file

@ -272,11 +272,15 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
)
# Only add tool_choice for models that explicitly support it
if supports_tool_choice(model=model, custom_llm_provider="fireworks_ai"):
if self._get_model_cost_capability_exact(
model=model, capability="supports_tool_choice"
) or supports_tool_choice(model=model, custom_llm_provider="fireworks_ai"):
supported_params.append("tool_choice")
# Only add reasoning params for models that support it
if supports_reasoning(model=model, custom_llm_provider="fireworks_ai"):
if self._get_model_cost_capability_exact(model=model, capability="supports_reasoning") or supports_reasoning(
model=model, custom_llm_provider="fireworks_ai"
):
supported_params.append("reasoning_effort")
supported_params.append("reasoning_history")
supported_params.append("thinking")

57
litellm/llms/fireworks_ai/cost_calculator.py
View file

@ -2,6 +2,8 @@
For calculating cost of fireworks ai serverless inference models.
"""
import math
from datetime import datetime
from typing import Final
from litellm.constants import (
@ -10,9 +12,12 @@ from litellm.constants import (
FIREWORKS_AI_56_B_MOE,
FIREWORKS_AI_176_B_MOE,
)
from litellm.types.utils import Usage
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import get_model_info
NO_CACHE_READ_RATE: Final = float("nan")
# Extract the number of billion parameters from the model name
# only used for together_computer LLMs
@ -54,44 +59,50 @@ def get_base_model_for_pricing(model_name: str) -> str:
return "fireworks-ai-default"
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
def _resolve_model_info(model: str) -> ModelInfo:
try:
return get_model_info(model=model, custom_llm_provider="fireworks_ai")
except Exception:
base_model: Final = get_base_model_for_pricing(model_name=model)
return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
Calculates the cost per token for a given model, prompt tokens, and completion tokens,
swapping in the model's off_peak_pricing rates while one of its windows is open.
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing anthropic caching information
- current_time: the moment the request is billed at; defaults to now, UTC
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
## check if model mapped, else use default pricing
try:
model_info = get_model_info(model=model, custom_llm_provider="fireworks_ai")
except Exception:
base_model: Final = get_base_model_for_pricing(model_name=model)
model_info: Final = _resolve_model_info(model)
standard_cache_read_rate: Final = model_info.get("cache_read_input_token_cost")
rates: Final = apply_off_peak_pricing(
model_info,
current_time,
TokenRates(
input_rate=model_info["input_cost_per_token"] or 0.0,
output_rate=model_info["output_cost_per_token"] or 0.0,
cache_read_rate=standard_cache_read_rate if standard_cache_read_rate is not None else NO_CACHE_READ_RATE,
cache_creation_rate=0.0,
reasoning_rate=None,
),
)
cache_read_rate: Final[float] = rates.input_rate if math.isnan(rates.cache_read_rate) else rates.cache_read_rate
## GET MODEL INFO
model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
## CALCULATE INPUT COST
prompt_tokens_details: Final = usage.prompt_tokens_details
cached_tokens: Final[int] = (
prompt_tokens_details.cached_tokens
if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
else 0
)
input_cost_per_token: Final[float] = model_info["input_cost_per_token"] or 0.0
cache_read_input_token_cost: Final = model_info.get("cache_read_input_token_cost")
cache_read_cost_per_token: Final[float] = (
cache_read_input_token_cost if cache_read_input_token_cost is not None else input_cost_per_token
)
non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)
prompt_cost: float = non_cached_prompt_tokens * input_cost_per_token + cached_tokens * cache_read_cost_per_token
## CALCULATE OUTPUT COST
output_cost_per_token: Final[float] = model_info["output_cost_per_token"] or 0.0
completion_cost: Final[float] = usage.completion_tokens * output_cost_per_token
prompt_cost: Final[float] = non_cached_prompt_tokens * rates.input_rate + cached_tokens * cache_read_rate
completion_cost: Final[float] = usage.completion_tokens * rates.output_rate
return prompt_cost, completion_cost

0
tests/rust-python-harness/strategies/e2e_parity/gateway/__init__.py → litellm/llms/mongodb/__init__.py
View file

303
litellm/llms/mongodb/common_utils.py Normal file
View file

@ -0,0 +1,303 @@
"""Shared helpers for the MongoDB integrations. pymongo lives in the optional ``mongodb`` extra,
so every import of it is deferred to call time."""
import asyncio
import threading
import weakref
from asyncio import AbstractEventLoop
from collections import OrderedDict
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar
from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout
if TYPE_CHECKING:
from pymongo import AsyncMongoClient, MongoClient
PYMONGO_INSTALL_HINT: Final = (
"The MongoDB vector store requires the 'pymongo' package. "
"Run 'pip install litellm[mongodb]' (or 'pip install pymongo') to install it."
)
MONGODB_PROVIDER: Final = "mongodb"
def config_error(message: str) -> BadRequestError:
"""400 rather than the 500 a bare ValueError becomes once litellm.exception_type wraps it."""
return BadRequestError(message=message, model=None, llm_provider=MONGODB_PROVIDER)
def timeout_error(message: str) -> Timeout:
return Timeout(message=message, model=None, llm_provider=MONGODB_PROVIDER)
def unavailable_error(message: str) -> ServiceUnavailableError:
"""litellm only retries 408, 409, 429 and 5xx, so a 400 here would make a failover permanent."""
return ServiceUnavailableError(message=message, model=None, llm_provider=MONGODB_PROVIDER)
DEFAULT_CONNECT_TIMEOUT_MS: Final = 10_000
DEFAULT_SOCKET_TIMEOUT_MS: Final = 30_000
DEFAULT_SERVER_SELECTION_TIMEOUT_MS: Final = 10_000
_MAX_CACHED_CLIENTS: Final = 32
_APP_NAME: Final = "litellm"
@dataclass(frozen=True, slots=True)
class MongoClientKey:
connection_string: str
connect_timeout_ms: int
socket_timeout_ms: int
server_selection_timeout_ms: int
SyncClientFactory: TypeAlias = Callable[..., "MongoClient"]
AsyncClientFactory: TypeAlias = Callable[..., "AsyncMongoClient"]
_K = TypeVar("_K")
_V = TypeVar("_V")
_AsyncClientCacheKey: TypeAlias = tuple[MongoClientKey, int]
# CPython recycles id() aggressively, so the id alone would hand a new loop a closed loop's client
_AsyncClientEntry: TypeAlias = tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"]
_SyncClientCache: TypeAlias = "OrderedDict[MongoClientKey, MongoClient]"
_AsyncClientCache: TypeAlias = "OrderedDict[_AsyncClientCacheKey, _AsyncClientEntry]"
_sync_clients: Final[_SyncClientCache] = OrderedDict() # mutable-ok: process-level client cache
_async_clients: Final[_AsyncClientCache] = OrderedDict() # mutable-ok: same cache, per loop
# async searches reach the sync client through executor threads, so both caches are shared state
_cache_lock: Final = threading.Lock()
def _store_bounded(cache: "OrderedDict[_K, _V]", cache_key: "_K", value: "_V") -> None:
"""Eviction only drops this cache's reference; an in-flight search keeps its client alive."""
with _cache_lock:
cache[cache_key] = value # mutable-ok: an LRU cache is mutable state by definition
cache.move_to_end(cache_key)
while len(cache) > _MAX_CACHED_CLIENTS:
cache.popitem(last=False)
def _mark_used(cache: "OrderedDict[_K, _V]", cache_key: "_K") -> None:
with _cache_lock:
if cache_key in cache:
cache.move_to_end(cache_key)
def import_sync_mongo_client() -> "type[MongoClient]":
try:
from pymongo import MongoClient as SyncMongoClient
except ImportError as e:
raise config_error(PYMONGO_INSTALL_HINT) from e
return SyncMongoClient
def import_async_mongo_client() -> "type[AsyncMongoClient]":
try:
from pymongo import AsyncMongoClient as AsyncMongoClientClass
except ImportError as e:
raise config_error(PYMONGO_INSTALL_HINT) from e
return AsyncMongoClientClass
def _client_kwargs(key: MongoClientKey) -> Mapping[str, object]:
return MappingProxyType(
{
"connectTimeoutMS": key.connect_timeout_ms,
"socketTimeoutMS": key.socket_timeout_ms,
"serverSelectionTimeoutMS": key.server_selection_timeout_ms,
"appname": _APP_NAME,
}
)
def get_sync_client(key: MongoClientKey, client_class: SyncClientFactory | None = None) -> "MongoClient":
cached: Final = _sync_clients.get(key)
if cached is not None:
_mark_used(_sync_clients, key)
return cached
build: Final = client_class if client_class is not None else import_sync_mongo_client()
client: Final = build(key.connection_string, **_client_kwargs(key))
_store_bounded(_sync_clients, key, client)
return client
def _purge_dead_loops() -> None:
"""A cached client holds its loop alive, so a closed loop's entry would pin that client and its
sockets for the life of the process."""
with _cache_lock:
for stale in tuple(
cache_key
for cache_key, (loop_ref, _) in _async_clients.items()
if (cached_loop := loop_ref()) is None or cached_loop.is_closed()
):
del _async_clients[stale]
def get_async_client(key: MongoClientKey, client_class: AsyncClientFactory | None = None) -> "AsyncMongoClient":
"""Async clients bind to the loop that created them, so the cache is keyed per loop."""
loop: Final = asyncio.get_running_loop()
loop_key: Final = (key, id(loop))
cached: Final = _async_clients.get(loop_key)
if cached is not None and cached[0]() is loop:
_mark_used(_async_clients, loop_key)
return cached[1]
_purge_dead_loops()
build: Final = client_class if client_class is not None else import_async_mongo_client()
client: Final = build(key.connection_string, **_client_kwargs(key))
_store_bounded(_async_clients, loop_key, (weakref.ref(loop), client))
return client
def reset_client_cache() -> None:
with _cache_lock:
_sync_clients.clear()
_async_clients.clear()
_AUTHENTICATION_FAILED_CODE: Final = 18
_UNAUTHORIZED_CODE: Final = 13
# Atlas reports a rejected user as code 8000 "AtlasError" where a self-managed mongod reports 18
_AUTHENTICATION_MESSAGE_MARKERS: Final = ("bad auth", "authentication failed", "not authorized")
_RESOLUTION_TIMEOUT_MARKERS: Final = ("resolution lifetime expired", "dns operation timed out")
_UNKNOWN_HOSTNAME_MARKERS: Final = ("dns query name does not exist", "name or service not known")
_CREDENTIAL_ESCAPING_MARKERS: Final = ("must be escaped according to rfc 3986", "bad database name")
def _index_hint(index_name: str, database: str, collection: str) -> str:
return (
f"No queryable MongoDB Vector Search index named '{index_name}' was found on "
f"'{database}.{collection}'. Confirm the index exists on that exact collection, that its "
"status is READY rather than still building, and that the vector store id matches the index name."
)
def missing_index_error(index_name: str, database: str, collection: str) -> BadRequestError:
"""$vectorSearch against a missing index, database or collection returns zero documents rather
than failing, so an empty result set is checked against the catalogue and reported as this."""
return config_error(
f"{_index_hint(index_name, database, collection)} A vector search against a database, "
"collection or index that does not exist returns no results rather than an error, so this "
"was reported as an empty result set by MongoDB."
)
def index_not_ready_error(index_name: str, database: str, collection: str, status: str) -> BadRequestError:
return config_error(
f"The MongoDB Vector Search index '{index_name}' on '{database}.{collection}' is not queryable "
f"yet; its status is {status}. Searches against it return no results until the build finishes."
)
def translate_mongo_error(error: Exception, index_name: str, database: str, collection: str) -> Exception:
"""Returns the exception to raise, so callers keep the driver error as ``__cause__``."""
try:
from pymongo.errors import (
ConfigurationError,
ConnectionFailure,
ExecutionTimeout,
InvalidOperation,
NetworkTimeout,
OperationFailure,
ServerSelectionTimeoutError,
)
except ImportError:
return error
if isinstance(error, ServerSelectionTimeoutError):
return timeout_error(
"Could not reach the MongoDB deployment before the timeout. On Atlas this is usually the "
"project's IP access list not containing this host, or a paused cluster. On a self-managed "
"deployment it is usually the host or port in the URI, or a firewall between this process "
f"and mongod. Either way it can also be an unresolvable hostname. Driver detail: {error}"
)
# ExecutionTimeout subclasses OperationFailure, so it has to be matched before it
if isinstance(error, (NetworkTimeout, ExecutionTimeout)):
return timeout_error(
f"The MongoDB vector search against '{database}.{collection}' timed out before returning. "
f"Driver detail: {error}"
)
# ServerSelectionTimeoutError and NetworkTimeout also subclass ConnectionFailure, so this only
# sees what those branches left
if isinstance(error, ConnectionFailure):
return unavailable_error(
f"The connection to '{database}.{collection}' was dropped or refused. That is usually a "
"replica set failover or a restarted node, so the search is worth retrying. If it keeps "
"happening: on Atlas the usual cause is a connection string with no username and password, "
"or a TLS failure, so confirm the URI is the one Atlas shows under Connect, Drivers; on a "
"self-managed deployment, check that mongod is listening on the host and port in the URI. "
f"Driver detail: {error}"
)
if isinstance(error, OperationFailure):
code: Final = error.code
detail: Final = str(error).lower()
if code in (_AUTHENTICATION_FAILED_CODE, _UNAUTHORIZED_CODE) or any(
marker in detail for marker in _AUTHENTICATION_MESSAGE_MARKERS
):
return config_error(
"MongoDB rejected the credentials in mongodb_connection_string, or the database user "
f"lacks read access to '{database}.{collection}'. Driver detail: {error.details}"
)
if "dimension" in detail:
return config_error(
"The query embedding does not match the vector dimensions the index was built for. "
"litellm_embedding_model must be the same model that produced the stored vectors. "
f"Driver detail: {error}"
)
if "is not indexed as vector" in detail:
return config_error(
"mongodb_embedding_field names a field the MongoDB Vector Search index does not cover. "
f"It must match the 'path' the index '{index_name}' was created on. Driver detail: {error}"
)
if "index" in detail and ("not found" in detail or "does not exist" in detail or "unknown" in detail):
return config_error(f"{_index_hint(index_name, database, collection)} Driver detail: {error}")
return config_error(
f"MongoDB rejected the vector search against '{database}.{collection}' using index "
f"'{index_name}'. Driver detail: {error}"
)
if isinstance(error, ConfigurationError):
configuration_detail: Final = str(error).lower()
if any(marker in configuration_detail for marker in _RESOLUTION_TIMEOUT_MARKERS):
return timeout_error(
"The DNS lookup for the cluster in mongodb_connection_string did not finish in time. "
"A mongodb+srv:// URI needs an SRV lookup before any connection is attempted, so this "
f"is DNS or the configured timeout, not MongoDB. Driver detail: {error}"
)
if any(marker in configuration_detail for marker in _UNKNOWN_HOSTNAME_MARKERS):
return config_error(
"The hostname in mongodb_connection_string does not exist in DNS. On Atlas, check the "
"cluster name against the URI shown under Connect, Drivers. On a self-managed deployment, "
f"check that the hostname resolves from this process. Driver detail: {error}"
)
if any(marker in configuration_detail for marker in _CREDENTIAL_ESCAPING_MARKERS):
return config_error(
"mongodb_connection_string could not be parsed. A username or password containing "
"'@', '/', ':' or '%' has to be percent-encoded per RFC 3986, so 'p@ss/word' becomes "
"'p%40ss%2Fword'. If the credentials are already encoded, check the database name in "
f"the URI path instead. Driver detail: {error}"
)
return config_error(
f"mongodb_connection_string is not a usable MongoDB connection string. Driver detail: {error}"
)
if isinstance(error, InvalidOperation):
return config_error(f"The MongoDB client was already closed or is unusable. Driver detail: {error}")
# An unreadable tlsCAFile or tlsCertificateKeyFile raises OSError, not a PyMongoError
if isinstance(error, OSError) and error.filename:
return config_error(
f"'{error.filename}', named by a TLS option in mongodb_connection_string, could not be read. "
"Check that tlsCAFile and tlsCertificateKeyFile point at files this process can open; inside "
f"a container that is the path in the container, not on the host. Driver detail: {error}"
)
# pymongo raises a plain ValueError, not a PyMongoError, for an unusable port
if isinstance(error, ValueError):
return config_error(
"The host and port in mongodb_connection_string could not be parsed. If the port is a "
"number between 0 and 65535, the cause is usually an unescaped ':' in the password, which "
f"has to be percent-encoded per RFC 3986 as '%3A'. Driver detail: {error}"
)
return error

0
tests/rust-python-harness/strategies/e2e_parity/sdk/__init__.py → litellm/llms/mongodb/vector_stores/__init__.py
View file

431
litellm/llms/mongodb/vector_stores/transformation.py Normal file
View file

@ -0,0 +1,431 @@
"""MongoDB Vector Search has no HTTP query API, so this is a direct provider that runs the
``$vectorSearch`` aggregation through pymongo. ``vector_store_id`` is the search index name."""
from collections.abc import Callable, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NoReturn
import httpx
from pydantic import BaseModel, ConfigDict
from litellm.llms.base_llm.vector_store.transformation import (
BaseDirectVectorStoreConfig,
LiteLLMVectorStoreEmbeddingExecutor,
VectorStoreEmbeddingExecutor,
)
from litellm.llms.mongodb.common_utils import (
DEFAULT_CONNECT_TIMEOUT_MS,
DEFAULT_SERVER_SELECTION_TIMEOUT_MS,
DEFAULT_SOCKET_TIMEOUT_MS,
MongoClientKey,
config_error,
get_async_client,
get_sync_client,
index_not_ready_error,
missing_index_error,
translate_mongo_error,
)
from litellm.types.utils import EmbeddingResponse
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
DEFAULT_EMBEDDING_FIELD_NAME: Final = "embedding"
DEFAULT_TEXT_FIELD_NAME: Final = "text"
SCORE_FIELD_NAME: Final = "score"
DEFAULT_MAX_NUM_RESULTS: Final = 10
MIN_MAX_NUM_RESULTS: Final = 1
MAX_MAX_NUM_RESULTS: Final = 50
NUM_CANDIDATES_MULTIPLIER: Final = 10
MIN_NUM_CANDIDATES: Final = 100
MAX_NUM_CANDIDATES: Final = 10_000
MAX_QUERY_CHARACTERS: Final = 32_000
_EMPTY_EMBEDDING_CONFIG: Final = MappingProxyType({})
_SEARCH_ONLY_MESSAGE: Final = (
"MongoDB vector store is search-only. Create the collection and its MongoDB Vector Search "
"index in MongoDB directly, then register it here by index name."
)
class _MongoDBSearchParams(BaseModel):
"""Typed view over the vector store's litellm_params; unrelated keys are ignored."""
model_config = ConfigDict(frozen=True, extra="ignore")
litellm_embedding_model: str | None = None
litellm_embedding_config: Mapping[str, object] | None = None
mongodb_connection_string: str | None = None
mongodb_database: str | None = None
mongodb_collection: str | None = None
mongodb_text_field: str | None = None
mongodb_embedding_field: str | None = None
mongodb_num_candidates: int | None = None
@property
def text_field(self) -> str:
return self.mongodb_text_field or DEFAULT_TEXT_FIELD_NAME
@property
def embedding_field(self) -> str:
return self.mongodb_embedding_field or DEFAULT_EMBEDDING_FIELD_NAME
def require_embedding_model(self) -> str:
if not self.litellm_embedding_model:
raise config_error(
"litellm_embedding_model is required in litellm_params for the MongoDB vector store. "
"It must be the same model that produced the vectors stored in "
f"'{self.mongodb_collection or '<collection>'}.{self.embedding_field}', or search results "
"will be meaningless. Example: litellm_embedding_model: openai/text-embedding-3-small"
)
return self.litellm_embedding_model
def require_connection_string(self) -> str:
if not self.mongodb_connection_string:
raise config_error(
"mongodb_connection_string is required in litellm_params for the MongoDB vector store. "
"Example: mongodb+srv://<user>:<password>@<cluster>.mongodb.net for Atlas, or "
"mongodb://<user>:<password>@<host>:27017 for a self-managed deployment"
)
scheme: Final = self.mongodb_connection_string.split("://", 1)[0].lower()
if scheme not in ("mongodb", "mongodb+srv"):
raise config_error(
"mongodb_connection_string must start with 'mongodb://' or 'mongodb+srv://', "
f"got '{self.mongodb_connection_string.split('://', 1)[0]}://'"
)
return self.mongodb_connection_string
def require_database(self) -> str:
if not self.mongodb_database:
raise config_error(
"mongodb_database is required in litellm_params for the MongoDB vector store. "
"Example: mongodb_database: sample_mflix"
)
return self.mongodb_database
def require_collection(self) -> str:
if not self.mongodb_collection:
raise config_error(
"mongodb_collection is required in litellm_params for the MongoDB vector store. "
"Example: mongodb_collection: embedded_movies"
)
return self.mongodb_collection
_MONGODB_PARAM_PREFIX: Final = "mongodb_"
_KNOWN_MONGODB_PARAMS: Final = frozenset(
name for name in _MongoDBSearchParams.model_fields if name.startswith(_MONGODB_PARAM_PREFIX)
)
class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
def __init__(
self,
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
sync_client_factory: Callable[[MongoClientKey], object] | None = None,
async_client_factory: Callable[[MongoClientKey], object] | None = None,
) -> None:
super().__init__()
self.embedding_executor: Final[VectorStoreEmbeddingExecutor] = (
embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor()
)
self.sync_client_factory: Final[Callable[[MongoClientKey], object]] = (
sync_client_factory if sync_client_factory is not None else get_sync_client
)
self.async_client_factory: Final[Callable[[MongoClientKey], object]] = (
async_client_factory if async_client_factory is not None else get_async_client
)
@staticmethod
def _reject_unknown_params(litellm_params: Mapping[str, object]) -> None:
"""Without this a mistyped mongodb_collection reads as 'mongodb_collection is required',
naming a key the reader can see they have set."""
unknown: Final = sorted(
key for key in litellm_params if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS
)
if unknown:
raise config_error(
f"Unrecognised MongoDB vector store parameter(s): {', '.join(unknown)}. "
f"Supported: {', '.join(sorted(_KNOWN_MONGODB_PARAMS))}."
)
@staticmethod
def _query_text(query: str | Sequence[str]) -> str:
text: Final = query if isinstance(query, str) else " ".join(query)
if not text.strip():
raise config_error("query must not be empty")
if len(text) > MAX_QUERY_CHARACTERS:
raise config_error(f"query must be at most {MAX_QUERY_CHARACTERS} characters, got {len(text)}")
return text
@staticmethod
def _limit(vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams) -> int:
requested: Final = vector_store_search_optional_params.get("max_num_results")
if requested is None:
return DEFAULT_MAX_NUM_RESULTS
if not MIN_MAX_NUM_RESULTS <= requested <= MAX_MAX_NUM_RESULTS:
raise config_error(
f"max_num_results must be between {MIN_MAX_NUM_RESULTS} and {MAX_MAX_NUM_RESULTS}, got {requested}"
)
return requested
@staticmethod
def _num_candidates(limit: int, configured: int | None) -> int:
if configured is not None:
if not limit <= configured <= MAX_NUM_CANDIDATES:
raise config_error(
f"mongodb_num_candidates must be between max_num_results ({limit}) and "
f"{MAX_NUM_CANDIDATES}, got {configured}"
)
return configured
return min(max(limit * NUM_CANDIDATES_MULTIPLIER, MIN_NUM_CANDIDATES), MAX_NUM_CANDIDATES)
@staticmethod
def _timeout_ms(timeout: float | httpx.Timeout | None) -> tuple[int, int]:
"""The connect and socket budgets pymongo is built with, in that order."""
if isinstance(timeout, httpx.Timeout):
return (
int((timeout.connect or DEFAULT_CONNECT_TIMEOUT_MS / 1000) * 1000),
int((timeout.read or DEFAULT_SOCKET_TIMEOUT_MS / 1000) * 1000),
)
if timeout is None:
return DEFAULT_CONNECT_TIMEOUT_MS, DEFAULT_SOCKET_TIMEOUT_MS
return min(int(float(timeout) * 1000), DEFAULT_CONNECT_TIMEOUT_MS), int(float(timeout) * 1000)
@classmethod
def _client_key(cls, params: _MongoDBSearchParams, timeout: float | httpx.Timeout | None) -> MongoClientKey:
connect_ms, socket_ms = cls._timeout_ms(timeout)
return MongoClientKey(
connection_string=params.require_connection_string(),
connect_timeout_ms=connect_ms,
socket_timeout_ms=socket_ms,
server_selection_timeout_ms=min(connect_ms, DEFAULT_SERVER_SELECTION_TIMEOUT_MS),
)
@classmethod
def _pipeline(
cls,
vector_store_id: str,
query_vector: Sequence[float],
params: _MongoDBSearchParams,
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
) -> Sequence[Mapping[str, object]]:
if vector_store_search_optional_params.get("filters") is not None:
raise config_error(
"MongoDB vector store does not support the filters parameter yet. "
"Restrict the collection or the MongoDB Vector Search index definition instead."
)
if vector_store_search_optional_params.get("ranking_options") is not None:
raise config_error(
"MongoDB vector store does not support the ranking_options parameter yet. "
"Every result already carries the vectorSearchScore, so filter or re-rank "
"on that rather than having the threshold silently ignored."
)
if vector_store_search_optional_params.get("rewrite_query") is not None:
raise config_error(
"MongoDB vector store does not support the rewrite_query parameter. The query is "
"embedded exactly as sent; rewrite it before calling if you need that."
)
limit: Final = cls._limit(vector_store_search_optional_params)
search: Final = MappingProxyType(
{
"index": vector_store_id,
"path": params.embedding_field,
"queryVector": tuple(query_vector),
"numCandidates": cls._num_candidates(limit, params.mongodb_num_candidates),
"limit": limit,
}
)
projection: Final = MappingProxyType(
{params.text_field: 1, SCORE_FIELD_NAME: MappingProxyType({"$meta": "vectorSearchScore"})}
)
return [ # mutable-ok: pymongo rejects any non-list pipeline in common.validate_list
MappingProxyType({"$vectorSearch": search}),
MappingProxyType({"$project": projection}),
]
@classmethod
def _field_value(cls, document: Mapping[str, object], dotted_path: str) -> str | None:
"""None means absent, which is what separates a mistyped field from genuinely empty text."""
head, _, rest = dotted_path.partition(".")
if head not in document:
return None
value: Final = document[head]
if not rest:
return None if value is None else str(value)
return cls._field_value(value, rest) if isinstance(value, Mapping) else None
@classmethod
def _to_result(cls, document: Mapping[str, object], text_field: str) -> VectorStoreSearchResult:
document_id: Final = document.get("_id")
identifier: Final = None if document_id is None else str(document_id)
content: Final = [ # mutable-ok: VectorStoreSearchResult declares a list of content parts
VectorStoreResultContent(text=cls._field_value(document, text_field) or "", type="text")
]
raw_score: Final = document.get(SCORE_FIELD_NAME)
return VectorStoreSearchResult(
score=float(raw_score) if isinstance(raw_score, (int, float)) else None,
content=content,
file_id=identifier,
filename=identifier,
)
@classmethod
def _raise_for_missing_text_field(
cls, documents: Sequence[Mapping[str, object]], text_field: str, database: str, collection: str
) -> None:
"""$vectorSearch matches documents carrying no text, so a mistyped mongodb_text_field
returns well-scored results with empty content instead of failing."""
if documents and all(cls._field_value(document, text_field) is None for document in documents):
raise config_error(
f"None of the {len(documents)} matched documents in '{database}.{collection}' has a "
f"'{text_field}' field, so every result would carry empty text. Set mongodb_text_field "
"to the field holding the readable text; it accepts a dotted path such as metadata.body."
)
@classmethod
def _to_response(
cls, documents: Sequence[Mapping[str, object]], query_text: str, text_field: str
) -> VectorStoreSearchResponse:
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query=query_text,
data=[ # mutable-ok: VectorStoreSearchResponse declares data as a list
cls._to_result(document, text_field) for document in documents
],
)
@staticmethod
def _raise_for_unusable_index(
catalogue: Sequence[Mapping[str, object]], index_name: str, database: str, collection: str
) -> None:
"""mongod returns zero documents both for a query that matched nothing and for a missing
database, collection or index, so the catalogue decides which one happened."""
if not catalogue:
raise missing_index_error(index_name, database, collection)
entry: Final = catalogue[0]
if not entry.get("queryable"):
raise index_not_ready_error(index_name, database, collection, str(entry.get("status") or "unknown"))
@staticmethod
def _embedding_vector(embedding_response: EmbeddingResponse) -> Sequence[float]:
data: Final = embedding_response.data
if not data:
raise config_error(
"The embedding model returned no embedding for the search query, so there is nothing "
"to search MongoDB with. Check the embedding deployment named by litellm_embedding_model."
)
return data[0]["embedding"]
def execute_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
self._reject_unknown_params(litellm_params)
params: Final = _MongoDBSearchParams.model_validate(litellm_params)
query_text: Final = self._query_text(query)
key: Final = self._client_key(params, timeout)
database: Final = params.require_database()
collection: Final = params.require_collection()
embedding_response: Final = (embedding_executor or self.embedding_executor).embed(
params.require_embedding_model(),
query_text,
params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
)
pipeline: Final = self._pipeline(
vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params
)
try:
client: Final = self.sync_client_factory(key)
target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted
documents: Final = tuple(target.aggregate(pipeline))
except Exception as e:
raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e
if not documents:
try:
catalogue: Final = tuple(target.list_search_indexes(vector_store_id))
except Exception as e:
raise translate_mongo_error(
e, index_name=vector_store_id, database=database, collection=collection
) from e
self._raise_for_unusable_index(catalogue, vector_store_id, database, collection)
self._raise_for_missing_text_field(documents, params.text_field, database, collection)
return self._to_response(documents, query_text, params.text_field)
async def aexecute_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
self._reject_unknown_params(litellm_params)
params: Final = _MongoDBSearchParams.model_validate(litellm_params)
query_text: Final = self._query_text(query)
key: Final = self._client_key(params, timeout)
database: Final = params.require_database()
collection: Final = params.require_collection()
embedding_response: Final = await (embedding_executor or self.embedding_executor).aembed(
params.require_embedding_model(),
query_text,
params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
)
pipeline: Final = self._pipeline(
vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params
)
try:
client: Final = self.async_client_factory(key)
target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted
cursor: Final = await target.aggregate(pipeline)
documents: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly
document async for document in cursor
]
except Exception as e:
raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e
if not documents:
try:
index_cursor: Final = await target.list_search_indexes(vector_store_id)
catalogue: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly
entry async for entry in index_cursor
]
except Exception as e:
raise translate_mongo_error(
e, index_name=vector_store_id, database=database, collection=collection
) from e
self._raise_for_unusable_index(catalogue, vector_store_id, database, collection)
self._raise_for_missing_text_field(documents, params.text_field, database, collection)
return self._to_response(documents, query_text, params.text_field)
def transform_create_vector_store_request(
self,
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
api_base: str,
) -> NoReturn:
raise config_error(_SEARCH_ONLY_MESSAGE)
def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn:
raise config_error(_SEARCH_ONLY_MESSAGE)

23
litellm/llms/perplexity/cost_calculator.py
View file

@ -3,19 +3,23 @@ Helper util for handling perplexity-specific cost calculation
- e.g.: citation tokens, search queries
"""
from datetime import datetime
from typing import Final
from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
from litellm.types.utils import Usage
from litellm.utils import get_model_info
def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
The manual fallback swaps in the model's off_peak_pricing rates while one of its windows is open.
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing perplexity-specific usage information
- current_time: the moment the request is billed at; defaults to now, UTC
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
@ -48,8 +52,21 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
except (ValueError, TypeError):
return default
rates: Final = apply_off_peak_pricing(
model_info,
current_time,
TokenRates(
input_rate=_safe_float_cast(model_info.get("input_cost_per_token")),
output_rate=_safe_float_cast(model_info.get("output_cost_per_token")),
cache_read_rate=0.0,
cache_creation_rate=0.0,
reasoning_rate=None,
),
)
input_cost_per_token: Final = rates.input_rate
output_cost_per_token: Final = rates.output_rate
## CALCULATE INPUT COST
input_cost_per_token: Final = _safe_float_cast(model_info.get("input_cost_per_token"))
prompt_cost: float = (usage.prompt_tokens or 0) * input_cost_per_token
## ADD CITATION TOKENS COST (if present)
@ -60,8 +77,6 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
prompt_cost += citation_tokens * citation_cost_per_token
## CALCULATE OUTPUT COST
output_cost_per_token: Final = _safe_float_cast(model_info.get("output_cost_per_token"))
reasoning_tokens = getattr(usage, "reasoning_tokens", 0) or 0
if reasoning_tokens == 0 and hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
reasoning_tokens = getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0

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