scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

Merge branch 'litellm_internal_staging' into fix/img-token-auto-detail-undercounted

Browse source
This commit is contained in:
ErenAta16 2026-08-01 16:17:32 +03:00 • committed by GitHub
commit 19b826edfb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2627 changed files with 171662 additions and 49207 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
.circleci/config.yml
View file

@ -2731,7 +2731,7 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
# The cimg/python:3.12-browsers image already ships the Chromium system
@ -2742,11 +2742,14 @@ jobs:
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
@ -2777,10 +2780,10 @@ jobs:
name: Seed database
command: |
PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \
-f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
-f tests/e2e/ui/fixtures/seed.sql
- run:
name: Start mock LLM server
command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start LiteLLM proxy
@ -2798,7 +2801,7 @@ jobs:
command: |
LITELLM_LICENSE="$LITELLM_LICENSE" \
uv run --no-sync python -m litellm.proxy.proxy_cli \
--config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \
--config tests/e2e/ui/fixtures/config.yml \
--port 4000
background: true
- run:
@ -2819,15 +2822,15 @@ jobs:
# Forward LITELLM_LICENSE so license.spec.ts can detect that the
# proxy was launched with a license and assert premium_user=true.
command: |
cd ui/litellm-dashboard
cd tests/e2e/ui
LITELLM_LICENSE="$LITELLM_LICENSE" \
npx playwright test --config e2e_tests/playwright.config.ts
npx playwright test --config playwright.config.ts
no_output_timeout: 10m
- store_artifacts:
path: ui/litellm-dashboard/test-results
path: tests/e2e/ui/test-results
destination: e2e-test-results
- store_artifacts:
path: ui/litellm-dashboard/playwright-report
path: tests/e2e/ui/playwright-report
destination: e2e-playwright-report
e2e_ui_testing_server_root_path:
@ -2870,17 +2873,20 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
@ -2902,10 +2908,10 @@ jobs:
name: Seed database
command: |
PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \
-f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
-f tests/e2e/ui/fixtures/seed.sql
- run:
name: Start mock LLM server
command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start LiteLLM proxy under a server root path
@ -2918,7 +2924,7 @@ jobs:
command: |
LITELLM_LICENSE="$LITELLM_LICENSE" \
uv run --no-sync python -m litellm.proxy.proxy_cli \
--config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \
--config tests/e2e/ui/fixtures/config.yml \
--port 4000
background: true
- run:
@ -2937,15 +2943,15 @@ jobs:
- run:
name: Run migration smoke under SERVER_ROOT_PATH
command: |
cd ui/litellm-dashboard
cd tests/e2e/ui
LITELLM_LICENSE="$LITELLM_LICENSE" \
npx playwright test --config e2e_tests/migration.serverRootPath.config.ts
npx playwright test --config migration.serverRootPath.config.ts
no_output_timeout: 10m
- store_artifacts:
path: ui/litellm-dashboard/test-results
path: tests/e2e/ui/test-results
destination: e2e-server-root-path-test-results
- store_artifacts:
path: ui/litellm-dashboard/playwright-report
path: tests/e2e/ui/playwright-report
destination: e2e-server-root-path-playwright-report
build_docker_database_image:

2
.circleci/scripts/classify_changes.sh
View file

@ -8,7 +8,7 @@ has_backend=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
ui/*) has_client=true ;;
ui/* | tests/e2e/ui/*) has_client=true ;;
docs/* | *.md | *.mdx) : ;;
*) has_backend=true ;;
esac

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

@ -30,7 +30,7 @@ body:
id: steps-to-reproduce
attributes:
label: Steps to Reproduce
description: Please provide detailed steps to reproduce this bug(A curl/python code to reproduce the bug)
description: Please provide a numbered list of the exact steps to reproduce this bug (include a curl/python snippet to reproduce it). Number each step (1., 2., 3., ...) in the order you performed them.
placeholder: |
1. config.yaml file/ .env file/ etc.
2. Run the following code...

16
.github/pull_request_template.md vendored
View file

@ -1,3 +1,18 @@
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
This section must be extremely human parsable, comprehensible, and readable: its target audience is humans, not AI agents -->
Problem this solves:
- <blah>
- ...
How it solves it:
- <blah>
- ...
## Relevant issues
<!-- e.g., "Fixes #000" -->
@ -25,6 +40,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
The proof must be completely e2e with no mocks, using, for example, actual LLM calls costing real $. `pytest` commands are not enough
For bug fixes: show reproduction before the fix and passing behavior after
Include the commit hash each proof was captured at, for both the before and the after runs
If the change applies to all three LLM endpoints (/v1/responses, /v1/chat/completions, /v1/messages), include proof for every single one of them, not just one
For new features: show the feature working end-to-end
For UI changes: include before/after screenshots -->

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

@ -80,7 +80,7 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'

5
.github/workflows/auto_update_price_and_context_window.yml vendored
View file

@ -24,9 +24,12 @@ jobs:
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/workflows/auto_update_price_and_context_window_file.py"
- name: Regenerate JSON Schema
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create Pull Request
run: |
git add model_prices_and_context_window.json
git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
gh pr create --title "Update model_prices_and_context_window.json file" \
--body "Automated update for model_prices_and_context_window.json" \

14
.github/workflows/codspeed.yml vendored
View file

@ -5,10 +5,24 @@ on:
branches:
- main
- litellm_internal_staging
paths:
- "litellm/**"
- "tests/benchmarks/**"
- "pyproject.toml"
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
pull_request:
branches:
- main
- litellm_internal_staging
paths:
- "litellm/**"
- "tests/benchmarks/**"
- "pyproject.toml"
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
# Allow CodSpeed to trigger backtest performance analysis
# in order to generate initial data
workflow_dispatch:

61
.github/workflows/create_daily_oss_branch.yml vendored
View file

@ -1,61 +0,0 @@
name: Create Daily OSS Branch
on:
schedule:
- cron: "0 16 * * 1-5" # 9am PT during daylight saving time, weekdays.
workflow_dispatch:
inputs:
date:
description: "Branch date in YYYY_MM_DD format. Defaults to today's UTC date."
required: false
type: string
permissions:
contents: write
jobs:
create-oss-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create dated OSS branch
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
REQUESTED_DATE: ${{ inputs.date }}
run: |
set -euo pipefail
if [ -n "${REQUESTED_DATE}" ]; then
if ! echo "${REQUESTED_DATE}" | grep -Eq '^[0-9]{4}_[0-9]{2}_[0-9]{2}$'; then
echo "::error::date must use YYYY_MM_DD format, got '${REQUESTED_DATE}'"
exit 1
fi
BRANCH_DATE="${REQUESTED_DATE}"
else
BRANCH_DATE="$(date -u +'%Y_%m_%d')"
fi
BRANCH_NAME="litellm_oss_daily_${BRANCH_DATE}"
echo "Creating branch: ${BRANCH_NAME}"
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
git fetch origin main "${BRANCH_NAME}" || true
if git show-ref --verify --quiet "refs/remotes/origin/${BRANCH_NAME}"; then
echo "Branch ${BRANCH_NAME} already exists. Skipping creation."
exit 0
fi
git checkout -b "${BRANCH_NAME}" origin/main
git push "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" "${BRANCH_NAME}"
echo "Successfully created and pushed branch: ${BRANCH_NAME}"

4
.github/workflows/guard-main-branch.yml vendored
View file

@ -31,12 +31,12 @@ jobs:
echo "PR head repo: $HEAD_REPO"
echo "PR head branch: $HEAD_REF"
if [ "$HEAD_REPO" != "$BASE_REPO" ]; then
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead."
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against 'litellm_internal_staging' instead."
exit 1
fi
if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then
echo "Allowed source branch."
exit 0
fi
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead."
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_internal_staging' instead."
exit 1

20
.github/workflows/image-scan.yml vendored
View file

@ -9,6 +9,7 @@ on:
- "litellm_**"
paths:
- docker/Dockerfile.non_root
- tests/proxy_migration_tests/test_offline_image_migration.py
- uv.lock
- ui/litellm-dashboard/package-lock.json
- .github/workflows/image-scan.yml
@ -51,6 +52,23 @@ jobs:
- name: Build runtime image
run: docker build -f docker/Dockerfile.non_root -t litellm-image-scan:${{ github.sha }} .
# The prisma bake must migrate a fresh DB with no egress as an arbitrary
# non-root uid (OpenShift restricted-v2 / air-gapped / readOnlyRootFilesystem).
# `docker run` as the default uid with network hides a broken bake because
# the migration entrypoint exits 0 even when it applied nothing; asserting
# the schema was created is what catches it.
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Verify offline migration as a non-root uid
env:
LITELLM_IMAGE: litellm-image-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
# Scans the whole shipped artifact: OS/apk plus every language package
# baked into the image, including ones no lockfile declares (e.g. prisma's
# vendored node engine) that osv-scan cannot see. osv-scan stays the fast
@ -58,6 +76,8 @@ jobs:
# free OSS, run as a pinned, checksum-verified binary; no GitHub Action
# dependency and no vendor SaaS callout.
- name: Scan image for fixable HIGH/CRITICAL CVEs
env:
GRYPE_MATCH_PYTHON_USING_CPES: "true"
run: |
"$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \
--only-fixed \

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

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

50
.github/workflows/oss_daily_guardrails.yml vendored
View file

@ -1,50 +0,0 @@
name: OSS Daily Guardrails
on:
push:
branches:
- "litellm_oss_daily_20*"
pull_request:
branches:
- "litellm_oss_daily_20*"
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
oss-safe-checks:
name: Run OSS daily safe checks
if: startsWith(github.ref_name, 'litellm_oss_daily_20') || startsWith(github.head_ref, 'litellm_oss_daily_20') || startsWith(github.base_ref, 'litellm_oss_daily_20')
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- 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: Run secret scan test
run: |
uv run --frozen --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
- name: Run Ruff
run: |
uv sync --frozen
cd litellm
uv run --no-sync ruff check .

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

@ -115,6 +115,9 @@ jobs:
- name: check_fastuuid_usage
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
- name: check_e2e_no_raw_requests
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
- name: memory_test
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py

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

@ -104,6 +104,7 @@ jobs:
- name: Check basedpyright budget (delta vs base)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
NODE_OPTIONS: --max-old-space-size=12288
run: |
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"

10
.github/workflows/test-litellm-ui-lint.yml vendored
View file

@ -29,7 +29,15 @@ jobs:
id: changed
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
# base.sha is the base branch tip from when the PR was opened, while
# actions/checkout leaves HEAD on a merge of the PR into the *current*
# base tip. "$BASE_SHA"...HEAD therefore spans every base-branch commit
# landed since, so a PR that touches no UI file still gets linted
# against hundreds of other people's files. Diff the PR head against its
# own merge base instead, which is exactly what this PR changed.
merge_base=$(git merge-base "$BASE_SHA" "$HEAD_SHA")
: > "$RUNNER_TEMP/prettier_files.txt"
: > "$RUNNER_TEMP/eslint_files.txt"
while IFS= read -r f; do
@ -41,7 +49,7 @@ jobs:
*.json | *.css | *.scss | *.md | *.mdx | *.yml | *.yaml | *.html)
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt" ;;
esac
done < <(git diff --name-only --diff-filter=ACMR --relative "$BASE_SHA"...HEAD -- .)
done < <(git diff --name-only --diff-filter=ACMR --relative "$merge_base" "$HEAD_SHA" -- .)
if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then
echo "has_files=true" >> "$GITHUB_OUTPUT"
else

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

@ -0,0 +1,57 @@
name: UI Unit Tests
permissions:
contents: read
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- litellm_internal_staging
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
ui-unit-tests:
runs-on: ubuntu-latest-16-cores
timeout-minutes: 20
defaults:
run:
working-directory: ui/litellm-dashboard
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Setup Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "20"
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dependencies
run: npm ci
- name: Run UI unit tests (Vitest)
env:
CI: "true"
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
if [ -n "$BASE_SHA" ]; then
echo "Pull request: running only tests related to changes since $BASE_SHA"
npm run test -- --run --changed "$BASE_SHA" --passWithNoTests \
--pool forks --poolOptions.forks.maxForks=14
else
echo "Push to $GITHUB_REF_NAME: running the full suite"
npm run test -- --run --pool forks --poolOptions.forks.maxForks=14
fi

9
.github/workflows/test-model-map.yaml vendored
View file

@ -22,3 +22,12 @@ jobs:
- name: Validate model_prices_and_context_window.json
run: |
jq empty model_prices_and_context_window.json
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Check model_prices_and_context_window.schema.json is in sync
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py --check

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

@ -61,5 +61,11 @@ jobs:
- name: Run Clippy
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- name: Run Rust tests
run: cargo test --workspace --locked
- name: Run core tests with Bedrock auth
run: cargo test -p litellm-core --features bedrock-auth --locked

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

@ -27,6 +27,7 @@ jobs:
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models

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

@ -46,6 +46,7 @@ jobs:
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
workers: 2
reruns: 2

151
.github/workflows/test_server_root_path.yml vendored
View file

@ -1,151 +0,0 @@
name: Test Proxy SERVER_ROOT_PATH Routing
permissions:
contents: read
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
jobs:
test-server-root-path:
runs-on: ubuntu-latest
timeout-minutes: 30
strategy:
fail-fast: false
matrix:
root_path: ["/api/v1", "/llmproxy"]
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Free up disk space
run: |
sudo rm -rf /usr/local/lib/android /usr/share/dotnet /opt/ghc /usr/local/share/boost
sudo apt-get clean
df -h /
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
- name: Build Docker image
uses: docker/build-push-action@0adf9959216b96bec444f325f1e493d4aa344497 # v6.14.0
with:
context: .
file: ./docker/Dockerfile.non_root
tags: litellm-test:${{ github.sha }}
load: true
push: false
- name: Start LiteLLM container with SERVER_ROOT_PATH
run: |
docker run -d \
--name litellm-test \
-p 4000:4000 \
-e SERVER_ROOT_PATH="${{ matrix.root_path }}" \
-e LITELLM_MASTER_KEY="sk-1234" \
litellm-test:${{ github.sha }} \
--detailed_debug
- name: Wait for container to be healthy
run: |
echo "Waiting for LiteLLM to start..."
max_attempts=30
attempt=0
while [ $attempt -lt $max_attempts ]; do
if docker logs litellm-test 2>&1 | grep -q "Uvicorn running"; then
echo "LiteLLM started successfully"
break
fi
attempt=$((attempt + 1))
echo "Attempt $attempt/$max_attempts - waiting for server to start..."
sleep 2
done
if [ $attempt -eq $max_attempts ]; then
echo "Server failed to start within timeout"
docker logs litellm-test
exit 1
fi
sleep 5
- name: Show container logs
if: always()
run: docker logs litellm-test
- name: Test UI endpoint with root path
run: |
ROOT_PATH="${{ matrix.root_path }}"
echo "Testing UI at: http://localhost:4000${ROOT_PATH}/ui/"
for i in 1 2 3; do
content=$(curl -sL --max-time 5 -H "Authorization: Bearer sk-1234" "http://localhost:4000${ROOT_PATH}/ui/")
if echo "$content" | grep -q -E "(html|<!DOCTYPE|<head|<body)"; then
echo "UI page contains valid HTML content"
exit 0
fi
echo "Attempt $i/3 - no valid HTML, retrying in 5s..."
sleep 5
done
echo "UI page does not contain expected HTML content"
echo "Response: $content"
docker logs litellm-test
exit 1
- name: Setup Node for Playwright
uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version: "20"
- name: Install UI deps and Chromium
working-directory: ui/litellm-dashboard
run: |
retry() {
local attempt=1
local max_attempts=4
until "$@"; do
if [ "$attempt" -ge "$max_attempts" ]; then
echo "Command failed after $attempt attempts: $*"
return 1
fi
echo "Attempt $attempt failed: $*. Retrying in $((attempt * 15))s..."
sleep $((attempt * 15))
attempt=$((attempt + 1))
done
}
npm config set fetch-retries 5
npm config set fetch-retry-mintimeout 20000
npm config set fetch-retry-maxtimeout 120000
retry npm ci
retry npx playwright install --with-deps chromium
- name: Run SERVER_ROOT_PATH redirect e2e
working-directory: ui/litellm-dashboard
env:
SERVER_ROOT_PATH: ${{ matrix.root_path }}
run: npx playwright test --config=e2e_tests/serverRootPath.config.ts
- name: Upload Playwright artifacts on failure
if: failure()
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
with:
name: playwright-trace-${{ strategy.job-index }}
path: ui/litellm-dashboard/test-results/
retention-days: 7
- name: Cleanup
if: always()
run: |
docker stop litellm-test || true
docker rm litellm-test || true

81
.github/workflows/weekly_load_anomaly.yml vendored Normal file
View file

@ -0,0 +1,81 @@
name: "Weekly Load Anomaly Check"
on:
schedule:
- cron: "0 12 * * 6"
workflow_dispatch:
permissions:
contents: read
jobs:
weekly-load-anomaly:
if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 45
services:
postgres:
image: postgres:16.6
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
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-weekly-anomaly-check
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- 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: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Start the proxy
run: |
nohup uv run --no-sync litellm --config tests/e2e/load/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
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 weekly session anomaly test
env:
E2E_WEEKLY_ANOMALY: "1"
run: |
uv run --no-sync pytest tests/e2e/load/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log

4
.gitignore vendored
View file

@ -15,6 +15,9 @@ litellm/rust_bridge/_native*.so
litellm/rust_bridge/_native*.pyd
litellm-rust/target/
# Python package build output
dist/
bun.lockb
**/.DS_Store
.aider*
@ -138,3 +141,4 @@ crash.*.log
.coverage
ui/litellm-dashboard/out/
litellm.log

