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 fix-skip-symlink-tests-on-windows

Browse source 
# Conflicts:
#	tests/test_litellm/proxy/client/cli/test_claude_settings.py
This commit is contained in:
SWAPI03 2026-09-12 12:33:09 +05:30
commit 98ba90d459
1249 changed files with 126930 additions and 18272 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/ISSUE_TEMPLATE/config.yml vendored
View file

@ -1,4 +1,4 @@
blank_issues_enabled: true
blank_issues_enabled: false
contact_links:
- name: Schedule Demo
url: https://enterprise.litellm.ai/demo

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

@ -106,11 +106,6 @@ dockerfiles:
and lint workflows already exercise that output, so building the image adds no signal about it
paths:
- ui/Dockerfile
- reason: >-
The Rust gateway ships as its own chart and package with a separate release pipeline, so its
image is not part of this repo's Python image set
paths:
- litellm-rust/crates/ai-gateway/Dockerfile
- reason: >-
An example image under cookbook/ that is documentation rather than a shipped artifact
paths:

2
.github/e2e-stack/down.sh vendored
View file

@ -10,7 +10,7 @@ for pid_file in "${STACK_DIR}"/pids/*.pid; do
rm -f "${pid_file}"
done
for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do
for container in e2e-nginx e2e-keycloak e2e-valkey e2e-jaeger e2e-postgres; do
docker rm -f "${container}" >/dev/null 2>&1
done

1
.github/e2e-stack/select_tests.py vendored
View file

@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile(
)
HARNESS: Final = re.compile(
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
r"|^tests/e2e/idp_realm\.json$"
r"|^tests/e2e/gateway/"
r"|^\.github/e2e-stack/"
r"|^\.github/workflows/test-e2e-changed\.yml$"

45
.github/e2e-stack/start-idp.sh vendored Normal file
View file

@ -0,0 +1,45 @@
#!/usr/bin/env bash
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
KEYCLOAK_IMAGE="${E2E_KEYCLOAK_IMAGE:-quay.io/keycloak/keycloak@sha256:ff4257d0d64efbe99ed1ddfaf07765cc3c36dc7518bf8324d41961327f441c54}"
KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}"
: "${DATABASE_HOST:?}" "${DATABASE_PORT:?}" "${DATABASE_USER:?}" "${DATABASE_PASSWORD:?}" "${DATABASE_NAME:?}"
DB_HOST="${DATABASE_HOST}"
DB_NETWORK_ARGS=(--network bridge)
IDP_NETWORK_ARGS=(-p "127.0.0.1:${KEYCLOAK_PORT}:${KEYCLOAK_PORT}")
if [[ "$(uname)" == "Linux" ]]; then
DB_NETWORK_ARGS=(--network host)
IDP_NETWORK_ARGS=(--network host)
elif [[ "${DB_HOST}" == "127.0.0.1" || "${DB_HOST}" == "localhost" ]]; then
DB_HOST=host.docker.internal
fi
docker run --rm "${DB_NETWORK_ARGS[@]}" -e "PGPASSWORD=${DATABASE_PASSWORD}" \
"${POSTGRES_IMAGE}" psql -h "${DB_HOST}" -p "${DATABASE_PORT}" \
-U "${DATABASE_USER}" -d "${DATABASE_NAME}" -v ON_ERROR_STOP=1 \
-c 'CREATE SCHEMA IF NOT EXISTS keycloak' >/dev/null
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
docker run -d --name e2e-keycloak "${IDP_NETWORK_ARGS[@]}" --memory 1536m \
-v "${REPO_ROOT}/tests/e2e/idp_realm.json:/opt/keycloak/data/import/realm.json:ro" \
-e KC_DB=postgres -e "KC_DB_URL_HOST=${DB_HOST}" -e "KC_DB_URL_PORT=${DATABASE_PORT}" \
-e "KC_DB_URL_DATABASE=${DATABASE_NAME}" -e KC_DB_SCHEMA=keycloak \
-e "KC_DB_USERNAME=${DATABASE_USER}" -e "KC_DB_PASSWORD=${DATABASE_PASSWORD}" \
-e KC_DB_POOL_INITIAL_SIZE=2 -e KC_DB_POOL_MIN_SIZE=2 -e KC_DB_POOL_MAX_SIZE=10 \
-e "KC_HTTP_PORT=${KEYCLOAK_PORT}" -e KC_BOOTSTRAP_ADMIN_USERNAME=admin \
-e KC_BOOTSTRAP_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret \
"${KEYCLOAK_IMAGE}" start-dev --import-realm >/dev/null
deadline=$((SECONDS + ${E2E_KEYCLOAK_STARTUP_TIMEOUT:-300}))
until curl -fsS --connect-timeout 2 --max-time 3 \
"http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/.well-known/openid-configuration" >/dev/null 2>&1; do
if ((SECONDS >= deadline)); then
echo 'e2e-stack: timed out waiting for the Keycloak realm' >&2
exit 1
fi
sleep 2
done
echo 'e2e-stack: Keycloak realm is up'

9
.github/e2e-stack/up.sh vendored
View file

@ -25,6 +25,7 @@ DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
KEYCLOAK_PORT="${E2E_KEYCLOAK_PORT:-8081}"
MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}"
@ -124,6 +125,9 @@ SERVER_ENV=(
"OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
"SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
"PYTHONPATH=${REPO_ROOT}"
"JWT_PUBLIC_KEY_URL=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e/protocol/openid-connect/certs"
"JWT_ISSUER=http://127.0.0.1:${KEYCLOAK_PORT}/realms/litellm-e2e"
"JWT_AUDIENCE=litellm-e2e"
)
if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then
printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json"
@ -132,6 +136,8 @@ fi
cd "${REPO_ROOT}"
env "${SERVER_ENV[@]}" "E2E_KEYCLOAK_PORT=${KEYCLOAK_PORT}" bash .github/e2e-stack/start-idp.sh
log "running migrations"
env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1
@ -200,6 +206,9 @@ LITELLM_MASTER_KEY=${MASTER_KEY}
REDIS_HOST=127.0.0.1
REDIS_PORT=${REDIS_PORT}
E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
E2E_KEYCLOAK_URL=http://127.0.0.1:${KEYCLOAK_PORT}
E2E_KEYCLOAK_ADMIN_USER=admin
E2E_KEYCLOAK_ADMIN_PASSWORD=e2e-ephemeral-idp-not-a-secret
SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem
DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME}
EOF

2
.github/pull_request_template.md vendored
View file

@ -127,6 +127,8 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
- Low: anything else worth noting: naming, cleanup, an edge case nobody hits
Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a
human reader
If you assumed something instead of testing it, e.g. "only reproduces with X on" or "no
user-observable behavior difference", list it here too with what breaks if it is wrong
Leave this section empty if there are none -->
## QA runbook

4
.github/scripts/verify_linux_native_wheel.py vendored
View file

@ -205,7 +205,7 @@ def main(
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 20_000_000
native_size_limit: Final = 25_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
@ -222,7 +222,7 @@ def main(
("Python extension entry point is present", extension_entry_point_present),
("Native module loads", native_module_loads),
("Production module omits the panic test hook", panic_test_hook_absent),
("Native extension does not exceed 20 MB", native_size_within_limit),
("Native extension does not exceed 25 MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)

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

@ -113,7 +113,7 @@ 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 --extra mongodb
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
- name: Cache Prisma binaries

73
.github/workflows/ai-gateway-image.yml vendored Normal file
View file

@ -0,0 +1,73 @@
name: ai-gateway image
on:
push:
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
ai-gateway-image:
name: ai-gateway release image
runs-on: ubuntu-latest
timeout-minutes: 60
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build the release image
run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} .
- name: Start the gateway and wait for readiness
env:
IMAGE: litellm-ai-gateway:${{ github.sha }}
run: |
docker run -d --name ai-gateway -p 4001:4001 \
-e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \
-e OPENAI_API_KEY=sk-ci-not-a-real-key \
"$IMAGE"
for _ in $(seq 1 60); do
if curl -fsS http://127.0.0.1:4001/health/readiness; then
echo "gateway is serving readiness"
exit 0
fi
sleep 2
done
echo "gateway never became ready" >&2
docker logs ai-gateway >&2
exit 1
- name: Assert the gateway loaded the baked config
run: |
docker logs ai-gateway 2>&1 | tee gateway.log
grep 'via python config reader' gateway.log
- name: Stop the gateway
if: always()
run: docker rm -f ai-gateway || true

2
.github/workflows/test-code-quality.yml vendored
View file

@ -81,7 +81,7 @@ jobs:
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

5
.github/workflows/test-e2e-changed.yml vendored
View file

@ -27,6 +27,8 @@ jobs:
sparse-checkout: |
.github/e2e-stack
tests/e2e/access_control
tests/e2e/management/test_jwt_management_e2e.py
tests/e2e/other/test_jwt_auth_e2e.py
persist-credentials: false
ref: ${{ github.sha }}
@ -45,7 +47,8 @@ jobs:
--jq '.[] | select(.status != "removed") | .filename')"
gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}"
tests="$(printf '%s\n' "${files}" \
| python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)"
| python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py \
tests/e2e/management/test_jwt_management_e2e.py tests/e2e/other/test_jwt_auth_e2e.py)"
echo "tests=${tests}" >> "${GITHUB_OUTPUT}"
if [ -n "${tests}" ]; then
echo "any=true" >> "${GITHUB_OUTPUT}"

103
.github/workflows/test-e2e-redis-chaos.yml vendored Normal file
View file

@ -0,0 +1,103 @@
name: "Redis Chaos E2E"
on:
workflow_dispatch:
workflow_call:
inputs:
ref:
description: "Commit SHA or ref to test. Defaults to the ref the workflow was triggered on"
required: false
type: string
permissions:
contents: read
jobs:
redis-chaos-e2e:
runs-on: ubuntu-latest-16-cores
timeout-minutes: 30
services:
postgres:
image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U llmproxy"
--health-interval 5s
--health-timeout 5s
--health-retries 10
valkey:
image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa
ports:
- 6379:6379
options: >-
--health-cmd "valkey-cli ping"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-redis-chaos-e2e
LITELLM_LOG: WARNING
JSON_LOGS: "true"
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
ref: ${{ inputs.ref || github.sha }}
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --group e2e-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Start a multi-worker proxy on the chaos config
run: |
nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_chaos_ci_config.yml --port 4000 --num_workers 4 > proxy.log 2>&1 &
echo "E2E_PROXY_PID=$!" >> "$GITHUB_ENV"
echo "E2E_PROXY_LOG=$(pwd)/proxy.log" >> "$GITHUB_ENV"
for _ in $(seq 1 90); do
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
exit 0
fi
sleep 2
done
echo "proxy never became live"
tail -n 100 proxy.log
exit 1
- name: Run the Redis chaos load test
env:
E2E_REDIS_CHAOS: "1"
LITELLM_PROXY_URL: http://localhost:4000
REDIS_HOST: 127.0.0.1
REDIS_PORT: "6379"
run: |
uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log

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

@ -113,7 +113,7 @@ jobs:
- name: Check ruff format
if: steps.changes.outputs.decision != 'skip'
run: |
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- ':(glob)litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
echo "No changed litellm Python files to check with ruff format."
exit 0
@ -172,7 +172,7 @@ jobs:
- name: Check tests/e2e basedpyright (zero errors)
if: steps.changes.outputs.decision != 'skip'
run: |
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- ':(glob)tests/e2e/**/*.py' | grep -q .; then
uv run --no-sync basedpyright tests/e2e
else
echo "No changed tests/e2e Python files; skipping."

4
.github/workflows/test-litellm-ui-unit.yml vendored
View file

@ -66,7 +66,7 @@ jobs:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
full_suite() { npm run test -- --run --pool forks --poolOptions.forks.maxForks=14; }
full_suite() { npm run test -- --run --pool forks --maxWorkers=14; }
if [ -z "$BASE_SHA" ]; then
echo "Push to $GITHUB_REF_NAME: running the full suite"
@ -95,4 +95,4 @@ jobs:
echo "Pull request: running tests related to ${#changed_files[@]} changed UI files"
npm run test -- related "${changed_files[@]}" --run --passWithNoTests \
--pool forks --poolOptions.forks.maxForks=14
--pool forks --maxWorkers=14

28
.github/workflows/test-rust.yml vendored
View file

@ -4,8 +4,22 @@ on:
push:
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
@ -20,8 +34,22 @@ on:
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/custom_logger.py"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"

40
.github/workflows/test-terraform-modules.yml vendored
View file

