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

Merge remote-tracking branch 'upstream/litellm_internal_staging'

Browse source
This commit is contained in:
mutnale_sushant 2026-07-20 10:09:28 +05:30
commit 61689220f9
2465 changed files with 149739 additions and 30782 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.circleci/config.yml
View file

@ -1029,6 +1029,8 @@ jobs:
- *python312_image
working_directory: ~/project
resource_class: large
environment:
REQUEST_TIMEOUT: "180"
steps:
- checkout
@ -1058,7 +1060,8 @@ jobs:
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 8"
-n 8 \
--reruns 1 --only-rerun Timeout"
no_output_timeout: 15m
# Store test results

3
.github/CODEOWNERS vendored Normal file
View file

@ -0,0 +1,3 @@
/ui/ @yuneng-jiang @ryan-crabbe-berri
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
/ui/litellm-dashboard/src/lib/http/schema.d.ts

48
.github/actions/detect-backend-changes/action.yml vendored Normal file
View file

@ -0,0 +1,48 @@
name: "Detect backend-relevant changes"
description: >-
Classify the pull request's changed files with .circleci/scripts/classify_changes.sh
and expose decision=run|skip. decision=skip means only ui/**, **.md or **.mdx files
changed, so callers can short-circuit expensive steps while the job still completes
successfully and satisfies its required status check. The decision defaults to run for
any non pull_request event or whenever the changed set cannot be resolved, so tests are
never skipped when the classification is uncertain.
outputs:
decision:
description: "run when backend-relevant files changed, otherwise skip"
value: ${{ steps.classify.outputs.decision }}
runs:
using: composite
steps:
- id: classify
shell: bash
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
set -uo pipefail
if [ -z "${BASE_SHA:-}" ]; then
echo "detect-backend-changes: not a pull_request event; running job"
echo "decision=run" >> "${GITHUB_OUTPUT}"
exit 0
fi
if ! git fetch --no-tags --depth=1 origin "${BASE_SHA}" >/dev/null 2>&1; then
echo "detect-backend-changes: could not fetch base ${BASE_SHA}; running job"
echo "decision=run" >> "${GITHUB_OUTPUT}"
exit 0
fi
changed="$(git diff --name-only "${BASE_SHA}" HEAD 2>/dev/null)" || {
echo "detect-backend-changes: git diff failed; running job"
echo "decision=run" >> "${GITHUB_OUTPUT}"
exit 0
}
if [ -z "${changed}" ]; then
echo "detect-backend-changes: no changed files vs ${BASE_SHA}; skipping job"
echo "decision=skip" >> "${GITHUB_OUTPUT}"
exit 0
fi
echo "detect-backend-changes: changed files vs ${BASE_SHA}:"
printf '%s\n' "${changed}" | sed 's/^/ /'
decision="$(printf '%s\n' "${changed}" | bash .circleci/scripts/classify_changes.sh backend)" || decision="run"
echo "detect-backend-changes: decision=${decision}"
echo "decision=${decision}" >> "${GITHUB_OUTPUT}"

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

@ -0,0 +1,47 @@
name: "Set up uv with retries"
description: >-
Install uv via astral-sh/setup-uv, retrying on transient failures. Even with
an exact pinned version, the action resolves the artifact URL by fetching
https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a
single request with no retry, timeout, or fallback, so one connection-level
network error ("fetch failed") fails the whole job before any test runs.
Retrying the full step covers the manifest fetch and the binary download.
inputs:
version:
description: "uv version to install"
required: true
runs:
using: composite
steps:
- name: Set up uv (attempt 1)
id: attempt-1
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
version: ${{ inputs.version }}
- name: Wait before attempt 2
if: steps.attempt-1.outcome == 'failure'
shell: bash
run: sleep 15
- name: Set up uv (attempt 2)
id: attempt-2
if: steps.attempt-1.outcome == 'failure'
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
version: ${{ inputs.version }}
- name: Wait before attempt 3
if: steps.attempt-2.outcome == 'failure'
shell: bash
run: sleep 30
- name: Set up uv (attempt 3)
if: steps.attempt-2.outcome == 'failure'
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
with:
version: ${{ inputs.version }}

24
.github/pull_request_template.md vendored
View file

@ -41,3 +41,27 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
✅ Test
## Changes
## QA runbook
<!-- Only needed when your PR edits tests/e2e; delete this section otherwise
For each e2e test you added or changed, list the manual steps a reviewer can follow to reproduce it by hand against a live proxy, mapping 1:1 to what the test asserts: one top-level bullet per test giving its pytest node id followed by what it proves in plain words, then a nested "- [ ]" checklist where each item is a concrete action (route, request body, expected response) and the final item is the sanity-check step shown in the examples. Note environment prerequisites (provider credentials, config flags) and any nuances a manual run will hit. See PRs #32914 and #32963 for full examples
Example checklists:
- tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py::TestKeyRateLimits::test_rpm_limit_blocks_over_limit - a key allowed 2 requests a minute serves exactly 2 and refuses the 3rd
- [ ] Generate a limited key: curl -X POST http://localhost:4000/key/generate -H "Authorization: Bearer sk-1234" -d '{"rpm_limit": 2}'
- [ ] Send three /v1/chat/completions requests with that key inside one minute
- [ ] Expect the first two to return 200 and the third to return 429 naming the rpm limit
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
- tests/e2e/management/test_management_e2e.py::TestModelRoutes::test_model_create_appears_in_ui - a deployment created through the API shows up on the Admin UI models page
- [ ] POST /model/new with the master key, a bedrock model, and aws_region_name (needs STORE_MODEL_IN_DB=True and AWS credentials)
- [ ] Open http://localhost:4000/ui/?page=models and expect a deployment row showing the returned model id
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
-->
### Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

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

@ -45,19 +45,25 @@ jobs:
name: Run tests
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.timeout-minutes }}
outputs:
decision: ${{ steps.changes.outputs.decision }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect backend-relevant changes
id: changes
uses: ./.github/actions/detect-backend-changes
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
@ -72,16 +78,19 @@ jobs:
${{ runner.os }}-uv-
- 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
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Run tests
if: steps.changes.outputs.decision != 'skip'
env:
TEST_PATH: ${{ inputs.test-path }}
MAX_FAILURES: ${{ inputs.max-failures }}
@ -114,7 +123,7 @@ jobs:
fi
- name: Save coverage report
if: always()
if: always() && steps.changes.outputs.decision != 'skip'
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
@ -124,7 +133,7 @@ jobs:
upload-coverage:
name: Upload coverage to Codecov
needs: run
if: always()
if: always() && needs.run.outputs.decision != 'skip'
runs-on: ubuntu-latest
permissions:
contents: read

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

@ -18,7 +18,7 @@ jobs:
with:
persist-credentials: false
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Update JSON Data

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

@ -31,7 +31,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

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

@ -37,7 +37,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

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

@ -0,0 +1,61 @@
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 'litellm_oss_staging' branch 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 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."
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 'litellm_oss_staging' 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 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."
exit 1

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

@ -39,7 +39,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

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

@ -0,0 +1,50 @@
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 .

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

@ -38,7 +38,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

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

@ -33,7 +33,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
@ -48,7 +48,7 @@ jobs:
- name: Install dependencies
run: |
uv sync --frozen --group proxy-dev
uv sync --frozen --group proxy-dev --group e2e-dev
# basedpyright resolves Prisma's generated client (litellm/proxy/schema.prisma)
# only after `prisma generate` writes prisma/client.py et al. Without this the
@ -107,6 +107,16 @@ jobs:
run: |
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"
- name: Check tests/e2e basedpyright (zero errors)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
if git diff --name-only --diff-filter=ACMRD "$BASE_SHA"...HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
uv run --no-sync basedpyright tests/e2e
else
echo "No changed tests/e2e Python files; skipping."
fi
- name: Check for circular imports
run: |
cd litellm
@ -162,7 +172,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

76
.github/workflows/test-litellm-ui-build.yml vendored
View file

@ -36,79 +36,3 @@ jobs:
- name: Build
run: npm run build
frontend-lint:
runs-on: ubuntu-latest
timeout-minutes: 8
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: Collect changed files
id: changed
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
: > "$RUNNER_TEMP/prettier_files.txt"
: > "$RUNNER_TEMP/eslint_files.txt"
while IFS= read -r f; do
[ -f "$f" ] || continue
case "$f" in
*.js | *.jsx | *.ts | *.tsx | *.mjs | *.cjs)
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt"
printf '%s\n' "$f" >> "$RUNNER_TEMP/eslint_files.txt" ;;
*.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 -- .)
if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then
echo "has_files=true" >> "$GITHUB_OUTPUT"
else
echo "has_files=false" >> "$GITHUB_OUTPUT"
echo "No lintable UI files changed in this PR; nothing to check."
fi
- name: Setup Node.js
if: steps.changed.outputs.has_files == 'true'
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
if: steps.changed.outputs.has_files == 'true'
run: npm ci
- name: Lint changed files (prettier + eslint)
if: steps.changed.outputs.has_files == 'true'
run: |
prettier_files=()
eslint_files=()
while IFS= read -r f; do prettier_files+=("$f"); done < "$RUNNER_TEMP/prettier_files.txt"
while IFS= read -r f; do eslint_files+=("$f"); done < "$RUNNER_TEMP/eslint_files.txt"
status=0
if [ ${#prettier_files[@]} -gt 0 ]; then
echo "::group::Prettier (${#prettier_files[@]} files)"
npx prettier --check "${prettier_files[@]}" || { status=1; echo "::error::Unformatted files. Fix with: npm run format"; }
echo "::endgroup::"
fi
if [ ${#eslint_files[@]} -gt 0 ]; then
echo "::group::ESLint (${#eslint_files[@]} files)"
npx eslint --no-warn-ignored --pass-on-unpruned-suppressions "${eslint_files[@]}" || status=1
echo "::endgroup::"
fi
exit $status
- name: Check lint budgets
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
run: |
npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json --check eslint-metrics.json

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

@ -0,0 +1,92 @@
name: UI Lint
permissions:
contents: read
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
jobs:
frontend-lint:
runs-on: ubuntu-latest
timeout-minutes: 8
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: Collect changed files
id: changed
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
: > "$RUNNER_TEMP/prettier_files.txt"
: > "$RUNNER_TEMP/eslint_files.txt"
while IFS= read -r f; do
[ -f "$f" ] || continue
case "$f" in
*.js | *.jsx | *.ts | *.tsx | *.mjs | *.cjs)
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt"
printf '%s\n' "$f" >> "$RUNNER_TEMP/eslint_files.txt" ;;
*.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 -- .)
if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then
echo "has_files=true" >> "$GITHUB_OUTPUT"
else
echo "has_files=false" >> "$GITHUB_OUTPUT"
echo "No lintable UI files changed in this PR; nothing to check."
fi
- name: Setup Node.js
if: steps.changed.outputs.has_files == 'true'
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
if: steps.changed.outputs.has_files == 'true'
run: npm ci
- name: Lint changed files (prettier + eslint)
if: steps.changed.outputs.has_files == 'true'
run: |
prettier_files=()
eslint_files=()
while IFS= read -r f; do prettier_files+=("$f"); done < "$RUNNER_TEMP/prettier_files.txt"
while IFS= read -r f; do eslint_files+=("$f"); done < "$RUNNER_TEMP/eslint_files.txt"
status=0
if [ ${#prettier_files[@]} -gt 0 ]; then
echo "::group::Prettier (${#prettier_files[@]} files)"
npx prettier --check "${prettier_files[@]}" || { status=1; echo "::error::Unformatted files. Fix with: npm run format"; }
echo "::endgroup::"
fi
if [ ${#eslint_files[@]} -gt 0 ]; then
echo "::group::ESLint (${#eslint_files[@]} files)"
npx eslint --no-warn-ignored --pass-on-unpruned-suppressions "${eslint_files[@]}" || status=1
echo "::endgroup::"
fi
exit $status
- name: Check lint budgets
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
run: |
npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json
- name: Check for dead code (knip)
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
run: npm run knip:ci

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

@ -32,7 +32,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

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

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

@ -31,7 +31,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

2
.github/workflows/test-terraform-provider.yml vendored
View file

@ -74,7 +74,7 @@ jobs:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"

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

@ -32,13 +32,17 @@ jobs:
path: docs/my-website
persist-credentials: false
- name: Detect backend-relevant changes
id: changes
uses: ./.github/actions/detect-backend-changes
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
@ -53,10 +57,12 @@ jobs:
${{ runner.os }}-uv-
- 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
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
@ -64,6 +70,7 @@ jobs:
# Run the same documentation tests that CircleCI ran (as direct Python scripts)
- name: Run documentation validation tests
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync python ./tests/documentation_tests/test_env_keys.py
uv run --no-sync python ./tests/documentation_tests/test_router_settings.py

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

@ -5,6 +5,8 @@ on:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read

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

@ -49,13 +49,17 @@ jobs:
with:
persist-credentials: false
- name: Detect backend-relevant changes
id: changes
uses: ./.github/actions/detect-backend-changes
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
@ -70,16 +74,19 @@ jobs:
${{ runner.os }}-uv-
- 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
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Run tests - ${{ matrix.test-group.name }}
if: steps.changes.outputs.decision != 'skip'
env:
TEST_PATH: ${{ matrix.test-group.path }}
run: |

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

@ -16,6 +16,7 @@ jobs:
timeout-minutes: 30
strategy:
fail-fast: false
matrix:
root_path: ["/api/v1", "/llmproxy"]
@ -108,8 +109,26 @@ jobs:
- name: Install UI deps and Chromium
working-directory: ui/litellm-dashboard
run: |
npm ci
npx playwright install --with-deps chromium
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

6
.github/workflows/zizmor.yml vendored
View file

@ -4,7 +4,11 @@ on:
push:
branches: [main, litellm_internal_staging]
pull_request:
branches: [main, litellm_internal_staging]
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}

10
.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*
@ -106,6 +109,13 @@ STABILIZATION_TODO.md
**/coverage
test-config
# Claude Code compatibility-matrix pytest artifact (CI-only output).
compat-results.json
compat-results.json.shards/
compat-rate-limit-summary.json
# Matrix JSON produced by the daily-cron publisher (pushed to litellm-docs).
compatibility-matrix.json
# ---------- Terraform ----------
# Provider binaries + module cache — regenerated by `terraform init`.
**/.terraform/

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
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 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
@ -39,6 +39,8 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
On a fresh worktree or clone, run `make bootstrap` before anything else. It provisions everything tests, `make pre-commit`, and a local proxy need
Run tests before you commit. Also, run `make pre-commit` right before each commit, which generates types (as needed) and formats/lints your code. Any errors found must be fixed. It only runs when there are staged frontend and/or backend changes and calculates violations, generates types, etc. based on the worktree, so stage what you need or stash/delete unwanted files in litellm/ or ui/ (where backend and frontend lint run, respectively) before running it. If it fails because dashboard api types are stale, it already regenerated them for you. You just need to stage the schema.d.ts, re-run `make pre-commit` to confirm it passes, and commit
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing

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 create a pull request
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`.
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

29
Dockerfile
View file

@ -86,7 +86,9 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra semantic-router \
--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 +102,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
@ -114,16 +120,19 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
# 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
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
# 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

28
Makefile
View file

@ -5,15 +5,16 @@
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
info lint lint-dev lint-checks format \
lint-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety pre-commit \
lint-install lint-fetch-base
lint-install lint-fetch-base bootstrap
# Default target
help:
@echo "Available commands:"
@echo " make bootstrap - Provision a fresh clone/worktree"
@echo " make install-dev - Install development dependencies"
@echo " make install-proxy-dev - Install proxy development dependencies"
@echo " make install-dev-ci - Install dev dependencies (CI-compatible, pins OpenAI)"
@ -27,6 +28,7 @@ help:
@echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
@echo " make lint-ruff - Run Ruff linting only"
@echo " make lint-basedpyright - Run basedpyright strict, gated by per-rule error counts"
@echo " make lint-e2e-basedpyright - Run basedpyright over tests/e2e (zero errors allowed)"
@echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed"
@echo " make lint-format - Check ruff format formatting (matches CI)"
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@ -54,6 +56,7 @@ UV := uv
UV_RUN := $(UV) run --no-sync
LINT_DEP_INSTALL ?= install-dev
LINT_E2E_DEP_INSTALL ?= lint-install
LINT_DEP_BASE ?= lint-fetch-base
LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4)
LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,)
@ -69,6 +72,18 @@ info:
install-dev:
$(UV) sync --inexact --frozen
bootstrap:
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
cd ui/litellm-dashboard && npm ci --no-audit --no-fund
@main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \
if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \
cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \
else \
echo "bootstrap: .env left untouched"; \
fi
@echo "bootstrap: done"
install-proxy-dev:
$(UV) sync --frozen --group proxy-dev --extra proxy
@ -111,7 +126,7 @@ lint-fetch-base:
# CI's). --inexact tops up the venv instead of pruning the proxy extras gen:api and the
# running proxy need.
lint-install:
$(UV) sync --inexact --frozen --group proxy-dev
$(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
# Diff-scoped format check, identical to test-linting.yml's "Check ruff format" step:
@ -164,6 +179,9 @@ lint-ruff-FULL-dev: install-dev
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
lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
$(UV_RUN) basedpyright tests/e2e
# Type-discipline budget (mutable collections / casts / type guards / kwargs /
# unexplained suppressions), the test-linting.yml step `make lint` used to omit.
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@ -208,9 +226,9 @@ check-import-safety: $(LINT_DEP_INSTALL)
# base fetch) runs once up front; the checks themselves are independent, so a sub-make
# fans them out with -j and the fast ones finish under basedpyright's shadow.
lint: lint-install lint-fetch-base
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_DEP_BASE= lint-checks
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright check-circular-imports check-import-safety
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety

13
README.md
View file

@ -552,17 +552,12 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws
2. Run dependent services `docker-compose up db prometheus`
#### Backend
1. (In root) create virtual environment `python -m venv .venv`
2. Activate virtual environment `source .venv/bin/activate`
3. Install dependencies `uv sync --all-extras --group proxy-dev`
4. `uv run prisma generate`
5. `prisma generate`
6. Start proxy backend `python litellm/proxy/proxy_cli.py`
1. Run `make bootstrap`
2. Start proxy backend: `uv run python litellm/proxy/proxy_cli.py`
#### Frontend
1. Navigate to `ui/litellm-dashboard`
2. Install dependencies `npm install`
3. Run `npm run dev` to start the dashboard
1. Navigate to `ui/litellm-dashboard` (dependencies were already installed w/ `make bootstrap`)
2. Start dashboard: `npm run dev`
### Verify Docker Image Signatures

1
backend/routes/allowlist.py
View file

@ -46,6 +46,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/fallback",
"/fallbacks",
"/cache_settings",
"/coordination_redis/",
"/cost_tracking",
"/cost/",
"/credentials",

12
basedpyright-code-budget.json
View file

@ -24,7 +24,7 @@
"limit": 42
},
"reportExplicitAny": {
"limit": 10397
"limit": 10389
},
"reportFunctionMemberAccess": {
"limit": 11
@ -57,7 +57,7 @@
"limit": 5900
},
"reportMissingTypeArgument": {
"limit": 15918
"limit": 15903
},
"reportMissingTypeStubs": {
"limit": 41
@ -105,13 +105,13 @@
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40541
"limit": 40539
},
"reportUnknownParameterType": {
"limit": 20418
"limit": 20403
},
"reportUnknownVariableType": {
"limit": 32151
"limit": 32141
},
"reportUnnecessaryCast": {
"limit": 177
@ -123,7 +123,7 @@
"limit": 7
},
"reportUnnecessaryIsInstance": {
"limit": 1212
"limit": 1209
},
"reportUntypedBaseClass": {
"limit": 165

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

Binary file not shown.

31
docker/Dockerfile.database
View file

@ -84,7 +84,9 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra semantic-router \
--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 +99,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
@ -111,16 +117,21 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
# 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
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
# 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

1
docker/Dockerfile.non_root
View file

@ -137,6 +137,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
# working directory on sys.path; litellm/proxy/hooks resolves
# 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
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets

2
docker/build_from_pip/Dockerfile.build_from_pip
View file

@ -36,7 +36,7 @@ RUN uv venv --python python && \
"opentelemetry-api==1.28.0" \
"opentelemetry-sdk==1.28.0" \
"opentelemetry-exporter-otlp==1.28.0" \
"ddtrace==2.19.0" \
"ddtrace==4.11.0" \
"sentry-sdk==2.21.0" \
"mangum==0.17.0" \
"azure-ai-contentsafety==1.0.0" \

2
enterprise/README.md
View file

@ -6,4 +6,4 @@ Code in this folder is licensed under a commercial license. Please review the [L
👉 **Using in an Enterprise / Need specific features ?** Meet with us [here](https://enterprise.litellm.ai/demo?month=2024-02)
See all Enterprise Features here 👉 [Docs](https://docs.litellm.ai/docs/proxy/enterprise)
See all Enterprise Features here 👉 [Docs](https://docs.litellm.ai/docs/enterprise)

141
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View file

@ -7,7 +7,8 @@
# Thank you users! We ❤️ you! - Krrish & Ishaan
## This provides an LLM Guard Integration for content moderation on the proxy
from typing import Literal, Optional
import asyncio
from typing import Optional
import aiohttp
from fastapi import HTTPException
@ -18,7 +19,6 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_formatted_prompt
class _ENTERPRISE_LLMGuard(CustomLogger):
@ -46,45 +46,44 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
except Exception:
pass
async def moderation_check(self, text: str):
async def moderation_check(self, text: str) -> str:
"""
Runs the LLM Guard moderation check on ``text``.
Raises an HTTPException when the content violates the safety policy;
otherwise returns the sanitized prompt from LLM Guard, falling back to
the original text when the API does not provide one.
[TODO] make this more performant for high-throughput scenario
"""
try:
async with aiohttp.ClientSession() as session:
if self.mock_redacted_text is not None:
redacted_text = self.mock_redacted_text
else:
# Make the first request to /analyze
analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
verbose_proxy_logger.debug("Making request to: %s", analyze_url)
analyze_payload = {"prompt": text}
redacted_text = None
if self.mock_redacted_text is not None:
redacted_text = self.mock_redacted_text
else:
analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
verbose_proxy_logger.debug("Making request to: %s", analyze_url)
async with aiohttp.ClientSession() as session:
async with session.post(
analyze_url, json=analyze_payload
analyze_url, json={"prompt": text}
) as response:
redacted_text = await response.json()
verbose_proxy_logger.debug(
f"LLM Guard: Received response - {redacted_text}"
verbose_proxy_logger.debug(
f"LLM Guard: Received response - {redacted_text}"
)
if redacted_text is None:
raise HTTPException(
status_code=500,
detail={
"error": f"Invalid content moderation response: {redacted_text}"
},
)
if redacted_text is not None:
if (
redacted_text.get("is_valid", None) is not None
and redacted_text["is_valid"] is False
):
raise HTTPException(
status_code=400,
detail={"error": "Violated content safety policy"},
)
else:
pass
else:
raise HTTPException(
status_code=500,
detail={
"error": f"Invalid content moderation response: {redacted_text}"
},
)
if redacted_text.get("is_valid", None) is False:
raise HTTPException(
status_code=400,
detail={"error": "Violated content safety policy"},
)
sanitized_prompt = redacted_text.get("sanitized_prompt")
return sanitized_prompt if isinstance(sanitized_prompt, str) else text
except Exception as e:
verbose_proxy_logger.exception(
"litellm.enterprise.enterprise_hooks.llm_guard::moderation_check - Exception occurred - {}".format(
@ -138,23 +137,75 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
return
self.print_verbose("Makes LLM Guard Check")
try:
assert call_type in [
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
]
except Exception:
if call_type not in [
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
]:
self.print_verbose(
f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']"
)
return data
formatted_prompt = get_formatted_prompt(data=data, call_type=call_type) # type: ignore
self.print_verbose(f"LLM Guard, formatted_prompt: {formatted_prompt}")
return await self.moderation_check(text=formatted_prompt)
return await self._moderate_request(data=data)
async def _moderate_request(self, data: dict) -> dict:
"""
Sanitizes the request in place using the prompt returned by LLM Guard so
the provider-bound request carries the redacted content, then returns it.
"""
messages = data.get("messages")
if messages is not None:
data["messages"] = list(
await asyncio.gather(
*(self._moderate_message(message) for message in messages)
)
)
return data
input_ = data.get("input")
if input_ is not None:
data["input"] = await self._moderate_input(input_)
return data
prompt = data.get("prompt")
if isinstance(prompt, str):
data["prompt"] = await self.moderation_check(text=prompt)
return data
async def _moderate_message(self, message: dict) -> dict:
content = message.get("content")
if isinstance(content, str):
return {**message, "content": await self.moderation_check(text=content)}
if isinstance(content, list):
return {
**message,
"content": list(
await asyncio.gather(
*(self._moderate_content_part(part) for part in content)
)
),
}
return message
async def _moderate_content_part(self, part: dict) -> dict:
if part.get("type") == "text" and isinstance(part.get("text"), str):
return {**part, "text": await self.moderation_check(text=part["text"])}
return part
async def _moderate_input(self, input_: object) -> object:
if isinstance(input_, str):
return await self.moderation_check(text=input_)
if isinstance(input_, list):
return [
await self.moderation_check(text=item)
if isinstance(item, str)
else item
for item in input_
]
return input_
async def async_post_call_streaming_hook(
self, user_api_key_dict: UserAPIKeyAuth, response: str

8
enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
View file

@ -113,6 +113,10 @@ class PagerDutyAlerting(SlackAlerting):
user_api_key_spend=_meta.get("user_api_key_spend"),
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
user_api_key_user_spend=_meta.get("user_api_key_user_spend"),
user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"),
user_api_key_team_spend=_meta.get("user_api_key_team_spend"),
user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"),
user_api_key_org_id=_meta.get("user_api_key_org_id"),
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
user_api_key_team_id=_meta.get("user_api_key_team_id"),
@ -196,6 +200,10 @@ class PagerDutyAlerting(SlackAlerting):
if user_api_key_dict.budget_reset_at
else None
),
user_api_key_user_spend=user_api_key_dict.user_spend,
user_api_key_user_max_budget=user_api_key_dict.user_max_budget,
user_api_key_team_spend=user_api_key_dict.team_spend,
user_api_key_team_max_budget=user_api_key_dict.team_max_budget,
user_api_key_org_id=user_api_key_dict.org_id,
user_api_key_org_alias=user_api_key_dict.organization_alias,
user_api_key_team_id=user_api_key_dict.team_id,

4
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
View file

@ -919,9 +919,9 @@ class BaseEmailLogger(CustomLogger):
"""
Construct invitation link for the user
# http://localhost:4000/ui?invitation_id=7a096b3a-37c6-440f-9dd1-ba22e8043f6b
# http://localhost:4000/ui/onboarding?invitation_id=7a096b3a-37c6-440f-9dd1-ba22e8043f6b
"""
return f"{base_url}/ui?invitation_id={invitation_id}"
return f"{base_url}/ui/onboarding?invitation_id={invitation_id}"
async def send_email(
self,

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

@ -29,7 +29,7 @@ class CheckBatchCost:
proxy_logging_obj: "ProxyLogging",
prisma_client: "PrismaClient",
llm_router: "Router",
track_unmanaged_vertex_batch_cost: bool = False,
track_unmanaged_batch_cost: bool = False,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -37,7 +37,7 @@ class CheckBatchCost:
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._track_unmanaged_vertex_batch_cost = track_unmanaged_vertex_batch_cost
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
# Cached after the first poll cycle. Once we know the column is absent we skip
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
@ -118,11 +118,11 @@ class CheckBatchCost:
Resolve (model_id, batch_id) for a managed-object row, where model_id is a router
deployment id and batch_id is the raw provider batch id.
Managed batches encode both in a base64 unified id. Unmanaged Vertex batches, created with
a raw gs:// input_file_id, store the raw provider job id as unified_object_id; when
track_unmanaged_vertex_batch_cost is enabled the model is derived from the gs:// path and
mapped to a configured vertex_ai deployment. Returns None (recording a metric) when the row
can't be routed.
Managed batches encode both in a base64 unified id. Unmanaged batches (created outside
LiteLLM's own /v1/batches with a raw input_file_id) store the raw provider job id as
unified_object_id instead; when track_unmanaged_batch_cost is enabled the model is derived
from the provider-specific input_file_id layout (Vertex gs:// or Bedrock s3://) and mapped
to a matching deployment. Returns None (recording a metric) when the row can't be routed.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
@ -142,8 +142,43 @@ class CheckBatchCost:
return None
return model_id, get_batch_id_from_unified_batch_id(decoded)
if self._track_unmanaged_vertex_batch_cost:
return self._resolve_unmanaged_vertex_routing(job, prom_logger)
if self._track_unmanaged_batch_cost:
from litellm.llms.bedrock.batches.transformation import (
BedrockBatchesConfig,
)
from litellm.llms.vertex_ai.batches.transformation import (
VertexAIBatchTransformation,
)
input_file_id = self._get_input_file_id(job)
if VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
input_file_id
):
assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id
return self._resolve_unmanaged_provider_routing(
job=job,
prom_logger=prom_logger,
llm_provider="vertex_ai",
bare_model_name=VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
input_file_id
),
)
if BedrockBatchesConfig.is_unmanaged_s3_batch_input_file_id(input_file_id):
assert input_file_id is not None # narrowed by is_unmanaged_s3_batch_input_file_id
return self._resolve_unmanaged_provider_routing(
job=job,
prom_logger=prom_logger,
llm_provider="bedrock",
bare_model_name=BedrockBatchesConfig.get_bare_model_name_from_s3_file(
input_file_id
),
)
verbose_proxy_logger.info(
f"Skipping job {unified_object_id}: not a recognized unmanaged batch "
"(no gs:// or s3:// input_file_id with an embedded model)"
)
self._record_error(prom_logger, "invalid_unified_id")
return None
verbose_proxy_logger.info(
f"Skipping job {unified_object_id} because it is not a valid unified object id"
@ -151,36 +186,17 @@ class CheckBatchCost:
self._record_error(prom_logger, "invalid_unified_id")
return None
def _resolve_unmanaged_vertex_routing(
def _resolve_unmanaged_provider_routing(
self,
job: "LiteLLM_ManagedObjectTable",
prom_logger: Optional["PrometheusLogger"],
llm_provider: str,
bare_model_name: str,
) -> Optional[Tuple[str, str]]:
from litellm.llms.vertex_ai.batches.transformation import (
VertexAIBatchTransformation,
)
input_file_id = self._get_input_file_id(job)
if not VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
input_file_id
):
verbose_proxy_logger.info(
f"Skipping job {job.unified_object_id}: not an unmanaged vertex batch "
"(no gs:// input_file_id with a publishers/ model path)"
)
self._record_error(prom_logger, "invalid_unified_id")
return None
assert input_file_id is not None # narrowed by is_unmanaged_gcs_batch_input_file_id
bare_model_name = VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
input_file_id
)
deployment_id = self._get_vertex_ai_deployment_id_for_bare_model(
bare_model_name
)
deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
if deployment_id is None:
verbose_proxy_logger.info(
f"Skipping unmanaged vertex batch {job.unified_object_id}: no vertex_ai "
f"Skipping unmanaged {llm_provider} batch {job.unified_object_id}: no {llm_provider} "
f"deployment configured for model {bare_model_name}"
)
self._record_error(prom_logger, "unmanaged_no_matching_deployment")
@ -188,22 +204,22 @@ class CheckBatchCost:
return deployment_id, job.unified_object_id
def _get_vertex_ai_deployment_id_for_bare_model(
self, bare_model_name: str
def _get_deployment_id_for_bare_model(
self, bare_model_name: str, llm_provider: str
) -> Optional[str]:
model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
deployment_id = (
self._get_vertex_ai_deployment_id(model_group) if model_group else None
self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None
)
if deployment_id is not None:
return deployment_id
return self._get_vertex_ai_deployment_id_from_matching_deployments(
bare_model_name
return self._get_deployment_id_from_matching_deployments(
bare_model_name, llm_provider
)
def _get_vertex_ai_deployment_id_from_matching_deployments(
self, bare_model_name: str
def _get_deployment_id_from_matching_deployments(
self, bare_model_name: str, llm_provider: str
) -> Optional[str]:
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
@ -215,13 +231,13 @@ class CheckBatchCost:
if not self._is_bare_model_match(actual_model, bare_model_name):
continue
try:
_, llm_provider, _, _ = get_llm_provider(
_, deployment_llm_provider, _, _ = get_llm_provider(
model=actual_model,
custom_llm_provider=litellm_params.get("custom_llm_provider"),
)
except Exception:
continue
if llm_provider != "vertex_ai":
if deployment_llm_provider != llm_provider:
continue
model_info = deployment.get("model_info") or {}
deployment_id = model_info.get("id")
@ -231,15 +247,21 @@ class CheckBatchCost:
@staticmethod
def _is_bare_model_match(actual_model: str, bare_model_name: str) -> bool:
# Bedrock model ids may have ":" replaced with "-" in the S3 object key (see
# BedrockBatchesConfig.get_bare_model_name_from_s3_file), so normalize both sides;
# a no-op for providers like vertex_ai whose model ids never contain a colon.
normalized_actual = actual_model.replace(":", "-")
normalized_bare = bare_model_name.replace(":", "-")
return (
actual_model == bare_model_name
or actual_model.endswith(f"/{bare_model_name}")
or actual_model.endswith(f":{bare_model_name}")
normalized_actual == normalized_bare
or normalized_actual.endswith(f"/{normalized_bare}")
)
def _get_vertex_ai_deployment_id(self, model_group: str) -> Optional[str]:
def _get_deployment_id_for_provider(
self, model_group: str, llm_provider: str
) -> Optional[str]:
"""
Returns the first deployment id for `model_group` whose provider is vertex_ai,
Returns the first deployment id for `model_group` whose provider is `llm_provider`,
skipping deployments from other providers that happen to share the model group name.
"""
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
@ -249,13 +271,13 @@ class CheckBatchCost:
if deployment_info is None:
continue
try:
_, llm_provider, _, _ = get_llm_provider(
_, deployment_llm_provider, _, _ = get_llm_provider(
model=deployment_info.litellm_params.model,
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
)
except Exception:
continue
if llm_provider == "vertex_ai":
if deployment_llm_provider == llm_provider:
return deployment_id
return None

4
enterprise/pyproject.toml
View file

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

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

@ -54,6 +54,12 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` |
| `pdb.labels` | Extra metadata labels to add to the PDB | `{}` |
| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` |
| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` |
| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` |
| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` |
| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` |
#### Example `proxy_config` ConfigMap from values (default):
```
@ -94,6 +100,21 @@ data:
type: Opaque
```
#### Enterprise billable-request metering
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
```
kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
```
```
billingMetrics:
enabled: true
```
Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
### Database Settings
| Name | Description | Value |

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

@ -50,6 +50,53 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
{{- end }}
{{/*
Enterprise billable-request metering. The client certificate identifies the
deployment to LiteLLM's collector, so it is mounted read-only from an existing
Secret rather than passed through the environment.
*/}}
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
{{- define "litellm.billingMetricsEnv" -}}
- name: LITELLM_BILLING_METRICS_ENDPOINT
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
{{- if .Values.billingMetrics.caSecretName }}
- name: LITELLM_BILLING_METRICS_CA_CERT
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
{{- end }}
{{- with .Values.billingMetrics.exportIntervalMs }}
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: {{ . | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumes" -}}
- name: billing-metrics-mtls
secret:
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
secret:
secretName: {{ .Values.billingMetrics.caSecretName }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumeMounts" -}}
- name: billing-metrics-mtls
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
readOnly: true
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
readOnly: true
{{- end }}
{{- end -}}
{{/*
Create the name of the service account to use
*/}}
@ -76,10 +123,13 @@ so fall back to "default" (or an explicit override) to avoid a cyclic dependency
{{- end }}
{{/*
Get redis service name
Get redis service name.
The bundled Redis subchart only serves sentinel in "replication" architecture
(it rejects standalone + sentinel outright), and in that mode the sentinel
Service is named "<release>-redis", not "<release>-redis-master".
*/}}
{{- define "litellm.redis.serviceName" -}}
{{- if and (eq .Values.redis.architecture "standalone") .Values.redis.sentinel.enabled -}}
{{- if .Values.redis.sentinel.enabled -}}
{{- printf "%s-%s" .Release.Name (default "redis" .Values.redis.nameOverride | trunc 63 | trimSuffix "-") -}}
{{- else -}}
{{- printf "%s-%s-master" .Release.Name (default "redis" .Values.redis.nameOverride | trunc 63 | trimSuffix "-") -}}

15
helm/litellm-helm/templates/configmap-litellm.yaml
View file

@ -1,9 +1,22 @@
{{- if .Values.proxyConfigMap.create }}
{{- $config := deepCopy .Values.proxy_config }}
{{- if and .Values.redis.enabled (dig "coordination" "enabled" true .Values.redis) }}
{{- $generalSettings := (get $config "general_settings") | default dict }}
{{- if not (hasKey $generalSettings "coordination_redis") }}
{{- $coordinationRedis := dict "host" "os.environ/REDIS_HOST" "port" "os.environ/REDIS_PORT" "password" "os.environ/REDIS_PASSWORD" }}
{{- if .Values.redis.sentinel.enabled }}
{{- $sentinelNode := list (include "litellm.redis.serviceName" .) (include "litellm.redis.port" . | int) }}
{{- $coordinationRedis = dict "sentinel_nodes" (list $sentinelNode) "service_name" (default "mymaster" .Values.redis.sentinel.masterSet) "password" "os.environ/REDIS_PASSWORD" }}
{{- end }}
{{- $_ := set $generalSettings "coordination_redis" $coordinationRedis }}
{{- $_ := set $config "general_settings" $generalSettings }}
{{- end }}
{{- end }}
apiVersion: v1
kind: ConfigMap
metadata:
name: {{ include "litellm.fullname" . }}-config
data:
config.yaml: |
{{ .Values.proxy_config | toYaml | indent 6 }}
{{ $config | toYaml | indent 6 }}
{{- end }}

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

@ -142,6 +142,9 @@ spec:
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
@ -220,6 +223,9 @@ spec:
- name: npm
mountPath: /.npm
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -252,6 +258,9 @@ spec:
items:
- key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
path: "config.yaml"
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}

297
helm/litellm-helm/tests/billing_metrics_tests.yaml Normal file
View file

@ -0,0 +1,297 @@
suite: test billingMetrics wiring on the proxy deployment
templates:
- deployment.yaml
- configmap-litellm.yaml
- migrations-job.yaml
tests:
- it: is off by default, adding no env, volume, or mount
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- notContains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- it: renders the endpoint and the mounted cert paths when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: /etc/litellm/billing-mtls/tls.crt
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: /etc/litellm/billing-mtls/tls.key
# The conventional Secret name is the default, so enabling the block is enough.
- it: mounts the default cert secret read-only alongside the config volume
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- it: honours a secretName override
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: my-billing-mtls
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: my-billing-mtls
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- it: honours an endpoint override
template: deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: https://collector.internal:4318
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://collector.internal:4318
# The production collector presents a public web-PKI certificate, so the CA
# override must stay absent unless a private collector is configured.
- it: omits the CA env, volume, and mount when no caSecretName is set
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- notContains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- it: mounts the CA secret when caSecretName is set
template: deployment.yaml
set:
billingMetrics:
enabled: true
caSecretName: billing-ca
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- it: passes the export interval through only when set
template: deployment.yaml
set:
billingMetrics:
enabled: true
exportIntervalMs: 5000
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "5000"
- it: omits the export interval when unset
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "60000"
# Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
# entries must render after .Values.envVars or a user could silently redirect
# the metering export. The three billing entries are the last ones emitted here
# (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
- it: renders the billing endpoint after envVars so it cannot be shadowed
template: deployment.yaml
set:
migrationJob:
enabled: false
billingMetrics:
enabled: true
envVars:
LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://shadowed.example
- equal:
path: spec.template.spec.containers[0].env[-3]
value:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- equal:
path: spec.template.spec.containers[0].env[-2].name
value: LITELLM_BILLING_METRICS_CLIENT_CERT
- equal:
path: spec.template.spec.containers[0].env[-1].name
value: LITELLM_BILLING_METRICS_CLIENT_KEY
- it: keeps user-supplied volumes and mounts alongside the billing secret
template: deployment.yaml
set:
billingMetrics:
enabled: true
volumes:
- name: custom-callbacks
configMap:
name: my-callbacks
volumeMounts:
- name: custom-callbacks
mountPath: /app/callbacks
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: custom-callbacks
configMap:
name: my-callbacks
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: custom-callbacks
mountPath: /app/callbacks
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- it: still mounts the proxy config when enabled
template: deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: litellm-config
mountPath: /etc/litellm/config.yaml
subPath: config.yaml
# Only the proxy serves billable traffic. The migrations Job must never mount
# the client certificate, and it renders its own env and volumes, so nothing
# stops a future edit from wiring the billing include into it by mistake.
- it: does not touch the migrations job when enabled
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- notExists:
path: spec.template.spec.containers[0].volumeMounts
- notExists:
path: spec.template.spec.volumes
- it: fails loudly when enabled with an emptied secretName
template: deployment.yaml
set:
billingMetrics:
enabled: true
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
- it: fails loudly when enabled without an endpoint
template: deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true

143
helm/litellm-helm/tests/coordination_redis_tests.yaml Normal file
View file

@ -0,0 +1,143 @@
suite: test coordination redis
templates:
- configmap-litellm.yaml
- deployment.yaml
tests:
- it: should not render coordination_redis when redis is disabled
template: configmap-litellm.yaml
set:
redis.enabled: false
asserts:
- notMatchRegex:
path: data["config.yaml"]
pattern: coordination_redis
- it: should not emit redis env vars when redis is disabled
template: deployment.yaml
set:
redis.enabled: false
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: RELEASE-NAME-redis-master
any: true
- it: should render coordination_redis pointing at the bundled redis when enabled
template: configmap-litellm.yaml
set:
redis.enabled: true
asserts:
- matchRegex:
path: data["config.yaml"]
pattern: "coordination_redis:\n host: os.environ/REDIS_HOST\n password: os.environ/REDIS_PASSWORD\n port: os.environ/REDIS_PORT\n"
- matchRegex:
path: data["config.yaml"]
pattern: "master_key: os.environ/PROXY_MASTER_KEY"
- it: should emit redis env vars backing the coordination_redis os.environ refs
template: deployment.yaml
set:
redis.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: RELEASE-NAME-redis-master
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6379"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: RELEASE-NAME-redis
key: redis-password
- it: should not render coordination_redis when coordination is opted out
template: configmap-litellm.yaml
set:
redis.enabled: true
redis.coordination.enabled: false
asserts:
- notMatchRegex:
path: data["config.yaml"]
pattern: coordination_redis
- it: should keep emitting redis env vars when coordination is opted out
template: deployment.yaml
set:
redis.enabled: true
redis.coordination.enabled: false
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: RELEASE-NAME-redis-master
- it: should not clobber a user supplied coordination_redis block
template: configmap-litellm.yaml
set:
redis.enabled: true
proxy_config.general_settings.coordination_redis:
url: os.environ/COORDINATION_REDIS_URL
asserts:
- matchRegex:
path: data["config.yaml"]
pattern: "coordination_redis:\n url: os.environ/COORDINATION_REDIS_URL\n"
- notMatchRegex:
path: data["config.yaml"]
pattern: "host: os.environ/REDIS_HOST"
- it: should render sentinel_nodes and service_name in sentinel mode
template: configmap-litellm.yaml
set:
redis.enabled: true
redis.architecture: replication
redis.sentinel.enabled: true
asserts:
# The sentinel Service the redis subchart renders is "<release>-redis", and a
# plain client cannot speak the sentinel protocol, so host/port must not appear
- matchRegex:
path: data["config.yaml"]
pattern: "coordination_redis:\n password: os.environ/REDIS_PASSWORD\n sentinel_nodes:\n - - RELEASE-NAME-redis\n - 26379\n service_name: mymaster\n"
- notMatchRegex:
path: data["config.yaml"]
pattern: "host: os.environ/REDIS_HOST"
- it: should carry a custom sentinel masterSet into service_name
template: configmap-litellm.yaml
set:
redis.enabled: true
redis.architecture: replication
redis.sentinel.enabled: true
redis.sentinel.masterSet: litellm-master
asserts:
- matchRegex:
path: data["config.yaml"]
pattern: "service_name: litellm-master"
- it: should point REDIS_HOST at the sentinel service in sentinel mode
template: deployment.yaml
set:
redis.enabled: true
redis.architecture: replication
redis.sentinel.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: RELEASE-NAME-redis
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "26379"

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

@ -139,6 +139,20 @@ masterkeySecretName: ""
# if set, use this secret key for the master key; otherwise, use the default key
masterkeySecretKey: ""
# Optional: enterprise billable-request metering. When enabled, the proxy counts
# successful requests to inference, MCP, and A2A endpoints and pushes them to
# LiteLLM's collector over mutual TLS. Requires an enterprise license.
# The client certificate identifies the deployment, so it is mounted read-only
# from an existing Secret and never passed through the environment.
billingMetrics:
enabled: false
endpoint: https://telemetry.litellm.ai # collector to push the counter to
secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
# Only for private or test collectors whose server certificate is not on the
# public web PKI. The production collector needs no CA override.
caSecretName: "" # existing Secret holding ca.crt
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
proxyConfigMap:
# when true, creates a new configmap
create: true
@ -331,12 +345,28 @@ postgresql:
# secretKeys:
# userPasswordKey: password
# requires cache: true in config file
# either enable this or pass a secret for REDIS_HOST, REDIS_PORT, REDIS_PASSWORD or REDIS_URL
# with cache: true to use existing redis instance
# Redis is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend
# tracking, and the pod lock manager. Enabling this deploys the bundled Redis
# subchart, wires REDIS_HOST / REDIS_PORT / REDIS_PASSWORD into the proxy, and
# renders a `general_settings.coordination_redis` block into the proxy config.
#
# To point at an existing Redis instead, leave `enabled: false` and pass a
# secret for REDIS_HOST, REDIS_PORT, REDIS_PASSWORD or REDIS_URL; the proxy
# falls back to those env vars for coordination. Set `cache: true` in the proxy
# config only if you also want LLM response caching, which is independent of
# coordination
#
# 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
redis:
enabled: false
architecture: standalone
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
# `coordination_redis` block you define yourself in `proxy_config` always wins
enabled: true
# Prisma migration job settings
migrationJob:

5
helm/litellm/templates/NOTES.txt
View file

@ -46,4 +46,9 @@ Reminders:
- gateway.config.proxy_config (rendered into a ConfigMap and mounted at
/app/config/config.yaml; gateway reads it via
CONFIG_FILE_PATH)
- {component}.pdb.{enabled,minAvailable,maxUnavailable} (per-component PodDisruptionBudget; disabled by
default — with hpa.minReplicas of 1, minAvailable: 1
would block node drains)
- {component}.topologySpreadConstraints (standard k8s list, e.g. spread replicas across
topology.kubernetes.io/zone)
- Enable ingress.enabled=true to dispatch / → ui, gateway data-plane prefixes → gateway, and the catch-all → backend.

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

@ -34,6 +34,57 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
{{- end -}}
{{/*
Enterprise billable-request metering. Wired into gateway and backend, not the
migrations job. The gateway serves nearly all billable traffic, but the backend
keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
SpendLogs row, so metering only the gateway would silently drop that traffic.
The client certificate identifies the deployment to LiteLLM's collector, so it is
mounted read-only from an existing Secret rather than passed through the
environment.
*/}}
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
{{- define "litellm.billingMetricsEnv" -}}
- name: LITELLM_BILLING_METRICS_ENDPOINT
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
{{- if .Values.billingMetrics.caSecretName }}
- name: LITELLM_BILLING_METRICS_CA_CERT
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
{{- end }}
{{- with .Values.billingMetrics.exportIntervalMs }}
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: {{ . | quote }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumes" -}}
- name: billing-metrics-mtls
secret:
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
secret:
secretName: {{ .Values.billingMetrics.caSecretName }}
{{- end }}
{{- end -}}
{{- define "litellm.billingMetricsVolumeMounts" -}}
- name: billing-metrics-mtls
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
readOnly: true
{{- if .Values.billingMetrics.caSecretName }}
- name: billing-metrics-mtls-ca
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
readOnly: true
{{- end }}
{{- end -}}
{{/*
Per-component selector labels — used in both Service selectors and Deployment matchLabels.
*/}}
@ -213,6 +264,10 @@ harmless no-op for the Job and authoritative for the app pods.
*/}}
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{/* These feed the proxy's coordination Redis (cross-pod rate limits, spend
tracking, pod lock manager) via its REDIS_* env fallback. An explicit
`general_settings.coordination_redis` block in proxy_config takes
precedence over anything emitted here. */}}
{{- if $root.Values.redis.host }}
- name: REDIS_HOST
value: {{ $root.Values.redis.host | quote }}
@ -226,10 +281,11 @@ harmless no-op for the Job and authoritative for the app pods.
key: {{ $root.Values.redis.passwordSecret.passwordKey | default "password" }}
{{- end }}
{{- if $root.Values.redis.cluster }}
{{/* The proxy's Cache() reads REDIS_CLUSTER_NODES as JSON and constructs a
RedisClusterCache when it's set (litellm/caching/caching.py:169-192).
We seed with the single configured endpoint — the cluster client
discovers the remaining nodes from CLUSTER SLOTS at startup. */}}
{{/* The proxy falls back to REDIS_CLUSTER_NODES (JSON) to build a cluster-mode
coordination client when `general_settings.coordination_redis` is absent
and no plain-Redis response cache is configured. We seed with the single
configured endpoint; the cluster client discovers the remaining nodes from
CLUSTER SLOTS at startup. */}}
- name: REDIS_CLUSTER_NODES
value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }}
{{- end }}
@ -239,6 +295,52 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end -}}
{{/*
PodDisruptionBudget shared by gateway, backend, and ui.
Invoke with a dict:
(dict "root" $ "component" .Values.gateway "componentName" "gateway"
"fullname" (include "litellm.gateway.fullname" .)
"selectorLabels" (include "litellm.gateway.selectorLabels" .))
Renders nothing unless both the component and its `pdb.enabled` are on.
Only one of minAvailable / maxUnavailable should be set; if both are,
minAvailable wins. If neither is set, falls back to `maxUnavailable: 1` so
an enabled-but-unconfigured PDB still permits node drains.
"Set" means non-nil and non-empty-string, so an explicit 0 (e.g.
`maxUnavailable: 0` to forbid all voluntary disruptions) is honored rather
than silently replaced by the fallback.
*/}}
{{- define "litellm.pdb" -}}
{{- $root := .root -}}
{{- $component := .component -}}
{{- $min := $component.pdb.minAvailable -}}
{{- $max := $component.pdb.maxUnavailable -}}
{{- $minSet := not (or (kindIs "invalid" $min) (eq (printf "%v" $min) "")) -}}
{{- $maxSet := not (or (kindIs "invalid" $max) (eq (printf "%v" $max) "")) -}}
{{- if and $component.enabled $component.pdb $component.pdb.enabled }}
apiVersion: policy/v1
kind: PodDisruptionBudget
metadata:
name: {{ .fullname }}
labels:
{{- include "litellm.commonLabels" $root | nindent 4 }}
app.kubernetes.io/component: {{ .componentName }}
spec:
selector:
matchLabels:
{{- .selectorLabels | nindent 6 }}
{{- if $minSet }}
minAvailable: {{ $min }}
{{- else if $maxSet }}
maxUnavailable: {{ $max }}
{{- else }}
maxUnavailable: 1
{{- end }}
{{- end }}
{{- end -}}
{{/*
Renders `envFrom:` block for a component's `envConfigMaps` / `envSecrets`
lists. Each entry is a resource name; the chart wires the whole ConfigMap /

17
helm/litellm/templates/backend/deployment.yaml
View file

@ -44,14 +44,20 @@ spec:
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.backend.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -66,13 +72,16 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.backend.resources | nindent 12 }}
{{- if or .Values.gateway.config.create .Values.backend.volumes }}
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.backend.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}
@ -89,4 +98,8 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.backend.topologySpreadConstraints }}
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

6
helm/litellm/templates/backend/poddisruptionbudget.yaml Normal file
View file

@ -0,0 +1,6 @@
{{- include "litellm.pdb" (dict
"root" $
"component" .Values.backend
"componentName" "backend"
"fullname" (include "litellm.backend.fullname" .)
"selectorLabels" (include "litellm.backend.selectorLabels" .)) }}

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

@ -46,14 +46,20 @@ spec:
- name: NUM_WORKERS
value: {{ .Values.gateway.numWorkers | quote }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
{{- with .Values.gateway.volumeMounts }}
{{- toYaml . | nindent 12 }}
{{- end }}
@ -68,13 +74,16 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}
{{- with .Values.gateway.volumes }}
{{- toYaml . | nindent 8 }}
{{- end }}
@ -91,4 +100,8 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.gateway.topologySpreadConstraints }}
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

6
helm/litellm/templates/gateway/poddisruptionbudget.yaml Normal file
View file

@ -0,0 +1,6 @@
{{- include "litellm.pdb" (dict
"root" $
"component" .Values.gateway
"componentName" "gateway"
"fullname" (include "litellm.gateway.fullname" .)
"selectorLabels" (include "litellm.gateway.selectorLabels" .)) }}

4
helm/litellm/templates/ui/deployment.yaml
View file

@ -76,4 +76,8 @@ spec:
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.ui.topologySpreadConstraints }}
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

6
helm/litellm/templates/ui/poddisruptionbudget.yaml Normal file
View file

@ -0,0 +1,6 @@
{{- include "litellm.pdb" (dict
"root" $
"component" .Values.ui
"componentName" "ui"
"fullname" (include "litellm.ui.fullname" .)
"selectorLabels" (include "litellm.ui.selectorLabels" .)) }}

249
helm/litellm/tests/billing_metrics_tests.yaml Normal file
View file

@ -0,0 +1,249 @@
suite: test billingMetrics wiring on gateway and backend
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
- it: is off by default, adding no env, volume, or mount
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- equal:
path: spec.template.spec.containers[0].volumeMounts
value:
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- it: renders the endpoint and the mounted cert paths when enabled
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_CERT
value: /etc/litellm/billing-mtls/tls.crt
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CLIENT_KEY
value: /etc/litellm/billing-mtls/tls.key
- it: mounts the cert secret read-only alongside the config volume
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
# The production collector presents a public web-PKI certificate, so the CA
# override must stay absent unless a private collector is configured.
- it: omits the CA env, volume, and mount when no caSecretName is set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- notContains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- it: mounts the CA secret when caSecretName is set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
caSecretName: billing-ca
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_CA_CERT
value: /etc/litellm/billing-mtls-ca/ca.crt
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls-ca
secret:
secretName: billing-ca
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls-ca
mountPath: /etc/litellm/billing-mtls-ca
readOnly: true
- it: passes the export interval through only when set
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
exportIntervalMs: 5000
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
value: "5000"
- it: keeps user-supplied gateway volumes alongside the billing secret
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
gateway.volumes:
- name: custom-callbacks
configMap:
name: my-callbacks
gateway.volumeMounts:
- name: custom-callbacks
mountPath: /app/callbacks
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: custom-callbacks
configMap:
name: my-callbacks
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
# The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
# which writes a SpendLogs row, so it must meter too or that traffic is lost.
- it: meters the backend as well, since it serves the MCP transport
template: backend/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: billing-metrics-mtls
mountPath: /etc/litellm/billing-mtls
readOnly: true
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: billing-mtls
- it: leaves the backend alone when metering is off
template: backend/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
# The migrations job runs prisma and serves no traffic; it must never receive
# the client key.
- it: never mounts the billing cert on the migrations job
template: migrations-job.yaml
set:
billingMetrics:
enabled: true
secretName: billing-mtls
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BILLING_METRICS_ENDPOINT
value: https://telemetry.litellm.ai
- isNull:
path: spec.template.spec.volumes
# The conventional Secret name is the default, so enabling metering needs no
# secretName at all; the guard below only fires on an explicitly blanked one.
- it: uses the conventional secret name by default
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
asserts:
- contains:
path: spec.template.spec.volumes
content:
name: billing-metrics-mtls
secret:
secretName: litellm-billing-metrics-mtls
- it: fails loudly when the secretName is explicitly blanked
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
secretName: ""
asserts:
- failedTemplate:
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
- it: fails loudly when enabled without an endpoint
template: gateway/deployment.yaml
set:
billingMetrics:
enabled: true
endpoint: ""
secretName: billing-mtls
asserts:
- failedTemplate:
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true

188
helm/litellm/tests/pdb_topology_spread_tests.yaml Normal file
View file

@ -0,0 +1,188 @@
suite: test pod disruption budgets and topology spread constraints
templates:
- gateway/poddisruptionbudget.yaml
- backend/poddisruptionbudget.yaml
- ui/poddisruptionbudget.yaml
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: renders no PDB by default
templates:
- gateway/poddisruptionbudget.yaml
- backend/poddisruptionbudget.yaml
- ui/poddisruptionbudget.yaml
asserts:
- hasDocuments:
count: 0
- it: gateway PDB uses minAvailable and matches the gateway selector labels
template: gateway/poddisruptionbudget.yaml
set:
gateway.pdb.enabled: true
gateway.pdb.minAvailable: 1
asserts:
- isKind:
of: PodDisruptionBudget
- equal:
path: apiVersion
value: policy/v1
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway
- equal:
path: spec.minAvailable
value: 1
- notExists:
path: spec.maxUnavailable
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
- it: backend PDB uses maxUnavailable when minAvailable is unset
template: backend/poddisruptionbudget.yaml
set:
backend.pdb.enabled: true
backend.pdb.maxUnavailable: 25%
asserts:
- equal:
path: spec.maxUnavailable
value: 25%
- notExists:
path: spec.minAvailable
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: backend
- it: minAvailable wins when both minAvailable and maxUnavailable are set
template: gateway/poddisruptionbudget.yaml
set:
gateway.pdb.enabled: true
gateway.pdb.minAvailable: 2
gateway.pdb.maxUnavailable: 1
asserts:
- equal:
path: spec.minAvailable
value: 2
- notExists:
path: spec.maxUnavailable
- it: an explicit maxUnavailable 0 is honored instead of the fallback
template: backend/poddisruptionbudget.yaml
set:
backend.pdb.enabled: true
backend.pdb.maxUnavailable: 0
asserts:
- equal:
path: spec.maxUnavailable
value: 0
- notExists:
path: spec.minAvailable
- it: an explicit minAvailable 0 is honored and beats a set maxUnavailable
template: gateway/poddisruptionbudget.yaml
set:
gateway.pdb.enabled: true
gateway.pdb.minAvailable: 0
gateway.pdb.maxUnavailable: 1
asserts:
- equal:
path: spec.minAvailable
value: 0
- notExists:
path: spec.maxUnavailable
- it: enabled PDB with neither knob set falls back to maxUnavailable 1
template: ui/poddisruptionbudget.yaml
set:
ui.pdb.enabled: true
asserts:
- equal:
path: spec.maxUnavailable
value: 1
- notExists:
path: spec.minAvailable
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: ui
- it: renders no PDB for a disabled component even when its pdb is enabled
template: gateway/poddisruptionbudget.yaml
set:
gateway.enabled: false
gateway.pdb.enabled: true
asserts:
- hasDocuments:
count: 0
- it: deployments omit topologySpreadConstraints by default
templates:
- gateway/deployment.yaml
- backend/deployment.yaml
- ui/deployment.yaml
asserts:
- notExists:
path: spec.template.spec.topologySpreadConstraints
- it: gateway deployment renders configured topologySpreadConstraints
template: gateway/deployment.yaml
set:
gateway.topologySpreadConstraints:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
labelSelector:
matchLabels:
app.kubernetes.io/component: gateway
asserts:
- equal:
path: spec.template.spec.topologySpreadConstraints
value:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
labelSelector:
matchLabels:
app.kubernetes.io/component: gateway
- it: backend deployment renders configured topologySpreadConstraints
template: backend/deployment.yaml
set:
backend.topologySpreadConstraints:
- maxSkew: 1
topologyKey: kubernetes.io/hostname
whenUnsatisfiable: DoNotSchedule
labelSelector:
matchLabels:
app.kubernetes.io/component: backend
asserts:
- equal:
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
value: kubernetes.io/hostname
- equal:
path: spec.template.spec.topologySpreadConstraints[0].whenUnsatisfiable
value: DoNotSchedule
- it: ui deployment renders configured topologySpreadConstraints
template: ui/deployment.yaml
set:
ui.topologySpreadConstraints:
- maxSkew: 1
topologyKey: topology.kubernetes.io/zone
whenUnsatisfiable: ScheduleAnyway
asserts:
- equal:
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
value: topology.kubernetes.io/zone

109
helm/litellm/tests/redis_env_tests.yaml Normal file
View file

@ -0,0 +1,109 @@
suite: test redis coordination env vars
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: gateway omits redis env vars when no host is configured
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: redis.example.com
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway emits host, port and password when redis is configured
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.port: 6380
redis.passwordSecret.name: redis-secret
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: redis.example.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6380"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: password
- it: backend emits the same redis env vars so both pods coordinate on one redis
template: backend/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.name: redis-secret
redis.passwordSecret.passwordKey: redis-password
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: redis.example.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: redis-password
- it: gateway omits REDIS_PASSWORD for an auth-less redis
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
value: redis.example.com
- it: gateway seeds REDIS_CLUSTER_NODES from host and port in cluster mode
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.port: 6380
redis.cluster: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
value: '[{"host":"redis.example.com","port":6380}]'
- it: gateway omits REDIS_CLUSTER_NODES when cluster mode is off
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
any: true

64
helm/litellm/values.yaml
View file

@ -73,6 +73,25 @@ masterKey:
secretName: litellm-master-key-secret # name of a Secret containing the master key
secretKey: master-key
# Optional: enterprise billable-request metering. When enabled, the gateway and
# backend count successful requests to inference, MCP, and A2A endpoints and push
# them to LiteLLM's collector over mutual TLS. Both components serve billable
# routes: the backend keeps the named-server MCP transport. Requires an
# enterprise license. The client certificate identifies the deployment, so it is
# mounted read-only from an existing Secret and never passed through the env.
billingMetrics:
enabled: false
endpoint: https://telemetry.litellm.ai # collector to push the counter to
# An existing Secret holding the client certificate under tls.crt and its key
# under tls.key, usually created from the onboarding artifact. The default is
# the conventional name, so the common path is to create that Secret and set
# enabled: true. Override only if yours is named differently.
secretName: litellm-billing-metrics-mtls
# Only for private or test collectors whose server certificate is not on the
# public web PKI. The production collector needs no CA override.
caSecretName: "" # existing Secret holding ca.crt
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
# External Postgres connection.
database:
writer:
@ -100,7 +119,18 @@ database:
usernameKey: username
passwordKey: password
# Optional Redis (caching, rate limiting). Leave host empty to disable.
# Optional Redis. Leave host empty to disable.
#
# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend
# tracking, and the pod lock manager. The chart emits REDIS_HOST / REDIS_PORT /
# REDIS_PASSWORD, which the proxy picks up through its coordination Redis env
# fallback. Response caching is separate and off unless you enable it in
# `proxy_config.litellm_settings.cache`.
#
# For full control, define `general_settings.coordination_redis` in
# `proxy_config` (host/port/password/username/url/ssl/startup_nodes/
# sentinel_nodes/sentinel_password/service_name, each accepting os.environ/VAR
# refs). An explicit block overrides these env vars.
#
# Set `cluster: true` for Redis Cluster mode (e.g. AWS ElastiCache Cluster,
# self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from
@ -160,10 +190,28 @@ gateway:
maxReplicas: 10
targetCPUUtilizationPercentage: 70
targetMemoryUtilizationPercentage: 80
# PodDisruptionBudget for the gateway pods. Set exactly one of
# `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
# enabling without either falls back to `maxUnavailable: 1`). Disabled by
# default: with the default hpa.minReplicas of 1, a `minAvailable: 1` PDB
# would block node drains entirely.
pdb:
enabled: false
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
# Standard k8s topologySpreadConstraints for the gateway pods, e.g. to
# spread replicas across zones:
# - maxSkew: 1
# topologyKey: topology.kubernetes.io/zone
# whenUnsatisfiable: ScheduleAnyway
# labelSelector:
# matchLabels:
# app.kubernetes.io/component: gateway
topologySpreadConstraints: []
# ---------- backend (UI / management API) ----------
backend:
@ -203,10 +251,17 @@ backend:
minReplicas: 1
maxReplicas: 4
targetCPUUtilizationPercentage: 70
# Same shape as gateway.pdb.
pdb:
enabled: false
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
# Same shape as gateway.topologySpreadConstraints.
topologySpreadConstraints: []
# ---------- ui (Next.js static dashboard) ----------
ui:
@ -249,7 +304,14 @@ ui:
minReplicas: 1
maxReplicas: 3
targetCPUUtilizationPercentage: 80
# Same shape as gateway.pdb.
pdb:
enabled: false
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
# Same shape as gateway.topologySpreadConstraints.
topologySpreadConstraints: []

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260710000000_add_dcr_bridge_to_mcp_server_table/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "dcr_bridge" BOOLEAN;