4
CLAUDE.md
View file

@ -19,7 +19,7 @@ Same thing for bug fixes. The tests should make it so that this specific bug can
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose for internal contributors; external / OSS contributions target the current daily OSS branch instead, named `litellm_oss_daily_YYYY_MM_DD` (a fresh one is cut each weekday, so use the most recent)
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
@ -61,6 +61,8 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages
When working on a PR, keep the PR description in sync with new commits being made
Replies/rebuttals to AI PR review bots must be 15-25 word human-readable replies
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.

2
CONTRIBUTING.md
View file

@ -322,7 +322,7 @@ npm run build
## Submitting Your PR
1. **Push your branch**: `git push origin your-feature-branch`
2. **Create a PR**: Go to GitHub and open a pull request against the current daily OSS branch, named `litellm_oss_daily_YYYY_MM_DD`. A fresh one is cut each weekday, so pick the most recent from the [branch list](https://github.com/BerriAI/litellm/branches/all?query=litellm_oss_daily). Do not target `main`.
2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
3. **Fill out the PR template**: Provide clear description of changes
4. **Wait for review**: Maintainers will review and provide feedback
5. **Address feedback**: Make requested changes and push updates

30
Dockerfile
View file

@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
# Copy full source tree
@ -84,9 +85,12 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
RUN prisma generate --schema=./schema.prisma
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
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
@ -100,7 +104,11 @@ USER root
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}"
ENV PATH="/app/.venv/bin:${PATH}" \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
PRISMA_OFFLINE_MODE=true
# 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
@ -115,16 +123,18 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
# whole /root/.cache drags in the uv build cache (~660 MB, includes a
# setuptools wheel that surfaces as a CVE finding even though it's not
# on the runtime sys.path).
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
# Prisma CLI + engines are baked under /opt/prisma, a fixed path every
# runtime uid can read and that no cache volume mount shadows. The paths are
# pinned via PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH and recorded into the
# generated client at build time, so `prisma migrate deploy` on a fresh
# database needs no npm and no network access (#33650, #24554).
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
EXPOSE 4000/tcp

2
Makefile
View file

@ -176,6 +176,8 @@ lint-ruff-FULL-dev: install-dev
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
else echo "No changed .py files to check."; fi
lint-basedpyright lint-basedpyright-budget-update: export NODE_OPTIONS := --max-old-space-size=12288
lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging

5
backend/routes/allowlist.py
View file

@ -18,6 +18,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/team/",
"/v2/team/",
"/organization/",
"/v2/organization/",
"/customer/",
"/end_user/",
"/sso/",
@ -69,6 +70,10 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/project/",
"/memory/",
"/mcp/",
# Control plane (see the List Endpoints + Tables standard). Every resource
# eventually moves under this prefix, so allowlist it once rather than
# per-resource.
"/management/v1/",
# Spend / analytics
"/spend/",
"/analytics/",

40
basedpyright-code-budget.json
View file

@ -1,30 +1,30 @@
{
"reportAny": {
"limit": 37484
"limit": 31256
},
"reportArgumentType": {
"limit": 2704
"limit": 2645
},
"reportAssignmentType": {
"limit": 330
"limit": 329
},
"reportAttributeAccessIssue": {
"limit": 516
},
"reportCallIssue": {
"limit": 124
"limit": 123
},
"reportConstantRedefinition": {
"limit": 59
},
"reportDeprecated": {
"limit": 326
"limit": 325
},
"reportDuplicateImport": {
"limit": 42
},
"reportExplicitAny": {
"limit": 10397
"limit": 10208
},
"reportFunctionMemberAccess": {
"limit": 11
@ -33,7 +33,7 @@
"limit": 227
},
"reportIncompatibleMethodOverride": {
"limit": 78
"limit": 77
},
"reportIncompatibleVariableOverride": {
"limit": 12
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5900
"limit": 5869
},
"reportMissingTypeArgument": {
"limit": 15903
"limit": 15861
},
"reportMissingTypeStubs": {
"limit": 41
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1085
"limit": 1079
},
"reportOptionalOperand": {
"limit": 0
@ -84,13 +84,13 @@
"limit": 77
},
"reportPrivateUsage": {
"limit": 2438
"limit": 2437
},
"reportRedeclaration": {
"limit": 12
},
"reportReturnType": {
"limit": 225
"limit": 219
},
"reportTypedDictNotRequiredAccess": {
"limit": 27
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45894
"limit": 45357
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40539
"limit": 40477
},
"reportUnknownParameterType": {
"limit": 20403
"limit": 20338
},
"reportUnknownVariableType": {
"limit": 32141
"limit": 32047
},
"reportUnnecessaryCast": {
"limit": 177
},
"reportUnnecessaryComparison": {
"limit": 1025
"limit": 1021
},
"reportUnnecessaryContains": {
"limit": 7
},
"reportUnnecessaryIsInstance": {
"limit": 1209
"limit": 1205
},
"reportUntypedBaseClass": {
"limit": 165
@ -135,10 +135,10 @@
"limit": 33
},
"reportUnusedFunction": {
"limit": 206
"limit": 204
},
"reportUnusedImport": {
"limit": 1005
"limit": 1003
},
"reportUnusedVariable": {
"limit": 1297

325
ci_cd/generate_model_prices_schema.py Normal file
View file

@ -0,0 +1,325 @@
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Optional
import jsonschema
REPO_ROOT = Path(__file__).parent.parent
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
JsonSchema = dict
NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
EXTRA_BOOLEAN_KEYS = frozenset(
{
"gemini_native_audio",
"gemini_audio_only_live",
"uses_embed_content",
"use_openai_responses_path",
"bedrock_converse_supports_strict_tools",
}
)
OBJECT_KEYS: dict[str, JsonSchema] = {
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
"properties": {
"search_context_size_low": NONNEG_NUMBER,
"search_context_size_medium": NONNEG_NUMBER,
"search_context_size_high": NONNEG_NUMBER,
},
"additionalProperties": False,
},
"metadata": {
"type": "object",
"description": "Free-form notes about the entry (e.g. pricing derivation).",
},
"provider_specific_entry": {
"type": "object",
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
},
}
ARRAY_KEYS: dict[str, JsonSchema] = {
"supported_endpoints": {
"type": "array",
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
"items": STRING,
},
"supported_modalities": {
"type": "array",
"description": "Input modalities the model accepts.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
},
"supported_output_modalities": {
"type": "array",
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
"items": STRING,
},
"tiered_pricing": {
"type": "array",
"description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
"items": {
"type": "object",
"properties": {
"range": {
"type": "array",
"description": "[min, max] prompt-token span this tier applies to.",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"max_results_range": {
"type": "array",
"description": "[min, max] result-count span this tier applies to (search models).",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"input_cost_per_query": NONNEG_NUMBER,
},
"additionalProperties": False,
},
},
}
INTEGER_KEYS: dict[str, JsonSchema] = {
"max_tokens": {
**NONNEG_INTEGER,
"description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
},
"max_input_tokens": {
**NONNEG_INTEGER,
"description": "Maximum prompt/context tokens the model accepts.",
},
"max_output_tokens": {
**NONNEG_INTEGER,
"description": "Maximum tokens the model can generate in one response.",
},
"output_vector_size": {
**NONNEG_INTEGER,
"description": "Embedding dimension for embedding models.",
},
"prompt_cache_min_tokens": {
**NONNEG_INTEGER,
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
},
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
}
NUMBER_KEYS: dict[str, JsonSchema] = {
"regional_processing_uplift_multiplier_eu": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
},
"regional_processing_uplift_multiplier_us": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
},
}
COST_DESCRIPTIONS: dict[str, str] = {
"input_cost_per_token": "USD per prompt token.",
"output_cost_per_token": "USD per generated token.",
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
"output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
}
def cost_description(key: str) -> Optional[str]:
if key in COST_DESCRIPTIONS:
return COST_DESCRIPTIONS[key]
if key.endswith("_flex"):
return "Flex service-tier rate for the same-named base field."
if key.endswith("_priority"):
return "Priority service-tier rate for the same-named base field."
if "_above_" in key:
return "Rate applied once the prompt exceeds the token threshold in the field name."
return None
def cost_schema(key: str) -> JsonSchema:
description = cost_description(key)
return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
return {
"litellm_provider": {
"type": "string",
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
},
"mode": {
"type": "string",
"description": "Primary API surface / task type of the model.",
"enum": list(modes),
},
"source": {
"type": "string",
"description": "URL of the provider pricing/model page this entry was taken from.",
},
"deprecation_date": {
"type": "string",
"description": "Date the provider deprecates the model, YYYY-MM-DD.",
"format": "date",
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
},
"web_search_billing_unit": {
"type": "string",
"description": "Whether web search is billed per query or per prompt.",
"enum": ["per_query", "per_prompt"],
},
"bedrock_output_config_effort_ceiling": {
"type": "string",
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"comment": STRING,
"audio_transcription_config": STRING,
}
def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
if key in curated:
return curated[key]
if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
return BOOLEAN
if "cost" in key:
return cost_schema(key)
return None
def build_schema(prices: dict) -> JsonSchema:
entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
if unclassified:
raise SystemExit(
f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
f"Add them to the key tables in {Path(__file__).name} and rerun it."
)
entry_properties = {key: classify(key, modes) for key in all_keys}
return {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "LiteLLM model_prices_and_context_window.json",
"description": (
"Schema for LiteLLM's model price and context window registry "
"(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
"Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
"optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
"All costs are USD per unit. New optional fields are added regularly, so consumers should "
"ignore unknown fields rather than reject them."
),
"type": "object",
"properties": {
"sample_spec": {
"type": "object",
"description": (
"Documentation placeholder illustrating the entry shape; not a real model and not "
"schema-conformant (several values are prose)."
),
},
"fallback_generalizations": {
"type": "object",
"description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
"properties": {
"rules": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": STRING,
"pattern": STRING,
"description": STRING,
},
"required": ["name", "pattern"],
"additionalProperties": True,
},
}
},
"additionalProperties": False,
},
},
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
"$defs": {
"modelEntry": {
"type": "object",
"description": (
"Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
"are optional; boolean capability flags are simply omitted when unknown or false."
),
"required": ["litellm_provider"],
"properties": entry_properties,
"additionalProperties": True,
}
},
}
def render(schema: JsonSchema) -> str:
return json.dumps(schema, indent=2) + "\n"
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(
schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER
)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)
)
def main() -> int:
check = "--check" in sys.argv[1:]
prices = json.loads(PRICES_PATH.read_text())
rendered = render(build_schema(prices))
errors = validation_errors(prices, json.loads(rendered))
if errors:
print(f"{PRICES_PATH.name} does not validate against the generated schema:")
print("\n".join(errors[:20]))
return 1
if not check:
SCHEMA_PATH.write_text(rendered)
print(f"wrote {SCHEMA_PATH}")
return 0
if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
print(
f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
)
return 1
print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
return 0
if __name__ == "__main__":
sys.exit(main())

523
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_genai_otel/grafana_dashboard.json Normal file
View file

@ -0,0 +1,523 @@
{
"annotations": {
"list": []
},
"editable": true,
"fiscalYearStartMonth": 0,
"graphTooltip": 0,
"links": [],
"panels": [
{
"type": "stat",
"title": "Requests",
"gridPos": {
"h": 4,
"w": 6,
"x": 0,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "short",
"decimals": 0,
"color": {
"mode": "fixed",
"fixedColor": "blue"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "sum(increase(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
}
],
"id": 1
},
{
"type": "stat",
"title": "Spend",
"description": "LiteLLM's computed cost for the selected window, from gen_ai.usage.cost",
"gridPos": {
"h": 4,
"w": 6,
"x": 6,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "currencyUSD",
"decimals": 4,
"color": {
"mode": "fixed",
"fixedColor": "green"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "sum(increase(gen_ai_usage_cost_USD_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
}
],
"id": 2
},
{
"type": "stat",
"title": "Tokens",
"gridPos": {
"h": 4,
"w": 6,
"x": 12,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "short",
"decimals": 0,
"color": {
"mode": "fixed",
"fixedColor": "purple"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "sum(increase(gen_ai_client_token_usage_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
}
],
"id": 3
},
{
"type": "stat",
"title": "p95 request duration",
"gridPos": {
"h": 4,
"w": 6,
"x": 18,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"decimals": 2,
"color": {
"mode": "fixed",
"fixedColor": "orange"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "histogram_quantile(0.95, sum by (le) (rate(gen_ai_client_operation_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range])))"
}
],
"id": 4
},
{
"type": "timeseries",
"title": "Request rate by model",
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 4
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "reqpm",
"custom": {
"lineWidth": 2,
"fillOpacity": 8,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "sum by (gen_ai_request_model) (rate(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 60"
}
],
"id": 5
},
{
"type": "timeseries",
"title": "Spend rate by model",
"description": "USD per hour, derived from the gen_ai.usage.cost histogram",
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 4
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "currencyUSD",
"custom": {
"lineWidth": 2,
"fillOpacity": 8,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "sum by (gen_ai_request_model) (rate(gen_ai_usage_cost_USD_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 3600"
}
],
"id": 6
},
{
"type": "timeseries",
"title": "Tokens per minute by model and type",
"description": "gen_ai.client.token.usage split by the gen_ai.token.type attribute",
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 12
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "short",
"custom": {
"lineWidth": 2,
"fillOpacity": 8,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}} {{gen_ai_token_type}}",
"expr": "sum by (gen_ai_request_model, gen_ai_token_type) (rate(gen_ai_client_token_usage_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 60"
}
],
"id": 7
},
{
"type": "timeseries",
"title": "p95 request duration by model",
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 12
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"custom": {
"lineWidth": 2,
"fillOpacity": 0,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_client_operation_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
}
],
"id": 8
},
{
"type": "timeseries",
"title": "p95 time to first token (streaming)",
"description": "gen_ai.server.time_to_first_token, recorded only for streaming requests",
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 20
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"custom": {
"lineWidth": 2,
"fillOpacity": 0,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_server_time_to_first_token_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
}
],
"id": 9
},
{
"type": "timeseries",
"title": "p95 provider generation time",
"description": "gen_ai.client.response.duration, upstream generation time excluding LiteLLM overhead",
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 20
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"custom": {
"lineWidth": 2,
"fillOpacity": 0,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_client_response_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
}
],
"id": 10
}
],
"preload": false,
"refresh": "30s",
"schemaVersion": 42,
"tags": [
"litellm",
"genai",
"opentelemetry"
],
"templating": {
"list": [
{
"name": "datasource",
"label": "Prometheus",
"type": "datasource",
"query": "prometheus",
"current": {},
"hide": 0
},
{
"name": "service",
"label": "Service",
"type": "query",
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"query": "label_values(gen_ai_client_operation_duration_seconds_count, service_name)",
"refresh": 2,
"includeAll": true,
"multi": true,
"current": {
"text": "All",
"value": "$__all"
}
},
{
"name": "model",
"label": "Model",
"type": "query",
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"query": "label_values(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\"}, gen_ai_request_model)",
"refresh": 2,
"includeAll": true,
"multi": true,
"current": {
"text": "All",
"value": "$__all"
}
}
]
},
"time": {
"from": "now-1h",
"to": "now"
},
"timepicker": {},
"timezone": "browser",
"title": "LiteLLM GenAI (OpenTelemetry)",
"uid": "litellm-genai-otel",
"weekStart": ""
}

