scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

Merge branch 'main' into bugfix/reasoning-not-redacted-in-provider-specific-fields

Browse source
This commit is contained in:
alchemicjon 2026-10-26 15:37:46 -04:00
commit 6b41dffcd9
No known key found for this signature in database
3423 changed files with 273420 additions and 41780 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -1291,7 +1291,7 @@ jobs:
name: Install MCP SDK1 peer
command: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- when:
condition:

2
.circleci/scripts/classify_changes.sh
View file

@ -33,7 +33,7 @@ while IFS= read -r file || [ -n "$file" ]; do
*) has_backend=true ;;
esac
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json)
has_cost_map=true ;;
tests/unit/proxy/*) : ;;
*) outside_cost_map_set=true ;;

90
.githooks/pre-push
View file

@ -1,90 +0,0 @@
#!/usr/bin/env bash
#
# pre-push — enforce Conventional Branches
# https://conventional-branch.github.io/
#
# Branch format: <type>/<description>
# <type> must be one of: feature, bugfix, hotfix, release, chore
#
# Protected branches (always allowed):
# - main
# - dependabot/*
# - gh-readonly-queue/*
#
# Tag pushes and branch deletions are skipped.
# Bypass: git push --no-verify.
set -eu
ZERO_OID="0000000000000000000000000000000000000000"
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
PROTECTED_NAMES="main"
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
is_protected() {
branch="$1"
for name in $PROTECTED_NAMES; do
if [ "$branch" = "$name" ]; then
return 0
fi
done
for prefix in $PROTECTED_PREFIXES; do
case "$branch" in "$prefix"*) return 0 ;; esac
done
return 1
}
invalid=""
while read -r local_ref local_oid remote_ref remote_oid; do
# Branch deletion (no local commit being pushed).
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
continue
fi
# Only validate branch pushes; ignore tags and other ref namespaces.
case "$remote_ref" in
refs/heads/*) ;;
*) continue ;;
esac
branch="${remote_ref#refs/heads/}"
if is_protected "$branch"; then
continue
fi
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
invalid="$invalid $branch"
fi
done
if [ -n "$invalid" ]; then
cat >&2 <<EOF
✗ Branch name does not follow Conventional Branches.
Invalid:$invalid
Expected: <type>/<description>
Allowed types: feature, bugfix, hotfix, release, chore
Examples:
feature/weighted-round-robin
bugfix/streaming-empty-chunks
chore/bump-deps
hotfix/auth-bypass
Protected (always allowed): main, dependabot/*, gh-readonly-queue/*.
See https://conventional-branch.github.io/
Rename with: git branch -m <new-name>
To bypass (use sparingly): git push --no-verify
EOF
exit 1
fi
exit 0

6
.github/CODEOWNERS vendored
View file

@ -1,2 +1,4 @@
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
/helm/ @yassin-berriai
/terraform/ @yassin-berriai

9
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -4,6 +4,13 @@
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
"## Affected release" with nothing under it must not appear in the final description -->
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
<details>
<summary>Details</summary>
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
@ -164,3 +171,5 @@ Example checklists:
## 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
</details>

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

@ -5,7 +5,7 @@ description: >-
test_paths:
- reason: >-
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
litellm.agent() end-to-end suite. It drives the real claude, codex, opencode and pi CLIs and
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and

69
.github/e2e-stack/assert_tests_ran.py vendored
View file

@ -1,69 +0,0 @@
import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
from coverage_registry.management_cases import MANAGEMENT_CASES
def main() -> int:
selected: Final = tuple(sys.argv[2:])
try:
report: Final = ET.parse(Path(sys.argv[1])).getroot()
except (ET.ParseError, OSError):
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
missing: Final = tuple(path for path in selected if path not in passed)
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
passed_nodes: Final = frozenset(
prop.get("value")
for case in cases
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
for prop in case.findall("./properties/property")
if prop.get("name") == "management_node"
)
missing_nodes: Final = required_nodes - passed_nodes
for node in sorted(missing_nodes):
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
for path in selected:
collected: Final = sum(case.get("file") == path for case in cases)
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
for case in cases:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
continue
outcome = "skipped" if case.find("skipped") is not None else "failed"
_ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
for prop in case.findall("./properties/property"):
name = prop.get("name", "")
value = prop.get("value", "")
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
r"[A-Za-z0-9_.:<>-]{1,240}", value
):
_ = sys.stdout.write(f" {name}: {value}\n")
if expected_count is not None and (
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
):
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
return 1
if (
selected
and not missing
and not missing_nodes
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
):
return 0
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
return 1
if __name__ == "__main__":
sys.exit(main())

58
.github/scripts/assert_ci_coverage.py vendored
View file

@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)")
DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
GLOB_CHARS = frozenset("*?")
DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"})
# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
# tests has to be named by some shard or it runs nowhere. A child listed here is
@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
return frozenset(ast.literal_eval(values))
def _owns(workflow: str, relative_path: str) -> bool:
owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow)
return owned_root is not None and relative_path.startswith(owned_root)
def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]:
return frozenset[str]().union(
*(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path))
)
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
runner: Final = repo_root / "tests/integration/run.py"
if not runner.exists():
@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
)
!= 1
)
gha_tokens: Final = _invoked_test_tokens(
scalar
for path in (repo_root / ".github/workflows").glob("*.y*ml")
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
workflow_tokens: Final = MappingProxyType(
{
path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name))
for path in (repo_root / ".github/workflows").glob("*.y*ml")
}
)
gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values())
circle_owned: Final = (
frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files
)
findings: Final = (
tuple(
Finding(path, "integration contract is also selected by GitHub Actions")
for path in paths
Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions")
for path in sorted(circle_owned)
if any(_token_covers(token, path) for token in gha_tokens)
)
+ tuple(
@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
for path in sorted(github_files)
if not (repo_root / path).is_file()
)
+ tuple(
Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow")
for path in sorted(github_files)
if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens))
)
)
browser_commands: Final = tuple(
scalar.value
@ -688,13 +710,9 @@ def main() -> int:
scalars = _all_scalars()
integration_paths, ownership_findings = _integration_ownership()
test_findings = (
_uncovered_tests(
allowlist,
_invoked_selections(scalars)
+ (Selection(included=integration_paths, ignored=frozenset()),),
)
+ ownership_findings
test_findings = _uncovered_tests(
allowlist,
_invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),),
)
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
@ -711,13 +729,25 @@ def main() -> int:
test_findings,
"Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if ownership_findings:
browser_remedy = (
" CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may "
"select them and the default Playwright config must exclude them."
if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings)
else ""
)
_report(
"integration contracts with the wrong CI owner",
ownership_findings,
"Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy,
)
if dockerfile_findings:
_report(
"Dockerfiles that no CI job builds",
dockerfile_findings,
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if stale_findings or test_findings or dockerfile_findings:
if stale_findings or test_findings or ownership_findings or dockerfile_findings:
return 1
_write(

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

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

91
.github/workflows/interactions-spec-canary.yml vendored Normal file
View file

@ -0,0 +1,91 @@
name: Interactions spec canary
on:
schedule:
- cron: "17 7 * * *"
pull_request:
paths:
- .github/workflows/interactions-spec-canary.yml
- tests/unit/interactions/**
- litellm/types/interactions/**
workflow_dispatch:
permissions: {}
concurrency: interactions-spec-canary
jobs:
interactions-spec-canary:
name: interactions-spec-canary
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 15
permissions:
contents: read
outputs:
compliance: ${{ steps.compliance.outcome }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project
- name: Run the compliance tests against Google's live spec
id: compliance
run: |
uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \
--interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \
--tb=short -vv -p no:rerunfailures -p no:pytest-retry \
--allow-hosts=127.0.0.1,::1,localhost,ai.google.dev
tracking-issue:
name: tracking-issue
needs: interactions-spec-canary
if: always() && github.event_name != 'pull_request' && github.ref_name == github.event.repository.default_branch && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
issues: write
steps:
- name: File the tracking issue unless the compliance step passed, close it when it did
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
env:
COMPLIANCE_OUTCOME: ${{ needs.interactions-spec-canary.outputs.compliance }}
RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
with:
script: |
const title = "Interactions spec canary is red";
const outcome = process.env.COMPLIANCE_OUTCOME;
const runUrl = process.env.RUN_URL;
const query = `repo:${context.repo.owner}/${context.repo.repo} is:issue is:open in:title "${title}"`;
const { data: found } = await github.rest.search.issuesAndPullRequests({ q: query, per_page: 20 });
const open = found.items.find((issue) => issue.title === title);
if (outcome === "success") {
if (open === undefined) return;
await github.rest.issues.createComment({ ...context.repo, issue_number: open.number, body: `Green again: ${runUrl}` });
await github.rest.issues.update({ ...context.repo, issue_number: open.number, state: "closed" });
return;
}
const state = outcome === "failure" ? "red" : `not run (compliance step outcome: ${outcome || "none"})`;
if (open !== undefined) {
await github.rest.issues.createComment({ ...context.repo, issue_number: open.number, body: `Still ${state}: ${runUrl}` });
return;
}
await github.rest.issues.create({
...context.repo,
title,
assignees: ["mateo-berri"],
body: `The scheduled run of tests/unit/interactions/test_openapi_compliance.py against Google's live Interactions spec is ${state}: ${runUrl}\n\nlitellm/types/interactions/README.md says how to refresh the snapshot and regenerate the types.`,
});

7
.github/workflows/sync-together-ai-models.yml vendored
View file

@ -41,10 +41,6 @@ jobs:
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
@ -57,8 +53,7 @@ jobs:
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json \
model_prices_and_context_window.schema.json
litellm/model_prices_and_context_window_backup.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"

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

@ -539,8 +539,8 @@ jobs:
- name: test_unit_passed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
- name: test_e2e_idp_stack
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py
- name: test_e2e_metadata
env:

4
.github/workflows/test-mcp-dependency-resolution.yml vendored
View file

@ -45,7 +45,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@ -132,7 +132,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:

180
.github/workflows/test-mcp-oauth-e2e.yml vendored
View file

@ -1,180 +0,0 @@
name: MCP OAuth happy path
on:
pull_request:
paths:
- '.github/workflows/test-mcp-oauth-e2e.yml'
- '.github/e2e-stack/**'
- 'tests/e2e/*.py'
- 'tests/e2e/pytest.ini'
- 'tests/e2e/idp_realm.json'
- 'tests/e2e/mcp/**'
- 'litellm/experimental_mcp_client/**'
- 'litellm/proxy/_experimental/mcp_server/**'
- 'litellm/proxy/auth/**'
- 'litellm/proxy/management_endpoints/mcp_management_endpoints.py'
- 'litellm/proxy/_types.py'
- 'litellm/types/mcp_server/mcp_server_manager.py'
- 'litellm/proxy/management_endpoints/*sso*.py'
- 'litellm/proxy/management_endpoints/sso/**'
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
- 'litellm/proxy/proxy_server.py'
- 'litellm/proxy/schema.prisma'
- 'ui/litellm-dashboard/src/app/connect/**'
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
- 'pyproject.toml'
- 'uv.lock'
workflow_dispatch:
permissions: {}
concurrency:
group: mcp-oauth-${{ github.ref }}
cancel-in-progress: true
jobs:
oauth:
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
environment: e2e-changed
timeout-minutes: 45
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_HOST: 127.0.0.1
DATABASE_PORT: '5432'
DATABASE_USER: litellm
DATABASE_PASSWORD: dbpassword9090
DATABASE_NAME: litellm
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
E2E_KEYCLOAK_ADMIN_USER: admin
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
E2E_FIXTURE_MODE: live
E2E_PROVIDER_CACHE: '0'
E2E_MCP_OAUTH_LIVE: '1'
E2E_REQUIRED_TEST_COUNT: '4'
steps:
- name: Checkout the tested source
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
persist-credentials: false
- name: Require and materialize the upstream login
env:
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
run: |
umask 077
python3 - <<'PY'
import base64
import json
import os
import secrets
from pathlib import Path
encoded = os.environ.get("STORAGE_STATE", "")
if not encoded:
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
state = json.loads(base64.b64decode(encoded, validate=True))
if not isinstance(state, dict) or not state.get("cookies"):
raise SystemExit("The captured login must contain browser cookies")
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
directory.mkdir(mode=0o700)
path = directory / "linear-state.json"
path.write_text(json.dumps(state))
with open(os.environ["GITHUB_ENV"], "a") as output:
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
value = "sk-e2e-" + secrets.token_hex(24)
print(f"::add-mask::{value}")
output.write(f"{name}={value}\n")
PY
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: '3.13'
- uses: ./.github/actions/setup-uv-with-retries
with:
version: '0.10.9'
- uses: ./.github/actions/cache-cargo-build
- name: Install the frozen E2E environment
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
uv run --no-sync python scripts/prisma_generate_if_needed.py
uv run --no-sync playwright install --with-deps chromium
- name: Configure license access
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: mcp-oauth-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Load the E2E license
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
test -n "${license}"
echo "::add-mask::${license}"
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
- name: Build the gateway consent UI at the tested commit
run: |
cd ui/litellm-dashboard
../../scripts/with_dashboard_node.sh npm ci
../../scripts/with_dashboard_node.sh npm run build
mkdir -p ../../litellm/proxy/_experimental/out
cp -r out/. ../../litellm/proxy/_experimental/out/
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
mkdir -p "${page%.html}"
mv "${page}" "${page%.html}/index.html"
done
- name: Prepare the isolated database and IdP
run: |
umask 077
bash .github/e2e-stack/start-idp.sh
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
- name: Run every required OAuth variant without retries
run: |
umask 077
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
- name: Report JUnit results and reject skipped or missing cases
if: always()
run: |
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
- name: Remove private login and logs
if: always()
run: |
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"

8
.github/workflows/test-merge-smoke.yml vendored
View file

@ -33,7 +33,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
steps:
@ -76,6 +76,12 @@ jobs:
- name: Import and CLI checks
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli
- name: Routing and UUID compatibility
run: >-
$IN_NETNS .venv/bin/python -m pytest
tests/unit/router_strategy/test_complexity_router.py::TestCustomDimensions
tests/unit/test_uuid_helper.py -q
- name: Proxy startup check
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics"

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

@ -16,6 +16,10 @@ on:
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/chat_completions/dispatch.py"
- "litellm/embeddings/dispatch.py"
- "litellm/messages/dispatch.py"
- "litellm/responses/dispatch.py"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "Makefile"
@ -46,6 +50,10 @@ on:
- "litellm/proxy/guardrails/**"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/chat_completions/dispatch.py"
- "litellm/embeddings/dispatch.py"
- "litellm/messages/dispatch.py"
- "litellm/responses/dispatch.py"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "Makefile"

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

@ -139,6 +139,7 @@ jobs:
test-path: |-
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
@ -215,10 +216,12 @@ jobs:
- shard: llms-providers
test-path: |-
tests/unit/llms/test_api_base_trailing_slash_routes.py
tests/unit/llms/test_cache_control_and_reasoning.py
tests/unit/llms/test_custom_llm.py
tests/unit/llms/test_file_content_block.py
tests/unit/llms/test_file_search_responses.py
tests/unit/llms/test_get_complete_url_trailing_slash.py
tests/unit/llms/test_lifecycle_fix.py
tests/unit/llms/test_oss_decision.py
tests/unit/llms/test_polling_url_origin_match.py
@ -266,6 +269,7 @@ jobs:
tests/unit/llms/featherless_ai
tests/unit/llms/firecrawl
tests/unit/llms/fireworks_ai
tests/unit/llms/friendliai
tests/unit/llms/gdc
tests/unit/llms/gemini
tests/unit/llms/gigachat
@ -321,6 +325,7 @@ jobs:
tests/unit/llms/pass_through
tests/unit/llms/perplexity
tests/unit/llms/pg_vector
tests/unit/llms/pi
tests/unit/llms/publicai
tests/unit/llms/ragflow
tests/unit/llms/recraft
@ -410,6 +415,7 @@ jobs:
tests/unit/embeddings
tests/unit/endpoints
tests/unit/files
tests/unit/fine_tuning
tests/unit/images
tests/unit/interactions
tests/unit/messages
@ -419,6 +425,7 @@ jobs:
tests/unit/realtime_api
tests/unit/rerank_api
tests/unit/search
tests/unit/text_completions
tests/unit/vector_stores
tests/unit/videos
python-version: "3.12"
@ -430,6 +437,13 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py
tests/test_litellm_rust/embeddings/test_dispatch.py
tests/test_litellm_rust/messages/test_dispatch.py
tests/test_litellm_rust/ocr/test_dispatch.py
tests/test_litellm_rust/responses/test_dispatch.py
tests/unit/compression
tests/unit/harness
tests/unit/rust_bridge
@ -658,6 +672,7 @@ jobs:
- shard: proxy-extras
test-path: >-
tests/unit/litellm_proxy_extras
tests/unit/db_scripts
python-version: "3.12"
workers: 2
reruns: 0
@ -696,7 +711,6 @@ jobs:
tests/unit/responses/litellm_completion_transformation
tests/unit/responses/test_additional_tools.py
tests/unit/responses/test_custom_tool_call.py
tests/unit/responses/test_dispatch.py
tests/unit/responses/test_metadata_codex_callback.py
tests/unit/responses/test_no_duplicate_spend_logs.py
tests/unit/responses/test_null_test_fix.py

4
AGENTS.md
View file

@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
`tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch

10
CONTRIBUTING.md
View file

@ -39,12 +39,12 @@ git clone https://github.com/YOUR_USERNAME/litellm.git
cd litellm
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
git checkout -b feature/your-feature
git checkout -b litellm_your_feature
# Install development dependencies
make install-dev
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
# Install git hooks that enforce commit conventions (one-time, opt-in)
make install-hooks
# Verify your setup works
@ -55,7 +55,7 @@ That's it! Your local development environment is ready.
## Commit and Branch Conventions
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/). Internal contributors name branches with the `litellm_` prefix and no `/`, and fork branches can use any name. Run `make install-hooks` once per clone to enable the local hook that enforces the commit convention. See the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for details
### 2. Development Workflow
@ -78,8 +78,8 @@ uv run pytest tests/unit/<your_test_file>.py -v
git add .
git commit -m "feat(scope): your descriptive commit message"
# Push and create a PR (branch must follow Conventional Branches — see above)
git push origin feature/your-feature
# Push and create a PR
git push origin litellm_your_feature
```
## Adding Testing

11
Dockerfile
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -61,6 +61,7 @@ RUN apk add --no-cache \
python-3.13-dev \
rust \
openssl \
libssl3 \
openssl-dev \
nodejs \
npm \
@ -120,8 +121,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder
COPY --from=uvbin /uv /usr/local/bin/uv
RUN apk add --no-cache python-3.13
ADD --checksum=sha256:2f7ae5cdd9d91731c0990e74a58239dc3e3fd2bf28dab23b55eafcdc47aaf87e \
https://github.com/BerriAI/litellm-admin-agent/archive/ef501e94bc9fbacb9233b922abf71427f030408c.tar.gz /tmp/liteadmin.tar.gz
ADD --checksum=sha256:4fb1847e28e2be30ba76dac94f02c4903291326516a0f4ba772f6ffb8b25d05c \
https://github.com/BerriAI/litellm-admin-agent/archive/5661c805ce479a0cf11b842ea133cf8d22cd7e34.tar.gz /tmp/liteadmin.tar.gz
RUN mkdir /tmp/liteadmin && tar xzf /tmp/liteadmin.tar.gz --strip-components=1 -C /tmp/liteadmin && \
uv venv /opt/liteadmin --python python3.13 && \
uv pip install --python /opt/liteadmin/bin/python --require-hashes -r /tmp/liteadmin/requirements.txt && \
@ -141,7 +142,7 @@ USER root
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app

4
Makefile
View file

@ -22,7 +22,7 @@ help:
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
@echo " make install-test-deps - Install the full local test environment"
@echo " make install-helm-unittest - Install helm unittest plugin"
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
@echo " make install-hooks - Install git hooks (Conventional Commits)"
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
@echo " make pre-commit - Legacy alias for make check"
@echo " make format - Apply ruff format code formatting"
@ -120,7 +120,7 @@ install-helm-unittest:
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.8.2; \
}
# Install git hooks that enforce Conventional Commits and Conventional Branches.
# Install git hooks that enforce Conventional Commits.
# Opt-in: not chained into install-dev.
install-hooks:
./scripts/install_git_hooks.sh

13
README.md
View file

@ -287,7 +287,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
</details>
<details>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode, pi or Deep Agents on any model (Python SDK)</summary>
### Python SDK - Agents
@ -296,7 +296,7 @@ import litellm
from litellm import Harness, sandbox
result = litellm.agent(
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.PI, Harness.DEEPAGENTS
"Find why tests/test_router.py is flaky and fix it.",
sandbox=sandbox.local("./repo"),
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
@ -305,7 +305,7 @@ result = litellm.agent(
print(result.text, result.cost, [f.path for f in result.files])
```
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents.
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
@ -384,6 +384,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | |
| [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | |
| [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | |
| [Nace (`nace`)](https://console.nace.ai/docs/drex) | | | | | | | | | | |
| [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | |
| [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ | | | | | | | |
| [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ | | | | | | | |
@ -615,15 +616,15 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws
### Verify Docker Image Signatures
All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Every release is signed with the same key introduced in [commit `0112e53`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0).
All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Releases are signed with the key introduced in [commit `e82a3c6`](https://github.com/BerriAI/litellm/commit/e82a3c646908ba0a6ee0beddc4310d7a6cac4065).
**Verify using the pinned commit hash (recommended):**
A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the original signing key:
A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the published signing key:
```bash
cosign verify \
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
--key https://raw.githubusercontent.com/BerriAI/litellm/e82a3c646908ba0a6ee0beddc4310d7a6cac4065/cosign.pub \
ghcr.io/berriai/litellm:<release-tag>
```

2
security.md → SECURITY.md
View file

@ -79,7 +79,7 @@ To qualify for a bounty, reports must include clear reproduction steps, a reprod
- Audit Logs with retention policy
- Control Allowed IP Addresses that can access your Cloud LiteLLM Instance
For security inquiries, please contact us at support@berri.ai
For security inquiries, please contact us at security@berri.ai
#### Supported data regions for LiteLLM Cloud

47
backend/Dockerfile
View file

@ -1,9 +1,36 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
# ---------- Rust extension ----------
# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
FROM $LITELLM_BUILD_IMAGE AS rust-native
WORKDIR /app
USER root
COPY --from=uvbin /uv /uvx /usr/local/bin/
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
ENV UV_PYTHON_DOWNLOADS=0
COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
python -c "import maturin; maturin.build_wheel('/wheels')"
# ---------- Builder ----------
FROM $LITELLM_BUILD_IMAGE AS builder
@ -16,7 +43,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -54,7 +81,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY . .
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-default-groups --no-editable \
uv sync --frozen --no-install-project --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -63,6 +90,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra saml \
--python python3.13
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@ -79,7 +109,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -94,7 +124,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \

1
backend/routes/allowlist.py
View file

@ -165,6 +165,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
BACKEND_MOUNT_PATHS: frozenset[str] = frozenset(
{
"/admin",
"/metrics",
"/swagger", # API documentation static assets belong to the backend
"/mcp", # lazily-mounted MCP sub-app serves on the backend component
}

41
ci_cd/cost_map_guard.py
View file

@ -1,11 +1,10 @@
"""Guard the cost map on pull requests.
Every pull request whose diff against its merge base touches one of the three cost map files gets the file
checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the
map. A pull request that leaves all three untouched skips them, since merging it keeps the base branch's copies
and its head tree only carries whatever state the branch was cut from. Pull requests from the cost map sync bot
(branches named litellm_cost_map_sync_*) always get the file checks and additionally may only touch those three
files and may only add or update models.
Every pull request whose diff against its merge base touches one of the two cost map files gets the file
checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
get the file checks and additionally may only touch those two files and may only add or update models.
"""
from __future__ import annotations
@ -19,13 +18,11 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
CostMap = dict[str, object]
@ -35,7 +32,6 @@ CostMap = dict[str, object]
class Snapshot:
cost_map: str
backup: str
schema: str
def _parse_object(text: str, path: str) -> CostMap | str:
@ -46,13 +42,6 @@ def _parse_object(text: str, path: str) -> CostMap | str:
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
def _rendered_schema(cost_map: CostMap) -> str:
try:
return render(build_schema(cost_map))
except SystemExit as error:
return str(error)
def _evaluation_output_price_failure(
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
) -> str | None:
@ -78,29 +67,13 @@ def _file_failures(
*,
evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
backup_failure: Final = (
()
if head.backup == head.cost_map
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
)
schema_failure: Final = (
()
if head.schema == schema_text
else (
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
)
)
return (
*backup_failure,
*schema_failure,
*(
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
*(
failure
for key, entry in head_map.items()

418
ci_cd/generate_model_prices_schema.py
View file

@ -1,418 +0,0 @@
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Final, Optional
import jsonschema
REPO_ROOT = Path(__file__).parent.parent
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
JsonSchema = dict
NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
WEEKDAY_PATTERN: Final = (
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
"gemini_native_audio",
"gemini_audio_only_live",
"uses_embed_content",
"use_openai_responses_path",
"bedrock_converse_supports_strict_tools",
"thinking_always_on",
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
"properties": {
"search_context_size_low": NONNEG_NUMBER,
"search_context_size_medium": NONNEG_NUMBER,
"search_context_size_high": NONNEG_NUMBER,
},
"additionalProperties": False,
},
"guardrail_cost_per_unit": {
"type": "object",
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
"additionalProperties": NONNEG_NUMBER,
},
"metadata": {
"type": "object",
"description": "Free-form notes about the entry (e.g. pricing derivation).",
},
"provider_specific_entry": {
"type": "object",
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
},
}
ARRAY_KEYS: dict[str, JsonSchema] = {
"supported_audio_formats": {
"type": "array",
"description": "Audio container formats the model can return.",
"items": {"type": "string", "enum": ["mp3", "wav"]},
},
"supported_endpoints": {
"type": "array",
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
"items": STRING,
},
"supported_modalities": {
"type": "array",
"description": "Input modalities the model accepts.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
},
"supported_output_modalities": {
"type": "array",
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"reasoning_effort_levels": {
"type": "array",
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
"items": STRING,
},
"tiered_pricing": {
"type": "array",
"description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
"items": {
"type": "object",
"properties": {
"range": {
"type": "array",
"description": "[min, max] prompt-token span this tier applies to.",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"max_results_range": {
"type": "array",
"description": "[min, max] result-count span this tier applies to (search models).",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
"input_cost_per_query": NONNEG_NUMBER,
},
"additionalProperties": False,
},
},
}
INTEGER_KEYS: dict[str, JsonSchema] = {
"max_tokens": {
**NONNEG_INTEGER,
"description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
},
"max_input_tokens": {
**NONNEG_INTEGER,
"description": "Maximum prompt/context tokens the model accepts.",
},
"max_output_tokens": {
**NONNEG_INTEGER,
"description": "Maximum tokens the model can generate in one response.",
},
"output_vector_size": {
**NONNEG_INTEGER,
"description": "Embedding dimension for embedding models.",
},
"prompt_cache_min_tokens": {
**NONNEG_INTEGER,
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
},
"token_count_framing": {
**NONNEG_INTEGER,
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
},
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
}
NUMBER_KEYS: dict[str, JsonSchema] = {
"regional_processing_uplift_multiplier_eu": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
},
"regional_processing_uplift_multiplier_us": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
},
"regional_endpoint_uplift_multiplier": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
},
}
COST_DESCRIPTIONS: dict[str, str] = {
"input_cost_per_token": "USD per prompt token.",
"output_cost_per_token": "USD per generated token.",
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
"google_maps_grounding_cost_per_query": (
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
),
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
"output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
}
def cost_description(key: str) -> Optional[str]:
if key in COST_DESCRIPTIONS:
return COST_DESCRIPTIONS[key]
if key.endswith("_flex"):
return "Flex service-tier rate for the same-named base field."
if key.endswith("_priority"):
return "Priority service-tier rate for the same-named base field."
if "_above_" in key:
return "Rate applied once the prompt exceeds the token threshold in the field name."
return None
def cost_schema(key: str) -> JsonSchema:
description = cost_description(key)
return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
return {
"litellm_provider": {
"type": "string",
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
},
"mode": {
"type": "string",
"description": "Primary API surface / task type of the model.",
"enum": list(modes),
},
"source": {
"type": "string",
"description": "URL of the provider pricing/model page this entry was taken from.",
},
"deprecation_date": {
"type": "string",
"description": "Date the provider deprecates the model, YYYY-MM-DD.",
"format": "date",
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
},
"web_search_billing_unit": {
"type": "string",
"description": "Whether web search is billed per query or per prompt.",
"enum": ["per_query", "per_prompt"],
},
"bedrock_output_config_effort_ceiling": {
"type": "string",
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"default_reasoning_effort": {
"type": "string",
"description": (
"Reasoning effort the provider applies when the request omits reasoning_effort. "
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
"accepted, which hold only when the effort resolves to 'none'."
),
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"token_count_fallback_model": {
"type": "string",
"minLength": 1,
"description": "Model that counts tokens for this one when the provider cannot count it.",
},
"audio_transcription_config": STRING,
"vertex_ai_audio_api": {
"type": "string",
"enum": ["lyria_predict", "lyria_interactions"],
},
}
def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
if key in curated:
return curated[key]
if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
return BOOLEAN
if "cost" in key:
return cost_schema(key)
return None
def build_schema(prices: dict) -> JsonSchema:
entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
if unclassified:
raise SystemExit(
f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
f"Add them to the key tables in {Path(__file__).name} and rerun it."
)
entry_properties = {key: classify(key, modes) for key in all_keys}
return {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "LiteLLM model_prices_and_context_window.json",
"description": (
"Schema for LiteLLM's model price and context window registry "
"(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
"Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
"optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
"All costs are USD per unit. New optional fields are added regularly, so consumers should "
"ignore unknown fields rather than reject them."
),
"type": "object",
"properties": {
"sample_spec": {
"type": "object",
"description": (
"Documentation placeholder illustrating the entry shape; not a real model and not "
"schema-conformant (several values are prose)."
),
},
"fallback_generalizations": {
"type": "object",
"description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
"properties": {
"rules": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": STRING,
"pattern": STRING,
"description": STRING,
},
"required": ["name", "pattern"],
"additionalProperties": True,
},
}
},
"additionalProperties": False,
},
},
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
"$defs": {
"modelEntry": {
"type": "object",
"description": (
"Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
"are optional; boolean capability flags are simply omitted when unknown or false."
),
"required": ["litellm_provider"],
"properties": entry_properties,
"additionalProperties": True,
}
},
}
def render(schema: JsonSchema) -> str:
return json.dumps(schema, indent=2) + "\n"
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)
)
def main() -> int:
check = "--check" in sys.argv[1:]
prices = json.loads(PRICES_PATH.read_text())
rendered = render(build_schema(prices))
errors = validation_errors(prices, json.loads(rendered))
if errors:
print(f"{PRICES_PATH.name} does not validate against the generated schema:")
print("\n".join(errors[:20]))
return 1
if not check:
SCHEMA_PATH.write_text(rendered)
print(f"wrote {SCHEMA_PATH}")
return 0
if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
print(
f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
)
return 1
print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
return 0
if __name__ == "__main__":
sys.exit(main())

41
cookbook/litellm_proxy_server/batch_api/bedrock/README.md Normal file
View file

@ -0,0 +1,41 @@
# Bedrock batch endpoints
Set `aws_bedrock_control_plane_endpoint` on a deployment when batch create, retrieve, and cancel requests must use a custom Bedrock control-plane endpoint. The proxy signs these requests for the configured URL. Without this setting, batch requests use the standard regional Bedrock endpoint
```yaml
model_list:
- model_name: bedrock-batch
litellm_params:
model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0
aws_region_name: us-east-1
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN
s3_output_bucket_name: os.environ/AWS_S3_OUTPUT_BUCKET_NAME
aws_bedrock_control_plane_endpoint: os.environ/BEDROCK_CONTROL_PLANE_ENDPOINT
```
Set `BEDROCK_CONTROL_PLANE_ENDPOINT` to your Bedrock control-plane VPC endpoint or compatible gateway base URL, without `/model-invocation-job`. The proxy adds the operation path. The endpoint must be reachable from the proxy, and the configured AWS region must match the endpoint's signing region
`api_base` and `aws_bedrock_runtime_endpoint` configure runtime inference separately. They do not override the batch control-plane endpoint. `s3_endpoint_url` configures S3 file access separately
The control-plane setting belongs in deployment configuration. Sending `aws_bedrock_control_plane_endpoint` in an API request body is rejected. There is no automatic `AWS_BEDROCK_CONTROL_PLANE_ENDPOINT` environment fallback, so reference the desired environment variable explicitly as shown above
Start the proxy with this configuration and set `LITELLM_API_BASE` to its URL. Set `BEDROCK_BATCH_INPUT_URI` to a valid input JSONL file in S3 that the batch service role can read
```bash
curl -sS "$LITELLM_API_BASE/v1/batches" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
-H 'Content-Type: application/json' \
-d "{\"model\":\"bedrock-batch\",\"input_file_id\":\"$BEDROCK_BATCH_INPUT_URI\",\"endpoint\":\"/v1/chat/completions\",\"completion_window\":\"24h\"}"
```
Use the returned `id` as `BATCH_ID` to retrieve or cancel the same job through the deployment's configured endpoint
```bash
curl -sS "$LITELLM_API_BASE/v1/batches/$BATCH_ID?model=bedrock-batch" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
curl -sS -X POST "$LITELLM_API_BASE/v1/batches/$BATCH_ID/cancel?model=bedrock-batch" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
```

4
cosign.pub
View file

@ -1,4 +1,4 @@
-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEKi4ivqGpE231OGH50PKbqy1Y1Kkb
POJC8+i2Wko82gBOUCe3M0Vw86H/4rhUhfoYEti4gdJ9wZbYmK0I2EE96g==
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEG06dpKM+vEDnmp2zcjyXgVCXxzRg
JA5H2LnmwBKQafZGka6OOrWujCEWATQewiAgyD4WQdJ0pvU9RC4mvNAjvQ==
-----END PUBLIC KEY-----

201
db_scripts/strip_vector_store_internal_metadata.py Normal file
View file

@ -0,0 +1,201 @@
import argparse
import os
import sys
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from itertools import accumulate, repeat, takewhile
from typing import Final, Literal, TextIO
import httpx
from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
from litellm.vector_stores.utils import is_litellm_written_metadata_key
DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
LIST_PAGE_SIZE: Final = "100"
class VectorStore(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
id: str
metadata: dict[str, JsonValue] | None = None
class VectorStorePage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
data: tuple[VectorStore, ...]
has_more: bool = False
class OpenAIErrorBody(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
message: str = ""
class OpenAIErrorResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
error: OpenAIErrorBody
@dataclass(frozen=True, slots=True)
class Options:
api_key: str
api_base: str
vector_store_ids: tuple[str, ...]
extra_keys: frozenset[str]
apply: bool
@dataclass(frozen=True, slots=True)
class StoreResult:
store_id: str
matched_keys: tuple[str, ...]
status: Literal["dry_run", "updated", "failed"]
error: str = ""
def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool:
return is_litellm_written_metadata_key(key) or key in extra_keys
def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]:
return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys)))
def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage:
params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after}
return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content)
def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None:
if page is None or not page.has_more or not page.data:
return None
return fetch_page(client, after=page.data[-1].id)
def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]:
pages: Final = accumulate(
repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None)
)
for page in takewhile(lambda page: page is not None, pages):
yield from page.data
def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]:
for store_id in store_ids:
yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content)
def describe_http_error(response: httpx.Response) -> str:
try:
message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message
except ValidationError:
return f"HTTP {response.status_code}"
return f"HTTP {response.status_code}: {message}"
def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult:
try:
response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)})
except httpx.TransportError as error:
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__)
if response.is_success:
return StoreResult(store_id=store_id, matched_keys=keys, status="updated")
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response))
def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None:
keys: Final = matching_keys(store, options.extra_keys)
if not keys:
return None
if not options.apply:
return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run")
return delete_keys(client, store.id, keys)
def format_result(result: StoreResult) -> str:
key_list: Final = ", ".join(result.matched_keys)
count: Final = len(result.matched_keys)
match result.status:
case "dry_run":
return f"{result.store_id}: would delete {count} keys: {key_list}"
case "updated":
return f"{result.store_id}: deleted {count} keys: {key_list}"
case "failed":
return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}"
def clean_and_report(
client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO
) -> Iterator[StoreResult | None]:
for store in stores:
result: Final = clean_store(client, store, options)
if result is not None:
print(format_result(result), file=out)
yield result
def run(options: Options, client: httpx.Client, out: TextIO) -> int:
stores: Final = (
retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client)
)
scanned: Final = tuple(clean_and_report(client, stores, options, out))
results: Final = tuple(result for result in scanned if result is not None)
updated: Final = sum(1 for result in results if result.status == "updated")
failed: Final = sum(1 for result in results if result.status == "failed")
mode: Final = "applied" if options.apply else "dry run, pass --apply to delete"
print(
f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})",
file=out,
)
return 1 if failed else 0
def parse_options(argv: Sequence[str] | None) -> Options:
parser: Final = argparse.ArgumentParser(
description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default."
)
parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY")
parser.add_argument("--api-base", default=DEFAULT_API_BASE)
parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store")
parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete")
parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them")
args: Final = parser.parse_args(argv)
if not args.api_key:
parser.error("--api-key or OPENAI_API_KEY is required")
return Options(
api_key=args.api_key,
api_base=args.api_base.rstrip("/"),
vector_store_ids=tuple(args.vector_store_id),
extra_keys=frozenset(args.extra_key),
apply=args.apply,
)
def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int:
options: Final = parse_options(argv)
with httpx.Client(
base_url=options.api_base,
headers={"Authorization": f"Bearer {options.api_key}"},
timeout=60.0,
transport=transport,
) as client:
try:
return run(options, client, sys.stdout)
except httpx.HTTPStatusError as error:
sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n")
return 2
except httpx.TransportError as error:
sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n")
return 2
except ValidationError:
sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n")
return 2
if __name__ == "__main__":
sys.exit(main())

7
docker/Dockerfile.database
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -59,6 +59,7 @@ RUN apk add --no-cache \
python-3.13 \
python-3.13-dev \
openssl \
libssl3 \
openssl-dev \
nodejs \
npm \
@ -122,7 +123,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
USER root
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app

7
docker/Dockerfile.non_root
View file

@ -1,8 +1,8 @@
# syntax=docker/dockerfile:1.7
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -60,6 +60,7 @@ RUN for i in 1 2 3; do \
coreutils \
curl \
openssl \
libssl3 \
libsndfile \
nodejs \
npm && break || sleep 5; \
@ -134,7 +135,7 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
apk add --no-cache python-3.13 bash openssl libssl3 tzdata libsndfile nodejs libevent && break || sleep 5; \
done
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer

2
docker/build_from_pip/Dockerfile.build_from_pip
View file

@ -30,7 +30,7 @@ RUN uv venv --python python && \
"google-cloud-aiplatform==1.133.0" \
"google-genai==1.37.0" \
"anthropic[vertex]==0.84.0" \
"grpcio==1.78.0" \
"grpcio==1.84.1" \
"prometheus-client==0.20.0" \
"langfuse==2.59.7" \
"opentelemetry-api==1.28.0" \

13
enterprise/enterprise_hooks/aporia_ai.py
View file

@ -15,9 +15,10 @@ sys.path.insert(
) # Adds the parent directory to the system path
import json
import sys
from typing import Any, List, Literal, Optional
from typing import Final, List, Literal, Optional
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -29,11 +30,13 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.llms.openai.common_utils import decoded_json_object
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.guardrails import GuardrailEventHooks
GUARDRAIL_NAME = "aporia"
_DECODED_JSON: Final = TypeAdapter(object)
class AporiaGuardrail(CustomGuardrail):
@ -68,7 +71,7 @@ class AporiaGuardrail(CustomGuardrail):
async def prepare_aporia_request(
self, new_messages: List[dict], response_string: Optional[str] = None
) -> dict:
data: dict[str, Any] = {}
data: dict[str, object] = {}
if new_messages is not None:
data["messages"] = new_messages
if response_string is not None:
@ -122,8 +125,10 @@ class AporiaGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Aporia AI response: %s", response.text)
if response.status_code == 200:
# check if the response was flagged
_json_response = response.json()
action: str = _json_response.get(
_json_response: Final = decoded_json_object(
_DECODED_JSON.validate_python(response.json())
)
action: Final = _json_response.get(
"action"
) # possible values are modify, passthrough, block, rephrase
if action == "block":

17
enterprise/enterprise_hooks/blocked_user_list.py
View file

@ -7,9 +7,10 @@
## This accepts a list of user id's for whom calls will be rejected
from typing import Literal, Optional
from typing import Final, Literal, Optional
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._internal_context import with_service_target
@ -19,6 +20,9 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import EndUserRepository
_REQUEST_BODY_VALUE: Final = TypeAdapter(object)
class ENTERPRISE_BlockedUserList(CustomLogger):
@ -94,15 +98,14 @@ class ENTERPRISE_BlockedUserList(CustomLogger):
)
if end_user_cache_obj is None and self.prisma_client is not None:
# check db
end_user_obj = (
await self.prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": user}
)
)
end_user_id: Final = _REQUEST_BODY_VALUE.validate_python(user)
end_user_obj = await EndUserRepository(
self.prisma_client
).table.find_unique(where={"user_id": end_user_id})
if end_user_obj is None: # user not in db - assume not blocked
end_user_obj = LiteLLM_EndUserTable(user_id=user, blocked=False)
cache.set_cache(key=cache_key, value=end_user_obj, ttl=60)
if end_user_obj is not None and end_user_obj.blocked is True:
if end_user_obj.blocked is True:
raise HTTPException(
status_code=400,
detail={

13
enterprise/enterprise_hooks/google_text_moderation.py
View file

@ -6,7 +6,10 @@
# +-----------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
from typing import Final
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -15,6 +18,8 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import iter_message_text
from litellm.types.utils import CallTypesLiteral
_THRESHOLD_VALUE: Final = TypeAdapter(object)
class ENTERPRISE_GoogleTextModeration(CustomLogger):
user_api_key_cache = None
@ -69,9 +74,11 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger):
f"{category}_confidence_threshold",
default_confidence_threshold,
)
set_confidence_value = getattr(
self,
f"{category}_confidence_threshold",
set_confidence_value = _THRESHOLD_VALUE.validate_python(
getattr(
self,
f"{category}_confidence_threshold",
)
)
verbose_proxy_logger.info(
f"Google Text Moderation: {category}_confidence_threshold: {set_confidence_value}"

7
enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py
View file

@ -1,15 +1,20 @@
# this is an example endpoint to receive data from litellm
from typing import Final
from fastapi import FastAPI, HTTPException, Request
from pydantic import TypeAdapter
app = FastAPI()
_LOGGED_EVENT: Final = TypeAdapter(object)
@app.post("/log-event")
async def log_event(request: Request):
try:
print("Received /log-event request") # noqa
# Assuming the incoming request has JSON data
data = await request.json()
data: Final = _LOGGED_EVENT.validate_python(await request.json())
print("Received request data:") # noqa
print(data) # noqa

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

@ -8,25 +8,30 @@
## This provides an LLM Guard Integration for content moderation on the proxy
import asyncio
from typing import Final, Optional
from collections.abc import Mapping
from typing import Final
import aiohttp
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.openai.common_utils import decoded_json_object
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import CallTypesLiteral
_DECODED_JSON: Final = TypeAdapter(object)
class _ENTERPRISE_LLMGuard(CustomLogger):
# Class variables or attributes
def __init__(
self,
mock_testing: bool = False,
mock_redacted_text: Optional[dict] = None,
mock_redacted_text: Mapping[str, object] | None = None,
):
self.mock_redacted_text = mock_redacted_text
self.llm_guard_mode = litellm.llm_guard_mode
@ -66,7 +71,9 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
async with session.post(
analyze_url, json={"prompt": text}
) as response:
redacted_text = await response.json()
redacted_text = _DECODED_JSON.validate_python(
await response.json()
)
verbose_proxy_logger.debug(
f"LLM Guard: Received response - {redacted_text}"
)
@ -77,12 +84,13 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
"error": f"Invalid content moderation response: {redacted_text}"
},
)
if redacted_text.get("is_valid", None) is False:
redacted_response: Final = decoded_json_object(redacted_text)
if redacted_response.get("is_valid", None) is False:
raise HTTPException(
status_code=400,
detail={"error": "Violated content safety policy"},
)
sanitized_prompt = redacted_text.get("sanitized_prompt")
sanitized_prompt: Final = redacted_response.get("sanitized_prompt")
return sanitized_prompt if isinstance(sanitized_prompt, str) else text
except Exception as e:
verbose_proxy_logger.exception(

16
enterprise/litellm_enterprise/integrations/custom_guardrail.py
View file

@ -20,9 +20,15 @@ class EnterpriseCustomGuardrailHelper:
the tag's value (the mode for that tag).
- If no request tag matches, fall back to default mode(s).
"""
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
)
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
)
from litellm.router_strategy.tag_based_routing import (
without_baggage_mapped_tags,
)
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import premium_user
@ -36,9 +42,13 @@ class EnterpriseCustomGuardrailHelper:
proxy_server_request = data.get("proxy_server_request", {})
request_tags = StandardLoggingPayloadSetup.get_request_tags(
litellm_params=data,
proxy_server_request=proxy_server_request,
request_tags = without_baggage_mapped_tags(
StandardLoggingPayloadSetup.get_request_tags(
litellm_params=data,
proxy_server_request=proxy_server_request,
),
data.get(get_metadata_variable_name_from_kwargs(data)),
keep_serving_deployment_tags=True,
)
# Check if any request tag matches a Mode tag key

4
enterprise/litellm_enterprise/proxy/auth/route_checks.py
View file

@ -2,6 +2,8 @@ import os
from fastapi import HTTPException, status
from litellm.proxy._types import MODEL_LIST_ROUTES
class EnterpriseRouteChecks:
@staticmethod
@ -43,7 +45,7 @@ class EnterpriseRouteChecks:
# Routes that should remain accessible even when LLM API endpoints are disabled.
# These are read-only model listing routes needed by the Admin UI.
LLM_API_EXEMPT_ROUTES = ["/models", "/v1/models"]
LLM_API_EXEMPT_ROUTES = MODEL_LIST_ROUTES
@staticmethod
def should_call_route(route: str):

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

@ -3,6 +3,8 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
"""
import asyncio
import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -12,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
row_metadata,
)
import litellm
@ -38,7 +41,12 @@ if TYPE_CHECKING:
from litellm.types.utils import LiteLLMBatch
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -83,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _row_metadata(row: object) -> dict[str, object]:
metadata: Final = getattr(row, "metadata", None)
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
class CheckBatchCost:
def __init__(
self,
@ -108,6 +111,11 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
self._line_item_store_turn: Final = asyncio.Lock()
self._poll_idle: Final = asyncio.Event()
@ -160,29 +168,9 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
return await self._find_jobs(
where={
"file_purpose": "batch",
"status": {
@ -196,10 +184,59 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
**where,
"created_at": {
"lte": before,
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
},
**(
{
"OR": [
{"created_at": {"gt": cursor[0]}},
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
]
}
if cursor is not None
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
if arrivals:
self._arrival_cursor = next_cursor
else:
self._poll_cursor = next_cursor
return page
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
@ -207,27 +244,22 @@ class CheckBatchCost:
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
return
if retired == 0:
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
"""
Atomically flip batch_processed from false to true, returning whether this pod won
the row. Every pod and uvicorn worker schedules its own poller against the shared
@ -241,10 +273,16 @@ class CheckBatchCost:
"""
if not self._has_batch_processed_column:
return True
from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json
try:
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": False},
data={"batch_processed": True},
data={
"batch_processed": True,
"status": response.status if response.status != "completed" else "complete",
"file_object": postgres_storable_json(response),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
@ -253,7 +291,9 @@ class CheckBatchCost:
return False
return claimed > 0
async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
async def _release_job_claim(
self, job: "ManagedObjectRow", original_status: str | None, original_file_object: object
) -> None:
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
Safe to match on batch_processed=True: while this poller is active the retrieve
@ -265,7 +305,11 @@ class CheckBatchCost:
try:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": True},
data={"batch_processed": False},
data={
"batch_processed": False,
"status": original_status,
"file_object": json.dumps(original_file_object),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
@ -299,13 +343,13 @@ class CheckBatchCost:
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
404 must not retire a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
@ -320,7 +364,7 @@ class CheckBatchCost:
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -342,18 +386,18 @@ class CheckBatchCost:
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
if finalized == 0:
return False
verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
return True
except Exception as db_err:
verbose_proxy_logger.error(
@ -564,8 +608,6 @@ class CheckBatchCost:
@staticmethod
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
file_object: Final = job.file_object
@ -608,7 +650,11 @@ class CheckBatchCost:
from litellm.files.types import FileContentCallOptions
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
mask_api_base_credentials,
resolve_vertex_location_for_cost,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -718,8 +764,7 @@ class CheckBatchCost:
)
except Exception as _e:
verbose_proxy_logger.warning(
f"CheckBatchCost: failed to create managed file ID for "
f"{_file_attr}={_raw_file_id!r}: {_e}"
f"CheckBatchCost: failed to create managed file ID for {_file_attr}={_raw_file_id!r}: {_e}"
)
# Pass the deployment's router-registered pricing (litellm_params custom
@ -738,6 +783,12 @@ class CheckBatchCost:
custom_llm_provider=batch_file_provider,
model_name=model_name,
model_info=deployment_model_info,
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
@ -755,8 +806,9 @@ class CheckBatchCost:
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
)
)
logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id
logging_obj = LiteLLMLogging(
model=batch_result.models[0],
model=logging_model,
messages=[{"role": "user", "content": "<retrieve_batch>"}],
stream=False,
call_type="aretrieve_batch",
@ -787,27 +839,35 @@ class CheckBatchCost:
custom_llm_provider=str(llm_provider) if llm_provider else None,
)
if not await self._claim_job_for_costing(job):
original_status: Final = job.status
original_file_object: Final = job.file_object
if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
"in this window, so its cost is already being tracked there"
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)
raise
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
finally:
if accounting.retryable_failure:
await self._release_job_claim(job, original_status, original_file_object)
if accounting.retryable_failure:
return None
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:
@ -876,7 +936,10 @@ class CheckBatchCost:
)
from litellm.batches.batch_utils import fetch_batch_managed_file_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
resolve_vertex_location_for_cost,
)
try:
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
@ -920,6 +983,12 @@ class CheckBatchCost:
cost_provider=cost_provider,
model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
input_content: Final = (
b""
@ -971,8 +1040,8 @@ class CheckBatchCost:
from litellm.proxy.proxy_server import proxy_config
from litellm.secret_managers.main import str_to_bool
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
api_key: Final = job.api_key if hasattr(job, "api_key") else None
team_id: Final = job.team_id if hasattr(job, "team_id") else None
has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
try:
key_row: Final = (
@ -990,7 +1059,9 @@ class CheckBatchCost:
if (has_key_row and key_row is None) or (team_id and team_row is None):
return True
settings: Final = get_dynamic_logging_metadata(
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
UserAPIKeyAuth(
metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id
),
proxy_config,
)
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
@ -1015,6 +1086,7 @@ class CheckBatchCost:
"""
try:
from litellm.integrations.prometheus import PrometheusLogger
prom_logger = PrometheusLogger.get_instance()
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
@ -1039,7 +1111,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await _managed_object_table(self.prisma_client).find_many(
jobs = await self._find_jobs(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1052,8 +1124,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:
@ -1099,10 +1169,7 @@ class CheckBatchCost:
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE
if (
response.status in PROVIDER_TERMINAL_BATCH_STATUSES
and response.output_file_id is not None
):
if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None:
try:
tracked = await self._track_completed_batch_cost(
job=job,
@ -1123,7 +1190,7 @@ class CheckBatchCost:
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
f"(job {job.id}): {tracking_err}"
)
self._record_error(prom_logger, "cost_tracking_error")
continue
@ -1133,24 +1200,19 @@ class CheckBatchCost:
# Track this job for the final metrics summary
processed_models.append(tracked)
# mark the job as complete
try:
update_data: dict = {
"status": response.status if response.status != "completed" else "complete",
"file_object": response.model_dump_json(),
}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data=update_data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
if not self._has_batch_processed_column:
# nothing claimed the row, so it stays in the poll page and its lines wait for that retry
if not self._has_batch_processed_column:
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data={
"status": response.status if response.status != "completed" else "complete",
"file_object": response.model_dump_json(),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
continue
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)

117
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
The spend event reaches the creating key's or team's own logging callbacks the way that creator's
requests do. Its message content is redacted when the creator's key or team logging turns message
logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before
the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or
whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only
and with its message content redacted.
A read that finds no row, because the create could not store one, stores it under the reading key
and bills it the same way, so a response the create left untracked is still billed once. A streamed
background create is free on its own; its stream stores the row under the creating key and bills
@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
ResponseCreator,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -54,12 +61,29 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
class _RoutedModelInfo(BaseModel):
model_config = ConfigDict(frozen=True)
id: str | None = None
class _RoutedDeployment(BaseModel):
model_config = ConfigDict(frozen=True)
model_info: _RoutedModelInfo = _RoutedModelInfo()
_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"})
_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True})
@dataclass(frozen=True, slots=True)
@ -76,6 +100,7 @@ class BilledBackgroundResponse:
custom_llm_provider: str | None
litellm_params: Mapping[str, object]
metadata: Mapping[str, object]
creator_logging: Mapping[str, object]
@dataclass(frozen=True, slots=True)
@ -94,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
return unpriced
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
def _string_callbacks(callbacks: object) -> tuple[str, ...]:
if not isinstance(callbacks, list):
return ()
return tuple(
callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str)
)
logging_obj: Final = LiteLLMLogging(
def _running_proxy_config() -> "ProxyConfig":
from litellm.proxy.proxy_server import proxy_config
return proxy_config
def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields
if not creator.settings_read:
return _REJECTED_CREATOR_LOGGING
try:
return key_and_team_logging_fields(
user_api_key_dict=UserAPIKeyAuth(
metadata=dict(creator.key_metadata),
team_id=creator.team_id,
team_metadata=dict(creator.team_metadata),
),
proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(),
)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _REJECTED_CREATOR_LOGGING
def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging":
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.utils import split_dynamic_success_callbacks
sync_callbacks, async_callbacks = split_dynamic_success_callbacks(
_string_callbacks(creator_logging.get("success_callback"))
)
return LiteLLMLogging(
model=billed.model,
messages=[{"role": "user", "content": "<retrieve_response>"}],
stream=False,
@ -106,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
start_time=datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id=str(uuid.uuid4()),
dynamic_success_callbacks=list(sync_callbacks),
dynamic_async_success_callbacks=list(async_callbacks) or None,
kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS},
)
def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging":
try:
return _billed_logging_obj(billed, billed.creator_logging)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING)
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj: Final = _creator_or_proxy_logging_obj(billed)
logging_obj.update_environment_variables(
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
optional_params={},
@ -174,6 +259,7 @@ class CheckResponsesCost:
prisma_client: "PrismaClient",
llm_router: "Router",
log_billed_response: BilledResponseLogger = log_billed_background_response,
proxy_config: "ProxyConfig | None" = None,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -183,6 +269,7 @@ class CheckResponsesCost:
self.llm_router: Router = llm_router
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
self._log_billed_response: Final = log_billed_response
self._proxy_config: Final = proxy_config
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
@ -294,11 +381,12 @@ class CheckResponsesCost:
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
)
async def _billed_response(
def _billed_response(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
creator: ResponseCreator,
) -> BilledBackgroundResponse:
deployment: Final = target.deployment
return BilledBackgroundResponse(
@ -310,12 +398,26 @@ class CheckResponsesCost:
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
metadata=MappingProxyType(
{
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
**creator.metadata,
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
**self._requested_model_group(target, job.team_id),
}
),
creator_logging=creator_logging_fields(creator, self._proxy_config),
)
def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]:
requested: Final = target.model_metadata.get("model_group")
if target.deployment is None or not isinstance(requested, str):
return MappingProxyType({})
routed: Final = _ROUTED_DEPLOYMENTS.validate_python(
self.llm_router.get_model_list(model_name=requested, team_id=team_id) or ()
)
serving_id: Final = target.deployment.model_info.id
return MappingProxyType(
{"model_group": requested} if any(row.model_info.id == serving_id for row in routed) else {}
)
async def _bill_once(
@ -328,7 +430,8 @@ class CheckResponsesCost:
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
return None
try:
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id)
charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator))
except BaseException:
await self._release_billing_claim(job)
raise

207
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -9,12 +9,17 @@ from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, cast
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLI_SESSION_KEY_PREFIX
from litellm.repositories.table_repositories import (
DeletedVerificationTokenRepository,
DeprecatedVerificationTokenRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -27,6 +32,12 @@ if TYPE_CHECKING:
class ManagedObjectRow(Protocol):
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def id(self) -> str: ...
@ -63,6 +74,7 @@ class ManagedObjectRow(Protocol):
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
_MAX_ROTATION_HOPS: Final = 16
def _string_tags(request_tags: object) -> list[str] | None:
@ -89,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL
return TeamRepository(prisma_client).table
def _rotated_token_table(
prisma_client: PrismaClient,
) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]:
return DeprecatedVerificationTokenRepository(prisma_client).table
def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]:
return DeletedVerificationTokenRepository(prisma_client).table
def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]:
return TeamRepository(prisma_client).deleted_table
@dataclass(frozen=True, slots=True)
class _CreatorRows:
key: prisma_models.LiteLLM_VerificationToken | None
@ -96,6 +122,15 @@ class _CreatorRows:
user: prisma_models.LiteLLM_UserTable | None
@dataclass(frozen=True, slots=True)
class ResponseCreator:
metadata: Mapping[str, object]
key_metadata: Mapping[str, object]
team_id: str | None
team_metadata: Mapping[str, object]
settings_read: bool = True
class ManagedObjectCreatorAttribution:
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
job's spend row is attributed like the creating request: key hash, team, organization,
@ -107,8 +142,15 @@ class ManagedObjectCreatorAttribution:
user's alias, because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
to increment organization spend, so leaving it out drops the cost from org accounting.
A background response also carries its end user and the key, team, and user budgets, the
way the read that billed it before did, and never names a key-scoped creator as its user.
A background response also carries its end user, the creating key's project and its alias,
and the key, team, and user budgets and per-model budgets, the way the read that billed it
before did, and never names a key-scoped creator as its user. A key that sets no budget of
its own takes its budget tier's, the same rule auth applies. Its creator's key and team
metadata come along too, so the billed event reaches the key's and team's logging callbacks
like that read did. A rotated key's metadata is read from the live key its rotations lead
to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the
creator's settings unread and the attribution to the rows that still load, so the
response still bills and still counts toward their budgets.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
@ -119,17 +161,70 @@ class ManagedObjectCreatorAttribution:
rows: Final = await self._rows(job, object_id)
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
rows: Final = await self._rows(job, object_id)
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
}
async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
try:
return await self._read_response_creator(job, object_id)
except Exception as e:
verbose_proxy_logger.error(
f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their "
f"settings: {e}"
)
return ResponseCreator(
metadata=_response_creator_metadata(job, await self._rows(job, object_id)),
key_metadata={},
team_id=getattr(job, "team_id", None),
team_metadata={},
settings_read=False,
)
async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
metadata=_response_creator_metadata(job, rows),
key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key),
team_id=team_id,
team_metadata=(
row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id)
),
)
async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]:
if not _names_a_key_row(api_key):
return {}
current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS)
current_key: Final = (
await _token_table(self._prisma_client).find_unique(where={"token": current_token})
if current_token != api_key
else None
)
if current_key is not None:
return row_metadata(current_key)
return row_metadata(
await _deleted_token_table(self._prisma_client).find_first(
where={"token": current_token}, order={"deleted_at": "desc"}
)
)
async def _token_rotated_to(self, token: str, hops_left: int) -> str:
rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token})
if rotation is None or hops_left == 0:
return token
return await self._token_rotated_to(rotation.active_token_id, hops_left - 1)
async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]:
if not team_id:
return {}
return row_metadata(
await _deleted_team_table(self._prisma_client).find_first(
where={"team_id": team_id}, order={"deleted_at": "desc"}
)
)
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
return _CreatorRows(
@ -152,23 +247,57 @@ class ManagedObjectCreatorAttribution:
async def _key_row(
self, object_id: str, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
try:
return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
return await self._live_key_row(api_key, created_by)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
return None
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
async def _live_key_row(
self, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
return await _token_table(self._prisma_client).find_unique(
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
)
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
try:
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
return await self._live_team_row(team_id)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
"user_api_key_team_spend",
"user_api_key_team_max_budget",
"user_api_key_team_model_max_budget",
),
**_budget_fields(
rows.user,
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_user_model_max_budget",
),
}
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
api_key: Final = getattr(job, "api_key", None)
@ -196,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str,
}
def row_metadata(row: object) -> Mapping[str, object]:
metadata: Final = getattr(row, "metadata", None)
if not isinstance(metadata, dict):
return {}
return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column
def _names_a_key_row(api_key: object) -> TypeGuard[str]:
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
return isinstance(api_key, str) and is_valid_sha256_hash(api_key)
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
@ -212,7 +354,28 @@ def _organization_id(row: object) -> str | None:
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
if row is None:
return {}
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}
return {
spend_key: getattr(row, "spend", None),
max_budget_key: getattr(row, "max_budget", None),
model_max_budget_key: getattr(row, "model_max_budget", None),
}
def _key_budget_fields(key: object) -> Mapping[str, object]:
if key is None:
return {}
tier: Final = getattr(key, "litellm_budget_table", None)
return {
"user_api_key_spend": getattr(key, "spend", None),
"user_api_key_max_budget": _own_or_tier_budget(key, tier, "max_budget"),
"user_api_key_model_max_budget": _own_or_tier_budget(key, tier, "model_max_budget"),
}
def _own_or_tier_budget(key: object, tier: object, field: str) -> object:
own: Final = getattr(key, field, None)
tier_value: Final = getattr(tier, field, None)
return own if tier_value is None or own not in (None, {}) else tier_value

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

@ -6,6 +6,7 @@ import base64
import json
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -43,6 +44,7 @@ from litellm.constants import (
)
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.internal_call_metadata import requests_background
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -132,7 +134,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
Span = Union[_Span, Any]
Span = _Span
InternalUsageCache = _InternalUsageCache
PrismaClient = _PrismaClient
else:
@ -416,7 +418,7 @@ def _storable_json_value(value: JsonValue) -> JsonValue:
return value
def _postgres_storable_json(file_object: BaseModel) -> str:
def postgres_storable_json(file_object: BaseModel) -> str:
dumped: Final = file_object.model_dump_json()
if "\\u0000" not in dumped:
return dumped
@ -479,9 +481,24 @@ def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str,
_MANAGED_FILES_TARGET: Final = "managed_files"
_PROVIDER_FILE_ID_KWARGS: Final = MappingProxyType(
{
CallTypes.acreate_batch: ("input_file_id",),
CallTypes.acreate_fine_tuning_job: ("training_file", "validation_file"),
}
)
_PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDER_FILE_ID_KWARGS.values())))
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
class _SupportsKeyedGet(Protocol):
def get(self, key: str, default: object, /) -> object: ...
def _keyed_get(source: object, key: str, default: object = None) -> object:
return require_attributes(source, _SupportsKeyedGet, "get").get(key, default)
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(
@ -650,7 +667,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
# FIX: Update status and file_object on every operation to keep state in sync
update_columns: Final = {
"file_object": _postgres_storable_json(stored_object),
"file_object": postgres_storable_json(stored_object),
"status": file_object.status,
"updated_by": user_api_key_dict.user_id,
}
@ -666,7 +683,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
data={
"create": {
"unified_object_id": unified_object_id,
"file_object": _postgres_storable_json(stored_object),
"file_object": postgres_storable_json(stored_object),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": resolve_resource_owner_id(user_api_key_dict),
@ -702,7 +719,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"unified_object_id": response_id,
"model_object_id": decrypted_response_id(response_id),
"file_purpose": "response",
"file_object": _postgres_storable_json(_without_model_output(response)),
"file_object": postgres_storable_json(_without_model_output(response)),
"status": PENDING_RESPONSE_STATUSES[0],
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
@ -1562,6 +1579,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if isinstance(file_id, str) and file_id
)
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
if fine_tuning_file_ids:
data["model_file_id_mapping"] = await self.get_model_file_id_mapping(
fine_tuning_file_ids, user_api_key_dict.parent_otel_span
)
return data
@ -1576,50 +1597,49 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if request_kwargs is None:
return healthy_deployments
input_file_id = cast(Optional[str], request_kwargs.get("input_file_id"))
model_file_id_mapping = cast(
Optional[Dict[str, Dict[str, str]]],
request_kwargs.get("model_file_id_mapping"),
request: Final = cast(Mapping[str, object], request_kwargs)
model_file_id_mapping: Final = cast(Mapping[str, Mapping[str, str]], request.get("model_file_id_mapping") or {})
requested_file_ids: Final = (request.get(key) for key in _PROVIDER_FILE_ID_KEYS)
model_ids_per_file: Final = tuple(
frozenset(model_file_id_mapping[file_id])
for file_id in requested_file_ids
if isinstance(file_id, str) and file_id in model_file_id_mapping
)
allowed_model_ids: Final[frozenset[str]] = (
model_ids_per_file[0].intersection(*model_ids_per_file[1:]) if model_ids_per_file else frozenset()
)
allowed_model_ids = []
if input_file_id and model_file_id_mapping:
model_id_dict = model_file_id_mapping.get(input_file_id, {})
allowed_model_ids = list(model_id_dict.keys())
if len(allowed_model_ids) == 0:
return healthy_deployments
return [
holders: Final[list[dict[str, object]]] = [
deployment
for deployment in healthy_deployments
if deployment.get("model_info", {}).get("id") in allowed_model_ids
]
return holders if holders or "input_file_id" in request else healthy_deployments
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
self, kwargs: dict[str, object], call_type: Optional[CallTypes]
) -> Optional[dict]:
"""
Allow modifying the request just before it's sent to the deployment.
"""
accessor_key: Optional[str] = None
if call_type and call_type == CallTypes.acreate_batch:
accessor_key = "input_file_id"
elif call_type and call_type == CallTypes.acreate_fine_tuning_job:
accessor_key = "training_file"
else:
accessor_keys: Final = _PROVIDER_FILE_ID_KWARGS.get(call_type) if call_type else None
if accessor_keys is None:
return kwargs
if accessor_key:
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
# model_info may be at top-level or nested under litellm_metadata
# (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], _keyed_get(kwargs.get("model_info", {}), "id"))
if model_id is None:
model_id = cast(
Optional[str],
_keyed_get(_keyed_get(kwargs.get("litellm_metadata", {}), "model_info", {}), "id"),
)
for accessor_key in accessor_keys:
input_file_id = cast(Optional[str], kwargs.get(accessor_key))
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
# model_info may be at top-level or nested under litellm_metadata
# (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
if model_id is None:
model_id = cast(
Optional[str],
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
)
mapped_file_id: Optional[str] = None
if input_file_id and model_file_id_mapping and model_id:
mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None)
@ -1813,7 +1833,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
)
async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict:
async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span | None) -> dict:
"""
Get model-specific file IDs for a list of proxy file IDs.
Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id
@ -1855,7 +1875,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
llm_router: Optional[Router],
_create_file_request: CreateFileRequest,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
) -> List[OpenAIFileObject]:
if llm_router is None:
raise Exception("LLM Router not initialized. Ensure models added to proxy.")
@ -1871,7 +1891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
create_file_request: CreateFileRequest,
llm_router: Router,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
) -> OpenAIFileObject:
responses = await self.create_file_for_each_model(
@ -1938,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_objects: List[OpenAIFileObject],
create_file_request: CreateFileRequest,
internal_usage_cache: InternalUsageCache,
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
target_model_names_list: List[str],
) -> OpenAIFileObject:
## GET THE FILE TYPE FROM THE CREATE FILE REQUEST
@ -2641,7 +2661,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Continue with other files even if one fails
continue
def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str:
def _get_content_type_from_file_object(self, file_object: object) -> str:
"""
Determine content type from file object.

6
enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
View file

@ -29,7 +29,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
Span = Union[_Span, Any]
Span = _Span
InternalUsageCache = _InternalUsageCache
PrismaClient = _PrismaClient
else:
@ -107,7 +107,7 @@ class _PROXY_LiteLLMManagedVectorStores(
llm_router: Router,
model: str,
request_data: dict[str, object] | VectorStoreCreateOptionalRequestParams,
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
) -> VectorStoreCreateResponse:
"""
Create a vector store for a specific model.
@ -134,7 +134,7 @@ class _PROXY_LiteLLMManagedVectorStores(
create_request: VectorStoreCreateOptionalRequestParams,
llm_router: Router,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
) -> VectorStoreCreateResponse:
"""

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

@ -36,6 +36,7 @@ from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import PrismaTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -73,6 +74,16 @@ def _verification_token_table(
return VerificationTokenRepository(prisma_client).table
class _WriterVerificationTokenRepository(PrismaTableRepository["prisma_models.LiteLLM_VerificationToken"]):
table_name = "litellm_verificationtoken"
def _writer_verification_token_table(
prisma_client: PrismaClient,
) -> TableActions["prisma_models.LiteLLM_VerificationToken"]:
return _WriterVerificationTokenRepository(prisma_client, use_writer=True).table
def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]:
return BudgetRepository(prisma_client).table
@ -849,7 +860,7 @@ async def update_project(
else None
)
if current_project is not None and data.team_id != current_project.team_id:
mismatched_key_count: Final = await prisma_client.writer_db.litellm_verificationtoken.count(
mismatched_key_count: Final = await _writer_verification_token_table(prisma_client).count(
where={
"project_id": data.project_id,
"OR": [{"team_id": {"not": data.team_id}}, {"team_id": None}],

8
enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py
View file

@ -14,6 +14,7 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, List, Optional, Protocol
from fastapi import APIRouter, Depends, HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -37,6 +38,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
router = APIRouter()
_DECODED_JSON: Final = TypeAdapter(object)
class ManagedVectorStoreRow(Protocol):
@ -348,11 +350,11 @@ async def update_vector_store(
raise HTTPException(status_code=500, detail="Database not connected")
try:
update_data = data.model_dump(exclude_unset=True)
vector_store_id: Final[str] = update_data.pop("vector_store_id")
update_data = data.model_dump(exclude_unset=True, exclude={"vector_store_id"})
vector_store_id: Final = data.vector_store_id
if update_data.get("vector_store_metadata") is not None:
update_data["vector_store_metadata"] = safe_dumps(
update_data["vector_store_metadata"]
_DECODED_JSON.validate_python(update_data["vector_store_metadata"])
)
updated = await managed_vector_store_table(prisma_client).update(

4
enterprise/pyproject.toml
View file

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

49
gateway/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.26.0
@ -20,6 +20,33 @@ RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# ---------- Rust extension ----------
# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
FROM $LITELLM_BUILD_IMAGE AS rust-native
WORKDIR /app
USER root
COPY --from=uvbin /uv /uvx /usr/local/bin/
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
ENV UV_PYTHON_DOWNLOADS=0
COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
python -c "import maturin; maturin.build_wheel('/wheels')"
# ---------- Builder ----------
FROM $LITELLM_BUILD_IMAGE AS builder
@ -32,7 +59,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -66,11 +93,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
# Stage 2 — copy source and install the workspace members.
COPY . .
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-default-groups --no-editable \
uv sync --frozen --no-install-project --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -79,6 +106,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3.13
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
# compiled Rust extension must live next to the source or it is never imported.
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
@ -95,7 +125,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic libevent && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -110,7 +140,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer

4
gateway/routes/allowlist.py
View file

@ -1,7 +1,7 @@
"""Path allowlist for the gateway component.
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image,
audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/decisions",
"/v1/systemone",
"/systemone",
"/v1/extraction",
"/extraction",
"/v1/ocr",
"/ocr",
"/v1/rag/",

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

@ -70,11 +70,18 @@ proxyConfigMap:
proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
model_list: []
```
The chart ships no models, so add them under `proxy_config.model_list` or from the Admin UI with `store_model_in_db: true`
```
proxy_config:
model_list:
- model_name: gpt-3.5-turbo
- model_name: gpt-5.6
litellm_params:
model: gpt-3.5-turbo
api_key: eXaMpLeOnLy
model: openai/gpt-5.6
api_key: os.environ/OPENAI_API_KEY
```
#### Example using existing `proxyConfigMap` instead of creating it:

26
helm/litellm-helm/tests/deployment_tests.yaml
View file

@ -270,6 +270,32 @@ tests:
- matchRegex:
path: data["config.yaml"]
pattern: '(?m)^\s*health_check_details:\s*false$'
- it: should not render any placeholder models when the user sets proxy_config without model_list
template: configmap-litellm.yaml
set:
proxy_config.general_settings.store_model_in_db: true
asserts:
- notMatchRegex:
path: data["config.yaml"]
pattern: 'model_name:'
- matchRegex:
path: data["config.yaml"]
pattern: 'store_model_in_db: true'
- matchRegex:
path: data["config.yaml"]
pattern: 'master_key: os.environ/PROXY_MASTER_KEY'
- it: should render a user-supplied model_list as the only models
template: configmap-litellm.yaml
set:
proxy_config.model_list:
- model_name: my-model
litellm_params:
model: openai/my-model
api_key: os.environ/OPENAI_API_KEY
asserts:
- matchRegex:
path: data["config.yaml"]
pattern: 'model_list:\n- litellm_params:\n api_key: os.environ/OPENAI_API_KEY\n model: openai/my-model\n model_name: my-model\n$'
- it: should allow overriding liveness, readiness, and startup probes
template: deployment.yaml
set:

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

@ -180,17 +180,7 @@ proxyConfigMap:
# Examples: https://github.com/BerriAI/litellm/tree/main/litellm/proxy/example_config_yaml
# Reference: https://docs.litellm.ai/docs/proxy/configs
proxy_config:
model_list:
# At least one model must exist for the proxy to start.
- model_name: gpt-3.5-turbo
litellm_params:
model: gpt-3.5-turbo
api_key: eXaMpLeOnLy
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
model_list: []
general_settings:
master_key: os.environ/PROXY_MASTER_KEY

23
litellm-proxy-extras/litellm_proxy_extras/migrations/20261015120000_add_environment_variables_revision_trigger/migration.sql Normal file
View file

@ -0,0 +1,23 @@
SET LOCAL lock_timeout = '10s';
CREATE SEQUENCE IF NOT EXISTS "litellm_environment_variables_revision_seq";
GRANT USAGE ON SEQUENCE "litellm_environment_variables_revision_seq" TO PUBLIC;
CREATE OR REPLACE FUNCTION litellm_bump_environment_variables_revision() RETURNS TRIGGER AS $$
DECLARE
next_revision BIGINT := nextval(format('%I.litellm_environment_variables_revision_seq', TG_TABLE_SCHEMA)::regclass);
BEGIN
IF TG_OP = 'UPDATE' THEN
NEW."reload_revision" := GREATEST(next_revision, OLD."reload_revision" + 1);
ELSE
NEW."reload_revision" := next_revision;
END IF;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
DROP TRIGGER IF EXISTS "litellm_environment_variables_revision" ON "LiteLLM_Config";
CREATE TRIGGER "litellm_environment_variables_revision"
BEFORE INSERT OR UPDATE ON "LiteLLM_Config"
FOR EACH ROW WHEN (NEW."param_name" = 'environment_variables')
EXECUTE PROCEDURE litellm_bump_environment_variables_revision();

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

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

9
litellm-rust/Cargo.lock generated
View file

@ -4083,6 +4083,7 @@ dependencies = [
name = "litellm-inference-responses"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-util",
"litellm-auth",
@ -8201,18 +8202,18 @@ dependencies = [
[[package]]
name = "webpki-root-certs"
version = "1.0.9"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b"
checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "webpki-roots"
version = "1.0.9"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a"
checksum = "f5f280bbb2b287bc8e7de7e3852945a0ff50d1bbcb5ed42b6c7a568c546df171"
dependencies = [
"rustls-pki-types",
]

2
litellm-rust/Cargo.toml
View file

@ -124,7 +124,7 @@ moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
percent-encoding = "2.3"
webpki-roots = "1"
webpki-roots = "1.0.10"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"

1
litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs
View file

@ -94,6 +94,7 @@ const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[
"docker_model_runner",
"ragflow",
"pinstripes",
"openinfer",
"darkbloom",
"meta",
"cognition",

8
litellm-rust/crates/gateway-inference/src/caching.rs
View file

@ -72,26 +72,26 @@ fn duration(seconds: f64) -> Result<Duration, Error> {
#[derive(Clone, Default)]
pub(crate) struct CacheHeaders(std::sync::Arc<std::sync::OnceLock<String>>);
impl litellm_host::interceptors::Interceptors<litellm_inference::RouteError> for CacheHeaders {
impl litellm_host::interceptors::Interceptors<litellm_inference::Error> for CacheHeaders {
async fn before_provider_request(
&self,
wire: litellm_host::interceptors::WireRequest,
_: litellm_host::interceptors::RequestContext,
) -> Result<litellm_host::interceptors::WireRequest, litellm_inference::RouteError> {
) -> Result<litellm_host::interceptors::WireRequest, litellm_inference::Error> {
Ok(wire)
}
async fn after_provider_response(
&self,
_: litellm_host::interceptors::RawResponse,
) -> Result<(), litellm_inference::RouteError> {
) -> Result<(), litellm_inference::Error> {
Ok(())
}
async fn result_ready(
&self,
facts: litellm_host::interceptors::ExecutionFacts,
) -> Result<(), litellm_inference::RouteError> {
) -> Result<(), litellm_inference::Error> {
if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source {
let _ = self.0.set(key);
}

26
litellm-rust/crates/gateway-inference/src/error.rs
View file

@ -4,7 +4,7 @@ use axum::{
response::{IntoResponse, Response},
};
use litellm_http::transport::Error as TransportError;
use litellm_inference::RouteError;
use litellm_inference::Error as InferenceError;
use litellm_llms::base_llm::ocr::error::Error as OcrError;
use serde_json::{Map, Value, json};
@ -19,7 +19,7 @@ pub enum Error {
)]
UnknownModel(String),
#[error(transparent)]
Route(#[from] RouteError),
Route(#[from] InferenceError),
#[error(transparent)]
Ocr(#[from] OcrError),
#[error("{0} is not implemented by the Rust gateway")]
@ -36,8 +36,8 @@ impl IntoResponse for Error {
}
}
impl From<litellm_host_http::Error<RouteError>> for Error {
fn from(error: litellm_host_http::Error<RouteError>) -> Self {
impl From<litellm_host_http::Error<InferenceError>> for Error {
fn from(error: litellm_host_http::Error<InferenceError>) -> Self {
match error {
litellm_host_http::Error::Call(error) => Self::Route(error),
litellm_host_http::Error::Protocol => Self::Internal(error.to_string()),
@ -50,7 +50,7 @@ impl Error {
match self {
Self::Auth(error) => error.status(),
Self::Unsupported(_)
| Self::Route(RouteError::Unsupported(_))
| Self::Route(InferenceError::Unsupported(_))
| Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED,
Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE,
Self::Ocr(
@ -64,10 +64,10 @@ impl Error {
.and_then(|status| StatusCode::from_u16(status).ok())
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST,
Self::Route(RouteError::Transport(TransportError::Http { status, .. })) => {
Self::Route(InferenceError::Transport(TransportError::Http { status, .. })) => {
StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY)
}
Self::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { .. })) => {
Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => {
StatusCode::UNAUTHORIZED
}
Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST,
@ -97,7 +97,9 @@ impl Error {
/// body already in that shape passes through, any other has its message extracted.
pub fn body(&self, request_id: Option<&str>) -> Value {
let raw = match self {
Self::Route(RouteError::Transport(TransportError::Http { body, .. })) => body.clone(),
Self::Route(InferenceError::Transport(TransportError::Http { body, .. })) => {
body.clone()
}
other => other.to_string(),
};
let parsed = serde_json::from_str::<Value>(&raw).ok();
@ -175,7 +177,7 @@ mod tests {
use super::*;
fn upstream(status: u16, body: &str) -> Error {
Error::Route(RouteError::Transport(TransportError::Http {
Error::Route(InferenceError::Transport(TransportError::Http {
status,
body: body.into(),
}))
@ -225,16 +227,16 @@ mod tests {
#[rstest]
#[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)]
#[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)]
#[case::rejected_request(Error::Route(InferenceError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)]
#[case::missing_key(
Error::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey {
Error::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey {
provider: "Anthropic",
environment_variable: "ANTHROPIC_API_KEY",
})),
StatusCode::UNAUTHORIZED,
)]
#[case::lost_connection(
Error::Route(RouteError::Transport(TransportError::Network("reset".into()))),
Error::Route(InferenceError::Transport(TransportError::Network("reset".into()))),
StatusCode::INTERNAL_SERVER_ERROR,
)]
fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) {

2
litellm-rust/crates/gateway-inference/src/messages.rs
View file

@ -24,7 +24,7 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [
"x-openrouter-experimental-metadata",
];
const ANTHROPIC_API_HEADER_PROVIDERS: &str =
"anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
"anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek";
pub async fn create(
State(gateway): State<Arc<Gateway>>,

2
litellm-rust/crates/gateway-inference/tests/messages.rs
View file

@ -16,7 +16,7 @@ use wiremock::{
#[rstest]
#[case::anthropic("anthropic/test-model", "/v1/messages", true)]
#[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)]
#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", false)]
#[case::deepseek("deepseek/test-model", "/anthropic/v1/messages", true)]
#[tokio::test]
async fn messages_reaches_the_provider_and_preserves_json_or_sse(
#[case] model: &str,

57
litellm-rust/crates/inference-chat/src/common_utils.rs
View file

@ -6,9 +6,11 @@ use litellm_llms::{
bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
openai_like::chat::transformation::OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG,
};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use serde_json::{Map, Value};
use super::Error;
use crate::types::{ChatCompletionsCall, ResolvedChatCompletionsRequest};
const HEADER_CONTEXT: &str = "chat completions";
@ -42,3 +44,58 @@ pub(super) fn string_headers(
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
}
pub(super) fn resolve_request(
call: ChatCompletionsCall,
config: &'static dyn BaseConfig,
) -> Result<ResolvedChatCompletionsRequest, Error> {
let messages = serde_json::from_value::<Vec<ChatMessage>>(call.messages).map_err(|err| {
Error::InvalidRequest(litellm_llms::ErrorDetail::invalid(
"chat completions messages",
err,
))
})?;
if messages.is_empty() {
return Err(Error::InvalidRequest(
"chat completions requires at least one message".into(),
));
}
if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) {
return Err(Error::Unsupported(reason.0));
}
Ok(ResolvedChatCompletionsRequest {
messages,
optional_params: call.optional_params,
api_key: call.api_key,
api_base: call.api_base,
extra_headers: call.extra_headers,
timeout: call.timeout,
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::chat_completions_provider;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
#[case::anthropic(LlmProviders::Anthropic, ("anthropic-version", "2023-06-01"))]
#[case::bedrock(LlmProviders::Bedrock, ("Content-Type", "application/json"))]
fn resolves_supported_provider_configs(
#[case] provider: LlmProviders,
#[case] expected_header: (&str, &str),
) {
let config = chat_completions_provider(provider)
.expect("provider has a chat completions config")
.config();
assert!(config.default_headers().contains(&expected_header));
}
#[rstest]
#[case::openai(LlmProviders::Openai)]
fn rejects_providers_without_a_chat_config(#[case] provider: LlmProviders) {
assert!(chat_completions_provider(provider).is_none());
}
}

116
litellm-rust/crates/inference-chat/src/handler.rs
View file

@ -3,44 +3,66 @@ use std::time::Duration;
use litellm_auth::AuthServices;
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body};
use litellm_http::{
Client,
outbound::OutboundRequest,
request::{truncate_error_body, with_default_headers},
};
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
chat::transformation::BaseConfig,
chat::transformation::ProviderChatResponseData,
};
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::Error;
use crate::{constants::CHAT_COMPLETIONS_TIMEOUT_SECS, types::ProviderChatCompletionsRequest};
use crate::{
common_utils::string_headers, constants::CHAT_COMPLETIONS_TIMEOUT_SECS,
types::ResolvedChatCompletionsRequest,
};
#[expect(clippy::too_many_arguments)] // mirrors BaseLLMHTTPHandler.completion, which takes each stage input separately
pub(super) async fn execute(
http: &Client,
auth: &AuthServices,
request: ProviderChatCompletionsRequest,
config: &'static dyn BaseConfig,
provider: litellm_inference::provider::ResolvedProvider<'_>,
call: ResolvedChatCompletionsRequest,
secrets: Secrets,
cache: Option<litellm_cache_response::ScopedCache>,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
let ProviderChatCompletionsRequest {
model,
custom_llm_provider,
config,
url,
body,
optional_params,
environment,
secrets,
timeout,
api_key,
} = request;
let validated = config.validate_environment(
string_headers(call.extra_headers)?,
call.api_key.as_deref(),
provider.model,
&call.optional_params,
&|key| secrets.get(key),
)?;
let url = config.get_complete_url(
call.api_base.as_deref(),
provider.model,
&call.optional_params,
&|key| secrets.get(key),
)?;
let body = config
.transform_request(provider.model, call.messages, call.optional_params.clone())?
.body;
let environment = ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
};
let provider_name = <&'static str>::from(provider.provider);
let context = RequestContext {
model: model.clone(),
custom_llm_provider,
optional_params: Value::Object(optional_params),
model: provider.model.to_owned(),
custom_llm_provider: provider_name.to_owned(),
optional_params: Value::Object(call.optional_params),
secret_fields: Vec::new(),
api_key,
api_key: call.api_key.map(litellm_auth::SecretValue::new),
};
let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?;
let identity = litellm_host::interceptors::ProviderIdentity {
@ -76,7 +98,7 @@ pub(super) async fn execute(
},
wire.url,
&wire.body,
timeout,
call.timeout,
)?;
let response = litellm_inference::outbound::send(outbound, http)
@ -121,7 +143,7 @@ pub(super) async fn execute(
))
})?;
config
.transform_response(&model, ProviderChatResponseData { body })
.transform_response(provider.model, ProviderChatResponseData { body })
.map_err(Error::from)
.map_err(as_response_error)
},
@ -134,7 +156,7 @@ pub(super) async fn execute(
///
/// A config reports the same variants on either side of the call: a missing
/// field or an unsupported block can mean "this request cannot be translated"
/// during prepare and "this response cannot be normalized" here. Only the
/// before the request is sent and "this response cannot be normalized" here. Only the
/// second kind has already been billed, and a host that keeps a reference
/// implementation must not retry those, so collapse them to one variant that
/// can only mean the provider was already called.
@ -178,9 +200,10 @@ mod tests {
use super::*;
use crate::{
prepare::{prepare_provider_request, resolve_request},
types::ChatCompletionsRequest,
common_utils::{chat_completions_provider, resolve_request},
types::{ChatCompletionsCall, ResolvedChatCompletionsRequest},
};
use litellm_inference::provider::resolve_llm_provider;
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
@ -217,22 +240,33 @@ mod tests {
}
}
fn prepared(api_base: &str) -> ProviderChatCompletionsRequest {
prepare_provider_request(
resolve_request(ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
fn prepared(
api_base: &str,
) -> (
litellm_inference::provider::ResolvedProvider<'static>,
&'static dyn BaseConfig,
ResolvedChatCompletionsRequest,
) {
let provider =
resolve_llm_provider("anthropic/claude-sonnet-4-5", None, "chat completions").unwrap();
let config = chat_completions_provider(provider.provider)
.expect("provider has a chat completions config")
.config();
let request = resolve_request(
ChatCompletionsCall {
model: "anthropic/claude-sonnet-4-5".into(),
messages: json!([{"role": "user", "content": "hi"}]),
optional_params: json!({"max_tokens": 16}).as_object().unwrap().clone(),
api_key: Some("sk-test"),
api_base: Some(api_base),
api_key: Some("sk-test".into()),
api_base: Some(api_base.into()),
custom_llm_provider: None,
extra_headers: None,
timeout: None,
})
.unwrap(),
std::sync::Arc::new(|_: &str| None),
},
config,
)
.unwrap()
.unwrap();
(provider, config, request)
}
#[rstest]
@ -247,10 +281,14 @@ mod tests {
.await;
let interceptors = RecordingHooks::default();
let (provider, config, request) = prepared(&upstream.uri());
execute(
&Client::plain_for_test(),
&AuthServices::default(),
prepared(&upstream.uri()),
config,
provider,
request,
std::sync::Arc::new(|_: &str| None),
None,
None,
&interceptors,
@ -288,10 +326,14 @@ mod tests {
.await;
let interceptors = RecordingHooks::default();
let (provider, config, request) = prepared(&upstream.uri());
let error = execute(
&Client::plain_for_test(),
&AuthServices::default(),
prepared(&upstream.uri()),
config,
provider,
request,
std::sync::Arc::new(|_: &str| None),
None,
None,
&interceptors,

55
litellm-rust/crates/inference-chat/src/lib.rs
View file

@ -1,18 +1,17 @@
use litellm_host::observation::ObservationSender;
pub mod route;
pub mod types;
pub use litellm_inference::RouteError as Error;
pub use litellm_inference::Error;
mod common_utils;
pub mod constants;
pub(crate) mod handler;
mod prepare;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use prepare::{prepare_provider_request, resolve_request};
use litellm_auth::AuthServices;
use litellm_inference::provider::resolve_llm_provider;
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
use types::ChatCompletionsRequest;
use types::{ChatCompletionsCall, ChatCompletionsRequest};
#[derive(Clone)]
pub struct ChatCompletionsRoute {
@ -55,7 +54,7 @@ impl ChatCompletionsRoute {
} = options.into();
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run_call(
self.run(
request.into(),
cache_options,
interceptors,
@ -65,30 +64,52 @@ impl ChatCompletionsRoute {
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "chat_completions",
model = %call.model,
provider,
resolved_model,
stream = false,
outcome
))]
async fn run(
&self,
request: ChatCompletionsRequest<'_>,
call: ChatCompletionsCall,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
let resolved = resolve_request(request)?;
let snapshot = self
.secrets
.resolve(&resolved.config.secret_names())
.await?;
let prepared = prepare_provider_request(resolved, snapshot)?;
litellm_inference::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
let execute: futures_util::future::BoxFuture<'_, Result<ChatCompletionsResponse, Error>> =
Box::pin(handler::execute(
litellm_inference::diagnostic::unary(async {
let model = call.model.clone();
let custom_llm_provider = call.custom_llm_provider.clone();
let provider =
resolve_llm_provider(&model, custom_llm_provider.as_deref(), "chat completions")?;
let config = common_utils::chat_completions_provider(provider.provider)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?
.config();
let call = common_utils::resolve_request(call, config)?;
litellm_inference::diagnostic::provider(
provider.model,
<&'static str>::from(provider.provider),
);
let secrets = self.secrets.resolve(&config.secret_names()).await?;
let execute: futures_util::future::BoxFuture<
'_,
Result<ChatCompletionsResponse, Error>,
> = Box::pin(handler::execute(
&self.http,
&self.auth,
prepared,
config,
provider,
call,
secrets,
self.cache.clone(),
cache_options,
interceptors,
observers,
));
execute.await
execute.await
})
.await
}
}

638
litellm-rust/crates/inference-chat/src/prepare.rs
View file

@ -1,638 +0,0 @@
use litellm_auth::SecretValue;
use litellm_core_utils::settings::Lookup;
use litellm_http::request::with_default_headers;
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::{
Error,
common_utils::{chat_completions_provider, string_headers},
};
use crate::types::{
ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest,
};
use litellm_inference::provider::resolve_llm_provider;
pub(super) struct ResolvedProvider {
pub(super) model: String,
pub(super) custom_llm_provider: String,
pub(super) config: &'static dyn BaseConfig,
}
pub(super) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<ResolvedProvider, Error> {
let provider_info = resolve_llm_provider(model, custom_llm_provider, "chat completions")?;
let config = chat_completions_provider(provider_info.provider)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?
.config();
Ok(ResolvedProvider {
model: provider_info.model.to_string(),
custom_llm_provider: <&str>::from(provider_info.provider).to_string(),
config,
})
}
pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error> {
serde_json::from_value(messages).map_err(|err| {
Error::InvalidRequest(litellm_llms::ErrorDetail::invalid(
"chat completions messages",
err,
))
})
}
pub(super) fn resolve_request(
request: ChatCompletionsRequest<'_>,
) -> Result<ResolvedChatCompletionsRequest<'_>, Error> {
let ResolvedProvider {
model,
custom_llm_provider,
config,
} = resolve_provider_config(request.model, request.custom_llm_provider)?;
let messages = parse_messages(request.messages)?;
if messages.is_empty() {
return Err(Error::InvalidRequest(
"chat completions requires at least one message".into(),
));
}
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
return Err(Error::Unsupported(reason.0));
}
Ok(ResolvedChatCompletionsRequest {
model,
custom_llm_provider,
config,
messages,
optional_params: request.optional_params,
api_key: request.api_key,
api_base: request.api_base,
extra_headers: request.extra_headers,
timeout: request.timeout,
})
}
fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn BaseConfig,
secrets: &dyn Lookup,
) -> Result<ValidatedEnvironment, Error> {
let env_lookup = |key: &str| secrets.get(key);
let forwarded = string_headers(request.extra_headers.clone())?;
let validated = config.validate_environment(
forwarded,
request.api_key,
model,
&request.optional_params,
&env_lookup,
)?;
Ok(ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
})
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub(super) fn prepare_provider_request(
request: ResolvedChatCompletionsRequest<'_>,
secrets: Secrets,
) -> Result<ProviderChatCompletionsRequest, Error> {
let environment =
validate_environment(&request, &request.model, request.config, secrets.as_ref())?;
let model = request.model;
let config = request.config;
let env_lookup = |key: &str| secrets.get(key);
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let transformed =
config.transform_request(&model, request.messages, request.optional_params.clone())?;
Ok(ProviderChatCompletionsRequest {
model,
custom_llm_provider: request.custom_llm_provider,
config,
url,
body: transformed.body,
optional_params: request.optional_params,
environment,
secrets,
timeout: request.timeout,
api_key: request.api_key.map(|key| SecretValue::new(key.to_string())),
})
}
#[cfg(test)]
mod tests {
use litellm_auth::CredentialPlacement;
use litellm_llms::base_llm::auth::{AuthScheme, resolve_auth};
use serde_json::{Map, Value, json};
use super::{prepare_provider_request, resolve_request};
use crate::{
Error,
types::{ChatCompletionsRequest, ProviderChatCompletionsRequest},
};
fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
) -> Result<ProviderChatCompletionsRequest, Error> {
prepare_provider_request(
resolve_request(request)?,
std::sync::Arc::new(|_: &str| None),
)
}
/// The headers as they go on the wire, credential applied.
fn wire_headers(prepared: &ProviderChatCompletionsRequest) -> Vec<(String, String)> {
tokio::runtime::Builder::new_current_thread()
.build()
.unwrap()
.block_on(resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment.clone(),
&|_| None,
))
.unwrap()
.headers
}
fn request<'a>(
model: &'a str,
provider: Option<&'a str>,
messages: Value,
optional_params: Value,
) -> ChatCompletionsRequest<'a> {
ChatCompletionsRequest {
model,
messages,
optional_params: match optional_params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: provider,
extra_headers: None,
timeout: None,
}
}
fn preparation_error(request: ChatCompletionsRequest<'_>) -> Error {
prepare_chat_completions_call(request)
.err()
.expect("request preparation should fail")
}
#[test]
fn resolves_the_provider_from_the_model_prefix() {
let prepared = prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
}
#[test]
fn strips_an_explicit_provider_prefix_from_the_model() {
let prepared = prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
}
#[test]
fn adds_the_auth_and_default_headers() {
let prepared = prepare_chat_completions_call(request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert!(
wire_headers(&prepared).contains(&("x-api-key".to_string(), "sk-test".to_string()))
);
assert!(
wire_headers(&prepared)
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
);
assert!(matches!(
prepared.environment.auth,
AuthScheme::Credential {
placement: CredentialPlacement::Header("x-api-key"),
..
}
));
}
#[test]
fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
// Python builds `{**headers, **anthropic_headers}`, so the deployment's key
// overwrites a forwarded one. Honouring the caller's would let whoever sends
// the request choose the Anthropic principal it bills to.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([(
"X-Api-Key".to_string(),
json!("sk-caller"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let keys: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", headers);
assert_eq!(keys[0].1, "sk-test");
}
#[test]
fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
// Anthropic's `validate_environment` pops `x-api-key` and sets `authorization`
// for an OAuth token, so re-adding the key here would put the credential into
// a header the host removed on purpose.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([
(
"Authorization".to_string(),
json!("Bearer sk-ant-oat01-token"),
),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
assert!(
!wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
"the resolved key must not be applied over an OAuth bearer, got {:?}",
wire_headers(&prepared)
);
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-ant-oat01-token")
);
}
#[test]
fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
// Only an OAuth bearer replaces the credential. Python sends the deployment's
// `x-api-key` alongside any other forwarded `authorization`, so deferring on
// the mere presence of that header would drop the deployment's auth.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer unrelated")),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let keys: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", headers);
assert_eq!(keys[0].1, "sk-test");
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer unrelated"),
"the unrelated authorization must survive, got {:?}",
wire_headers(&prepared)
);
}
#[rstest::rstest]
fn rejects_empty_messages_before_resolving_credentials() {
let call = ChatCompletionsRequest {
api_key: None,
..request("claude-sonnet-4-5", Some("anthropic"), json!([]), json!({}))
};
assert!(matches!(preparation_error(call), Error::InvalidRequest(_)));
}
#[test]
fn rejects_an_unknown_provider() {
assert_eq!(
preparation_error(request(
"openai/gpt-4o",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)),
Error::InvalidProvider("openai".to_string())
);
}
#[test]
fn rejects_a_model_with_no_resolvable_provider() {
assert!(matches!(
preparation_error(request(
"claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)),
Error::InvalidProvider(_)
));
}
#[test]
fn rejects_an_empty_or_malformed_message_list() {
assert_eq!(
preparation_error(request(
"anthropic/claude-sonnet-4-5",
None,
json!([]),
json!({}),
)),
Error::InvalidRequest(
"chat completions requires at least one message"
.to_string()
.into()
)
);
assert!(matches!(
preparation_error(request(
"anthropic/claude-sonnet-4-5",
None,
json!("not a list"),
json!({}),
)),
Error::InvalidRequest(_)
));
}
#[test]
fn rejects_non_string_extra_headers() {
let mut call = request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
preparation_error(call),
Error::Headers(litellm_http::request::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
})
);
}
#[test]
fn prepares_a_bedrock_call_without_resolving_credentials() {
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.api_key = None;
let prepared = prepare_chat_completions_call(call).expect("prepares");
assert_eq!(
prepared.url,
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
assert!(matches!(
&prepared.environment.auth,
AuthScheme::AwsSigV4 { region, service: "bedrock", .. } if region == "us-east-1"
));
// SigV4 signs the serialized body, so prepare must not have added an
// Authorization header; the signer does it over the bytes sent.
assert!(
!prepared
.environment
.headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
);
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
}
#[tokio::test]
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
// Python signs only the AWS header set and reattaches the rest, so a header
// the caller forwarded rides along without joining the canonical request.
// Signing it makes Converse 403 on a deployment that works on Python.
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
);
// A key would resolve to a bearer token and never reach the signer.
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(
"x-request-id".to_string(),
json!("abc-123"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let authenticated = resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment,
&|_| None,
)
.await
.expect("resolves");
let signed = crate::handler::outbound_request(
authenticated,
prepared.url,
&prepared.body,
prepared.timeout,
)
.expect("signs");
let authorization = signed
.header("authorization")
.expect("carries an authorization header")
.to_string();
assert!(
authorization.starts_with("AWS4-HMAC-SHA256"),
"expected a SigV4 signature, got {authorization}"
);
assert!(
!authorization.contains("x-request-id"),
"forwarded header reached SignedHeaders: {authorization}"
);
// It still goes on the wire, it is just not part of the signature.
assert!(
signed
.headers()
.iter()
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
"forwarded header was dropped instead of reattached"
);
}
#[rstest::rstest]
#[case::authorization("Authorization")]
#[case::amz_date("x-amz-date")]
#[case::security_token("x-amz-security-token")]
#[case::date("Date")]
#[tokio::test]
async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) {
let call = ChatCompletionsRequest {
api_key: None,
extra_headers: Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
)
};
let prepared = prepare_chat_completions_call(call).expect("prepares");
let authenticated = resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment,
&|_| None,
)
.await
.expect("resolves");
let error = crate::handler::outbound_request(
authenticated,
prepared.url,
&prepared.body,
prepared.timeout,
)
.expect_err("conflicting signing headers must fail");
assert!(
matches!(error, Error::Unsupported(_)),
"{forwarded} returned {error:?}"
);
}
#[test]
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
// once a bearer token resolves, so the deployment's identity wins on
// Python. Keeping the caller's would authorize and bill the call as a
// different principal, and only when the deployment carries `rust: true`.
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.extra_headers = Some(Map::from_iter([(
"Authorization".to_string(),
json!("Bearer caller-supplied"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let authorizations: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
.collect();
assert_eq!(
authorizations,
vec!["Bearer sk-test"],
"the deployment token must be the only authorization on the wire"
);
}
#[test]
fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
// The opposite precedence, and deliberate: Anthropic's own transform
// honours a forwarded OAuth bearer, so the Bedrock fix above must not be
// generalized into a rule that the configured key always wins.
//
// An OAuth bearer is the whole of that exception. This forwarded a plain
// `x-api-key` until round 17, which read as the same claim and was not:
// Python overwrites a forwarded `x-api-key` with the deployment's.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([(
"authorization".to_string(),
json!("Bearer sk-ant-oat01-forwarded"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let keys: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.map(|(_, value)| value.as_str())
.collect();
assert!(keys.is_empty(), "got {:?}", headers);
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-ant-oat01-forwarded")
);
}
#[test]
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
// The configured bearer identity has its own account and quota boundary,
// so a request carrying one must not be signed as whatever principal the
// host's AWS credentials resolve to.
let prepared = prepare_chat_completions_call(request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
))
.expect("prepares");
assert!(matches!(
&prepared.environment.auth,
AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret }
if secret.expose() == "sk-test"
));
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-test"),
"prepare did not carry the bearer token"
);
}
}

40
litellm-rust/crates/inference-chat/src/route.rs
View file

@ -1,4 +1,3 @@
use litellm_host::observation::ObservationSender;
use std::convert::Infallible;
use litellm_host::{
@ -7,10 +6,7 @@ use litellm_host::{
};
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use super::{
ChatCompletionsRoute, Error,
types::{ChatCompletionsCall, ChatCompletionsRequest},
};
use super::{ChatCompletionsRoute, Error, types::ChatCompletionsCall};
pub struct ChatCompletions;
@ -37,44 +33,12 @@ impl ChatCompletionsRoute {
call,
observers,
move |call, _, interceptors, observers| async move {
self.run_call(call, cache_options, &interceptors, observers.as_ref())
self.run(call, cache_options, &interceptors, observers.as_ref())
.await
.map(CallOutput::Complete)
},
)
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "chat_completions",
model = %call.model,
provider,
resolved_model,
stream = false,
outcome
))]
pub(super) async fn run_call(
&self,
call: ChatCompletionsCall,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
litellm_inference::diagnostic::unary(async {
let request = ChatCompletionsRequest {
model: &call.model,
messages: call.messages,
optional_params: call.optional_params,
api_key: call.api_key.as_deref(),
api_base: call.api_base.as_deref(),
custom_llm_provider: call.custom_llm_provider.as_deref(),
extra_headers: call.extra_headers,
timeout: call.timeout,
};
self.run(request, cache_options, interceptors, observers)
.await
})
.await
}
}
impl litellm_inference::caching::Cachable for ChatCompletions {

29
litellm-rust/crates/inference-chat/src/types.rs
View file

@ -1,8 +1,5 @@
use litellm_secrets::source::Secrets;
use std::time::Duration;
use litellm_auth::SecretValue;
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use serde_json::{Map, Value};
@ -49,31 +46,11 @@ impl From<ChatCompletionsRequest<'_>> for ChatCompletionsCall {
}
}
pub struct ResolvedChatCompletionsRequest<'a> {
pub model: String,
pub custom_llm_provider: String,
pub config: &'static dyn BaseConfig,
pub struct ResolvedChatCompletionsRequest {
pub messages: Vec<ChatMessage>,
pub optional_params: Map<String, Value>,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
}
pub struct ProviderChatCompletionsRequest {
pub model: String,
pub custom_llm_provider: String,
pub config: &'static dyn BaseConfig,
pub url: String,
pub body: Value,
/// The route's parameters before the provider transformation, reported to the host
/// beside the wire request.
pub optional_params: Map<String, Value>,
/// The forwarded and default headers plus how the call authenticates; the credential
/// itself is applied when the request is sent.
pub environment: ValidatedEnvironment,
pub secrets: Secrets,
pub timeout: Option<Duration>,
pub api_key: Option<SecretValue>,
}

10
litellm-rust/crates/inference-chat/tests/caching.rs
View file

@ -21,7 +21,7 @@ use litellm_host::{
lifecycle::{CallEvent, ExecutionEvent},
observation::observation_channel,
};
use litellm_inference::RouteError;
use litellm_inference::Error;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
@ -166,12 +166,12 @@ struct ChangingHooks {
facts: std::sync::Mutex<Vec<ExecutionFacts>>,
}
impl Interceptors<RouteError> for ChangingHooks {
impl Interceptors<Error> for ChangingHooks {
async fn before_provider_request(
&self,
mut wire: WireRequest,
_: RequestContext,
) -> Result<WireRequest, RouteError> {
) -> Result<WireRequest, Error> {
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if self.rewrite {
wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 });
@ -179,11 +179,11 @@ impl Interceptors<RouteError> for ChangingHooks {
Ok(wire)
}
async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> {
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Ok(())
}
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> {
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
self.facts.lock().unwrap().push(facts);
Ok(())
}

439
litellm-rust/crates/inference-chat/tests/request.rs Normal file
View file

@ -0,0 +1,439 @@
use std::time::Duration;
use litellm_inference_chat::{Error, types::ChatCompletionsRequest};
use rstest::rstest;
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
mod support;
use support::*;
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
fn request<'a>(
model: &'a str,
provider: Option<&'a str>,
messages: Value,
optional_params: Value,
) -> ChatCompletionsRequest<'a> {
let Value::Object(optional_params) = optional_params else {
panic!("params must be an object");
};
ChatCompletionsRequest {
model,
messages,
optional_params,
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: provider,
extra_headers: None,
timeout: Some(Duration::from_secs(10)),
}
}
async fn complete(
request: ChatCompletionsRequest<'_>,
) -> Result<litellm_llms_types::formats::chat_completions::ChatCompletionsResponse, Error> {
chat_completions_route().execute(request, &(), None).await
}
fn object(value: Value) -> Map<String, Value> {
let Value::Object(map) = value else {
panic!("expected an object");
};
map
}
fn anthropic_response() -> ResponseTemplate {
ResponseTemplate::new(200).set_body_raw(ANTHROPIC_MESSAGE, "application/json")
}
fn bedrock_response() -> ResponseTemplate {
json_response(json!({
"output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
}))
}
#[rstest]
#[tokio::test]
async fn resolves_the_provider_from_the_model_prefix() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/v1/messages");
assert_eq!(sent.json()["model"], "claude-sonnet-4-5");
}
#[rstest]
#[tokio::test]
async fn strips_an_explicit_provider_prefix_from_the_model() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"anthropic/claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.json()["model"],
"claude-sonnet-4-5"
);
}
#[rstest]
#[tokio::test]
async fn adds_the_auth_and_default_headers() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
assert_eq!(sent.header("anthropic-version"), Some("2023-06-01"));
}
#[rstest]
#[tokio::test]
async fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(json!({"X-Api-Key": "sk-caller"}))),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.header_values("x-api-key"),
["sk-test"]
);
}
#[rstest]
#[tokio::test]
async fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(json!({
"Authorization": "Bearer sk-ant-oat01-token",
"X-Api-Key": "sk-caller"
}))),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-caller"]);
assert_eq!(
sent.header_values("authorization"),
["Bearer sk-ant-oat01-token"]
);
}
#[rstest]
#[tokio::test]
async fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(json!({
"Authorization": "Bearer unrelated",
"X-Api-Key": "sk-caller"
}))),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
assert_eq!(sent.header_values("authorization"), ["Bearer unrelated"]);
}
#[rstest]
#[tokio::test]
async fn rejects_empty_messages_before_resolving_credentials() {
let error = complete(ChatCompletionsRequest {
api_key: None,
..request("claude-sonnet-4-5", Some("anthropic"), json!([]), json!({}))
})
.await
.expect_err("empty messages are rejected");
assert!(matches!(error, Error::InvalidRequest(_)));
}
#[rstest]
#[tokio::test]
async fn rejects_an_unknown_provider() {
assert_eq!(
complete(request(
"openai/gpt-4o",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.await
.unwrap_err(),
Error::InvalidProvider("openai".into())
);
}
#[rstest]
#[tokio::test]
async fn rejects_a_model_with_no_resolvable_provider() {
assert!(matches!(
complete(request(
"claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.await,
Err(Error::InvalidProvider(_))
));
}
#[rstest]
#[case::empty(json!([]), Some("chat completions requires at least one message"))]
#[case::malformed(json!("not a list"), None)]
#[tokio::test]
async fn rejects_an_empty_or_malformed_message_list(
#[case] messages: Value,
#[case] expected: Option<&str>,
) {
let error = complete(request(
"anthropic/claude-sonnet-4-5",
None,
messages,
json!({}),
))
.await
.unwrap_err();
if let Some(expected) = expected {
assert_eq!(error, Error::InvalidRequest(expected.to_string().into()));
} else {
assert!(matches!(error, Error::InvalidRequest(_)));
}
}
#[rstest]
#[tokio::test]
async fn rejects_non_string_extra_headers() {
assert_eq!(
complete(ChatCompletionsRequest {
extra_headers: Some(Map::from_iter([("x-trace".to_string(), json!(7))])),
..request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.unwrap_err(),
Error::Headers(litellm_http::request::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
})
);
}
#[rstest]
#[tokio::test]
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
let upstream = upstream([bedrock_response()]).await;
let base = upstream.uri();
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
);
call.api_key = None;
call.api_base = Some(&base);
call.extra_headers = Some(Map::from_iter([(
"x-request-id".to_string(),
json!("abc-123"),
)]));
complete(call).await.expect("call succeeds");
let sent = only_request(&upstream).await;
let authorization = sent.header("authorization").expect("signed request");
assert!(authorization.starts_with("AWS4-HMAC-SHA256"));
assert!(!authorization.contains("x-request-id"));
assert_eq!(sent.header("x-request-id"), Some("abc-123"));
assert_eq!(sent.url.path(), "/model/anthropic.claude-v2/converse");
let body: Value = serde_json::from_slice(&sent.body).unwrap();
assert_eq!(body["inferenceConfig"], json!({"maxTokens": 16}));
assert_eq!(
body["messages"],
json!([{"role": "user", "content": [{"text": "hi"}]}])
);
}
#[rstest]
#[case::authorization("Authorization")]
#[case::amz_date("x-amz-date")]
#[case::security_token("x-amz-security-token")]
#[case::date("Date")]
#[tokio::test]
async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) {
let upstream = wiremock::MockServer::start().await;
let base = upstream.uri();
let call = ChatCompletionsRequest {
api_key: None,
api_base: Some(&base),
extra_headers: Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
)
};
let error = complete(call)
.await
.expect_err("conflicting signing headers fail");
assert!(matches!(error, Error::Unsupported(_)), "{error:?}");
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
let upstream = upstream([bedrock_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(Map::from_iter([(
"Authorization".to_string(),
json!("Bearer caller-supplied"),
)])),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer sk-test"]
);
}
#[rstest]
#[tokio::test]
async fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(Map::from_iter([(
"authorization".to_string(),
json!("Bearer sk-ant-oat01-forwarded"),
)])),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert!(sent.header_values("x-api-key").is_empty());
assert_eq!(
sent.header_values("authorization"),
["Bearer sk-ant-oat01-forwarded"]
);
}
#[rstest]
#[tokio::test]
async fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
let upstream = upstream([bedrock_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer sk-test"]
);
}

4
litellm-rust/crates/inference-messages/AGENTS.md
View file

@ -7,8 +7,8 @@
- delegate authentication policy, beta selection, payload rewriting, and response interpretation to those adapters
- keep provider policy out of request preparation and transport handlers
- calling a concrete provider helper for every provider is still a policy dependency
- Route types such as `MessagesCall`, prepared requests, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here
- Route types such as `MessagesCall`, provider calls, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here
- The route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks
- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `src/lib.rs` explicitly sequences cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage
- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `handler::execute` explicitly sequences the provider call build in Python's handler order, cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage
- Preserve the order of validation, normalization, caller-requested parameter removal, and provider transformation when that order affects observable behavior
- Test provider dispatch, auth precedence, header handling, transformations, and responses through behavior, not source structure

22
litellm-rust/crates/inference-messages/src/common_utils.rs
View file

@ -82,17 +82,17 @@ mod tests {
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
#[case::tencent("tencent")]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[case::bedrock("bedrock")]
#[case::deepseek("deepseek")]
#[case::edenai("edenai")]
#[case::minimax("minimax")]
#[case::vertex_ai("vertex_ai")]
fn provider_keeps_its_python_name(#[case] name: &str) {
let provider =
messages_provider(name.parse::<LlmProviders>().unwrap(), "claude-sonnet-4-5").unwrap();
#[case::tencent(LlmProviders::Tencent, "tencent")]
#[case::anthropic(LlmProviders::Anthropic, "anthropic")]
#[case::azure_ai(LlmProviders::AzureAi, "azure_ai")]
#[case::bedrock(LlmProviders::Bedrock, "bedrock")]
#[case::deepseek(LlmProviders::Deepseek, "deepseek")]
#[case::edenai(LlmProviders::Edenai, "edenai")]
#[case::minimax(LlmProviders::Minimax, "minimax")]
#[case::vertex_ai(LlmProviders::VertexAi, "vertex_ai")]
#[case::openrouter(LlmProviders::Openrouter, "openrouter")]
fn provider_keeps_its_python_name(#[case] provider: LlmProviders, #[case] name: &str) {
let provider = messages_provider(provider, "claude-sonnet-4-5").unwrap();
assert_eq!(provider.as_str(), name);
}

821
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -1,148 +1,271 @@
use std::time::Duration;
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
use litellm_host::interceptors::{Interceptors, ProviderIdentity, RequestContext, WireRequest};
use litellm_http::transport::Error as TransportError;
use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream};
use litellm_auth::{AuthServices, SecretValue};
use litellm_core_utils::{
dot_notation_indexing::delete_nested_value,
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
};
use litellm_host::interceptors::{
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest,
};
use litellm_http::{request::with_default_headers, transport::Error as TransportError};
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
messages::{
context::MessagesTransformContext,
streaming::{self, ByteStream, StreamDecoder},
transformation::BaseMessagesConfig,
},
};
use litellm_llms_types::formats::messages::MessagesResponse;
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use litellm_secrets::source::Secrets;
use litellm_tracing::ByteChunk;
use serde_json::Value;
use super::{
Error, MessagesCallResponse, MessagesRoute, common_utils::truncate_error_body,
prepare::ProviderMessagesRequest,
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, string_headers, truncate_error_body},
constants::MESSAGES_TIMEOUT_SECS,
route::Messages,
types::invalid_request,
};
use crate::constants::MESSAGES_TIMEOUT_SECS;
use litellm_inference::{context::CallContext, outbound::outbound_request};
use litellm_inference::{caching::CallCache, context::CallContext, outbound::outbound_request};
pub(super) struct ProviderCall {
pub identity: ProviderIdentity,
pub wire: WireRequest,
pub(super) provider: super::common_utils::MessagesProvider,
struct ProviderCall {
identity: ProviderIdentity,
wire: WireRequest,
provider: MessagesProvider,
signer: Option<litellm_auth_aws::SigV4Signer>,
timeout: Option<Duration>,
stream: bool,
}
impl ProviderCall {
pub fn cacheable(&self) -> bool {
fn cacheable(&self) -> bool {
self.signer.is_none()
}
}
impl MessagesRoute {
pub(super) async fn prepare_outbound(
&self,
request: ProviderMessagesRequest,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<ProviderCall, Error> {
let ProviderMessagesRequest {
provider,
url,
body,
environment,
timeout,
api_key,
} = request;
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
secret_fields: Vec::new(),
api_key,
};
let authenticated =
resolve_auth(&self.auth, environment, &|key| std::env::var(key).ok()).await?;
let identity = ProviderIdentity {
model: request_context.model.clone(),
provider: request_context.custom_llm_provider.clone(),
};
let wire = context
.interceptors
.before_provider_request(
WireRequest {
url,
headers: authenticated.headers,
body: provider
.config()
.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?),
},
request_context,
)
.await?;
let stream = match wire.body.get("stream") {
None | Some(Value::Null) => false,
Some(Value::Bool(stream)) => *stream,
Some(value) => {
return Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "stream",
expected: "a boolean",
actual: value.clone(),
},
));
}
};
Ok(ProviderCall {
identity,
wire,
provider,
signer: authenticated.signer,
timeout,
stream,
pub(super) async fn execute(
route: &MessagesRoute,
provider: MessagesProvider,
model: String,
call: MessagesCall,
secrets: Secrets,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let request = provider_call(
&route.auth,
provider,
model,
call,
secrets.as_ref(),
context,
)
.boxed()
.await?;
let cache = CallCache::<Messages>::from_wire(
route.cache.as_ref().filter(|_| request.cacheable()),
context.cache,
&request.identity,
&request.wire,
);
let identity = request.identity.clone();
let (output, source) = match cache.lookup().await {
Some(hit) => hit,
None => (
call_provider(&route.http, request, context).await?,
ResultSource::Provider,
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
MessagesCallResponse::Stream { .. } => None,
};
context
.result_ready(ExecutionFacts {
provider: identity,
cost: Cost::settle(&source, reported),
source: source.clone(),
})
}
.await?;
Ok(cache.finish(output, &source).await)
}
pub(super) async fn call_provider(
&self,
request: ProviderCall,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let ProviderCall {
identity,
wire,
provider,
signer,
timeout,
stream,
} = request;
let provider_name = provider.as_str();
log_request_body(provider_name, stream, &wire.body);
let response = send(
&self.http,
Authenticated {
headers: wire.headers,
signer,
async fn provider_call(
auth: &AuthServices,
provider: MessagesProvider,
model: String,
call: MessagesCall,
secrets: &(dyn Lookup + Send + Sync),
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<ProviderCall, Error> {
let MessagesCall {
body,
api_key,
api_base,
litellm_params,
extra_headers,
provider_specific_header,
timeout,
shaping,
..
} = call;
let connection_params = litellm_router_types::LitellmParams {
api_base: api_base.clone().or(litellm_params.api_base.clone()),
..litellm_params
};
let config = provider.config();
let env_lookup = |key: &str| secrets.get(key);
let sanitized = config.shape_request(
MessagesRequest { model, ..body },
shaping.settings.reasoning_auto_summary,
)?;
let trimmed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
let forwarded = string_headers(Some(
extra_headers.into_iter().flatten().chain(scoped).collect(),
))?;
let validated = config.validate_environment(
forwarded,
api_key.as_deref(),
&trimmed.model,
&trimmed.params,
&connection_params,
&env_lookup,
)?;
let body = config.transform_anthropic_messages_request(
trimmed,
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
)?;
let environment = ValidatedEnvironment {
headers: config.request_headers(
with_default_headers(validated.headers, config.default_headers()),
&body,
),
auth: validated.auth,
};
let url = config.get_complete_url(
api_base.as_deref(),
&body.model,
&connection_params,
body.params.stream == Some(true),
&env_lookup,
)?;
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
secret_fields: Vec::new(),
api_key: api_key.map(SecretValue::new),
};
let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?;
let identity = ProviderIdentity {
model: request_context.model.clone(),
provider: request_context.custom_llm_provider.clone(),
};
let wire = context
.interceptors
.before_provider_request(
WireRequest {
url,
headers: authenticated.headers,
body: config.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?),
},
&wire.url,
&wire.body,
timeout,
request_context,
)
.await?;
if !response.status().is_success() {
return Err(provider_error(response).await);
}
let config = provider.config();
if stream {
return Ok(streaming_response(
response,
config.stream_decoder(),
provider_name,
let stream = match wire.body.get("stream") {
None | Some(Value::Null) => false,
Some(Value::Bool(stream)) => *stream,
Some(value) => {
return Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "stream",
expected: "a boolean",
actual: value.clone(),
},
));
}
let text = response.text().await.map_err(network)?;
log_response_body(&text);
context.response_received(&text).await?;
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
};
Ok(ProviderCall {
identity,
wire,
provider,
signer: authenticated.signer,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
stream,
})
}
fn without_additional_drop_params(
request: MessagesRequest,
paths: &[String],
) -> Result<MessagesRequest, Error> {
if paths.is_empty() {
return Ok(request);
}
let params = serde_json::to_value(request.params).map_err(invalid_request)?;
let trimmed = paths
.iter()
.fold(params, |params, path| delete_nested_value(params, path));
Ok(MessagesRequest {
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
..request
})
}
async fn call_provider(
http: &litellm_http::Client,
request: ProviderCall,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let ProviderCall {
identity,
wire,
provider,
signer,
timeout,
stream,
} = request;
let provider_name = provider.as_str();
log_request_body(provider_name, stream, &wire.body);
let response = send(
http,
Authenticated {
headers: wire.headers,
signer,
},
&wire.url,
&wire.body,
timeout,
)
.await?;
if !response.status().is_success() {
return Err(provider_error(response).await);
}
let config = provider.config();
if stream {
return Ok(streaming_response(
response,
config.stream_decoder(),
provider_name,
));
}
let text = response.text().await.map_err(network)?;
log_response_body(&text);
context.response_received(&text).await?;
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -271,10 +394,482 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
#[cfg(test)]
mod tests {
use rstest::rstest;
use litellm_llms_types::headers::ProviderSpecificHeaders;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
use super::*;
use crate::{MessagesSettings, MessagesShaping, common_utils::messages_provider};
use litellm_inference::provider::resolve_llm_provider;
#[fixture]
fn shaping() -> MessagesShaping {
MessagesShaping::default()
}
fn body(value: Value) -> MessagesRequest {
serde_json::from_value(value).unwrap()
}
async fn provider_call_for(call: MessagesCall) -> Result<ProviderCall, Error> {
provider_call_with_secrets(call, &|_: &str| None).await
}
async fn provider_call_with_secrets(
call: MessagesCall,
secrets: &(dyn Lookup + Send + Sync),
) -> Result<ProviderCall, Error> {
let provider = resolve_llm_provider(
&call.body.model,
call.custom_llm_provider.as_deref(),
"messages",
)?;
let config = messages_provider(provider.provider, provider.model)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?;
let model = provider.model.to_owned();
let interceptors = ();
let context = CallContext::new(&interceptors, litellm_inference::CallOptions::default());
provider_call(
&AuthServices::default(),
config,
model,
call,
secrets,
&context,
)
.await
}
async fn wire_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
provider_call_for(MessagesCall {
body: body(fields),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
})
.await
.map(|call| call.wire.body)
}
#[rstest]
#[case::api_key(
&[("ANTHROPIC_API_KEY", "sk-secret")],
&[("x-api-key", "sk-secret")],
"https://api.anthropic.com/v1/messages"
)]
#[case::auth_token(
&[("ANTHROPIC_AUTH_TOKEN", "token")],
&[("authorization", "Bearer token")],
"https://api.anthropic.com/v1/messages"
)]
#[case::api_base(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")],
&[("x-api-key", "sk-secret")],
"https://gateway.test/v1/messages"
)]
#[case::sdk_base_url(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")],
&[("x-api-key", "sk-secret")],
"https://sdk.test/v1/messages"
)]
#[tokio::test]
async fn credentials_and_base_come_from_the_resolved_secrets(
shaping: MessagesShaping,
#[case] secrets: &[(&str, &str)],
#[case] expected_auth: &[(&str, &str)],
#[case] expected_url: &str,
) {
let lookup = |name: &str| {
secrets
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
};
let call = provider_call_with_secrets(
MessagesCall {
body: body(
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
},
&lookup,
)
.await
.unwrap();
let auth: Vec<(&str, &str)> = call
.wire
.headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(
(auth.as_slice(), call.wire.url.as_str()),
(expected_auth, expected_url)
);
}
#[rstest]
#[tokio::test]
async fn missing_credentials_are_reported_before_the_request_is_transformed(
shaping: MessagesShaping,
) {
let result = provider_call_for(MessagesCall {
body: body(
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}]}),
),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
})
.await;
assert!(
matches!(
result,
Err(Error::Auth(litellm_auth::Error::MissingApiKey { .. }))
),
"{:?}",
result.err()
);
}
#[rstest]
#[case::nothing_configured(None, None, None)]
#[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
#[case::call_timeout_wins(
Some(Duration::from_secs(30)),
Some(1800.0),
Some(Duration::from_secs(30))
)]
#[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
#[case::unusable_setting_ignored(None, Some(-1.0), None)]
#[case::zero_setting_ignored(None, Some(0.0), None)]
#[tokio::test]
async fn timeout_falls_back_to_the_configured_request_timeout(
shaping: MessagesShaping,
#[case] call_timeout: Option<Duration>,
#[case] request_timeout: Option<f64>,
#[case] expected: Option<Duration>,
) {
let call = provider_call_for(MessagesCall {
body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: call_timeout,
shaping: MessagesShaping {
settings: MessagesSettings { request_timeout, ..shaping.settings },
..shaping
},
})
.await
.unwrap();
assert_eq!(call.timeout, expected);
}
#[rstest]
#[case::top_level_and_nested_paths(
json!({
"max_tokens": 1024,
"thinking": {"type": "enabled", "budget_tokens": 2048},
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
}),
&["thinking", "context_management", "tools[*].input_examples"],
json!({
"max_tokens": 1024,
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}]
}),
)]
#[case::no_paths(
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
&[],
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
)]
#[case::model_and_messages_are_never_dropped(
json!({"max_tokens": 16}),
&["model", "messages", "messages[0].content"],
json!({"max_tokens": 16}),
)]
#[tokio::test]
async fn wire_body_drops_configured_paths(
shaping: MessagesShaping,
#[case] fields: Value,
#[case] additional_drop_params: &[&str],
#[case] expected_fields: Value,
) {
let with_messages = |fields: Value| -> Value {
let Value::Object(fields) = fields else {
unreachable!()
};
Value::Object(
[
("model".to_string(), json!("claude-test")),
(
"messages".to_string(),
json!([{"role": "user", "content": "hi"}]),
),
]
.into_iter()
.chain(fields)
.collect(),
)
};
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: additional_drop_params
.iter()
.map(ToString::to_string)
.collect(),
..shaping.settings
},
..shaping
};
assert_eq!(
wire_body(with_messages(fields), shaping).await,
Ok(with_messages(expected_fields))
);
}
#[rstest]
#[case::model_prefix_picks_the_provider(
"azure_ai/claude-test",
None,
&[("x-priority", "extra"), ("x-scoped", "azure_ai")]
)]
#[case::explicit_provider(
"claude-test",
Some("anthropic"),
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
#[case::provider_prefix_on_an_anthropic_model(
"anthropic/claude-test",
None,
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
#[tokio::test]
async fn provider_specific_headers_follow_the_resolved_provider(
shaping: MessagesShaping,
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] expected: &[(&str, &str)],
) {
let configured: ProviderSpecificHeaders = serde_json::from_value(json!([
{"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}},
{"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}}
]))
.unwrap();
let call = provider_call_for(MessagesCall {
body: body(
json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: Some("sk-test".into()),
api_base: Some("https://resource.services.ai.azure.com".into()),
custom_llm_provider: custom_llm_provider.map(Into::into),
litellm_params: Default::default(),
extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])),
provider_specific_header: Some(configured),
timeout: None,
shaping,
})
.await
.unwrap();
let caller_headers: Vec<(&str, &str)> = call
.wire
.headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(caller_headers, expected);
}
#[rstest]
#[tokio::test]
async fn wire_body_carries_the_provider_stripped_model(shaping: MessagesShaping) {
assert_eq!(
wire_body(
json!({
"model": "anthropic/claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}),
shaping,
)
.await,
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}))
);
}
#[rstest]
#[tokio::test]
async fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
let shaping = MessagesShaping {
settings: MessagesSettings {
reasoning_auto_summary: true,
additional_drop_params: vec!["thinking.display".to_string()],
..shaping.settings
},
..shaping
};
assert_eq!(
wire_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}),
shaping,
)
.await,
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}))
);
}
#[rstest]
#[tokio::test]
async fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(
shaping: MessagesShaping,
) {
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["metadata.user_id".to_string()],
..shaping.settings
},
..shaping
};
assert!(matches!(
wire_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
)
.await,
Err(Error::InvalidRequest(_))
));
}
#[rstest]
#[tokio::test]
async fn wire_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) {
assert_eq!(
wire_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
)
.await,
Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "metadata.user_id",
expected: "a string",
actual: json!(123),
}
))
);
}
#[rstest]
#[case::mantle(
"bedrock_mantle/claude-test",
"https://bedrock-mantle.us-west-2.api.aws",
"/anthropic/v1/messages",
"bedrock"
)]
#[case::platform(
"bedrock/claude_platform/claude-test",
"https://aws-external-anthropic.us-west-2.api.aws",
"/v1/messages",
"aws-external-anthropic"
)]
#[tokio::test]
async fn explicit_endpoint_region_is_used_by_the_signer(
shaping: MessagesShaping,
#[case] model: &str,
#[case] base: &str,
#[case] path: &str,
#[case] service: &str,
) {
let call = provider_call_with_secrets(MessagesCall {
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
api_key: None,
api_base: Some(base.into()),
custom_llm_provider: None,
litellm_params: crate::litellm_params(Map::from_iter([
("model".into(), json!(model)),
("aws_access_key_id".into(), json!("AKIDEXAMPLE")),
("aws_secret_access_key".into(), json!("secret")),
]))
.unwrap(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
}, &|name: &str| match name {
"AWS_REGION_NAME" => Some("us-east-1".into()),
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
_ => None,
})
.await
.unwrap();
assert_eq!(call.wire.url, format!("{base}{path}"));
let signed = outbound_request(
Authenticated {
headers: call.wire.headers,
signer: call.signer,
},
call.wire.url,
&call.wire.body,
None,
)
.unwrap();
let authorization = signed.header("authorization").unwrap_or_default();
assert!(
authorization.contains(&format!("/us-west-2/{service}/aws4_request")),
"{authorization}"
);
}
#[rstest]
#[case::event(

57
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -1,19 +1,20 @@
mod common_utils;
mod constants;
mod handler;
mod prepare;
pub mod route;
mod types;
use futures_util::FutureExt;
use litellm_auth::AuthServices;
use litellm_host::interceptors::{Cost, ExecutionFacts, Interceptors, ResultSource};
use litellm_inference::{caching::CallCache, context::CallContext};
use litellm_host::interceptors::Interceptors;
use litellm_inference::context::CallContext;
use litellm_inference::provider::resolve_llm_provider;
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
pub use litellm_inference::RouteError as Error;
use crate::common_utils::messages_provider;
pub use litellm_inference::Error;
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
messages_body,
@ -74,39 +75,23 @@ impl MessagesRoute {
context: CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
litellm_inference::diagnostic::call(async {
let prepared = prepare::prepare(call, self.secrets.as_ref()).await?;
let provider = resolve_llm_provider(
&call.body.model,
call.custom_llm_provider.as_deref(),
"messages",
)?;
let config = messages_provider(provider.provider, provider.model).ok_or_else(|| {
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
})?;
litellm_inference::diagnostic::provider(
&prepared.body.model,
prepared.provider.as_str(),
provider.model,
<&'static str>::from(provider.provider),
);
let request = self.prepare_outbound(prepared, &context).boxed().await?;
let cache = CallCache::<route::Messages>::from_wire(
self.cache.as_ref().filter(|_| request.cacheable()),
context.cache,
&request.identity,
&request.wire,
);
let identity = request.identity.clone();
let config = request.provider.config();
let (output, source) = match cache.lookup().await {
Some(hit) => hit,
None => (
self.call_provider(request, &context).await?,
ResultSource::Provider,
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => config.reported_cost(response),
MessagesCallResponse::Stream { .. } => None,
};
context
.result_ready(ExecutionFacts {
provider: identity,
cost: Cost::settle(&source, reported),
source: source.clone(),
})
.await?;
Ok(cache.finish(output, &source).await)
let model = provider.model.to_owned();
let secrets = self.secrets.resolve(config.config().secret_names()).await?;
handler::execute(self, config, model, call, secrets, &context)
.boxed()
.await
})
.await
}

554
litellm-rust/crates/inference-messages/src/prepare.rs
View file

@ -1,554 +0,0 @@
use std::time::Duration;
use litellm_auth::SecretValue;
use litellm_core_utils::{
dot_notation_indexing::delete_nested_value,
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
};
use litellm_http::request::with_default_headers;
use litellm_llms::base_llm::{
auth::ValidatedEnvironment, messages::context::MessagesTransformContext,
};
use litellm_llms_types::formats::messages::MessagesRequest;
use litellm_secrets::source::SecretSource;
use super::{
Error, MessagesCall,
common_utils::{MessagesProvider, messages_provider, string_headers},
types::invalid_request,
};
use litellm_inference::provider::resolve_llm_provider;
struct ResolvedProvider {
model: String,
provider: MessagesProvider,
}
pub(super) struct ProviderMessagesRequest {
pub(super) provider: MessagesProvider,
pub(super) url: String,
pub(super) body: MessagesRequest,
pub(super) environment: ValidatedEnvironment,
pub(super) timeout: Option<Duration>,
/// The caller's own credential, reported to the host beside the wire request.
pub(super) api_key: Option<SecretValue>,
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub(super) async fn prepare(
call: MessagesCall,
secrets: &dyn SecretSource,
) -> Result<ProviderMessagesRequest, Error> {
let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?;
let secrets = secrets
.resolve(resolved.provider.config().secret_names())
.await?;
prepare_provider_request(call, resolved, secrets.as_ref())
}
fn resolve_provider(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<ResolvedProvider, Error> {
let resolved = resolve_llm_provider(model, custom_llm_provider, "messages")?;
let provider = messages_provider(resolved.provider, resolved.model)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(resolved.provider).to_string()))?;
Ok(ResolvedProvider {
model: resolved.model.to_string(),
provider,
})
}
fn prepare_provider_request(
call: MessagesCall,
resolved: ResolvedProvider,
secrets: &dyn Lookup,
) -> Result<ProviderMessagesRequest, Error> {
let ResolvedProvider { model, provider } = resolved;
let MessagesCall {
body,
api_key,
api_base,
litellm_params,
extra_headers,
provider_specific_header,
timeout,
shaping,
..
} = call;
let connection_params = litellm_router_types::LitellmParams {
api_base: api_base.clone().or(litellm_params.api_base.clone()),
..litellm_params
};
let config = provider.config();
let env_lookup = |key: &str| secrets.get(key);
let sanitized = config.shape_request(
MessagesRequest { model, ..body },
shaping.settings.reasoning_auto_summary,
)?;
let trimmed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let transformed = config.transform_anthropic_messages_request(
trimmed,
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
let forwarded = string_headers(Some(
extra_headers.into_iter().flatten().chain(scoped).collect(),
))?;
let forwarded = config.request_headers(
with_default_headers(forwarded, config.default_headers()),
&transformed,
);
let environment = config.validate_environment(
forwarded,
api_key.as_deref(),
&transformed.model,
&connection_params,
&env_lookup,
)?;
let url = config.get_complete_url(
api_base.as_deref(),
&transformed.model,
&connection_params,
transformed.params.stream == Some(true),
&env_lookup,
)?;
Ok(ProviderMessagesRequest {
provider,
url,
body: transformed,
environment,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
api_key: api_key.map(SecretValue::new),
})
}
fn without_additional_drop_params(
request: MessagesRequest,
paths: &[String],
) -> Result<MessagesRequest, Error> {
if paths.is_empty() {
return Ok(request);
}
let params = serde_json::to_value(request.params).map_err(invalid_request)?;
let trimmed = paths
.iter()
.fold(params, |params, path| delete_nested_value(params, path));
Ok(MessagesRequest {
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
..request
})
}
#[cfg(test)]
mod tests {
use litellm_llms::base_llm::auth::resolve_auth;
use litellm_llms_types::headers::ProviderSpecificHeaders;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use super::*;
use crate::{MessagesSettings, MessagesShaping};
#[fixture]
fn shaping() -> MessagesShaping {
MessagesShaping::default()
}
fn body(value: Value) -> MessagesRequest {
serde_json::from_value(value).unwrap()
}
fn prepare(call: MessagesCall) -> Result<ProviderMessagesRequest, Error> {
prepare_with_secrets(call, &|_: &str| None)
}
fn prepare_with_secrets(
call: MessagesCall,
secrets: &dyn Lookup,
) -> Result<ProviderMessagesRequest, Error> {
let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?;
prepare_provider_request(call, resolved, secrets)
}
/// The headers as they go on the wire, credential applied.
fn wire_headers(prepared: &ProviderMessagesRequest) -> Vec<(String, String)> {
tokio::runtime::Builder::new_current_thread()
.build()
.unwrap()
.block_on(resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment.clone(),
&|_| None,
))
.unwrap()
.headers
}
#[rstest]
#[case::api_key(
&[("ANTHROPIC_API_KEY", "sk-secret")],
&[("x-api-key", "sk-secret")],
"https://api.anthropic.com/v1/messages"
)]
#[case::auth_token(
&[("ANTHROPIC_AUTH_TOKEN", "token")],
&[("authorization", "Bearer token")],
"https://api.anthropic.com/v1/messages"
)]
#[case::api_base(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")],
&[("x-api-key", "sk-secret")],
"https://gateway.test/v1/messages"
)]
#[case::sdk_base_url(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")],
&[("x-api-key", "sk-secret")],
"https://sdk.test/v1/messages"
)]
fn credentials_and_base_come_from_the_resolved_secrets(
shaping: MessagesShaping,
#[case] secrets: &[(&str, &str)],
#[case] expected_auth: &[(&str, &str)],
#[case] expected_url: &str,
) {
let lookup = |name: &str| {
secrets
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
};
let prepared = prepare_with_secrets(
MessagesCall {
body: body(
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
},
&lookup,
)
.unwrap();
let headers = wire_headers(&prepared);
let auth: Vec<(&str, &str)> = headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(
(auth.as_slice(), prepared.url.as_str()),
(expected_auth, expected_url)
);
}
#[rstest]
#[case::nothing_configured(None, None, None)]
#[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
#[case::call_timeout_wins(
Some(Duration::from_secs(30)),
Some(1800.0),
Some(Duration::from_secs(30))
)]
#[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
#[case::unusable_setting_ignored(None, Some(-1.0), None)]
#[case::zero_setting_ignored(None, Some(0.0), None)]
fn timeout_falls_back_to_the_configured_request_timeout(
shaping: MessagesShaping,
#[case] call_timeout: Option<Duration>,
#[case] request_timeout: Option<f64>,
#[case] expected: Option<Duration>,
) {
let prepared = prepare(MessagesCall {
body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: call_timeout,
shaping: MessagesShaping {
settings: MessagesSettings { request_timeout, ..shaping.settings },
..shaping
},
})
.unwrap();
assert_eq!(prepared.timeout, expected);
}
fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
prepare(MessagesCall {
body: body(fields),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
})
.map(|prepared| serde_json::to_value(prepared.body).unwrap())
}
#[rstest]
#[case::top_level_and_nested_paths(
json!({
"max_tokens": 1024,
"thinking": {"type": "enabled", "budget_tokens": 2048},
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
}),
&["thinking", "context_management", "tools[*].input_examples"],
json!({
"max_tokens": 1024,
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}]
}),
)]
#[case::no_paths(
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
&[],
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
)]
#[case::model_and_messages_are_never_dropped(
json!({"max_tokens": 16}),
&["model", "messages", "messages[0].content"],
json!({"max_tokens": 16}),
)]
fn prepared_body_drops_configured_paths(
shaping: MessagesShaping,
#[case] fields: Value,
#[case] additional_drop_params: &[&str],
#[case] expected_fields: Value,
) {
let with_messages = |fields: Value| -> Value {
let Value::Object(fields) = fields else {
unreachable!()
};
Value::Object(
[
("model".to_string(), json!("claude-test")),
(
"messages".to_string(),
json!([{"role": "user", "content": "hi"}]),
),
]
.into_iter()
.chain(fields)
.collect(),
)
};
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: additional_drop_params
.iter()
.map(ToString::to_string)
.collect(),
..shaping.settings
},
..shaping
};
assert_eq!(
prepared_body(with_messages(fields), shaping),
Ok(with_messages(expected_fields))
);
}
#[rstest]
#[case::model_prefix_picks_the_provider(
"azure_ai/claude-test",
None,
&[("x-priority", "extra"), ("x-scoped", "azure_ai")]
)]
#[case::explicit_provider(
"claude-test",
Some("anthropic"),
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
#[case::provider_prefix_on_an_anthropic_model(
"anthropic/claude-test",
None,
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
fn provider_specific_headers_follow_the_resolved_provider(
shaping: MessagesShaping,
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] expected: &[(&str, &str)],
) {
let configured: ProviderSpecificHeaders = serde_json::from_value(json!([
{"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}},
{"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}}
]))
.unwrap();
let prepared = prepare(MessagesCall {
body: body(
json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: Some("sk-test".into()),
api_base: Some("https://resource.services.ai.azure.com".into()),
custom_llm_provider: custom_llm_provider.map(Into::into),
litellm_params: Default::default(),
extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])),
provider_specific_header: Some(configured),
timeout: None,
shaping,
})
.unwrap();
let caller_headers: Vec<(&str, &str)> = prepared
.environment
.headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(caller_headers, expected);
}
#[rstest]
fn prepared_body_carries_the_provider_stripped_model(shaping: MessagesShaping) {
assert_eq!(
prepared_body(
json!({
"model": "anthropic/claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}),
shaping,
),
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}))
);
}
#[rstest]
fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
let shaping = MessagesShaping {
settings: MessagesSettings {
reasoning_auto_summary: true,
additional_drop_params: vec!["thinking.display".to_string()],
..shaping.settings
},
..shaping
};
assert_eq!(
prepared_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}),
shaping,
),
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}))
);
}
#[rstest]
fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(shaping: MessagesShaping) {
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["metadata.user_id".to_string()],
..shaping.settings
},
..shaping
};
assert!(matches!(
prepared_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
),
Err(Error::InvalidRequest(_))
));
}
#[rstest]
fn prepared_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) {
assert_eq!(
prepared_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
),
Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "metadata.user_id",
expected: "a string",
actual: json!(123),
}
))
);
}
#[rstest]
#[case::mantle(
"bedrock_mantle/claude-test",
"https://bedrock-mantle.us-west-2.api.aws",
"/anthropic/v1/messages",
"bedrock"
)]
#[case::platform(
"bedrock/claude_platform/claude-test",
"https://aws-external-anthropic.us-west-2.api.aws",
"/v1/messages",
"aws-external-anthropic"
)]
fn explicit_endpoint_region_is_used_by_the_signer(
shaping: MessagesShaping,
#[case] model: &str,
#[case] base: &str,
#[case] path: &str,
#[case] service: &str,
) {
let prepared = prepare_with_secrets(MessagesCall {
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
api_key: None,
api_base: Some(base.into()),
custom_llm_provider: None,
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
}, &|name: &str| match name {
"AWS_REGION_NAME" => Some("us-east-1".into()),
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
_ => None,
}).unwrap();
assert_eq!(prepared.url, format!("{base}{path}"));
assert!(
matches!(prepared.environment.auth, litellm_llms::base_llm::auth::AuthScheme::AwsSigV4 { region, service: actual, .. } if region == "us-west-2" && actual == service)
);
}
}

10
litellm-rust/crates/inference-messages/tests/caching.rs
View file

@ -17,7 +17,7 @@ use litellm_host::interceptors::{
ResultSource, WireRequest,
};
use litellm_inference::{
RouteError,
Error,
caching::{CacheRequest, execute_unary},
};
use rstest::{fixture, rstest};
@ -133,12 +133,12 @@ struct ChangingHooks {
facts: std::sync::Mutex<Vec<ExecutionFacts>>,
}
impl Interceptors<RouteError> for ChangingHooks {
impl Interceptors<Error> for ChangingHooks {
async fn before_provider_request(
&self,
mut wire: WireRequest,
_: RequestContext,
) -> Result<WireRequest, RouteError> {
) -> Result<WireRequest, Error> {
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if self.rewrite {
wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 });
@ -146,11 +146,11 @@ impl Interceptors<RouteError> for ChangingHooks {
Ok(wire)
}
async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> {
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Ok(())
}
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> {
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
self.facts.lock().unwrap().push(facts);
Ok(())
}

71
litellm-rust/crates/inference-messages/tests/messages/request.rs
View file

@ -343,6 +343,77 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
assert_eq!(sent, expected);
}
#[rstest]
#[case::anthropic("anthropic", &[AnthropicBeta::DangerousToolUse20260903])]
#[case::deepseek("deepseek", &[])]
#[case::minimax("minimax", &[])]
#[tokio::test]
async fn safeguards_add_the_dangerous_tool_use_beta_only_where_the_beta_map_has_a_column(
call: MessagesCall,
#[case] provider: &str,
#[case] derived: &[AnthropicBeta],
) {
let upstream = upstream([message_response()]).await;
let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]);
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
..call
},
json!({"safeguards": safeguards}),
))
.await;
let request = only_request(&upstream).await;
let expected: BetaSet = derived
.iter()
.cloned()
.chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
.collect();
assert_eq!(sent_betas(&request), expected);
assert_eq!(
request.body_json::<Value>().unwrap()["safeguards"],
safeguards
);
}
#[rstest]
#[tokio::test]
async fn null_safeguards_are_dropped_instead_of_sent_without_their_beta(
call: MessagesCall,
#[values("anthropic", "azure_ai", "bedrock_mantle", "deepseek", "minimax")] provider: &str,
) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
..call
},
json!({"safeguards": null}),
))
.await;
let request = only_request(&upstream).await;
assert_eq!(
sent_betas(&request),
BetaSet::from_iter([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
);
assert!(
!request
.body_json::<Map<String, Value>>()
.unwrap()
.contains_key("safeguards")
);
}
#[rstest]
#[tokio::test]
async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) {

2
litellm-rust/crates/inference-ocr/AGENTS.md
View file

@ -3,4 +3,4 @@
- The route returns its completed response directly
- Route closures supply OCR-specific host capabilities, such as the caller's Azure AD token provider (`src/route.rs`)
- Provider code reaches the caller's hooks mid-call only through `litellm_llms::base_llm::ocr::handler::CallHooks`, which this crate implements over its host until it folds into `litellm_host::interceptors::Interceptors`
- OCR failures use `base_llm/ocr/error.rs`, the recorded exception to `litellm_inference::RouteError` until OCR folds into `litellm_llms::Error`
- OCR failures use `base_llm/ocr/error.rs`, the recorded exception to `litellm_inference::Error` until OCR folds into `litellm_llms::Error`

4
litellm-rust/crates/inference-ocr/src/arguments.rs
View file

@ -44,8 +44,8 @@ pub fn consumed_optional_param_names(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Vec<&'static str>, Error> {
let (model, config) = resolve_provider_config(model, custom_llm_provider)?;
let provider_fields = config.get_supported_ocr_params(&model);
let (provider, config) = resolve_provider_config(model, custom_llm_provider)?;
let provider_fields = config.get_supported_ocr_params(provider.model);
let auth_fields: &[&str] = match config {
OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS,
OcrConfigKind::AzureAi

81
litellm-rust/crates/inference-ocr/src/client.rs
View file

@ -1,81 +0,0 @@
use litellm_host::observation::ObservationSender;
use std::sync::Arc;
use litellm_host::interceptors::Interceptors;
use litellm_llms::base_llm::ocr::{error::Error, handler::OcrClient};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use super::{
handler::perform_ocr_request,
types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest},
};
#[derive(Clone)]
pub struct OcrRoute {
client: OcrClient,
}
impl OcrRoute {
pub fn new(client: OcrClient) -> Self {
Self { client }
}
pub async fn execute(
&self,
request: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run(request, interceptors, observers.as_ref()),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "ocr",
model = %request.model,
resolved_model = %request.model,
provider = <&str>::from(request.config.provider()),
stream = false,
outcome
))]
pub(super) async fn run(
&self,
request: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_inference::diagnostic::unary(async {
let caller_document = matches!(&request.document, OcrDocumentInput::Document(_));
let prepared = prepare_request_document(request).await?;
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
Box::pin(perform_ocr_request(
&self.client,
prepared,
interceptors,
caller_document,
observers,
));
execute.await
})
.await
}
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
async fn prepare_request_document(
request: LiteLLMOcrRequest<OcrDocumentInput>,
) -> Result<ResolvedOcrRequest, Error> {
if let OcrDocumentInput::Document(_) = &request.document {
return request.map_document(super::document::prepare_document);
}
let logger = litellm_tracing::Logger::current();
let span = tracing::Span::current();
tokio::task::spawn_blocking(move || {
logger.scope(|| span.in_scope(|| request.map_document(super::document::prepare_document)))
})
.await
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
}

19
litellm-rust/crates/inference-ocr/src/document.rs
View file

@ -1,10 +1,10 @@
use std::{collections::BTreeMap as Map, io::Read, path::Path};
use std::{collections::BTreeMap as Map, io::Read, path::Path, sync::Arc};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_llms::base_llm::ocr::{error::Error, transformation::OCR_INLINE_MAX_BYTES};
use litellm_llms_types::formats::ocr::OcrDocument;
use crate::types::OcrDocumentInput;
use crate::types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest};
pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
match input {
@ -24,6 +24,21 @@ pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
}
}
pub(crate) async fn resolve_document(
request: LiteLLMOcrRequest<OcrDocumentInput>,
) -> Result<ResolvedOcrRequest, Error> {
if let OcrDocumentInput::Document(_) = &request.document {
return request.map_document(prepare_document);
}
let logger = litellm_tracing::Logger::current();
let span = tracing::Span::current();
tokio::task::spawn_blocking(move || {
logger.scope(|| span.in_scope(|| request.map_document(prepare_document)))
})
.await
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
}
pub fn read_path_document(path: &Path, mime_type: Option<&str>) -> Result<OcrDocument, Error> {
let mut bytes = Vec::new();
std::fs::File::open(path)

296
litellm-rust/crates/inference-ocr/src/handler.rs
View file

@ -4,29 +4,45 @@ use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender};
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::PreparedOcrRequest,
transformation::{OcrConnection, PreparedOcrRequest},
};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::{arguments::is_secret_param, prepare::prepare_request, provider_config::OcrConfigKind};
use crate::types::ResolvedOcrRequest;
use super::{arguments::is_secret_param, provider_config::OcrConfigKind};
use crate::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) async fn perform_ocr_request(
pub(crate) async fn execute(
client: &OcrClient,
request: ResolvedOcrRequest,
config: OcrConfigKind,
call: ResolvedOcrRequest,
secrets: Secrets,
host: &impl Interceptors<Error>,
caller_document: bool,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
request.response_format()?;
let config = request.config;
let secrets = client
.secret_source()
.resolve(&config.secret_names())
.await
.map_err(|error| Error::Secret(std::sync::Arc::new(error)))?;
let request = prepare_request(request, caller_document, client, secrets);
call.response_format()?;
let LiteLLMOcrRequest {
model,
document,
credentials,
transport,
optional_params,
input_sources,
azure_ad_token_provider,
..
} = call;
let resolved = config.resolve_credentials(credentials, &secrets);
let request = PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
caller_document,
optional_params,
input_sources,
azure_ad_token_provider,
};
let interceptors = OcrCallHooks::new(host, &request, config, observers);
config.ocr(client, &request, &interceptors).await
}
@ -86,3 +102,257 @@ impl<H: Interceptors<Error>> CallHooks<Error> for OcrCallHooks<'_, H> {
Box::pin(self.interceptors.after_provider_response(raw))
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
use litellm_llms_types::formats::ocr::OcrResponseFormat;
use rstest::rstest;
use serde_json::{Value, json};
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
use super::*;
use crate::{
OcrRoute,
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
fn client() -> OcrClient {
OcrClient::for_test(
litellm_http::Client::plain_for_test(),
litellm_http::Client::no_redirect_for_test(),
)
.with_secrets(litellm_inference_testing::no_secrets())
}
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
async fn run_request(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
OcrRoute::new(client()).execute(request, &(), None).await
}
async fn upstream(request_count: u64, response: Value) -> MockServer {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(ResponseTemplate::new(200).set_body_json(response))
.up_to_n_times(request_count)
.mount(&server)
.await;
server
}
fn image(url: &str) -> Value {
json!({"type": "image_url", "image_url": url})
}
fn body(request: &wiremock::Request) -> Value {
serde_json::from_slice(&request.body).unwrap()
}
fn assert_mistral_request(request: &wiremock::Request) {
assert_eq!(
request
.headers
.get("authorization")
.unwrap()
.to_str()
.unwrap(),
"Bearer test-key"
);
assert_eq!(
request
.headers
.get("content-type")
.unwrap()
.to_str()
.unwrap(),
"application/json"
);
assert_eq!(
body(request),
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true,
"unknown": "preserved"
})
);
}
#[rstest]
#[tokio::test]
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
let upstream = upstream(
1,
json!({"pages": [{"index": 0, "markdown": {"content": "hello"}}]}),
)
.await;
let request = request(
"cohere/parse",
&upstream.uri(),
image("https://example.com/original.png"),
json!({
"output_format": "markdown", "timeout": 30,
"extra_body": {
"output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
}
}),
);
run_request(request).await.unwrap();
let [request] = upstream
.received_requests()
.await
.unwrap()
.try_into()
.unwrap();
assert_eq!(
body(&request),
json!({
"model": "parse", "output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
})
);
}
#[rstest]
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let upstream = upstream(
1,
json!({"pages": [{"index": 0, "markdown": {"content": "hello"}}]}),
)
.await;
let request = request(
"cohere/parse",
&upstream.uri(),
image("https://example.com/a.png"),
json!({"output_format": null, "req_format": null}),
);
assert_eq!(
request.response_format().unwrap(),
OcrResponseFormat::Litellm
);
run_request(request).await.unwrap();
let [request] = upstream
.received_requests()
.await
.unwrap()
.try_into()
.unwrap();
let body = body(&request);
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[rstest]
#[tokio::test]
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
let upstream = upstream(
2,
json!({
"pages": [{"index": 0, "markdown": "hello"}],
"extra": "preserved"
}),
)
.await;
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "preserved"
});
let document =
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
let direct_call = request(
"mistral/mistral-ocr-maas",
&upstream.uri(),
document.clone(),
options.clone(),
);
let vertex_call = request(
"vertex_ai/mistral-ocr-maas",
&upstream.uri(),
document,
options,
);
assert_eq!(direct_call.transport.timeout, Some(Duration::from_secs(2)));
assert_eq!(vertex_call.transport.timeout, Some(Duration::from_secs(2)));
let direct_response = run_request(direct_call).await.unwrap().into_json();
let vertex_response = run_request(vertex_call).await.unwrap().into_json();
let [direct_http, vertex_http] = upstream
.received_requests()
.await
.unwrap()
.try_into()
.unwrap();
assert_eq!(direct_http.url.path(), "/v1/ocr");
assert_eq!(
vertex_http.url.path(),
"/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
assert_mistral_request(&direct_http);
assert_mistral_request(&vertex_http);
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
#[derive(serde::Deserialize)]
struct KnownParams {
pages: Option<Vec<i64>>,
}
#[rstest]
fn parsed_provider_params_separates_known_and_extra_params() {
let arguments: CallArguments = serde_json::from_value(json!({
"pages": [0, 2],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}))
.unwrap();
let known: KnownParams = parse_options(&arguments).unwrap();
assert_eq!(known.pages, Some(vec![0, 2]));
assert_eq!(arguments["future_ocr_option"], true);
assert_eq!(arguments["extra_body"], json!({"provider_option": "value"}));
assert_eq!(
arguments
.iter()
.filter(|(name, _)| name.as_str() != "pages")
.count(),
2
);
assert_eq!(
compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(),
json!({
"pages": [0, 2], "future_ocr_option": true, "provider_option": "value"
})
);
}
}

78
litellm-rust/crates/inference-ocr/src/lib.rs
View file

@ -1,10 +1,82 @@
pub mod arguments;
mod client;
pub use client::OcrRoute;
pub mod document;
pub(crate) mod handler;
pub(crate) mod prepare;
pub mod provider_config;
pub mod route;
pub mod types;
pub mod wire;
use std::sync::Arc;
use litellm_host::interceptors::Interceptors;
use litellm_host::observation::ObservationSender;
use litellm_llms::base_llm::ocr::{error::Error, handler::OcrClient};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use crate::{
handler::execute,
types::{LiteLLMOcrRequest, OcrDocumentInput},
};
#[derive(Clone)]
pub struct OcrRoute {
client: OcrClient,
}
impl OcrRoute {
pub fn new(client: OcrClient) -> Self {
Self { client }
}
pub async fn execute(
&self,
request: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run(request, interceptors, observers.as_ref()),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "ocr",
model = %call.model,
resolved_model = %call.model,
provider = <&str>::from(call.config.provider()),
stream = false,
outcome
))]
async fn run(
&self,
call: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_inference::diagnostic::unary(async {
let caller_document = matches!(&call.document, OcrDocumentInput::Document(_));
let call = document::resolve_document(call).await?;
let config = call.config;
let secrets = self
.client
.secret_source()
.resolve(&config.secret_names())
.await
.map_err(|error| Error::Secret(Arc::new(error)))?;
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
Box::pin(execute(
&self.client,
config,
call,
secrets,
interceptors,
caller_document,
observers,
));
execute.await
})
.await
}
}

309
litellm-rust/crates/inference-ocr/src/prepare.rs
View file

@ -1,309 +0,0 @@
use litellm_auth::{InputSource, SecretValue, Sourced};
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_llms::base_llm::ocr::{
handler::OcrClient,
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
};
use litellm_secrets::source::Secrets;
use crate::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) fn prepare_request(
request: ResolvedOcrRequest,
caller_document: bool,
client: &OcrClient,
secrets: Secrets,
) -> PreparedOcrRequest {
let credentials = request.credentials.clone();
let (preferred_api_key_env, api_base_env) = match request.config.provider() {
LlmProviders::Mistral => (
Some("MISTRAL_AZURE_API_KEY"),
Some("MISTRAL_AZURE_API_BASE"),
),
LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")),
_ => (None, None),
};
let secret = |name: &str| secrets.truthy(name);
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
preferred_api_key_env
.into_iter()
.chain(request.config.get_api_key_env_var())
.find_map(secret)
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
});
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
credentials.api_base.clone().or_else(|| {
api_base_env
.and_then(secret)
.map(|value| Sourced::new(value, InputSource::Environment))
})
});
let resolved = request
.config
.resolve_connection_params(OcrCredentialInputs {
dynamic_api_key,
dynamic_api_base,
..credentials
});
let LiteLLMOcrRequest {
model,
document,
transport,
optional_params,
input_sources,
azure_ad_token_provider,
..
} = request;
PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
caller_document,
optional_params,
input_sources,
azure_ad_token_provider,
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use futures_util::future::BoxFuture;
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
use litellm_host::interceptors::WireRequest;
use litellm_llms::{
base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::BaseOcrConfig,
},
cohere::ocr::transformation::CohereParseConfig,
mistral::ocr::transformation::MistralOcrConfig,
vertex_ai::ocr::transformation::VertexAiOcrConfig,
};
use litellm_llms_types::formats::ocr::OcrResponseFormat;
use serde_json::{Value, json};
use super::*;
use crate::{
document::prepare_document,
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
/// Stands in for a host with no interceptors registered.
struct NoHooks;
impl CallHooks<Error> for NoHooks {
fn before_provider_request(
&self,
wire: WireRequest,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
Box::pin(async move { Ok(wire) })
}
fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(async { Ok(()) })
}
}
fn client() -> OcrClient {
OcrClient::for_test(
litellm_http::Client::plain_for_test(),
litellm_http::Client::no_redirect_for_test(),
)
}
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
fn prepared(request: LiteLLMOcrRequest) -> PreparedOcrRequest {
prepare_request(
request.map_document(prepare_document).unwrap(),
true,
&client(),
std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
)
}
fn image(url: &str) -> Value {
json!({"type": "image_url", "image_url": url})
}
#[rstest::rstest]
#[tokio::test]
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
let request = request(
"cohere/parse",
"https://example.com",
image("https://example.com/original.png"),
json!({
"output_format": "markdown", "timeout": 30,
"extra_body": {
"output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
}
}),
);
let http = CohereParseConfig
.prepare_request(&prepared(request), &client(), &NoHooks)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
"model": "parse", "output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
})
);
}
#[rstest::rstest]
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let request = request(
"cohere/parse",
"https://example.com",
image("https://example.com/a.png"),
json!({"output_format": null, "req_format": null}),
);
assert_eq!(
request.response_format().unwrap(),
OcrResponseFormat::Litellm
);
let http = CohereParseConfig
.prepare_request(&prepared(request), &client(), &NoHooks)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[rstest::rstest]
#[tokio::test]
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "preserved"
});
let document =
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
let direct = prepared(request(
"mistral/mistral-ocr-maas",
"https://mistral.test",
document.clone(),
options.clone(),
));
let vertex = prepared(request(
"vertex_ai/mistral-ocr-maas",
"https://vertex.test",
document,
options,
));
let direct_http = MistralOcrConfig
.prepare_request(&direct, &client(), &NoHooks)
.await
.unwrap();
let vertex_http = VertexAiOcrConfig
.prepare_request(&vertex, &client(), &NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
for http in [&direct_http, &vertex_http] {
assert_eq!(http.header("authorization").unwrap(), "Bearer test-key");
assert_eq!(http.header("content-type").unwrap(), "application/json");
assert_eq!(http.timeout(), Some(Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true,
"unknown": "preserved"
})
);
}
let payload = serde_json::to_vec(
&json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}),
)
.unwrap();
let direct_response = MistralOcrConfig
.transform_ocr_response(&direct.model, &payload, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
let vertex_response = VertexAiOcrConfig
.transform_ocr_response(&vertex.model, &payload, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
#[derive(serde::Deserialize)]
struct KnownParams {
pages: Option<Vec<i64>>,
}
#[rstest::rstest]
fn parsed_provider_params_separates_known_and_extra_params() {
let arguments: CallArguments = serde_json::from_value(json!({
"pages": [0, 2],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}))
.unwrap();
let known: KnownParams = parse_options(&arguments).unwrap();
assert_eq!(known.pages, Some(vec![0, 2]));
assert_eq!(arguments["future_ocr_option"], true);
assert_eq!(arguments["extra_body"], json!({"provider_option": "value"}));
assert_eq!(
arguments
.iter()
.filter(|(name, _)| name.as_str() != "pages")
.count(),
2
);
assert_eq!(
compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(),
json!({
"pages": [0, 2], "future_ocr_option": true, "provider_option": "value"
})
);
}
}

106
litellm-rust/crates/inference-ocr/src/provider_config.rs
View file

@ -1,5 +1,6 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_auth::{InputSource, SecretValue, Sourced};
use litellm_core_utils::get_llm_provider_logic::{LlmProviders, get_custom_llm_provider};
use litellm_inference::provider::{ResolvedProvider, resolve_llm_provider};
use litellm_llms::{
aws_textract::ocr::{
analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation,
@ -25,6 +26,7 @@ use litellm_llms::{
},
};
use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat};
use litellm_secrets::source::Secrets;
macro_rules! with_config {
($kind:expr, $config:ident => $body:expr) => {
@ -115,7 +117,18 @@ impl OcrConfigKind {
}
pub(crate) fn secret_names(self) -> Vec<&'static str> {
with_config!(self, config => config.secret_names())
let names = with_config!(self, config => config.secret_names());
let fallback_names: &[&str] = match self.provider() {
LlmProviders::Mistral => &["MISTRAL_AZURE_API_KEY", "MISTRAL_AZURE_API_BASE"],
LlmProviders::AzureAi => &["AZURE_AI_API_BASE"],
_ => &[],
};
let additional_names = fallback_names
.iter()
.copied()
.filter(|name| !names.contains(name))
.collect::<Vec<_>>();
names.into_iter().chain(additional_names).collect()
}
pub(crate) fn get_health_check_document(self) -> OcrDocument {
@ -129,6 +142,43 @@ impl OcrConfigKind {
with_config!(self, config => config.resolve_connection_params(inputs))
}
pub(crate) fn resolve_credentials(
self,
credentials: OcrCredentialInputs,
secrets: &Secrets,
) -> ResolvedOcrCredentials {
let (preferred_api_key_env, api_base_env) = match self.provider() {
LlmProviders::Mistral => (
Some("MISTRAL_AZURE_API_KEY"),
Some("MISTRAL_AZURE_API_BASE"),
),
LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")),
_ => (None, None),
};
let secret = |name: &str| secrets.truthy(name);
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
preferred_api_key_env
.into_iter()
.chain(self.get_api_key_env_var())
.find_map(secret)
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
});
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
credentials.api_base.clone().or_else(|| {
api_base_env
.and_then(secret)
.map(|value| Sourced::new(value, InputSource::Environment))
})
});
self.resolve_connection_params(OcrCredentialInputs {
dynamic_api_key,
dynamic_api_base,
..credentials
})
}
pub(crate) async fn ocr(
self,
client: &OcrClient,
@ -164,7 +214,7 @@ pub fn passthrough_response(
endpoint: &str,
body: &[u8],
) -> Result<Option<LiteLLMOcrResponse>, Error> {
let (model, config) = resolve_provider_config(model, Some("azure_ai"))?;
let (provider, config) = resolve_provider_config(model, Some("azure_ai"))?;
let segments: Vec<&str> = endpoint
.split('/')
.filter(|segment| !segment.is_empty())
@ -173,7 +223,7 @@ pub fn passthrough_response(
OcrConfigKind::AzureAi => segments == AZURE_AI_OCR_PATH,
OcrConfigKind::AzureCohere => segments == AZURE_COHERE_PARSE_PATH,
OcrConfigKind::AzureDocumentIntelligence => {
segments == AzureDocumentIntelligenceOcrConfig::analyze_path(&model)?
segments == AzureDocumentIntelligenceOcrConfig::analyze_path(provider.model)?
}
other => {
let provider: &'static str = other.provider().into();
@ -183,24 +233,28 @@ pub fn passthrough_response(
if !is_ocr_endpoint {
return Ok(None);
}
with_config!(config, config => config.transform_ocr_response(&model, body, OcrResponseFormat::Litellm))
with_config!(config, config => config.transform_ocr_response(provider.model, body, OcrResponseFormat::Litellm))
.map(Some)
}
pub(crate) fn resolve_provider_config(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<(String, OcrConfigKind), Error> {
let provider =
get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider {
model,
custom_llm_provider: LlmProviders::Mistral.into(),
});
let llm_provider = provider
.custom_llm_provider
.parse::<LlmProviders>()
.map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
let config = match llm_provider {
pub(crate) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<(ResolvedProvider<'a>, OcrConfigKind), Error> {
let provider = match resolve_llm_provider(model, custom_llm_provider, "ocr") {
Ok(provider) => provider,
Err(_) if get_custom_llm_provider(model, custom_llm_provider).is_none() => {
ResolvedProvider {
model,
provider: LlmProviders::Mistral,
}
}
Err(litellm_inference::Error::InvalidProvider(provider)) => {
return Err(Error::InvalidProvider(provider));
}
Err(error) => return Err(Error::InvalidProvider(error.to_string())),
};
let config = match provider.provider {
LlmProviders::AwsTextract => match TextractOperation::from_model(provider.model)? {
TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract,
TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze,
@ -227,11 +281,11 @@ pub(crate) fn resolve_provider_config(
LlmProviders::VertexAi => OcrConfigKind::VertexAi,
_ => {
return Err(Error::InvalidProvider(
provider.custom_llm_provider.to_string(),
<&'static str>::from(provider.provider).to_string(),
));
}
};
Ok((provider.model.to_string(), config))
Ok((provider, config))
}
fn is_document_intelligence_model(model: &str) -> bool {
@ -243,7 +297,7 @@ fn is_document_intelligence_model(model: &str) -> bool {
mod tests {
use std::collections::HashSet;
use litellm_auth::{InputSource, Sourced};
use litellm_auth::Sourced;
use litellm_llms::{
base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document,
};
@ -484,8 +538,8 @@ mod tests {
#[case] expected_config: OcrConfigKind,
) {
let expected_model = qualified_model.split_once('/').unwrap().1;
let (model, config) = resolve_provider_config(qualified_model, None).unwrap();
assert_eq!(model, expected_model);
let (provider, config) = resolve_provider_config(qualified_model, None).unwrap();
assert_eq!(provider.model, expected_model);
assert_eq!(config, expected_config);
}
@ -528,7 +582,7 @@ mod tests {
expected_config
);
assert_eq!(
resolve_provider_config(model, None).unwrap().0,
resolve_provider_config(model, None).unwrap().0.model,
model.split_once('/').unwrap().1
);
}

13
litellm-rust/crates/inference-ocr/src/types.rs
View file

@ -75,6 +75,8 @@ impl OcrConnectionInputs {
pub struct LiteLLMOcrRequest<D = OcrDocumentInput> {
pub model: String,
pub(crate) requested_model: String,
pub(crate) custom_llm_provider: Option<String>,
pub document: D,
pub credentials: OcrCredentialInputs,
pub transport: OcrTransportConfig,
@ -91,7 +93,8 @@ impl LiteLLMOcrRequest {
custom_llm_provider: Option<&str>,
optional_params: CallArguments,
) -> Result<Self, Error> {
let (model, config) = resolve_provider_config(&model, custom_llm_provider)?;
let (provider, config) = resolve_provider_config(&model, custom_llm_provider)?;
let resolved_model = provider.model.to_owned();
let default_transport = OcrTransportConfig::default();
let max_response_bytes = optional_params
.get("max_response_bytes")
@ -116,7 +119,9 @@ impl LiteLLMOcrRequest {
.collect();
Ok(Self {
model,
model: resolved_model,
requested_model: model,
custom_llm_provider: custom_llm_provider.map(str::to_owned),
document: document.into(),
credentials: OcrCredentialInputs::default(),
transport,
@ -135,6 +140,8 @@ impl<D> LiteLLMOcrRequest<D> {
) -> Result<LiteLLMOcrRequest<T>, E> {
Ok(LiteLLMOcrRequest {
model: self.model,
requested_model: self.requested_model,
custom_llm_provider: self.custom_llm_provider,
document: map(self.document)?,
credentials: self.credentials,
transport: self.transport,
@ -148,6 +155,8 @@ impl<D> LiteLLMOcrRequest<D> {
pub fn with_document<T>(self, document: T) -> LiteLLMOcrRequest<T> {
LiteLLMOcrRequest {
model: self.model,
requested_model: self.requested_model,
custom_llm_provider: self.custom_llm_provider,
document,
credentials: self.credentials,
transport: self.transport,

13
litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs
View file

@ -66,6 +66,19 @@ async fn interception_runs_once_and_observers_receive_ordered_success_events() {
assert_eq!(received(&upstream).await.len(), 1);
}
#[rstest::rstest]
#[case::mutated_model("model-b")]
#[tokio::test]
async fn provider_request_uses_the_request_model_after_it_is_mutated(#[case] model: &str) {
let upstream = upstream([pages_response()]).await;
let mut request = ocr_request("mistral/model-a", &upstream.uri(), json!({}));
request.model = model.into();
perform(request).await.unwrap();
assert_eq!(only_request(&upstream).await.json()["model"], json!(model));
}
#[rstest::rstest]
#[tokio::test]
async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {

1
litellm-rust/crates/inference-responses/AGENTS.md
View file

@ -2,3 +2,4 @@
- This crate owns Responses API call orchestration. Payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/responses` and `llms/src/<provider>/responses`
- The HTTP route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks, cached through `litellm_inference::caching::execute_streaming`
- WebSocket sessions (`src/websocket.rs`) stay separate from the HTTP call driver because a connection can accept multiple requests while receiving events
- `src/ids.rs` owns the LiteLLM-managed ids handed back to callers (`resp_`, `encitem_`, `litellm_enc:`, `cntr_`) and must stay fully compatible with `litellm/responses/utils.py`

1
litellm-rust/crates/inference-responses/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
base64.workspace = true
bytes.workspace = true
futures-util.workspace = true
litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }

61
litellm-rust/crates/inference-responses/src/handler.rs
View file

@ -3,39 +3,61 @@ use litellm_host::observation::ObservationSender;
use std::time::Duration;
use futures_util::StreamExt;
use litellm_host::interceptors::{Interceptors, RawResponse, WireRequest};
use litellm_llms::base_llm::auth::{Authenticated, resolve_auth};
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
use litellm_llms::base_llm::auth::{AuthScheme, Authenticated, resolve_auth};
use litellm_secrets::source::Secrets;
use super::{
Error,
types::{ProviderResponsesRequest, ResponsesOutput, ResponsesStreamHead},
Error, ResponsesRoute,
types::{ResponsesCall, ResponsesOutput, ResponsesProvider, ResponsesStreamHead},
};
pub(super) async fn execute(
http: &litellm_http::Client,
auth: &litellm_auth::AuthServices,
request: ProviderResponsesRequest,
cache: Option<litellm_cache_response::ScopedCache>,
route: &ResponsesRoute,
ResponsesProvider { identity, config }: ResponsesProvider,
call: ResponsesCall,
secrets: Secrets,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
let authenticated = resolve_auth(auth, request.environment, &|_| None).await?;
let identity = litellm_host::interceptors::ProviderIdentity {
model: request.context.model.clone(),
provider: request.context.custom_llm_provider.clone(),
let lookup = |name: &str| secrets.get(name);
let environment = config.validate_environment(
litellm_http::request::string_headers("responses", call.extra_headers)?,
call.api_key.as_deref(),
&lookup,
)?;
let url = config.get_complete_url(call.api_base.as_deref(), &lookup);
let context = RequestContext {
model: identity.model.clone(),
custom_llm_provider: identity.provider.clone(),
optional_params: serde_json::Value::Object(call.optional_params.clone()),
secret_fields: Vec::new(),
api_key: match &environment.auth {
AuthScheme::Credential { secret, .. } => Some(secret.clone()),
_ => None,
},
};
let body = config.transform_responses_api_request(
&identity.model,
call.input,
call.optional_params,
)?;
let authenticated = resolve_auth(&route.auth, environment, &|_| None).await?;
let wire = interceptors
.before_provider_request(
WireRequest {
url: request.url,
url,
headers: authenticated.headers,
body: request.body,
body,
},
request.context,
context,
)
.await?;
let cache = cache.filter(|_| authenticated.signer.is_none());
let cache = route
.cache
.clone()
.filter(|_| authenticated.signer.is_none());
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
identity,
cache.as_ref().map(|_| &wire),
@ -59,9 +81,9 @@ pub(super) async fn execute(
},
wire.url,
&wire.body,
Some(request.timeout.unwrap_or(Duration::from_secs(600))),
Some(call.timeout.unwrap_or(Duration::from_secs(600))),
)?;
let response = litellm_inference::outbound::send(outbound, http)
let response = litellm_inference::outbound::send(outbound, &route.http)
.await
.map_err(network)?;
let status = response.status().as_u16();
@ -103,8 +125,7 @@ pub(super) async fn execute(
.map_err(Error::post_call)?;
let value = serde_json::from_str(&body)
.map_err(|error| Error::InvalidResponse(error.to_string().into()))?;
request
.config
config
.transform_response_api_response(value)
.map(ResponsesOutput::Complete)
.map_err(Error::from)

499
litellm-rust/crates/inference-responses/src/ids.rs Normal file
View file

@ -0,0 +1,499 @@
//! LiteLLM-managed identifiers on Responses payloads, fully compatible with `litellm/responses/utils.py`
//!
//! LiteLLM rewrites some ids in a Responses API result before handing it back, so that a later
//! request carrying those ids can be routed to the provider and deployment that produced them,
//! and then restored to what the provider originally issued. The native and Python paths serve
//! the same traffic, so an id minted by either one must decode on the other.
//!
//! Each managed id is a fixed prefix followed by standard, padded base64 of a `;`-separated
//! plaintext:
//!
//! | Id | Applied to | Decoded plaintext |
//! |---|---|---|
//! | `resp_<b64>` | the response `id` | `litellm:custom_llm_provider:{provider};model_id:{model_id};response_id:{id}` |
//! | `encitem_<b64>` | the `id` of an output item with `encrypted_content` | `litellm:model_id:{model_id};item_id:{id}` |
//! | `litellm_enc:<b64>;{content}` | that item's `encrypted_content` | `<b64>` is `model_id:{model_id}`, and the original content follows the `;` unencoded |
//! | `cntr_<b64>` | `container_id` on output items, `code_interpreter_call`, and message annotations | `litellm:custom_llm_provider:{provider};model_id:{model_id};container_id:{id}` |
//!
//! For example, OpenAI's `resp_0a1b2c3d4e5f` served without a deployment id becomes
//! `resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpOb25lO3Jlc3BvbnNlX2lkOnJlc3BfMGExYjJjM2Q0ZTVm`,
//! which decodes to `litellm:custom_llm_provider:openai;model_id:None;response_id:resp_0a1b2c3d4e5f`.
//!
//! How they fit together:
//! - `resp_` lets a follow-up's `previous_response_id` find the provider and deployment again.
//! [`encode_response`] stamps it on every response, unless the id is already managed.
//! - `encitem_` and `litellm_enc:` carry encrypted-content affinity: encrypted reasoning only
//! decrypts under the API key that produced it, so the router pins replays to that deployment.
//! They are stamped only when the router sets `encrypted_content_affinity_enabled` and a
//! deployment `model_id` is known (see [`IdScope`]). Items without an `id`, such as Codex's,
//! carry the deployment in the wrapped content alone.
//! - `cntr_` lets container calls find the provider and deployment that own a container. A
//! container id that already names a provider is left as it is.
//! - On the way in, [`restore_input`] turns `encitem_` ids and `litellm_enc:` content in the
//! request input back into the provider's own bytes before the request is sent. The routing
//! decision itself stays in Python's `EncryptedContentAffinityCheck`.
//!
//! Python's quirks are kept on purpose:
//! - A missing provider or `model_id` is written as `None` in `resp_` ids (an f-string of
//! `None`), but as an empty string in `cntr_` ids. [`decode_container_id`] reads both as absent.
//! - Decoding follows Python's non-strict `base64.b64decode`: characters outside the alphabet
//! are skipped, trailing bits are allowed, and padding stripped in transit is added back.
//! - A `resp_` id counts as managed when it decodes to at least three `;`-separated parts.
use base64::{
Engine, alphabet,
engine::{GeneralPurpose, general_purpose::PAD},
};
use litellm_llms_types::formats::responses::ResponsesApiResponse;
use serde_json::{Map, Value};
const RESPONSE_PREFIX: &str = "resp_";
const ITEM_PREFIX: &str = "encitem_";
const CONTAINER_PREFIX: &str = "cntr_";
const WRAPPED_CONTENT_PREFIX: &str = "litellm_enc:";
/// Python's non-strict `base64.b64decode`: characters outside the alphabet are skipped and
/// non-zero trailing bits are accepted
const LENIENT: GeneralPurpose = GeneralPurpose::new(
&alphabet::STANDARD,
PAD.with_decode_allow_trailing_bits(true),
);
/// Where a response came from, as encoded into the ids LiteLLM hands back
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct IdScope {
pub provider: Option<String>,
pub model_id: Option<String>,
pub encrypted_content_affinity: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DecodedContainerId {
pub provider: Option<String>,
pub model_id: Option<String>,
pub container_id: String,
}
/// Rewrites the response id, encrypted-content items and container ids the way the Python
/// Responses path does before returning to the caller. An already managed id is left untouched
pub fn encode_response(response: &mut ResponsesApiResponse, scope: &IdScope) {
if is_managed_response_id(&response.id) {
return;
}
response.id = response_id(scope, &response.id);
if let Some(model_id) = scope
.model_id
.as_deref()
.filter(|_| scope.encrypted_content_affinity)
.filter(|model_id| !model_id.is_empty())
{
response
.output
.iter_mut()
.filter_map(Value::as_object_mut)
.for_each(|item| encode_encrypted_item(item, model_id));
}
response
.output
.iter_mut()
.for_each(|item| encode_containers(item, scope));
}
/// Restores encrypted-content item ids and contents that LiteLLM encoded on an earlier response,
/// so the provider receives what it originally issued
pub fn restore_input(input: &mut Value) {
let Some(items) = input.as_array_mut() else {
return;
};
for item in items.iter_mut().filter_map(Value::as_object_mut) {
if let Some(Value::String(id)) = item.get_mut("id")
&& let Some((_, original)) = decode_encrypted_item_id(id)
{
*id = original;
}
if let Some(Value::String(content)) = item.get_mut("encrypted_content")
&& let Some(original) = unwrap_encrypted_content(content)
{
*content = original;
}
}
}
pub fn response_id(scope: &IdScope, id: &str) -> String {
let assembled = format!(
"litellm:custom_llm_provider:{};model_id:{};response_id:{id}",
python_str(scope.provider.as_deref()),
python_str(scope.model_id.as_deref()),
);
format!("{RESPONSE_PREFIX}{}", LENIENT.encode(assembled))
}
pub fn is_managed_response_id(id: &str) -> bool {
decode_text(&id.replace(RESPONSE_PREFIX, ""))
.is_some_and(|decoded| decoded.split(';').count() >= 3)
}
pub fn encrypted_item_id(model_id: &str, id: &str) -> String {
format!(
"{ITEM_PREFIX}{}",
LENIENT.encode(format!("litellm:model_id:{model_id};item_id:{id}"))
)
}
pub fn decode_encrypted_item_id(id: &str) -> Option<(String, String)> {
let decoded = decode_text(&padded(id.strip_prefix(ITEM_PREFIX)?))?;
let (model_id, item_id) = decoded.split_once(';')?;
Some((
model_id.replace("litellm:model_id:", ""),
item_id.replace("item_id:", ""),
))
}
pub fn wrap_encrypted_content(content: &str, model_id: &str) -> String {
format!(
"{WRAPPED_CONTENT_PREFIX}{};{content}",
LENIENT.encode(format!("model_id:{model_id}"))
)
}
/// The original content of a wrapped `encrypted_content`, or `None` when it is not wrapped
pub fn unwrap_encrypted_content(content: &str) -> Option<String> {
if !content.starts_with(WRAPPED_CONTENT_PREFIX) {
return None;
}
let (metadata, original) = content.split_once(';')?;
decode_text(&padded(&metadata.replace(WRAPPED_CONTENT_PREFIX, "")))?;
Some(original.to_owned())
}
pub fn container_id(scope: &IdScope, id: &str) -> String {
let assembled = format!(
"litellm:custom_llm_provider:{};model_id:{};container_id:{id}",
scope.provider.as_deref().unwrap_or_default(),
scope.model_id.as_deref().unwrap_or_default(),
);
format!("{CONTAINER_PREFIX}{}", LENIENT.encode(assembled))
}
pub fn decode_container_id(id: &str) -> DecodedContainerId {
let unmanaged = || DecodedContainerId {
provider: None,
model_id: None,
container_id: id.to_owned(),
};
if !id.starts_with(CONTAINER_PREFIX) {
return unmanaged();
}
let Some(decoded) = decode_text(&id.replace(CONTAINER_PREFIX, "")) else {
return unmanaged();
};
let parsed = decoded
.strip_prefix("litellm:custom_llm_provider:")
.and_then(|rest| rest.split_once(';'))
.and_then(|(provider, rest)| {
let (model_id, rest) = rest.strip_prefix("model_id:")?.split_once(';')?;
let container_id = rest.strip_prefix("container_id:")?;
(!container_id.is_empty() && !container_id.contains('\n')).then_some((
provider,
model_id,
container_id,
))
});
let Some((provider, model_id, container_id)) = parsed else {
return unmanaged();
};
let present = |value: &str| (!matches!(value, "" | "None")).then(|| value.to_owned());
DecodedContainerId {
provider: present(provider),
model_id: present(model_id),
container_id: container_id.to_owned(),
}
}
fn encode_encrypted_item(item: &mut Map<String, Value>, model_id: &str) {
let Some(Value::String(content)) = item.get_mut("encrypted_content") else {
return;
};
if content.is_empty() {
return;
}
*content = wrap_encrypted_content(content, model_id);
if let Some(Value::String(id)) = item.get_mut("id")
&& !id.is_empty()
{
*id = encrypted_item_id(model_id, id);
}
}
fn encode_containers(item: &mut Value, scope: &IdScope) {
let Some(item) = item.as_object_mut() else {
return;
};
encode_container_field(item, scope);
if let Some(nested) = item
.get_mut("code_interpreter_call")
.and_then(Value::as_object_mut)
{
encode_container_field(nested, scope);
}
if item.get("type").and_then(Value::as_str) != Some("message") {
return;
}
let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else {
return;
};
content
.iter_mut()
.filter_map(|part| part.get_mut("annotations").and_then(Value::as_array_mut))
.flatten()
.for_each(|annotation| encode_containers(annotation, scope));
}
fn encode_container_field(item: &mut Map<String, Value>, scope: &IdScope) {
if let Some(Value::String(id)) = item.get_mut("container_id")
&& decode_container_id(id).provider.is_none()
{
*id = container_id(scope, id);
}
}
fn python_str(value: Option<&str>) -> &str {
value.unwrap_or("None")
}
fn padded(value: &str) -> String {
let missing = value.len() % 4;
let padding = if missing == 0 { 0 } else { 4 - missing };
format!("{value}{}", "=".repeat(padding))
}
fn decode_text(value: &str) -> Option<String> {
let alphabet: String = value
.chars()
.filter(|character| {
character.is_ascii_alphanumeric() || matches!(character, '+' | '/' | '=')
})
.collect();
String::from_utf8(LENIENT.decode(alphabet).ok()?).ok()
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
fn scope(provider: Option<&str>, model_id: Option<&str>, affinity: bool) -> IdScope {
IdScope {
provider: provider.map(str::to_owned),
model_id: model_id.map(str::to_owned),
encrypted_content_affinity: affinity,
}
}
fn response(id: &str, output: Value) -> ResponsesApiResponse {
serde_json::from_value(json!({"id": id, "model": "gpt-6.1-sol", "output": output})).unwrap()
}
#[rstest]
#[case::openai_without_model_id(
Some("openai"),
None,
"resp_0a1b2c3d4e5f",
"resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpOb25lO3Jlc3BvbnNlX2lkOnJlc3BfMGExYjJjM2Q0ZTVm"
)]
#[case::openrouter_with_model_id(
Some("openrouter"),
Some("deploy-1"),
"gen-123",
"resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5yb3V0ZXI7bW9kZWxfaWQ6ZGVwbG95LTE7cmVzcG9uc2VfaWQ6Z2VuLTEyMw=="
)]
fn response_ids_match_python(
#[case] provider: Option<&str>,
#[case] model_id: Option<&str>,
#[case] id: &str,
#[case] expected: &str,
) {
assert_eq!(response_id(&scope(provider, model_id, false), id), expected);
}
#[rstest]
#[case::raw_short_id("resp_0a1b2c3d4e5f", false)]
#[case::raw_long_id("resp_68a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8a9b0c1d2e3", false)]
#[case::two_segments("resp_YTti", false)]
#[case::three_segments("resp_YTtiO2M=", true)]
#[case::managed(
"resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpOb25lO3Jlc3BvbnNlX2lkOng=",
true
)]
fn managed_response_ids_match_python(#[case] id: &str, #[case] expected: bool) {
assert_eq!(is_managed_response_id(id), expected);
}
#[test]
fn encrypted_items_match_python() {
assert_eq!(
encrypted_item_id("deploy-1", "rs_0a1b"),
"encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI="
);
assert_eq!(
wrap_encrypted_content("gAAAA", "deploy-1"),
"litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA"
);
}
#[rstest]
#[case::padded("encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=", Some(("deploy-1", "rs_0a1b")))]
#[case::padding_stripped_in_transit("encitem_bGl0ZWxsbTptb2RlbF9pZDptO2l0ZW1faWQ6cnNfMQ", Some(("m", "rs_1")))]
#[case::not_base64("encitem_notbase64!", None)]
#[case::unmanaged("rs_0a1b", None)]
fn encrypted_item_ids_decode_like_python(
#[case] id: &str,
#[case] expected: Option<(&str, &str)>,
) {
assert_eq!(
decode_encrypted_item_id(id),
expected.map(|(model_id, item_id)| (model_id.to_owned(), item_id.to_owned()))
);
}
#[rstest]
#[case::unpadded_metadata("litellm_enc:bW9kZWxfaWQ6bQ;gAAAA", Some("gAAAA"))]
#[case::no_separator("litellm_enc:nosemicolon", None)]
#[case::unwrapped("gAAAA", None)]
fn encrypted_content_unwraps_like_python(
#[case] content: &str,
#[case] expected: Option<&str>,
) {
assert_eq!(unwrap_encrypted_content(content).as_deref(), expected);
}
#[rstest]
#[case::with_model_id(
Some("deploy-1"),
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpkZXBsb3ktMTtjb250YWluZXJfaWQ6Y250cl9hYmM="
)]
#[case::without_model_id(
None,
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDo7Y29udGFpbmVyX2lkOmNudHJfYWJj"
)]
fn container_ids_match_python(#[case] model_id: Option<&str>, #[case] expected: &str) {
assert_eq!(
container_id(&scope(Some("openai"), model_id, false), "cntr_abc"),
expected
);
}
#[rstest]
#[case::raw("cntr_abc", None, None, "cntr_abc")]
#[case::managed(
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDptO2NvbnRhaW5lcl9pZDpjbnRyXzE=",
Some("openai"),
Some("m"),
"cntr_1"
)]
#[case::empty_segments(
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOjttb2RlbF9pZDo7Y29udGFpbmVyX2lkOmNudHJfMQ==",
None,
None,
"cntr_1"
)]
#[case::python_none_provider(
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOk5vbmU7bW9kZWxfaWQ6O2NvbnRhaW5lcl9pZDp4",
None,
None,
"x"
)]
fn container_ids_decode_like_python(
#[case] id: &str,
#[case] provider: Option<&str>,
#[case] model_id: Option<&str>,
#[case] container_id: &str,
) {
assert_eq!(
decode_container_id(id),
DecodedContainerId {
provider: provider.map(str::to_owned),
model_id: model_id.map(str::to_owned),
container_id: container_id.to_owned(),
}
);
}
#[test]
fn encodes_response_items_under_affinity() {
let mut encoded = response(
"resp_1",
json!([
{"type": "reasoning", "id": "rs_0a1b", "encrypted_content": "gAAAA", "summary": []},
{"type": "reasoning", "encrypted_content": "gBBBB", "summary": []},
{"type": "code_interpreter_call", "id": "ci_1", "container_id": "cntr_abc"},
{"type": "message", "id": "msg_1", "content": [{"type": "output_text", "text": "", "annotations": [
{"type": "container_file_citation", "container_id": "cntr_abc"}
]}]}
]),
);
let scope = scope(Some("openai"), Some("deploy-1"), true);
encode_response(&mut encoded, &scope);
let container = container_id(&scope, "cntr_abc");
assert_eq!(encoded.id, response_id(&scope, "resp_1"));
assert_eq!(
Value::Array(encoded.output),
json!([
{
"type": "reasoning",
"id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=",
"encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA",
"summary": []
},
{"type": "reasoning", "encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gBBBB", "summary": []},
{"type": "code_interpreter_call", "id": "ci_1", "container_id": container},
{"type": "message", "id": "msg_1", "content": [{"type": "output_text", "text": "", "annotations": [
{"type": "container_file_citation", "container_id": container}
]}]}
])
);
}
#[rstest]
#[case::affinity_disabled(scope(Some("openai"), Some("deploy-1"), false))]
#[case::no_model_id(scope(Some("openai"), None, true))]
fn leaves_encrypted_items_without_affinity(#[case] scope: IdScope) {
let output = json!([{"type": "reasoning", "id": "rs_1", "encrypted_content": "gAAAA", "summary": []}]);
let mut encoded = response("resp_1", output.clone());
encode_response(&mut encoded, &scope);
assert_eq!(Value::Array(encoded.output), output);
}
#[test]
fn leaves_managed_responses_untouched() {
let scope = scope(Some("openai"), Some("deploy-1"), true);
let id = response_id(&scope, "resp_1");
let output = json!([{"type": "code_interpreter_call", "container_id": "cntr_abc"}]);
let mut encoded = response(&id, output.clone());
encode_response(&mut encoded, &scope);
assert_eq!((encoded.id, Value::Array(encoded.output)), (id, output));
}
#[test]
fn restores_encrypted_items_in_input() {
let mut input = json!([
{
"type": "reasoning",
"id": "encitem_bGl0ZWxsbTptb2RlbF9pZDpkZXBsb3ktMTtpdGVtX2lkOnJzXzBhMWI=",
"encrypted_content": "litellm_enc:bW9kZWxfaWQ6ZGVwbG95LTE=;gAAAA",
"summary": []
},
{"role": "user", "content": "hi"}
]);
restore_input(&mut input);
assert_eq!(
input,
json!([
{"type": "reasoning", "id": "rs_0a1b", "encrypted_content": "gAAAA", "summary": []},
{"role": "user", "content": "hi"}
])
);
}
}

111
litellm-rust/crates/inference-responses/src/lib.rs
View file

@ -1,19 +1,36 @@
use litellm_host::observation::ObservationSender;
pub use litellm_inference::RouteError as Error;
pub use litellm_inference::Error;
pub mod ids;
pub mod route;
pub mod types;
pub mod websocket;
mod handler;
mod prepare;
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, LlmProviders, get_custom_llm_provider,
};
use litellm_host::interceptors::Interceptors;
use litellm_host::interceptors::ProviderIdentity;
use litellm_llms::{
base_llm::responses::transformation::BaseResponsesApiConfig,
openai::responses::transformation::OpenAiResponsesApiConfig,
openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
};
use litellm_secrets::source::SecretSource;
use types::{ResponsesCall, ResponsesOutput};
use types::{ResponsesCall, ResponsesOutput, ResponsesProvider};
fn responses_provider(provider: LlmProviders) -> Option<&'static dyn BaseResponsesApiConfig> {
match provider {
LlmProviders::Openai => Some(&OpenAiResponsesApiConfig),
LlmProviders::Openrouter => Some(&OPENROUTER_RESPONSES_CONFIG),
_ => None,
}
}
#[derive(Clone)]
pub struct ResponsesRoute {
@ -77,34 +94,68 @@ impl ResponsesRoute {
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
litellm_inference::diagnostic::call(async {
self.run_provider(call, cache_options, interceptors, observers)
.await
let explicit = call
.custom_llm_provider
.as_deref()
.filter(|provider| !provider.is_empty());
let (provider_name, resolved_model) =
match get_custom_llm_provider(&call.model, explicit) {
Some(CustomLlmProvider {
model,
custom_llm_provider,
}) => (custom_llm_provider, model),
None => ("openai", call.model.as_str()),
};
if resolved_model.is_empty() {
return Err(Error::InvalidProvider(call.model));
}
let provider = provider_name
.parse::<LlmProviders>()
.map_err(|_| Error::Unsupported("native HTTP responses provider"))?;
let config = responses_provider(provider)
.ok_or(Error::Unsupported("native HTTP responses provider"))?;
let wire_model = match provider {
LlmProviders::Openai if resolved_model.contains('/') => {
return Err(Error::InvalidProvider(call.model));
}
LlmProviders::Openai => resolved_model.into(),
LlmProviders::Openrouter if explicit.is_some() && !resolved_model.contains('/') => {
format!("openrouter/{resolved_model}")
}
LlmProviders::Openrouter => resolved_model.into(),
_ => return Err(Error::Unsupported("native HTTP responses provider")),
};
let provider = ResponsesProvider {
identity: ProviderIdentity {
model: wire_model,
provider: provider_name.into(),
},
config,
};
litellm_inference::diagnostic::provider(
&provider.identity.model,
&provider.identity.provider,
);
let secrets = self
.secrets
.resolve(
provider
.config
.secret_names(call.api_key.as_deref(), call.api_base.as_deref()),
)
.await?;
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
Box::pin(handler::execute(
self,
provider,
call,
secrets,
cache_options,
interceptors,
observers,
));
execute.await
})
.await
}
async fn run_provider(
&self,
call: ResponsesCall,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
litellm_inference::diagnostic::provider(
&request.context.model,
&request.context.custom_llm_provider,
);
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
Box::pin(handler::execute(
&self.http,
&self.auth,
request,
self.cache.clone(),
cache_options,
interceptors,
observers,
));
execute.await
}
}

162
litellm-rust/crates/inference-responses/src/prepare.rs
View file

@ -1,162 +0,0 @@
use litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, LlmProviders, get_custom_llm_provider,
};
use litellm_host::interceptors::{ProviderIdentity, RequestContext};
use litellm_llms::{
base_llm::responses::transformation::BaseResponsesApiConfig,
openai::responses::transformation::OpenAiResponsesApiConfig,
openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
};
use litellm_secrets::source::SecretSource;
use serde_json::Value;
use super::{
Error,
types::{ProviderResponsesRequest, ResponsesCall},
};
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub(super) async fn prepare(
call: ResponsesCall,
secrets: &dyn SecretSource,
) -> Result<ProviderResponsesRequest, Error> {
let ResponsesProvider { identity, config } =
resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
let provider = identity.provider.as_str();
let model = identity.model.as_str();
let snapshot = secrets
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
.await?;
let lookup = |name: &str| snapshot.get(name);
let environment = config.validate_environment(
litellm_http::request::string_headers("responses", call.extra_headers)?,
call.api_key.as_deref(),
&lookup,
)?;
let context = RequestContext {
model: model.into(),
custom_llm_provider: provider.into(),
optional_params: Value::Object(call.optional_params.clone()),
secret_fields: Vec::new(),
api_key: match &environment.auth {
litellm_llms::base_llm::auth::AuthScheme::Credential { secret, .. } => {
Some(secret.clone())
}
_ => None,
},
};
let body = config.transform_responses_api_request(model, call.input, call.optional_params)?;
Ok(ProviderResponsesRequest {
url: config.get_complete_url(call.api_base.as_deref(), &lookup),
config,
environment,
body,
context,
timeout: call.timeout,
})
}
pub(super) struct ResponsesProvider {
pub identity: ProviderIdentity,
pub config: &'static dyn BaseResponsesApiConfig,
}
/// Python's `get_llm_provider`: a bare model is OpenAI's, and with an explicit OpenRouter
/// provider a model without a vendor segment keeps the `openrouter/` prefix as its wire id.
pub(super) fn resolve_provider(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<ResponsesProvider, Error> {
let explicit = custom_llm_provider.filter(|provider| !provider.is_empty());
let (provider, resolved) = match get_custom_llm_provider(model, explicit) {
Some(CustomLlmProvider {
model,
custom_llm_provider,
}) => (custom_llm_provider, model),
None => ("openai", model),
};
if resolved.is_empty() {
return Err(Error::InvalidProvider(model.into()));
}
let (config, wire_model): (&'static dyn BaseResponsesApiConfig, String) =
match provider.parse::<LlmProviders>() {
Ok(LlmProviders::Openai) if resolved.contains('/') => {
return Err(Error::InvalidProvider(model.into()));
}
Ok(LlmProviders::Openai) => (&OpenAiResponsesApiConfig, resolved.into()),
Ok(LlmProviders::Openrouter) if explicit.is_some() && !resolved.contains('/') => (
&OPENROUTER_RESPONSES_CONFIG,
format!("openrouter/{resolved}"),
),
Ok(LlmProviders::Openrouter) => (&OPENROUTER_RESPONSES_CONFIG, resolved.into()),
_ => return Err(Error::Unsupported("native HTTP responses provider")),
};
Ok(ResponsesProvider {
identity: ProviderIdentity {
model: wire_model,
provider: provider.into(),
},
config,
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
/// Pinned against Python's `get_llm_provider` on 2026-10-11.
#[rstest]
#[case::bare("gpt-5", None, "openai", "gpt-5")]
#[case::openai_prefix("openai/gpt-5", None, "openai", "gpt-5")]
#[case::openai_explicit("gpt-5", Some("openai"), "openai", "gpt-5")]
#[case::openrouter_prefix(
"openrouter/anthropic/claude",
None,
"openrouter",
"anthropic/claude"
)]
#[case::openrouter_explicit(
"openrouter/anthropic/claude",
Some("openrouter"),
"openrouter",
"anthropic/claude"
)]
#[case::openrouter_explicit_unprefixed(
"anthropic/claude",
Some("openrouter"),
"openrouter",
"anthropic/claude"
)]
#[case::openrouter_native_id(
"openrouter/auto",
Some("openrouter"),
"openrouter",
"openrouter/auto"
)]
#[case::openrouter_native_id_implicit("openrouter/auto", None, "openrouter", "auto")]
fn providers_and_wire_models_match_python(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] provider: &str,
#[case] wire_model: &str,
) {
let resolved = resolve_provider(model, custom_llm_provider).unwrap();
assert_eq!(resolved.identity.provider, provider);
assert_eq!(resolved.identity.model, wire_model);
}
#[rstest]
#[case::other_prefix("other/model", None)]
#[case::other_explicit("model", Some("other"))]
#[case::openai_with_vendor_segment("openai/other/model", None)]
#[case::empty_model("openai/", None)]
#[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
fn everything_else_is_rejected(#[case] model: &str, #[case] custom_llm_provider: Option<&str>) {
assert!(matches!(
resolve_provider(model, custom_llm_provider),
Err(Error::InvalidProvider(_) | Error::Unsupported(_))
));
}
}

19
litellm-rust/crates/inference-responses/src/types.rs
View file

@ -2,9 +2,8 @@ use std::time::Duration;
use bytes::Bytes;
use litellm_host::call::CallOutput;
use litellm_llms::base_llm::{
auth::ValidatedEnvironment, responses::transformation::BaseResponsesApiConfig,
};
use litellm_host::interceptors::ProviderIdentity;
use litellm_llms::base_llm::responses::transformation::BaseResponsesApiConfig;
use litellm_llms_types::formats::responses::ResponsesApiResponse;
use serde_json::{Map, Value};
@ -21,17 +20,13 @@ pub struct ResponsesCall {
pub timeout: Option<Duration>,
}
pub(crate) struct ResponsesProvider {
pub identity: ProviderIdentity,
pub config: &'static dyn BaseResponsesApiConfig,
}
pub struct ResponsesStreamHead {
pub headers: Vec<(String, String)>,
}
pub type ResponsesOutput = CallOutput<ResponsesApiResponse, ResponsesStreamHead, Bytes, Error>;
pub(super) struct ProviderResponsesRequest {
pub config: &'static dyn BaseResponsesApiConfig,
pub environment: ValidatedEnvironment,
pub url: String,
pub body: Value,
pub context: litellm_host::interceptors::RequestContext,
pub timeout: Option<Duration>,
}

10
litellm-rust/crates/inference-responses/tests/caching.rs
View file

@ -18,7 +18,7 @@ use litellm_host::interceptors::{
WireRequest,
};
use litellm_inference::{
RouteError,
Error,
caching::{CacheRequest, execute_unary},
};
use rstest::{fixture, rstest};
@ -124,12 +124,12 @@ struct ChangingHooks {
facts: std::sync::Mutex<Vec<ExecutionFacts>>,
}
impl Interceptors<RouteError> for ChangingHooks {
impl Interceptors<Error> for ChangingHooks {
async fn before_provider_request(
&self,
mut wire: WireRequest,
_: RequestContext,
) -> Result<WireRequest, RouteError> {
) -> Result<WireRequest, Error> {
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if self.rewrite {
wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 });
@ -137,11 +137,11 @@ impl Interceptors<RouteError> for ChangingHooks {
Ok(wire)
}
async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> {
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Ok(())
}
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> {
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
self.facts.lock().unwrap().push(facts);
Ok(())
}

105
litellm-rust/crates/inference-responses/tests/responses.rs
View file

@ -7,6 +7,7 @@ use std::sync::Arc;
use futures_util::TryStreamExt;
use litellm_host::{call::HostedCompletion, lifecycle::CallEvent};
use litellm_inference_responses::{
Error,
route::Responses,
types::{ResponsesCall, ResponsesOutput},
};
@ -32,6 +33,110 @@ fn call() -> ResponsesCall {
}
}
#[rstest]
#[case::bare_openai("gpt-5", None, "gpt-5")]
#[case::openai_prefix("openai/gpt-5", None, "gpt-5")]
#[case::empty_explicit_provider("openai/gpt-5", Some(""), "gpt-5")]
#[case::explicit_openai("gpt-5", Some("openai"), "gpt-5")]
#[case::openrouter_prefix("openrouter/anthropic/claude", None, "anthropic/claude")]
#[case::explicit_openrouter_prefix(
"openrouter/anthropic/claude",
Some("openrouter"),
"anthropic/claude"
)]
#[case::explicit_openrouter_unprefixed("anthropic/claude", Some("openrouter"), "anthropic/claude")]
#[case::explicit_openrouter_native_id("auto", Some("openrouter"), "openrouter/auto")]
#[case::implicit_openrouter_native_id("openrouter/auto", None, "auto")]
#[tokio::test]
async fn provider_resolution_preserves_the_wire_model(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] expected_wire_model: &str,
) {
let response = json!({
"id": "response-1",
"model": expected_wire_model,
"output": [{"type": "message", "content": []}],
"usage": {"total_tokens": 7}
});
let upstream = upstream([json_response(response)]).await;
responses_route(no_secrets())
.execute(
ResponsesCall {
model: model.into(),
api_key: Some("test-key".into()),
api_base: Some(upstream.uri()),
custom_llm_provider: custom_llm_provider.map(str::to_owned),
..call()
},
&(),
None,
)
.await
.unwrap();
assert_eq!(
only_request(&upstream).await.json()["model"],
expected_wire_model
);
}
#[rstest]
#[case::explicit_unknown_provider("test-model", Some("unknown-provider"))]
#[case::unsupported_responses_provider("test-model", Some("anthropic"))]
#[tokio::test]
async fn unsupported_providers_fail_before_transport(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
) {
let upstream = upstream([]).await;
let result = responses_route(no_secrets())
.execute(
ResponsesCall {
model: model.into(),
api_base: Some(upstream.uri()),
custom_llm_provider: custom_llm_provider.map(str::to_owned),
..call()
},
&(),
None,
)
.await;
assert!(matches!(
result,
Err(Error::Unsupported("native HTTP responses provider"))
));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::openai_with_vendor_segment("openai/other/model", None)]
#[case::empty_model("openai/", None)]
#[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
#[tokio::test]
async fn invalid_provider_models_fail_before_transport(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
) {
let upstream = upstream([]).await;
let result = responses_route(no_secrets())
.execute(
ResponsesCall {
model: model.into(),
api_base: Some(upstream.uri()),
custom_llm_provider: custom_llm_provider.map(str::to_owned),
..call()
},
&(),
None,
)
.await;
assert!(matches!(result, Err(Error::InvalidProvider(_))));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::direct(false)]
#[case::hosted(true)]

60
litellm-rust/crates/inference-transcription/src/handler.rs
View file

@ -1,28 +1,57 @@
use std::time::Duration;
use litellm_http::{Client, request::truncate_error_body};
use litellm_llms::base_llm::auth::resolve_auth;
use litellm_http::{
Client,
request::{string_headers, truncate_error_body, with_default_headers},
};
use litellm_llms::base_llm::{
audio_transcription::transformation::BaseAudioTranscriptionConfig,
auth::{ValidatedEnvironment, resolve_auth},
};
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::Error;
use crate::{
constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, types::ProviderAudioTranscriptionRequest,
};
use crate::{constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, types::AudioTranscriptionRequest};
use litellm_inference::provider::ResolvedProvider;
pub async fn execute_audio_transcription_provider_call(
pub(super) async fn execute(
http: &Client,
auth: &litellm_auth::AuthServices,
request: ProviderAudioTranscriptionRequest,
config: &'static dyn BaseAudioTranscriptionConfig,
provider: ResolvedProvider<'_>,
call: AudioTranscriptionRequest<'_>,
secrets: Secrets,
) -> Result<Value, Error> {
let env_lookup = |key: &str| request.secrets.get(key);
let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?;
let env_lookup = |key: &str| secrets.get(key);
let forwarded = string_headers("audio transcription", call.extra_headers)?;
let validated = config.validate_environment(
forwarded,
provider.model,
&call.optional_params,
&env_lookup,
)?;
let environment = ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
};
let url = config.get_complete_url(
call.api_base,
provider.model,
&call.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&call.optional_params);
let body = config
.transform_audio_transcription_request(provider.model, call.audio, filtered_params)?
.body;
let authenticated = resolve_auth(auth, environment, &env_lookup).await?;
let outbound = litellm_inference::outbound::outbound_request(
authenticated,
request.url.clone(),
&request.body,
url,
&body,
Some(
request
.timeout
call.timeout
.unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)),
),
)?;
@ -47,8 +76,7 @@ pub async fn execute_audio_transcription_provider_call(
error,
))
})?;
Ok(request
.config
.transform_audio_transcription_response(&request.model, response_json)?
Ok(config
.transform_audio_transcription_response(provider.model, response_json)?
.into_json())
}

45
litellm-rust/crates/inference-transcription/src/lib.rs
View file

@ -1,17 +1,30 @@
pub mod types;
pub use litellm_inference::RouteError as Error;
pub use litellm_inference::Error;
mod constants;
mod handler;
mod prepare;
pub use handler::execute_audio_transcription_provider_call;
use litellm_auth::AuthServices;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_inference::provider::resolve_llm_provider;
use litellm_llms::{
base_llm::audio_transcription::transformation::BaseAudioTranscriptionConfig,
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
};
use litellm_secrets::source::SecretSource;
pub use prepare::prepare_audio_transcription_provider_call;
use serde_json::Value;
use std::sync::Arc;
use crate::types::AudioTranscriptionRequest;
fn audio_transcription_provider(
provider: LlmProviders,
) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
match provider {
LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
#[derive(Clone)]
pub struct AudioTranscriptionRoute {
http: litellm_http::Client,
@ -32,6 +45,10 @@ impl AudioTranscriptionRoute {
}
}
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
self.run(request).await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "audio_transcription",
model = request.model,
@ -40,13 +57,23 @@ impl AudioTranscriptionRoute {
stream = false,
outcome
))]
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
async fn run(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
litellm_inference::diagnostic::unary(async {
let request =
prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?;
litellm_inference::diagnostic::provider(&request.model, &request.custom_llm_provider);
let provider = resolve_llm_provider(
request.model,
request.custom_llm_provider,
"audio transcription",
)?;
let config = audio_transcription_provider(provider.provider).ok_or_else(|| {
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
})?;
litellm_inference::diagnostic::provider(
provider.model,
<&'static str>::from(provider.provider),
);
let secrets = self.secrets.resolve(&config.secret_names()).await?;
let execute: futures_util::future::BoxFuture<'_, Result<Value, Error>> = Box::pin(
execute_audio_transcription_provider_call(&self.http, &self.auth, request),
handler::execute(&self.http, &self.auth, config, provider, request, secrets),
);
execute.await
})

65
litellm-rust/crates/inference-transcription/src/prepare.rs
View file

@ -1,65 +0,0 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers;
use litellm_http::request::with_default_headers;
use litellm_llms::{
base_llm::{
audio_transcription::transformation::BaseAudioTranscriptionConfig,
auth::ValidatedEnvironment,
},
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
};
use litellm_secrets::source::SecretSource;
use super::Error;
use crate::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use litellm_inference::provider::resolve_llm_provider;
fn provider_config(provider: LlmProviders) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
match provider {
LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub async fn prepare_audio_transcription_provider_call(
request: AudioTranscriptionRequest<'_>,
secrets: &dyn SecretSource,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let provider_info = resolve_llm_provider(
request.model,
request.custom_llm_provider,
"audio transcription",
)?;
let model = provider_info.model.to_string();
let config = provider_config(provider_info.provider)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?;
let snapshot = secrets.resolve(&config.secret_names()).await?;
let env_lookup = |key: &str| snapshot.get(key);
let forwarded = string_headers("audio transcription", request.extra_headers)?;
let validated =
config.validate_environment(forwarded, &model, &request.optional_params, &env_lookup)?;
let environment = ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
};
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_audio_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: <&str>::from(provider_info.provider).to_string(),
config,
url,
body: transformed.body,
environment,
secrets: snapshot,
timeout: request.timeout,
})
}

38
litellm-rust/crates/inference-transcription/src/types.rs
View file

@ -1,9 +1,5 @@
use litellm_secrets::source::Secrets;
use std::time::Duration;
use litellm_llms::base_llm::{
audio_transcription::transformation::BaseAudioTranscriptionConfig, auth::ValidatedEnvironment,
};
use serde_json::{Map, Value};
pub struct AudioTranscriptionRequest<'a> {
@ -16,37 +12,3 @@ pub struct AudioTranscriptionRequest<'a> {
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}
#[derive(Clone)]
pub struct ProviderAudioTranscriptionRequest {
pub model: String,
pub custom_llm_provider: String,
pub config: &'static dyn BaseAudioTranscriptionConfig,
pub url: String,
pub body: Value,
pub environment: ValidatedEnvironment,
pub secrets: Secrets,
pub timeout: Option<Duration>,
}
impl ProviderAudioTranscriptionRequest {
pub fn model(&self) -> &str {
&self.model
}
pub fn custom_llm_provider(&self) -> &str {
&self.custom_llm_provider
}
pub fn url(&self) -> &str {
&self.url
}
pub fn body(&self) -> &Value {
&self.body
}
pub fn with_body(self, body: Value) -> Self {
Self { body, ..self }
}
}

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