@ -25,13 +25,17 @@ concurrency:
cancel-in-progress: true
jobs:
aws-module:
name: fmt, validate, test (aws)
module:
name: fmt, validate, test (${{ matrix.module }})
runs-on: ubuntu-latest
timeout-minutes: 15
strategy:
fail-fast: false
matrix:
module: [aws, gcp]
defaults:
run:
working-directory: terraform/litellm/aws
working-directory: terraform/litellm/${{ matrix.module }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -51,35 +55,7 @@ jobs:
- name: validate
run: terraform validate
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
# Plan-only, mock_provider-backed: no cloud credentials, no API calls.
- name: test
run: terraform test
gcp-module:
name: fmt, validate, test (gcp)
runs-on: ubuntu-latest
timeout-minutes: 15
defaults:
run:
working-directory: terraform/litellm/gcp
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
with:
terraform_version: 1.13.3
terraform_wrapper: false
- name: fmt
run: terraform fmt -recursive -check -diff
- name: init
run: terraform init -backend=false -input=false
- name: validate
run: terraform validate
- name: test
run: terraform test

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

@ -200,6 +200,7 @@ jobs:
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
tests/test_gateway
workers: 4
reruns: 2
timeout-minutes: 20

2
AGENTS.md
View file

@ -1 +1,3 @@
Read @CLAUDE.md for coding guidelines
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check

21
Dockerfile
View file

@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -67,7 +83,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Copy full source tree
@ -90,7 +105,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@ -112,7 +126,8 @@ USER root
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

4
Makefile
View file

@ -150,8 +150,8 @@ lint-install:
# Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step:
# only the litellm Python files changed vs the base are checked, so a pre-existing
# format issue elsewhere doesn't block an unrelated commit. Git pathspecs match
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
# recursively, so 'litellm/*.py' covers top-level files and nested modules alike,
# the same set CI's ':(glob)litellm/**/*.py' selects.
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@base_ref=$$($(RESOLVE_BASE)) && \
changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \

2
README.md
View file

@ -262,6 +262,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
}
```
For MCP OAuth, an upstream may advertise dynamic client registration but refuse requests with HTTP 401 or 403. If the provider requires a pre-registered OAuth app, configure its `credentials.client_id` and, when required, `credentials.client_secret` on the MCP server. This skips dynamic registration in the gateway sign-in flow. The provider must approve the app for MCP access; reaching its authorization page does not establish that login or tool calls will succeed
[**Docs: MCP Gateway**](https://docs.litellm.ai/docs/mcp)
</details>

2
basedpyright-code-budget.json
View file

@ -105,7 +105,7 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38271
"limit": 38269
},
"reportUnknownParameterType": {
"limit": 19584

21
docker/Dockerfile.database
View file

@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -65,7 +81,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Copy full source tree
@ -88,7 +103,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@ -103,7 +117,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

22
docker/Dockerfile.non_root
View file

@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -71,7 +87,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Copy full source tree
@ -100,7 +115,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
@ -111,7 +125,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13; \
fi
@ -131,8 +144,9 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
done
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
# Copy only what runtime needs. The application is installed inside the venv;
# the rest of the builder's /app is source and build metadata that must not

6
docker/component_entrypoint.sh
View file

@ -1,5 +1,11 @@
#!/bin/sh
# stale samples from a previous container incarnation would be summed into the aggregate
if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then
mkdir -p "$PROMETHEUS_MULTIPROC_DIR"
rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"

136
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -11,10 +11,14 @@ import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import functools
import configparser
import contextlib
import itertools
import re
import tempfile
from collections.abc import Generator, Iterator, Sequence
from contextvars import ContextVar
from typing import TYPE_CHECKING, ClassVar, Literal, Optional
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -433,12 +437,101 @@ _default_detect_secrets_config = {
"name": "ZendeskSecretKeyDetector",
"path": _custom_plugins_path + "/zendesk_secret_key.py",
},
{
"name": "CredentialKeywordDetector",
"path": _custom_plugins_path + "/credential_keyword.py",
},
{"name": "Base64HighEntropyString", "limit": 4.5},
{"name": "HexHighEntropyString", "limit": 3.0},
],
}
_CONFIG_SECTION: Final = "litellm-prompt"
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
_SHELL_OPERATORS: Final = ";&|"
_SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
_SCAN_SUFFIX: Final = ".py"
@contextlib.contextmanager
def _temp_file(text: str) -> Generator[str, None, None]:
temp_file: Final = tempfile.NamedTemporaryFile(suffix=_SCAN_SUFFIX, delete=False)
try:
temp_file.write(text.encode("utf-8"))
temp_file.close()
yield temp_file.name
finally:
temp_file.close()
os.remove(temp_file.name)
def _scan_lines(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
from detect_secrets import SecretsCollection
secrets: Final = SecretsCollection()
with _temp_file("\n".join(lines)) as path:
secrets.scan_file(path)
return frozenset(
(found_secret.secret_value, found_secret.type)
for file in secrets.files
for found_secret in secrets[file]
if found_secret.secret_value is not None
)
def _classify_line(state: tuple[bool, str | None], numbered: tuple[int, str]) -> tuple[bool, str | None]:
open_option: Final = state[0]
number, line = numbered
stripped: Final = line.strip()
if not stripped or stripped[0] in "#;":
return open_option, None
shell_assignment: Final = _SHELL_ASSIGNMENT.match(stripped)
if shell_assignment is not None:
return True, f"{shell_assignment['key']}_{number}={shell_assignment['value']}"
assignment: Final = _ASSIGNMENT_LINE.match(stripped)
if assignment is not None:
return True, f"{assignment.group()[:-1].strip()}_{number}{stripped[assignment.end() - 1 :]}"
if line[0].isspace() and open_option:
return True, line
return False, None
def _parseable_lines(text: str) -> Iterator[str]:
states: Final = itertools.accumulate(enumerate(text.splitlines()), _classify_line, initial=(False, None))
return (line for _, line in states if line is not None)
def _lone_value(line: str) -> str | None:
tokens: Final = line.split()
if not tokens or '"' in tokens[0]:
return None
value: Final = tokens[0].rstrip(_SHELL_OPERATORS)
if len(tokens) == 1 or value != tokens[0] or _SHELL_TRAILER.fullmatch(tokens[1]) is not None:
return value
return None
def _quoted_assignments(text: str) -> tuple[str, ...]:
parser: Final = configparser.ConfigParser(interpolation=None)
parser.optionxform = str # pyright: ignore[reportAttributeAccessIssue] # configparser types optionxform as a method
parser.read_string(f"[{_CONFIG_SECTION}]\n" + "\n".join(_parseable_lines(text)))
return tuple(
f'{key} = "{value}"'
for section in parser
for key, values in parser.items(section)
for line in values.splitlines()
if (value := _lone_value(line)) is not None
)
class _ENTERPRISE_SecretDetection(CustomGuardrail):
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
# path skips should_run_check and never sees data["prompt"]).
@ -449,35 +542,21 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
super().__init__(**kwargs)
def scan_message_for_secrets(self, message_content: str):
from detect_secrets import SecretsCollection
from detect_secrets.settings import transient_settings
temp_file = tempfile.NamedTemporaryFile(delete=False)
temp_file.write(message_content.encode("utf-8"))
temp_file.close()
secrets = SecretsCollection()
detect_secrets_config = (
self.user_defined_detect_secrets_config or _default_detect_secrets_config
)
with transient_settings(detect_secrets_config):
secrets.scan_file(temp_file.name)
os.remove(temp_file.name)
found: Final = _scan_lines(
(*message_content.splitlines(), *_quoted_assignments(message_content))
)
return [
{"type": found_secret.type, "value": found_secret.secret_value}
for file in sorted(secrets.files)
for found_secret in sorted(
secrets[file],
key=lambda secret: (
-len(secret.secret_value or ""),
secret.type,
secret.secret_value or "",
),
{"type": secret_type, "value": value}
for value, secret_type in sorted(
found, key=lambda pair: (-len(pair[0]), pair[1], pair[0])
)
if found_secret.secret_value is not None
]
def redact_text(self, text: str, source: str = "message") -> str:
@ -490,15 +569,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
if counts is not None:
for secret in detected_secrets:
counts[secret["type"]] = counts.get(secret["type"], 0) + 1
secret_types = [secret["type"] for secret in detected_secrets]
secret_types: Final = sorted(
dict.fromkeys(secret["type"] for secret in detected_secrets)
)
verbose_proxy_logger.warning(
f"Detected and redacted secrets in {source}: {secret_types}"
"Detected and redacted secrets in %s: %s", source, secret_types
)
return functools.reduce(
lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"),
detected_secrets,
text,
pattern: Final = re.compile(
"|".join(re.escape(secret["value"]) for secret in detected_secrets)
)
return pattern.sub("[REDACTED]", text)
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.permissions is not None:

63
enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/credential_keyword.py Normal file
View file

@ -0,0 +1,63 @@
import re
from collections.abc import Generator, Mapping
from string import punctuation
from typing import Final
from detect_secrets.plugins.keyword import (
QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP,
KeywordDetector,
)
_CREDENTIAL_VALUE: Final = re.compile(r"[^\s()\[\]]+")
_ENVIRONMENT_REFERENCE: Final = re.compile(r"os\.environ/\w+", re.IGNORECASE)
_ENVIRONMENT_VARIABLE_NAME: Final = re.compile(r"[A-Z][A-Z0-9]*(?:_[A-Z0-9]+)+")
_LOWERCASE_WORD_SEQUENCE: Final = re.compile(r"[a-z]+(?:[-._/][a-z]+)+")
_ISO_8601_TIMESTAMP: Final = re.compile(
r"\d{4}-\d{2}-\d{2}(?:T\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?)?"
)
_URL_WITHOUT_USERINFO_OR_QUERY: Final = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s@?]*")
_BENIGN_VALUES: Final = (
_ENVIRONMENT_REFERENCE,
_ENVIRONMENT_VARIABLE_NAME,
_LOWERCASE_WORD_SEQUENCE,
_ISO_8601_TIMESTAMP,
_URL_WITHOUT_USERINFO_OR_QUERY,
)
class CredentialKeywordDetector(KeywordDetector): # pyright: ignore[reportUntypedBaseClass] # detect_secrets ships no type information
secret_type = "Credential Keyword"
def __init__(self, minimum_length: int = 12, keyword_exclude: str | None = None) -> None:
if (
not isinstance(minimum_length, int) # pyright: ignore[reportUnnecessaryIsInstance] # the value comes from an operator's YAML
or minimum_length < 1
):
raise ValueError(f"minimum_length must be a positive integer, got {minimum_length!r}")
super().__init__(keyword_exclude=keyword_exclude)
self.minimum_length = minimum_length
def _is_credential(self, value: str) -> bool:
core: Final = value.strip(punctuation)
return (
len(value) >= self.minimum_length
and _CREDENTIAL_VALUE.fullmatch(value) is not None
and all(benign.fullmatch(core) is None for benign in _BENIGN_VALUES)
)
def analyze_string(
self,
string: str,
denylist_regex_to_group: Mapping[re.Pattern[str], int] | None = None,
) -> Generator[str, None, None]:
if self.keyword_exclude is not None and self.keyword_exclude.search(string):
return
regex_to_group: Final = (
QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP if denylist_regex_to_group is None else denylist_regex_to_group
)
yield from (
match.group(group)
for regex, group in regex_to_group.items()
for match in regex.finditer(string)
if self._is_credential(match.group(group))
)

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

@ -142,6 +142,40 @@ class CheckBatchCost:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None:
org_id = getattr(job, "org_id", None)
if org_id:
return org_id
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
)
key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
if key_org_id:
return key_org_id
except Exception as e:
verbose_proxy_logger.error(
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
f"still trying the team's: {e}"
)
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
)
return getattr(team_row, "organization_id", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
return None
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
) -> dict[str, object]:
@ -153,6 +187,10 @@ class CheckBatchCost:
user_api_key_alias; when it has no alias, or the key has since been rotated or
deleted, the field keeps the creating user's alias that _get_user_info filled in,
because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id must be resolved here too: the spend update writer reads it
off this metadata to increment organization spend, so leaving it out silently
drops batch cost from org accounting for keys and teams that belong to one.
"""
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
@ -172,6 +210,9 @@ class CheckBatchCost:
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
org_id: Final = await self._get_org_id(job, batch_id)
if org_id is not None:
metadata["user_api_key_org_id"] = org_id
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
@ -641,7 +682,7 @@ class CheckBatchCost:
from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -805,6 +846,7 @@ class CheckBatchCost:
function_id=str(uuid.uuid4()),
)
deployment_api_base: Final = deployment_info.litellm_params.api_base
logging_obj.update_environment_variables(
litellm_params={
# set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@ -813,9 +855,17 @@ class CheckBatchCost:
"user-agent": CHECK_BATCH_COST_USER_AGENT,
}
},
"metadata": await self._build_creator_attribution_metadata(job, batch_id),
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
"metadata": {
**(await self._build_creator_attribution_metadata(job, batch_id)),
# spend logs read the deployment identity off these metadata keys, so
# without them the batch cost row carries no model_id or model_group
"model_info": {"id": model_id},
"model_group": deployment_info.model_name,
},
},
optional_params={},
custom_llm_provider=str(llm_provider) if llm_provider else None,
)
if not await self._claim_job_for_costing(job):
@ -833,6 +883,8 @@ class CheckBatchCost:
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)

35
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -280,6 +280,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
if user_api_key_dict.org_id:
return user_api_key_dict.org_id
if not user_api_key_dict.team_id:
return None
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
try:
team: Final = await get_team_object(
team_id=user_api_key_dict.team_id,
prisma_client=self.prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
return team.organization_id
except Exception as e:
verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
return None
async def store_unified_object_id(
self,
unified_object_id: str,
@ -352,6 +373,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_purpose": file_purpose,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"org_id": await self._resolve_creator_org_id(user_api_key_dict),
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
**attribution_columns,
@ -1779,7 +1801,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1790,7 +1821,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
if stored_file_object:
return stored_file_object
return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id})
elif delete_response:
delete_response.id = file_id
return delete_response

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.65"
version = "0.1.66"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.65"
version = "0.1.66"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

30
gateway/Dockerfile
View file

@ -1,9 +1,25 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.25.2
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
FROM $UV_IMAGE AS uvbin
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
ARG PGBOUNCER_VERSION
ARG PGBOUNCER_SHA256
USER root
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
WORKDIR /build
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
tar xzf pgbouncer.tar.gz --strip-components=1 && \
./configure --prefix=/usr/local --with-openssl=/usr && \
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# ---------- Builder ----------
FROM $LITELLM_BUILD_IMAGE AS builder
@ -47,7 +63,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
@ -60,9 +75,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
# compiled Rust extension must live next to the source or it is never imported.
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma
@ -75,7 +93,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -92,15 +110,17 @@ ENV HOME=/home/nonroot \
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
USER nonroot
EXPOSE 4000/tcp
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
CMD ["--host", "0.0.0.0", "--port", "4000"]

77
gateway/launch.py Normal file
View file

@ -0,0 +1,77 @@
"""Gateway supervisor: assemble DATABASE_URL, start the in-container PgBouncer, then run uvicorn.
``gateway/main.py`` assembles ``DATABASE_URL`` inside every uvicorn worker, which
is fine for a plain Postgres URL but not for the pooler: PgBouncer must be
started exactly once per pod, before the workers fork, and the workers must be
handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
``DatabaseURLSettings.apply_to_env`` under password auth, and one marked pooled
wins under token auth too, so exporting it here is enough for every worker to
pick the pooled URL up unchanged.
Run with:
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
"""
import os
import sys
from collections.abc import Callable, Mapping, Sequence
from typing import Final
from uvicorn.main import main as uvicorn_main
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerSettings,
export_pooled_database_url,
start_in_container_pgbouncer,
)
GATEWAY_APP: Final = "gateway.main:app"
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str, ...]:
"""Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
return (GATEWAY_APP, *argv)
return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive)
def pool_database_url(
settings: DatabaseURLSettings,
pgbouncer: PgBouncerSettings,
environ: Mapping[str, str],
) -> str | PgBouncerError | None:
"""Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
auth the pooler mints and renews the upstream token itself.
"""
if not pgbouncer.enabled:
return None
upstream_url: Final = environ.get("DATABASE_URL")
if upstream_url is None:
return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled")
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())
def _serve(argv: Sequence[str]) -> None:
uvicorn_main(tuple(argv), prog_name="uvicorn")
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
settings: Final = DatabaseURLSettings.from_env()
settings.apply_to_env()
pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
if isinstance(pooled_url, PgBouncerError):
sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}")
if pooled_url is not None:
export_pooled_database_url(pooled_url)
serve(uvicorn_argv(argv, os.environ))
if __name__ == "__main__":
main(sys.argv[1:])

4
gateway/routes/allowlist.py
View file

@ -3,7 +3,8 @@
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics).
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
Any path not listed here is dropped from the gateway process so management/UI
endpoints don't ride on the same pods.
@ -121,6 +122,7 @@ GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
"/docs/oauth2-redirect",
"/redoc",
"/test",
"/debug/memory/summary",
}
)

160
helm/litellm-helm/templates/_helpers.tpl
View file

@ -161,3 +161,163 @@ taken before the change, which by that point no longer exists.
{{- fail (printf "postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got %q). An unpinned tag can start a different PostgreSQL major against the existing data directory, which makes the database unreadable and is not recoverable in place. Crossing a major version requires a dump and restore." $tag) -}}
{{- end -}}
{{- end -}}
{{/*
Environment shared by the proxy container and the opt-in collector sidecar:
database, pgbouncer, master key, redis, user envVars. Both containers must see
the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
and the same spend transaction buffer.
*/}}
{{- define "litellm.proxyEnv" -}}
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
value: {{ .Values.service.port | quote}}
{{- if .Values.db.deployStandalone }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: username
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: password
- name: DATABASE_HOST
value: {{ .Release.Name }}-postgresql
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaUrlKey }}
{{- else if .Values.db.readReplicaUrl }}
- name: DATABASE_URL_READ_REPLICA
value: {{ .Values.db.readReplicaUrl | quote }}
{{- end }}
{{- if .Values.db.connectionPool.enabled }}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: {{ .Values.db.connectionPool.maxDbConnections | quote }}
- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: {{ .Values.db.connectionPool.maxClientConn | quote }}
{{- end }}
- name: PROXY_MASTER_KEY
valueFrom:
secretKeyRef:
name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }}
key: {{ .Values.masterkeySecretKey | default "masterkey" }}
{{- if .Values.redis.enabled }}
- name: REDIS_HOST
value: {{ include "litellm.redis.serviceName" . }}
- name: REDIS_PORT
value: {{ include "litellm.redis.port" . | quote }}
- name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{ toYaml . }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- end }}
{{- end -}}
{{/*
Proxy-only metering and metrics env. The collector sidecar serves no HTTP
traffic, so it gets neither.
*/}}
{{- define "litellm.proxyMetricsEnv" -}}
{{- if .Values.billingMetrics.enabled }}
{{ include "litellm.billingMetricsEnv" . }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- end -}}
{{/*
Directory of the collector's unix socket, shared between the two containers
through an emptyDir. Empty when the sidecar is off or uses 127.0.0.1 TCP.
*/}}
{{- define "litellm.collector.socketDir" -}}
{{- if and .Values.collector.enabled (hasPrefix "unix://" .Values.collector.address) -}}
{{- dir (trimPrefix "unix://" .Values.collector.address) -}}
{{- end -}}
{{- end -}}
{{- define "litellm.collectorEnv" -}}
- name: LITELLM_COLLECTOR_ENABLED
value: "true"
- name: LITELLM_COLLECTOR_ADDRESS
value: {{ .Values.collector.address | quote }}
- name: LITELLM_COLLECTOR_BUFFER_SIZE
value: {{ .Values.collector.bufferSize | quote }}
- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: {{ .Values.collector.onUnavailable | quote }}
- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
value: {{ .Values.collector.drainTimeoutSeconds | quote }}
{{- end -}}

170
helm/litellm-helm/templates/deployment.yaml
View file

@ -56,111 +56,10 @@ spec:
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.image.pullPolicy }}
env:
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
value: {{ .Values.service.port | quote}}
{{- if .Values.db.deployStandalone }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: username
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "litellm.fullname" . }}-dbcredentials
key: password
- name: DATABASE_HOST
value: {{ .Release.Name }}-postgresql
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaUrlKey }}
{{- else if .Values.db.readReplicaUrl }}
- name: DATABASE_URL_READ_REPLICA
value: {{ .Values.db.readReplicaUrl | quote }}
{{- end }}
- name: PROXY_MASTER_KEY
valueFrom:
secretKeyRef:
name: {{ .Values.masterkeySecretName | default (printf "%s-masterkey" (include "litellm.fullname" .)) }}
key: {{ .Values.masterkeySecretKey | default "masterkey" }}
{{- if .Values.redis.enabled }}
- name: REDIS_HOST
value: {{ include "litellm.redis.serviceName" . }}
- name: REDIS_PORT
value: {{ include "litellm.redis.port" . | quote }}
- name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- include "litellm.proxyEnv" . | nindent 12 }}
{{- include "litellm.proxyMetricsEnv" . | nindent 12 }}
{{- if .Values.collector.enabled }}
{{- include "litellm.collectorEnv" . | nindent 12 }}
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
@ -189,6 +88,11 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.metricsServer.enabled }}
- name: metrics
containerPort: {{ .Values.metricsServer.port }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}
@ -233,6 +137,10 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -240,6 +148,53 @@ spec:
lifecycle:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.collector.enabled }}
- name: {{ include "litellm.name" . }}-collector
securityContext:
{{- toYaml .Values.securityContext | nindent 12 }}
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.image.pullPolicy }}
command: {{ toYaml .Values.collector.command | nindent 12 }}
env:
{{- include "litellm.proxyEnv" . | nindent 12 }}
{{- include "litellm.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector
{{- if not (hasKey (default dict .Values.envVars) "CONFIG_FILE_PATH") }}
- name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
- secretRef:
name: {{ . }}
{{- end }}
{{- range .Values.environmentConfigMaps }}
- configMapRef:
name: {{ . }}
{{- end }}
resources:
{{- toYaml .Values.collector.resources | nindent 12 }}
volumeMounts:
- name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.collector.socketDir" . }}
{{- end }}
{{ if .Values.securityContext.readOnlyRootFilesystem }}
- name: tmp
mountPath: /tmp
- name: cache
mountPath: /.cache
- name: npm
mountPath: /.npm
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- end }}
{{- with .Values.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
@ -268,6 +223,11 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- if include "litellm.collector.socketDir" . }}
- name: collector-socket
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- with .Values.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

28
helm/litellm-helm/templates/hpa.yaml
View file

@ -18,6 +18,15 @@ spec:
{{- end }}
metrics:
{{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- if and .Values.collector.enabled .Values.collector.scaleOnProxyContainerCpu }}
- type: ContainerResource
containerResource:
name: cpu
container: {{ include "litellm.name" . }}
target:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- else }}
- type: Resource
resource:
name: cpu
@ -25,6 +34,7 @@ spec:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
{{- end }}
{{- end }}
{{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
- type: Resource
resource:
@ -33,4 +43,22 @@ spec:
type: Utilization
averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
{{- end }}
{{- with .Values.autoscaling.targetRequestsPerSecond }}
- type: Pods
pods:
metric:
name: litellm_requests_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- with .Values.autoscaling.targetTokensPerSecond }}
- type: Pods
pods:
metric:
name: litellm_tokens_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- end }}

21
helm/litellm-helm/templates/keda.yaml
View file

@ -23,6 +23,27 @@ spec:
triggers:
{{- with .Values.keda.triggers }}
{{- toYaml . | nindent 2 }}
{{- end }}
{{- $prom := .Values.keda.prometheus }}
{{- if or $prom.requestsPerSecond $prom.tokensPerSecond }}
{{- if not $prom.serverAddress }}
{{- fail "keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set" }}
{{- end }}
{{- $selector := printf "namespace=%q,job=%q" .Release.Namespace (printf "%s%s" (include "litellm.fullname" .) (ternary "-metrics" "" .Values.metricsServer.enabled)) }}
{{- with $prom.requestsPerSecond }}
- type: prometheus
metadata:
serverAddress: {{ $prom.serverAddress | quote }}
threshold: {{ toJson . | trimAll "\"" | quote }}
query: {{ printf "sum(rate(litellm_proxy_total_requests_metric_total{%s}[1m]))" $selector | quote }}
{{- end }}
{{- with $prom.tokensPerSecond }}
- type: prometheus
metadata:
serverAddress: {{ $prom.serverAddress | quote }}
threshold: {{ toJson . | trimAll "\"" | quote }}
query: {{ printf "sum(rate(litellm_total_tokens_metric_total{%s}[1m]))" $selector | quote }}
{{- end }}
{{- end }}
advanced:
restoreToOriginalReplicaCount: {{ .Values.keda.restoreToOriginalReplicaCount }}

17
helm/litellm-helm/templates/service-metrics.yaml Normal file
View file

@ -0,0 +1,17 @@
{{- if .Values.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.fullname" . }}-metrics
labels:
{{- include "litellm.labels" . | nindent 4 }}
spec:
type: ClusterIP
ports:
- port: {{ .Values.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.selectorLabels" . | nindent 4 }}
{{- end }}

2
helm/litellm-helm/templates/servicemonitor.yaml
View file

@ -26,7 +26,7 @@ spec:
{{- toYaml .namespaceSelector.matchNames | nindent 4 }}
{{- end }}
endpoints:
- port: http
- port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
path: /metrics/
interval: {{ .interval }}
scrapeTimeout: {{ .scrapeTimeout }}

272
helm/litellm-helm/tests/collector_tests.yaml Normal file
View file

@ -0,0 +1,272 @@
suite: test collector sidecar
templates:
- deployment.yaml
- hpa.yaml
- configmap-litellm.yaml
tests:
- it: should run the proxy alone with no collector env by default
template: deployment.yaml
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
- it: should add the sidecar on the same image and point both containers at the unix socket
template: deployment.yaml
set:
image.tag: test
db.connectionPool.enabled: true
collector.enabled: true
collector.resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "1"
memory: 2Gi
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm:test
- equal:
path: spec.template.spec.containers[1].command
value: [python, -m, litellm.proxy.collector]
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 500m
- equal:
path: spec.template.spec.containers[1].resources.limits.memory
value: 2Gi
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "1000"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: fallback
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_JOB_ROLE
value: collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_JOB_ROLE
value: collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_HOST
value: RELEASE-NAME-postgresql
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: RELEASE-NAME-litellm-dbcredentials
key: password
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
- contains:
path: spec.template.spec.volumes
content:
name: collector-socket
emptyDir:
sizeLimit: 1Mi
- it: should skip the socket volume and pass the policy through on tcp transport
template: deployment.yaml
set:
collector.enabled: true
collector.address: tcp://127.0.0.1:4100
collector.onUnavailable: drop
collector.bufferSize: 50
envVars:
CONFIG_FILE_PATH: /custom/config.yaml
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- notContains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /etc/litellm/config.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /custom/config.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: tcp://127.0.0.1:4100
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: drop
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "50"
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
- it: should keep metrics and billing env on the proxy container only
template: deployment.yaml
set:
collector.enabled: true
metricsServer.enabled: true
metricsServer.port: 9090
billingMetrics.enabled: true
billingMetrics.endpoint: https://metering.example.com
billingMetrics.secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "9090"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://metering.example.com
- notContains:
path: spec.template.spec.containers[1].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
any: true
- notContains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: billing-metrics-mtls
any: true
- it: should give the sidecar the same scratch mounts as the proxy on a read-only root
template: deployment.yaml
set:
collector.enabled: true
securityContext.readOnlyRootFilesystem: true
asserts:
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: npm
mountPath: /.npm
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: cache
mountPath: /.cache
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: tmp
mountPath: /tmp
- it: should keep the pod-wide cpu metric unless asked to scale on the proxy container
template: hpa.yaml
set:
autoscaling.enabled: true
collector.enabled: true
asserts:
- equal: { path: "spec.metrics[0].type", value: Resource }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- it: should scale on the proxy container's cpu only when opted in
template: hpa.yaml
set:
autoscaling.enabled: true
collector.enabled: true
collector.scaleOnProxyContainerCpu: true
asserts:
- equal: { path: "spec.metrics[0].type", value: ContainerResource }
- equal: { path: "spec.metrics[0].containerResource.name", value: cpu }
- equal: { path: "spec.metrics[0].containerResource.container", value: litellm }
- equal: { path: "spec.metrics[0].containerResource.target.averageUtilization", value: 60 }
- isNull: { path: "spec.metrics[0].resource" }
- it: should not switch to the container metric while the sidecar is off
template: hpa.yaml
set:
autoscaling.enabled: true
collector.scaleOnProxyContainerCpu: true
asserts:
- equal: { path: "spec.metrics[0].type", value: Resource }

112
helm/litellm-helm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,112 @@
suite: test in-container connection pool
templates:
- deployment.yaml
- configmap-litellm.yaml
tests:
- it: should not emit pgbouncer env vars by default
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- it: should enable the pool with the default sizing when connectionPool.enabled is set
template: deployment.yaml
set:
db.connectionPool.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "1000"
- it: should pass custom sizing through as strings next to the worker count
template: deployment.yaml
set:
numWorkers: 4
db.connectionPool.enabled: true
db.connectionPool.maxDbConnections: 8
db.connectionPool.maxClientConn: 400
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "400"
- contains:
path: spec.template.spec.containers[0].args
content: "4"
- it: should give the collector sidecar the same pool env as the proxy container
template: deployment.yaml
set:
collector.enabled: true
db.connectionPool.enabled: true
db.connectionPool.maxDbConnections: 8
db.connectionPool.maxClientConn: 400
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "400"
- it: should give the collector sidecar no pool env when the pool is off
template: deployment.yaml
set:
collector.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true

78
helm/litellm-helm/tests/hpa_tests.yaml
View file

@ -61,6 +61,84 @@ tests:
- equal: { path: "spec.metrics[1].resource.name", value: memory }
- equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 }
- it: "renders no workload metrics by default"
set:
autoscaling.enabled: true
autoscaling.targetMemoryUtilizationPercentage: 80
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- notContains: { path: spec.metrics, content: { type: Pods }, any: true }
- it: "adds a requests-per-second Pods metric after the cpu metric"
set:
autoscaling.enabled: true
autoscaling.targetRequestsPerSecond: 90
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- equal:
path: "spec.metrics[1]"
value:
type: Pods
pods:
metric: { name: litellm_requests_per_second }
target: { type: AverageValue, averageValue: "90" }
- it: "adds a tokens-per-second Pods metric on its own"
set:
autoscaling.enabled: true
autoscaling.targetTokensPerSecond: 6M
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- equal:
path: "spec.metrics[1]"
value:
type: Pods
pods:
metric: { name: litellm_tokens_per_second }
target: { type: AverageValue, averageValue: "6M" }
- notContains:
path: spec.metrics
content: { type: Pods, pods: { metric: { name: litellm_requests_per_second } } }
any: true
- it: "renders requests, tokens, cpu and memory metrics together"
set:
autoscaling.enabled: true
autoscaling.targetMemoryUtilizationPercentage: 80
autoscaling.targetRequestsPerSecond: 90
autoscaling.targetTokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.metrics, count: 4 }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- equal: { path: "spec.metrics[1].resource.name", value: memory }
- equal: { path: "spec.metrics[2].pods.metric.name", value: litellm_requests_per_second }
- equal: { path: "spec.metrics[2].pods.target.averageValue", value: "90" }
- equal: { path: "spec.metrics[3].pods.metric.name", value: litellm_tokens_per_second }
- equal: { path: "spec.metrics[3].pods.target.averageValue", value: "6000000" }
- it: "scales on workload metrics alone when the cpu target is cleared"
set:
autoscaling.enabled: true
autoscaling.targetCPUUtilizationPercentage: null
autoscaling.targetRequestsPerSecond: 90
autoscaling.targetTokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- notContains: { path: spec.metrics, content: { type: Resource }, any: true }
- equal: { path: "spec.metrics[0].pods.metric.name", value: litellm_requests_per_second }
- equal: { path: "spec.metrics[1].pods.metric.name", value: litellm_tokens_per_second }
- notMatchRegexRaw: { pattern: per_minute }
- it: "ignores the per-minute keys, which the chart never shipped"
set:
autoscaling.enabled: true
autoscaling.targetRequestsPerMinute: 5400
autoscaling.targetTokensPerMinute: 360000000
asserts:
- lengthEqual: { path: spec.metrics, count: 1 }
- notContains: { path: spec.metrics, content: { type: Pods }, any: true }
- it: "renders no hpa when autoscaling is disabled"
asserts:
- hasDocuments: { count: 0 }

106
helm/litellm-helm/tests/keda_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: "keda"
templates:
- keda.yaml
release:
name: rel
namespace: llm
tests:
- it: "renders no scaled object by default"
asserts:
- hasDocuments: { count: 0 }
- it: "passes user triggers through and adds no prometheus triggers by default"
set:
keda.enabled: true
keda.triggers:
- type: cpu
metricType: Utilization
metadata: { value: "60" }
asserts:
- isKind: { of: ScaledObject }
- equal:
path: spec.triggers
value:
- type: cpu
metricType: Utilization
metadata: { value: "60" }
- it: "scales on release-wide requests per second divided by the per-replica target"
set:
keda.enabled: true
keda.prometheus.serverAddress: http://prometheus-operated.monitoring.svc:9090
keda.prometheus.requestsPerSecond: 90
asserts:
- lengthEqual: { path: spec.triggers, count: 1 }
- equal:
path: "spec.triggers[0]"
value:
type: prometheus
metadata:
serverAddress: http://prometheus-operated.monitoring.svc:9090
threshold: "90"
query: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm"}[1m]))
- it: "scales on tokens per second on its own"
set:
keda.enabled: true
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.tokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.triggers, count: 1 }
- equal: { path: "spec.triggers[0].type", value: prometheus }
- equal: { path: "spec.triggers[0].metadata.threshold", value: "6000000" }
- equal:
path: "spec.triggers[0].metadata.query"
value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm"}[1m]))
- it: "appends requests and tokens triggers after user triggers and selects the metrics service job"
set:
keda.enabled: true
metricsServer.enabled: true
keda.triggers:
- type: cpu
metricType: Utilization
metadata: { value: "60" }
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.requestsPerSecond: 90
keda.prometheus.tokensPerSecond: 6000000
asserts:
- lengthEqual: { path: spec.triggers, count: 3 }
- equal: { path: "spec.triggers[0].type", value: cpu }
- equal: { path: "spec.triggers[1].metadata.threshold", value: "90" }
- equal:
path: "spec.triggers[1].metadata.query"
value: sum(rate(litellm_proxy_total_requests_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m]))
- equal: { path: "spec.triggers[2].metadata.threshold", value: "6000000" }
- equal:
path: "spec.triggers[2].metadata.query"
value: sum(rate(litellm_total_tokens_metric_total{namespace="llm",job="rel-litellm-metrics"}[1m]))
- notMatchRegexRaw: { pattern: "\\* *60|per_minute|PerMinute" }
- it: "ignores the per-minute keys, which the chart never shipped"
set:
keda.enabled: true
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.requestsPerMinute: 5400
keda.prometheus.tokensPerMinute: 360000000
asserts:
- isKind: { of: ScaledObject }
- isNullOrEmpty: { path: spec.triggers }
- it: "refuses a workload target without a prometheus server address"
set:
keda.enabled: true
keda.prometheus.requestsPerSecond: 90
asserts:
- failedTemplate:
errorMessage: keda.prometheus.serverAddress is required when keda.prometheus.requestsPerSecond or tokensPerSecond is set
- it: "yields to the hpa when both autoscalers are enabled"
set:
autoscaling.enabled: true
keda.enabled: true
keda.prometheus.serverAddress: http://prom:9090
keda.prometheus.requestsPerSecond: 90
asserts:
- hasDocuments: { count: 0 }

106
helm/litellm-helm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: separate metrics server
templates:
- configmap-litellm.yaml
- deployment.yaml
- service.yaml
- service-metrics.yaml
- servicemonitor.yaml
tests:
- it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default
asserts:
- notContains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
any: true
template: deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- hasDocuments:
count: 0
template: service-metrics.yaml
- it: should scrape the proxy port when the metrics server is disabled
template: servicemonitor.yaml
set:
serviceMonitor.enabled: true
asserts:
- equal:
path: spec.endpoints[0].port
value: http
- it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor
set:
metricsServer.enabled: true
metricsServer.port: 4101
serviceMonitor.enabled: true
service.type: LoadBalancer
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "4101"
template: deployment.yaml
- contains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
containerPort: 4101
protocol: TCP
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-metrics
template: service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
template: service-metrics.yaml
- equal:
path: spec.endpoints[0].port
value: metrics
template: servicemonitor.yaml
- equal:
path: spec.endpoints[0].path
value: /metrics/
template: servicemonitor.yaml
- it: should reject a metrics port equal to the proxy port
template: deployment.yaml
set:
metricsServer.enabled: true
metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: metricsServer.port must differ from service.port

102
helm/litellm-helm/values.yaml
View file

@ -180,6 +180,58 @@ proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT)
# so a scrape never runs on an inference worker. Adds a `metrics` port to the
# container and a dedicated ClusterIP `<release>-metrics` Service, and the
# ServiceMonitor scrapes it instead of the proxy port. The separate port has
# no virtual-key auth: keep it off public ingress. Needs the proxy image
# v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
# Opt-in sidecar that runs the post-response spend pipeline (cost calculation,
# spend logs, spend counters, budget reservation reconciliation) so the proxy's
# uvicorn workers only serialise a compact typed event and go back to serving
# inference. Same image and tag as the proxy, second container in the same pod,
# fed over loopback (a unix socket on a shared emptyDir, or 127.0.0.1 TCP). It
# reuses the pod's in-container pgbouncer (db.connectionPool) and the same Redis
# spend transaction buffer, so the per-pod DB connection budget is unchanged.
# Delivery is at-most-once inside the pod: events already handed to the sidecar
# are lost if it crashes before writing them; events the workers could not hand
# over follow onUnavailable. Both containers drain on SIGTERM within
# terminationGracePeriodSeconds
collector:
enabled: false
# unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or tcp://127.0.0.1:<port>
address: unix:///var/run/litellm/collector.sock
# Events each uvicorn worker holds in memory while the sidecar is slow or restarting
bufferSize: 1000
# fallback: run the pipeline in the worker when the sidecar is unreachable or the
# buffer is full (spend stays exact, that request costs proxy CPU again)
# drop: count and discard the event instead (spend under-reports)
onUnavailable: fallback
# How long the workers keep pushing buffered events on shutdown, and how long the
# sidecar keeps serving its open connections after SIGTERM
drainTimeoutSeconds: 10
command:
- python
- -m
- litellm.proxy.collector
# Sized independently of the proxy container; the pipeline is CPU bound
resources: {}
# requests:
# cpu: 500m
# memory: 1Gi
# limits:
# cpu: "1"
# memory: 2Gi
# When autoscaling.enabled, swap the pod-wide cpu Resource metric for an
# autoscaling/v2 ContainerResource metric on the proxy container only, so the
# sidecar's CPU never scales inference replicas. Needs Kubernetes 1.30+ (or the
# HPAContainerMetrics feature gate on 1.27 to 1.29)
scaleOnProxyContainerCpu: false
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an
@ -212,6 +264,25 @@ autoscaling:
# Memory is a floor to provision under 'resources', not a signal to scale on.
# targetMemoryUtilizationPercentage: 80
# behavior: {}
# Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
# named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
# AverageValue target, alongside whichever resource targets are set (the HPA
# follows the metric asking for the most replicas). A Prometheus Adapter must
# serve those two names on custom.metrics.k8s.io from the proxy's counters,
# grouped by the scrape target's `pod` label (enable serviceMonitor below so
# every pod is scraped on its own):
# litellm_requests_per_second:
# sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# litellm_tokens_per_second:
# sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# rate() over [1m] is already per second, so no `* 60`. How fast the HPA
# reacts is set by that window, the scrape interval and the HPA sync period
# (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
# faster so a 1m window holds at least 4 samples. averageValue takes SI
# suffixes, so "6M" is six million tokens per second per pod. Tokens are
# counted when a response completes, so TPS trails long streams.
targetRequestsPerSecond: ""
targetTokensPerSecond: ""
# Autoscaling with keda is mutually exclusive with hpa
keda:
@ -233,6 +304,23 @@ keda:
# metricName: http_requests_total
# threshold: '100'
# query: sum(rate(http_requests_total{deployment="my-deployment"}[2m]))
# First-class Prometheus triggers on the proxy's own request and token
# counters, appended to `triggers`. Each target is the per-second load one
# replica should carry: KEDA divides the release-wide
# `sum(rate(<counter>[1m]))` by it to pick the replica count. Thresholds
# are plain numbers (KEDA parses them as floats, no SI suffixes). The
# queries select samples by the release namespace and the `job` label the
# chart's ServiceMonitor produces (the metrics Service name), so enable
# serviceMonitor below together with metricsServer: the http port serves
# /metrics/ behind virtual-key auth and answers an unauthenticated scrape
# with 401. Reaction time comes from the [1m] window, the scrape interval
# and pollingInterval above, so keep both at 15s or faster. Tokens are
# counted at completion, so TPS trails long streams. serverAddress is
# required once either target is set.
prometheus:
serverAddress: ""
requestsPerSecond: ""
tokensPerSecond: ""
behavior: {}
# scaleDown:
# stabilizationWindowSeconds: 300
@ -309,6 +397,20 @@ db:
# only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
readReplicaUrl: ""
# In-container connection pool (PgBouncer, transaction mode) shared by every
# worker in the pod. Without it each --num_workers worker opens its own
# connection_limit connections to Postgres, so a pod's footprint against the
# database's connection ceiling is workers x connection_limit and grows with
# every replica. With it, the pod holds at most maxDbConnections upstream
# connections no matter how many workers run; the workers connect to the pool
# over loopback, with no extra network hop. Migrations still go straight to
# Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so
# a database with a 5000-connection ceiling fits roughly 200 replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
# The Stackgres Operator must already be installed within the target
# Kubernetes cluster.

55
helm/litellm/templates/_helpers.tpl
View file

@ -257,6 +257,14 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
- name: DATABASE_SCHEMA
value: {{ .schema | quote }}
{{- end }}
{{- if .sslMode }}
- name: DATABASE_SSLMODE
value: {{ .sslMode | quote }}
{{- end }}
{{- if .sslRootCert }}
- name: DATABASE_SSLROOTCERT
value: {{ .sslRootCert | quote }}
{{- end }}
{{- if and .useIAMAuth .useAzureEntraAuth }}
{{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }}
{{- end }}
@ -360,6 +368,20 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end -}}
{{/*
In-container PgBouncer env for the gateway container. Under IAM or Entra auth the pooler mints and renews the database token itself.
*/}}
{{- define "litellm.connectionPoolEnv" -}}
{{- with .Values.database.connectionPool -}}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }}
- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }}
{{- end }}
{{- end -}}
{{/*
PodDisruptionBudget shared by gateway, backend, and ui.
@ -441,3 +463,36 @@ ImplementationSpecific
{{- .pathType -}}
{{- end -}}
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}
{{/*
Directory of the collector's unix socket, shared by the gateway and
collector containers through an emptyDir. Empty when the sidecar is off
or gateway.collector.address is a tcp://127.0.0.1:<port> address.
*/}}
{{- define "litellm.gateway.collectorSocketDir" -}}
{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}}
{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}}
{{- end -}}
{{- end -}}
{{/*
LITELLM_COLLECTOR_* env shared by the producer (gateway container) and the
consumer (collector container), so both agree on the transport and the
shutdown drain window.
*/}}
{{- define "litellm.gateway.collectorEnv" -}}
{{- with .Values.gateway.collector }}
- name: LITELLM_COLLECTOR_ENABLED
value: "true"
- name: LITELLM_COLLECTOR_ADDRESS
value: {{ .address | quote }}
- name: LITELLM_COLLECTOR_BUFFER_SIZE
value: {{ .bufferSize | quote }}
- name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: {{ .onUnavailable | quote }}
- name: LITELLM_COLLECTOR_DRAIN_TIMEOUT_SECONDS
value: {{ .drainTimeoutSeconds | quote }}
{{- end }}
{{- end -}}

111
helm/litellm/templates/gateway/deployment.yaml
View file

@ -61,17 +61,38 @@ spec:
- name: NUM_WORKERS
value: {{ .Values.gateway.numWorkers | quote }}
{{- end }}
{{- if .Values.database.connectionPool.enabled }}
{{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
{{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
{{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
{{- end }}
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.gateway.collector.enabled }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -97,16 +118,102 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: metrics
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- {{ .Values.gateway.metricsServer.port | quote }}
env:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
ports:
- name: metrics
containerPort: {{ .Values.gateway.metricsServer.port }}
protocol: TCP
volumeMounts:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
readinessProbe:
tcpSocket: { port: metrics }
periodSeconds: 10
livenessProbe:
tcpSocket: { port: metrics }
periodSeconds: 15
failureThreshold: 6
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- if .Values.gateway.collector.enabled }}
- name: collector
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.collector
env:
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- if .Values.database.connectionPool.enabled }}
{{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
{{- end }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts (include "litellm.gateway.collectorSocketDir" .) }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
mountPath: {{ include "litellm.gateway.collectorSocketDir" . }}
{{- end }}
{{- with .Values.gateway.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- end }}
resources:
{{- toYaml .Values.gateway.collector.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled (include "litellm.gateway.collectorSocketDir" .) }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if include "litellm.gateway.collectorSocketDir" . }}
- name: collector-socket
emptyDir:
sizeLimit: 1Mi
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

28
helm/litellm/templates/gateway/hpa.yaml
View file

@ -15,6 +15,15 @@ spec:
maxReplicas: {{ .Values.gateway.hpa.maxReplicas }}
metrics:
{{- if .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- if and .Values.gateway.collector.enabled .Values.gateway.collector.scaleOnGatewayContainerCpu }}
- type: ContainerResource
containerResource:
name: cpu
container: gateway
target:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- else }}
- type: Resource
resource:
name: cpu
@ -22,6 +31,7 @@ spec:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetCPUUtilizationPercentage }}
{{- end }}
{{- end }}
{{- if .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
- type: Resource
resource:
@ -30,6 +40,24 @@ spec:
type: Utilization
averageUtilization: {{ .Values.gateway.hpa.targetMemoryUtilizationPercentage }}
{{- end }}
{{- with .Values.gateway.hpa.targetRequestsPerSecond }}
- type: Pods
pods:
metric:
name: litellm_requests_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- with .Values.gateway.hpa.targetTokensPerSecond }}
- type: Pods
pods:
metric:
name: litellm_tokens_per_second
target:
type: AverageValue
averageValue: {{ toJson . | trimAll "\"" | quote }}
{{- end }}
{{- with .Values.gateway.hpa.behavior }}
behavior:
{{- toYaml . | nindent 4 }}

18
helm/litellm/templates/gateway/service-metrics.yaml Normal file
View file

@ -0,0 +1,18 @@
{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.gateway.fullname" . }}-metrics
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
spec:
type: ClusterIP
ports:
- port: {{ .Values.gateway.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
{{- end }}

28
helm/litellm/templates/gateway/servicemonitor.yaml Normal file
View file

@ -0,0 +1,28 @@
{{- if and .Values.gateway.enabled .Values.gateway.serviceMonitor.enabled }}
{{- if not .Values.gateway.metricsServer.enabled }}
{{- fail "gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled: the http port serves /metrics/ behind virtual-key auth, so an unauthenticated scrape gets 401" }}
{{- end }}
apiVersion: monitoring.coreos.com/v1
kind: ServiceMonitor
metadata:
name: {{ include "litellm.gateway.fullname" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
{{- with .Values.gateway.serviceMonitor.labels }}
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
selector:
matchLabels:
{{- include "litellm.gateway.selectorLabels" . | nindent 6 }}
namespaceSelector:
matchNames:
- {{ .Release.Namespace | quote }}
endpoints:
- port: metrics
path: /metrics/
interval: {{ .Values.gateway.serviceMonitor.interval }}
scrapeTimeout: {{ .Values.gateway.serviceMonitor.scrapeTimeout }}
scheme: http
{{- end }}

11
helm/litellm/templates/ingress.yaml
View file

@ -89,7 +89,7 @@
at "/" Prefix would swallow the whole backend management API) instead of
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}}
{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -129,6 +129,8 @@ spec:
# --- Gateway data plane ---
# Exact /test only (see the $gatewayPrefixes comment above);
# /test/* MCP management endpoints fall to the backend catch-all.
# Exact /debug/memory/summary reads a serving worker's RSS (the e2e memory
# gate); the rest of /debug/* stays on the backend.
- path: /test
pathType: Exact
backend:
@ -136,6 +138,13 @@ spec:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
- path: /debug/memory/summary
pathType: Exact
backend:
service:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
{{- range $gatewayPrefixes }}
{{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" . "pathType" "Prefix") }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" . $pathType) }}

216
helm/litellm/tests/collector_tests.yaml Normal file
View file

@ -0,0 +1,216 @@
suite: test gateway collector sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/hpa.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, env, volume or container metric when the collector is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ENABLED
value: "true"
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0].type
value: Resource
template: gateway/hpa.yaml
- it: runs the collector as a sidecar sharing env, config, the pod pool and a unix socket emptyDir, and scales on the gateway container only
set:
gateway.collector.enabled: true
gateway.collector.bufferSize: 250
gateway.collector.onUnavailable: drop
gateway.image.tag: v1.102.0
gateway.numWorkers: 4
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
gateway.envSecrets:
- litellm-license
gateway.volumes:
- name: redis-ca
secret:
secretName: redis-ca
gateway.volumeMounts:
- name: redis-ca
mountPath: /etc/litellm/redis-ca
readOnly: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_BUFFER_SIZE
value: "250"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_COLLECTOR_ON_UNAVAILABLE
value: drop
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: collector
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.102.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_JOB_ROLE
value: collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: CONFIG_FILE_PATH
value: /app/config/config.yaml
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_HOST
value: postgres.example.com
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: unix:///var/run/litellm/collector.sock
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[1].env
content:
name: NUM_WORKERS
any: true
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].envFrom
value:
- secretRef:
name: litellm-license
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
mountPath: /var/run/litellm
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: redis-ca
mountPath: /etc/litellm/redis-ca
readOnly: true
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.limits.cpu
value: "1"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: collector-socket
emptyDir:
sizeLimit: 1Mi
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0]
value:
type: ContainerResource
containerResource:
name: cpu
container: gateway
target:
type: Utilization
averageUtilization: 70
template: gateway/hpa.yaml
- it: uses loopback tcp without a socket volume and keeps the pod-wide cpu metric when asked
set:
gateway.collector.enabled: true
gateway.collector.address: tcp://127.0.0.1:4010
gateway.collector.scaleOnGatewayContainerCpu: false
asserts:
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_COLLECTOR_ADDRESS
value: tcp://127.0.0.1:4010
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[1].volumeMounts
content:
name: collector-socket
any: true
template: gateway/deployment.yaml
- equal:
path: spec.metrics[0].type
value: Resource
template: gateway/hpa.yaml

204
helm/litellm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,204 @@
suite: test in-container connection pool env vars
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
values:
- ./values/required.yaml
tests:
- it: renders no pool env by default
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true
- it: enabled pool renders the three pgbouncer vars with the configured sizes
template: gateway/deployment.yaml
set:
gateway.numWorkers: 4
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
- contains:
path: spec.template.spec.containers[0].env
content:
name: NUM_WORKERS
value: "4"
- it: enabled pool uses the chart default sizes
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "1000"
- it: backend never gets the pool env
template: backend/deployment.yaml
set:
database.connectionPool.enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- it: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
gateway.metricsServer.enabled: true
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: metrics
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- equal:
path: spec.template.spec.containers[2].name
value: collector
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
- it: collector sidecar gets no pool env when the pool is off
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: collector
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true
- it: pool with IAM auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
database.connectionPool.enabled: true
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: pool with Entra auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useAzureEntraAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
- it: IAM auth without the pool still renders
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true

77
helm/litellm/tests/database_auth_tests.yaml
View file

@ -4,6 +4,7 @@ templates:
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
@ -67,6 +68,82 @@ tests:
value: "true"
any: true
- it: emits no TLS env by default
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
any: true
- it: writer sslMode and sslRootCert reach gateway and backend as DATABASE_SSLMODE and DATABASE_SSLROOTCERT
templates:
- gateway/deployment.yaml
- backend/deployment.yaml
set:
database.writer.useIAMAuth: true
database.writer.sslMode: verify-full
database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
value: verify-full
any: true
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
value: /etc/ssl/certs/ca-certificates.crt
any: true
- it: writer sslMode and sslRootCert reach the collector sidecar and the migrations job, which dial Postgres themselves
set:
gateway.collector.enabled: true
database.connectionPool.enabled: true
database.writer.sslMode: verify-full
database.writer.sslRootCert: /etc/ssl/certs/ca-certificates.crt
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: collector
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_SSLMODE
value: verify-full
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_SSLROOTCERT
value: /etc/ssl/certs/ca-certificates.crt
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
value: verify-full
any: true
template: migrations-job.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
value: /etc/ssl/certs/ca-certificates.crt
any: true
template: migrations-job.yaml
- it: writer rejects both token sources at once
template: gateway/deployment.yaml
set:

213
helm/litellm/tests/hpa_workload_metrics_tests.yaml Normal file
View file

@ -0,0 +1,213 @@
suite: test gateway HPA per-pod requests-per-second and tokens-per-second targets
templates:
- gateway/hpa.yaml
- gateway/servicemonitor.yaml
values:
- ./values/required.yaml
tests:
- it: scales on CPU and memory only by default
template: gateway/hpa.yaml
asserts:
- equal:
path: spec.metrics
value:
- type: Resource
resource:
name: cpu
target:
type: Utilization
averageUtilization: 70
- type: Resource
resource:
name: memory
target:
type: Utilization
averageUtilization: 80
- it: adds a requests-per-second Pods metric next to the resource metrics
template: gateway/hpa.yaml
set:
gateway.hpa.targetRequestsPerSecond: 90
asserts:
- lengthEqual:
path: spec.metrics
count: 3
- contains:
path: spec.metrics
content:
type: Resource
resource:
name: cpu
target:
type: Utilization
averageUtilization: 70
- equal:
path: spec.metrics[2]
value:
type: Pods
pods:
metric:
name: litellm_requests_per_second
target:
type: AverageValue
averageValue: "90"
- notContains:
path: spec.metrics
content:
type: Pods
pods:
metric:
name: litellm_tokens_per_second
any: true
- it: adds a tokens-per-second Pods metric on its own
template: gateway/hpa.yaml
set:
gateway.hpa.targetTokensPerSecond: 6M
asserts:
- lengthEqual:
path: spec.metrics
count: 3
- equal:
path: spec.metrics[2]
value:
type: Pods
pods:
metric:
name: litellm_tokens_per_second
target:
type: AverageValue
averageValue: "6M"
- notContains:
path: spec.metrics
content:
type: Pods
pods:
metric:
name: litellm_requests_per_second
any: true
- it: renders requests and tokens targets together and keeps CPU and memory
template: gateway/hpa.yaml
set:
gateway.hpa.targetRequestsPerSecond: 90
gateway.hpa.targetTokensPerSecond: 6000000
asserts:
- lengthEqual:
path: spec.metrics
count: 4
- equal:
path: spec.metrics[0].resource.name
value: cpu
- equal:
path: spec.metrics[1].resource.name
value: memory
- equal:
path: spec.metrics[2].pods.metric.name
value: litellm_requests_per_second
- equal:
path: spec.metrics[3].pods.metric.name
value: litellm_tokens_per_second
- equal:
path: spec.metrics[3].pods.target.averageValue
value: "6000000"
- it: scales on workload metrics alone when the resource targets are cleared
template: gateway/hpa.yaml
set:
gateway.hpa.targetCPUUtilizationPercentage: null
gateway.hpa.targetMemoryUtilizationPercentage: null
gateway.hpa.targetRequestsPerSecond: 90
gateway.hpa.targetTokensPerSecond: 6000000
asserts:
- lengthEqual:
path: spec.metrics
count: 2
- notContains:
path: spec.metrics
content:
type: Resource
any: true
- equal:
path: spec.metrics[0].pods.metric.name
value: litellm_requests_per_second
- equal:
path: spec.metrics[1].pods.metric.name
value: litellm_tokens_per_second
- notMatchRegexRaw:
pattern: per_minute
- it: ignores the per-minute keys, which the chart never shipped
template: gateway/hpa.yaml
set:
gateway.hpa.targetRequestsPerMinute: 5400
gateway.hpa.targetTokensPerMinute: 360000000
asserts:
- lengthEqual:
path: spec.metrics
count: 2
- notContains:
path: spec.metrics
content:
type: Pods
any: true
- it: renders no ServiceMonitor by default
template: gateway/servicemonitor.yaml
asserts:
- hasDocuments:
count: 0
- it: refuses a ServiceMonitor without the metrics server, whose http port needs a bearer token
template: gateway/servicemonitor.yaml
set:
gateway.serviceMonitor.enabled: true
asserts:
- failedTemplate:
errorPattern: gateway.serviceMonitor.enabled requires gateway.metricsServer.enabled
- it: scrapes each gateway pod through the metrics port
template: gateway/servicemonitor.yaml
release:
name: rel
namespace: llm
set:
gateway.serviceMonitor.enabled: true
gateway.metricsServer.enabled: true
gateway.serviceMonitor.labels:
release: kube-prometheus-stack
asserts:
- isKind:
of: ServiceMonitor
- equal:
path: metadata.labels.release
value: kube-prometheus-stack
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: rel
app.kubernetes.io/component: gateway
- equal:
path: spec.namespaceSelector.matchNames
value:
- llm
- equal:
path: spec.endpoints
value:
- port: metrics
path: /metrics/
interval: 15s
scrapeTimeout: 10s
scheme: http
- it: honours a custom scrape interval
template: gateway/servicemonitor.yaml
set:
gateway.serviceMonitor.enabled: true
gateway.serviceMonitor.interval: 30s
gateway.metricsServer.enabled: true
asserts:
- equal:
path: spec.endpoints[0].interval
value: 30s

10
helm/litellm/tests/ingress_controller_tests.yaml
View file

@ -97,6 +97,16 @@ tests:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /debug/memory/summary
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- equal:
path: spec.rules[0].http.paths[-1]
value:

11
helm/litellm/tests/ingress_extra_paths_tests.yaml
View file

@ -288,6 +288,17 @@ tests:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /test with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: rejects an entry that would take over the exact /debug/memory/summary route
set:
ingress.enabled: true
ingress.extraPaths:
- path: /debug/memory/summary
pathType: Exact
service: backend
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /debug/memory/summary with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: allows a built-in path under a different pathType, which is a distinct rule
set:
ingress.enabled: true

148
helm/litellm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,148 @@
suite: test gateway metrics sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/service.yaml
- gateway/service-metrics.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, volume, env or service port when the metrics server is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
any: true
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- hasDocuments:
count: 0
template: gateway/service-metrics.yaml
- it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4101
gateway.service.type: LoadBalancer
gateway.image.tag: v1.101.0
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.101.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- "4101"
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].env
value:
- name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].ports
value:
- name: metrics
containerPort: 4101
protocol: TCP
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].volumeMounts
value:
- name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 50m
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
emptyDir: {}
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: gateway/service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway-metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: gateway/service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
template: gateway/service-metrics.yaml
- it: rejects a metrics port equal to the gateway port
template: gateway/deployment.yaml
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: gateway.metricsServer.port must differ from the gateway port 4000

106
helm/litellm/values.yaml
View file

@ -208,6 +208,11 @@ database:
name: litellm-writer-secret
usernameKey: username
passwordKey: password
# libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
# in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
# sslRootCert on its own implies sslMode verify-full
sslMode: ""
sslRootCert: ""
# Optional read-replica routing. When `reader.host` is set, the proxy routes
# reads (find_*, count, group_by, query_raw/_first) to this endpoint while
@ -225,6 +230,26 @@ database:
usernameKey: username
passwordKey: password
# In-container connection pool (PgBouncer, transaction mode) shared by every
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
# workers opens its own Prisma pool straight to Postgres, so a pod's
# footprint against the database's connection ceiling is
# numWorkers x connection_limit and grows with every replica. With it, the
# pod holds at most maxDbConnections upstream connections no matter how many
# workers run; the workers connect to the pool over loopback, with no extra
# network hop. The chart emits LITELLM_PGBOUNCER_ENABLED /
# LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on
# the gateway container and its collector sidecar only: the backend runs a
# single worker and the migrations Job must keep a direct connection. With
# `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and
# renews the database token itself, so the workers never see it. Starting profile for
# `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
# 5000-connection ceiling fits roughly 200 gateway replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Optional Redis. Leave host empty to disable.
#
# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend
@ -268,6 +293,68 @@ gateway:
config:
create: true
proxy_config: {}
# Serve Prometheus /metrics from a `metrics` sidecar container (same image,
# `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the
# workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a
# scrape never runs on an inference worker. Adds a `metrics` port to the pod
# and a dedicated ClusterIP `<gateway>-metrics` Service; point your scrape
# config at it. The port has no virtual-key auth: keep it off public ingress.
# Needs the gateway image v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
requests:
cpu: 50m
memory: 128Mi
limits:
memory: 512Mi
# Prometheus Operator ServiceMonitor for the gateway pods. Scrapes the
# `<gateway>-metrics` Service, so it requires metricsServer above (the http
# port serves /metrics/ behind virtual-key auth). Every pod is its own scrape
# target, so the samples carry the `pod` label the per-pod autoscaling
# queries below group by.
serviceMonitor:
enabled: false
labels: {}
interval: 15s
scrapeTimeout: 10s
# Opt-in `collector` sidecar (same image, `python -m litellm.proxy.collector`)
# that runs the post-response spend pipeline (cost calculation, spend logs,
# spend counters, budget reservation reconciliation) so the uvicorn workers
# only serialise a compact event over loopback and go back to serving
# requests. It shares the pod's env, proxy config, in-container pgbouncer and
# Redis spend buffer, so the per-pod DB connection budget is unchanged.
# Delivery is at-most-once inside the pod: events already handed over are
# lost if the sidecar dies before writing them; events the workers cannot
# hand over follow `onUnavailable`.
collector:
enabled: false
# unix:///<dir>/<file>.sock (the <dir> becomes a shared emptyDir) or
# tcp://127.0.0.1:<port>
address: unix:///var/run/litellm/collector.sock
# Events each uvicorn worker holds in memory while the sidecar is slow or
# restarting.
bufferSize: 1000
# fallback: run the pipeline in the worker when the sidecar is unreachable
# or the buffer is full (spend stays exact, that request costs gateway CPU
# again). drop: count and discard the event instead (spend under-reports).
onUnavailable: fallback
# How long the workers keep pushing buffered events on shutdown, and how
# long the sidecar keeps serving open connections after SIGTERM.
drainTimeoutSeconds: 10
# Sized independently of the gateway container; the pipeline is CPU bound.
resources:
requests:
cpu: 500m
memory: 1Gi
limits:
cpu: "1"
memory: 2Gi
# With hpa.targetCPUUtilizationPercentage set, scale on an autoscaling/v2
# ContainerResource metric of the `gateway` container only, so the
# sidecar's CPU never drives inference replicas. Needs Kubernetes 1.30+.
scaleOnGatewayContainerCpu: true
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion
@ -324,6 +411,25 @@ gateway:
# policies:
# - { type: Percent, value: 100, periodSeconds: 30 }
behavior: {}
# Opt-in per-pod workload targets, rendered as autoscaling/v2 `Pods` metrics
# named `litellm_requests_per_second` and `litellm_tokens_per_second` with an
# AverageValue target. They coexist with the CPU/memory targets above: the
# HPA scales on whichever metric asks for the most replicas. Kubernetes has
# no idea what a token is, so a Prometheus Adapter must serve those two
# names on custom.metrics.k8s.io from the proxy's counters, grouped by the
# scrape target's `pod` label (enable serviceMonitor above):
# litellm_requests_per_second:
# sum(rate(litellm_proxy_total_requests_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# litellm_tokens_per_second:
# sum(rate(litellm_total_tokens_metric_total{<<.LabelMatchers>>}[1m])) by (<<.GroupBy>>)
# rate() over [1m] is already per second, so no `* 60`. How fast the HPA
# reacts is set by that window, the scrape interval and the HPA sync period
# (15s by default), not by the unit: keep serviceMonitor.interval at 15s or
# faster so a 1m window holds at least 4 samples. averageValue takes SI
# suffixes, so "6M" is six million tokens per second per pod. Tokens are
# counted when a response completes, so TPS trails long streams.
targetRequestsPerSecond: ""
targetTokensPerSecond: ""
# PodDisruptionBudget for the gateway pods. Set exactly one of
# `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
# enabling without either falls back to `maxUnavailable: 1`). Disabled by