6
litellm-proxy-extras/litellm_proxy_extras/migrations/20260713230852_add_key_type_to_litellm_verification_token/migration.sql Normal file
View file

@ -0,0 +1,6 @@
-- AlterTable
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "key_type" TEXT;
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "key_type" TEXT;

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260715000000_add_issuer_to_mcp_server_table/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "issuer" TEXT;

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;

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

@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable {
command String?
args String[] @default([])
env Json? @default("{}")
issuer String?
authorization_url String?
token_url String?
registration_url String?
@ -339,6 +340,7 @@ model LiteLLM_MCPServerTable {
available_on_public_internet Boolean @default(true)
delegate_auth_to_upstream Boolean @default(false)
oauth_passthrough Boolean @default(false)
dcr_bridge Boolean?
is_byok Boolean @default(false)
byok_description String[] @default([])
byok_api_key_help_url String?
@ -394,6 +396,13 @@ 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")
}
// Generate Tokens for Proxy
model LiteLLM_VerificationToken {
token String @id
@ -421,6 +430,7 @@ model LiteLLM_VerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
@ -515,6 +525,7 @@ model LiteLLM_DeletedVerificationToken {
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
key_type String?
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
@ -725,6 +736,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)
@ -756,6 +770,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)
@ -787,6 +804,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)
@ -817,6 +837,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)
@ -847,6 +870,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)
@ -879,6 +905,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)

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

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