35
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_genai_otel/readme.md Normal file
View file

@ -0,0 +1,35 @@
# LiteLLM GenAI dashboard (OpenTelemetry metrics)
Dashboard for the `gen_ai.*` metrics the OpenTelemetry v2 integration emits, as opposed to the `litellm_*` Prometheus metrics the other dashboards in this folder chart.
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source. Panels: request count, spend, token count, p95 duration, request rate by model, spend rate per hour by model, tokens per minute split by input and output, p95 duration by model, p95 time to first token, and p95 provider generation time. Template variables for data source, service, and model.
## Pre-requisites
Metrics are off by default. In the proxy environment:
```shell
LITELLM_OTEL_V2=true
LITELLM_OTEL_INTEGRATION_ENABLE_METRICS=true
OTEL_EXPORTER="otlp_http"
OTEL_ENDPOINT="<your OTLP endpoint>"
```
You also need the metric attribute filter, or the panels will plot flat lines at zero. LiteLLM's default attribute set includes per-request fields, so nearly every request lands in its own time series with a single sample, and `rate()` has nothing to compute over:
```yaml title="config.yaml"
callback_settings:
otel:
attributes:
include_list:
- gen_ai.operation.name
- gen_ai.system
- gen_ai.request.model
- gen_ai.framework
```
See [Grafana Cloud](https://docs.litellm.ai/docs/observability/grafana_cloud) for the full setup, and [OpenTelemetry v2](https://docs.litellm.ai/docs/observability/opentelemetry_v2#metrics) for the metric reference.
## Note on Grafana's AI Observability integration
Grafana Cloud ships prebuilt GenAI dashboards that query these same metric names, so they look like a drop-in alternative to this one. They are not: twenty of their twenty-two panels filter on `telemetry_sdk_name="openlit"`, a label LiteLLM does not carry and cannot be configured to add, so those panels stay empty.

4
cookbook/litellm_proxy_server/grafana_dashboard/readme.md
View file

@ -2,6 +2,10 @@
This folder contains the `json` for creating Grafana Dashboards
## [LiteLLM GenAI Dashboard (OpenTelemetry)](./dashboard_genai_otel)
Charts the `gen_ai.*` metrics from the OpenTelemetry v2 integration: spend, tokens, request rate, and latency percentiles by model. Separate from the dashboards below, which chart the `litellm_*` Prometheus metrics.
## [LiteLLM v2 Dashboard](./dashboard_v2)
<img width="1316" alt="grafana_1" src="https://github.com/user-attachments/assets/d0df802d-0cb9-4906-a679-941c547789ab">

46
db_scripts/backfill_daily_tool_spend.sql Normal file
View file

@ -0,0 +1,46 @@
-- One-shot backfill of the LiteLLM_DailyToolSpend rollup from the per-request
-- LiteLLM_SpendLogToolIndex x LiteLLM_SpendLogs tables.
--
-- This is an opt-in, manual operation. New deployments do not need it: the
-- rollup is written at request time from the moment the release is deployed.
-- Run it only if you want the Cost Optimization "Spend by tool" card to show
-- history from before the deploy, and only once.
--
-- IMPORTANT caveats before running:
--
-- 1. Pre-deploy index rows may include tools that were merely DECLARED in a
-- request body but never invoked (the release this ships with stops
-- recording those). For agentic clients that declare many tools per
-- request, backfilled history attributes each request's full spend to
-- every declared tool, overstating per-tool spend. Post-deploy rows do not
-- have this problem. If your traffic is mostly such clients, consider not
-- backfilling.
--
-- 2. Coverage is bounded by spend-log retention: rows older than
-- maximum_spend_logs_retention_period are already gone.
--
-- 3. Replace the cutover timestamp below with the time you deployed the
-- release, so backfilled per-request rows cannot double-count on top of
-- rollup rows the new writer already created. ON CONFLICT DO NOTHING is a
-- second guard for (date, tool_name) buckets the writer already touched:
-- such buckets keep the writer's numbers and skip the backfill's.
--
-- Usage:
-- psql "$DATABASE_URL" -v cutover="'2026-07-25T00:00:00Z'" -f db_scripts/backfill_daily_tool_spend.sql
SET TIME ZONE 'UTC';
INSERT INTO "LiteLLM_DailyToolSpend" (date, tool_name, spend, total_tokens, request_count, created_at, updated_at)
SELECT
to_char(ti.start_time, 'YYYY-MM-DD') AS date,
ti.tool_name,
COALESCE(SUM(sl.spend), 0) AS spend,
COALESCE(SUM(sl.total_tokens), 0) AS total_tokens,
COUNT(*) AS request_count,
now() AS created_at,
now() AS updated_at
FROM "LiteLLM_SpendLogToolIndex" ti
JOIN "LiteLLM_SpendLogs" sl ON sl.request_id = ti.request_id
WHERE ti.start_time < :cutover::timestamptz
GROUP BY 1, 2
ON CONFLICT (date, tool_name) DO NOTHING;

BIN
dist/litellm-1.79.1.tar.gz vendored
View file

Binary file not shown.

32
docker/Dockerfile.database
View file

@ -62,6 +62,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
# Copy full source tree
@ -82,9 +83,12 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
RUN prisma generate --schema=./schema.prisma
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
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
@ -97,7 +101,11 @@ USER root
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}"
ENV PATH="/app/.venv/bin:${PATH}" \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
PRISMA_OFFLINE_MODE=true
# 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
@ -112,16 +120,20 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
# Only the Prisma subdirs: the whole /root/.cache drags in the uv build cache.
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
# Prisma CLI + engines are baked under /opt/prisma, a fixed path every
# runtime uid can read and that no cache volume mount shadows (unlike
# /app/.cache or $HOME/.cache in readOnlyRootFilesystem + emptyDir setups).
# The paths are pinned via PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH and
# recorded into the generated client at build time, so `prisma migrate
# deploy` on a fresh database needs no npm and no network access
# (#33650, #24554).
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
EXPOSE 4000/tcp

40
docker/Dockerfile.non_root
View file

@ -54,7 +54,6 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
PATH="/app/.venv/bin:${PATH}" \
LITELLM_NON_ROOT=true \
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
XDG_CACHE_HOME=/app/.cache
# Copy dependency metadata first for layer caching
@ -69,6 +68,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
# Copy full source tree
@ -95,6 +95,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3 \
--no-sources-package litellm-proxy-extras; \
else \
@ -103,10 +104,13 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3; \
fi
RUN prisma generate --schema=./schema.prisma
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
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
@ -127,8 +131,6 @@ RUN for i in 1 2 3; do \
# the rest of the builder's /app is source and build metadata that must not
# ship (manifest-scanning tools attribute everything in it to this image).
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
# Prisma caches live under /app/.cache here (XDG_CACHE_HOME /
# PRISMA_BINARY_CACHE_DIR) so the runtime prisma generate finds them.
COPY --from=builder /app/.venv /app/.venv
COPY --from=builder /app/docker /app/docker
COPY --from=builder /app/schema.prisma /app/schema.prisma
@ -138,21 +140,35 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
COPY --from=builder /app/.cache /app/.cache
# Prisma CLI + engines are baked under /opt/prisma, a fixed path every runtime
# uid can read and that no cache volume mount shadows (unlike /app/.cache or
# $HOME/.cache under readOnlyRootFilesystem + emptyDir or arbitrary-uid setups).
# PRISMA_CLI_QUERY_ENGINE_TYPE=binary makes the CLI use the baked binary query
# engine directly, so `prisma migrate deploy` on a fresh database needs no npm
# and no network access; without it the CLI looks for the library engine, which
# prisma stopped baking, and falls back to a download that fails offline or as a
# non-writable uid (#33650, #24554).
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
# XDG_CACHE_HOME is intentionally left unset so it falls back to $HOME/.cache
# (/app/.cache, writable by the runtime uid). The prisma bake at the read-only
# /opt/prisma is anchored by PRISMA_BINARY_CACHE_DIR / PRISMA_CLI_PATH, so
# nothing needs XDG to point there; pointing it at the read-only bake would
# deny any XDG-aware library that writes a cache at runtime.
ENV PATH="/app/.venv/bin:${PATH}" \
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
PRISMA_CLI_PATH=/opt/prisma/binaries/node_modules/.bin/prisma \
PRISMA_CLI_QUERY_ENGINE_TYPE=binary \
HOME=/app \
LITELLM_NON_ROOT=true \
XDG_CACHE_HOME=/app/.cache \
PRISMA_SKIP_POSTINSTALL_GENERATE=1 \
PRISMA_HIDE_UPDATE_MESSAGE=1 \
PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \
PRISMA_OFFLINE_MODE=true
RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
RUN mkdir -p /nonexistent /app/.cache /var/lib/litellm/assets /var/lib/litellm/ui && \
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
chown -R nobody:nogroup "$PRISMA_PATH" && \
@ -165,12 +181,14 @@ RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g=u "$LITELLM_PROXY_EXTRAS_PATH" || true && \
chmod -R g+w "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w "$LITELLM_PROXY_EXTRAS_PATH" || true && \
chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets /app/.cache
chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets && \
chmod -R a+rX /opt/prisma && \
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1
USER 65534
RUN prisma generate --schema=./schema.prisma
EXPOSE 4000/tcp
ENTRYPOINT ["/app/docker/prod_entrypoint.sh"]

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