15
litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql Normal file
View file

@ -0,0 +1,15 @@
-- DropForeignKey
DO $$
BEGIN
IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
END IF;
END $$;
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END $$;

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql Normal file
View file

@ -0,0 +1,4 @@
-- Add org_id column to LiteLLM_ManagedObjectTable
-- Snapshots the creating key's organization at submission time, like team_id,
-- so CheckBatchCost can bill organization spend hours later without re-resolving
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT;

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_add_autorouter_session_baseline_models/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "baseline_models" JSONB NOT NULL DEFAULT '{}';

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260910000000_skills_on_object_permission/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "skills" TEXT[] DEFAULT ARRAY[]::TEXT[];

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

@ -282,6 +282,7 @@ model LiteLLM_ObjectPermissionTable {
mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user
search_tools String[] @default([]) // search_tool_name values this key/team/user may call
mcp_tool_search_enabled Boolean?
skills String[] @default([]) // Claude Code plugin names granted to this key/team beyond the public (enabled) set
teams LiteLLM_TeamTable[]
projects LiteLLM_ProjectTable[]
verification_tokens LiteLLM_VerificationToken[]
@ -492,7 +493,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@ -1036,6 +1037,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
created_at DateTime @default(now())
created_by String?
team_id String?
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
updated_at DateTime @updatedAt
@ -1512,6 +1514,7 @@ model LiteLLM_AutoRouterSession {
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
baseline_models Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")

130
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -8,14 +8,10 @@ import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
from typing import TYPE_CHECKING, Final, Optional
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
prisma_command_timeout,
prisma_migrate_deploy_timeout,
)
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
if TYPE_CHECKING:
import psycopg
import psycopg.sql
def str_to_bool(value: Optional[str]) -> bool:
@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
_INVALID_LITELLM_INDEXES_SQL: Final = (
"SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
"FROM pg_index i "
"JOIN pg_class c ON c.oid = i.indexrelid "
"JOIN pg_class t ON t.oid = i.indrelid "
"JOIN pg_namespace n ON n.oid = t.relnamespace "
"WHERE NOT i.indisvalid "
" AND c.relkind = 'i' "
" AND n.nspname = %s "
" AND t.relname LIKE %s "
" AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
"ORDER BY c.relname"
)
@dataclass(frozen=True, slots=True)
class _InvalidIndex:
schema: str
name: str
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
schema, etc.) from DATABASE_URL so psycopg can parse it."""
from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
parsed = urlparse(url)
if not parsed.query:
@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
"target_session_attrs",
}
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
return urlunparse(parsed._replace(query=urlencode(kept)))
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
)
@staticmethod
def _invalid_litellm_indexes(
conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
) -> tuple[_InvalidIndex, ...]:
rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
return tuple(_InvalidIndex(*row) for row in rows)
@staticmethod
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
from psycopg import sql
target: Final = sql.Identifier(index.schema, index.name)
if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
@staticmethod
def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
import psycopg
statement, action = ProxyExtrasDBManager._index_repair(index)
try:
conn.execute(statement)
except psycopg.Error as e:
logger.warning(
"Could not repair invalid index %s.%s, will retry on the next startup. "
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
index.schema,
index.name,
statement.as_string(conn),
e,
)
return
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
@staticmethod
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
INVALID (a migration deadlock between replicas is the usual cause; the
retried migration skips them because of IF NOT EXISTS). Never raises:
returns True when no invalid index remains, False when the repair was
skipped or failed and will be retried on the next startup. Looks in the
schema DATABASE_URL names, the only URL Prisma migrates through, but
connects over DIRECT_URL when set: the session settings, the advisory
lock and REINDEX CONCURRENTLY all need one server session, which a
transaction pooler does not give."""
prisma_url: Final = os.getenv("DATABASE_URL")
if not prisma_url:
return False
try:
import psycopg
from psycopg import sql
except ImportError:
logger.warning(
"psycopg is not installed; skipping the invalid index check. "
"Install the litellm[extra_proxy] extra, which includes psycopg."
)
return False
schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public"
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url)
try:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
conn.execute("SET statement_timeout = 0")
conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
if not found:
return True
logger.warning(
"Found %d invalid index(es) left by an interrupted CREATE INDEX "
"CONCURRENTLY, rebuilding: %s",
len(found),
", ".join(f"{index.name} (table size {index.table_size})" for index in found),
)
lock_row: Final = conn.execute(
"SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
).fetchone()
if lock_row is None or not lock_row[0]:
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
return False
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
ProxyExtrasDBManager._repair_index(conn, index)
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
except psycopg.Error as e:
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
return False
return not remaining
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
"""
@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.94"
version = "0.4.96"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.94"
version = "0.4.96"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

5
litellm-rust/AGENTS.md
View file

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

1176
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

10
litellm-rust/Cargo.toml
View file

@ -1,6 +1,7 @@
[workspace]
members = [
"crates/core",
"crates/token-counter",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
@ -18,6 +19,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-token-counter = { path = "crates/token-counter" }
litellm-config = { path = "crates/config" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }
@ -39,6 +41,14 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
gcp_auth = "0.12.7"
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
criterion = "0.8.2"
veil = "0.3.0"
[profile.release]
opt-level = 3

41
litellm-rust/crates/ai-gateway/Dockerfile
View file

@ -14,15 +14,20 @@
# ---- Chef -------------------------------------------------------------------
# cargo-chef caches the dependency build so only the gateway crate recompiles on
# a source-only change. python3-dev is present in every rust stage because the
# `python-config` feature links libpython via pyo3 (even in the cook step).
FROM rust:1.90-slim-bookworm AS chef
# `python-config` feature links libpython via pyo3 (even in the cook step), and
# python3-pip builds the litellm wheel in the builder stage.
FROM rust:1.98-slim-bookworm AS chef
ENV PYO3_PYTHON=python3.11
# rustup reads rust-toolchain.toml from any parent of the working directory, so
# copying it in is what keeps every cargo call below on the repo's pinned
# channel rather than on whatever the base image happens to ship.
COPY rust-toolchain.toml /build/rust-toolchain.toml
WORKDIR /build/litellm-rust
RUN apt-get update \
&& apt-get install -y --no-install-recommends \
python3 python3-dev pkg-config libssl-dev clang \
python3 python3-dev python3-pip pkg-config libssl-dev clang \
&& rm -rf /var/lib/apt/lists/* \
&& cargo install cargo-chef --locked --version 0.1.77
WORKDIR /build/litellm-rust
# ---- Planner ----------------------------------------------------------------
# Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
@ -43,6 +48,19 @@ RUN cargo chef cook --locked --release \
COPY litellm-rust/ .
RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
# The root pyproject builds with maturin against litellm-rust/crates/python-bridge,
# so the wheel is built here, next to the crate sources and the cargo toolchain,
# and the runtime stage installs the artifact instead of compiling anything.
# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions
# in this repo, and those hit PyPI hours after every version bump merges, so both
# wheels are built from the repo too instead of being resolved from PyPI.
COPY pyproject.toml README.md LICENSE /build/
COPY litellm/ /build/litellm/
COPY enterprise/ /build/enterprise/
COPY litellm-proxy-extras/ /build/litellm-proxy-extras/
RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \
/build /build/enterprise /build/litellm-proxy-extras
# ---- Runtime ----------------------------------------------------------------
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
# 3.11 ABI so the embedded interpreter links and imports cleanly.
@ -56,11 +74,16 @@ RUN apt-get update \
WORKDIR /app
# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
# package + packaging metadata, then pip install the proxy extra.
COPY pyproject.toml README.md LICENSE ./
COPY litellm/ ./litellm/
RUN pip install --no-cache-dir ".[proxy]"
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two
# sibling wheels come from the builder as well, so the pins in litellm[proxy]
# resolve against them and never wait on a PyPI publish.
COPY --from=builder /build/dist/*.whl /tmp/wheels/
RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \
&& pip install --no-cache-dir \
/tmp/wheels/litellm_enterprise-*.whl \
/tmp/wheels/litellm_proxy_extras-*.whl \
"${wheel}[proxy]" \
&& rm -rf /tmp/wheels
# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
# only).

11
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
View file

@ -9,19 +9,28 @@
# Strategy: ignore everything, then re-include only what the build needs:
# - litellm/ (pip install . needs the full package + proxy reader)
# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources)
# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install)
# - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it)
# - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy])
# - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build)
# - rust-toolchain.toml (the pinned channel every cargo call in the build uses)
*
# --- re-include the build inputs ---
!litellm/
!litellm-rust/
!enterprise/
!litellm-proxy-extras/
!pyproject.toml
!rust-toolchain.toml
!README.md
!LICENSE
# --- prune heavy / irrelevant subpaths back out of the re-included trees ---
# Rust build artifacts (huge; regenerated in the builder).
**/target/
# Committed python distribution artifacts; the wheel build does not read them.
enterprise/dist/
litellm-proxy-extras/dist/
# Python caches and compiled bytecode.
**/__pycache__/
**/*.pyc

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

@ -6,17 +6,18 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
## Crates
`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route:
`litellm-rust` has six 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-token-counter | Standalone input token counting shared by host integrations without pulling in the full SDK. |
| 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: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers, token counter, and Python interop.
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)

6
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -269,7 +269,11 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) => "AuthError",
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",

6
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -15,6 +15,8 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::auth::error::MissingCredential;
use litellm_core::error::Error;
use litellm_core::realtime::transformation::RealtimeProviderConfig;
use litellm_core::realtime::types::RealtimeEvent;
@ -32,8 +34,6 @@ use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
/// Default **idle** timeout: if neither side sends a frame for this long, the
/// session is reaped. It resets on any activity, so it does not cap a healthy
/// (continuously streaming) session — it only frees a stalled one (e.g. a
@ -59,7 +59,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.ok()
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey)))
}
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.

6
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -4,7 +4,9 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::Error;
use litellm_core::auth::error::MissingCredential;
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
@ -23,8 +25,6 @@ use crate::constants::{
};
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
@ -120,7 +120,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey)))
}
async fn dial_upstream(

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

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

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

@ -1,84 +0,0 @@
use litellm_core::error::Error;
use litellm_core::http_utils::http_request;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::Value;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::hooks::OcrLifecycleHooks;
use super::types::PreparedOcrRequest;
use crate::client::http_client;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn execute_ocr_provider_call(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
) -> Result<Value, Error> {
let request = hooks.prepare_provider_request(request).await?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|err| Error::Network(err.to_string()))?;
let status = response.status();
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
&& status.as_u16() == 202
{
let operation_url = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.ok_or_else(|| {
Error::InvalidResponse(
"Azure Document Intelligence returned 202 but no Operation-Location header found"
.to_string(),
)
})?;
let response_json = poll_document_intelligence(
&operation_url,
&request.url,
&request.upstream_headers,
request.timeout,
)
.await?;
return Ok(request
.config
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json());
}
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text)
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
Ok(request
.config
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

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

@ -1,401 +0,0 @@
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, truncate_error_body};
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::client::http_client;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
use crate::integrations::custom_logger::{
CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::{
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
};
pub(crate) struct OcrLifecycleHooks {
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl OcrLifecycleHooks {
pub(crate) fn new(
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
) -> Self {
Self {
logger_runner,
guardrail_runner,
request_metadata,
}
}
async fn run_pre_call_guardrails(
&self,
request: PreparedOcrRequest,
) -> Result<PreparedOcrRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let context = guardrail_context(&self.request_metadata);
let guardrail_request = GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": request.custom_llm_provider,
"document": request.document,
"optional_params": request.optional_params,
}));
let (guardrail_request, _) = self
.guardrail_runner
.run_pre_call(&context, guardrail_request)
.await
.map_err(guardrail_error_to_core_error)?;
let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
let optional_params = match &request.config {
Ok(config) => config.map_ocr_params(&optional_params),
Err(_) => optional_params,
};
Ok(PreparedOcrRequest {
document,
optional_params,
..request
})
}
pub(crate) async fn prepare_provider_request(
&self,
request: PreparedOcrRequest,
) -> Result<ProviderOcrRequest, Error> {
let config = request.config?;
let env_lookup = |key: &str| std::env::var(key).ok();
let upstream_headers = config.validate_environment(
string_headers(request.extra_headers)?,
request.api_key.as_deref(),
&env_lookup,
)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
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, optional_params.clone())?
.data;
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,
})
}
async fn run_during_call_guardrails(
&self,
model: &str,
custom_llm_provider: &str,
url: &str,
body: Value,
) -> Result<Value, Error> {
if self.guardrail_runner.is_empty() {
return Ok(body);
}
let context = guardrail_context(&self.request_metadata);
let guardrail_request = GuardrailRequest::new(json!({
"model": model,
"custom_llm_provider": custom_llm_provider,
"url": url,
"body": body,
}));
let (guardrail_request, _) = self
.guardrail_runner
.run_during_call(&context, guardrail_request)
.await
.map_err(guardrail_error_to_core_error)?;
parse_ocr_during_call_guardrail_request(guardrail_request)
}
fn standard_logging_payload(
&self,
context: &CallLifecycleContext,
timing: &CallLifecycleTiming,
) -> StandardLoggingPayload {
StandardLoggingPayload {
id: context.litellm_call_id.clone(),
litellm_call_id: context.litellm_call_id.clone(),
call_type: context.call_type.clone(),
model: context.model.clone(),
custom_llm_provider: context.custom_llm_provider.clone(),
response_cost: 0.0,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
start_time: timing.start_time,
end_time: timing.end_time,
stream: false,
metadata: StandardLoggingMetadata {
user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
..Default::default()
},
messages: None,
}
}
}
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>;
type SuccessFuture<'a> = OcrLogFuture<'a>;
type FailureFuture<'a> = OcrLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move { self.run_pre_call_guardrails(request).await })
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a Value,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let response_obj = CallbackValue::new("ocr", response.clone());
self.logger_runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(
self.standard_logging_payload(context, timing),
),
&response_obj,
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let logging_error = LoggingError {
message: error.to_string(),
kind: core_error_kind(error).to_string(),
};
let response_obj = CallbackValue::new(
"error",
json!({
"message": logging_error.message,
"kind": logging_error.kind,
}),
);
self.logger_runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(
self.standard_logging_payload(context, timing),
)
.with_failure_error(logging_error),
Some(&response_obj),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Ocr,
selected_guardrails: Vec::new(),
metadata: std::collections::HashMap::new(),
user_api_key_hash: metadata.user_api_key_hash.clone(),
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
trace_parent: None,
}
}
fn parse_ocr_pre_call_guardrail_request(
request: GuardrailRequest,
) -> Result<(Value, Map<String, Value>), Error> {
let Value::Object(mut data) = request.data else {
return Err(Error::InvalidRequest(
"OCR pre_call guardrail must return an object".to_string(),
));
};
let document = data.remove("document").ok_or_else(|| {
Error::InvalidRequest("OCR pre_call guardrail removed document".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(params)) => params,
Some(_) => {
return Err(Error::InvalidRequest(
"OCR pre_call guardrail optional_params must be an object".to_string(),
));
}
None => Map::new(),
};
Ok((document, optional_params))
}
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<Value, Error> {
let Value::Object(mut data) = request.data else {
return Err(Error::InvalidRequest(
"OCR during_call guardrail must return an object".to_string(),
));
};
data.remove("body")
.ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string()))
}
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

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

@ -1,174 +1,127 @@
use litellm_core::Error;
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::ocr::{
OcrClient,
wire::{OcrWireRequest, decode_request},
};
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::OcrRequest;
use handler::execute_ocr_provider_call;
use prepare::{PreparedOcrCall, prepare_ocr_call};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
CallLifecycle::default()
.run_request(request, &hooks, |request| {
execute_ocr_provider_call(request, &hooks)
})
core_ocr(request).await
}
async fn core_ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
validate_host_hooks(&request)?;
let client = OcrClient::new(crate::client::http_client().clone())?;
let core_request = decode_request(OcrWireRequest {
model: request.model.to_string(),
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
custom_llm_provider: request.custom_llm_provider.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
input_sources: Default::default(),
timeout_seconds: request.timeout.map(|timeout| timeout.as_secs_f64()),
})?;
client
.perform(core_request)
.await
.map(|response| response.into_json())
}
fn validate_host_hooks(request: &OcrRequest<'_>) -> Result<(), Error> {
if !request.guardrails.is_empty() {
return Err(Error::Unsupported(
"OCR host guardrails are not wired to the core path",
));
}
if !request.callbacks.is_empty() {
return Err(Error::Unsupported(
"OCR host callbacks are not wired to the core path",
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use litellm_core::ocr::wire::is_supported_request;
use serde_json::{Map, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::{OcrRequest, ocr};
use crate::integrations::types::RequestMetadata;
use super::{OcrRequest, validate_host_hooks};
use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook};
use crate::integrations::custom_logger::CustomLogger;
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]);
struct TestGuardrail;
impl CustomGuardrail for TestGuardrail {
fn guardrail_name(&self) -> &str {
"test"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&[]
}
String::from_utf8(request).expect("request is utf8")
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
struct TestLogger;
impl CustomLogger for TestLogger {}
fn request() -> OcrRequest<'static> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
model: "model",
document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}),
api_key: None,
api_base: None,
custom_llm_provider: None,
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
request_metadata: Default::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");
#[test]
fn core_activation_includes_migrated_providers() {
assert!(is_supported_request("model", Some("mistral")));
assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
assert!(is_supported_request(
"doc-intelligence/prebuilt-layout",
Some("azure_ai")
));
assert!(is_supported_request("parse-v3", Some("reducto")));
assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
}
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"})),
]);
#[test]
fn core_path_rejects_unwired_guardrails() {
let request = OcrRequest {
guardrails: vec![Arc::new(TestGuardrail)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("guardrails are not wired"));
}
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""#));
#[test]
fn core_path_rejects_unwired_callbacks() {
let request = OcrRequest {
callbacks: vec![Arc::new(TestLogger)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("callbacks are not wired"));
}
}

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

@ -1,163 +0,0 @@
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;
use super::types::{OcrRequest, PreparedOcrRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
pub(crate) struct PreparedOcrCall {
pub(crate) request: PreparedOcrRequest,
pub(crate) hooks: OcrLifecycleHooks,
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let call_id = request
.litellm_call_id
.map(str::to_string)
.unwrap_or_else(new_ocr_call_id);
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "mistral",
});
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()))
.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();
let mut mapped = config.map_ocr_params(
&request
.optional_params
.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,
};
PreparedOcrCall {
request: PreparedOcrRequest {
config,
model,
custom_llm_provider,
litellm_call_id: call_id,
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params,
timeout: request.timeout,
},
hooks: OcrLifecycleHooks::new(
CustomLoggerRunner::new(request.callbacks),
CustomGuardrailRunner::new(request.guardrails),
request.request_metadata,
),
}
}
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);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.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`"))
);
}
}

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

@ -1,8 +1,6 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use litellm_core::ocr::transformation::OcrProviderConfig;
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
@ -23,37 +21,3 @@ pub struct OcrRequest<'a> {
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}
pub(crate) struct PreparedOcrRequest {
pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>,
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
pub(crate) document: Value,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) optional_params: Map<String, Value>,
pub(crate) timeout: Option<Duration>,
}
impl CallLifecycleRequest for PreparedOcrRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"ocr",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}
pub(crate) struct ProviderOcrRequest {
pub(crate) model: String,
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>,
}

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

@ -105,7 +105,11 @@ impl IntoResponse for MessagesRouteError {
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
Error::Auth(_) => (
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),

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

@ -47,10 +47,10 @@ pub async fn messages_request(
.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 response = match routes::app(state).oneshot(request).await {
Ok(response) => response,
Err(error) => match error {},
};
let status: StatusCode = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX)
.await

641
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -1,641 +0,0 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use litellm_ai_gateway::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
};
use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
String::from_utf8(request).expect("request is utf8")
}
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")
}
#[derive(Clone, Debug, PartialEq)]
struct RecordedLogEvent {
hook: &'static str,
model: String,
call_type: String,
user_id: Option<String>,
response_object: Option<String>,
error_kind: Option<String>,
}
#[derive(Default)]
struct RecordingOcrLogger {
events: Mutex<Vec<RecordedLogEvent>>,
}
impl RecordingOcrLogger {
fn events(&self) -> Vec<RecordedLogEvent> {
self.events.lock().unwrap().clone()
}
}
impl CustomLogger for RecordingOcrLogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_success_event",
model: model_call_details.model.clone(),
call_type: model_call_details.call_type.to_string(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: Some(response_obj.object.clone()),
error_kind: None,
});
Ok(())
})
}
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_failure_event",
model: model_call_details.model.clone(),
call_type: model_call_details.call_type.to_string(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: response_obj.map(|value| value.object.clone()),
error_kind: model_call_details
.failure_error
.as_ref()
.map(|error| error.kind.clone()),
});
Ok(())
})
}
}
struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
}
impl RecordingOcrGuardrail {
fn new(hooks: Vec<GuardrailEventHook>) -> Self {
Self {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
}
}
fn blocking_pre_call() -> Self {
Self {
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,
}
}
fn events(&self) -> Vec<&'static str> {
self.events.lock().unwrap().clone()
}
}
impl CustomGuardrail for RecordingOcrGuardrail {
fn guardrail_name(&self) -> &str {
"recording-ocr-guardrail"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&self.hooks
}
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_pre_call_hook");
if self.block_pre_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["document"]["guarded_pre"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
fn async_moderation_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
mut request: GuardrailRequest,
) -> 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))
})
}
}
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_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()];
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");
}
#[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");
assert!(
matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)"))
);
server.await.expect("server task completes");
}
#[tokio::test]
async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: vec![logger.clone()],
guardrails: vec![guardrail.clone()],
request_metadata: RequestMetadata {
user_api_key_user_id: Some("user-1".to_string()),
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
guardrail.events(),
vec!["async_pre_call_hook", "async_moderation_hook"]
);
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_success_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: Some("user-1".to_string()),
response_object: Some("ocr".to_string()),
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
assert!(request.contains(r#""guarded_during":true"#), "{request}");
}
#[tokio::test]
async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let _request = read_http_request(&mut socket).await;
let response_body = "provider failed";
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
});
let logger = Arc::new(RecordingOcrLogger::default());
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: vec![logger.clone()],
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_failure_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: None,
response_object: Some("error".to_string()),
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]
async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
let err = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_millis(100)),
callbacks: vec![logger.clone()],
guardrails: vec![guardrail.clone()],
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-3"),
})
.await
.expect_err("guardrail blocks request");
assert!(matches!(err, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_failure_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: None,
response_object: Some("error".to_string()),
error_kind: Some("InvalidRequest".to_string()),
}]
);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "provider socket should not be touched");
}
#[tokio::test]
async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_headers(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer sk-from-python".to_string()),
);
headers.insert(
"x-trace-id".to_string(),
Value::String("trace-1".to_string()),
);
let response = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-for-rust-fallback"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: Some(headers),
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
})
.await
.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let request = server.await.expect("server task completes");
let authorization_count = request
.lines()
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
.count();
assert_eq!(authorization_count, 1, "{request}");
assert!(
request.contains("authorization: Bearer sk-from-python")
|| request.contains("Authorization: Bearer sk-from-python"),
"{request}"
);
}
#[tokio::test]
async fn document_intelligence_poll_uses_resolved_subscription_key() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let operation_url = format!("http://{addr}/operations/1");
let server = tokio::spawn(async move {
let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
let post_request = read_http_headers(&mut post_socket).await;
let post_response = format!(
"HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
);
post_socket
.write_all(post_response.as_bytes())
.await
.expect("writes post response");
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
let poll_request = read_http_headers(&mut poll_socket).await;
let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
let poll_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
poll_socket
.write_all(poll_response.as_bytes())
.await
.expect("writes poll response");
(post_request, poll_request)
});
let response = ocr(OcrRequest {
model: "doc-intelligence/prebuilt-read",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("di-key"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
})
.await
.expect("document intelligence request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let (post_request, poll_request) = server.await.expect("server task completes");
assert!(
post_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{post_request}"
);
assert!(
poll_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{poll_request}"
);
}

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

@ -4,17 +4,33 @@ version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
autotests = false
[[test]]
name = "workspace_crate_allowlist"
path = "tests/workspace_crate_allowlist.rs"
[dependencies]
base64.workspace = true
azure_core.workspace = true
azure_identity.workspace = true
data-url = "0.3.2"
gcp_auth.workspace = true
moka.workspace = true
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_path_to_error = "0.1"
strum.workspace = true
subtle.workspace = true
tokio.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
url.workspace = true
veil.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 }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
@ -36,5 +52,4 @@ observability = ["dep:tracing-subscriber"]
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber.workspace = true

168
litellm-rust/crates/core/src/auth/credential.rs Normal file
View file

@ -0,0 +1,168 @@
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use veil::Redact;
use crate::AuthError;
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialFileRef {
Path(PathBuf),
EnvironmentVariable(String),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialRef {
Explicit(SecretValue),
Env(String),
File(CredentialFileRef),
Request(String),
Host(String),
None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialLookup {
Found(SecretValue),
Missing,
Declined,
}
pub type CredentialLookupFuture<'a> =
Pin<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + Send + 'a>>;
pub trait CredentialResolver: std::fmt::Debug + Send + Sync {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>;
}
#[derive(Clone, Redact)]
pub struct CredentialResolverHandle(#[redact(with = "[REDACTED]")] Arc<dyn CredentialResolver>);
impl CredentialResolverHandle {
pub fn new(resolver: Arc<dyn CredentialResolver>) -> Self {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
self.0.resolve(reference).await
}
}
#[derive(Clone, Debug)]
pub enum CredentialPlan {
Static(CredentialRef),
Caller(TokenProviderHandle),
None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialPlanResolution {
Resolved(ResolvedCredential),
Unavailable,
}
impl CredentialPlan {
pub async fn resolve(
&self,
resolver: &CredentialResolverHandle,
) -> Result<CredentialPlanResolution, AuthError> {
match self {
Self::Static(CredentialRef::Explicit(secret)) => Ok(
CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())),
),
Self::Static(CredentialRef::None) | Self::None => {
Ok(CredentialPlanResolution::Unavailable)
}
Self::Static(reference) => match resolver.resolve(reference).await? {
CredentialLookup::Found(secret) => Ok(CredentialPlanResolution::Resolved(
ResolvedCredential::Static(secret),
)),
CredentialLookup::Missing | CredentialLookup::Declined => {
Ok(CredentialPlanResolution::Unavailable)
}
},
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyCallerCredential);
}
Ok(CredentialPlanResolution::Resolved(credential))
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::{
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::AuthError;
use crate::auth::SecretValue;
#[derive(Debug)]
struct HostResolver;
impl CredentialResolver for HostResolver {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async move {
Ok(match reference {
CredentialRef::Host(name) if name == "rotating-token" => {
CredentialLookup::Found(SecretValue::new("resolved"))
}
_ => CredentialLookup::Declined,
})
})
}
}
#[tokio::test]
async fn static_host_reference_resolves_at_acquisition_time() {
let resolver = CredentialResolverHandle::new(Arc::new(HostResolver));
let plan = CredentialPlan::Static(CredentialRef::Host("rotating-token".to_string()));
let resolved = plan.resolve(&resolver).await.unwrap();
assert!(matches!(resolved, CredentialPlanResolution::Resolved(_)));
}
#[tokio::test]
async fn declined_reference_is_available_for_pre_acquisition_fallback() {
let resolver = CredentialResolverHandle::new(Arc::new(HostResolver));
let plan = CredentialPlan::Static(CredentialRef::Request("api-key".to_string()));
assert_eq!(
plan.resolve(&resolver).await.unwrap(),
CredentialPlanResolution::Unavailable
);
}
#[derive(Debug)]
struct FailingResolver;
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async { Err(AuthError::UnresolvedOidcReference) })
}
}
#[tokio::test]
async fn acquisition_failure_is_terminal() {
let resolver = CredentialResolverHandle::new(Arc::new(FailingResolver));
let plan = CredentialPlan::Static(CredentialRef::Host("token".to_string()));
let error = plan
.resolve(&resolver)
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(error, AuthError::UnresolvedOidcReference);
}
}

128
litellm-rust/crates/core/src/auth/error.rs Normal file
View file

@ -0,0 +1,128 @@
use thiserror::Error;
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthError {
#[error("invalid authentication configuration: {0}")]
Configuration(#[from] AuthConfigurationError),
#[error("credential acquisition failed: {0}")]
AzureTokenAcquisition(String),
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
VertexTokenAcquisition(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<AuthError>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[error("{0}")]
MissingCredential(#[from] MissingCredential),
#[error("{0}")]
Aws(#[from] AwsAuthError),
#[error("invalid authentication header")]
InvalidHeader,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthConfigurationError {
#[error("credential header already exists")]
ExistingCredentialHeader,
#[error("credential plan is not allowed by the provider auth policy")]
DisallowedCredentialPlan,
#[error("credential cannot be empty")]
EmptyCredential,
#[error("invalid Azure credential selector")]
InvalidAzureSelector,
#[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")]
MissingClientSecretFields,
#[error("WorkloadIdentityCredential requires tenant_id")]
MissingWorkloadTenant,
#[error("WorkloadIdentityCredential requires client_id")]
MissingWorkloadClient,
#[error("WorkloadIdentityCredential requires azure_federated_token_file")]
MissingWorkloadTokenFile,
#[error("credential reference requires a host credential resolver")]
MissingHostResolver,
#[error("caller credential plan requires provider-specific inputs")]
MissingCallerInputs,
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{0} must be a string or null")]
InvalidFieldType(String),
#[error("unsupported OIDC reference")]
UnsupportedOidcReference,
#[error("{0} cannot be empty")]
EmptyReference(String),
#[error("Azure credential initialization failed: {0}")]
AzureCredentialInitialization(String),
#[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")]
InvalidAzureAuthority,
#[error("request-controlled Azure auth inputs cannot be combined with host credentials")]
MixedAzureCredentialSources,
#[error("request-controlled Azure credential references are not allowed")]
RequestAzureCredentialReference,
#[error("host credentials cannot be sent to a request-controlled Azure endpoint")]
RequestAzureCredentialDestination,
#[error("credentials cannot be sent to a request-controlled Vertex AI endpoint")]
RequestVertexCredentialDestination,
#[error(
"request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
)]
RequestVertexTokenEndpoint,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum MissingCredential {
#[error(
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
)]
AnthropicApiKey,
#[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")]
AzureApiKey,
#[error(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
)]
AzureApiBase,
#[error(
"Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiRealtimeApiKey,
#[error(
"Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiResponsesApiKey,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AwsAuthError {
#[error("AWS profile credentials failed: {0}")]
Profile(String),
#[error("AWS default credentials failed: {0}")]
DefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
WebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
WebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
SigningParameters(String),
#[error("AWS signable request failed: {0}")]
SignableRequest(String),
#[error("AWS request signing failed: {0}")]
Signing(String),
#[error("AWS web identity response had no credentials")]
MissingWebIdentityCredentials,
}

86
litellm-rust/crates/core/src/auth/http.rs Normal file
View file

@ -0,0 +1,86 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlacement {
Bearer,
Header(&'static str),
}
impl CredentialPlacement {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "Authorization",
Self::Header(name) => name,
}
}
}
pub(crate) fn apply_credential(
headers: Vec<(String, String)>,
credential: &str,
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, AuthError> {
if credential.trim().is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyCredential,
));
}
if headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name()))
{
return Err(AuthError::Configuration(
AuthConfigurationError::DuplicateHeader(placement.header_name()),
));
}
let value = match placement {
CredentialPlacement::Bearer => format!("Bearer {credential}"),
CredentialPlacement::Header(_) => credential.to_string(),
};
Ok(
std::iter::once((placement.header_name().to_string(), value))
.chain(headers)
.collect(),
)
}
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RequestAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
}
#[cfg(test)]
mod tests {
use super::{CredentialPlacement, apply_credential};
#[test]
fn bearer_uses_authorization_header() {
let headers = apply_credential(Vec::new(), "key", CredentialPlacement::Bearer)
.expect("credential applies");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer key".to_string())]
);
}
#[test]
fn named_header_rejects_existing_value() {
let error = apply_credential(
vec![(
"ocp-apim-subscription-key".to_string(),
"caller-key".to_string(),
)],
"configured-key",
CredentialPlacement::Header("Ocp-Apim-Subscription-Key"),
)
.expect_err("provider policy must handle existing credentials");
assert!(error.to_string().contains("already exists"));
}
}

56
litellm-rust/crates/core/src/auth/mod.rs Normal file
View file

@ -0,0 +1,56 @@
mod credential;
pub mod error;
pub(crate) mod vertex;
pub use error::AuthError;
pub(crate) mod http;
mod policy;
mod secret;
mod token;
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum InputSource {
Request,
#[default]
Deployment,
Environment,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Sourced<T> {
value: T,
source: InputSource,
}
impl<T> Sourced<T> {
pub fn new(value: T, source: InputSource) -> Self {
Self { value, source }
}
pub fn value(&self) -> &T {
&self.value
}
pub fn source(&self) -> InputSource {
self.source
}
pub fn into_value(self) -> T {
self.value
}
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
Sourced::new(map(self.value), self.source)
}
}
pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
};
pub use http::{CredentialPlacement, RequestAuth};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

114
litellm-rust/crates/core/src/auth/policy.rs Normal file
View file

@ -0,0 +1,114 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use super::http::apply_credential;
use super::{CredentialPlacement, ResolvedCredential};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlanKind {
Static,
Entra,
Caller,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CredentialRule {
pub kind: CredentialPlanKind,
pub placement: CredentialPlacement,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExistingHeaderBehavior {
Preserve,
Reject,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ProviderAuthPolicy {
pub rules: &'static [CredentialRule],
pub accepted_existing_headers: &'static [&'static str],
pub existing_header_behavior: ExistingHeaderBehavior,
pub scope: Option<&'static str>,
pub audience: Option<&'static str>,
}
impl ProviderAuthPolicy {
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, _)| {
self.accepted_existing_headers
.iter()
.any(|accepted| name.eq_ignore_ascii_case(accepted))
})
}
pub fn apply(
&self,
headers: Vec<(String, String)>,
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, AuthError> {
if self.has_existing_credential(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(AuthError::Configuration(
AuthConfigurationError::ExistingCredentialHeader,
)),
};
}
let rule =
self.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(AuthError::Configuration(
AuthConfigurationError::DisallowedCredentialPlan,
))?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
#[cfg(test)]
mod tests {
use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue};
const RULES: &[CredentialRule] = &[CredentialRule {
kind: CredentialPlanKind::Static,
placement: CredentialPlacement::Header("x-api-key"),
}];
const POLICY: ProviderAuthPolicy = ProviderAuthPolicy {
rules: RULES,
accepted_existing_headers: &["x-api-key"],
existing_header_behavior: ExistingHeaderBehavior::Preserve,
scope: None,
audience: None,
};
#[test]
fn rules_define_allowed_plans_and_credential_placement() {
let headers = POLICY
.apply(
Vec::new(),
CredentialPlanKind::Static,
&ResolvedCredential::Static(SecretValue::new("secret")),
)
.unwrap();
assert_eq!(
headers,
vec![("x-api-key".to_string(), "secret".to_string())]
);
}
#[test]
fn unsupported_plan_is_rejected() {
let error = POLICY
.apply(
Vec::new(),
CredentialPlanKind::Entra,
&ResolvedCredential::Static(SecretValue::new("secret")),
)
.unwrap_err();
assert!(error.to_string().contains("not allowed"));
}
}

41
litellm-rust/crates/core/src/auth/secret.rs Normal file
View file

@ -0,0 +1,41 @@
use veil::Redact;
#[derive(Redact, Clone)]
pub struct SecretValue(#[redact(with = "[REDACTED]")] String);
impl SecretValue {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn expose(&self) -> &str {
&self.0
}
}
impl PartialEq for SecretValue {
fn eq(&self, other: &Self) -> bool {
subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into()
}
}
impl Eq for SecretValue {}
#[cfg(test)]
mod tests {
use super::SecretValue;
#[test]
fn debug_redacts_plaintext() {
let debug = format!("{:?}", SecretValue::new("credential-value"));
assert!(!debug.contains("credential-value"));
assert!(debug.contains("REDACTED"));
}
#[test]
fn equality_compares_plaintext_values() {
assert_eq!(SecretValue::new("same"), SecretValue::new("same"));
assert_ne!(SecretValue::new("same"), SecretValue::new("different"));
}
}

47
litellm-rust/crates/core/src/auth/token.rs Normal file
View file

@ -0,0 +1,47 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::SystemTime;
use veil::Redact;
use crate::AuthError;
use super::secret::SecretValue;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ResolvedCredential {
Static(SecretValue),
AccessToken {
token: SecretValue,
expires_on: Option<SystemTime>,
},
}
impl ResolvedCredential {
pub fn secret(&self) -> &SecretValue {
match self {
Self::Static(secret) | Self::AccessToken { token: secret, .. } => secret,
}
}
}
pub type TokenFuture<'a> =
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
fn acquire(&self) -> TokenFuture<'_>;
}
#[derive(Clone, Redact)]
pub struct TokenProviderHandle(#[redact(with = "[REDACTED]")] Arc<dyn TokenProvider>);
impl TokenProviderHandle {
pub fn new(caller: Arc<dyn TokenProvider>) -> Self {
Self(caller)
}
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
self.0.acquire().await
}
}

592
litellm-rust/crates/core/src/auth/vertex.rs Normal file
View file

@ -0,0 +1,592 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use gcp_auth::{CustomServiceAccount, TokenProvider};
use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::auth::error::AuthConfigurationError;
use crate::auth::http::apply_credential;
use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[derive(Clone, Debug, Default)]
pub(crate) struct VertexConfig {
credentials: Option<Sourced<SecretValue>>,
project_id: Option<String>,
location: Option<String>,
}
impl VertexConfig {
pub(crate) fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
Ok(Self {
credentials: optional_credentials(
params,
sources,
&["vertex_credentials", "vertex_ai_credentials"],
)?,
project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?,
location: optional_string(params, &["vertex_location", "vertex_ai_location"])?,
})
}
pub(crate) fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
pub(crate) fn location(&self) -> Option<&str> {
self.location.as_deref()
}
}
pub(crate) struct VertexEnvironment {
pub headers: Vec<(String, String)>,
pub project_id: String,
}
struct VertexAccessToken {
token: String,
project_id: String,
}
pub(crate) fn get_vertex_ai_project(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
config
.project_id()
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV))
}
pub(crate) fn get_vertex_ai_location(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
config
.location()
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEXAI_LOCATION_ENV))
.or_else(|| non_empty_env(env_lookup, VERTEX_LOCATION_ENV))
}
#[derive(Clone)]
pub(crate) struct VertexAuth {
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
impl Default for VertexAuth {
fn default() -> Self {
Self::new(Arc::new(GcpProviderLoader))
}
}
impl VertexAuth {
fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
Self {
providers: Cache::builder().max_capacity(64).build(),
loader,
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexEnvironment, AuthError> {
let has_authorization = headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("Authorization"));
let static_token = api_key
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV))
.or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV));
let project_id = get_vertex_ai_project(config, env_lookup);
if !has_authorization && static_token.is_none() {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
let project_id = match project_id {
Some(project_id) => project_id,
None => {
self.load_provider(config, env_lookup)
.await?
.project_id()
.await?
}
};
let headers = if has_authorization {
headers
} else {
apply_credential(
headers,
static_token.as_deref().expect("static token was checked"),
CredentialPlacement::Bearer,
)?
};
Ok(VertexEnvironment {
headers,
project_id,
})
}
async fn get_access_token(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexAccessToken, AuthError> {
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
Ok(VertexAccessToken { token, project_id })
}
async fn load_provider(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, AuthError> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
.try_get_with(key, self.loader.load(source))
.await
.map_err(|error| (*error).clone())
}
}
trait VertexTokenSource: Send + Sync {
fn project_id(&self) -> VertexAuthFuture<'_, String>;
fn token(&self) -> VertexAuthFuture<'_, String>;
}
trait VertexProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
}
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
impl VertexTokenSource for GcpTokenSource {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.project_id()
.await
.map(|project| project.to_string())
.map_err(auth_acquisition_error)
})
}
fn token(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.token(&[CLOUD_PLATFORM_SCOPE])
.await
.map(|token| token.as_str().to_string())
.map_err(auth_acquisition_error)
})
}
}
struct GcpProviderLoader;
impl VertexProviderLoader for GcpProviderLoader {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = match source {
CredentialSource::Inline(configured) => Arc::new(
CustomServiceAccount::from_json(validate_request_credentials(
configured.expose(),
)?)
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
let service_account = if Path::new(configured).is_file() {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)?;
Arc::new(service_account)
}
CredentialSource::ApplicationCredentials(path) => {
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
}
CredentialSource::Adc => {
gcp_auth::provider().await.map_err(auth_acquisition_error)?
}
};
Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
credentials
.get("token_uri")
.and_then(Value::as_str)
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into());
}
Ok(configured)
}
#[derive(Clone, Debug)]
enum CredentialSource {
Inline(SecretValue),
Trusted(SecretValue),
ApplicationCredentials(String),
Adc,
}
impl CredentialSource {
fn cache_key(&self) -> CredentialCacheKey {
match self {
Self::Inline(configured) => {
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
}
Self::Trusted(configured) => {
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
}
Self::ApplicationCredentials(path) => {
CredentialCacheKey::ApplicationCredentials(path.clone())
}
Self::Adc => CredentialCacheKey::Adc,
}
}
}
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
enum CredentialCacheKey {
Inline([u8; 32]),
Trusted([u8; 32]),
ApplicationCredentials(String),
Adc,
}
fn credential_source(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CredentialSource {
if let Some(configured) = config.credentials.clone() {
return match configured.source() {
InputSource::Request => CredentialSource::Inline(configured.into_value()),
InputSource::Deployment | InputSource::Environment => {
CredentialSource::Trusted(configured.into_value())
}
};
}
if let Some(configured) = non_empty_env(env_lookup, VERTEXAI_CREDENTIALS_ENV) {
return CredentialSource::Trusted(SecretValue::new(configured));
}
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
.map(CredentialSource::ApplicationCredentials)
.unwrap_or(CredentialSource::Adc)
}
fn optional_credentials(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
names: &[&str],
) -> Result<Option<Sourced<SecretValue>>, AuthError> {
for name in names {
let source = source_for(sources, name);
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => {
return Ok(Some(Sourced::new(SecretValue::new(value), source)));
}
Some(Value::Object(value)) if value.is_empty() => continue,
Some(Value::Object(value)) => {
return serde_json::to_string(value)
.map(SecretValue::new)
.map(|value| Sourced::new(value, source))
.map(Some)
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!(
"{}: {error}",
names[0]
)))
});
}
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
}
}
}
Ok(None)
}
fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
sources.get(name).copied().unwrap_or_default()
}
fn optional_string(
params: &Map<String, Value>,
names: &[&str],
) -> Result<Option<String>, AuthError> {
for name in names {
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => return Ok(Some(value.clone())),
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
}
}
}
Ok(None)
}
fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option<String>, name: &str) -> Option<String> {
env_lookup(name)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError {
AuthError::VertexTokenAcquisition(error.to_string())
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use serde_json::json;
use super::*;
struct FakeProvider {
calls: Arc<AtomicUsize>,
}
impl VertexTokenSource for FakeProvider {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-project".into()) })
}
fn token(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-token".into()) })
}
}
struct FakeLoader {
loads: Arc<AtomicUsize>,
provider: Arc<dyn VertexTokenSource>,
}
impl VertexProviderLoader for FakeLoader {
fn load(
&self,
_source: CredentialSource,
) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
let loads = self.loads.clone();
let provider = self.provider.clone();
Box::pin(async move {
loads.fetch_add(1, Ordering::SeqCst);
Ok(provider)
})
}
}
fn config(value: Value) -> VertexConfig {
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
.unwrap()
}
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
}
#[test]
fn config_is_typed_and_secrets_are_redacted() {
let config = config(json!({
"vertex_credentials":{"private_key":"secret-key"},
"vertex_project":"project-1",
"vertex_location":"europe-west4"
}));
assert_eq!(config.project_id(), Some("project-1"));
assert_eq!(config.location(), Some("europe-west4"));
assert!(!format!("{config:?}").contains("secret-key"));
assert!(
VertexConfig::from_sourced_optional_params(
json!({"vertex_credentials":true}).as_object().unwrap(),
&BTreeMap::new()
)
.is_err()
);
}
#[test]
fn empty_primary_values_fall_back_to_python_aliases() {
let config = config(json!({
"vertex_credentials": null,
"vertex_ai_credentials": "alias-credentials",
"vertex_project": " ",
"vertex_ai_project": "alias-project",
"vertex_location": null,
"vertex_ai_location": "alias-location"
}));
assert_eq!(
config.credentials.as_ref().unwrap().value().expose(),
"alias-credentials"
);
assert_eq!(config.project_id(), Some("alias-project"));
assert_eq!(config.location(), Some("alias-location"));
}
#[test]
fn project_and_location_prefer_input_then_environment() {
let configured =
config(json!({"vertex_project":"input-project","vertex_location":"input-location"}));
let env = |name: &str| Some(format!("env-{name}"));
assert_eq!(
get_vertex_ai_project(&configured, &env).as_deref(),
Some("input-project")
);
assert_eq!(
get_vertex_ai_location(&configured, &env).as_deref(),
Some("input-location")
);
let empty = VertexConfig::default();
assert_eq!(
get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(),
Some("env-project")
);
assert_eq!(
get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV)
.then(|| "fallback-location".into()))
.as_deref(),
Some("fallback-location")
);
}
#[test]
fn credential_discovery_prefers_input_then_environment_then_adc() {
let params = json!({"vertex_credentials":"input-json"});
let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]);
let configured =
VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap();
assert!(
matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json")
);
let empty = VertexConfig::default();
assert!(
matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json")
);
assert!(
matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
);
assert!(matches!(
credential_source(&empty, &|_| None),
CredentialSource::Adc
));
assert_ne!(
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
);
}
#[test]
fn request_credentials_require_canonical_token_endpoint() {
assert!(
validate_request_credentials(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#)
.is_ok()
);
assert!(matches!(
validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
));
assert!(matches!(
validate_request_credentials("{}"),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
));
}
#[tokio::test]
async fn explicit_token_and_header_do_not_acquire_adc() {
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
let configured = config(json!({"vertex_project":"project-1"}));
let explicit = auth
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
.await
.unwrap();
assert_eq!(explicit.headers[0].1, "Bearer access-token");
let existing = auth
.validate_environment(
vec![("authorization".into(), "Bearer existing".into())],
None,
&configured,
&|_| None,
)
.await
.unwrap();
assert_eq!(existing.headers[0].1, "Bearer existing");
assert_eq!(loads.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn provider_is_reused_across_authentication_calls() {
let calls = Arc::new(AtomicUsize::new(0));
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(calls.clone(), loads.clone());
for _ in 0..2 {
let environment = auth
.validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
.await
.unwrap();
assert_eq!(environment.project_id, "adc-project");
assert_eq!(environment.headers[0].1, "Bearer adc-token");
}
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
}

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

@ -43,3 +43,23 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10;
pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120;
pub(crate) const OCR_POLL_RETRY_SECS: u64 = 2;
pub(crate) const AZURE_DI_API_VERSION: &str = "2024-11-30";
pub(crate) const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
pub(crate) const AZURE_DI_DEFAULT_DPI: i64 = 96;
pub(crate) const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
pub(crate) const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
pub(crate) const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://";
pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr";
pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";

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

@ -17,6 +17,22 @@ pub enum Error {
InvalidRequest(String),
#[error("{0}")]
Auth(String),
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
@ -36,6 +52,90 @@ pub enum Error {
Unsupported(&'static str),
}
#[derive(Debug, ThisError)]
pub(crate) enum MediaError {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
DownloadDisabled,
#[error("media download exceeds the maximum size")]
DownloadTooLarge,
#[error("too many redirects while fetching media")]
TooManyRedirects,
#[error("media redirect is missing a Location header")]
MissingRedirectLocation,
#[error("invalid media redirect")]
InvalidRedirect,
#[error("media download failed with status {0}")]
Http(u16),
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] TransportError),
}
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum TransportError {
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
#[error("could not reach the provider: {0}")]
Connect(String),
}
impl TransportError {
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
let message = error.without_url().to_string();
if before_dispatch {
Self::Connect(message)
} else {
Self::Network(message)
}
}
}
impl From<reqwest::Error> for TransportError {
fn from(error: reqwest::Error) -> Self {
Self::Network(error.without_url().to_string())
}
}
impl From<crate::ocr::error::OcrRequestError> for Error {
fn from(error: crate::ocr::error::OcrRequestError) -> Self {
match error {
crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
error => Self::InvalidRequest(error.to_string()),
}
}
}
impl From<crate::ocr::error::OcrResponseError> for Error {
fn from(error: crate::ocr::error::OcrResponseError) -> Self {
Self::InvalidResponse(error.to_string())
}
}
impl From<TransportError> for Error {
fn from(error: TransportError) -> Self {
match error {
TransportError::Http { status, body } => Self::Http { status, body },
TransportError::Network(message) => Self::Network(message),
TransportError::Connect(message) => Self::Connect(message),
}
}
}
impl From<crate::AuthError> for Error {
fn from(error: crate::AuthError) -> Self {
match error {
crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider },
error => Self::Auth(error.to_string()),
}
}
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
@ -46,3 +146,61 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod transport_tests {
use super::*;
#[test]
fn missing_auth_key_preserves_provider_in_public_error() {
assert_eq!(
Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }),
Error::MissingApiKey { provider: "Vertex" }
);
}
#[tokio::test]
async fn transport_errors_remove_urls_and_keep_dispatch_context() {
let error = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get("http://localhost:invalid/private?api_key=secret")
.send()
.await
.expect_err("invalid port");
let error = TransportError::from_reqwest_before_dispatch(error);
assert!(matches!(error, TransportError::Connect(_)));
assert!(!error.to_string().contains("secret"));
assert!(!error.to_string().contains("private"));
}
#[tokio::test]
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
use std::time::Duration;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let address = listener.local_addr().expect("address");
let request = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}"))
.timeout(Duration::from_millis(200))
.send();
let (response, accepted) = tokio::join!(
request,
tokio::time::timeout(Duration::from_secs(2), listener.accept())
);
let _connection = accepted
.expect("accept deadline")
.expect("accepted connection");
let error = response.expect_err("server does not respond");
assert!(error.is_timeout());
assert!(matches!(
TransportError::from_reqwest_before_dispatch(error),
TransportError::Network(_)
));
}
}

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

@ -1,10 +1,43 @@
//! Header and upstream-body helpers shared by every route module.
use serde_json::{Map, Value};
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
use crate::error::{Error, json_type_name};
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) enum HeaderPolicy<'a> {
All,
Only(&'a [&'a str]),
Except(&'a [&'a str]),
}
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) fn with_headers(
builder: reqwest::RequestBuilder,
headers: &[(String, String)],
policy: HeaderPolicy<'_>,
) -> reqwest::RequestBuilder {
headers
.iter()
.filter(|(name, _)| match policy {
HeaderPolicy::All => true,
HeaderPolicy::Only(names) => names
.iter()
.any(|allowed| name.eq_ignore_ascii_case(allowed)),
HeaderPolicy::Except(names) => !names
.iter()
.any(|excluded| name.eq_ignore_ascii_case(excluded)),
})
.fold(builder, |builder, (name, value)| {
builder.header(name, value)
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn http_request(
request: reqwest::RequestBuilder,
@ -12,8 +45,6 @@ pub async fn http_request(
request.send().await
}
/// Bound an upstream error body before it crosses a host boundary, so provider
/// bodies stay data-minimized.
pub fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
return body.to_string();
@ -61,11 +92,77 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
})
}
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub(crate) fn deserialize_optional_param<'de, D, T>(
deserializer: D,
) -> Result<Option<Option<T>>, D::Error>
where
D: serde::Deserializer<'de>,
T: serde::Deserialize<'de>,
{
<Option<T> as serde::Deserialize>::deserialize(deserializer).map(Some)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[rstest::rstest]
#[case(HeaderPolicy::All, true, true)]
#[case(HeaderPolicy::Only(&["authorization"]), true, false)]
#[case(HeaderPolicy::Except(&["authorization"]), false, true)]
fn forwarding_policy_preserves_matching_headers_and_duplicates(
#[case] policy: HeaderPolicy<'_>,
#[case] auth: bool,
#[case] trace: bool,
) {
let request = with_headers(
reqwest::Client::new().get("https://example.com"),
&[
("AuThOrIzAtIoN".into(), "Bearer token".into()),
("X-Trace".into(), "first".into()),
("x-trace".into(), "second".into()),
],
policy,
)
.build()
.unwrap();
assert_eq!(request.headers().contains_key("authorization"), auth);
let traces: Vec<_> = request.headers().get_all("x-trace").iter().collect();
if trace {
assert_eq!(traces, ["first", "second"]);
} else {
assert!(traces.is_empty());
}
}
#[test]
fn multipart_policy_leaves_content_headers_to_reqwest() {
let request = with_headers(
reqwest::Client::new()
.post("https://example.com")
.multipart(reqwest::multipart::Form::new().text("file", "abc")),
&[
("Content-Type".into(), "application/json".into()),
("CONTENT-LENGTH".into(), "0".into()),
],
HeaderPolicy::Except(&["content-type", "content-length"]),
)
.build()
.unwrap();
assert!(
request.headers()["content-type"]
.to_str()
.unwrap()
.starts_with("multipart/form-data; boundary=")
);
assert_ne!(request.headers()["content-length"], "0");
}
#[test]
fn truncate_leaves_short_bodies_untouched() {
assert_eq!(truncate_error_body("short"), "short");

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

@ -1,10 +1,12 @@
pub mod audio_transcription;
pub mod auth;
pub mod caching;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
mod media;
pub mod messages;
#[cfg(any(feature = "observability", test))]
pub mod observability;
@ -14,5 +16,7 @@ pub mod realtime;
pub mod responses;
pub mod router;
pub mod routing_utils;
mod url_utils;
pub use auth::AuthError;
pub use error::Error;

528
litellm-rust/crates/core/src/media.rs Normal file
View file

@ -0,0 +1,528 @@
use std::future::Future;
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use reqwest::Url;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS;
use crate::error::{MediaError, TransportError};
#[derive(Clone)]
pub(crate) struct MediaFetcher {
client: reqwest::Client,
address_resolver: Arc<dyn AddressResolver>,
allow_private_network: bool,
}
type AddressResolution<'a> = Pin<Box<dyn Future<Output = io::Result<Vec<SocketAddr>>> + Send + 'a>>;
trait AddressResolver: Send + Sync {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a>;
}
#[derive(Clone, Copy)]
pub(crate) struct DownloadPolicy {
pub(crate) timeout: Duration,
pub(crate) max_bytes: u64,
pub(crate) max_redirects: usize,
}
#[derive(Debug)]
pub(crate) struct DownloadedMedia {
pub(crate) bytes: Vec<u8>,
pub(crate) content_type: String,
}
impl MediaFetcher {
pub(crate) fn new() -> Result<Self, reqwest::Error> {
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
}
fn with_resolvers<R>(
transport_resolver: Arc<R>,
address_resolver: Arc<dyn AddressResolver>,
) -> Result<Self, reqwest::Error>
where
R: Resolve + 'static,
{
let client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.no_proxy()
.dns_resolver(transport_resolver)
.build()?;
Ok(Self {
client,
address_resolver,
allow_private_network: false,
})
}
#[cfg(test)]
pub(crate) fn for_test(client: reqwest::Client) -> Self {
Self {
client,
address_resolver: Arc::new(AllowPrivateResolver),
allow_private_network: true,
}
}
pub(crate) async fn fetch(
&self,
url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
if policy.max_bytes == 0 {
return Err(MediaError::DownloadDisabled);
}
tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy))
.await
.map_err(|_| MediaError::Timeout)?
}
async fn fetch_before_deadline(
&self,
mut url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
let mut redirects_followed = 0;
loop {
self.validate_url(&url).await?;
let mut response = self
.client
.get(url.clone())
.send()
.await
.map_err(TransportError::from)?;
if response.status().is_redirection() {
if redirects_followed == policy.max_redirects {
return Err(MediaError::TooManyRedirects);
}
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or(MediaError::MissingRedirectLocation)?;
url = url
.join(location)
.map_err(|_| MediaError::InvalidRedirect)?;
redirects_followed += 1;
continue;
}
if !response.status().is_success() {
return Err(MediaError::Http(response.status().as_u16()));
}
enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?;
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let mut bytes = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(TransportError::from)? {
enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?;
bytes.extend_from_slice(&chunk);
}
return Ok(DownloadedMedia {
bytes,
content_type,
});
}
}
async fn validate_url(&self, url: &Url) -> Result<(), MediaError> {
if !matches!(url.scheme(), "http" | "https")
|| !url.username().is_empty()
|| url.password().is_some()
{
return Err(MediaError::BlockedUrl);
}
let host = url.host_str().ok_or(MediaError::BlockedUrl)?;
if self.allow_private_network {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return (!is_blocked_ip(ip))
.then_some(())
.ok_or(MediaError::BlockedUrl);
}
let port = url.port_or_known_default().ok_or(MediaError::BlockedUrl)?;
let addresses = self
.address_resolver
.resolve(host, port)
.await
.map_err(|error| TransportError::Network(error.to_string()))?;
validate_addresses(&addresses)
}
}
fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), MediaError> {
if length > max_bytes {
return Err(MediaError::DownloadTooLarge);
}
Ok(())
}
fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), MediaError> {
if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) {
return Err(MediaError::BlockedUrl);
}
Ok(())
}
fn is_blocked_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(ip) => {
let [first, second, third, _] = ip.octets();
first == 0
|| first == 10
|| first == 127
|| (first == 100 && (64..=127).contains(&second))
|| (first == 169 && second == 254)
|| (first == 172 && (16..=31).contains(&second))
|| (first == 192 && second == 0 && (third == 0 || third == 2))
|| (first == 192 && second == 168)
|| (first == 192 && second == 88 && third == 99)
|| (first == 198 && (second == 18 || second == 19))
|| (first == 198 && second == 51 && third == 100)
|| (first == 203 && second == 0 && third == 113)
|| first >= 224
}
IpAddr::V6(ip) => {
let segments = ip.segments();
ip.is_loopback()
|| ip.is_unspecified()
|| ip.is_multicast()
|| (segments[0] & 0xfe00) == 0xfc00
|| (segments[0] & 0xffc0) == 0xfe80
|| (segments[0] & 0xffc0) == 0xfec0
|| (segments[0] == 0x2001 && segments[1] == 0x0db8)
|| ip
.to_ipv4_mapped()
.or_else(|| ip.to_ipv4())
.map(|ipv4| is_blocked_ip(IpAddr::V4(ipv4)))
.unwrap_or(false)
}
}
}
#[derive(Default)]
struct PublicDnsResolver;
struct SystemAddressResolver;
impl AddressResolver for SystemAddressResolver {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move {
Ok(tokio::net::lookup_host((host, port))
.await?
.collect::<Vec<_>>())
})
}
}
#[cfg(test)]
struct AllowPrivateResolver;
#[cfg(test)]
impl AddressResolver for AllowPrivateResolver {
fn resolve<'a>(&'a self, _host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) })
}
}
impl Resolve for PublicDnsResolver {
fn resolve(&self, name: Name) -> Resolving {
let host = name.as_str().to_string();
Box::pin(async move {
let addresses = tokio::net::lookup_host((host.as_str(), 0))
.await
.map_err(|error| Box::new(error) as Box<dyn std::error::Error + Send + Sync>)?
.collect::<Vec<_>>();
validate_addresses(&addresses).map_err(|_| {
Box::new(io::Error::other("destination rejected by network policy"))
as Box<dyn std::error::Error + Send + Sync>
})?;
Ok(Box::new(addresses.into_iter()) as Addrs)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let task = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = [0_u8; 1024];
let bytes_read = socket.read(&mut request).await.expect("reads request");
assert!(bytes_read > 0);
socket.write_all(response).await.expect("writes response");
});
(
Url::parse(&format!("http://{address}/document")).expect("valid test URL"),
task,
)
}
async fn serve_named(
host: &str,
responses: Vec<&'static [u8]>,
) -> (Url, tokio::task::JoinHandle<Vec<String>>, SocketAddr) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let task = tokio::spawn(async move {
let mut requests = Vec::with_capacity(responses.len());
for response in responses {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = [0_u8; 4096];
let bytes_read = socket.read(&mut request).await.expect("reads request");
requests.push(String::from_utf8_lossy(&request[..bytes_read]).into_owned());
socket.write_all(response).await.expect("writes response");
}
requests
});
(
Url::parse(&format!("http://{host}:{}/document", address.port()))
.expect("valid test URL"),
task,
address,
)
}
struct LoopbackDnsResolver(SocketAddr);
impl Resolve for LoopbackDnsResolver {
fn resolve(&self, _name: Name) -> Resolving {
let address = self.0;
Box::pin(async move { Ok(Box::new(vec![address].into_iter()) as Addrs) })
}
}
struct TestAddressResolver {
blocked_hosts: HashSet<&'static str>,
}
impl AddressResolver for TestAddressResolver {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> {
let blocked = self.blocked_hosts.contains(host);
Box::pin(async move {
let ip = if blocked {
IpAddr::from([127, 0, 0, 1])
} else {
IpAddr::from([8, 8, 8, 8])
};
Ok(vec![SocketAddr::new(ip, port)])
})
}
}
fn policy_checked_fetcher(
address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
) -> MediaFetcher {
MediaFetcher::with_resolvers(
Arc::new(LoopbackDnsResolver(address)),
Arc::new(TestAddressResolver { blocked_hosts }),
)
.expect("test fetcher builds")
}
fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy {
DownloadPolicy {
timeout: Duration::from_secs(1),
max_bytes,
max_redirects,
}
}
#[test]
fn blocks_non_public_addresses() {
for address in [
"0.0.0.1",
"10.0.0.1",
"100.64.0.1",
"127.0.0.1",
"169.254.1.1",
"172.16.0.1",
"192.168.0.1",
"198.18.0.1",
"198.51.100.1",
"203.0.113.1",
"224.0.0.1",
"::1",
"fc00::1",
"fe80::1",
"2001:db8::1",
"::ffff:127.0.0.1",
] {
assert!(is_blocked_ip(address.parse().expect("valid test address")));
}
assert!(!is_blocked_ip(
"8.8.8.8".parse().expect("valid public address")
));
}
#[tokio::test]
async fn fetches_exact_limit_and_normalizes_content_type() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let media = MediaFetcher::for_test(client)
.fetch(url, policy(3, 0))
.await
.expect("download succeeds at exact limit");
server.await.expect("server completes");
assert_eq!(media.bytes, b"abc");
assert_eq!(media.content_type, "application/pdf");
}
#[tokio::test]
async fn rejects_declared_oversize_body() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let error = MediaFetcher::for_test(client)
.fetch(url, policy(2, 0))
.await
.expect_err("oversize body is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
}
#[tokio::test]
async fn rejects_streamed_oversize_body() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let error = MediaFetcher::for_test(client)
.fetch(url, policy(3, 0))
.await
.expect_err("stream crossing limit is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
}
#[tokio::test]
async fn follows_allowed_redirects_and_revalidates_each_destination() {
let (url, server, address) = serve_named(
"public.test",
vec![
b"HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\n\r\n",
b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok",
],
)
.await;
let media = policy_checked_fetcher(address, HashSet::new())
.fetch(url, policy(2, 1))
.await
.expect("redirected fetch succeeds");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 2);
assert!(requests[1].starts_with("GET /final "));
assert_eq!(media.bytes, b"ok");
}
#[tokio::test]
async fn blocks_redirected_private_destination_before_second_request() {
let (url, server, address) = serve_named(
"public.test",
vec![b"HTTP/1.1 302 Found\r\nLocation: http://blocked.test/document\r\nContent-Length: 0\r\n\r\n"],
)
.await;
let error = policy_checked_fetcher(address, HashSet::from(["blocked.test"]))
.fetch(url, policy(10, 1))
.await
.expect_err("private redirect is rejected");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 1);
assert!(matches!(error, MediaError::BlockedUrl));
}
#[tokio::test]
async fn enforces_total_timeout() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (_socket, _) = listener.accept().await.expect("accepts request");
tokio::time::sleep(Duration::from_millis(100)).await;
});
let url = Url::parse(&format!("http://public.test:{}/document", address.port()))
.expect("valid test URL");
let error = policy_checked_fetcher(address, HashSet::new())
.fetch(
url,
DownloadPolicy {
timeout: Duration::from_millis(20),
max_bytes: 10,
max_redirects: 0,
},
)
.await
.expect_err("fetch times out");
server.await.expect("server completes");
assert!(matches!(error, MediaError::Timeout));
}
#[tokio::test]
async fn document_client_does_not_send_ambient_credentials() {
let (url, server, address) = serve_named(
"public.test",
vec![b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"],
)
.await;
policy_checked_fetcher(address, HashSet::new())
.fetch(url, policy(2, 0))
.await
.expect("fetch succeeds");
let requests = server.await.expect("server completes");
assert!(!requests[0].to_ascii_lowercase().contains("authorization:"));
assert!(!requests[0].to_ascii_lowercase().contains("api-key:"));
}
#[tokio::test]
async fn rejects_url_credentials_before_network_access() {
let fetcher = MediaFetcher::new().expect("media fetcher builds");
let url =
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
assert!(matches!(
fetcher.validate_url(&url).await,
Err(MediaError::BlockedUrl)
));
}
}

214
litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs Normal file
View file

@ -0,0 +1,214 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER};
use crate::ocr::OcrClient;
use crate::ocr::codecs::document_intelligence::{
self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams,
};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
use crate::ocr::wire::DecodedOcrResponse;
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
mod polling;
const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
#[derive(Clone, Debug)]
pub(crate) struct AzureDocumentIntelligenceAdapter;
impl OcrAdapter for AzureDocumentIntelligenceAdapter {
type ProviderResponse = AzureDocumentIntelligenceOperation;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = map_ocr_params(request)?;
let config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let endpoint = nonblank(request.connection.api_base.clone())
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
.ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?;
let url = get_complete_url(&endpoint, &request.model, &params)?;
let body = document_intelligence::transform_ocr_request(request.document.clone())?;
transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
document_intelligence::transform_ocr_response(&request.model, response)
}
async fn read_response(
&self,
client: &OcrClient,
response: reqwest::Response,
url: &str,
headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError> {
polling::read_operation_response(
client.polling_http(),
response,
url,
headers,
&request.connection,
request.response_format()? == OcrResponseFormat::Native,
)
.await
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(
request: &LiteLLMOcrRequest,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let params = document_intelligence::decode_input_params(
request.optional_params.clone(),
"optional_params",
)?;
let crate::ocr::prepare::ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = params;
document_intelligence::map_ocr_params(params)
}
fn get_complete_url(
endpoint: &str,
model: &str,
params: &DocumentIntelligenceParams,
) -> Result<String, OcrError> {
let model = format!("{}:analyze", model_id(model)?);
ApiUrl::parse(endpoint)
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
.map(|url| {
url.append_query_pairs(
[("api-version", AZURE_DI_API_VERSION)]
.into_iter()
.chain(params.pages.iter().map(|pages| ("pages", pages.as_str())))
.chain(
params
.features
.iter()
.map(|features| ("features", features.as_str())),
),
)
.into_string()
})
.map_err(|_| OcrRequestError::RequestField {
path: "api_base".into(),
})
.map_err(OcrError::from)
}
async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization")
|| crate::http_utils::has_header(&connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER)
{
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_DI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(
std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value()))
.chain(connection.extra_headers.clone())
.collect(),
);
}
let token = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?;
super::validate_destination(connection, token.source())?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {}", token.value())))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn model_id(model: &str) -> Result<&str, OcrRequestError> {
let model = model.rsplit('/').next().unwrap_or(model);
if matches!(model, "." | "..") {
return Err(OcrRequestError::DotModel);
}
Ok(model)
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
(AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into())
);
}
}

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