19
litellm-rust/ADDING_A_PROVIDER.md
View file

@ -6,4 +6,23 @@ Three layers, same for every route (see `ocr` and `realtime` as references):
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.
## 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:** 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`.

9
litellm-rust/AGENTS.md
View file

@ -15,3 +15,12 @@ Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-
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.

44
litellm-rust/CLAUDE.md
View file

@ -2,6 +2,25 @@
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`
@ -20,6 +39,11 @@ Route-level Rust structure mirrors LiteLLM's Python responsibilities:
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
never inside `core`.
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
@ -77,6 +101,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

999
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

4
litellm-rust/Cargo.toml
View file

@ -15,8 +15,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"] }

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

@ -0,0 +1,53 @@
# Provider coding standards (litellm-rust)
Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MISTRAL_OCR_CONFIG`) is the reference; `messages` (`ANTHROPIC_MESSAGES_CONFIG`) is the next port.
## 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` = pure transforms/types (no network, env, secrets, auth, logging, global mutable state); `ai-gateway` = all I/O, auth headers, HTTP/SSE, lifecycle hooks; `python-bridge` = thin PyO3 adapter.
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: `<route>()` -> `prepare_*` -> `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing. Handlers validate and delegate; no business logic in them.
10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Env reads happen only at the host/config layer, with the `DEFAULT_*` fallback defined in `constants.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. Host 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. Rust paths stay off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven.
## Checks before push
22. 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
```

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

@ -41,3 +41,4 @@ python-config = ["dep:pyo3"]
[dev-dependencies]
futures-channel = "0.3"
tower = { version = "0.5.3", features = ["util"] }

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

@ -28,3 +28,31 @@ 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";
/// Full-request timeout ceiling for Anthropic Messages provider calls, in
/// seconds. Mirrors the Python Anthropic Messages default. The per-request
/// timeout from `litellm_params` still overrides this on the request builder.
pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
/// Connect timeout for Anthropic Messages provider calls, in seconds.
pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
/// Max characters of an upstream error body echoed across the host boundary
/// before truncation, so provider bodies are bounded and data-minimized.
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
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";
/// Provider name used by the Anthropic Messages route when a deployment's
/// provider model does not carry an explicit provider prefix.
pub(crate) const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
/// 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/messages.rs Normal file
View file

@ -0,0 +1 @@
pub use crate::messages::{messages, MessagesRequest};

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

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

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

@ -0,0 +1,548 @@
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::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::header::{HeaderName, AUTHORIZATION};
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
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 = self.socket.lock().await;
let Some(socket) = socket.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");
}
}

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

@ -12,6 +12,7 @@
//! for the load-time config reader.
pub mod io;
pub mod messages;
pub mod ocr;
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
@ -25,9 +26,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")]

15
litellm-rust/crates/ai-gateway/src/messages/client.rs Normal file
View file

@ -0,0 +1,15 @@
use std::sync::OnceLock;
use std::time::Duration;
use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS};
pub(super) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(MESSAGES_TIMEOUT_SECS))
.connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS))
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
}

52
litellm-rust/crates/ai-gateway/src/messages/common_utils.rs Normal file
View file

@ -0,0 +1,52 @@
use litellm_core::error::{json_type_name, CoreError};
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use litellm_core::CoreResult;
use serde_json::{Map, Value};
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
pub(super) fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
}
pub(super) fn messages_provider_config(
provider: &str,
) -> Option<&'static dyn AnthropicMessagesProviderConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG),
"azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG),
_ => None,
}
}
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"messages extra_headers.{key} must be a string, got {}",
json_type_name(&value)
))
})
})
.collect()
}
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
headers
.iter()
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}

83
litellm-rust/crates/ai-gateway/src/messages/handler.rs Normal file
View file

@ -0,0 +1,83 @@
use litellm_core::error::CoreError;
use litellm_core::CoreResult;
use serde_json::Value;
use super::client::http_client;
use super::common_utils::truncate_error_body;
use super::types::ProviderMessagesRequest;
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
pub(super) async fn execute_messages_provider_call(
request: ProviderMessagesRequest,
) -> CoreResult<Value> {
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response = serde_json::from_str(&text).map_err(|err| {
CoreError::InvalidResponse(format!("invalid messages response JSON: {err}"))
})?;
let transformed = request
.config
.transform_response(&request.model, response)?;
serde_json::to_value(transformed).map_err(|err| {
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
})
}
pub(super) async fn execute_messages_provider_stream(
request: ProviderMessagesRequest,
) -> CoreResult<reqwest::Response> {
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(CoreError::InvalidRequest(
"streaming messages is not supported for this provider".to_string(),
));
}
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
let status = response.status();
if !status.is_success() {
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
Ok(response)
}

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

@ -0,0 +1,49 @@
use litellm_core::CoreResult;
use serde_json::Value;
mod client;
mod common_utils;
mod handler;
mod prepare;
mod types;
pub use types::MessagesRequest;
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
use prepare::prepare_messages_call;
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<Value> {
match execute_messages(request, false).await? {
MessagesResponse::Json(body) => Ok(body),
MessagesResponse::Stream(response) => {
drop(response);
Err(litellm_core::CoreError::InvalidResponse(
"non-streaming messages execution returned a stream".to_string(),
))
}
}
}
pub(crate) enum MessagesResponse {
Json(Value),
Stream(reqwest::Response),
}
pub(crate) async fn execute_messages(
request: MessagesRequest<'_>,
stream: bool,
) -> CoreResult<MessagesResponse> {
let prepared = prepare_messages_call(request)?;
if stream {
execute_messages_provider_stream(prepared)
.await
.map(MessagesResponse::Stream)
} else {
execute_messages_provider_call(prepared)
.await
.map(MessagesResponse::Json)
}
}
#[cfg(test)]
mod tests;

73
litellm-rust/crates/ai-gateway/src/messages/prepare.rs Normal file
View file

@ -0,0 +1,73 @@
use litellm_core::messages::transformation::MessagesAuthStrategy;
use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
use litellm_core::CoreError;
use litellm_core::CoreResult;
use super::common_utils::{has_header, messages_provider_config, string_headers};
use super::types::{MessagesRequest, ProviderMessagesRequest};
pub(super) fn prepare_messages_call(
request: MessagesRequest<'_>,
) -> CoreResult<ProviderMessagesRequest> {
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.or_else(|| {
request
.custom_llm_provider
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
})
})
.ok_or_else(|| {
CoreError::InvalidProvider(
"unable to resolve custom_llm_provider for messages request".to_string(),
)
})?;
let model = provider_info.model.to_string();
let provider = provider_info.custom_llm_provider;
let config = messages_provider_config(provider)
.ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?;
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers(request.extra_headers)?;
let auth_strategy = config.auth_strategy();
if !has_header(&headers, auth_strategy.header_name()) {
let api_key = config.resolve_api_key(request.api_key, &env_lookup)?;
let auth_header = match auth_strategy {
MessagesAuthStrategy::Bearer => {
("authorization".to_string(), format!("Bearer {api_key}"))
}
MessagesAuthStrategy::Header(name) => (name.to_string(), api_key),
};
headers.push(auth_header);
}
for (name, value) in config.default_headers() {
if !has_header(&headers, name) {
headers.push((name.to_string(), value.to_string()));
}
}
let url = config.complete_url(request.api_base, &model, &env_lookup)?;
let typed_request = serde_json::from_value(request.body).map_err(|err| {
CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
})?;
let transformed = config.transform_request(typed_request)?;
let body = serde_json::to_value(transformed).map_err(|err| {
CoreError::InvalidRequest(format!(
"failed to serialize Anthropic messages request: {err}"
))
})?;
Ok(ProviderMessagesRequest {
provider: provider.to_string(),
model,
config,
url,
body,
upstream_headers: headers,
timeout: request.timeout,
})
}

305
litellm-rust/crates/ai-gateway/src/messages/tests.rs Normal file
View file

@ -0,0 +1,305 @@
use std::time::Duration;
use litellm_core::error::CoreError;
use serde_json::{json, Map, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{
has_header, messages_provider_config, string_headers, truncate_error_body,
};
use super::{messages, MessagesRequest};
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
fn write_response(body: &str) -> String {
format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
)
}
#[test]
fn provider_config_resolves_anthropic_and_azure_ai() {
assert!(messages_provider_config("anthropic").is_some());
assert!(messages_provider_config("azure_ai").is_some());
assert!(messages_provider_config("openai").is_none());
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(400);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert!(matches!(err, CoreError::InvalidRequest(_)));
}
#[test]
fn has_header_is_case_insensitive() {
let headers = vec![("X-Api-Key".to_string(), "secret".to_string())];
assert!(has_header(&headers, "x-api-key"));
assert!(!has_header(&headers, "authorization"));
}
#[tokio::test]
async fn messages_round_trip_builds_azure_request_and_passes_response_through() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let response = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({
"model": "claude-sonnet-4-5",
"max_tokens": 1024,
"messages": [{
"role": "user",
"content": [{
"type": "text",
"text": "hi",
"cache_control": {"type": "ephemeral", "scope": "global"}
}]
}]
}),
api_key: Some("sk-azure"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("messages request succeeds");
assert_eq!(response["content"][0]["text"], "hi");
assert_eq!(response["stop_reason"], "end_turn");
let request = server.await.expect("server task completes");
let (head, body) = request.split_once("\r\n\r\n").expect("has body");
assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}");
let head_lower = head.to_ascii_lowercase();
assert!(head_lower.contains("x-api-key: sk-azure"), "{head}");
assert!(
head_lower.contains("anthropic-version: 2023-06-01"),
"{head}"
);
assert!(
head_lower.contains("content-type: application/json"),
"{head}"
);
let sent_body: Value = serde_json::from_str(body).expect("body is json");
assert_eq!(
sent_body["messages"][0]["content"][0]["cache_control"],
json!({"type": "ephemeral"})
);
}
#[tokio::test]
async fn messages_round_trip_builds_native_anthropic_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let response = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({
"model": "claude-sonnet-4-5",
"max_tokens": 1024,
"messages": [{"role": "user", "content": "hi"}]
}),
api_key: Some("sk-ant"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("anthropic"),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("messages request succeeds");
assert_eq!(response["content"][0]["text"], "hi");
assert_eq!(response["stop_reason"], "end_turn");
let request = server.await.expect("server task completes");
let (head, _) = request.split_once("\r\n\r\n").expect("has body");
assert!(head.starts_with("POST /v1/messages "), "{head}");
let head_lower = head.to_ascii_lowercase();
assert!(head_lower.contains("x-api-key: sk-ant"), "{head}");
assert!(
head_lower.contains("anthropic-version: 2023-06-01"),
"{head}"
);
}
#[tokio::test]
async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body =
r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"x-api-key".to_string(),
Value::String("from-python".to_string()),
);
headers.insert(
"anthropic-beta".to_string(),
Value::String("token-efficient-tools-2025-02-19".to_string()),
);
messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: Some("rust-fallback-key"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("messages request succeeds");
let request = server.await.expect("server task completes");
let head = request
.split_once("\r\n\r\n")
.expect("has body")
.0
.to_ascii_lowercase();
let api_key_count = head
.lines()
.filter(|line| line.starts_with("x-api-key:"))
.count();
assert_eq!(api_key_count, 1, "{head}");
assert!(head.contains("x-api-key: from-python"), "{head}");
assert!(
head.contains("anthropic-beta: token-efficient-tools-2025-02-19"),
"{head}"
);
assert!(!head.contains("rust-fallback-key"), "{head}");
}
#[tokio::test]
async fn messages_maps_provider_error_status_to_http_error() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let _ = read_http_request(&mut socket).await;
let body = "unauthorized";
let response = format!(
"HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
});
let err = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: Some("sk-azure"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect_err("provider error propagates");
assert!(matches!(err, CoreError::Http { status: 401, .. }));
}
#[tokio::test]
async fn messages_rejects_unsupported_provider() {
let err = messages(MessagesRequest {
model: "claude-3-5-sonnet",
body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
api_key: Some("sk"),
api_base: Some("http://127.0.0.1:1"),
custom_llm_provider: Some("openai"),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
})
.await
.expect_err("unsupported provider errors");
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai"));
}

24
litellm-rust/crates/ai-gateway/src/messages/types.rs Normal file
View file

@ -0,0 +1,24 @@
use std::time::Duration;
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
use serde_json::{Map, Value};
pub struct MessagesRequest<'a> {
pub model: &'a str,
pub body: 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 timeout: Option<Duration>,
}
pub(crate) struct ProviderMessagesRequest {
pub(crate) provider: String,
pub(crate) model: String,
pub(crate) config: &'static dyn AnthropicMessagesProviderConfig,
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/python/config.rs
View file

@ -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"))

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::body::Body;
use axum::extract::{Json, State};
use axum::http::header::{HeaderMap, HeaderValue, CACHE_CONTROL, CONTENT_TYPE};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use axum::Router;
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::header::{CACHE_CONTROL, CONTENT_TYPE};
use axum::http::Request;
use axum::http::StatusCode;
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);
}
}

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

@ -0,0 +1,64 @@
use std::sync::Arc;
use litellm_core::router::Router;
use litellm_core::{CoreError, CoreResult};
use serde_json::{Map, Value};
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
use crate::messages::{execute_messages, MessagesRequest};
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,
};
let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
execute_messages(request, stream)
.await
.map(|response| match response {
crate::messages::MessagesResponse::Json(body) => MessagesResponse::Json(body),
crate::messages::MessagesResponse::Stream(upstream) => {
MessagesResponse::Stream(upstream)
}
})
}

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

@ -7,7 +7,9 @@
pub mod gil;
pub mod health;
pub mod messages;
pub mod realtime;
pub mod responses;
use axum::Router;
@ -18,6 +20,8 @@ pub fn app(state: AppState) -> Router {
Router::new()
.merge(health::router())
.merge(gil::router())
.merge(messages::router())
.merge(realtime::router())
.merge(responses::router())
.with_state(state)
}

348
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs Normal file
View file

@ -0,0 +1,348 @@
mod service;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::StatusCode;
use axum::response::Response;
use axum::routing::get;
use axum::Router;
use futures_util::{Sink, SinkExt, StreamExt};
use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, ResponsesWsEventType};
use litellm_core::router::Router as ModelRouter;
use serde::Deserialize;
use crate::auth::RequireMasterKey;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
use crate::state::AppState;
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
fn new_call_id() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or(0);
let sequence = CALL_SEQ.fetch_add(1, Ordering::Relaxed);
format!("respws-{nanos:x}-{sequence:x}")
}
pub fn router() -> Router<AppState> {
Router::new()
.route("/v1/responses", get(handle))
.route("/responses", get(handle))
}
#[derive(Debug, Deserialize)]
struct ResponsesQuery {
model: Option<String>,
}
async fn handle(
_auth: RequireMasterKey,
ws: WebSocketUpgrade,
State(state): State<AppState>,
Query(query): Query<ResponsesQuery>,
) -> Result<Response, (StatusCode, String)> {
if let Some(model) = query.model.as_deref() {
validate_model(&state.router, model)?;
}
let router = state.router.clone();
let loggers = state.loggers.clone();
let master_key = state.master_key.clone();
Ok(ws.on_upgrade(move |socket| bridge(socket, router, loggers, master_key, query.model)))
}
fn validate_model(router: &ModelRouter, model: &str) -> Result<(), (StatusCode, String)> {
if model.trim().is_empty() {
return Err((
StatusCode::BAD_REQUEST,
"missing 'model' query param".to_string(),
));
}
let Some(deployment) = router.get_available_deployment(model) else {
return Err((
StatusCode::NOT_FOUND,
format!("no deployment for model '{model}'"),
));
};
if deployment.litellm_params.model.contains('/')
&& !deployment.litellm_params.model.starts_with("openai/")
{
return Err((
StatusCode::BAD_REQUEST,
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}
Ok(())
}
async fn send_error_and_close<S>(sink: &mut S, message: String)
where
S: futures_util::Sink<Message> + Unpin,
S::Error: std::fmt::Display,
{
if let Ok(payload) = serde_json::to_string(&ResponsesErrorFrame::invalid_request(message)) {
let _ = sink.send(Message::Text(payload)).await;
}
let _ = sink
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
code: 1008,
reason: "Pre-call error".into(),
})))
.await;
let _ = sink.close().await;
}
struct ResponseClientSink {
sink: futures_util::stream::SplitSink<WebSocket, Message>,
}
impl Sink<ResponsesWsEvent> for ResponseClientSink {
type Error = axum::Error;
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_ready(context)
}
fn start_send(
mut self: std::pin::Pin<&mut Self>,
item: ResponsesWsEvent,
) -> Result<(), Self::Error> {
let payload = serde_json::to_string(&item).map_err(axum::Error::new)?;
std::pin::Pin::new(&mut self.sink).start_send(Message::Text(payload))
}
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_flush(context)
}
fn poll_close(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_close(context)
}
}
impl ResponseClientSink {
async fn close_with_code(&mut self, code: u16, reason: &'static str) {
let _ = self
.sink
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
code,
reason: reason.into(),
})))
.await;
let _ = self.sink.close().await;
}
}
async fn bridge(
socket: WebSocket,
router: Arc<ModelRouter>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
master_key: Option<Arc<str>>,
requested_model: Option<String>,
) {
let (mut ws_sink, ws_stream) = socket.split();
let (model, first_frame, stream) = if let Some(model) = requested_model {
(model, None, ws_stream)
} else {
let mut stream = ws_stream;
let first = match stream.next().await {
Some(Ok(Message::Text(text))) => {
match serde_json::from_str::<ResponsesWsEvent>(&text) {
Ok(event) => event,
Err(_) => {
send_error_and_close(
&mut ws_sink,
"Invalid JSON in response.create event".to_string(),
)
.await;
return;
}
}
}
_ => {
send_error_and_close(&mut ws_sink, "Missing response.create event".to_string())
.await;
return;
}
};
let Some(model) = first.model().filter(|value| !value.trim().is_empty()) else {
send_error_and_close(
&mut ws_sink,
"Missing model in response.create event".to_string(),
)
.await;
return;
};
if first.event_type != ResponsesWsEventType::ResponseCreate {
send_error_and_close(
&mut ws_sink,
"First frame must be a response.create event".to_string(),
)
.await;
return;
}
(model.to_string(), Some(first), stream)
};
if let Err((status, message)) = validate_model(&router, &model) {
let _ = status;
let _ = message;
send_error_and_close(&mut ws_sink, "Unknown model deployment".to_string()).await;
return;
}
let call_id = new_call_id();
let metadata = RequestMetadata {
user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token),
..RequestMetadata::default()
};
let client_in = Box::pin(stream.filter_map(|message| async move {
match message {
Ok(Message::Text(text)) => serde_json::from_str::<ResponsesWsEvent>(&text).ok(),
_ => None,
}
}));
let mut client_out = ResponseClientSink { sink: ws_sink };
let result = service::run(
&router,
&model,
first_frame,
None,
loggers,
call_id,
metadata,
client_in,
&mut client_out,
)
.await;
if result.is_err() {
client_out
.close_with_code(1011, "Internal server error")
.await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::io::realtime_pool::RealtimePool;
use crate::state::AppState;
use axum::body::Body;
use axum::http::Request;
use litellm_core::router::Router as ModelRouter;
use serde_json::json;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tower::ServiceExt;
struct RecordingSink {
messages: Vec<Message>,
}
impl Sink<Message> for RecordingSink {
type Error = std::convert::Infallible;
fn poll_ready(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.messages.push(item);
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn pre_call_error_matches_python_frame_and_close() {
let mut sink = RecordingSink {
messages: Vec::new(),
};
send_error_and_close(&mut sink, "missing model".to_string()).await;
let Message::Text(payload) = &sink.messages[0] else {
panic!("expected error text frame");
};
assert_eq!(
serde_json::from_str::<serde_json::Value>(payload).expect("error json"),
json!({
"type": "error",
"error": {
"type": "invalid_request_error",
"message": "missing model"
}
})
);
assert_eq!(
sink.messages[1],
Message::Close(Some(axum::extract::ws::CloseFrame {
code: 1008,
reason: "Pre-call error".into(),
}))
);
}
fn state() -> AppState {
AppState {
router: Arc::new(ModelRouter::default()),
master_key: Some(Arc::from("master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
}
}
#[tokio::test]
async fn auth_rejects_responses_upgrade_before_handler() {
let request = Request::builder()
.uri("/responses?model=known")
.body(Body::empty())
.expect("request");
let response = router()
.with_state(state())
.oneshot(request)
.await
.expect("response");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[test]
fn unknown_query_model_is_rejected_before_upgrade() {
assert_eq!(
validate_model(&ModelRouter::default(), "unknown").expect_err("unknown model"),
(
StatusCode::NOT_FOUND,
"no deployment for model 'unknown'".to_string()
)
);
}
}

156
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs Normal file
View file

@ -0,0 +1,156 @@
use std::sync::Arc;
use std::time::Duration;
use futures_util::{Sink, Stream};
use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
use litellm_core::responses::instrumentation::{
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
ResponsesWsMetadata,
};
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::{CoreError, CoreResult};
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::RequestMetadata;
#[allow(clippy::too_many_arguments)]
pub async fn run<In, Out>(
router: &litellm_core::router::Router,
model: &str,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
call_id: String,
metadata: RequestMetadata,
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 deployment = router.get_available_deployment(model).ok_or_else(|| {
CoreError::Routing(format!("no deployment available for model '{model}'"))
})?;
let params = &deployment.litellm_params;
let provider_model = params
.model
.strip_prefix("openai/")
.unwrap_or(&params.model);
if params.model.contains('/') && !params.model.starts_with("openai/") {
return Err(CoreError::InvalidProvider(
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}
let instrumentation = Arc::new(ResponsesWsInstrumentation::new(
call_id.clone(),
model,
ResponsesWsMetadata {
user_api_key_hash: metadata.user_api_key_hash,
user_api_key_user_id: metadata.user_api_key_user_id,
user_api_key_team_id: metadata.user_api_key_team_id,
},
));
let observer_instrumentation = Arc::clone(&instrumentation);
let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id);
let result = CallLifecycle::default()
.run(context, (), instrumentation.as_ref(), |_| async move {
crate::io::responses_ws::async_responses_websocket(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
first_frame,
idle_timeout,
move |event| {
observer_instrumentation.observe(event);
},
client_in,
client_out,
)
.await
})
.await;
let outcome = instrumentation.take_or_build_outcome(result.is_ok());
dispatch_outcome(loggers, outcome).await;
result
}
async fn dispatch_outcome(
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
outcome: ResponsesWsLogOutcome,
) {
let runner = CustomLoggerRunner::new(loggers.as_ref().clone());
match outcome {
ResponsesWsLogOutcome::Success { payload, callback } => {
let (details, response, start_time, end_time) = logging_values(payload, callback, None);
let _ = runner
.async_log_success_event(
&details,
&response,
CallbackTiming::new(start_time, end_time),
)
.await;
}
ResponsesWsLogOutcome::Failure {
payload,
callback,
error_message,
error_kind,
} => {
let error = LoggingError {
message: error_message,
kind: error_kind,
};
let (details, response, start_time, end_time) =
logging_values(payload, callback, Some(error));
let _ = runner
.async_log_failure_event(
&details,
Some(&response),
CallbackTiming::new(start_time, end_time),
)
.await;
}
}
}
fn logging_values(
payload: litellm_core::responses::instrumentation::ResponsesWsLogPayload,
callback: ResponsesWsCallbackPayload,
error: Option<LoggingError>,
) -> (ModelCallDetails, CallbackValue, f64, f64) {
let start_time = payload.start_time;
let end_time = payload.end_time;
let callback = CallbackValue::new(callback.object, callback.value);
let details = ModelCallDetails::from_standard_logging_payload(
crate::integrations::types::StandardLoggingPayload {
id: payload.id,
litellm_call_id: payload.litellm_call_id,
call_type: payload.call_type,
model: payload.model,
custom_llm_provider: payload.custom_llm_provider,
response_cost: payload.response_cost,
prompt_tokens: payload.usage.prompt_tokens,
completion_tokens: payload.usage.completion_tokens,
total_tokens: payload.usage.total_tokens,
start_time: payload.start_time,
end_time: payload.end_time,
stream: payload.stream,
metadata: crate::integrations::types::StandardLoggingMetadata {
user_api_key_hash: payload.metadata.user_api_key_hash,
user_api_key_user_id: payload.metadata.user_api_key_user_id,
user_api_key_team_id: payload.metadata.user_api_key_team_id,
..Default::default()
},
messages: None,
},
);
let details = match error {
Some(error) => details.with_failure_error(error),
None => details,
};
(details, callback, start_time, end_time)
}

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

@ -10,6 +10,25 @@ rand.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
sha2.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-sigv4 = { version = "1.5.1", optional = true }
aws-types = { version = "1.4.0", optional = true }
aws-smithy-runtime-api = { version = "1.13.0", optional = true }
[features]
default = []
bedrock-auth = [
"dep:aws-config",
"dep:aws-credential-types",
"dep:aws-sdk-sts",
"dep:aws-sigv4",
"dep:aws-types",
"dep:aws-smithy-runtime-api",
]
[dev-dependencies]
reqwest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

258
litellm-rust/crates/core/src/caching/in_memory_cache.rs Normal file
View file

@ -0,0 +1,258 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
pub struct InMemoryCache<V: Clone> {
pub cache_dict: HashMap<String, V>,
pub ttl_dict: HashMap<String, Duration>,
pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
pub max_size_in_memory: usize,
pub default_ttl: Duration,
now: Box<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
cache_dict: HashMap::new(),
ttl_dict: HashMap::new(),
expiration_heap: BinaryHeap::new(),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
now: Box::new(now),
}
}
pub fn evict_cache(&mut self) {
if self.max_size_in_memory == 0 {
return;
}
let current_time = (self.now)();
while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
self.expiration_heap.pop();
} else if expiration_time <= current_time {
self.expiration_heap.pop();
self.remove_key(&key);
} else {
break;
}
}
while self.cache_dict.len() >= self.max_size_in_memory {
let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
break;
};
if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
self.remove_key(&key);
}
}
}
pub fn allow_ttl_override(&self, key: &str) -> bool {
match self.ttl_dict.get(key).copied() {
None => true,
Some(expiration_time) => expiration_time < (self.now)(),
}
}
pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
if self.max_size_in_memory == 0 {
return;
}
self.evict_cache();
let key = key.into();
self.cache_dict.insert(key.clone(), value);
if self.allow_ttl_override(&key) {
let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
self.ttl_dict.insert(key.clone(), expiration_time);
self.expiration_heap.push(Reverse((expiration_time, key)));
}
}
// Generic values intentionally omit Python's per-item size check.
pub fn get_cache(&mut self, key: &str) -> Option<V> {
if self.cache_dict.contains_key(key) {
if self.is_key_expired(key) {
self.remove_key(key);
return None;
}
return self.cache_dict.get(key).cloned();
}
None
}
pub fn get_ttl(&self, key: &str) -> Option<Duration> {
self.ttl_dict.get(key).copied()
}
pub fn delete_cache(&mut self, key: &str) {
self.remove_key(key);
}
pub fn flush_cache(&mut self) {
self.cache_dict.clear();
self.ttl_dict.clear();
self.expiration_heap.clear();
}
fn is_key_expired(&self, key: &str) -> bool {
self.ttl_dict
.get(key)
.is_some_and(|expiration_time| *expiration_time < (self.now)())
}
fn remove_key(&mut self, key: &str) {
self.cache_dict.remove(key);
self.ttl_dict.remove(key);
}
}
#[cfg(test)]
mod tests {
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};
use super::InMemoryCache;
use std::time::Duration;
fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
Duration::from_secs(now.load(Ordering::Relaxed))
})
}
#[test]
fn ttl_expiry_is_deterministic() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), Some("value".to_string()));
now.store(161, Ordering::Relaxed);
assert_eq!(cache.get_cache("key"), None);
assert_eq!(cache.get_ttl("key"), None);
}
#[test]
fn default_and_per_set_ttl_are_applied() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("default", "value".to_string(), None);
cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
}
#[test]
fn unexpired_entries_do_not_allow_ttl_override() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_cache("key"), Some("second".to_string()));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
now.store(121, Ordering::Relaxed);
cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
}
#[test]
fn max_size_evicts_earliest_expiration() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 2, Duration::from_secs(60));
cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("early"), None);
assert!(cache.get_cache("late").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn expired_entries_are_evicted_before_live_entries() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
cache.set_cache(
"expired-one",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.set_cache(
"expired-two",
"value".to_string(),
Some(Duration::from_secs(20)),
);
cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
now.store(121, Ordering::Relaxed);
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
assert_eq!(cache.get_cache("expired-one"), None);
assert_eq!(cache.get_cache("expired-two"), None);
assert!(cache.get_cache("live").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn stale_heap_entries_are_skipped() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 1, Duration::from_secs(60));
cache.set_cache(
"removed",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.delete_cache("removed");
cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("removed"), None);
assert_eq!(cache.get_cache("kept"), None);
assert!(cache.get_cache("new").is_some());
}
#[test]
fn delete_and_flush_remove_values_and_ttls() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 10, Duration::from_secs(60));
cache.set_cache("one", "value".to_string(), None);
cache.set_cache("two", "value".to_string(), None);
cache.delete_cache("one");
assert_eq!(cache.get_cache("one"), None);
cache.flush_cache();
assert!(cache.cache_dict.is_empty());
assert!(cache.ttl_dict.is_empty());
assert!(cache.expiration_heap.is_empty());
}
#[test]
fn zero_max_size_does_not_cache() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 0, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), None);
assert!(cache.cache_dict.is_empty());
}
}

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

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

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

@ -0,0 +1,3 @@
pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com";
pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1";
pub const OPENAI_RESPONSES_PATH: &str = "/responses";

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