@ -215,7 +215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
if result:
return LiteLLM_ManagedFileTable(**result)
return LiteLLM_ManagedFileTable.model_validate(result)
## CHECK DB
db_object = await self.prisma_client.db.litellm_managedfiletable.find_first(
@ -223,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
if db_object:
return LiteLLM_ManagedFileTable(**db_object.model_dump())
return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump())
return None
async def delete_unified_file_id(
@ -316,32 +316,40 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter}
if after:
where_clause["id"] = {"gt": after}
cursor_row = (
await self.prisma_client.db.litellm_managedobjecttable.find_first(
where={**where_clause, "unified_object_id": after}
)
)
if cursor_row is None:
raise HTTPException(
status_code=400,
detail=f"Invalid 'after' cursor: no batch found with id '{after}'.",
)
fetch_limit = limit or 20
if target_model_names:
# Oversample so post-fetch model-name filtering still has enough rows.
fetch_limit = max(fetch_limit * 3, 100)
page_size = limit or 20
cursor_args: Dict[str, Any] = (
{"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
)
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where=where_clause,
take=fetch_limit,
order={"created_at": "desc"},
take=page_size + 1,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
)
batch_objects: List[LiteLLMBatch] = []
for batch in batches:
try:
# Stop once we have enough after filtering
if len(batch_objects) >= (limit or 20):
break
has_more = len(batches) > page_size
batch_objects: List[LiteLLMBatch] = []
for batch in batches[:page_size]:
try:
batch_data = (
json.loads(batch.file_object)
if isinstance(batch.file_object, str)
else batch.file_object
)
batch_obj = LiteLLMBatch(**batch_data)
batch_obj = LiteLLMBatch.model_validate(batch_data)
batch_obj.id = batch.unified_object_id
batch_objects.append(batch_obj)
@ -351,9 +359,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
continue
return build_list_page(
batch_objects, has_more=len(batch_objects) == (limit or 20)
)
return build_list_page(batch_objects, has_more=has_more)
async def get_user_created_file_ids(
self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str]
@ -376,7 +382,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"flat_model_file_ids": {"hasSome": model_object_ids},
}
)
return [OpenAIFileObject(**file_object.file_object) for file_object in file_ids]
return [OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids]
async def check_managed_file_id_access(
self, data: Dict, user_api_key_dict: UserAPIKeyAuth

191
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -11,7 +11,8 @@ Endpoints for /project operations
#### PROJECT MANAGEMENT ####
import json
from typing import List, Optional, Union
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING
from fastapi import APIRouter, Depends, HTTPException, Request
@ -25,15 +26,24 @@ from litellm.proxy.management_helpers.utils import (
)
from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma.actions import LiteLLM_TeamTableActions
router = APIRouter()
def _team_table(prisma_client: PrismaClient) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]":
team_table: LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable] = prisma_client.db.litellm_teamtable
return team_table
async def _check_user_permission_for_project(
user_api_key_dict: UserAPIKeyAuth,
team_id: Optional[str],
team_id: str | None,
prisma_client: PrismaClient,
require_admin: bool = False,
team_object: Optional[LiteLLM_TeamTable] = None,
team_object: LiteLLM_TeamTable | None = None,
) -> bool:
"""
Check if user has permission to manage a project.
@ -57,9 +67,7 @@ async def _check_user_permission_for_project(
team = team_object
if team is None:
team = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
team = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
if team and team.admins:
return user_api_key_dict.user_id in team.admins
@ -70,9 +78,9 @@ async def _check_user_permission_for_project(
async def _validate_team_exists(
team_id: str,
prisma_client: PrismaClient,
):
) -> "prisma_models.LiteLLM_TeamTable":
"""Validate that a team exists. Returns the team row."""
team = await prisma_client.db.litellm_teamtable.find_unique(
team = await _team_table(prisma_client).find_unique(
where={"team_id": team_id},
)
@ -89,7 +97,7 @@ async def _validate_team_exists(
def _check_team_project_limits(
team_object: LiteLLM_TeamTable,
data: Union[NewProjectRequest, UpdateProjectRequest],
data: NewProjectRequest | UpdateProjectRequest,
) -> None:
"""
Check that project limits respect its parent Team's limits.
@ -108,16 +116,12 @@ def _check_team_project_limits(
if data.max_budget is not None and data.max_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"max_budget cannot be negative. Received: {data.max_budget}"
},
detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"},
)
if data.soft_budget is not None and data.soft_budget < 0:
raise HTTPException(
status_code=400,
detail={
"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"
},
detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"},
)
# --- soft_budget < max_budget ---
@ -131,7 +135,7 @@ def _check_team_project_limits(
)
# --- Validate project models are a subset of team models ---
project_models = getattr(data, "models", None)
project_models = data.models
team_models = team_object.models or []
if project_models and len(team_models) > 0:
# If team has 'all-proxy-models', skip validation as it allows all models
@ -148,11 +152,7 @@ def _check_team_project_limits(
# --- Validate project max_budget <= team max_budget ---
# Team stores budget fields directly (max_budget, tpm_limit, rpm_limit)
# unlike Project which uses a separate LiteLLM_BudgetTable relation
if (
data.max_budget is not None
and team_object.max_budget is not None
and data.max_budget > team_object.max_budget
):
if data.max_budget is not None and team_object.max_budget is not None and data.max_budget > team_object.max_budget:
raise HTTPException(
status_code=400,
detail={
@ -161,11 +161,7 @@ def _check_team_project_limits(
)
# --- Validate project tpm_limit <= team tpm_limit ---
if (
data.tpm_limit is not None
and team_object.tpm_limit is not None
and data.tpm_limit > team_object.tpm_limit
):
if data.tpm_limit is not None and team_object.tpm_limit is not None and data.tpm_limit > team_object.tpm_limit:
raise HTTPException(
status_code=400,
detail={
@ -174,11 +170,7 @@ def _check_team_project_limits(
)
# --- Validate project rpm_limit <= team rpm_limit ---
if (
data.rpm_limit is not None
and team_object.rpm_limit is not None
and data.rpm_limit > team_object.rpm_limit
):
if data.rpm_limit is not None and team_object.rpm_limit is not None and data.rpm_limit > team_object.rpm_limit:
raise HTTPException(
status_code=400,
detail={
@ -189,19 +181,19 @@ def _check_team_project_limits(
async def _create_budget_for_project(
data: NewProjectRequest,
user_id: Optional[str],
user_id: str | None,
litellm_proxy_admin_name: str,
prisma_client: PrismaClient,
) -> str:
"""Create a budget for the project and return budget_id."""
budget_params = LiteLLM_BudgetTable.model_fields.keys()
_json_data = data.json(exclude_none=True)
_json_data: Mapping[str, object] = data.json(exclude_none=True)
_budget_data = {k: v for k, v in _json_data.items() if k in budget_params}
budget_row = LiteLLM_BudgetTable(**_budget_data)
budget_row = LiteLLM_BudgetTable.model_validate(_budget_data)
new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
_budget = await prisma_client.db.litellm_budgettable.create(
_budget: prisma_models.LiteLLM_BudgetTable = await prisma_client.db.litellm_budgettable.create(
data={
**new_budget,
"created_by": user_id or litellm_proxy_admin_name,
@ -214,8 +206,8 @@ async def _create_budget_for_project(
async def _set_project_object_permission(
data: NewProjectRequest,
prisma_client: Optional[PrismaClient],
) -> Optional[str]:
prisma_client: PrismaClient | None,
) -> str | None:
"""
Creates the LiteLLM_ObjectPermissionTable record for the project.
Returns the object_permission_id if created, otherwise None.
@ -224,7 +216,7 @@ async def _set_project_object_permission(
return None
if data.object_permission is not None:
created_object_permission = (
created_object_permission: prisma_models.LiteLLM_ObjectPermissionTable = (
await prisma_client.db.litellm_objectpermissiontable.create(
data=data.object_permission.model_dump(exclude_none=True),
)
@ -344,8 +336,7 @@ async def new_project(
raise HTTPException(
status_code=403,
detail={
"error": "Only premium users can add tags to projects. "
+ CommonProxyErrors.not_premium_user.value
"error": "Only premium users can add tags to projects. " + CommonProxyErrors.not_premium_user.value
},
)
@ -353,8 +344,7 @@ async def new_project(
raise HTTPException(
status_code=403,
detail={
"error": "Project management is an enterprise feature. "
+ CommonProxyErrors.not_premium_user.value
"error": "Project management is an enterprise feature. " + CommonProxyErrors.not_premium_user.value
},
)
@ -375,13 +365,11 @@ async def new_project(
)
# Validate team exists and get team object with budget
team_object = await _validate_team_exists(
team_id=data.team_id, prisma_client=prisma_client
)
team_object = await _validate_team_exists(team_id=data.team_id, prisma_client=prisma_client)
# Validate project limits against team limits
_check_team_project_limits(
team_object=LiteLLM_TeamTable(**team_object.model_dump()),
team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()),
data=data,
)
@ -391,7 +379,7 @@ async def new_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
team_object=LiteLLM_TeamTable(**team_object.model_dump()),
team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()),
)
if not has_permission:
@ -449,17 +437,13 @@ async def new_project(
value=getattr(data, field),
)
new_project_row = prisma_client.jsonify_object(
project_row.json(exclude_none=True)
)
new_project_row = prisma_client.jsonify_object(project_row.json(exclude_none=True))
# Remove budget fields (following organization_endpoints.py pattern)
new_project_row = _remove_budget_fields_from_project_data(new_project_row)
verbose_proxy_logger.info(
f"new_project_row: {json.dumps(new_project_row, indent=2)}"
)
response = await prisma_client.db.litellm_projecttable.create(
verbose_proxy_logger.info(f"new_project_row: {json.dumps(new_project_row, indent=2)}")
response: prisma_models.LiteLLM_ProjectTable = await prisma_client.db.litellm_projecttable.create(
data={
**new_project_row, # type: ignore
},
@ -469,9 +453,7 @@ async def new_project(
return response
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.project_endpoints.new_project(): Exception occured - {}".format(
str(e)
)
"litellm.proxy.management_endpoints.project_endpoints.new_project(): Exception occured - {}".format(str(e))
)
raise handle_exception_on_proxy(e)
@ -539,8 +521,7 @@ async def update_project(
raise HTTPException(
status_code=403,
detail={
"error": "Only premium users can add tags to projects. "
+ CommonProxyErrors.not_premium_user.value
"error": "Only premium users can add tags to projects. " + CommonProxyErrors.not_premium_user.value
},
)
@ -548,8 +529,7 @@ async def update_project(
raise HTTPException(
status_code=403,
detail={
"error": "Project management is an enterprise feature. "
+ CommonProxyErrors.not_premium_user.value
"error": "Project management is an enterprise feature. " + CommonProxyErrors.not_premium_user.value
},
)
@ -576,9 +556,9 @@ async def update_project(
)
# Fetch existing project
existing_project = await prisma_client.db.litellm_projecttable.find_unique(
where={"project_id": data.project_id}
)
existing_project: (
prisma_models.LiteLLM_ProjectTable | None
) = await prisma_client.db.litellm_projecttable.find_unique(where={"project_id": data.project_id})
if existing_project is None:
raise ProxyException(
@ -595,9 +575,7 @@ async def update_project(
target_team_id = data.team_id or existing_project.team_id
target_team_obj = None
if target_team_id is not None:
target_team_obj = await _validate_team_exists(
team_id=target_team_id, prisma_client=prisma_client
)
target_team_obj = await _validate_team_exists(team_id=target_team_id, prisma_client=prisma_client)
has_permission = await _check_user_permission_for_project(
user_api_key_dict=user_api_key_dict,
@ -620,32 +598,26 @@ async def update_project(
team_id=data.team_id,
prisma_client=prisma_client,
team_object=(
LiteLLM_TeamTable(**target_team_obj.model_dump())
if target_team_obj
else None
LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None
),
)
if not can_assign_to_target:
raise HTTPException(
status_code=403,
detail={
"error": "Cannot reassign project to a team you are not an admin of"
},
detail={"error": "Cannot reassign project to a team you are not an admin of"},
)
# Validate project limits against team limits
if target_team_obj is not None:
_check_team_project_limits(
team_object=LiteLLM_TeamTable(**target_team_obj.model_dump()),
team_object=LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()),
data=data,
)
# Prepare update data
update_data = data.json(exclude_none=True, exclude={"project_id"})
update_data = prisma_client.jsonify_object(update_data)
update_data["updated_by"] = (
user_api_key_dict.user_id or litellm_proxy_admin_name
)
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
# Handle budget updates
budget_fields = LiteLLM_BudgetTable.model_fields.keys()
@ -671,21 +643,17 @@ async def update_project(
if existing_project.object_permission_id:
# Update existing permission
await prisma_client.db.litellm_objectpermissiontable.update(
where={
"object_permission_id": existing_project.object_permission_id
},
where={"object_permission_id": existing_project.object_permission_id},
data=object_permission_data,
)
else:
# Create new permission
created_permission = (
created_permission: prisma_models.LiteLLM_ObjectPermissionTable = (
await prisma_client.db.litellm_objectpermissiontable.create(
data=object_permission_data,
)
)
update_data["object_permission_id"] = (
created_permission.object_permission_id
)
update_data["object_permission_id"] = created_permission.object_permission_id
# Handle metadata fields
for field in LiteLLM_ManagementEndpoint_MetadataFields:
@ -698,7 +666,7 @@ async def update_project(
update_data = _remove_budget_fields_from_project_data(update_data)
# Update project
updated_project = await prisma_client.db.litellm_projecttable.update(
updated_project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.update(
where={"project_id": data.project_id},
data=update_data,
include={"litellm_budget_table": True, "object_permission": True},
@ -718,7 +686,7 @@ async def update_project(
"/project/delete",
tags=["project management"],
dependencies=[Depends(user_api_key_auth)],
response_model=List[LiteLLM_ProjectTable],
response_model=list[LiteLLM_ProjectTable],
)
@management_endpoint_wrapper
async def delete_project(
@ -749,8 +717,7 @@ async def delete_project(
raise HTTPException(
status_code=403,
detail={
"error": "Project management is an enterprise feature. "
+ CommonProxyErrors.not_premium_user.value
"error": "Project management is an enterprise feature. " + CommonProxyErrors.not_premium_user.value
},
)
@ -778,9 +745,7 @@ async def delete_project(
for project_id in data.project_ids:
# Check if project exists
existing_project = await prisma_client.db.litellm_projecttable.find_unique(
where={"project_id": project_id}
)
existing_project = await prisma_client.db.litellm_projecttable.find_unique(where={"project_id": project_id})
if existing_project is None:
raise ProxyException(
@ -791,11 +756,9 @@ async def delete_project(
)
# Check if there are any keys associated with this project
associated_keys = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"project_id": project_id}
)
)
associated_keys: Sequence[
prisma_models.LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(where={"project_id": project_id})
if len(associated_keys) > 0:
raise ProxyException(
@ -806,9 +769,9 @@ async def delete_project(
)
# Delete the project
deleted_project = await prisma_client.db.litellm_projecttable.delete(
where={"project_id": project_id}
)
deleted_project: (
prisma_models.LiteLLM_ProjectTable | None
) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id})
deleted_projects.append(deleted_project)
@ -854,7 +817,7 @@ async def project_info(
)
# Fetch project
project = await prisma_client.db.litellm_projecttable.find_unique(
project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.find_unique(
where={"project_id": project_id},
include={"litellm_budget_table": True, "object_permission": True},
)
@ -872,17 +835,11 @@ async def project_info(
is_team_member = False
if project.team_id and user_api_key_dict.user_id:
team = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": project.team_id}
)
team = await _team_table(prisma_client).find_unique(where={"team_id": project.team_id})
if team:
caller_user_id = user_api_key_dict.user_id
for m in team.members_with_roles or []:
m_user_id = (
m.get("user_id")
if isinstance(m, dict)
else getattr(m, "user_id", None)
)
m_user_id = m.get("user_id") if isinstance(m, dict) else getattr(m, "user_id", None)
if m_user_id == caller_user_id:
is_team_member = True
break
@ -896,9 +853,7 @@ async def project_info(
return project
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.project_endpoints.project_info(): Exception occured - {}".format(
str(e)
)
"litellm.proxy.management_endpoints.project_endpoints.project_info(): Exception occured - {}".format(str(e))
)
raise handle_exception_on_proxy(e)
@ -907,7 +862,7 @@ async def project_info(
"/project/list",
tags=["project management"],
dependencies=[Depends(user_api_key_auth)],
response_model=List[LiteLLM_ProjectTable],
response_model=list[LiteLLM_ProjectTable],
)
async def list_projects(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -932,21 +887,19 @@ async def list_projects(
# If proxy admin, get all projects
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
projects = await prisma_client.db.litellm_projecttable.find_many(
projects: Sequence[
prisma_models.LiteLLM_ProjectTable
] = await prisma_client.db.litellm_projecttable.find_many(
include={"litellm_budget_table": True, "object_permission": True}
)
else:
# Look up the user's team memberships via the reverse-index on
# LiteLLM_UserTable.teams (maintained by team_member_add alongside
# members_with_roles). This avoids a full scan of all team rows.
user_record = await prisma_client.db.litellm_usertable.find_unique(
user_record: prisma_models.LiteLLM_UserTable | None = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
)
user_team_ids = (
user_record.teams
if user_record is not None and user_record.teams
else []
)
user_team_ids: Sequence[str] = user_record.teams if user_record is not None and user_record.teams else []
projects = await prisma_client.db.litellm_projecttable.find_many(
where={"team_id": {"in": user_team_ids}},

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.51"
version = "0.1.52"
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.51"
version = "0.1.52"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

2
gateway/Dockerfile
View file

@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--python python3
# Stage 2 — copy source and install the project + workspace members.
@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--python python3
RUN mkdir -p /home/nonroot && \

1
gateway/routes/allowlist.py
View file

@ -54,6 +54,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/messages",
"/v1/skills",
"/v1/a2a/",
"/a2a/",
# LiteLLM-native LLM surface
"/v1/rerank",
"/v2/rerank",

4
helm/litellm-helm/Chart.lock
View file

@ -5,5 +5,5 @@ dependencies:
- name: redis
repository: oci://registry-1.docker.io/bitnamicharts
version: 18.19.1
digest: sha256:8660fe6287f9941d08c0902f3f13731079b8cecd2a5da2fbc54e5b7aae4a6f62
generated: "2024-03-10T02:28:52.275022+05:30"
digest: sha256:38962e231f6596b93f82a8412bbe4cf5de696caecf5775dfbbd163383eb1c009
generated: "2026-07-28T10:21:22.511401-07:00"

6
helm/litellm-helm/Chart.yaml
View file

@ -18,7 +18,7 @@ type: application
# This is the chart version. This version number should be incremented each time you make changes
# to the chart and its templates, including the app version.
# Versions are expected to follow Semantic Versioning (https://semver.org/)
version: 1.1.0
version: 1.1.1
# This is the version number of the application being deployed. This version number should be
# incremented each time you make changes to the application. Versions are not expected to
@ -32,10 +32,10 @@ annotations:
dependencies:
- name: "postgresql"
version: ">=13.3.0"
version: "14.3.1"
repository: oci://registry-1.docker.io/bitnamicharts
condition: db.deployStandalone
- name: redis
version: ">=18.0.0"
version: "18.19.1"
repository: oci://registry-1.docker.io/bitnamicharts
condition: redis.enabled

10
helm/litellm-helm/README.md
View file

@ -130,6 +130,16 @@ Set `billingMetrics.caSecretName` only when the collector is a private or test o
| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` |
| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) |
| `postgresql.auth.*` | If `db.deployStandalone` is `true`, care should be taken to ensure the default `password` and `postgres-password` values are **NOT** used. | `NoTaGrEaTpAsSwOrD` |
| `postgresql.image.*` | If `db.deployStandalone` is `true`, the image for the bundled Postgres. Pinned to a `docker.io/bitnamilegacy` build because Bitnami retired the versioned tags under `docker.io/bitnami`. | `bitnamilegacy/postgresql:16.2.0-debian-12-r6` |
| `redis.image.*` | If `redis.enabled` is `true`, the image for the bundled Redis. Pinned to a `docker.io/bitnamilegacy` build for the same reason. | `bitnamilegacy/redis:7.2.4-debian-12-r9` |
#### Bundled Postgres image
Bitnami removed the versioned tags from `docker.io/bitnami` and republished the archived builds under `docker.io/bitnamilegacy`, so the image defaults that ship inside the `postgresql` and `redis` subcharts no longer pull. The chart pins both to the `bitnamilegacy` copies of the exact builds those subchart versions were released with, which keeps the on-disk data directory layout unchanged for existing installs.
Keep `postgresql.image.tag` pinned. `docker.io/bitnami/postgresql` still publishes a floating `latest`, and pointing the bundled Postgres at a different major version starts the server against a data directory it cannot read (`database files are incompatible with server`). There is no in-place way back, so crossing a major version means dumping the database with the old image and restoring it into the new one. The chart refuses to render when the tag is empty or `latest`.
Those images no longer receive updates. For anything beyond getting started, run Postgres outside the chart and point at it with `db.useExisting`.
#### Example Postgres `db.useExisting` Secret

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

@ -146,3 +146,18 @@ Get redis service port
{{ .Values.redis.master.service.ports.redis }}
{{- end -}}
{{- end -}}
{{/*
Reject an unpinned image tag for the bundled PostgreSQL.
A floating tag lets a chart upgrade start a newer PostgreSQL major against the
existing PersistentVolumeClaim. The server then refuses to start on a data
directory written by another major version, and the only way back is a dump
taken before the change, which by that point no longer exists.
*/}}
{{- define "litellm.validateBundledPostgresImageTag" -}}
{{- $tag := .Values.postgresql.image.tag | default "" | toString -}}
{{- $digest := .Values.postgresql.image.digest | default "" | toString -}}
{{- if and (eq $digest "") (or (eq $tag "") (eq $tag "latest")) -}}
{{- 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 -}}

1
helm/litellm-helm/templates/secret-dbcredentials.yaml
View file

@ -1,4 +1,5 @@
{{- if .Values.db.deployStandalone -}}
{{- include "litellm.validateBundledPostgresImageTag" . -}}
apiVersion: v1
kind: Secret
metadata:

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

@ -10,7 +10,7 @@ metadata:
spec:
containers:
- name: test
image: bitnami/kubectl:latest
image: docker.io/bitnamilegacy/kubectl:1.29.2-debian-12-r3
command: ['sh', '-c']
args:
- |

94
helm/litellm-helm/tests/bundled_db_images_tests.yaml Normal file
View file

@ -0,0 +1,94 @@
suite: test bundled database images
templates:
- charts/postgresql/templates/primary/statefulset.yaml
- charts/redis/templates/master/application.yaml
- charts/redis/templates/configmap.yaml
- charts/redis/templates/health-configmap.yaml
- charts/redis/templates/scripts-configmap.yaml
- charts/redis/templates/secret.yaml
- secret-dbcredentials.yaml
- templates/tests/test-servicemonitor.yaml
tests:
- it: should pull the bundled postgres from a repository that still publishes the pinned tag
template: charts/postgresql/templates/primary/statefulset.yaml
set:
db.deployStandalone: true
asserts:
- equal:
path: spec.template.spec.containers[0].image
value: docker.io/bitnamilegacy/postgresql:16.2.0-debian-12-r6
- it: should pull the bundled postgres metrics exporter from the same repository
template: charts/postgresql/templates/primary/statefulset.yaml
set:
db.deployStandalone: true
postgresql.metrics.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].image
value: docker.io/bitnamilegacy/postgres-exporter:0.15.0-debian-12-r14
- it: should run the bundled postgres init container from the same repository
template: charts/postgresql/templates/primary/statefulset.yaml
set:
db.deployStandalone: true
postgresql.volumePermissions.enabled: true
asserts:
- equal:
path: spec.template.spec.initContainers[0].image
value: docker.io/bitnamilegacy/os-shell:12-debian-12-r16
- it: should pull the bundled redis from a repository that still publishes the pinned tag
template: charts/redis/templates/master/application.yaml
set:
redis.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[0].image
value: docker.io/bitnamilegacy/redis:7.2.4-debian-12-r9
- it: should reject a floating postgres tag that could cross a major version on an existing volume
template: secret-dbcredentials.yaml
set:
db.deployStandalone: true
postgresql.image.tag: latest
asserts:
- failedTemplate:
errorMessage: 'postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got "latest"). 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.'
- it: should reject an empty postgres tag
template: secret-dbcredentials.yaml
set:
db.deployStandalone: true
postgresql.image.tag: ""
asserts:
- failedTemplate:
errorMessage: 'postgresql.image.tag must be pinned to an explicit version when db.deployStandalone is true (got ""). 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.'
- it: should accept an empty postgres tag when the image is pinned by digest
template: secret-dbcredentials.yaml
set:
db.deployStandalone: true
postgresql.image.tag: ""
postgresql.image.digest: sha256:0d0e2f1a5b3c4d6e7f8091a2b3c4d5e6f708192a3b4c5d6e7f8091a2b3c4d5e6
asserts:
- hasDocuments:
count: 1
- it: should run the servicemonitor test pod from a pinned image
template: templates/tests/test-servicemonitor.yaml
set:
serviceMonitor.enabled: true
asserts:
- equal:
path: spec.containers[0].image
value: docker.io/bitnamilegacy/kubectl:1.29.2-debian-12-r3
- it: should not constrain the postgres tag when the bundled database is not deployed
template: secret-dbcredentials.yaml
set:
db.deployStandalone: false
postgresql.image.tag: latest
asserts:
- hasDocuments:
count: 0

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

@ -328,8 +328,32 @@ lifecycle: {}
# Settings for Bitnami postgresql chart (if db.deployStandalone is true, ignored
# otherwise)
#
# Bitnami retired the versioned tags under docker.io/bitnami and republished the
# archived builds under docker.io/bitnamilegacy, so the subchart's own image
# defaults no longer resolve. The repository below points at the same build the
# subchart was released with, which keeps the on-disk data directory layout
# identical for existing installs.
#
# Keep the tag pinned. docker.io/bitnami still publishes a floating `latest`,
# and starting a newer PostgreSQL major against an existing data directory
# leaves the server refusing to boot ("database files are incompatible with
# server") with no way back other than a dump taken beforehand. Crossing a major
# version is a dump-and-restore, not an image bump. The chart refuses to render
# an unpinned tag for this reason
postgresql:
architecture: standalone
image:
repository: bitnamilegacy/postgresql
tag: 16.2.0-debian-12-r6
volumePermissions:
image:
repository: bitnamilegacy/os-shell
tag: 12-debian-12-r16
metrics:
image:
repository: bitnamilegacy/postgres-exporter
tag: 0.15.0-debian-12-r14
auth:
username: litellm
database: litellm
@ -359,9 +383,36 @@ postgresql:
# When `redis.sentinel.enabled` is set, the coordination block is rendered with
# `sentinel_nodes` and `service_name` (from `redis.sentinel.masterSet`) instead
# of host/port, because a plain Redis client cannot talk to the sentinel port
#
# The image repositories carry the same bitnamilegacy repoint as postgresql
# above; the versioned tags the subchart ships with are gone from
# docker.io/bitnami
redis:
enabled: false
architecture: standalone
image:
repository: bitnamilegacy/redis
tag: 7.2.4-debian-12-r9
sentinel:
image:
repository: bitnamilegacy/redis-sentinel
tag: 7.2.4-debian-12-r7
metrics:
image:
repository: bitnamilegacy/redis-exporter
tag: 1.58.0-debian-12-r4
volumePermissions:
image:
repository: bitnamilegacy/os-shell
tag: 12-debian-12-r16
sysctl:
image:
repository: bitnamilegacy/os-shell
tag: 12-debian-12-r16
kubectl:
image:
repository: bitnamilegacy/kubectl
tag: 1.29.2-debian-12-r3
coordination:
# Set to false to keep the bundled Redis for response caching only and leave
# `general_settings.coordination_redis` out of the rendered config. A

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

@ -19,7 +19,7 @@
"/v1/fine-tuning" "/fine-tuning" "/v1/responses" "/responses" "/v1/threads" "/threads"
"/v1/assistants" "/assistants" "/v1/vector_stores" "/vector_stores" "/v1/indexes"
"/v1/models" "/models" "/openai" "/engines"
"/v1/messages" "/messages" "/v1/skills" "/v1/a2a"
"/v1/messages" "/messages" "/v1/skills" "/v1/a2a" "/a2a"
"/v1/rerank" "/v2/rerank" "/rerank" "/v1/ocr" "/ocr" "/v1/rag" "/rag"
"/v1/video" "/v1/videos" "/video" "/videos" "/v1/search" "/search"
"/v1/containers" "/containers" "/v1/evals" "/v1/memory" "/queue/chat"

17
litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_compression_saved_tokens/migration.sql Normal file
View file

@ -0,0 +1,17 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql Normal file
View file

@ -0,0 +1,9 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_MCPServerOAuthClient" (
"server_id" TEXT NOT NULL,
"credentials" JSONB,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_MCPServerOAuthClient_pkey" PRIMARY KEY ("server_id")
);

23
litellm-proxy-extras/litellm_proxy_extras/migrations/20260718000000_add_savings_spend/migration.sql Normal file
View file

@ -0,0 +1,23 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20260721000000_add_sso_identity_assertion/migration.sql Normal file
View file

@ -0,0 +1,9 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_SSOIdentityAssertion" (
"user_id" TEXT NOT NULL,
"assertion_b64" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_SSOIdentityAssertion_pkey" PRIMARY KEY ("user_id")
);

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260724000000_add_spend_log_tool_index_start_time_idx/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogToolIndex_start_time_idx" ON "LiteLLM_SpendLogToolIndex"("start_time");

12
litellm-proxy-extras/litellm_proxy_extras/migrations/20260725000000_add_daily_tool_spend/migration.sql Normal file
View file

@ -0,0 +1,12 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyToolSpend" (
"date" TEXT NOT NULL,
"tool_name" TEXT NOT NULL,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"total_tokens" BIGINT NOT NULL DEFAULT 0,
"request_count" BIGINT NOT NULL DEFAULT 0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_DailyToolSpend_pkey" PRIMARY KEY ("date","tool_name")
);

106
litellm-proxy-extras/litellm_proxy_extras/replica_identity.py Normal file
View file

@ -0,0 +1,106 @@
"""Optional post-migration step that raises Postgres REPLICA IDENTITY to FULL.
Logical-replication consumers (Neon / lakehouse sync and similar) need FULL
replica identity to reconstruct the old row of an UPDATE or DELETE. Prisma
leaves every table it creates at the Postgres default, so the setting has to be
re-applied by hand after each migration run. Setting
``LITELLM_SET_REPLICA_IDENTITY_FULL`` makes every migration run re-assert it.
The statement goes through the Prisma CLI rather than a Postgres driver because
``litellm-proxy-extras`` has no runtime dependencies, while the CLI is already
required for the migrations themselves.
"""
import subprocess
import tempfile
from pathlib import Path
from litellm_proxy_extras._logging import logger
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
REPLICA_IDENTITY_FULL_SQL = r"""
DO $$
DECLARE
target regclass;
BEGIN
SET LOCAL lock_timeout = '5s';
FOR target IN
SELECT c.oid::regclass
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind = 'r'
AND c.relreplident <> 'f'
AND n.nspname = ANY (current_schemas(false))
AND c.relname LIKE 'LiteLLM\_%'
LOOP
BEGIN
EXECUTE format('ALTER TABLE %s REPLICA IDENTITY FULL', target);
EXCEPTION WHEN lock_not_available THEN
RAISE WARNING 'REPLICA IDENTITY FULL skipped for %: table busy, retrying next run', target;
END;
END LOOP;
END
$$;
"""
def apply_replica_identity_full(
schema_path: str,
prisma_command: str,
prisma_env: dict[str, str],
) -> bool:
"""Set REPLICA IDENTITY FULL on every LiteLLM table that is not already FULL.
Never raises. Replication metadata is not needed to serve requests, so
every failure mode is reported and stepped over rather than taking down a
migration run that already succeeded: a database that refuses the ALTER
(most often because the runtime user does not own the tables), a missing
or unrunnable Prisma CLI, a read-only temp directory, or a timeout.
Returns True when the statement was applied, False when it failed.
"""
logger.info("Applying REPLICA IDENTITY FULL to LiteLLM tables")
try:
with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir:
sql_path = Path(tmp_dir) / "replica_identity_full.sql"
sql_path.write_text(REPLICA_IDENTITY_FULL_SQL)
subprocess.run(
[
prisma_command,
"db",
"execute",
"--file",
str(sql_path),
"--schema",
schema_path,
],
timeout=60,
check=True,
capture_output=True,
text=True,
env=prisma_env,
)
except subprocess.CalledProcessError as e:
logger.error(
"Failed to set REPLICA IDENTITY FULL. Logical replication "
"consumers may reject updates to these tables. Grant table "
"ownership to the migration user, or apply "
"`ALTER TABLE ... REPLICA IDENTITY FULL` by hand. Error: %s",
e.stderr,
)
return False
except subprocess.TimeoutExpired:
logger.error("Timed out setting REPLICA IDENTITY FULL on LiteLLM tables")
return False
except OSError as e:
logger.error(
"Could not run the REPLICA IDENTITY FULL statement. Logical "
"replication consumers may reject updates to these tables. "
"Error: %s",
e,
)
return False
logger.info("REPLICA IDENTITY FULL applied to LiteLLM tables")
return True

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

@ -396,6 +396,22 @@ model LiteLLM_MCPUserEnvVars {
@@index([server_id])
}
model LiteLLM_MCPServerOAuthClient {
server_id String @id
credentials Json?
created_at DateTime @default(now()) @map("created_at")
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
}
// The enterprise IdP identity assertion captured at SSO login, one row per user.
// assertion_b64 is an encrypted JSON payload: {id_token, refresh_token?, issuer?, expires_at?}.
model LiteLLM_SSOIdentityAssertion {
user_id String @id
assertion_b64 String
created_at DateTime @default(now()) @map("created_at")
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
}
// Generate Tokens for Proxy
model LiteLLM_VerificationToken {
token String @id
@ -729,6 +745,9 @@ model LiteLLM_DailyUserSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -760,6 +779,9 @@ model LiteLLM_DailyOrganizationSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -791,6 +813,9 @@ model LiteLLM_DailyEndUserSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -821,6 +846,9 @@ model LiteLLM_DailyAgentSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -851,6 +879,9 @@ model LiteLLM_DailyTeamSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -883,6 +914,9 @@ model LiteLLM_DailyTagSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -1060,6 +1094,20 @@ model LiteLLM_SpendLogToolIndex {
@@id([request_id, tool_name])
@@index([tool_name, start_time])
@@index([start_time])
}
// Daily tool spend rollup (one row per tool per day) – the Cost Optimization card reads this, never SpendLogs
model LiteLLM_DailyToolSpend {
date String
tool_name String
spend Float @default(0.0)
total_tokens BigInt @default(0)
request_count BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@id([date, tool_name])
}
// Prompt table for storing prompt configurations

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

@ -10,6 +10,10 @@ from pathlib import Path
from typing import Optional
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
def str_to_bool(value: Optional[str]) -> bool:
@ -676,6 +680,39 @@ class ProxyExtrasDBManager:
finally:
os.chdir(original_dir)
@staticmethod
def apply_replica_identity_full_if_requested() -> bool:
"""
Re-assert REPLICA IDENTITY FULL on LiteLLM's tables when the operator
opted in via LITELLM_SET_REPLICA_IDENTITY_FULL.
Prisma leaves new tables at the Postgres default, which logical
replication consumers reject, so the setting has to be re-applied after
every migration run rather than once by hand.
Returns:
bool: True if the setting was applied, False if it was not
requested or could not be applied.
"""
if not str_to_bool(os.getenv(REPLICA_IDENTITY_FULL_ENV_VAR)):
return False
try:
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
prisma_command = _get_prisma_command()
prisma_env = _get_prisma_env()
except OSError as e:
logger.error(
"Could not resolve the migrations directory for the REPLICA "
"IDENTITY FULL step, skipping it. Error: %s",
e,
)
return False
return apply_replica_identity_full(
schema_path=schema_path,
prisma_command=prisma_command,
prisma_env=prisma_env,
)
@staticmethod
def setup_database(
use_migrate: bool = False, use_v2_resolver: bool = False
@ -694,6 +731,15 @@ class ProxyExtrasDBManager:
Returns:
bool: True if setup was successful, False otherwise
"""
migrated = ProxyExtrasDBManager._run_migrations(
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated
@staticmethod
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
if use_v2_resolver:
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)

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

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.78"
version = "0.4.81"
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.78"
version = "0.4.81"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

30
litellm-rust/ADDING_A_PROVIDER.md
View file

@ -1,9 +1,29 @@
# Adding a provider / route to litellm-rust
Three layers, same for every route (see `ocr` and `realtime` as references):
Everything for a route lives in `crates/core/src/<route>/`; `crates/core/src/messages` is the reference. A host (the axum gateway, the Python bridge) only calls the route's entrypoint.
1. **Transform contract (pure)** — `crates/core/src/<route>/transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth.
2. **Provider config (pure)** — `crates/providers/src/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
3. **HTTP / transport (the host)** — `crates/providers/src/<route>.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O.
1. **Entrypoint** — `mod.rs`: `pub async fn <route>(request) -> CoreResult<Response>`, the Rust equivalent of `litellm.<route>()`, plus a `<route>_stream` variant when the route streams. It is the only thing a host touches.
2. **Transform contract** — `transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) with types in `types.rs`.
3. **Provider config** — `crates/core/src/providers/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
4. **Prepare + handler** — `prepare.rs` resolves provider/model, credentials, auth headers, and URL, then transforms the request; `handler.rs` performs the provider call through the shared client in `client.rs` and transforms the response.
**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
## Coding standards
Before writing new logic, look for an existing base to extend. When a change is
“the same behavior for one more provider/endpoint/integration”, the codebase
almost always already has a shared abstraction for it (for example, provider
`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
helpers in `litellm_core_utils/`, typed request/response models, or factory
functions). Find it first with a search, then add the new variant by inheriting
from or composing that base, overriding only what genuinely differs (model
name, parameter mapping, or auth).
Never copy an existing implementation and edit it in place, and never hand-roll
a parallel version of logic a base already provides. If you catch yourself
writing a second copy of a pattern that exists twice already, stop and extract a
base instead: put the shared shape in one place and make both call sites thin
variants of it. The test for a good abstraction is that adding the next provider
is a few declarative lines, not a new file of duplicated flow. Only diverge from
the base when behavior is genuinely different, and say so explicitly in the PR.
**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.

35
litellm-rust/AGENTS.md
View file

@ -4,14 +4,39 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (
## Crates
| Crate | Role | Pure / I/O |
|-------|------|------------|
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
| 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-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-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
## Where a route lives
A top-level LiteLLM call is a module under `crates/core/src/<route>/`, shaped like `messages`:
```
core/src/messages/
mod.rs # pub async fn messages(..) -> CoreResult<..> (+ messages_stream for SSE)
types.rs # request/response types, MessagesRequest
transformation.rs # the provider template trait
prepare.rs # provider resolution, auth headers, URL
handler.rs # the provider call
client.rs # the shared reqwest client
```
Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
## Style
All Rust in `litellm-rust/` follows the official Rust Style Guide:
https://doc.rust-lang.org/style-guide/
`rustfmt` implements its formatting by default, so run `cargo fmt` before committing; CI gates every PR on `cargo fmt --check`. Do not hand-format against rustfmt or add a `rustfmt.toml` that diverges from the default style.
Beyond formatting, follow the guide's naming and idiom conventions rustfmt cannot auto-apply: `snake_case` items/functions/modules, `UpperCamelCase` types/traits/variants, `SCREAMING_SNAKE_CASE` constants/statics (acronyms as one word, e.g. `HttpClient`), and the import grouping and item ordering it prescribes. See CLAUDE.md for the detailed version.

112
litellm-rust/CLAUDE.md
View file

@ -2,42 +2,96 @@
This file defines the rules for Rust work in LiteLLM.
## Provider Coding Standards
Before writing new logic, look for an existing base to extend. When a change is
“the same behavior for one more provider/endpoint/integration”, the codebase
almost always already has a shared abstraction for it (for example, provider
`BaseConfig` transformation classes in `litellm/llms/base_llm/`, shared
helpers in `litellm_core_utils/`, typed request/response models, or factory
functions). Find it first with a search, then add the new variant by inheriting
from or composing that base, overriding only what genuinely differs (model
name, parameter mapping, or auth).
Never copy an existing implementation and edit it in place, and never hand-roll
a parallel version of logic a base already provides. If you catch yourself
writing a second copy of a pattern that exists twice already, stop and extract a
base instead: put the shared shape in one place and make both call sites thin
variants of it. The test for a good abstraction is that adding the next provider
is a few declarative lines, not a new file of duplicated flow. Only diverge from
the base when behavior is genuinely different, and say so explicitly in the PR.
## Crates (exactly three — see AGENTS.md)
`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge`
exposes it to the Python SDK. A crate is a **layer**, not a route — add modules, not crates.
`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not
a route — add modules, not crates.
## Core Boundary
`litellm-core` is the pure translation layer; the `litellm-ai-gateway` host executes work.
`litellm-core` owns the whole call. The Rust equivalent of `litellm.messages()`
is `litellm_core::messages::messages(request).await`: you call it, it does the
provider call, and you get a typed non-streaming response back.
Route-level Rust structure mirrors LiteLLM's Python responsibilities:
- `core/src/<route>/` owns the route contract, shared types, and provider
template traits. For OCR, this means `core/src/ocr`.
- `core/src/<route>/` owns the route end to end: the public entrypoint fn named
after the route in `mod.rs`, the request/response types (`types.rs`), the
provider template trait (`transformation.rs`), the provider/auth/URL
resolution (`prepare.rs`), the HTTP client (`client.rs`), and the handler that
performs the call (`handler.rs`). `core/src/messages` is the reference.
- `core/src/providers/<provider>/<route>/transformation.rs` owns the
provider-specific transform. For Mistral OCR, this means
`core/src/providers/mistral/ocr/transformation.rs`.
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
never inside `core`.
provider-specific transform. For Anthropic Messages, this means
`core/src/providers/anthropic/messages/transformation.rs`.
- Handlers live in `core`, never in a host. `ai-gateway` must not contain a
route handler that talks to a provider; its axum route reads the HTTP request,
picks a deployment, and calls the `core` entrypoint. `python-bridge` marshals
Python objects and calls the same entrypoint.
Streaming keeps the same shape: the route entrypoint has a `<route>_stream`
variant in `core` that returns the upstream response so a host can splice it to
its own caller; the host still owns no provider logic.
Call-hook and lifecycle instrumentation, including phase timing, usage
accumulation, and callback payload construction, always lives in `core`.
Hosts feed observed events into core and dispatch the completed payloads through
their I/O logger; hosts must not own callback orchestration.
Allowed in `core`:
- Pure request transforms
- Pure response transforms
- Pure stream chunk normalization
- The public entrypoint for a top-level LiteLLM call
- Request/response transforms and stream chunk normalization
- Provider resolution, auth header construction, and URL building
- The provider HTTP call itself, through a shared reused client with connect and
request timeouts
- Shared data types and validation errors
- Deterministic token/cost helper logic
Not allowed in `core`:
- Network calls
- Environment variable or secret reads
- Serving HTTP: axum routes, extractors, and transport concerns stay in the host
- Filesystem access
- Database or cache access
- Provider SDK signing or auth flows
- Database access
- Config file reading and rollout state
- Logging callbacks, spend writes, or custom callbacks
- Global mutable runtime state
Env reads in `core` are limited to credential fallback inside a route's
`prepare.rs` (the `env_lookup` closure), mirroring what the Python SDK does when
no key is passed. Everything else config-shaped is resolved by the host and
passed in.
Routes still hosted in `ai-gateway` (`ocr`, `audio_transcription`, `realtime`)
predate this rule and are being moved into `core` route modules; do not add new
ones there, and prefer moving one when you touch it.
Python owns rollout state and fallback while Rust is being introduced. Rust
paths must be off by default until parity tests prove equivalence with Python.
A new provider/route may instead be implemented rust-only with no Python
reference; then the Python interface is a thin dispatch that calls Rust with no
fallback, and you state the rust-only choice explicitly in the PR. Either way
the Python side stays minimal (it only marshals inputs and calls the Rust
interface), never add a per-route feature flag, and never push provider
dispatch into `litellm/main.py`; put it in a thin dispatch class under
`litellm/llms/<provider>/<route>/`.
## Production Bar
@ -62,10 +116,10 @@ the first PR:
- Preserve Python output shape intentionally. If a field is always serialized as
`null` for Python parity, leave a short comment explaining that parity choice.
## Host I/O Rules
## Network I/O Rules
These rules apply when adding future crates or modules that execute network I/O,
such as `ai-gateway`, router hosts, or standalone servers:
These rules apply to every module that executes network I/O, whether it is a
`core` route handler or a host such as `ai-gateway`:
- Set connect and full-request timeouts. No unbounded waits.
- Reuse HTTP clients; do not construct clients per request.
@ -77,6 +131,26 @@ such as `ai-gateway`, router hosts, or standalone servers:
- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
impossible by construction and documented.
## Rust Style Guide
All Rust in `litellm-rust/` follows the official Rust Style Guide:
https://doc.rust-lang.org/style-guide/
`rustfmt` implements the guide's formatting rules by default, so the mechanical
side is enforced for you: run `cargo fmt` before committing and CI gates every
PR on `cargo fmt --check` (see Checks). Do not hand-format against rustfmt or add
a `rustfmt.toml` that diverges from the default style; the default style *is* the
guide.
The guide also covers conventions rustfmt cannot auto-apply; follow these too:
- Naming: `snake_case` for items, functions, and modules; `UpperCamelCase` for
types, traits, and enum variants; `SCREAMING_SNAKE_CASE` for constants and
statics; acronyms count as one word (`HttpClient`, not `HTTPClient`).
- Ordering and grouping the guide prescribes: imports grouped std / external /
crate-local, derives before other attributes, and consistent item order.
- Idioms the guide recommends over the formatter fighting you (e.g. prefer
restructuring an over-long expression rather than forcing an awkward wrap).
## Constants
Magic numbers and fixed strings go in a crate-level `constants.rs`, never

1373
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

7
litellm-rust/Cargo.toml
View file

@ -7,7 +7,8 @@ members = [
resolver = "2"
[workspace.package]
edition = "2021"
edition = "2024"
rust-version = "1.88"
license = "MIT"
repository = "https://github.com/BerriAI/litellm"
@ -15,8 +16,8 @@ repository = "https://github.com/BerriAI/litellm"
litellm-core = { path = "crates/core" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
axum = "0.7"
pyo3 = "0.23.5"
pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] }
pyo3 = "0.29.0"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }

41
litellm-rust/README.md
View file

@ -2,18 +2,31 @@
This workspace contains the staged Rust implementation for LiteLLM.
Rust starts as a pure transform core used by the existing Python host. Python
continues to own auth, configuration, network I/O, retries, routing, logging,
`litellm-core` is the LiteLLM SDK in Rust: one entrypoint per top-level call
that makes the LLM call and hands back a typed response, the same shape as
`litellm.messages()` in Python.
```rust
let response = litellm_core::messages::messages(MessagesRequest {
model: "claude-sonnet-4-5",
body,
api_key: Some(key),
..
})
.await?;
```
Python continues to own configuration, retries, routing policy, logging,
callbacks, spend tracking, and customer plugins until each Rust path has parity
coverage and production evidence.
## Crates
| Crate | Role | Pure / I/O |
|-------|------|------------|
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
| Crate | Role |
|-------|------|
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
@ -21,16 +34,16 @@ Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-
```text
crates/
core/ Route contracts, shared pure types, errors, and templates.
src/ocr/
providers/ Provider-specific pure transforms.
src/mistral/ocr/transformation.rs
core/ The SDK: route modules + provider transforms.
src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client
src/providers/anthropic/messages/transformation.rs
ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints.
python-bridge/ PyO3 bridge for Python LiteLLM.
```
The folder shape should follow the Python provider tree:
`providers/src/<provider>/<route>/transformation.rs`. The bridge should expose
one function per top-level route, starting with `ocr(payload)`.
The folder shape follows the Python provider tree:
`core/src/providers/<provider>/<route>/transformation.rs`. The bridge exposes one
function per top-level route, mirroring the core entrypoints.
## Checks

59
litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md Normal file
View file

@ -0,0 +1,59 @@
# Provider coding standards (litellm-rust)
Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` (`core/src/messages`, `ANTHROPIC_MESSAGES_CONFIG`) is the reference: a route is a `core` module with a public entrypoint that makes the call and returns a typed response.
## Provider resolution
1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string.
2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers.
## Transforms and the base config
3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/transformation.rs`, mirroring the Python provider tree.
5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it.
6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers.
## Boundaries
7. Layers never cross: `core` = the call itself (entrypoint, types, transforms, provider resolution, auth headers, provider HTTP, lifecycle hooks); `ai-gateway` = serving HTTP/WS (routing, extractors, auth of *our* callers, streaming to the client); `python-bridge` = thin PyO3 adapter. Hosts call the core entrypoint; they never build a provider request.
8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers/<provider>/<route>/`; a route is a module, never a new crate.
9. Route entry point stays thin: `core::<route>::<route>()` -> `prepare_*` -> handler (or `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing). Axum handlers validate and delegate to a service that calls the entrypoint; no business logic in them.
10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Config-shaped env reads happen at the host/config layer with the `DEFAULT_*` fallback defined in `constants.rs`; the only env read in `core` is the credential fallback in a route's `prepare.rs`.
## Types and errors
11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec<String>` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string.
12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input.
13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating.
14. Early returns over deep nesting; small focused files over god modules.
15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test.
## Safety and data minimization
16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary.
17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer.
18. Network I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS.
## Tests and rollout
19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity.
20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping.
21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR.
## Python bridge (SDK side)
22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
## Checks before push
25. Run, and keep green:
```bash
cd litellm-rust
cargo fmt --check
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
cargo test --workspace
```

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

@ -1,7 +1,9 @@
# ai-gateway — folder architecture
The Axum server that fronts the Rust gateway. It owns transport + config + auth
only; deployment selection lives in `core::router`, transforms in `core`/`providers`.
only; deployment selection lives in `core::router`, and the LLM call itself
(transforms, auth headers, provider HTTP) lives behind a `core` route entrypoint
such as `litellm_core::messages::messages`. No provider handler lives here.
```
src/
@ -32,6 +34,11 @@ src/
args; it runs during extraction. Never re-implement the check per route.
- **Handlers are thin.** A handler validates and delegates to its `service`. No
business logic, no provider calls, no transforms in handlers.
- **Services call `core`, they don't reimplement it.** A `service` picks the
deployment and calls the `core` route entrypoint. Provider resolution, auth
headers, URL building, and the HTTP call are `core`'s job; a service that
builds a provider request itself is a bug (`routes/messages/service.rs` is
the reference).
- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
`state.rs`; read env/config only in `main.rs` when building state.

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

@ -14,7 +14,7 @@ path = "src/main.rs"
required-features = ["server"]
[dependencies]
litellm-core.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
@ -41,3 +41,4 @@ python-config = ["dep:pyo3"]
[dev-dependencies]
futures-channel = "0.3"
tower = { version = "0.5.3", features = ["util"] }

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

@ -8,11 +8,11 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
`litellm-rust` is exactly three crates (a crate is a **layer**, not a route):
| Crate | Role | Pure / I/O |
|-------|------|------------|
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. Builds requests/responses; no network. | Pure |
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under `io/`) plus the Axum server binary (behind the `server` feature). | I/O |
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
| 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-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.

48
litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs Normal file
View file

@ -0,0 +1,48 @@
use std::collections::BTreeMap;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::error::CoreError;
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use serde_json::{Map, Value};
pub(super) fn audio_transcription_provider_config(
provider: &str,
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
match provider {
"bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
pub(super) fn string_headers(
headers: Option<Map<String, Value>>,
) -> CoreResult<BTreeMap<String, String>> {
headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"audio transcription extra_headers.{key} must be a string"
))
})
})
.collect()
}
pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
headers.keys().any(|key| key.eq_ignore_ascii_case(name))
}
pub(super) fn truncate_error_body(body: &str) -> String {
let truncated: String = body.chars().take(256).collect();
if truncated.chars().count() == body.chars().count() {
truncated
} else {
format!("{truncated}... (truncated)")
}
}

89
litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs Normal file
View file

@ -0,0 +1,89 @@
use std::time::SystemTime;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::error::CoreError;
use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use serde_json::Value;
use super::common_utils::truncate_error_body;
use super::types::ProviderAudioTranscriptionRequest;
use crate::client::http_client;
pub(crate) async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> CoreResult<Value> {
let body = serde_json::to_vec(&request.body).map_err(|error| {
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
})?;
let mut request_builder = http_client().post(&request.url).body(body.clone());
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 = request_builder
.send()
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text).map_err(|error| {
CoreError::InvalidResponse(format!("invalid audio response JSON: {error}"))
})?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.into_json())
}
pub(crate) async fn sign_request(
request: &ProviderAudioTranscriptionRequest,
optional_params: &serde_json::Map<String, Value>,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
let env_lookup = environment_lookup;
let auth = request
.config
.auth_strategy(&request.model, optional_params, &env_lookup)?;
let body = serde_json::to_vec(&request.body).map_err(|error| {
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
})?;
let mut headers = super::common_utils::string_headers(None)?;
headers.insert("Content-Type".to_string(), "application/json".to_string());
headers.extend(request.upstream_headers.iter().cloned());
match auth {
AudioTranscriptionAuth::Bearer => {}
AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
let credentials =
resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
.await?;
headers.extend(sign_bedrock_post(
&request.url,
&body,
&headers,
&region,
&credentials,
SystemTime::now(),
)?);
}
}
Ok(ProviderAudioTranscriptionRequest {
upstream_headers: headers.into_iter().collect(),
..request.clone()
})
}
pub(super) fn environment_lookup(key: &str) -> Option<String> {
std::env::var(key).ok()
}

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

@ -0,0 +1,300 @@
use std::future::Future;
use std::pin::Pin;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::CoreError;
use serde_json::{Map, Value, json};
use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
use super::handler::sign_request;
use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
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 AudioTranscriptionLifecycleHooks {
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl AudioTranscriptionLifecycleHooks {
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: PreparedAudioTranscriptionRequest,
) -> CoreResult<PreparedAudioTranscriptionRequest> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let (guardrail_request, _) = self
.guardrail_runner
.run_pre_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": request.custom_llm_provider,
"audio": request.audio,
"optional_params": request.optional_params,
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(CoreError::InvalidRequest(
"audio transcription pre_call guardrail must return an object".to_string(),
));
};
let audio = data.remove("audio").ok_or_else(|| {
CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(value)) => value,
Some(_) => {
return Err(CoreError::InvalidRequest(
"audio transcription optional_params must be an object".to_string(),
));
}
None => Map::new(),
};
Ok(PreparedAudioTranscriptionRequest {
audio,
optional_params,
..request
})
}
async fn prepare_provider_request(
&self,
request: PreparedAudioTranscriptionRequest,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
let config = audio_transcription_provider_config(&request.custom_llm_provider)
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
let env_lookup = super::handler::environment_lookup;
let headers = string_headers(request.extra_headers)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let body = config.transform_transcription_request(
&request.model,
request.audio,
filtered_params,
)?;
let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(
&upstream_headers
.iter()
.cloned()
.collect::<std::collections::BTreeMap<_, _>>(),
"authorization",
)
&& let Some(api_key) = request.api_key.as_deref()
{
upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
let provider_request = ProviderAudioTranscriptionRequest {
model: request.model,
config,
url,
body: body.body,
upstream_headers,
timeout: request.timeout,
};
let provider_request = self.run_during_call_guardrails(provider_request).await?;
sign_request(&provider_request, &request.optional_params).await
}
async fn run_during_call_guardrails(
&self,
request: ProviderAudioTranscriptionRequest,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let (guardrail_request, _) = self
.guardrail_runner
.run_during_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": "bedrock",
"url": request.url,
"body": request.body,
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(CoreError::InvalidRequest(
"audio transcription during_call guardrail must return an object".to_string(),
));
};
let body = data.remove("body").ok_or_else(|| {
CoreError::InvalidRequest("audio transcription guardrail removed body".to_string())
})?;
Ok(ProviderAudioTranscriptionRequest { body, ..request })
}
fn 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,
}
}
}
impl CallLifecycleHooks<PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, Value>
for AudioTranscriptionLifecycleHooks
{
type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>;
type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>;
type SuccessFuture<'a> = AudioLogFuture<'a>;
type FailureFuture<'a> = AudioLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedAudioTranscriptionRequest,
) -> 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: PreparedAudioTranscriptionRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { self.prepare_provider_request(request).await })
}
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;
}
self.logger_runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(
self.logging_payload(context, timing),
),
&CallbackValue::new("audio_transcription", response.clone()),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a CoreError,
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(),
};
self.logger_runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(
self.logging_payload(context, timing),
)
.with_failure_error(logging_error.clone()),
Some(&CallbackValue::new(
"error",
json!({"message": logging_error.message, "kind": logging_error.kind}),
)),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Other("audio_transcription".to_string()),
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 guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &CoreError) -> &'static str {
match error {
CoreError::Auth(_) => "AuthError",
CoreError::InvalidProvider(_) => "InvalidProvider",
CoreError::InvalidRequest(_) => "InvalidRequest",
CoreError::InvalidType { .. } => "InvalidType",
CoreError::MissingField(_) => "MissingField",
CoreError::Http { .. } => "HttpError",
CoreError::InvalidResponse(_) => "InvalidResponse",
CoreError::Network(_) => "NetworkError",
CoreError::Routing(_) => "RoutingError",
}
}

25
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs Normal file
View file

@ -0,0 +1,25 @@
use litellm_core::CoreResult;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::AudioTranscriptionRequest;
use handler::execute_audio_transcription_provider_call;
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult<Value> {
let PreparedAudioTranscriptionCall { request, hooks } =
prepare_audio_transcription_call(request);
CallLifecycle::default()
.run_request(request, &hooks, execute_audio_transcription_provider_call)
.await
}
#[cfg(test)]
mod tests;

55
litellm-rust/crates/ai-gateway/src/audio_transcription/prepare.rs Normal file
View file

@ -0,0 +1,55 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::hooks::AudioTranscriptionLifecycleHooks;
use super::types::{AudioTranscriptionRequest, PreparedAudioTranscriptionRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
pub(crate) struct PreparedAudioTranscriptionCall {
pub(crate) request: PreparedAudioTranscriptionRequest,
pub(crate) hooks: AudioTranscriptionLifecycleHooks,
}
pub(crate) fn prepare_audio_transcription_call(
request: AudioTranscriptionRequest<'_>,
) -> PreparedAudioTranscriptionCall {
let call_id = request
.litellm_call_id
.map(str::to_string)
.unwrap_or_else(new_audio_transcription_call_id);
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "bedrock",
});
PreparedAudioTranscriptionCall {
request: PreparedAudioTranscriptionRequest {
model: provider_info.model.to_string(),
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
litellm_call_id: call_id,
audio: request.audio,
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: request.optional_params,
timeout: request.timeout,
},
hooks: AudioTranscriptionLifecycleHooks::new(
CustomLoggerRunner::new(request.callbacks),
CustomGuardrailRunner::new(request.guardrails),
request.request_metadata,
),
}
}
fn new_audio_transcription_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_or(0, |duration| duration.as_nanos());
format!("audio-transcription-{timestamp}-{sequence}")
}

53
litellm-rust/crates/ai-gateway/src/audio_transcription/tests.rs Normal file
View file

@ -0,0 +1,53 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use serde_json::{Map, json};
use super::{AudioTranscriptionRequest, audio_transcription};
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("connection");
let mut request = Vec::new();
let mut buffer = [0_u8; 16_384];
let count = stream.read(&mut buffer).expect("request");
request.extend_from_slice(&buffer[..count]);
let request = String::from_utf8_lossy(&request);
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
assert!(request.contains("x-amz-date:"));
assert!(request.contains("\"bytes\":\"AQI=\""));
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
stream.write_all(response).expect("response");
});
let optional_params = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let api_base = format!("http://{address}");
let response = audio_transcription(AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
api_key: None,
api_base: Some(&api_base),
custom_llm_provider: Some("bedrock"),
extra_headers: None,
optional_params,
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
.expect("transcription");
assert_eq!(response, json!({"text": "hello"}));
server.join().expect("server");
}

58
litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs Normal file
View file

@ -0,0 +1,58 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
pub callbacks: Vec<Arc<dyn CustomLogger>>,
pub guardrails: Vec<Arc<dyn CustomGuardrail>>,
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}
pub(crate) struct PreparedAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
pub(crate) audio: 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 PreparedAudioTranscriptionRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"audio_transcription",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}
#[derive(Clone)]
pub(crate) struct ProviderAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

2
litellm-rust/crates/ai-gateway/src/auth/mod.rs
View file

@ -9,9 +9,9 @@
//! runs during extraction, before the handler body. Routes never re-implement it.
use axum::extract::FromRequestParts;
use axum::http::StatusCode;
use axum::http::header::AUTHORIZATION;
use axum::http::request::Parts;
use axum::http::StatusCode;
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq;

6
litellm-rust/crates/ai-gateway/src/ocr/client.rs → litellm-rust/crates/ai-gateway/src/client.rs
View file

@ -1,13 +1,13 @@
use std::sync::OnceLock;
use std::time::Duration;
const OCR_TIMEOUT_SECS: u64 = 600;
const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600;
pub(super) fn http_client() -> &'static reqwest::Client {
pub(crate) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
.timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS))
.build()
.expect("failed to build reqwest client")
})

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

@ -28,3 +28,15 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
/// Provider attributed to realtime sessions in the logging payload.
#[cfg(feature = "server")]
pub(crate) const DEFAULT_PROVIDER: &str = "openai";
pub(crate) const DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300;
/// HTTP path for the non-streaming Anthropic Messages route.
#[cfg(feature = "server")]
pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages";
/// Request headers owned by the gateway and never forwarded upstream.
#[cfg(feature = "server")]
pub(crate) const MESSAGES_HEADERS_NOT_FORWARDED: &[&str] =
&["authorization", "connection", "content-length", "host"];

1
litellm-rust/crates/ai-gateway/src/io/audio_transcription.rs Normal file
View file

@ -0,0 +1 @@
pub use crate::audio_transcription::{AudioTranscriptionRequest, audio_transcription};

2
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -1,3 +1,5 @@
pub mod audio_transcription;
pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;

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

@ -1 +1 @@
pub use crate::ocr::{ocr, OcrRequest};
pub use crate::ocr::{OcrRequest, ocr};

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

@ -15,16 +15,16 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::realtime::transformation::RealtimeProviderConfig;
use litellm_core::realtime::types::RealtimeEvent;
use litellm_core::CoreResult;
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
@ -113,7 +113,7 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<Realt
Message::Close(_) => {
return Err(CoreError::Network(
"upstream closed before first event".to_string(),
))
));
}
_ => continue,
}

24
litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs
View file

@ -28,11 +28,11 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use futures_util::StreamExt;
use litellm_core::realtime::types::RealtimeEvent;
use litellm_core::CoreResult;
use litellm_core::realtime::types::RealtimeEvent;
use crate::io::realtime::{
dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs,
UpstreamRx, UpstreamTx, UpstreamWs, dial_upstream, read_event, resolve_api_key,
};
/// Default target warm sockets per key when pooling is enabled.
@ -473,8 +473,8 @@ pub fn upstream_key(
/// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an
/// unexpected state. `Pending` (the healthy case) returns `false`.
fn is_dead(rx: &mut UpstreamRx) -> bool {
use futures_util::task::noop_waker_ref;
use futures_util::Stream;
use futures_util::task::noop_waker_ref;
use std::pin::Pin;
use std::task::{Context, Poll};
@ -523,15 +523,15 @@ mod tests {
))
.await;
while let Some(Ok(msg)) = ws.next().await {
if let Message::Text(text) = msg {
if text.contains("response.create") {
for frame in [
r#"{"type":"response.created"}"#,
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
r#"{"type":"response.done"}"#,
] {
let _ = ws.send(Message::Text(frame.to_string())).await;
}
if let Message::Text(text) = msg
&& text.contains("response.create")
{
for frame in [
r#"{"type":"response.created"}"#,
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
r#"{"type":"response.done"}"#,
] {
let _ = ws.send(Message::Text(frame.to_string())).await;
}
}
}

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

@ -0,0 +1,549 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
use litellm_core::{CoreError, CoreResult};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
};
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>;
#[derive(Clone)]
pub struct ResponsesWebSocketConnection {
socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
}
impl ResponsesWebSocketConnection {
pub async fn connect_url(
url: &str,
headers: &HashMap<String, String>,
timeout: Option<Duration>,
) -> CoreResult<Self> {
let mut request = url
.into_client_request()
.map_err(|error| CoreError::Network(error.to_string()))?;
for (name, value) in headers {
let header_name = name
.parse::<HeaderName>()
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
let header_value = HeaderValue::from_str(value)
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
CoreError::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => CoreError::Network(other.to_string()),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
})
}
pub async fn send_text(&self, text: String) -> CoreResult<()> {
let mut socket = self.socket.lock().await;
let Some(socket) = socket.as_mut() else {
return Err(CoreError::Network(
"Responses WebSocket is closed".to_string(),
));
};
socket
.send(Message::Text(text))
.await
.map_err(|error| CoreError::Network(error.to_string()))
}
pub async fn recv_text(&self) -> CoreResult<Option<String>> {
let mut socket_guard = self.socket.lock().await;
let Some(socket) = socket_guard.as_mut() else {
return Ok(None);
};
match socket.next().await {
Some(Ok(Message::Text(text))) => Ok(Some(text)),
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
.map(Some)
.map_err(|error| CoreError::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(CoreError::Network(error.to_string())),
}
}
pub async fn close(&self) -> CoreResult<()> {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket
.close(None)
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
*socket = None;
Ok(())
}
}
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
api_key
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| {
std::env::var(OPENAI_API_KEY_ENV)
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
}
async fn dial_upstream(
model: &str,
api_key: &str,
api_base: Option<&str>,
) -> CoreResult<ResponsesUpstreamWs> {
let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model);
let mut request = url
.as_str()
.into_client_request()
.map_err(|error| CoreError::Network(error.to_string()))?;
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}"))
.map_err(|error| CoreError::Auth(error.to_string()))?,
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
)
.await
.map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => CoreError::Network(other.to_string()),
})
}
pub struct ResponsesWebSocketStreaming;
impl ResponsesWebSocketStreaming {
pub async fn bidirectional_forward<In, Out>(
model: &str,
upstream_tx: UpstreamTx,
upstream_rx: UpstreamRx,
idle_timeout: Option<Duration>,
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
splice(
model,
upstream_tx,
upstream_rx,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
}
pub(crate) async fn splice<In, Out>(
model: &str,
mut upstream_tx: UpstreamTx,
mut upstream_rx: UpstreamRx,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
mut client_in: In,
mut client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let idle =
idle_timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS));
loop {
tokio::select! {
event = client_in.next() => {
let Some(event) = event else { break };
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&event, model)?
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
upstream_tx.send(Message::Text(payload))
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
}
message = upstream_rx.next() => {
let Some(message) = message else { break };
match message.map_err(|error| CoreError::Network(error.to_string()))? {
Message::Text(text) => {
let event = serde_json::from_str::<ResponsesWsEvent>(&text)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
observe(&event);
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_response(&event, model)?
.events
{
client_out.send(outbound)
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
}
Message::Close(_) => break,
_ => {}
}
}
_ = tokio::time::sleep(idle) => break,
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub async fn async_responses_websocket<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let key = resolve_api_key(api_key)?;
let upstream = dial_upstream(model, &key, api_base).await?;
let (mut upstream_tx, upstream_rx) = upstream.split();
if let Some(first_frame) = first_frame {
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&first_frame, model)?
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
upstream_tx
.send(Message::Text(payload))
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
}
ResponsesWebSocketStreaming::bidirectional_forward(
model,
upstream_tx,
upstream_rx,
idle_timeout,
&mut observe,
client_in,
client_out,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn responses_ws<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
async_responses_websocket(
model,
api_key,
api_base,
first_frame,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
#[cfg(test)]
mod tests {
use super::*;
use futures_channel::mpsc;
use futures_util::{SinkExt, StreamExt};
use litellm_core::responses::types::ResponsesWsEventType;
use serde_json::json;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");
let task = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let mut socket = accept_async(stream).await.expect("websocket handshake");
while let Some(Ok(Message::Text(text))) = socket.next().await {
let request: serde_json::Value = serde_json::from_str(&text).expect("request json");
let model = request
.get("model")
.and_then(serde_json::Value::as_str)
.or_else(|| {
request
.get("response")
.and_then(serde_json::Value::as_object)
.and_then(|response| {
response.get("model").and_then(serde_json::Value::as_str)
})
})
.expect("enforced model");
socket
.send(Message::Text(
json!({
"type": "response.created",
"response": {
"id": format!("resp-{model}"),
"model": model,
"extra": "preserved"
}
})
.to_string(),
))
.await
.expect("created event");
socket
.send(Message::Text(
json!({
"type": "response.completed",
"response": {
"id": format!("resp-{model}"),
"model": model,
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}
})
.to_string(),
))
.await
.expect("completed event");
}
});
(format!("http://{address}"), task)
}
fn event(value: serde_json::Value) -> ResponsesWsEvent {
serde_json::from_value(value).expect("event")
}
#[test]
fn explicit_nonblank_key_wins() {
assert_eq!(
resolve_api_key(Some(" explicit ")).expect("key"),
"explicit"
);
}
#[test]
fn blank_key_is_not_accepted_without_environment_key() {
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
assert!(resolve_api_key(Some(" ")).is_err());
}
}
#[tokio::test]
async fn forwards_events_sequentially_and_enforces_model() {
let (api_base, server) = websocket_base().await;
let (client_tx, client_rx) = mpsc::unbounded();
let (output_tx, mut output_rx) = mpsc::unbounded();
let (observed_tx, observed_rx) = mpsc::unbounded();
client_tx
.unbounded_send(event(json!({
"type": "response.create",
"model": "wrong"
})))
.expect("first request");
client_tx
.unbounded_send(event(json!({
"type": "response.create",
"response": {"model": "also-wrong"}
})))
.expect("second request");
let task = tokio::spawn(async move {
responses_ws(
"authorized-model",
Some("test-key"),
Some(&api_base),
None,
Some(Duration::from_secs(1)),
move |event| {
observed_tx
.unbounded_send(event.clone())
.expect("observe event");
},
client_rx,
output_tx,
)
.await
});
let first = output_rx.next().await.expect("first output");
let second = output_rx.next().await.expect("second output");
let third = output_rx.next().await.expect("third output");
let fourth = output_rx.next().await.expect("fourth output");
drop(client_tx);
task.await.expect("splice task").expect("successful splice");
server.await.expect("server task");
assert_eq!(first.event_type, ResponsesWsEventType::ResponseCreated);
assert_eq!(first.model(), Some("authorized-model"));
assert_eq!(first.data["response"]["extra"], "preserved");
assert_eq!(second.event_type, ResponsesWsEventType::ResponseCompleted);
assert_eq!(third.event_type, ResponsesWsEventType::ResponseCreated);
assert_eq!(fourth.event_type, ResponsesWsEventType::ResponseCompleted);
let observed: Vec<_> = observed_rx.collect().await;
assert_eq!(observed.len(), 4);
assert!(
observed
.iter()
.all(|event| event.event_type != ResponsesWsEventType::ResponseCreate)
);
}
#[tokio::test]
async fn idle_timeout_ends_without_upstream_events() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let _socket = accept_async(stream).await.expect("handshake");
tokio::time::sleep(Duration::from_secs(1)).await;
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, mut output_rx) = mpsc::unbounded();
let result = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await;
assert!(result.is_ok());
assert!(output_rx.next().await.is_none());
server.abort();
}
#[tokio::test]
async fn dial_http_status_is_preserved() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
stream
.write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n")
.await
.expect("response");
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, _output_rx) = mpsc::unbounded();
let error = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await
.expect_err("status error");
assert!(matches!(error, CoreError::Http { status: 401, .. }));
server.await.expect("server task");
}
#[tokio::test]
async fn dial_http_500_status_is_preserved() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
stream
.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n")
.await
.expect("response");
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, _output_rx) = mpsc::unbounded();
let error = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await
.expect_err("status error");
assert!(matches!(error, CoreError::Http { status: 500, .. }));
server.await.expect("server task");
}
}

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

@ -4,13 +4,17 @@
//! without pulling in the HTTP server:
//!
//! - Call-type modules such as [`ocr`]: provider transforms, lifecycle hooks,
//! and provider I/O. Always available — no feature required.
//! and provider I/O. Always available — no feature required. These predate the
//! rule that a route's entrypoint and handler live in `litellm-core` (see
//! `litellm_core::messages`) and move there as they are touched.
//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
//! for the load-time config reader.
pub mod audio_transcription;
mod client;
pub mod io;
pub mod ocr;
@ -25,9 +29,6 @@ pub mod routes;
#[cfg(feature = "server")]
pub mod state;
// Realtime request logging. Only the server serves realtime, so these are
// `server`-gated; `io::realtime` exposes the generic `observe` hook while the
// collector and callback fan-out live here.
mod constants;
pub mod integrations;
#[cfg(feature = "server")]

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

@ -11,7 +11,7 @@
use std::sync::Arc;
use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key};
use litellm_ai_gateway::routes;
use litellm_ai_gateway::state::AppState;
use litellm_core::router::{Deployment, LiteLLMParams, Router};

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

@ -1,11 +1,11 @@
use std::net::IpAddr;
use std::time::{Duration, Instant};
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrProviderConfig;
use litellm_core::CoreResult;
use reqwest::Url;
use serde_json::{Map, Value};
@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
};
use super::client::http_client;
use crate::client::http_client;
const ERROR_BODY_MAX_CHARS: usize = 256;
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;

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

@ -1,11 +1,11 @@
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrResponseHandling;
use litellm_core::CoreResult;
use serde_json::Value;
use super::client::http_client;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::types::ProviderOcrRequest;
use crate::client::http_client;
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
let mut request_builder = http_client().post(&request.url).json(&request.body);

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

@ -1,11 +1,11 @@
use std::future::Future;
use std::pin::Pin;
use litellm_core::CoreResult;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrAuthStrategy;
use litellm_core::CoreResult;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use super::common_utils::{
convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
@ -292,7 +292,7 @@ fn parse_ocr_pre_call_guardrail_request(
Some(_) => {
return Err(CoreError::InvalidRequest(
"OCR pre_call guardrail optional_params must be an object".to_string(),
))
));
}
None => Map::new(),
};

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

@ -1,8 +1,7 @@
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::CoreResult;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod client;
mod common_utils;
mod handler;
mod hooks;
@ -12,7 +11,7 @@ mod types;
pub use types::OcrRequest;
use handler::execute_ocr_provider_call;
use prepare::{prepare_ocr_call, PreparedOcrCall};
use prepare::{PreparedOcrCall, prepare_ocr_call};
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);

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

@ -1,7 +1,7 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::hooks::OcrLifecycleHooks;
use super::types::{OcrRequest, PreparedOcrRequest};

22
litellm-rust/crates/ai-gateway/src/ocr/tests.rs
View file

@ -3,12 +3,12 @@ use std::time::Duration;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body};
use super::{ocr, OcrRequest};
use super::{OcrRequest, ocr};
use crate::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
@ -228,19 +228,23 @@ fn truncate_error_body_does_not_split_multibyte_chars() {
#[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!(
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("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}

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

@ -7,9 +7,9 @@
//!
//! Compiled only under the `python-config` feature.
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::router::{Deployment, Router};
use litellm_core::CoreResult;
use pyo3::prelude::*;
use crate::gil;
@ -17,7 +17,7 @@ use crate::gil;
/// Load the router's `model_list` from `config_path` via the Python reader.
pub fn load_router_from_config(config_path: &str) -> CoreResult<Router> {
gil::record_acquisition();
Python::with_gil(|py| {
Python::attach(|py| {
let model_list = py
.import("litellm.proxy.read_model_list")
.and_then(|module| module.getattr("read_model_list"))

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

@ -19,7 +19,10 @@ async fn handle(...) -> impl IntoResponse { ... }
When a route has business logic worth testing without axum, put it in a sibling
`service` (a file, or a folder if the route grows). The route file stays the
**axum surface** (router + handler + any socket/SSE adapter); `service` is plain
Rust with **no axum types**. `realtime/` is the example:
Rust with **no axum types**, and its job is to pick the deployment and call the
`core` route entrypoint (see `messages/service.rs` calling
`litellm_core::messages::messages`). Never build a provider request, resolve a
key, or perform the provider call here. `realtime/` is the older example:
```
realtime/
mod.rs # axum surface: router() + handler + the WS<->events adapter
@ -33,6 +36,8 @@ genuinely gets hard to read.
`crate::auth::RequireMasterKey` to its arguments; it runs during extraction.
Never re-implement the check per route.
- **Handlers contain no business logic; `service` contains no axum types.**
- **No provider handlers in this crate.** Transforms, auth headers, and the
provider HTTP call live in `core/src/<route>/`.
- A route owns its paths in its own `router()`; `mod.rs` only merges.
- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`,
not duplicated in handlers.

2
litellm-rust/crates/ai-gateway/src/routes/health.rs
View file

@ -1,8 +1,8 @@
//! Health probes. Simple-route template: a `router()` plus its handlers, in one file.
use axum::Router;
use axum::http::StatusCode;
use axum::routing::get;
use axum::Router;
use crate::state::AppState;

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

@ -0,0 +1,513 @@
//! `POST /v1/messages`, the Anthropic Messages HTTP surface.
mod service;
use axum::Router;
use axum::body::Body;
use axum::extract::{Json, State};
use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue};
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use litellm_core::CoreError;
use serde_json::{Map, Value};
use crate::auth::RequireMasterKey;
use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH};
use crate::state::AppState;
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
}
async fn handle(
_auth: RequireMasterKey,
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response, MessagesRouteError> {
let extra_headers = forwarded_headers(&headers)?;
match service::run(&state.router, body, extra_headers)
.await
.map_err(MessagesRouteError::from)?
{
service::MessagesResponse::Json(body) => Ok(Json(body).into_response()),
service::MessagesResponse::Stream(upstream) => stream_response(upstream),
}
}
fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRouteError> {
let content_type = upstream
.headers()
.get(CONTENT_TYPE)
.cloned()
.unwrap_or_else(|| HeaderValue::from_static("text/event-stream"));
let mut response = Response::builder()
.status(
StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
MessagesRouteError(CoreError::InvalidResponse(format!(
"invalid upstream response status: {error}"
)))
})?,
)
.header(CONTENT_TYPE, content_type);
if let Some(value) = upstream.headers().get(CACHE_CONTROL) {
response = response.header(CACHE_CONTROL, value);
}
response
.body(Body::from_stream(upstream.bytes_stream()))
.map_err(|error| {
MessagesRouteError(CoreError::InvalidResponse(format!(
"failed to build streaming response: {error}"
)))
})
}
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, CoreError> {
let forwarded = headers
.iter()
.filter(|(name, _)| {
!MESSAGES_HEADERS_NOT_FORWARDED
.iter()
.any(|excluded| name.as_str().eq_ignore_ascii_case(excluded))
})
.map(|(name, value)| {
let value = value.to_str().map_err(|_| {
CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str()))
})?;
Ok((name.to_string(), Value::String(value.to_string())))
})
.collect::<Result<Map<_, _>, CoreError>>()?;
Ok((!forwarded.is_empty()).then_some(forwarded))
}
#[derive(Debug)]
struct MessagesRouteError(CoreError);
impl From<CoreError> for MessagesRouteError {
fn from(error: CoreError) -> Self {
Self(error)
}
}
impl IntoResponse for MessagesRouteError {
fn into_response(self) -> Response {
let (status, message) = match self.0 {
CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
CoreError::InvalidProvider(_) | CoreError::Routing(_) => (
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
CoreError::Auth(_) => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),
CoreError::Http { .. }
| CoreError::Network(_)
| CoreError::InvalidResponse(_)
| CoreError::InvalidType { .. }
| CoreError::MissingField(_) => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),
};
(
status,
Json(serde_json::json!({"error": {"message": message}})),
)
.into_response()
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use axum::body::Body;
use axum::http::Request;
use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE};
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde_json::json;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tower::ServiceExt;
use super::super::app;
use crate::io::realtime_pool::RealtimePool;
use crate::state::AppState;
fn state(model: &str, api_base: String, master_key: Option<&str>) -> AppState {
state_with_provider(model, model, api_base, master_key)
}
fn state_with_provider(
model_alias: &str,
provider_model: &str,
api_base: String,
master_key: Option<&str>,
) -> AppState {
AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias.to_string(),
litellm_params: LiteLLMParams {
model: format!("anthropic/{provider_model}"),
api_key: Some("upstream-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: master_key.map(Arc::from),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
}
}
async fn upstream(listener: TcpListener) -> (String, tokio::task::JoinHandle<String>) {
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = socket.read(&mut buffer).await.expect("reads request");
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let request = String::from_utf8(request).expect("request is utf8");
let content_length = request
.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);
let header_end = request.find("\r\n\r\n").expect("request has headers") + 4;
let mut full_request = request.into_bytes();
while full_request.len().saturating_sub(header_end) < content_length {
let read = socket.read(&mut buffer).await.expect("reads body");
full_request.extend_from_slice(&buffer[..read]);
}
let request = String::from_utf8(full_request).expect("request is utf8");
let body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-test"}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
(format!("http://{address}"), server)
}
async fn streaming_upstream(
listener: TcpListener,
status: u16,
content_type: &'static str,
body: &'static str,
) -> (String, tokio::task::JoinHandle<String>) {
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = socket.read(&mut buffer).await.expect("reads request");
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let request_text = String::from_utf8(request).expect("request is utf8");
let content_length = request_text
.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);
let header_end = request_text.find("\r\n\r\n").expect("request has headers") + 4;
let mut full_request = request_text.into_bytes();
while full_request.len().saturating_sub(header_end) < content_length {
let read = socket.read(&mut buffer).await.expect("reads body");
full_request.extend_from_slice(&buffer[..read]);
}
let response = format!(
"HTTP/1.1 {status} OK\r\ncontent-type: {content_type}\r\ncache-control: no-cache\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
String::from_utf8(full_request).expect("request is utf8")
});
(format!("http://{address}"), server)
}
#[tokio::test]
async fn route_constructs_anthropic_upstream_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = upstream(listener).await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("x-api-key", "request-upstream-key")
.header("anthropic-beta", "beta-feature")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("json")["id"],
"msg_1"
);
let upstream_request = server.await.expect("upstream task completes");
let (head, body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
let head = head.to_ascii_lowercase();
assert!(head.contains("x-api-key: request-upstream-key"));
assert!(head.contains("anthropic-beta: beta-feature"));
assert!(!head.contains("authorization: bearer master-key"));
let body: serde_json::Value = serde_json::from_str(body).expect("upstream body is json");
assert_eq!(body["model"], "claude-test");
assert_eq!(body["messages"][0]["content"], "hello");
}
#[tokio::test]
async fn route_substitutes_model_alias_with_provider_model_upstream() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = upstream(listener).await;
let app = app(state_with_provider(
"production",
"claude-sonnet-4-5",
api_base,
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "production",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
let upstream_request = server.await.expect("upstream task completes");
let (_, upstream_body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
let upstream_body: serde_json::Value =
serde_json::from_str(upstream_body).expect("upstream body is json");
assert_eq!(upstream_body["model"], "claude-sonnet-4-5");
assert_ne!(upstream_body["model"], "production");
}
#[tokio::test]
async fn route_streams_anthropic_events_without_buffering_or_reordering() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let events = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let (api_base, server) =
streaming_upstream(listener, 200, "text/event-stream", events).await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTENT_TYPE)
.unwrap()
.to_str()
.unwrap(),
"text/event-stream"
);
assert_eq!(
response
.headers()
.get(CACHE_CONTROL)
.unwrap()
.to_str()
.unwrap(),
"no-cache"
);
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(response_body, events.as_bytes());
let upstream_request = server.await.expect("upstream task completes");
let (_, upstream_body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
assert_eq!(
serde_json::from_str::<serde_json::Value>(upstream_body)
.expect("upstream body is json")["stream"],
true
);
}
#[tokio::test]
async fn route_maps_streaming_upstream_errors_before_starting_response() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = streaming_upstream(
listener,
429,
"application/json",
r#"{"error":"rate limited"}"#,
)
.await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&response_body).expect("error is json")["error"]
["message"],
"messages provider request failed"
);
server.await.expect("upstream task completes");
}
#[tokio::test]
async fn route_rejects_missing_master_key() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("content-type", "application/json")
.body(Body::from("{}"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn route_rejects_invalid_master_key() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer wrong-key")
.header("content-type", "application/json")
.body(Body::from("{}"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn route_rejects_malformed_json_without_panicking() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from("{not-json"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
}

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

@ -0,0 +1,65 @@
use std::sync::Arc;
use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER;
use litellm_core::messages::types::MessagesRequest;
use litellm_core::messages::{messages, messages_stream};
use litellm_core::router::Router;
use litellm_core::{CoreError, CoreResult};
use serde_json::{Map, Value};
pub(crate) enum MessagesResponse {
Json(Value),
Stream(reqwest::Response),
}
pub async fn run(
router: &Arc<Router>,
body: Value,
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<MessagesResponse> {
let model = body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|model| !model.is_empty())
.ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?;
let deployment = router.get_available_deployment(model).ok_or_else(|| {
CoreError::Routing(format!("no deployment available for model '{model}'"))
})?;
let provider_model = deployment.litellm_params.model.as_str();
let upstream_model = provider_model
.split_once('/')
.map_or(provider_model, |(_, model)| model);
let custom_llm_provider = if provider_model.contains('/') {
None
} else {
Some(ANTHROPIC_MESSAGES_PROVIDER)
};
let mut body = body;
body.as_object_mut()
.ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))?
.insert(
"model".to_string(),
Value::String(upstream_model.to_string()),
);
let request = MessagesRequest {
model: provider_model,
body,
api_key: deployment.litellm_params.api_key.as_deref(),
api_base: deployment.litellm_params.api_base.as_deref(),
custom_llm_provider,
extra_headers,
timeout: None,
};
if request.body.get("stream").and_then(Value::as_bool) == Some(true) {
return messages_stream(request).await.map(MessagesResponse::Stream);
}
let response = messages(request).await?;
serde_json::to_value(response)
.map(MessagesResponse::Json)
.map_err(|err| {
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
})
}

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