scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-25 05:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 125

Merge branch 'main' into fix-chat-cli-utf8

Browse source 
# Conflicts:
#	tests/unit/proxy/client/cli/test_chat_commands.py
This commit is contained in:
SWAPI03 2026-10-18 09:53:59 +05:30
commit 7c8da00db3
1999 changed files with 213390 additions and 19853 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
.circleci/config.yml
View file

@ -1293,6 +1293,19 @@ jobs:
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'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- when:
condition:
equal: [sdk, << parameters.suite >>]
steps:
- run:
name: Install redis-server 7.2.16
command: |
mkdir -p /tmp/redis-bin
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
docker cp "${cid}:/usr/local/bin/redis-server" /tmp/redis-bin/redis-server
docker rm "$cid"
/tmp/redis-bin/redis-server --version | grep -q 'v=7.2.16'
echo 'export PATH="/tmp/redis-bin:$PATH"' >> "$BASH_ENV"
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"

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" \

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"

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

@ -21,7 +21,7 @@ jobs:
pull-requests: read
actions: read
runs-on: ubuntu-latest
timeout-minutes: 15
timeout-minutes: 20
services:
redis:
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
@ -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:

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

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

@ -346,6 +346,7 @@ jobs:
tests/unit/llms/tinyfish
tests/unit/llms/together_ai
tests/unit/llms/tool_loop
tests/unit/llms/topaz
tests/unit/llms/triton
tests/unit/llms/v0
tests/unit/llms/valkey
@ -455,11 +456,11 @@ jobs:
tests/unit/telemetry tests/unit/proxy/telemetry
tests/unit/tracing
python-version: "3.12"
workers: 2
workers: 4
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
dist: loadscope
dist: load
- shard: proxy-auth
test-path: |-
@ -564,7 +565,6 @@ jobs:
- shard: mcp-elicitation
test-path: |-
--override-ini=pythonpath=tests
tests/unit/experimental_mcp_client/test_mcp_client.py
tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
tests/unit/proxy/_experimental/mcp_server/test_interactions.py
@ -572,9 +572,8 @@ jobs:
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
tests/unit/proxy/_experimental/mcp_server/test_operations.py
tests/integration/mcp/test_interactions.py
python-version: "3.12"
workers: 2
workers: 4
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
@ -659,6 +658,7 @@ jobs:
- shard: proxy-extras
test-path: >-
tests/unit/litellm_proxy_extras
tests/unit/db_scripts
python-version: "3.12"
workers: 2
reruns: 0
@ -697,7 +697,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
@ -1452,7 +1451,6 @@ jobs:
tests/unit/caching/test_redis_connection_pool.py \
tests/unit/caching/test_redis_cluster_cache.py \
tests/unit/caching/test_evicted_client_closer.py \
tests/integration/sdk/test_redis_cluster_iam_auth.py \
--tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
--cov=./litellm --cov-report=xml:coverage.xml

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

2
Makefile
View file

@ -290,7 +290,7 @@ test-rust-extension:
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
--mypy-config-file tests/unit/rust_bridge/stubtest.ini \
litellm.rust_bridge._native && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
rust-sqlx-prepare:

1
README.md
View file

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

32
backend/Dockerfile
View file

@ -4,6 +4,33 @@ ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2
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
@ -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 \

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
}

45
ci_cd/cost_map_guard.py
View file

@ -1,11 +1,11 @@
"""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 or the legacy schema
path and may only add or update models.
"""
from __future__ import annotations
@ -19,13 +19,12 @@ 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)
LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
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 +34,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 +44,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 +69,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()
@ -135,7 +110,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS
if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),

409
ci_cd/generate_model_prices_schema.py
View file

@ -1,409 +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.",
},
"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,
"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())

81
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json
View file

@ -6563,6 +6563,87 @@
],
"title": "litellm_project_model_rate_limit_used_metric",
"type": "timeseries"
},
{
"collapsed": false,
"gridPos": {
"h": 1,
"w": 24,
"x": 0,
"y": 447
},
"id": 115,
"panels": [],
"title": "Agent tracing export (needs LITELLM_LENS_URL)",
"type": "row"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"description": "Request records the Lens exporter delivered to the Lens service against the records it dropped before the service accepted them",
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"drawStyle": "line",
"fillOpacity": 10,
"lineWidth": 1,
"showPoints": "never",
"spanNulls": false
},
"unit": "short"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 24,
"x": 0,
"y": 448
},
"id": 116,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_lens_export_rows_written_total[$__rate_interval]))",
"legendFormat": "written",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_lens_export_rows_dropped_total[$__rate_interval]))",
"legendFormat": "dropped",
"range": true,
"refId": "B"
}
],
"title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate",
"type": "timeseries"
}
],
"preload": false,

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

@ -1,10 +1,10 @@
# LiteLLM All Prometheus Metrics dashboard
Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
Every `litellm_*` metric family the proxy can expose on `/metrics` (143 families across 98 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes), project model rate limits, and the agent tracing export to Lens. Panel titles are the metric names so you can grep the JSON for the metric you care about
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
The first eleven rows and the project model rate limits row need only `callbacks: ["prometheus"]`. The other rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, the agent tracing export row needs `LITELLM_LENS_URL`, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected
## Pre-requisites

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())

2
docker/docker-compose.quickstart.yml
View file

@ -11,7 +11,7 @@
# beyond local evaluation, pin the image to a specific release tag.
services:
litellm:
image: docker.litellm.ai/berriai/litellm:main-stable
image: docker.litellm.ai/berriai/litellm:latest
ports:
# LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart
# script sets it to "127.0.0.1:" so new installs listen on this machine only.

10
enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py
View file

@ -1,14 +1,16 @@
from typing import Any, Optional
from collections.abc import Awaitable, Callable
from fastapi import Request
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy._types import ProxyException
async def enterprise_custom_auth(
request: Request, api_key: str, user_custom_auth: Optional[Any]
) -> Optional[UserAPIKeyAuth]:
request: Request,
api_key: str,
user_custom_auth: Callable[[Request, str], Awaitable[object]] | None,
) -> object:
from litellm_enterprise.proxy.proxy_server import custom_auth_settings
if user_custom_auth is None:

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

@ -2,27 +2,35 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
from collections.abc import Sequence
import asyncio
import json
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
)
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
CLI_SESSION_KEY_PREFIX,
BATCH_LINE_ITEMS_MAX_PART_BYTES,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
)
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.batches.batch_line_items import BatchFileProvider
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
@ -47,29 +55,6 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
)
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@property
def org_id(self) -> str | None: ...
@property
def api_key(self) -> str | None: ...
@property
def team_id(self) -> str | None: ...
class _ReadableFileContent(Protocol):
async def read(self) -> bytes: ...
@ -87,14 +72,10 @@ async def _file_content_bytes(file_content: object) -> bytes:
return cast(bytes, file_content)
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
return ManagedObjectRepository(prisma_client).table
def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]":
return UserRepository(prisma_client).table
def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
return VerificationTokenRepository(prisma_client).table
@ -103,6 +84,11 @@ 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,
@ -118,10 +104,15 @@ class CheckBatchCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
# Cached after the first poll cycle. Once we know the column is absent we skip
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: 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()
self._poll_idle.set()
@staticmethod
def _is_missing_batch_processed_column_error(err: Exception) -> bool:
@ -149,133 +140,6 @@ class CheckBatchCost:
return
self.batch_processed_support_confirmed = True
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
"""
Look up user email and key alias by user_id for enriching the S3 callback metadata.
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
Returns an empty dict when user_id is None: batches created by a team or service
account key carry no user id, and find_unique(where={"user_id": None}) raises.
"""
if not user_id:
return {}
try:
user_row: prisma_models.LiteLLM_UserTable | None = await _user_table(self.prisma_client).find_unique(
where={"user_id": user_id}
)
if user_row is None:
return {}
return {
"user_api_key_user_email": getattr(user_row, "user_email", None),
"user_api_key_alias": getattr(user_row, "user_alias", None),
}
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
return {}
async def _get_key_alias(self, batch_id: str, api_key: str | None, created_by: str | None) -> str | None:
"""Resolve the creating virtual key's alias from its hashed token."""
if not api_key:
return None
if created_by and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}":
return api_key
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
self.prisma_client
).find_unique(where={"token": api_key})
return getattr(key_row, "key_alias", None) if key_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
return None
async def _get_team_alias(self, team_id: str | None) -> str | None:
"""Resolve a team's alias from its id."""
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
where={"team_id": team_id}
)
return getattr(team_row, "team_alias", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
org_id: Final = getattr(job, "org_id", None)
if org_id:
return org_id
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
self.prisma_client
).find_unique(where={"token": api_key})
key_org_id: Final = (
cast(str | None, getattr(key_row, "organization_id", None))
if key_row is not None
else None
)
if key_org_id:
return key_org_id
except Exception as e:
verbose_proxy_logger.error(
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
f"still trying the team's: {e}"
)
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
where={"team_id": team_id}
)
return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
return None
async def _build_creator_attribution_metadata(
self, job: "_ManagedObjectRow", batch_id: str
) -> dict[str, object]:
"""
Rebuild the spend-tracking metadata for the key, team, and tags that created the
batch so the batch-cost spend log is attributed the same way a non-batch request
is. Rows created before api_key and request_tags were persisted carry only
created_by and team_id, and fall back to those. A named creating key owns
user_api_key_alias; when it has no alias, or the key has since been rotated or
deleted, the field keeps the creating user's alias that _get_user_info filled in,
because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id must be resolved here too: the spend update writer reads it
off this metadata to increment organization spend, so leaving it out silently
drops batch cost from org accounting for keys and teams that belong to one.
"""
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
request_tags = getattr(job, "request_tags", None)
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}
key_alias = await self._get_key_alias(batch_id, api_key, job.created_by)
if key_alias is not None:
metadata["user_api_key_alias"] = key_alias
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
org_id: Final = await self._get_org_id(job, batch_id)
if org_id is not None:
metadata["user_api_key_org_id"] = org_id
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
return metadata
async def _cleanup_stale_managed_objects(self) -> None:
"""
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@ -317,7 +181,7 @@ class CheckBatchCost:
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]":
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(
where={
@ -337,7 +201,7 @@ class CheckBatchCost:
order={"created_at": "asc"},
)
async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None:
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
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
@ -364,7 +228,7 @@ class CheckBatchCost:
"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
@ -378,10 +242,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(
@ -390,7 +260,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
@ -402,7 +274,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(
@ -411,7 +287,7 @@ class CheckBatchCost:
)
@staticmethod
def _has_unified_id_without_model(job: "_ManagedObjectRow") -> bool:
def _has_unified_id_without_model(job: "ManagedObjectRow") -> bool:
"""A unified id that decodes but carries no model_id can never be routed."""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
@ -460,10 +336,10 @@ class CheckBatchCost:
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
async def _finalize_unbilled_terminal_job(
self, job: "_ManagedObjectRow", response: "LiteLLMBatch"
) -> None:
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
) -> bool:
"""Persist a terminal batch that has nothing billable, converting any raw
provider file ids to managed ids, and take it out of the poll page."""
provider file ids to managed ids, and take it out of the poll page. Returns whether it was persisted."""
try:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
@ -491,10 +367,12 @@ class CheckBatchCost:
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(
f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
)
return False
@staticmethod
def _record_error(
@ -505,7 +383,7 @@ class CheckBatchCost:
def _resolve_job_routing(
self,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
prom_logger: Optional["PrometheusLogger"],
) -> Optional[Tuple[str, str]]:
"""
@ -582,7 +460,7 @@ class CheckBatchCost:
def _resolve_unmanaged_provider_routing(
self,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
prom_logger: Optional["PrometheusLogger"],
llm_provider: str,
bare_model_name: str,
@ -678,7 +556,7 @@ class CheckBatchCost:
@classmethod
def _get_managed_file_model_name(
cls,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
deployment_info: "Deployment",
) -> Optional[str]:
"""
@ -698,9 +576,7 @@ class CheckBatchCost:
)
@staticmethod
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
import json
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
from litellm.types.utils import LiteLLMBatch
file_object: Final = job.file_object
@ -719,7 +595,7 @@ class CheckBatchCost:
async def _track_completed_batch_cost(
self,
job: "_ManagedObjectRow",
job: "ManagedObjectRow",
response: "LiteLLMBatch",
model_id: str,
batch_id: str,
@ -735,9 +611,9 @@ class CheckBatchCost:
"""
from litellm.batches.batch_utils import (
calculate_batch_cost_and_usage,
count_error_file_failed_requests,
get_file_content_as_dictionary,
calculate_batch_cost_and_usage,
)
from litellm.files.main import afile_content
from litellm.files.types import FileContentCallOptions
@ -853,8 +729,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
@ -890,8 +765,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",
@ -911,7 +787,7 @@ class CheckBatchCost:
},
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
"metadata": {
**(await self._build_creator_attribution_metadata(job, batch_id)),
**(await self._creator_attribution.batch_metadata(job, batch_id)),
# spend logs read the deployment identity off these metadata keys, so
# without them the batch cost row carries no model_id or model_group
"model_info": {"id": model_id},
@ -922,7 +798,9 @@ 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"
@ -941,7 +819,7 @@ class CheckBatchCost:
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)
await self._release_job_claim(job, original_status, original_file_object)
raise
# Record batch duration (completed_at - created_at)
@ -956,7 +834,191 @@ class CheckBatchCost:
return model_name, str(llm_provider) if llm_provider else None
async def check_batch_cost(self):
async def _retire_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str
) -> None:
"""Persist a terminal batch that has nothing billable, then store its lines, reading its input file by the id
the provider reported rather than the managed id that was persisted."""
provider_input_file_id: Final = response.input_file_id
if await self._finalize_unbilled_terminal_job(job, response):
response.input_file_id = provider_input_file_id
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=False)
def _store_batch_line_items_in_background(
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
) -> None:
"""Store a retired batch's lines off the poll loop, one batch at a time and only between polls, so neither
pricing the lines nor a slow or hung bucket holds up billing or retiring the other batches. While a poll
cycle's worth of batches is still waiting, the lines of the next one are dropped with an error."""
if not litellm.batch_line_item_callbacks:
return
if len(self._line_item_stores) >= MAX_OBJECTS_PER_POLL_CYCLE:
verbose_proxy_logger.error(
f"CheckBatchCost: {len(self._line_item_stores)} batches are still waiting to store their line items, "
f"so the line items of batch {batch_id} (job {job.id}) were not stored"
)
return
store: Final = asyncio.create_task(
self._store_batch_line_items_in_turn(job, response, model_id, batch_id, billed=billed)
)
self._line_item_stores.add(store)
store.add_done_callback(self._line_item_stores.discard)
async def _store_batch_line_items_in_turn(
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
) -> None:
async with self._line_item_store_turn:
await self._poll_idle.wait()
await self._store_batch_line_items(job, response, model_id, batch_id, billed=billed)
async def _store_batch_line_items(
self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool
) -> None:
"""Send every line of a finished batch to batch_line_item_callbacks, one part at a time. A billed batch's
output file is read again here, so nothing waiting for its turn holds a batch's files. Failures, including an
output file that can no longer be read, are logged and never raised, so the batch stays billed and retired
whether or not its lines were stored."""
if not litellm.batch_line_item_callbacks:
return
from litellm.batches.batch_line_items import (
BatchFileProvider,
BatchLineItemSource,
batch_line_item_parts,
dispatch_batch_line_items,
line_item_metadata,
)
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
try:
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
if deployment_info is None:
verbose_proxy_logger.error(
f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
f"{batch_id} (job {job.id}) were not stored"
)
return
litellm_model_name: Final = deployment_info.litellm_params.model
model_name, llm_provider, _, _ = get_llm_provider(
model=litellm_model_name,
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
)
cost_provider: Final = cast(BatchFileProvider, llm_provider)
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
file_access: Final = {
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
}
output_content: Final = (
await fetch_batch_managed_file_content(
response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access
)
if billed and response.output_file_id
else b""
)
error_content: Final = await self._batch_file_content(response.error_file_id, cost_provider, file_access)
if not output_content and not error_content:
return
redact_bodies: Final = await self._creator_turned_off_message_logging(job)
source: Final = BatchLineItemSource(
batch_id=job.unified_object_id,
finished_at=response.completed_at
or response.failed_at
or response.expired_at
or response.cancelled_at
or response.created_at,
model=model_name,
custom_llm_provider=llm_provider,
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)),
)
input_content: Final = (
b""
if redact_bodies
else await self._batch_file_content(response.input_file_id, cost_provider, file_access)
)
parts: Final = batch_line_item_parts(
source, input_content, (output_content, error_content), BATCH_LINE_ITEMS_MAX_PART_BYTES
)
handed = 0
# pricing a part's lines is CPU-bound, so each part is built off the event loop
while (payload := await asyncio.to_thread(next, parts, None)) is not None:
await dispatch_batch_line_items(payload, redact_bodies=redact_bodies)
handed += len(payload["lines"])
if handed:
verbose_proxy_logger.info(
f"CheckBatchCost: handed the {handed} line items of batch {batch_id} to batch_line_item_callbacks"
)
except Exception as error: # noqa: BLE001 # storing the lines is best effort, the batch is already billed
verbose_proxy_logger.exception(
f"CheckBatchCost: the line items of batch {batch_id} (job {job.id}) were not stored: {error}"
)
@staticmethod
async def _batch_file_content(
file_id: str | None, custom_llm_provider: "BatchFileProvider", file_access: Mapping[str, object]
) -> bytes:
"""A batch file's raw JSONL, or nothing when the batch has no such file or the provider will not serve it."""
from litellm.batches.batch_utils import fetch_batch_managed_file_content
if not file_id:
return b""
try:
return await fetch_batch_managed_file_content(
file_id, custom_llm_provider=custom_llm_provider, litellm_params=file_access
)
except Exception as error: # noqa: BLE001 # a missing input or error file leaves those lines out, not the rest
verbose_proxy_logger.warning(f"CheckBatchCost: could not read batch file {file_id}: {error}")
return b""
async def _creator_turned_off_message_logging(self, job: "ManagedObjectRow") -> bool:
"""Whether the creating key's or team's logging settings turn off message logging, falling back to the
proxy-wide setting. A batch whose saved credential is not a virtual key's hash, such as one created under JWT
or OAuth2 auth or with the master key, has only its team's and the proxy's settings. A key or team that can no
longer be read counts as turned off."""
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata
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)
has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
try:
key_row: Final = (
await _token_table(self.prisma_client).find_unique(where={"token": api_key}) if has_key_row else None
)
team_row: Final = (
await _team_table(self.prisma_client).find_unique(where={"team_id": team_id}) if team_id else None
)
except Exception as error:
verbose_proxy_logger.error(
f"CheckBatchCost: could not read the logging settings of batch job {job.id}'s creator, "
f"so its line items are stored without request and response: {error}"
)
return True
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),
proxy_config,
)
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
creator_setting: Final = str_to_bool(callback_vars.get("turn_off_message_logging"))
return creator_setting if creator_setting is not None else litellm.turn_off_message_logging is True
async def check_batch_cost(self) -> None:
"""Run one poll cycle. Batch line items wait until it has billed and retired every batch it found."""
self._poll_idle.clear()
try:
await self._poll_batches()
finally:
self._poll_idle.set()
async def _poll_batches(self):
"""
Check if the batch JOB has been tracked.
- get all status="validating" and file_purpose="batch" jobs
@ -1050,10 +1112,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,
@ -1070,7 +1129,7 @@ class CheckBatchCost:
f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
f"does not exist at the provider; retiring job {job.id} unbilled"
)
await self._finalize_unbilled_terminal_job(job, response)
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
@ -1084,22 +1143,21 @@ 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:
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)
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
from litellm.proxy.openai_files_endpoints.common_utils import (
@ -1112,7 +1170,7 @@ class CheckBatchCost:
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
)
continue
await self._finalize_unbilled_terminal_job(job, response)
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
# Record polling run metrics (always, even if nothing was processed)
if prom_logger:

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

@ -1,12 +1,36 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled by the get-responses call, which prices normally only because the
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
same route are non-inference and free.
A background response is billed exactly once: whichever client read or poll first sees it
finished moves its row out of a pending status, then logs the body that read already fetched as
one spend event stamped with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN and the creating key's
attribution, priced against the deployment that served it. The poll's status checks are stamped
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.
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
the finished response through it before the client sees response.completed.
A read the caller retargets with client-side credentials, and a delete of a still-pending response,
settle the bill from a fetch through the configured deployment instead: the client controls what a
retargeted read returns, and a deleted response leaves the poll nothing to fetch.
A row stored before this billing model keeps the encrypted create id as its model_object_id. The
first lookup by that id rekeys it by the provider id, and the poll retires such a row without
billing it once a read has already tracked the response under the provider id.
"""
import uuid
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Dict, Final, Protocol, cast
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -14,36 +38,124 @@ from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
PENDING_RESPONSE_STATUSES,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
TERMINAL_RESPONSE_STATUSES,
)
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
from litellm.types.router import Deployment
from litellm.types.utils import (
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
class _ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@dataclass(frozen=True, slots=True)
class _PollTarget:
response_id: str
model_metadata: Mapping[str, object]
deployment: Deployment | None
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
@dataclass(frozen=True, slots=True)
class BilledBackgroundResponse:
response: ResponsesAPIResponse
model: str
custom_llm_provider: str | None
litellm_params: Mapping[str, object]
metadata: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class BilledCharge:
cost: float | None
logging_obj: "LiteLLMLogging"
BilledResponseLogger = Callable[[BilledBackgroundResponse], Awaitable[BilledCharge]]
UntrackedResponseTracker = Callable[[ResponsesAPIResponse], Awaitable[bool]]
def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
unpriced: Final = response.model_copy(deep=True)
unpriced.hidden_params = {key: value for key, value in response.hidden_params.items() if key != "response_cost"}
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
logging_obj: Final = LiteLLMLogging(
model=billed.model,
messages=[{"role": "user", "content": "<retrieve_response>"}],
stream=False,
call_type=_RETRIEVE_RESPONSE_CALL_TYPE,
start_time=datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id=str(uuid.uuid4()),
)
logging_obj.update_environment_variables(
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
optional_params={},
custom_llm_provider=billed.custom_llm_provider,
)
unpriced: Final = _without_read_price(billed.response)
cost: Final = logging_obj.response_cost_calculator(result=unpriced)
end_time: Final = datetime.now()
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time
)
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=unpriced, start_time=logging_obj.start_time, end_time=end_time
)
return BilledCharge(cost=cost, logging_obj=logging_obj)
class _DeploymentBaseModel(BaseModel):
model_config = ConfigDict(frozen=True)
base_model: str | None = None
def _deployment_litellm_params(deployment: Deployment) -> Mapping[str, object]:
from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
api_base: Final = deployment.litellm_params.api_base
base_model: Final = _DeploymentBaseModel.model_validate(deployment.litellm_params, from_attributes=True).base_model
return MappingProxyType(
{
**({"api_base": mask_api_base_credentials(api_base)} if api_base else {}),
**({"base_model": base_model} if base_model else {}),
}
)
def _deployment_metadata(deployment: Deployment) -> Mapping[str, object]:
return MappingProxyType(
{
"model_info": deployment.model_info.model_dump(exclude_none=True),
"model_group": deployment.model_name,
"deployment": deployment.litellm_params.model,
}
)
def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[ManagedObjectRow]":
return ManagedObjectRepository(prisma_client).table
@ -51,12 +163,17 @@ def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -
return getattr(error, "status_code", None) == 404 and provider_response_id in str(error)
def _keyed_by_issued_id(job: ManagedObjectRow) -> bool:
return decrypted_response_id(job.model_object_id) != job.model_object_id
class CheckResponsesCost:
def __init__(
self,
proxy_logging_obj: "ProxyLogging",
prisma_client: "PrismaClient",
llm_router: "Router",
log_billed_response: BilledResponseLogger = log_billed_background_response,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -64,15 +181,17 @@ class CheckResponsesCost:
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
self._log_billed_response: Final = log_billed_response
def _resolve_deployment(self, response_id: str) -> bool:
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
return model_id is not None and self.llm_router.get_deployment(model_id=model_id) is not None
return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
async def _get_response(
self,
response_id: str,
litellm_metadata: Dict[str, str],
litellm_metadata: dict[str, object],
via_router: bool,
) -> ResponsesAPIResponse:
"""Fetch the upstream response, using deployment credentials when available.
@ -91,9 +210,17 @@ class CheckResponsesCost:
)
return cast(ResponsesAPIResponse, router_response)
async def _expire_stale_rows(
self, cutoff: datetime, batch_size: int
) -> int:
async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
return await self._get_response(
response_id=target.response_id,
litellm_metadata={
**target.model_metadata,
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN,
},
via_router=target.deployment is not None,
)
async def _expire_stale_rows(self, cutoff: datetime, batch_size: int) -> int:
"""Execute the bounded UPDATE that marks stale rows as 'stale_expired'.
Isolated so it can be swapped / mocked in tests without touching the
@ -140,14 +267,154 @@ class CheckResponsesCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
def _target(self, response_id: str, model_name: str | None) -> _PollTarget:
return _PollTarget(
response_id=response_id,
model_metadata=MappingProxyType({"model": model_name, "model_group": model_name} if model_name else {}),
deployment=self._resolve_deployment(response_id),
)
def _poll_target(self, job: ManagedObjectRow) -> _PollTarget:
stored_model: Final = _STORED_RESPONSE_ADAPTER.validate_python(job.file_object).get("model")
return self._target(
decrypted_response_id(job.unified_object_id), stored_model if isinstance(stored_model, str) else None
)
async def _claim_for_billing(self, job: ManagedObjectRow) -> bool:
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
# bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
where={"id": job.id, "status": {"in": list(PENDING_RESPONSE_STATUSES)}},
data={"status": "completed"},
)
return claimed > 0
async def _release_billing_claim(self, job: ManagedObjectRow) -> None:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "status": "completed"},
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
)
async def _billed_response(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
) -> BilledBackgroundResponse:
deployment: Final = target.deployment
return BilledBackgroundResponse(
response=response.model_copy(update={"id": job.unified_object_id}),
model=deployment.litellm_params.model if deployment is not None else (response.model or ""),
custom_llm_provider=ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
"custom_llm_provider"
),
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)),
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
}
),
)
async def _bill_once(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
) -> BilledCharge | None:
if not await self._claim_for_billing(job):
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))
except BaseException:
await self._release_billing_claim(job)
raise
verbose_proxy_logger.info(f"Response {job.unified_object_id} reached a terminal status and was billed")
return charge
async def _job_for_provider_id(self, provider_response_id: str) -> ManagedObjectRow | None:
return await _managed_object_table(self.prisma_client).find_first(
where={"model_object_id": provider_response_id, "file_purpose": "response"}
)
async def _rekey_by_provider_id(self, job: ManagedObjectRow) -> None:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id}, data={"model_object_id": decrypted_response_id(job.model_object_id)}
)
async def _tracked_job(self, response_id: str) -> ManagedObjectRow | None:
if tracked := await self._job_for_provider_id(decrypted_response_id(response_id)):
return tracked
stored_by_issued_id: Final = await _managed_object_table(self.prisma_client).find_first(
where={"unified_object_id": response_id, "file_purpose": "response"}
)
if stored_by_issued_id is not None and _keyed_by_issued_id(stored_by_issued_id):
await self._rekey_by_provider_id(stored_by_issued_id)
return stored_by_issued_id
async def _superseded_by_a_provider_id_row(self, job: ManagedObjectRow) -> bool:
if not _keyed_by_issued_id(job):
return False
if await self._job_for_provider_id(decrypted_response_id(job.model_object_id)) is None:
await self._rekey_by_provider_id(job)
return False
await self._claim_for_billing(job)
return True
async def _track_from_deployment(
self, response_id: str, track_untracked: UntrackedResponseTracker | None
) -> BilledCharge | None:
target: Final = self._target(decrypted_response_id(response_id), None)
if track_untracked is None or target.deployment is None:
return None
return await self.bill_if_finished(response_id, await self._fetch_unbilled(target), track_untracked)
async def _start_tracking(
self, response_id: str, response: ResponsesAPIResponse, track_untracked: UntrackedResponseTracker | None
) -> ManagedObjectRow | None:
if track_untracked is None or not await track_untracked(response):
return None
return await self._tracked_job(response_id)
async def bill_if_finished(
self,
response_id: str,
response: ResponsesAPIResponse,
track_untracked: UntrackedResponseTracker | None = None,
) -> BilledCharge | None:
if response.status not in TERMINAL_RESPONSE_STATUSES:
return None
job: Final = await self._tracked_job(response_id) or await self._start_tracking(
response_id, response, track_untracked
)
if job is None or job.status not in PENDING_RESPONSE_STATUSES:
return None
return await self._bill_once(job, self._poll_target(job), response)
async def bill_from_deployment(
self, response_id: str, track_untracked: UntrackedResponseTracker | None = None
) -> BilledCharge | None:
job: Final = await self._tracked_job(response_id)
if job is None:
return await self._track_from_deployment(response_id, track_untracked)
if job.status not in PENDING_RESPONSE_STATUSES:
return None
target: Final = self._poll_target(job)
response: Final = await self._fetch_unbilled(target)
if response.status not in TERMINAL_RESPONSE_STATUSES:
return None
return await self._bill_once(job, target, response)
async def check_responses_cost(self):
"""
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is tracked by the get-responses call, billed because the poll is stamped
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
- Mark responses in a terminal state as complete in the database
- Query the provider, unbilled, to check if the response is complete
- Claim each terminal job by moving its row to completed, then log the body that check fetched
as its one billed event, so two pollers racing over the same job bill it once
- Mark responses the provider no longer has (404 through a resolved
deployment) as stale_expired
"""
@ -160,99 +427,53 @@ class CheckResponsesCost:
jobs = await _managed_object_table(self.prisma_client).find_many(
where={
"status": {"in": ["queued", "in_progress"]},
# bounded-ok: PENDING_RESPONSE_STATUSES is a two-value constant tuple
"status": {"in": list(PENDING_RESPONSE_STATUSES)},
"file_purpose": "response",
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
completed_jobs: Final[list[_ManagedObjectRow]] = []
expired_jobs: Final[list[_ManagedObjectRow]] = []
expired_jobs: Final[list[ManagedObjectRow]] = []
for job in jobs:
unified_object_id = job.unified_object_id
try:
from litellm.proxy.hooks.responses_id_security import (
ResponsesIDSecurity,
)
# Get the stored response object to extract model information
stored_response = job.file_object
model_name = stored_response.get("model", None)
# Decrypt the response ID
responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
# Prepare metadata with model information for cost tracking
litellm_metadata = {
"user_api_key_user_id": job.created_by or "default-user-id",
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
}
# Add model information if available
if model_name:
litellm_metadata["model"] = model_name
litellm_metadata["model_group"] = model_name # Use same value for model_group
via_router = self._resolve_deployment(responses_id_security)
if await self._superseded_by_a_provider_id_row(job):
verbose_proxy_logger.info(f"Response {unified_object_id} is billed under its provider id row")
continue
target = self._poll_target(job)
except Exception as e:
verbose_proxy_logger.warning(
f"Skipping job {unified_object_id} due to error: {e}"
)
verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
continue
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(
responses_id_security
).get("response_id", responses_id_security)
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(target.response_id).get(
"response_id", target.response_id
)
try:
response = await self._get_response(
response_id=responses_id_security,
litellm_metadata=litellm_metadata,
via_router=via_router,
)
response = await self._fetch_unbilled(target)
verbose_proxy_logger.debug(
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
)
verbose_proxy_logger.debug(f"Response {unified_object_id} status: {response.status}")
if response.status in TERMINAL_RESPONSE_STATUSES:
await self._bill_once(job, target, response)
except Exception as e:
if via_router and _is_response_gone_at_provider(e, provider_response_id):
if target.deployment is not None and _is_response_gone_at_provider(e, provider_response_id):
verbose_proxy_logger.info(
f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}"
)
expired_jobs.append(job)
continue
verbose_proxy_logger.warning(
f"Skipping job {unified_object_id} due to error: {e}"
)
verbose_proxy_logger.warning(f"Skipping job {unified_object_id} due to error: {e}")
continue
if response.status in TERMINAL_RESPONSE_STATUSES:
verbose_proxy_logger.info(
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
)
completed_jobs.append(job)
# Mark completed jobs in the database
if len(completed_jobs) > 0:
await _managed_object_table(self.prisma_client).update_many(
# bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
where={"id": {"in": [job.id for job in completed_jobs]}},
data={"status": "completed"},
)
verbose_proxy_logger.info(
f"Marked {len(completed_jobs)} response jobs as completed"
)
if len(expired_jobs) > 0:
await _managed_object_table(self.prisma_client).update_many(
# bounded-ok: at most MAX_OBJECTS_PER_POLL_CYCLE rows per cycle, the find_many take above
where={"id": {"in": [job.id for job in expired_jobs]}},
data={"status": "stale_expired"},
)
verbose_proxy_logger.info(
f"Marked {len(expired_jobs)} response jobs as stale_expired"
)
verbose_proxy_logger.info(f"Marked {len(expired_jobs)} response jobs as stale_expired")

218
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py Normal file
View file

@ -0,0 +1,218 @@
"""Creator attribution for the spend row a managed-object cost poll writes.
A batch or a background response is billed by a poll job that runs long after the creating
request returned, so the managed object row's creator columns are the only record of which
key, team, organization, user, end user, and tags carry the spend.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, cast
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLI_SESSION_KEY_PREFIX
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.proxy.utils import PrismaClient
from litellm.repositories.prisma_protocols import TableActions
class ManagedObjectRow(Protocol):
@property
def id(self) -> str: ...
@property
def unified_object_id(self) -> str: ...
@property
def model_object_id(self) -> str: ...
@property
def created_by(self) -> str | None: ...
@property
def status(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@property
def org_id(self) -> str | None: ...
@property
def api_key(self) -> str | None: ...
@property
def team_id(self) -> str | None: ...
@property
def request_tags(self) -> object: ...
@property
def end_user_id(self) -> str | None: ...
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
def _string_tags(request_tags: object) -> list[str] | None:
if not isinstance(request_tags, list):
return None
return [tag for tag in _TAG_LIST_ADAPTER.validate_python(request_tags) if isinstance(tag, str)] or None
def _creator_user_id(created_by: str | None) -> str | None:
if created_by is None or created_by.startswith(_KEY_OWNER_SCOPE_PREFIX):
return None
return created_by
def _user_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_UserTable]:
return UserRepository(prisma_client).table
def _token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_VerificationToken]:
return VerificationTokenRepository(prisma_client).table
def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_TeamTable]:
return TeamRepository(prisma_client).table
@dataclass(frozen=True, slots=True)
class _CreatorRows:
key: prisma_models.LiteLLM_VerificationToken | None
team: prisma_models.LiteLLM_TeamTable | None
user: prisma_models.LiteLLM_UserTable | None
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,
user, aliases, and tags.
Rows created before api_key and request_tags were persisted carry only created_by and
team_id, and fall back to those. A named creating key owns user_api_key_alias; when it
has no alias, or the key has since been rotated or deleted, the field keeps the creating
user's alias, 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.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
self._prisma_client: Final = prisma_client
self._poller_name: Final = poller_name
async def batch_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
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 _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
return _CreatorRows(
key=await self._key_row(object_id, getattr(job, "api_key", None), job.created_by),
team=await self._team_row(object_id, getattr(job, "team_id", None)),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
async def _user_row(self, object_id: str, user_id: str | None) -> prisma_models.LiteLLM_UserTable | None:
"""Objects created by a team or service account key carry no user id, and
find_unique(where={"user_id": None}) raises, so those resolve no row."""
if not user_id:
return None
try:
return await _user_table(self._prisma_client).find_unique(where={"user_id": user_id})
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up user {user_id} for {object_id}: {e}")
return None
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})
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:
return None
try:
return await _team_table(self._prisma_client).find_unique(where={"team_id": 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
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
key_alias: Final = _key_alias(api_key, job.created_by, rows.key)
team_alias: Final = getattr(rows.team, "team_alias", None)
org_id: Final = getattr(job, "org_id", None) or _organization_id(rows.key) or _organization_id(rows.team)
tags: Final = _string_tags(getattr(job, "request_tags", None))
return {
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(
{
"user_api_key_user_email": getattr(rows.user, "user_email", None),
"user_api_key_alias": getattr(rows.user, "user_alias", None),
}
if rows.user is not None
else {}
),
**({"user_api_key_alias": key_alias} if key_alias is not None else {}),
**({"user_api_key_team_alias": team_alias} if team_alias is not None else {}),
**({"user_api_key_org_id": org_id} if org_id is not None else {}),
**({"tags": tags} if tags else {}),
}
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}"
def _key_alias(
api_key: str | None, created_by: str | None, key_row: prisma_models.LiteLLM_VerificationToken | None
) -> str | None:
if api_key and _is_cli_session_key(api_key, created_by):
return api_key
return getattr(key_row, "key_alias", None) if key_row is not None else None
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]:
if row is None:
return {}
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}

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

@ -24,7 +24,9 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import BaseModel, ValidationError
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -35,17 +37,17 @@ from litellm.caching.caching import DualCache
from litellm.constants import (
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
MAX_FILE_LIST_LIMIT,
PENDING_RESPONSE_STATUSES,
)
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
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 (
extract_file_metadata,
)
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import (
@ -88,6 +90,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
request_tags_from_metadata,
)
from litellm.repositories.managed_file_repository import ManagedFileRepository
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
AllMessageValues,
AsyncCursorPage,
@ -100,6 +103,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
ResponsesAPIResponse,
)
from litellm.types.utils import (
ADDRESSED_RESPONSE_ID_FIELD,
CallTypesLiteral,
LiteLLMBatch,
LiteLLMFineTuningJob,
@ -346,6 +350,8 @@ class _ManagedObjectTableActions(Protocol):
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
async def create(self, data: Mapping[str, object]) -> "PrismaManagedObjectRow": ...
class _ManagedResourceDatabase(Protocol):
@property
@ -392,6 +398,39 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
return prisma_client.db.litellm_managedobjecttable
_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
def _storable_text(text: str) -> str:
return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
def _storable_json_value(value: JsonValue) -> JsonValue:
if isinstance(value, str):
return _storable_text(value)
if isinstance(value, list):
return [_storable_json_value(item) for item in value]
if isinstance(value, dict):
return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
return value
def postgres_storable_json(file_object: BaseModel) -> str:
dumped: Final = file_object.model_dump_json()
if "\\u0000" not in dumped:
return dumped
return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
def _without_model_output(
file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
return (
file_object.model_copy(update={"output": []}) if isinstance(file_object, ResponsesAPIResponse) else file_object
)
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@ -405,6 +444,40 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
)
_MANAGED_RESPONSE_READ_CALL_TYPES: Final = frozenset({"aget_responses", "acancel_responses"})
_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object])
def _addresses_a_proxy_issued_response(data: Mapping[str, object]) -> bool:
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
addressed_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD) or data.get("response_id")
return (
isinstance(addressed_id, str)
and ResponsesAPIRequestUtils.get_model_id_from_response_id(decrypted_response_id(addressed_id)) is not None
)
def _is_streamed_background_create(request: Mapping[str, object]) -> bool:
return (
request.get("stream") is True
and isinstance(request.get("litellm_metadata"), dict)
and requests_background(request)
)
def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str, object]:
metadata_field: Final = "litellm_metadata" if isinstance(request.get("litellm_metadata"), dict) else "metadata"
existing_metadata: Final = request.get(metadata_field)
return {
**request,
metadata_field: {
**(_REQUEST_FIELDS.validate_python(existing_metadata) if isinstance(existing_metadata, dict) else {}),
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: True,
},
}
_MANAGED_FILES_TARGET: Final = "managed_files"
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
@ -529,11 +602,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) -> None:
"""Persist a managed object row, caching it and upserting it in the DB.
persist_attribution is set only by the batch create, which is the one caller
that can speak for the creator; it gates the api_key and request_tags columns
that CheckBatchCost bills against, so a later poll or retrieve of the same
batch cannot record itself as the paying key. Like created_by and team_id,
both are written only in the upsert create branch, never on update.
persist_attribution is set only by the batch and background response creates,
the callers that can speak for the creator; it gates the api_key, request_tags, and,
for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
retrieve of the same object cannot record itself as the paying key. Like created_by and team_id,
they are written only in the upsert create branch, never on update.
A stored response keeps no model output: billing reads only its status and model from
the row, so the row never holds generated content.
create_if_missing is cleared by callers that observe a batch they did not
create, such as a poll. They still refresh status and file_object, but a
@ -546,11 +622,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
It is written only in the upsert create branch.
"""
verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
stored_object: Final = _without_model_output(file_object)
litellm_managed_object = LiteLLM_ManagedObjectTable(
unified_object_id=unified_object_id,
model_object_id=model_object_id,
file_purpose=file_purpose,
file_object=file_object,
file_object=stored_object,
)
await self.internal_usage_cache.async_set_cache(
key=unified_object_id,
@ -561,32 +638,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
from prisma import Json
api_key = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
end_user_id: Final = user_api_key_dict.end_user_id if file_purpose == "response" else None
attribution_columns = (
{
**({"api_key": api_key} if api_key is not None else {}),
**({"request_tags": Json(list(request_tags))} if request_tags else {}),
**({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
**({"end_user_id": end_user_id} if end_user_id else {}),
}
if persist_attribution
else {}
)
# FIX: Update status and file_object on every operation to keep state in sync
update_columns: Final = {
"file_object": file_object.model_dump_json(),
"file_object": postgres_storable_json(stored_object),
"status": file_object.status,
"updated_by": user_api_key_dict.user_id,
}
managed_object_table: Final = _managed_object_table(self.prisma_client)
if not create_if_missing:
await _managed_object_table(self.prisma_client).update_many(
await managed_object_table.update_many(
where={"unified_object_id": unified_object_id},
data=update_columns,
)
return
await _managed_object_table(self.prisma_client).upsert(
await managed_object_table.upsert(
where={"unified_object_id": unified_object_id},
data={
"create": {
"unified_object_id": unified_object_id,
"file_object": file_object.model_dump_json(),
"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),
@ -601,6 +681,42 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
},
)
async def track_untracked_background_response(
self,
response_id: str,
response: ResponsesAPIResponse,
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
) -> bool:
if not _addresses_a_proxy_issued_response({"response_id": response_id}):
return False
from prisma import Json
from prisma.errors import UniqueViolationError
from litellm.proxy.hooks.responses_id_security import decrypted_response_id
api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) or None
try:
await _managed_object_table(self.prisma_client).create(
data={
"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)),
"status": PENDING_RESPONSE_STATUSES[0],
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"org_id": await self._resolve_creator_org_id(user_api_key_dict),
"updated_by": user_api_key_dict.user_id,
**({"api_key": api_key} if api_key is not None else {}),
**({"end_user_id": user_api_key_dict.end_user_id} if user_api_key_dict.end_user_id else {}),
**({"request_tags": Json([_storable_text(tag) for tag in request_tags])} if request_tags else {}),
}
)
except UniqueViolationError:
return False
return True
@with_service_target(_MANAGED_FILES_TARGET)
async def get_unified_file_id(
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
@ -1375,6 +1491,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools)
if unified_vs_ids:
await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict)
create_request: Final = _REQUEST_FIELDS.validate_python(data)
if _is_streamed_background_create(create_request):
return _with_managed_background_marker(create_request)
elif call_type == CallTypes.afile_content.value:
retrieve_file_id = cast(Optional[str], data.get("file_id"))
potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False
@ -1432,6 +1551,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
elif retrieve_object_id and accessor_key == "batch_id":
await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
elif call_type in _MANAGED_RESPONSE_READ_CALL_TYPES:
request: Final = _REQUEST_FIELDS.validate_python(data)
if _addresses_a_proxy_issued_response(request):
return _with_managed_background_marker(request)
elif call_type == CallTypes.acreate_fine_tuning_job.value:
fine_tuning_file_ids: Final = tuple(
file_id

4
enterprise/pyproject.toml
View file

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

34
gateway/Dockerfile
View file

@ -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
@ -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/

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

@ -126,6 +126,11 @@ Set `billingMetrics.caSecretName` only when the collector is a private or test o
| `db.secret.name` | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials. | `postgres` |
| `db.secret.usernameKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance. | `username` |
| `db.secret.passwordKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user. | `password` |
| `db.username` | If `db.useExisting` is `true`, the database user as a plain value. When set, `db.secret.usernameKey` is not read. | `""` |
| `db.port` | If `db.useExisting` is `true` with token auth, the Postgres port used to build the connection URL. Password auth takes the port from `db.url`. | `5432` |
| `db.schema` | With token auth, the Postgres schema holding the LiteLLM tables, rendered as `DATABASE_SCHEMA` when the tables are not in `public`. Ignored with password auth, where the schema stays in `db.url` as `?schema=`. | `""` |
| `db.useIAMAuth` | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with an AWS RDS IAM token instead of a password. See [Token auth](#token-auth-to-an-existing-postgres-rds-iam-or-microsoft-entra-id). | `false` |
| `db.useAzureEntraAuth` | If `db.useExisting` is `true` and `db.deployStandalone` is `false`, authenticate with a Microsoft Entra ID token instead of a password. Mutually exclusive with `db.useIAMAuth`. | `false` |
| `db.useStackgresOperator` | Not yet implemented. | `false` |
| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` |
| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) |
@ -156,6 +161,22 @@ data:
type: Opaque
```
#### Token auth to an existing Postgres (RDS IAM or Microsoft Entra ID)
When the database accepts short-lived tokens instead of a password (AWS RDS or Aurora with IAM database authentication, or Azure Database for PostgreSQL with Microsoft Entra ID), there is no password to put in a Secret. Turn on `db.useIAMAuth` or `db.useAzureEntraAuth` and the chart stops referencing `db.secret.passwordKey` and `db.url`. It sets `IAM_TOKEN_DB_AUTH` (or `AZURE_POSTGRESQL_AUTH`) and emits `DATABASE_HOST`, `DATABASE_PORT`, `DATABASE_USERNAME` and `DATABASE_NAME` for the proxy, the migrations job and the collector sidecar. Each of them mints its own token at startup and builds `DATABASE_URL` from those parts, and the in-container pgbouncer (`db.connectionPool`) refreshes the token before it expires.
```yaml
db:
useExisting: true
deployStandalone: false
useIAMAuth: true
endpoint: litellm.cluster-abc123.us-west-1.rds.amazonaws.com
database: litellm
username: litellm_iam # the role granted rds_iam; no db.secret needed
```
The pod needs cloud credentials for the token: on EKS use IRSA or Pod Identity on the chart's `serviceAccount` (and `migrationJob.serviceAccountName` when the Helm hook runs before the ServiceAccount exists), on AKS use workload identity. Set `AWS_RDS_REGION` under `envVars` when the RDS hostname is not a standard regional endpoint (the token must be signed for the region the database lives in). TLS options that would otherwise go in `db.url` are set as `DATABASE_SSLMODE` and `DATABASE_SSLROOTCERT` under `envVars`. Leaving `db.username` empty reads the user from `db.secret.usernameKey` as before.
#### Examples for `environmentSecrets` and `environemntConfigMaps`
```yaml

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

@ -162,6 +162,77 @@ taken before the change, which by that point no longer exists.
{{- end -}}
{{- end -}}
{{/*
Whether the existing Postgres is reached with a short-lived token (AWS RDS IAM
or Microsoft Entra ID) instead of a password.
*/}}
{{- define "litellm.db.tokenAuth" -}}
{{- if or .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}true{{- end -}}
{{- end -}}
{{- define "litellm.validateDbTokenAuth" -}}
{{- if and .Values.db.useIAMAuth .Values.db.useAzureEntraAuth -}}
{{- fail "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" -}}
{{- end -}}
{{- if and (include "litellm.db.tokenAuth" .) (or (not .Values.db.useExisting) .Values.db.deployStandalone) -}}
{{- fail "db.useIAMAuth and db.useAzureEntraAuth need db.useExisting: true and db.deployStandalone: false (token auth targets an existing RDS or Azure Database for PostgreSQL; the bundled postgresql subchart only supports password auth)" -}}
{{- end -}}
{{- end -}}
{{/*
Database env for an existing Postgres, shared by the proxy container, the
collector sidecar and the migrations Job. Under token auth every entrypoint
mints its own token and assembles DATABASE_URL itself, so db.url is left out.
*/}}
{{- define "litellm.existingDbEnv" -}}
{{- $tokenAuth := include "litellm.db.tokenAuth" . -}}
- name: DATABASE_USERNAME
{{- if .Values.db.username }}
value: {{ .Values.db.username | quote }}
{{- else }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
{{- end }}
{{- if .Values.db.useIAMAuth }}
- name: IAM_TOKEN_DB_AUTH
value: "true"
{{- else if .Values.db.useAzureEntraAuth }}
- name: AZURE_POSTGRESQL_AUTH
value: "true"
{{- else }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
{{- end }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
{{- if $tokenAuth }}
- name: DATABASE_PORT
value: {{ .Values.db.port | quote }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
{{- if and $tokenAuth .Values.db.schema }}
- name: DATABASE_SCHEMA
value: {{ .Values.db.schema | quote }}
{{- end }}
{{- if not $tokenAuth }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- end -}}
{{/*
Environment shared by the proxy container and the opt-in collector sidecar:
database, pgbouncer, master key, redis, user envVars. Both containers must see
@ -169,6 +240,7 @@ the same DATABASE_URL and REDIS_* so the sidecar reaches the pod's pgbouncer
and the same spend transaction buffer.
*/}}
{{- define "litellm.proxyEnv" -}}
{{- include "litellm.validateDbTokenAuth" . -}}
- name: HOST
value: "{{ .Values.listen | default "0.0.0.0" }}"
- name: PORT
@ -189,29 +261,7 @@ and the same spend transaction buffer.
- name: DATABASE_NAME
value: litellm
{{- else if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{ include "litellm.existingDbEnv" . }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST

25
helm/litellm-helm/templates/migrations-job.yaml
View file

@ -68,30 +68,9 @@ spec:
{{- end }}
workingDir: "/app"
env:
{{- include "litellm.validateDbTokenAuth" . -}}
{{- if .Values.db.useExisting }}
- name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.usernameKey }}
- name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.passwordKey }}
- name: DATABASE_HOST
{{- if .Values.db.secret.endpointKey }}
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.endpointKey }}
{{- else }}
value: {{ .Values.db.endpoint }}
{{- end }}
- name: DATABASE_NAME
value: {{ .Values.db.database }}
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- include "litellm.existingDbEnv" . | nindent 12 }}
{{- else if .Values.db.deployStandalone }}
- name: DATABASE_URL
value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }}

288
helm/litellm-helm/tests/database_auth_tests.yaml Normal file
View file

@ -0,0 +1,288 @@
suite: test existing database token auth
templates:
- deployment.yaml
- migrations-job.yaml
- configmap-litellm.yaml
tests:
- it: should keep password auth from the secret when no token auth is set
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: postgres
key: username
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: postgres
key: password
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
- notContains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT
any: true
- it: should render RDS IAM auth without a password secret on the proxy and the collector sidecar
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
db.useIAMAuth: true
db.username: litellm_iam
db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
db.port: 5433
collector.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USERNAME
value: litellm_iam
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
value: writer.cluster-abc.us-west-1.rds.amazonaws.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT
value: "5433"
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME
value: litellm
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
any: true
- contains:
path: spec.template.spec.containers[1].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[1].env
content:
name: DATABASE_PASSWORD
any: true
- it: should give the migrations job the same token auth env as the proxy
template: migrations-job.yaml
set:
migrationJob.enabled: true
db.useExisting: true
db.deployStandalone: false
db.useIAMAuth: true
db.username: litellm_iam
db.endpoint: writer.cluster-abc.us-west-1.rds.amazonaws.com
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USERNAME
value: litellm_iam
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
value: writer.cluster-abc.us-west-1.rds.amazonaws.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT
value: "5432"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
any: true
- it: should carry db.schema as DATABASE_SCHEMA under token auth and omit it when empty
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
db.useIAMAuth: true
db.username: litellm_iam
db.schema: tenant1
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA
value: tenant1
- it: should leave the schema to db.url under password auth even when db.schema is set
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
db.schema: tenant1
db.url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA
any: true
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)?schema=tenant1
- it: should not emit DATABASE_SCHEMA when db.schema is empty
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
db.useIAMAuth: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA
any: true
- it: should read the user from the secret when db.username is empty under token auth
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
db.useIAMAuth: true
db.secret.name: rds-identity
db.secret.usernameKey: user
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USERNAME
valueFrom:
secretKeyRef:
name: rds-identity
key: user
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- it: should render Entra auth with AZURE_POSTGRESQL_AUTH and no password
template: deployment.yaml
set:
db.useExisting: true
db.deployStandalone: false
db.useAzureEntraAuth: true
db.username: litellm@example.onmicrosoft.com
db.endpoint: litellm.postgres.database.azure.com
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USERNAME
value: litellm@example.onmicrosoft.com
- notContains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
any: true
- it: should fail when both token sources are enabled
template: deployment.yaml
set:
migrationJob.enabled: false
db.useExisting: true
db.deployStandalone: false
db.useIAMAuth: true
db.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "db.useIAMAuth and db.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
- it: should fail when token auth is set while the bundled postgres is still deployed
template: deployment.yaml
set:
migrationJob.enabled: false
db.useExisting: true
db.useIAMAuth: true
asserts:
- failedTemplate:
errorPattern: "need db.useExisting: true and db.deployStandalone: false"
- it: should fail when token auth is set without useExisting
template: migrations-job.yaml
set:
migrationJob.enabled: true
db.deployStandalone: false
db.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorPattern: "need db.useExisting: true and db.deployStandalone: false"

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

@ -376,6 +376,30 @@ db:
endpoint: localhost
database: litellm
url: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)
# Authenticate with a short-lived token instead of a password. useIAMAuth is
# AWS RDS / Aurora IAM database authentication (the pod needs AWS credentials,
# e.g. IRSA or EKS Pod Identity, and the database user needs GRANT rds_iam);
# useAzureEntraAuth is Azure Database for PostgreSQL with a Microsoft Entra ID
# identity (workload identity). Mutually exclusive, and both need
# useExisting: true with deployStandalone: false. The chart then sets
# IAM_TOKEN_DB_AUTH or AZURE_POSTGRESQL_AUTH, emits DATABASE_HOST, DATABASE_PORT,
# DATABASE_USERNAME and DATABASE_NAME and skips DATABASE_PASSWORD and db.url:
# the proxy, the migrations job and the collector sidecar each mint their own
# token and build DATABASE_URL from those parts. Set DATABASE_SSLMODE /
# DATABASE_SSLROOTCERT under envVars for TLS options you would otherwise put
# in db.url.
useIAMAuth: false
useAzureEntraAuth: false
# Postgres port, used to build the token-auth URL (password auth takes the
# port from db.url)
port: 5432
# Postgres schema for the LiteLLM tables, rendered as DATABASE_SCHEMA under
# token auth only (password auth carries it as ?schema= in db.url)
schema: ""
# Database user as a plain value. Empty means read it from
# db.secret.usernameKey. With token auth the user name is the only credential
# the chart needs, so setting it here removes the need for db.secret entirely.
username: ""
secret:
name: postgres
usernameKey: username

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

@ -475,6 +475,27 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end }}
{{- end }}
{{- $sslCertReqs := $root.Values.redis.sslCertReqs }}
{{- if and $sslCertReqs (not (has $sslCertReqs (list "none" "optional" "required"))) }}
{{- fail (printf "redis.sslCertReqs must be one of none, optional, required (got %q)" $sslCertReqs) }}
{{- end }}
{{- if and $urlKey (or $root.Values.redis.ssl $sslCertReqs $root.Values.redis.sslCaCerts) }}
{{- fail "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead" }}
{{- end }}
{{- if or $root.Values.redis.host $hostKey }}
{{- if $root.Values.redis.ssl }}
- name: REDIS_SSL
value: "true"
{{- end }}
{{- if $sslCertReqs }}
- name: REDIS_SSL_CERT_REQS
value: {{ $sslCertReqs | quote }}
{{- end }}
{{- if $root.Values.redis.sslCaCerts }}
- name: REDIS_SSL_CA_CERTS
value: {{ $root.Values.redis.sslCaCerts | quote }}
{{- end }}
{{- end }}
{{- with $component.extraEnv }}
{{ toYaml . }}
{{- end }}

166
helm/litellm/tests/redis_env_tests.yaml
View file

@ -3,6 +3,7 @@ templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
@ -320,3 +321,168 @@ tests:
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
- it: redis omits TLS env vars by default
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CERT_REQS
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CA_CERTS
any: true
- it: redis ssl false omits REDIS_SSL
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.ssl: false
redis.sslCertReqs: required
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
any: true
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CERT_REQS
value: "required"
- it: gateway renders redis TLS env vars
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.ssl: true
redis.sslCertReqs: none
redis.sslCaCerts: /etc/redis-tls/ca.crt
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CERT_REQS
value: "none"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CA_CERTS
value: "/etc/redis-tls/ca.crt"
- it: backend renders redis TLS env vars
template: backend/deployment.yaml
set:
redis.host: redis.example.com
redis.ssl: true
redis.sslCertReqs: required
redis.sslCaCerts: /etc/redis-tls/ca.crt
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CERT_REQS
value: "required"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CA_CERTS
value: "/etc/redis-tls/ca.crt"
- it: migration job renders redis TLS env vars
template: migrations-job.yaml
set:
db.deployStandalone: false
db.useExisting: true
db.endpoint: postgres.example.com
db.secret.name: pg-secret
redis.host: redis.example.com
redis.ssl: true
redis.sslCertReqs: optional
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CERT_REQS
value: "optional"
- it: redis TLS values reject a secret-backed URL
template: gateway/deployment.yaml
set:
redis.passwordSecret.name: redis-secret
redis.passwordSecret.urlKey: connection-url
redis.sslCaCerts: /etc/redis-tls/ca.crt
asserts:
- failedTemplate:
errorMessage: "redis.ssl, redis.sslCertReqs and redis.sslCaCerts cannot be combined with redis.passwordSecret.urlKey, use a rediss:// URL with ssl_cert_reqs / ssl_ca_certs query parameters instead"
- it: hostKey mode renders redis TLS env vars
template: gateway/deployment.yaml
set:
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.ssl: true
redis.sslCaCerts: /etc/redis-tls/ca.crt
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CA_CERTS
value: "/etc/redis-tls/ca.crt"
- it: redis TLS env vars are omitted when redis is not configured
template: gateway/deployment.yaml
set:
redis.ssl: true
redis.sslCertReqs: none
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_SSL_CERT_REQS
any: true
- it: redis sslCertReqs rejects unknown values
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.ssl: true
redis.sslCertReqs: strict
asserts:
- failedTemplate:
errorMessage: 'redis.sslCertReqs must be one of none, optional, required (got "strict")'

11
helm/litellm/values.yaml
View file

@ -275,10 +275,21 @@ database:
# self-hosted Redis Cluster). The chart emits REDIS_CLUSTER_NODES from
# `host` / `port` as the single seed; the cluster client discovers the
# remaining nodes from CLUSTER SLOTS at startup.
#
# For a TLS Redis endpoint set `ssl: true`, which the chart emits as REDIS_SSL.
# `sslCertReqs` (none / optional / required) becomes REDIS_SSL_CERT_REQS and
# `sslCaCerts` (path to a CA bundle mounted in the pod, e.g. via
# `extraVolumeMounts`) becomes REDIS_SSL_CA_CERTS. All three are omitted when
# unset. They apply to `host` / `hostKey` mode only. With `passwordSecret.urlKey`
# the proxy reads TLS settings from the URL itself, so use `rediss://` and the
# `ssl_cert_reqs` / `ssl_ca_certs` query parameters there instead.
redis:
cluster: false
host: ""
port: 6379
ssl: false
sslCertReqs: ""
sslCaCerts: ""
passwordSecret:
name: "" # Leave empty for auth-less Redis
passwordKey: password

50
litellm-proxy-extras/litellm_proxy_extras/_redaction.py Normal file
View file

@ -0,0 +1,50 @@
import functools
import os
import re
from collections.abc import Callable
from typing import Final
from urllib.parse import unquote, urlsplit
_REDACTED: Final = "REDACTED"
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
@functools.cache
def _secret_shape_redactor() -> Callable[[str], str]:
try:
from litellm._logging import redact_secrets
except ImportError:
return lambda text: text
return redact_secrets
def _url_passwords(url: str) -> frozenset[str]:
try:
parts: Final = urlsplit(url)
except ValueError:
return frozenset()
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
raw_query_passwords: Final = tuple(
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
)
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
def _configured_database_passwords() -> frozenset[str]:
database_url: Final = os.getenv("DATABASE_URL")
direct_url: Final = os.getenv("DIRECT_URL")
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
return database_passwords | direct_passwords
def _redact_credentials(text: str) -> str:
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
alternation: Final = "|".join(re.escape(password) for password in passwords)
password_pattern: Final = (
re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
)
result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
return _secret_shape_redactor()(result)

74
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py
View file

@ -1,4 +1,5 @@
import random
import re
import time
from collections.abc import Generator, Mapping
from contextlib import contextmanager
@ -7,6 +8,7 @@ from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
migration_lock_timeout,
@ -14,6 +16,26 @@ from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
_SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
@dataclass(frozen=True, slots=True)
class AdvisoryLock:
key: int
name: str
holder: str
COORDINATOR_LOCK: Final = AdvisoryLock(
MIGRATION_LOCK_KEY,
"v2 migration coordinator lock",
"another v2 migration resolver",
)
PRISMA_LOCK: Final = AdvisoryLock(
PRISMA_ADVISORY_LOCK_KEY,
"Prisma migration lock",
"another Prisma migration",
)
if TYPE_CHECKING:
import psycopg
@ -49,6 +71,28 @@ def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = M
return row is not None and row.acquired
def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting: bool) -> str:
import psycopg
prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
raw: Final = str(exc)
unquoted: Final = re.sub(r"\"[^\"]*\"|'[^']*'", "", raw)
detail: Final = _redact_credentials(raw).strip()
if (
not connecting
or not isinstance(exc, psycopg.OperationalError)
or any(marker in unquoted for marker in _SERVER_ANSWERED_MARKERS)
):
return prefix + detail
parsed_url: Final = urlsplit(database_url)
try:
port: Final = parsed_url.port
except ValueError:
return prefix + f"P1001: Can't reach database server. {detail}"
host: Final = parsed_url.hostname or "localhost"
return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
@dataclass(frozen=True, slots=True)
class MigrationCoordinator:
connection: "psycopg.Connection[tuple[object, ...]]"
@ -75,39 +119,45 @@ class MigrationCoordinator:
@contextmanager
def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
def migration_lock(
database_url: str,
lock: AdvisoryLock = COORDINATOR_LOCK,
) -> Generator[MigrationCoordinator, None, None]:
import psycopg
wait_seconds: Final = migration_lock_timeout()
deadline: Final = time.monotonic() + wait_seconds
try:
with psycopg.connect(database_url, connect_timeout=10, autocommit=True) as connection:
connection: Final = psycopg.connect(database_url, connect_timeout=10, autocommit=True)
except psycopg.Error as exc:
raise RuntimeError(_coordination_failure(database_url, exc, connecting=True)) from None
try:
with connection:
coordinator: Final = MigrationCoordinator(connection)
logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds)
while time.monotonic() < deadline:
with connection.transaction():
if _try_lock(connection):
logger.info("Acquired the v2 migration coordinator lock")
if _try_lock(connection, lock.key):
logger.info("Acquired the %s", lock.name)
yield coordinator
coordinator.check_connection()
return
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
except psycopg.Error as exc:
raise RuntimeError(f"Lost or could not establish v2 migration coordination with the database: {exc}") from exc
raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
raise RuntimeError(
f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
f"Timed out waiting for {lock.holder} after {wait_seconds}s. "
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
)
@contextmanager
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
"""Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
pinned connection while a Prisma command runs with its session advisory lock disabled, so a
transaction pooler can never strand that lock on an idle backend after the command exits."""
with migration_lock(database_url) as coordinator:
coordinator.acquire_prisma_lock()
"""Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled.
This keeps deploy independent of the coordinator key held by request-log index builds."""
with migration_lock(database_url, PRISMA_LOCK) as coordinator:
coordinator.release_statement_snapshot()
yield

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20261015000000_managed_object_end_user_id/migration.sql Normal file
View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "end_user_id" TEXT;

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();

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql Normal file
View file

@ -0,0 +1,2 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at"
ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at");

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

@ -1179,6 +1179,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
end_user_id String?
updated_at DateTime @updatedAt
updated_by String?
@ -1980,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
outcome String?
@@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
@@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
}
model LiteLLM_Lens {

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

@ -1,4 +1,3 @@
import functools
import glob
import os
import random
@ -13,10 +12,10 @@ from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
from urllib.parse import unquote, urlsplit
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@ -212,51 +211,6 @@ def _max_migration_timestamp(names) -> int:
return max(_migration_timestamp(n) for n in names)
_REDACTED: Final = "REDACTED"
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
@functools.cache
def _secret_shape_redactor() -> Callable[[str], str]:
try:
from litellm._logging import redact_secrets
except ImportError:
return lambda text: text
return redact_secrets
def _url_passwords(url: str) -> frozenset[str]:
try:
parts: Final = urlsplit(url)
except ValueError:
return frozenset()
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
raw_query_passwords: Final = tuple(
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
)
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
def _configured_database_passwords() -> frozenset[str]:
database_url: Final = os.getenv("DATABASE_URL")
direct_url: Final = os.getenv("DIRECT_URL")
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
return database_passwords | direct_passwords
def _redact_credentials(text: str) -> str:
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
alternation: Final = "|".join(re.escape(password) for password in passwords)
password_pattern: Final = (
re.compile(rf"(?<![A-Za-z0-9_])(?:{alternation})(?![A-Za-z0-9_])", re.IGNORECASE) if passwords else None
)
result: Final = password_pattern.sub(_REDACTED, text) if password_pattern is not None else text
return _secret_shape_redactor()(result)
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
if isinstance(command, tuple):
return tuple(_redact_credentials(str(argument)) for argument in command)
@ -789,16 +743,23 @@ class ProxyExtrasDBManager:
(connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
client certificate where Prisma reads it as the CA."""
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlparse, urlunparse
parsed: Final = urlparse(url)
netloc: Final = parsed.netloc
userinfo, separator, hostport = netloc.rpartition("@")
user, password_separator, password = userinfo.partition(":")
canonical_userinfo: Final = quote(unquote(user), safe="") + (
password_separator + quote(unquote(password), safe="") if password_separator else ""
)
canonical_netloc: Final = f"{canonical_userinfo}{separator}{hostport}" if separator else netloc
if not parsed.query:
return url
return url if canonical_netloc == netloc else urlunparse(parsed._replace(netloc=canonical_netloc))
pairs: Final = tuple(parse_qsl(parsed.query))
kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
return urlunparse(parsed._replace(netloc=canonical_netloc, query=urlencode(libpq_pairs, quote_via=quote)))
@staticmethod
def _libpq_tls_params(

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

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

1
litellm-rust/AGENTS.md
View file

@ -15,6 +15,7 @@ Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `st
- `crates/<crate>/tests/` is only for tests that drive several modules together end to end, such as a full request through a route. A test of one type or function never goes there
- A helper shared by several inline test modules goes in a `#[cfg(test)] mod test_support;` at the crate root, and only once a second module needs it
- A test for another crate's item belongs in that crate, not in a downstream one
- Tests that call a real provider go in `crates/<crate>/tests/live/<provider>.rs` behind a `tests/live/main.rs`. Each reads the provider's key from the environment through an `Option` fixture and returns early when it is unset, never `#[ignore]`, so `cargo test` stays offline until the key is exported
- Never set `autotests = false` or hand-list `[[test]]` targets. Every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests/<name>/mod.rs` or `tests/<subject>/support.rs` so it is not picked up as a test crate of its own
## Test fixtures and cases

2
litellm-rust/Cargo.lock generated
View file

@ -3406,6 +3406,7 @@ dependencies = [
name = "litellm-auth-gcp"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"gcp_auth",
"google-cloud-auth",
"http 1.4.2",
@ -4088,6 +4089,7 @@ dependencies = [
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-host",
"litellm-host-native",
"litellm-http",

1
litellm-rust/crates/auth-gcp/Cargo.toml
View file

@ -11,6 +11,7 @@ google-sdk = ["dep:google-cloud-auth", "dep:http"]
[dependencies]
litellm-auth-types.workspace = true
base64.workspace = true
moka.workspace = true
serde.workspace = true
serde_json.workspace = true

98
litellm-rust/crates/auth-gcp/src/auth.rs
View file

@ -1,4 +1,4 @@
use std::{future::Future, path::Path, pin::Pin, sync::Arc};
use std::{future::Future, pin::Pin, sync::Arc};
use gcp_auth::{CustomServiceAccount, TokenProvider};
use litellm_auth_types::{CredentialPlacement, Error, ErrorDetail, http::apply_credential};
@ -8,7 +8,8 @@ use sha2::{Digest, Sha256};
use crate::{
config::{
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project, non_empty_env,
CredentialSource, VertexConfig, credential_source, get_vertex_ai_project,
is_base64_encoded_json, names_a_credentials_file, non_empty_env,
},
constants::{
CLOUD_PLATFORM_SCOPE, GOOGLE_OAUTH_TOKEN_ENDPOINT, VERTEX_AI_API_KEY_ENV,
@ -176,14 +177,7 @@ impl VertexProviderLoader for GcpProviderLoader {
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
let service_account = if Path::new(configured).is_file() {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)?;
Arc::new(service_account)
Arc::new(load_trusted_service_account(configured.expose())?)
}
CredentialSource::ApplicationCredentials(path) => {
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
@ -240,6 +234,23 @@ enum CredentialCacheKey {
Adc,
}
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
Set it to a readable file path, or to the credentials JSON itself.";
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
BASE64_CREDENTIALS.into(),
)));
}
if names_a_credentials_file(configured) {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)
}
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
Error::CredentialAcquisition(litellm_auth_types::ErrorDetail::failed(
"Vertex AI credentials",
@ -264,6 +275,73 @@ mod tests {
assert_eq!(validate_request_credentials(credentials).is_ok(), accepted);
}
#[rstest::rstest]
#[case::missing_file(
"/nonexistent/vertex-key.json",
"failed to open application credentials file"
)]
#[case::malformed_json(
" {\"type\": \"service_account\",",
"failed to deserialize ApplicationCredentials"
)]
#[case::quoted_json(
"'{\"type\": \"service_account\"}'",
"failed to deserialize ApplicationCredentials"
)]
#[tokio::test]
async fn trusted_credentials_are_read_as_a_file_unless_they_are_json(
#[case] configured: &str,
#[case] expected: &str,
) {
let Err(error) = GcpProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
else {
panic!("{configured:?} should not load");
};
assert!(error.to_string().contains(expected), "{error}");
}
#[rstest::rstest]
#[case::padded("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9")]
#[case::quoted_prefixed_unpadded("\"base64:eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCJ9\"")]
#[case::url_safe_wrapped("eyJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsICJr\nZXkiOiAiPz8_In0=")]
#[tokio::test]
async fn base64_encoded_credentials_are_named_as_base64(#[case] configured: &str) {
assert!(
trusted_load_error(configured)
.await
.contains("is base64-encoded JSON, not the JSON itself")
);
}
async fn trusted_load_error(configured: &str) -> String {
match GcpProviderLoader
.load(CredentialSource::Trusted(SecretValue::new(configured)))
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => error.to_string(),
}
}
#[rstest::rstest]
#[tokio::test]
async fn a_relative_key_file_named_like_json_is_read_as_a_file() {
let key = tempfile::Builder::new()
.prefix("{vertex")
.suffix(".json")
.tempfile_in(".")
.unwrap();
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
let absolute_path = std::fs::canonicalize(key.path()).unwrap();
assert_eq!(
trusted_load_error(relative_name).await,
trusted_load_error(absolute_path.to_str().unwrap()).await
);
}
#[rstest::rstest]
fn inline_and_trusted_credentials_never_share_a_cache_entry() {
assert_ne!(

37
litellm-rust/crates/auth-gcp/src/config.rs
View file

@ -1,5 +1,9 @@
use std::{collections::BTreeMap, path::Path};
use base64::{
Engine, alphabet,
engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
};
use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced};
use serde_json::{Map, Value};
@ -80,9 +84,10 @@ pub fn get_vertex_ai_project_from_credentials(
CredentialSource::Inline(configured) => configured.expose().to_string(),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
match std::fs::read_to_string(configured) {
Ok(contents) if Path::new(configured).is_file() => contents,
_ => configured.to_string(),
if names_a_credentials_file(configured) {
std::fs::read_to_string(configured).ok()?
} else {
configured.to_string()
}
}
CredentialSource::ApplicationCredentials(path) => std::fs::read_to_string(path).ok()?,
@ -105,6 +110,32 @@ pub fn get_vertex_ai_location(
.or_else(|| VertexParams::LOCATION.resolve(&|_| None, env_lookup))
}
const INLINE_JSON_LEADING_NOISE: [char; 7] = [' ', '\t', '\r', '\n', '\u{feff}', '"', '\''];
fn is_inline_json(configured: &str) -> bool {
configured
.trim_start_matches(INLINE_JSON_LEADING_NOISE)
.starts_with('{')
}
pub(crate) fn names_a_credentials_file(configured: &str) -> bool {
!is_inline_json(configured) || Path::new(configured).is_file()
}
const BASE64_ANY_PADDING: GeneralPurpose = GeneralPurpose::new(
&alphabet::STANDARD,
GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
);
pub(crate) fn is_base64_encoded_json(configured: &str) -> bool {
let compact: String = configured.split_whitespace().collect();
let quoted = compact.trim_matches(['"', '\'']);
let encoded = quoted.strip_prefix("base64:").unwrap_or(quoted);
BASE64_ANY_PADDING
.decode(encoded.replace('-', "+").replace('_', "/"))
.is_ok_and(|decoded| serde_json::from_slice::<Map<String, Value>>(&decoded).is_ok())
}
#[derive(Clone, Debug)]
pub enum CredentialSource {
Inline(SecretValue),

19
litellm-rust/crates/auth-gcp/tests/vertex_config.rs
View file

@ -93,3 +93,22 @@ fn project_is_read_from_the_credentials_that_would_authenticate() {
None
);
}
#[rstest]
fn project_is_read_from_a_relative_key_file_named_like_json() {
let key = tempfile::Builder::new()
.prefix("{vertex")
.suffix(".json")
.tempfile_in(".")
.unwrap();
std::fs::write(key.path(), r#"{"project_id":"file-project"}"#).unwrap();
let relative_name = key.path().file_name().unwrap().to_str().unwrap();
let from_file = VertexConfig::from_params(&VertexParams {
vertex_credentials: Some(relative_name.into()),
..VertexParams::default()
});
assert_eq!(
get_vertex_ai_project_from_credentials(&from_file, &|_| None).as_deref(),
Some("file-project")
);
}

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",

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,

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

@ -124,7 +124,7 @@ fn prepare_provider_request(
url,
body: transformed,
environment,
timeout,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
api_key: api_key.map(SecretValue::new),
})
}
@ -253,6 +253,41 @@ mod tests {
);
}
#[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),

14
litellm-rust/crates/inference-messages/src/types.rs
View file

@ -60,6 +60,16 @@ pub struct MessagesSettings {
pub reasoning_auto_summary: bool,
#[serde(default)]
pub additional_drop_params: Vec<String>,
#[serde(default)]
pub request_timeout: Option<f64>,
}
impl MessagesSettings {
pub fn request_timeout(&self) -> Option<Duration> {
self.request_timeout
.and_then(|secs| Duration::try_from_secs_f64(secs).ok())
.filter(|timeout| !timeout.is_zero())
}
}
#[cfg(test)]
@ -120,13 +130,15 @@ mod tests {
},
"drop_params": true,
"reasoning_auto_summary": true,
"additional_drop_params": ["metadata.user_id", "thinking"]
"additional_drop_params": ["metadata.user_id", "thinking"],
"request_timeout": 1800.0
}),
MessagesShaping {
settings: MessagesSettings {
drop_params: true,
reasoning_auto_summary: true,
additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()],
request_timeout: Some(1800.0),
},
capabilities: MessagesModelCapabilities {
supports_reasoning: true,

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) {

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

@ -10,6 +10,7 @@ bytes.workspace = true
futures-util.workspace = true
litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
litellm-cache-response.workspace = true
litellm-core-utils.workspace = true
litellm-host.workspace = true
litellm-http.workspace = true
litellm-inference.workspace = true

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

@ -1,7 +1,11 @@
use litellm_host::interceptors::RequestContext;
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;
@ -16,10 +20,10 @@ pub(super) async fn prepare(
call: ResponsesCall,
secrets: &dyn SecretSource,
) -> Result<ProviderResponsesRequest, Error> {
let identity = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
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 config: &'static dyn BaseResponsesApiConfig = &OpenAiResponsesApiConfig;
let snapshot = secrets
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
.await?;
@ -52,20 +56,107 @@ pub(super) async fn prepare(
})
}
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<litellm_host::interceptors::ProviderIdentity, Error> {
let provider = custom_llm_provider.unwrap_or("openai");
if provider != "openai" {
return Err(Error::Unsupported("native HTTP responses provider"));
}
let resolved = model.strip_prefix("openai/").unwrap_or(model);
if resolved.is_empty() || resolved.contains('/') {
) -> 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()));
}
Ok(litellm_host::interceptors::ProviderIdentity {
model: resolved.into(),
provider: provider.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(_))
));
}
}

7
litellm-rust/crates/inference-responses/tests/live/main.rs Normal file
View file

@ -0,0 +1,7 @@
//! Real provider calls, one module per provider. Each test skips when its provider's key is not
//! in the environment, so `cargo test` stays offline by default.
#[path = "../support/mod.rs"]
mod support;
mod openrouter;

98
litellm-rust/crates/inference-responses/tests/live/openrouter.rs Normal file
View file

@ -0,0 +1,98 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_responses::types::{ResponsesCall, ResponsesOutput};
use litellm_inference_testing::RecordingSecrets;
use litellm_llms::base_llm::responses::streaming::decode;
use litellm_llms_types::formats::responses::streaming::ResponsesEventType;
use rstest::{fixture, rstest};
use serde_json::json;
use crate::support::responses_route;
const DEFAULT_MODEL: &str = "openrouter/openai/gpt-5.5";
/// `OPENROUTER_API_KEY` as a secret source; tests skip when it is unset.
#[fixture]
fn secrets() -> Option<Arc<RecordingSecrets>> {
let key = std::env::var("OPENROUTER_API_KEY").ok()?;
Some(Arc::new(RecordingSecrets::new([(
"OPENROUTER_API_KEY",
key.as_str(),
)])))
}
fn call(stream: bool) -> ResponsesCall {
ResponsesCall {
model: std::env::var("OPENROUTER_LIVE_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()),
input: json!("Reply with the single word: pong"),
optional_params: json!({"max_output_tokens": 32, "stream": stream})
.as_object()
.unwrap()
.clone(),
api_key: None,
api_base: None,
custom_llm_provider: None,
extra_headers: None,
timeout: Some(Duration::from_secs(60)),
}
}
#[rstest]
#[tokio::test]
async fn a_complete_response_carries_output_and_usage(secrets: Option<Arc<RecordingSecrets>>) {
let Some(secrets) = secrets else { return };
let output = responses_route(secrets)
.execute(call(false), &(), None)
.await
.unwrap();
let ResponsesOutput::Complete(response) = output else {
panic!("expected a complete response");
};
let body = serde_json::to_value(response).unwrap();
let text = body["output"]
.as_array()
.into_iter()
.flatten()
.filter_map(|item| item["content"].as_array())
.flatten()
.filter_map(|part| part["text"].as_str())
.collect::<String>();
assert!(text.to_ascii_lowercase().contains("pong"), "{body}");
assert!(
body["usage"]["total_tokens"].as_u64().unwrap() > 0,
"{body}"
);
}
#[rstest]
#[tokio::test]
async fn a_stream_ends_with_response_completed(secrets: Option<Arc<RecordingSecrets>>) {
let Some(secrets) = secrets else { return };
let output = responses_route(secrets)
.execute(call(true), &(), None)
.await
.unwrap();
let ResponsesOutput::Stream { chunks, .. } = output else {
panic!("expected SSE");
};
let wire = chunks.try_collect::<Vec<_>>().await.unwrap().concat();
let events = decode(&wire).unwrap();
assert_eq!(
events.first().map(|event| &event.event_type),
Some(&ResponsesEventType::ResponseCreated)
);
assert!(
events
.iter()
.any(|event| event.event_type == ResponsesEventType::OutputTextDelta)
);
let last = events.last().unwrap();
assert_eq!(last.event_type, ResponsesEventType::ResponseCompleted);
assert!(
last.data["response"]["usage"]["total_tokens"]
.as_u64()
.unwrap()
> 0
);
}

150
litellm-rust/crates/inference-responses/tests/openrouter.rs Normal file
View file

@ -0,0 +1,150 @@
use std::sync::Arc;
use futures_util::TryStreamExt;
use litellm_inference_responses::{
Error,
types::{ResponsesCall, ResponsesOutput},
};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use rstest::{fixture, rstest};
use serde_json::json;
use wiremock::ResponseTemplate;
mod support;
use support::*;
#[fixture]
fn call() -> ResponsesCall {
ResponsesCall {
model: "openrouter/anthropic/test-model".into(),
input: json!("hello"),
optional_params: Default::default(),
api_key: None,
api_base: None,
custom_llm_provider: None,
extra_headers: None,
timeout: None,
}
}
fn response_body() -> serde_json::Value {
json!({"id": "resp_1", "model": "anthropic/test-model", "output": [{"type": "message", "content": []}], "usage": {"total_tokens": 7, "cost": 0.00017}})
}
#[rstest]
#[case::explicit_key(Some("explicit"), &[], "Bearer explicit")]
#[case::managed_key(None, &[("OPENROUTER_API_KEY", "managed")], "Bearer managed")]
#[case::legacy_name(None, &[("OR_API_KEY", "legacy")], "Bearer legacy")]
#[tokio::test]
async fn the_payload_goes_to_the_openrouter_responses_endpoint_with_a_bearer(
call: ResponsesCall,
#[case] api_key: Option<&str>,
#[case] secrets: &[(&str, &str)],
#[case] expected: &str,
) {
let upstream = upstream([json_response(response_body())]).await;
let output = responses_route(Arc::new(RecordingSecrets::new(secrets.iter().copied())))
.execute(
ResponsesCall {
api_key: api_key.map(str::to_owned),
api_base: Some(format!("{}/api/v1", upstream.uri())),
optional_params: json!({"max_output_tokens": 16})
.as_object()
.unwrap()
.clone(),
..call
},
&(),
None,
)
.await
.unwrap();
let ResponsesOutput::Complete(response) = output else {
panic!("expected a complete response");
};
assert_eq!(serde_json::to_value(response).unwrap(), response_body());
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), "/api/v1/responses");
assert_eq!(request.header("authorization"), Some(expected));
assert_eq!(
request.json(),
json!({"max_output_tokens": 16, "model": "anthropic/test-model", "input": "hello"})
);
}
#[rstest]
#[tokio::test]
async fn the_secret_manager_supplies_the_base_and_key(call: ResponsesCall) {
let upstream = upstream([json_response(response_body())]).await;
let base = format!("{}/api/v1/", upstream.uri());
let secrets = Arc::new(RecordingSecrets::new([
("OPENROUTER_API_KEY", "managed"),
("OPENROUTER_API_BASE", base.as_str()),
]));
responses_route(secrets.clone())
.execute(call, &(), None)
.await
.unwrap();
assert_eq!(
secrets.requested(),
["OPENROUTER_API_KEY", "OR_API_KEY", "OPENROUTER_API_BASE"]
);
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), "/api/v1/responses");
assert_eq!(request.header("authorization"), Some("Bearer managed"));
}
#[rstest]
#[tokio::test]
async fn a_missing_key_fails_before_transport(call: ResponsesCall) {
let upstream = upstream([json_response(response_body())]).await;
let result = responses_route(no_secrets())
.execute(
ResponsesCall {
api_base: Some(upstream.uri()),
..call
},
&(),
None,
)
.await;
assert!(matches!(
result,
Err(Error::Auth(litellm_auth::Error::MissingApiKey {
provider: "OpenRouter",
..
}))
));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn streaming_relays_every_openrouter_event_untouched(call: ResponsesCall) {
let sse = "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\ndata: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"usage\":{\"cost\":0.0002}}}\n\n";
let upstream =
upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await;
let output = responses_route(no_secrets())
.execute(
ResponsesCall {
api_key: Some("key".into()),
api_base: Some(upstream.uri()),
optional_params: json!({"stream": true}).as_object().unwrap().clone(),
..call
},
&(),
None,
)
.await
.unwrap();
let ResponsesOutput::Stream { chunks, .. } = output else {
panic!("expected SSE");
};
assert_eq!(
chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
sse.as_bytes()
);
let request = only_request(&upstream).await;
assert_eq!(request.json()["stream"], true);
assert_eq!(request.json()["model"], "anthropic/test-model");
}

11
litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs
View file

@ -195,9 +195,10 @@ impl AnthropicBeta {
fn rejected_by(&self) -> &'static [BetaProvider] {
match self {
Self::AdvancedToolUse20251120 | Self::ContextManagement20250627 => {
&[BetaProvider::BedrockConverse]
Self::AdvancedToolUse20251120 => {
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
}
Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse],
Self::AdvisorTool20260301 | Self::Compact20260904 => &[
BetaProvider::AzureAi,
BetaProvider::BedrockConverse,
@ -231,12 +232,11 @@ impl AnthropicBeta {
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse],
Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[],
Self::DangerousToolUse20260903 => {
&[BetaProvider::BedrockConverse, BetaProvider::Databricks]
}
Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse],
Self::Effort20251124 => &[BetaProvider::VertexAi],
Self::FastMode20260201 | Self::Oauth20250420 => &[
BetaProvider::AzureAi,
@ -244,6 +244,7 @@ impl AnthropicBeta {
BetaProvider::Bedrock,
BetaProvider::BedrockMantle,
BetaProvider::VertexAi,
BetaProvider::Databricks,
],
Self::FineGrainedToolStreaming20250514 => {
&[BetaProvider::AzureAi, BetaProvider::VertexAi]

1
litellm-rust/crates/llms/AGENTS.md
View file

@ -58,6 +58,7 @@
- `AnthropicConfig`, `AmazonConverseConfig` and the non-OCR base traits are partial. Preserve their acceptance gates, passthrough and host fallback contracts
- `OpenAiResponsesApiConfig` implements WebSocket transformations and direct HTTP Responses, without Python's model-specific param rewriting or Responses-to-Chat emulation
- OpenAI-compatible Responses hosts (`openrouter`) reuse the OpenAI request and response transforms as free functions and own only their credentials and base URL
- Chat `supported_openai_param_mappings` lists OpenAI/provider name pairs, unlike Python's name-only list. Audio `map_transcription_params` is a Rust filtering helper
- Auth resolution is split between configs and route preparation in `inference-<format>`

2
litellm-rust/crates/llms/src/anthropic/beta_headers.rs
View file

@ -76,7 +76,7 @@ mod tests {
#[case::vertex_ai(BetaProvider::VertexAi, Some("tool-search-tool-2025-10-19"))]
#[case::anthropic(BetaProvider::Anthropic, Some("advanced-tool-use-2025-11-20"))]
#[case::azure_ai(BetaProvider::AzureAi, Some("advanced-tool-use-2025-11-20"))]
#[case::databricks(BetaProvider::Databricks, Some("advanced-tool-use-2025-11-20"))]
#[case::databricks(BetaProvider::Databricks, None)]
#[case::bedrock_converse(BetaProvider::BedrockConverse, None)]
fn filter_renames_per_host(
#[case] provider: BetaProvider,

91
litellm-rust/crates/llms/src/anthropic/messages/transformation.rs
View file

@ -109,7 +109,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
}
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
update_headers_with_anthropic_beta(headers, request)
merge_beta_headers(
headers,
provider_feature_betas(request, BetaProvider::Anthropic),
)
}
}
@ -136,12 +139,20 @@ pub(crate) fn transform_messages_request(
messages: strip_encrypted_reasoning_blocks(messages),
params: MessagesOptionalParams {
context_management,
extra: without_null_safeguards(request.params.extra),
..request.params
},
..request
})
}
fn without_null_safeguards(extra: Map<String, Value>) -> Map<String, Value> {
extra
.into_iter()
.filter(|(key, value)| key != "safeguards" || !value.is_null())
.collect()
}
pub(crate) fn update_headers_with_anthropic_beta(
headers: Headers,
request: &MessagesRequest,
@ -149,12 +160,34 @@ pub(crate) fn update_headers_with_anthropic_beta(
merge_beta_headers(headers, feature_betas(request))
}
/// The betas a request needs on a host other than the first-party API: the Anthropic set
/// filtered and renamed through the host's `BetaProvider` policy.
pub(crate) fn update_headers_with_safeguards_beta(
headers: Headers,
request: &MessagesRequest,
provider: BetaProvider,
) -> Headers {
merge_beta_headers(
headers,
feature_betas(request).union(on_provider(mapped_host_betas(&request.params), provider)),
)
}
/// The betas a request needs on a host with a column in the beta map, filtered and renamed
/// through that column.
pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet {
feature_betas(request)
.iter()
.filter_map(|beta| beta.on(provider))
on_provider(
feature_betas(request).union(mapped_host_betas(&request.params)),
provider,
)
}
fn on_provider(betas: BetaSet, provider: BetaProvider) -> BetaSet {
betas.iter().filter_map(|beta| beta.on(provider)).collect()
}
fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet {
carries_safeguards(params)
.then_some(AnthropicBeta::DangerousToolUse20260903)
.into_iter()
.collect()
}
@ -206,6 +239,13 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool {
.is_some_and(|config| config.format.is_some())
}
fn carries_safeguards(params: &MessagesOptionalParams) -> bool {
params
.extra
.get("safeguards")
.is_some_and(|value| !value.is_null())
}
fn messages_carry_output_config(messages: &[Message]) -> bool {
messages
.iter()
@ -1025,10 +1065,44 @@ mod tests {
json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}),
&["per-turn-control-2026-07-01"]
)]
#[case::safeguards(
json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}),
&[]
)]
fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) {
assert_eq!(feature_betas(&request(fields)), betas(expected));
}
#[rstest]
#[case::anthropic(BetaProvider::Anthropic, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])]
#[case::vertex_ai(BetaProvider::VertexAi, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])]
#[case::empty_list(BetaProvider::Anthropic, json!([]), &["dangerous-tool-use-2026-09-03"])]
#[case::not_a_list(BetaProvider::Anthropic, json!(1), &["dangerous-tool-use-2026-09-03"])]
#[case::null(BetaProvider::Anthropic, Value::Null, &[])]
#[case::bedrock_converse_has_a_null_cell(BetaProvider::BedrockConverse, json!([{"type": "dangerous_tool_use"}]), &[])]
fn safeguards_add_the_dangerous_tool_use_beta_where_the_beta_map_has_a_cell(
#[case] provider: BetaProvider,
#[case] safeguards: Value,
#[case] expected: &[&str],
) {
assert_eq!(
provider_feature_betas(&request(json!({"safeguards": safeguards})), provider),
betas(expected)
);
}
#[rstest]
fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() {
let client = headers(&[("anthropic-beta", "claude-code-20250219")]);
assert_eq!(
update_headers_with_anthropic_beta(
client.clone(),
&request(json!({"safeguards": [{"type": "dangerous_tool_use"}]}))
),
client
);
}
#[rstest]
#[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])]
#[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])]
@ -1067,6 +1141,11 @@ mod tests {
json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
&[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")],
)]
#[case::safeguards_add_the_dangerous_tool_use_beta_to_the_client_betas(
&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")],
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
&[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14")],
)]
fn request_headers_merge_the_feature_betas(
#[case] input: &[(&str, &str)],
#[case] fields: Value,

12
litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs
View file

@ -2,7 +2,7 @@ use litellm_auth::SecretValue;
use litellm_http::request::{has_bearer_auth, has_header};
use litellm_llms_types::{
formats::messages::MessagesRequest,
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
};
use litellm_router_types::LitellmParams;
@ -12,7 +12,7 @@ use crate::{
common_utils::API_KEY_PLACEMENT,
messages::{
handler::shape_anthropic_messages_request,
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
transformation::{transform_messages_request, update_headers_with_safeguards_beta},
},
},
azure_ai::common_utils::{
@ -104,7 +104,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
}
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
update_headers_with_anthropic_beta(headers, request)
update_headers_with_safeguards_beta(headers, request, BetaProvider::AzureAi)
}
}
@ -544,6 +544,12 @@ mod tests {
&[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")]
)]
#[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])]
#[case::safeguards_need_the_dangerous_tool_use_beta(
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
&[("anthropic-beta", "claude-code-20250219")],
&[("x-api-key", "k"), ("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")]
)]
#[case::null_safeguards_need_no_beta(json!({"safeguards": null}), &[], &[("x-api-key", "k")])]
fn request_headers_carry_the_anthropic_feature_betas(
#[case] fields: serde_json::Value,
#[case] forwarded: &[(&str, &str)],

21
litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs
View file

@ -6,7 +6,7 @@ use litellm_core_utils::settings::resolve_non_empty;
use litellm_http::request::{header_value, with_header, without_headers};
use litellm_llms_types::{
formats::messages::MessagesRequest,
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
};
use litellm_router_types::LitellmParams;
use serde_json::Value;
@ -15,7 +15,7 @@ use crate::{
Error, ErrorDetail,
anthropic::messages::{
handler::shape_anthropic_messages_request,
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
transformation::{transform_messages_request, update_headers_with_safeguards_beta},
},
base_llm::{
auth::{AuthScheme, Headers, ValidatedEnvironment},
@ -174,7 +174,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()),
None => headers,
};
update_headers_with_anthropic_beta(headers, request)
update_headers_with_safeguards_beta(headers, request, BetaProvider::Anthropic)
}
fn wire_body(&self, body: Value) -> Value {
@ -356,6 +356,21 @@ mod tests {
);
}
#[rstest]
#[case::safeguards(json!([{"type": "dangerous_tool_use"}]), Some("dangerous-tool-use-2026-09-03"))]
#[case::null_safeguards(Value::Null, None)]
fn safeguards_need_the_dangerous_tool_use_beta(
#[case] safeguards: Value,
#[case] expected: Option<&str>,
) {
let request = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": safeguards})).unwrap();
let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &request);
assert_eq!(
header_value(&headers, "anthropic-beta").map(|value| value.to_string()),
expected.map(str::to_string)
);
}
#[rstest]
#[case::default(json!({}), false)]
#[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)]

23
litellm-rust/crates/llms/src/deepseek/messages/transformation.rs
View file

@ -451,4 +451,27 @@ mod tests {
json!({"type": "enabled", "budget_tokens": 1024})
);
}
#[rstest]
#[case::without_client_betas(&[])]
#[case::with_client_betas(&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")])]
fn safeguards_and_client_betas_reach_deepseek_as_sent(#[case] client: &[(&str, &str)]) {
let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]);
let body = json!({
"model": "deepseek-v4-pro",
"max_tokens": 100,
"messages": [{"role": "user", "content": "hi"}],
"safeguards": safeguards
});
let headers: Headers = client
.iter()
.map(|(name, value)| (name.to_string(), value.to_string()))
.collect();
assert_eq!(
DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG
.request_headers(headers.clone(), &request_from(body.clone())),
headers
);
assert_eq!(transformed(body)["safeguards"], safeguards);
}
}

56
litellm-rust/crates/llms/src/openai/responses/transformation.rs
View file

@ -165,34 +165,46 @@ impl BaseResponsesApiConfig for OpenAiResponsesApiConfig {
input: Value,
params: Map<String, Value>,
) -> Result<Value, Error> {
if !input.is_string() && !input.is_array() {
return Err(Error::InvalidRequest(
"responses input must be a string or an array".into(),
));
}
if params
.get("stream")
.is_some_and(|stream| !stream.is_boolean())
{
return Err(Error::InvalidRequest("stream must be a boolean".into()));
}
Ok(Value::Object(
params
.into_iter()
.chain([
("model".into(), Value::String(model.into())),
("input".into(), input),
])
.collect(),
))
responses_request(model, input, params)
}
fn transform_response_api_response(&self, body: Value) -> Result<ResponsesApiResponse, Error> {
serde_json::from_value(body)
.map_err(|error| Error::InvalidResponse(error.to_string().into()))
responses_response(body)
}
}
/// The OpenAI Responses request body, shared by every OpenAI-compatible host.
pub fn responses_request(
model: &str,
input: Value,
params: Map<String, Value>,
) -> Result<Value, Error> {
if !input.is_string() && !input.is_array() {
return Err(Error::InvalidRequest(
"responses input must be a string or an array".into(),
));
}
if params
.get("stream")
.is_some_and(|stream| !stream.is_boolean())
{
return Err(Error::InvalidRequest("stream must be a boolean".into()));
}
Ok(Value::Object(
params
.into_iter()
.chain([
("model".into(), Value::String(model.into())),
("input".into(), input),
])
.collect(),
))
}
pub fn responses_response(body: Value) -> Result<ResponsesApiResponse, Error> {
serde_json::from_value(body).map_err(|error| Error::InvalidResponse(error.to_string().into()))
}
#[cfg(test)]
mod tests {
use super::*;

1
litellm-rust/crates/llms/src/openrouter/mod.rs
View file

@ -1 +1,2 @@
pub mod messages;
pub mod responses;

10
litellm-rust/crates/llms/src/openrouter/responses/AGENTS.md Normal file
View file

@ -0,0 +1,10 @@
# rules
- Send the OpenAI Responses payload to `/api/v1/responses` and relay the OpenAI SSE events
- Authenticate with a bearer token from `OPENROUTER_API_KEY`, falling back to `OR_API_KEY`
- Only `OPENROUTER_API_BASE` overrides the default base, never the OpenAI base variables
- No native WebSocket mode
# references
- https://openrouter.ai/docs/api_reference/responses/overview

1
litellm-rust/crates/llms/src/openrouter/responses/mod.rs Normal file
View file

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

174
litellm-rust/crates/llms/src/openrouter/responses/transformation.rs Normal file
View file

@ -0,0 +1,174 @@
use litellm_auth::{CredentialPlacement, SecretValue};
use litellm_core_utils::settings::resolve_non_empty;
use litellm_llms_types::formats::responses::ResponsesApiResponse;
use serde_json::{Map, Value};
use crate::{
Error,
base_llm::{
auth::{AuthScheme, ValidatedEnvironment},
responses::transformation::BaseResponsesApiConfig,
},
openai::responses::transformation::{responses_request, responses_response},
};
const API_KEY_ENVS: [&str; 2] = ["OPENROUTER_API_KEY", "OR_API_KEY"];
const API_BASE_ENV: &str = "OPENROUTER_API_BASE";
const DEFAULT_BASE: &str = "https://openrouter.ai/api/v1";
const RESPONSES_PATH: &str = "/responses";
pub struct OpenRouterResponsesApiConfig;
pub const OPENROUTER_RESPONSES_CONFIG: OpenRouterResponsesApiConfig = OpenRouterResponsesApiConfig;
impl BaseResponsesApiConfig for OpenRouterResponsesApiConfig {
fn secret_names(
&self,
api_key: Option<&str>,
api_base: Option<&str>,
) -> &'static [&'static str] {
match (
api_key.is_some_and(|key| !key.trim().is_empty()),
api_base.is_some_and(|base| !base.trim().is_empty()),
) {
(true, true) => &[],
(true, false) => &[API_BASE_ENV],
(false, true) => &API_KEY_ENVS,
(false, false) => &["OPENROUTER_API_KEY", "OR_API_KEY", API_BASE_ENV],
}
}
fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
let key = resolve_non_empty(api_key, lookup, &API_KEY_ENVS).ok_or(
litellm_auth::Error::MissingApiKey {
provider: "OpenRouter",
environment_variable: API_KEY_ENVS[0],
},
)?;
Ok(ValidatedEnvironment {
headers: litellm_http::request::with_default_headers(
headers,
&[("content-type", "application/json")],
),
auth: AuthScheme::Credential {
placement: CredentialPlacement::Bearer,
secret: SecretValue::new(key),
},
})
}
fn get_complete_url(
&self,
api_base: Option<&str>,
lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
let base = resolve_non_empty(api_base, lookup, &[API_BASE_ENV])
.unwrap_or_else(|| DEFAULT_BASE.into());
format!("{}{RESPONSES_PATH}", base.trim_end_matches('/'))
}
fn transform_responses_api_request(
&self,
model: &str,
input: Value,
params: Map<String, Value>,
) -> Result<Value, Error> {
responses_request(model, input, params)
}
fn transform_response_api_response(&self, body: Value) -> Result<ResponsesApiResponse, Error> {
responses_response(body)
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn env(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option<String> {
move |name| {
values
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| (*value).to_owned())
}
}
#[rstest]
#[case::default(None, &[], "https://openrouter.ai/api/v1/responses")]
#[case::explicit(Some("https://proxy.test/api/v1"), &[], "https://proxy.test/api/v1/responses")]
#[case::trailing_slash(Some("https://proxy.test/api/v1/"), &[], "https://proxy.test/api/v1/responses")]
#[case::blank_falls_to_secret(Some(" "), &[("OPENROUTER_API_BASE", "https://secret.test/v1")], "https://secret.test/v1/responses")]
#[case::openai_base_is_ignored(None, &[("OPENAI_BASE_URL", "https://openai.test/v1")], "https://openrouter.ai/api/v1/responses")]
fn url_is_the_base_plus_responses(
#[case] api_base: Option<&str>,
#[case] secrets: &'static [(&'static str, &'static str)],
#[case] expected: &str,
) {
assert_eq!(
OPENROUTER_RESPONSES_CONFIG.get_complete_url(api_base, &env(secrets)),
expected
);
}
#[rstest]
#[case::explicit(Some("explicit"), &[("OPENROUTER_API_KEY", "managed")], "explicit")]
#[case::managed(None, &[("OPENROUTER_API_KEY", "managed"), ("OR_API_KEY", "legacy")], "managed")]
#[case::legacy_name(None, &[("OR_API_KEY", "legacy")], "legacy")]
#[case::blank_explicit(Some(" "), &[("OR_API_KEY", "legacy")], "legacy")]
fn bearer_comes_from_the_call_then_either_secret_name(
#[case] api_key: Option<&str>,
#[case] secrets: &'static [(&'static str, &'static str)],
#[case] expected: &str,
) {
let environment = OPENROUTER_RESPONSES_CONFIG
.validate_environment(vec![], api_key, &env(secrets))
.unwrap();
let AuthScheme::Credential { placement, secret } = environment.auth else {
panic!("expected a credential");
};
assert_eq!(placement, CredentialPlacement::Bearer);
assert_eq!(secret.expose(), expected);
assert_eq!(
environment.headers,
vec![("content-type".to_owned(), "application/json".to_owned())]
);
}
#[rstest]
#[case::nothing(&[])]
#[case::openai_key_only(&[("OPENAI_API_KEY", "openai")])]
fn a_missing_key_names_the_primary_variable(
#[case] secrets: &'static [(&'static str, &'static str)],
) {
assert!(matches!(
OPENROUTER_RESPONSES_CONFIG.validate_environment(vec![], None, &env(secrets)),
Err(Error::Auth(litellm_auth::Error::MissingApiKey {
provider: "OpenRouter",
environment_variable: "OPENROUTER_API_KEY",
}))
));
}
#[rstest]
#[case::both_given(Some("key"), Some("base"), &[])]
#[case::key_given(Some("key"), None, &["OPENROUTER_API_BASE"])]
#[case::base_given(None, Some("base"), &["OPENROUTER_API_KEY", "OR_API_KEY"])]
#[case::neither(None, None, &["OPENROUTER_API_KEY", "OR_API_KEY", "OPENROUTER_API_BASE"])]
fn only_missing_settings_are_looked_up(
#[case] api_key: Option<&str>,
#[case] api_base: Option<&str>,
#[case] expected: &[&str],
) {
assert_eq!(
OPENROUTER_RESPONSES_CONFIG.secret_names(api_key, api_base),
expected
);
}
}

19
litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs
View file

@ -213,23 +213,14 @@ fn has_web_search_tool(tools: Option<&[Recognized<MessagesTool>]>) -> bool {
}
/// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature
/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled
/// out.
/// betas under Vertex's policy, plus the web search beta Vertex needs spelled out.
fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet {
let params = &request.params;
provider_feature_betas(request, BetaProvider::VertexAi).union(
[
has_web_search_tool(params.tools.as_deref())
.then_some(AnthropicBeta::WebSearch20250305),
params
.extra
.get("safeguards")
.is_some_and(|value| !value.is_null())
.then_some(AnthropicBeta::DangerousToolUse20260903),
]
.into_iter()
.flatten()
.collect(),
has_web_search_tool(params.tools.as_deref())
.then_some(AnthropicBeta::WebSearch20250305)
.into_iter()
.collect(),
)
}

1
litellm-rust/crates/model-catalog/AGENTS.md
View file

@ -2,5 +2,4 @@
For `model_prices_and_context_window.json` validation, we should eventually:
- Remove any schema file like `model_prices_and_context_window.schema.json`
- Stop skipping this crate's tests

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

@ -23,7 +23,7 @@ fn run_chat_completions(
) -> PyResult<Py<PyAny>> {
let host = InferenceHost::new(
call.view()?,
"litellm.rust_bridge.chat_completions.route_host",
"litellm.rust_bridge.host.routes.chat_completions",
);
run_inference::<ChatCompletionsRoute, _>(
py,

10
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -21,10 +21,7 @@ use serde_json::{Map, Value};
use super::{NativeCall, RequestView};
use crate::{
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{
RouteOptions, optional_timeout, project_optional_fields, public_response,
python_timeout_seconds,
},
marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds},
};
pub(super) struct InferenceHost {
@ -115,10 +112,7 @@ impl InferenceHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<Map<String, Value>> {
let names: Vec<String> = py.import(self.module)?.getattr("PARAMETERS")?.extract()?;
project_optional_fields(names.iter().map(String::as_str), |name| {
self.argument(py, arguments, name)
})
self.view.body_params(py, arguments, self.module)
}
pub fn response(&self, py: Python<'_>, response: &impl Serialize) -> PyResult<Py<PyAny>> {

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

@ -26,34 +26,9 @@ use crate::{
routes::RequestView,
};
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host";
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages";
const REQUEST_ERROR_MARKER: &str = "messages_request_error";
const BODY_FIELDS: [&str; 22] = [
"max_tokens",
"metadata",
"stop_sequences",
"stream",
"system",
"temperature",
"thinking",
"tool_choice",
"tools",
"top_k",
"inference_geo",
"top_p",
"mcp_servers",
"context_management",
"compaction",
"container",
"output_format",
"speed",
"output_config",
"cache_control",
"reasoning_effort",
"safeguards",
];
fn merge_headers(
forwarded: Option<Map<String, Value>>,
extra_headers: Option<Map<String, Value>>,
@ -143,7 +118,7 @@ impl MessagesPythonHost {
let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?;
let messages =
argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?;
let fields = project_optional_fields(BODY_FIELDS, argument)?;
let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?;
let body = [
("model".to_string(), Value::String(model.clone())),
("messages".to_string(), from_py(&messages)?),

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

@ -13,6 +13,7 @@ use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCal
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
use litellm_host_python::{
CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
from_py,
};
use litellm_host_python::{Error as HostError, present};
use pyo3::{
@ -20,6 +21,7 @@ use pyo3::{
prelude::*,
types::{PyDict, PyMapping, PyTuple},
};
use serde_json::{Map, Value};
/// The public call as Python bound it: `base` holds the positional arguments by name plus
/// the signature defaults, `kwargs` the caller's keyword dict.
@ -92,6 +94,22 @@ impl RequestView {
present(arguments, self.request(py)?, name)
}
/// The body params `module.body_params` keeps from the keywords the call would send, so
/// each route's Python host decides which params reach the provider.
pub(crate) fn body_params(
&self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
module: &str,
) -> PyResult<Map<String, Value>> {
let request = effective_py_args(self.request(py)?, arguments)?;
from_py(
&py.import(module)?
.getattr("body_params")?
.call1((request,))?,
)
}
pub(crate) fn close(&mut self) {
self.0 = None;
}

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

@ -77,7 +77,7 @@ impl OcrPythonHost {
_ => "",
};
let mapped = py
.import("litellm.rust_bridge.ocr.route_host")
.import("litellm.rust_bridge.host.routes.ocr")
.and_then(|module| module.getattr("map_failure"))
.and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider)))
.and_then(|mapped| mapped.extract::<Py<PyBaseException>>().map_err(PyErr::from));
@ -115,7 +115,7 @@ impl PythonBinding for OcrPythonHost {
py: Python<'_>,
response: LiteLLMOcrResponse,
) -> PyResult<Py<PyAny>> {
public_response(py, "litellm.rust_bridge.ocr.route_host", &response)
public_response(py, "litellm.rust_bridge.host.routes.ocr", &response)
}
fn encode_stream_head(

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

@ -21,7 +21,7 @@ use super::{
};
use crate::errors::RustBridgeDeclined;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host";
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let resolved = call.resolved()?;

46
litellm/__init__.py
View file

@ -1,34 +1,13 @@
### Hide pydantic namespace conflict warnings globally ###
from __future__ import annotations
import warnings
from litellm._process_defaults import apply_process_defaults as _apply_process_defaults
_apply_process_defaults()
warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*")
# Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances
# This warning can accumulate during streaming and cause memory leaks
warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*")
# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it
# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked
warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*")
### INIT VARIABLES #########################
import threading
import os
# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available
import dotenv as _dotenv
def _dev_env_hot_reload_enabled() -> bool:
"""The proxy exports this flag when started with ``--reload``. A reloaded
worker is a fresh process that inherits the reloader's environment, so an
edited ``.env`` value stays masked by the stale inherited one unless we
let the file win; overriding makes the edit take effect on reload."""
return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
from collections.abc import Mapping, Sequence
from typing import (
Any,
@ -118,6 +97,7 @@ success_callback: List[CALLBACK_TYPES] = []
failure_callback: List[CALLBACK_TYPES] = []
service_callback: List[CALLBACK_TYPES] = []
audit_log_callbacks: List[CALLBACK_TYPES] = []
batch_line_item_callbacks: Sequence[CALLBACK_TYPES] = ()
# logging_callback_manager is lazy-loaded via __getattr__
_custom_logger_compatible_callbacks_literal = Literal[
"lago",
@ -206,7 +186,7 @@ _async_failure_callback: List[Union[str, Callable[..., object], "CustomLogger"]]
[]
) # internal variable - async custom callbacks are routed here.
pre_call_rules: List[Callable] = []
post_call_rules: List[Callable] = []
post_call_rules: List[Callable[..., object]] = []
turn_off_message_logging: Optional[bool] = False
standard_logging_payload_excluded_fields: Optional[List[str]] = (
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
@ -392,6 +372,7 @@ key_alias_pattern: str | None = None
enable_gemini_default_thinking_level_low: bool = (
False # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping
)
gemini_include_thoughts_default: bool = True
####################
logging: bool = True
enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
@ -695,6 +676,7 @@ qwen_ai_platform_models: Set = set()
moonshot_models: Set = set()
publicai_models: Set = set()
darkbloom_models: Set = set()
openinfer_models: Final[Set[str]] = set()
v0_models: Set = set()
morph_models: Set = set()
lambda_ai_models: Set = set()
@ -761,7 +743,7 @@ def is_openai_finetune_model(key: str) -> bool:
def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None:
for key, value in model_cost_map.items():
for key, value in tuple(model_cost_map.items()):
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
elif value.get("litellm_provider") == "text-completion-openai":
@ -964,6 +946,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
publicai_models.add(key)
elif value.get("litellm_provider") == "darkbloom":
darkbloom_models.add(key)
elif value.get("litellm_provider") == "openinfer":
openinfer_models.add(key)
elif value.get("litellm_provider") == "v0":
v0_models.add(key)
elif value.get("litellm_provider") == "morph":
@ -1132,6 +1116,7 @@ model_list = list(
| moonshot_models
| publicai_models
| darkbloom_models
| openinfer_models
| v0_models
| morph_models
| lambda_ai_models
@ -1175,7 +1160,9 @@ def _build_models_by_provider() -> dict:
"edenai": edenai_models,
"datarobot": datarobot_models,
"vertex_ai": vertex_chat_models
| vertex_code_chat_models
| vertex_text_models
| vertex_code_text_models
| vertex_anthropic_models
| vertex_vision_models
| vertex_language_models
@ -1183,7 +1170,7 @@ def _build_models_by_provider() -> dict:
| vertex_minimax_models
| vertex_moonshot_models
| vertex_zai_models,
"ai21": ai21_models,
"ai21": ai21_models | ai21_chat_models,
"bedrock": bedrock_models | bedrock_converse_models,
"petals": petals_models,
"ollama": ollama_models,
@ -1246,6 +1233,7 @@ def _build_models_by_provider() -> dict:
"moonshot": moonshot_models,
"publicai": publicai_models,
"darkbloom": darkbloom_models,
"openinfer": openinfer_models,
"v0": v0_models,
"morph": morph_models,
"lambda_ai": lambda_ai_models,
@ -1414,6 +1402,7 @@ from .exceptions import (
InvalidRequestError,
BadRequestError,
ImageFetchError,
MalformedRequestFieldError as MalformedRequestFieldError,
VectorStoreSearchError,
NotFoundError,
PermissionDeniedError,
@ -1689,6 +1678,9 @@ if TYPE_CHECKING:
from .llms.perplexity.decisions.transformation import (
PerplexityDecisionsConfig as PerplexityDecisionsConfig,
)
from .llms.nace.decisions.transformation import (
NaceDecisionsConfig as NaceDecisionsConfig,
)
from .llms.typesafe.decisions.transformation import (
TypeSafeDecisionsConfig as TypeSafeDecisionsConfig,
)

10
litellm/_internal_context.py
View file

@ -37,6 +37,16 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families"
_service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
@contextmanager
def internal_sub_call() -> Generator[None]:
"""A nested call the parent call logs and bills, so the parent must fold this call's cost into its own."""
token: Final = is_internal_call.set(True)
try:
yield
finally:
is_internal_call.reset(token)
@contextmanager
def post_response_phase() -> Generator[None]:
"""Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns."""

2
litellm/_lazy_imports_registry.py
View file

@ -162,6 +162,7 @@ LLM_CONFIG_NAMES: Final = (
"HostedVLLMRerankConfig",
"PerplexityDecisionsConfig",
"TypeSafeDecisionsConfig",
"NaceDecisionsConfig",
"OpenRouterDecisionsConfig",
"FireworksAIDecisionsConfig",
"CloudflareDecisionsConfig",
@ -729,6 +730,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"HostedVLLMRerankConfig",
),
"PerplexityDecisionsConfig": (".llms.perplexity.decisions.transformation", "PerplexityDecisionsConfig"),
"NaceDecisionsConfig": (".llms.nace.decisions.transformation", "NaceDecisionsConfig"),
"TypeSafeDecisionsConfig": (".llms.typesafe.decisions.transformation", "TypeSafeDecisionsConfig"),
"OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
"FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"),

19
litellm/_logging.py
View file

@ -6,12 +6,14 @@ import logging
import os
import re
import sys
from collections.abc import Iterator
from collections.abc import Iterator, Mapping
from datetime import datetime
from logging import Formatter
from typing import Final, TextIO
from urllib.parse import unquote
from pydantic import TypeAdapter
import litellm
from litellm.constants import (
LITELLM_TRUNCATED_PAYLOAD_FIELD,
@ -84,6 +86,7 @@ _redact_string = redact_string
_REDACTED_RECORD_ATTR: Final = "litellm_redacted"
_REDACTED_STAMP: Final = object()
_DECODED_LOG_PAYLOAD: Final = TypeAdapter(object)
_UNREDACTED_SCALAR_TYPES: Final = (bool, int, float, type(None))
@ -91,6 +94,10 @@ def _is_redacted(record: logging.LogRecord) -> bool:
return getattr(record, _REDACTED_RECORD_ATTR, None) is _REDACTED_STAMP
def _record_attributes(record: logging.LogRecord) -> Mapping[str, object]:
return record.__dict__
def _scrubbing_changed_nothing(scrubbed: object, original: object) -> bool:
try:
return bool(scrubbed == original)
@ -438,7 +445,7 @@ def _process_record(record: logging.LogRecord, *, base64_limit: int, text_limit:
key, substituted_color if key == "color_message" and substituted_color is not None else value
),
)
for key, value in record.__dict__.items()
for key, value in _record_attributes(record).items()
if key not in _STANDARD_RECORD_ATTRS
and key != _REDACTED_RECORD_ATTR
and not isinstance(value, _UNREDACTED_SCALAR_TYPES)
@ -669,7 +676,7 @@ def _try_parse_json_message(message: str) -> dict[str, object] | None:
msg_stripped: Final = message.strip()
if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")):
return None
parsed: Final[object] = safe_json_loads(message, default=None)
parsed: Final = _DECODED_LOG_PAYLOAD.validate_python(safe_json_loads(message, default=None))
if parsed is None or not isinstance(parsed, dict):
return None
return parsed
@ -699,7 +706,7 @@ def _try_parse_embedded_python_dict(message: str) -> dict[str, object] | None:
if depth == 0:
substr = message[start : j + 1]
try:
result: object = ast.literal_eval(substr)
result: object = _DECODED_LOG_PAYLOAD.validate_python(ast.literal_eval(substr))
if isinstance(result, dict) and len(result) > 0:
return result
except (ValueError, SyntaxError, TypeError):
@ -757,7 +764,7 @@ class JsonFormatter(Formatter):
json_record[key] = value
# Include extra attributes passed via logger.debug("msg", extra={...})
for key, value in record.__dict__.items():
for key, value in _record_attributes(record).items():
if key not in _NON_EXTRA_RECORD_ATTRS and key not in json_record:
json_record[key] = value
@ -993,6 +1000,8 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
lg.handlers.clear() # remove any existing handlers
lg.addHandler(handler) # add JSON formatter handler
lg.propagate = False # prevent bubbling to parent/root
if (configured_level := os.getenv("LITELLM_LOG")) is not None:
logging.getLogger().setLevel(resolve_log_level(configured_level))
def _get_uvicorn_json_log_config():

24
litellm/_process_defaults.py Normal file
View file

@ -0,0 +1,24 @@
import os
import warnings
import dotenv
def _dev_env_hot_reload_enabled() -> bool:
"""The proxy exports this flag when started with ``--reload``. A reloaded
worker is a fresh process that inherits the reloader's environment, so an
edited ``.env`` value stays masked by the stale inherited one unless we
let the file win; overriding makes the edit take effect on reload."""
return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
def apply_process_defaults() -> None:
warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*")
# Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances
# This warning can accumulate during streaming and cause memory leaks
warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*")
# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it
# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked
warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*")
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())

14
litellm/_redis.py
View file

@ -19,6 +19,7 @@ from urllib.parse import urlsplit, urlunsplit
import redis
import redis.asyncio as async_redis
from pydantic import TypeAdapter
from redis.credentials import CredentialProvider
from litellm import get_secret, get_secret_str
@ -40,6 +41,7 @@ from ._logging import verbose_logger
_generate_gcp_iam_access_token = generate_gcp_iam_access_token
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
_UNWRAPPED_INIT: Final = TypeAdapter(object)
_AWS_IAM_KWARG_NAMES: Final = (
"aws_iam_auth",
@ -64,7 +66,7 @@ def _unwrapped_init_args(cls: type) -> frozenset[str]:
``@deprecated_args`` too, which the same class of bug as ``_init_arg_names``
would otherwise silently empty this allowlist through (see its docstring).
"""
spec: Final = inspect.getfullargspec(inspect.unwrap(cls.__init__))
spec: Final = inspect.getfullargspec(_UNWRAPPED_INIT.validate_python(inspect.unwrap(cls.__init__)))
return frozenset(spec.args + spec.kwonlyargs)
@ -112,7 +114,7 @@ def _init_arg_names(cls: type) -> frozenset[str]:
name
for klass in inspect.getmro(cls)
if klass is not object
for spec in (inspect.getfullargspec(inspect.unwrap(klass.__init__)),)
for spec in (inspect.getfullargspec(_UNWRAPPED_INIT.validate_python(inspect.unwrap(klass.__init__))),)
for name in spec.args + spec.kwonlyargs
)
@ -935,6 +937,10 @@ def get_redis_connection_pool(
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
def _mask_for_display(masker: SensitiveDataMasker, redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
return masker.mask_dict(redis_kwargs)
def _redis_kwargs_for_logging(redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
return {
key: "<credential provider>"
@ -966,7 +972,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
masker = SensitiveDataMasker()
# Mask sensitive data in redis_kwargs
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
masked_redis_kwargs = _mask_for_display(masker, redis_kwargs_for_logging)
# Create main panel title
title: Final = Text("Redis Configuration", style="bold blue")
@ -1029,7 +1035,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
except ImportError:
# Fallback to simple logging if rich is not available
masker = SensitiveDataMasker()
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
masked_redis_kwargs = _mask_for_display(masker, redis_kwargs_for_logging)
verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
except Exception as e:
verbose_logger.error("Error pretty printing Redis configuration: %s", e)

47
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -11,6 +11,7 @@ A2A Streaming Events (in order):
"""
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Any, Final
import litellm
@ -21,6 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
A2AStreamingContext,
)
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS
from litellm.interactions.agents.utils import merge_agent_headers
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import ModelResponse
@ -44,6 +46,32 @@ _AGENT_ONLY_PARAMS: Final = frozenset(
)
def bridge_model_name(litellm_params: Mapping[str, object]) -> str:
"""The litellm model string a bridge agent completes with: ``<custom_llm_provider>/<model>``.
A provider bridged on chat names a real upstream endpoint with ``model``, so a missing one stays empty for that
provider to reject or ignore instead of becoming a made-up ``agent`` endpoint.
"""
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
default_model: Final = "" if custom_llm_provider in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS else "agent"
configured_model: Final = litellm_params.get("model")
model: Final = configured_model if isinstance(configured_model, str) else default_model
if not custom_llm_provider or model.startswith(f"{custom_llm_provider}/"):
return model
return f"{custom_llm_provider}/{model}"
def agent_completion_kwargs(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
"""The agent's litellm_params that are valid litellm.acompletion kwargs (api_key, api_base, ...)."""
return MappingProxyType(
{
key: value
for key, value in litellm_params.items()
if key not in ("model", "custom_llm_provider") and key not in _AGENT_ONLY_PARAMS
}
)
class A2ACompletionBridgeHandler:
"""
Static methods for handling A2A requests via LiteLLM completion.
@ -64,16 +92,7 @@ class A2ACompletionBridgeHandler:
# Transform A2A message to OpenAI format
openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
model: Final[str] = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
full_model: Final = bridge_model_name(litellm_params)
if stream:
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
@ -87,13 +106,7 @@ class A2ACompletionBridgeHandler:
"api_base": api_base,
"stream": stream,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add: Final = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
completion_params.update(agent_completion_kwargs(litellm_params))
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.

10
litellm/a2a_protocol/main.py
View file

@ -14,7 +14,9 @@ import datetime
import uuid
from collections.abc import AsyncIterator, Coroutine, Mapping
from types import MappingProxyType, ModuleType
from typing import TYPE_CHECKING, Any, Final, Optional, cast
from typing import TYPE_CHECKING, Final, Optional, cast
from typing_extensions import TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
@ -89,6 +91,10 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
A2ACardResolver: Final = LiteLLMA2ACardResolver
class _SendMessageOptions(TypedDict, total=False, extra_items=object):
pass
def _set_usage_on_logging_obj(
kwargs: Mapping[str, object],
prompt_tokens: int,
@ -563,7 +569,7 @@ async def asend_message(
def send_message(
a2a_client: "A2AClientType",
request: "SendMessageRequest",
**kwargs: Any,
**kwargs: Unpack[_SendMessageOptions],
) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]:
"""
Sync: Send a message to an A2A agent.

12
litellm/a2a_protocol/streaming_iterator.py
View file

@ -7,16 +7,21 @@ from collections.abc import AsyncIterator
from datetime import datetime
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
if TYPE_CHECKING:
from a2a.compat.v0_3.types import SendStreamingMessageRequest, SendStreamingMessageResponse
_OPAQUE_PAYLOAD: Final = TypeAdapter(object)
class A2AStreamingIterator:
"""
@ -160,8 +165,11 @@ class A2AStreamingIterator:
# Add final chunk result if available
if self.final_chunk:
try:
chunk_dict: Final = self.final_chunk.model_dump(mode="json", exclude_none=True)
result["result"] = chunk_dict.get("result", {})
chunk_dict: Final = as_str_mapping(
_OPAQUE_PAYLOAD.validate_python(self.final_chunk.model_dump(mode="json", exclude_none=True))
)
if chunk_dict is not None:
result["result"] = chunk_dict.get("result", {})
except Exception:
pass

22
litellm/anthropic_beta_headers_config.json
View file

@ -220,36 +220,36 @@
},
"databricks": {
"advisor-tool-2026-03-01": null,
"advanced-tool-use-2025-11-20": "advanced-tool-use-2025-11-20",
"advanced-tool-use-2025-11-20": null,
"bash_20241022": null,
"bash_20250124": null,
"code-execution-2025-08-25": "code-execution-2025-08-25",
"code-execution-2025-08-25": null,
"compact-2026-01-12": "compact-2026-01-12",
"computer-use-2025-01-24": "computer-use-2025-01-24",
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",
"context-management-2025-06-27": "context-management-2025-06-27",
"dangerous-tool-use-2026-09-03": null,
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
"files-api-2025-04-14": "files-api-2025-04-14",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"structured-output-2024-03-01": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14",
"mcp-client-2025-11-20": "mcp-client-2025-11-20",
"mcp-client-2025-04-04": "mcp-client-2025-04-04",
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,
"mcp-servers-2025-12-04": null,
"oauth-2025-04-20": "oauth-2025-04-20",
"oauth-2025-04-20": null,
"output-128k-2025-02-19": "output-128k-2025-02-19",
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
"skills-2025-10-02": "skills-2025-10-02",
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-fetch-2025-09-10": null,
"web-search-2025-03-05": "web-search-2025-03-05"
}
}

45
litellm/anthropic_beta_headers_manager.py
View file

@ -22,17 +22,28 @@ Environment Variables:
- LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional)
"""
import asyncio
import json
import os
import threading
from collections.abc import Mapping
from concurrent.futures import ThreadPoolExecutor
from functools import cache
from importlib.resources import files
from typing import Final
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter
from litellm._logging import verbose_logger
# Cache for the loaded configuration
_BETA_HEADERS_CONFIG: dict | None = None
_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock()
_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config")
BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]]
_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig)
class GetAnthropicBetaHeadersConfig:
@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict:
if _BETA_HEADERS_CONFIG is not None:
return _BETA_HEADERS_CONFIG
# Get the URL from environment or use default
from litellm import anthropic_beta_headers_url
with _BETA_HEADERS_CONFIG_LOCK:
if _BETA_HEADERS_CONFIG is not None:
return _BETA_HEADERS_CONFIG
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
verbose_logger.debug("Loaded and cached beta headers config")
# Get the URL from environment or use default
from litellm import anthropic_beta_headers_url
return _BETA_HEADERS_CONFIG
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
verbose_logger.debug("Loaded and cached beta headers config")
return _BETA_HEADERS_CONFIG
async def aload_beta_headers_config() -> None:
if _BETA_HEADERS_CONFIG is None:
await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config)
@cache
def _bundled_beta_headers_config() -> BetaHeadersConfig:
return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config())
def reload_beta_headers_config() -> dict:
@ -288,6 +313,14 @@ def is_beta_header_supported(
return beta_header in provider_mapping and provider_mapping[beta_header] is not None
def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool:
config: Final = _bundled_beta_headers_config()
aliases: Final = config.get("provider_aliases", {})
canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider
provider_mapping: Final = config.get(canonical_provider, {})
return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None
def get_provider_beta_header(
anthropic_beta_header: str,
provider: str,

5
litellm/anthropic_interface/messages/__init__.py
View file

@ -11,7 +11,6 @@ This is an __init__.py file to allow the following interface
"""
from collections.abc import AsyncIterator, Coroutine, Iterator
from typing import Any
from litellm.messages import (
anthropic_messages as _async_anthropic_messages,
@ -102,8 +101,8 @@ def create(
) -> (
AnthropicMessagesResponse
| Iterator[bytes]
| AsyncIterator[Any]
| Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]]
| AsyncIterator[object]
| Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[object] | Iterator[bytes]]
):
"""
Async wrapper for Anthropic's messages API

8
litellm/assistants/main.py
View file

@ -129,6 +129,7 @@ def get_assistants(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
client=client,
aget_assistants=aget_assistants,
)
@ -312,6 +313,7 @@ def create_assistants(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
create_assistant_data=create_assistant_data,
client=client,
async_create_assistants=async_create_assistants,
@ -461,6 +463,7 @@ def delete_assistant(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
assistant_id=assistant_id,
client=client,
async_delete_assistants=async_delete_assistants,
@ -637,6 +640,7 @@ def create_thread(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
client=client,
acreate_thread=acreate_thread,
)
@ -783,6 +787,7 @@ def get_thread(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
client=client,
aget_thread=aget_thread,
)
@ -960,6 +965,7 @@ def add_message(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
client=client,
a_add_message=a_add_message,
)
@ -1111,6 +1117,7 @@ def get_messages(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
client=client,
aget_messages=aget_messages,
)
@ -1308,6 +1315,7 @@ def run_thread(
timeout=timeout,
max_retries=optional_params.max_retries,
organization=organization,
litellm_params=litellm_params_dict,
client=client,
arun_thread=arun_thread,
event_handler=event_handler,

250
litellm/batches/batch_line_items.py Normal file
View file

@ -0,0 +1,250 @@
"""Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks."""
import asyncio
import hashlib
import re
from collections import Counter
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from itertools import accumulate, chain, count
from typing import Final, Literal, TypeAlias
import litellm
from litellm._logging import verbose_logger
from litellm.batches.batch_utils import batch_cost_and_usage, iter_batch_input_lines, parse_batch_output_line
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class, is_valid_sha256_hash
from litellm.llms.vertex_ai.files.transformation import vertex_batch_row_custom_id
from litellm.types.utils import (
ModelInfo,
StandardBatchLineItem,
StandardBatchLineItemMetadata,
StandardBatchLineItemsPayload,
StandardBatchLineItemUsage,
)
BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"]
_VERTEX_AI_UNKNOWN_CUSTOM_ID: Final = "unknown" # what a Vertex AI batch row that carries no custom_id is read as
_LINE: Final = re.compile(rb"[^\n]+")
@dataclass(frozen=True, slots=True)
class BatchLineItemSource:
batch_id: str
finished_at: int
model: str
custom_llm_provider: str | None
cost_provider: BatchFileProvider
model_info: ModelInfo | None
metadata: StandardBatchLineItemMetadata
def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata:
return StandardBatchLineItemMetadata(
user_api_key_hash=_creator_key_hash(attribution),
user_api_key_alias=_text(attribution, "user_api_key_alias"),
user_api_key_team_id=_text(attribution, "user_api_key_team_id"),
user_api_key_team_alias=_text(attribution, "user_api_key_team_alias"),
user_api_key_user_id=_text(attribution, "user_api_key_user_id"),
user_api_key_org_id=_text(attribution, "user_api_key_org_id"),
)
def _text(attribution: Mapping[str, object], name: str) -> str | None:
value: Final = attribution.get(name)
return value if isinstance(value, str) else None
def _creator_key_hash(attribution: Mapping[str, object]) -> str | None:
"""The creating credential as a SHA-256 hash. A virtual key is saved hashed already, while a credential saved as
is, such as an OAuth2 access token, is hashed here so the object never holds it."""
credential: Final = _text(attribution, "user_api_key_hash")
if credential is None or is_valid_sha256_hash(credential):
return credential
return hashlib.sha256(credential.encode()).hexdigest()
def _line_key(entry: Mapping[str, object]) -> str | None:
"""A line's custom_id. Bedrock calls it recordId, and native Vertex rows carry it in the row key or request labels."""
key: Final = entry.get("custom_id", entry.get("recordId"))
return key if isinstance(key, str) else vertex_batch_row_custom_id(entry)
def batch_line_item_parts(
source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes], max_part_bytes: int
) -> Iterator[StandardBatchLineItemsPayload]:
"""Pair every output and error line with the input line that shares its custom_id, priced one line at a time, in
parts that each close once their raw request and result lines reach max_part_bytes. Input lines stay unparsed until
their result line is reached, so only one part's records are held at a time."""
sized: Final = _sized_line_items(source, input_content, result_contents)
for part in count():
lines = tuple(_until_full(sized, max_part_bytes))
if not lines:
return
yield StandardBatchLineItemsPayload(
batch_id=source.batch_id, finished_at=source.finished_at, part=part, lines=lines
)
def _sized_line_items(
source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes]
) -> Iterator[tuple[StandardBatchLineItem, int]]:
"""Each output and error line's record, with the raw bytes of that line and of the input line it pairs with."""
request_spans: Final = _request_spans_by_custom_id(input_content)
for raw_result in chain.from_iterable(map(iter_batch_input_lines, result_contents)):
entry = parse_batch_output_line(raw_result)
if entry is None:
continue
custom_id = _line_key(entry)
span = _request_span_for(source, request_spans, custom_id)
request = None if span is None else parse_batch_output_line(input_content[span[0] : span[1]])
yield (
_line_item(source, custom_id, request, entry),
len(raw_result) + (0 if span is None else span[1] - span[0]),
)
def _until_full(
sized: Iterator[tuple[StandardBatchLineItem, int]], max_part_bytes: int
) -> Iterator[StandardBatchLineItem]:
"""Take records until their raw bytes reach max_part_bytes, leaving the rest in sized for the next part."""
for line, part_bytes in accumulate(sized, _with_bytes_so_far):
yield line
if part_bytes >= max_part_bytes:
return
def _with_bytes_so_far(
so_far: tuple[StandardBatchLineItem, int], line: tuple[StandardBatchLineItem, int]
) -> tuple[StandardBatchLineItem, int]:
return line[0], so_far[1] + line[1]
def _request_spans_by_custom_id(input_content: bytes) -> Mapping[str | None, tuple[int, int]]:
"""Where each input line starts and ends, by its custom_id. A custom_id that several input lines share, as when a
Vertex AI embeddings entry is split into one row per input, is paired with none of them."""
keyed: Final = tuple(_keyed_spans(input_content))
counts: Final = Counter(key for key, _ in keyed)
return {key: span for key, span in keyed if counts[key] == 1}
def _keyed_spans(input_content: bytes) -> Iterator[tuple[str | None, tuple[int, int]]]:
for line in _LINE.finditer(input_content):
entry = parse_batch_output_line(line[0]) if line[0].strip() else None
if entry is not None:
yield _line_key(entry), line.span()
def _line_item(
source: BatchLineItemSource, custom_id: str | None, request: Mapping[str, object] | None, entry: dict[str, object]
) -> StandardBatchLineItem:
priced: Final = batch_cost_and_usage(
(entry,), source.cost_provider, model_name=source.model, model_info=source.model_info
)
succeeded: Final = priced.successful_requests > 0
return StandardBatchLineItem(
batch_id=source.batch_id,
custom_id=custom_id,
status="success" if succeeded else "failure",
request=request,
response=entry,
usage=StandardBatchLineItemUsage(
prompt_tokens=priced.usage.prompt_tokens,
completion_tokens=priced.usage.completion_tokens,
total_tokens=priced.usage.total_tokens,
)
if succeeded
else None,
response_cost=priced.cost,
model=source.model,
custom_llm_provider=source.custom_llm_provider,
metadata=source.metadata,
)
def _request_span_for(
source: BatchLineItemSource, request_spans: Mapping[str | None, tuple[int, int]], custom_id: str | None
) -> tuple[int, int] | None:
"""Where the input line that shares a result line's custom_id is. A Vertex AI line read as "unknown" pairs with
none."""
if custom_id is None or (source.cost_provider == "vertex_ai" and custom_id == _VERTEX_AI_UNKNOWN_CUSTOM_ID):
return None
return request_spans.get(custom_id)
def without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload:
return StandardBatchLineItemsPayload(
batch_id=payload["batch_id"],
finished_at=payload["finished_at"],
part=payload["part"],
lines=tuple(
StandardBatchLineItem(
batch_id=line["batch_id"],
custom_id=line["custom_id"],
status=line["status"],
request=None,
response=None,
usage=line["usage"],
response_cost=line["response_cost"],
model=line["model"],
custom_llm_provider=line["custom_llm_provider"],
metadata=line["metadata"],
)
for line in payload["lines"]
),
)
async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None:
"""Hand one part's lines to every batch_line_item_callbacks logger. A failing logger is logged and never stops the
rest."""
await asyncio.gather(
*(_deliver(callback, payload, redact_bodies) for callback in litellm.batch_line_item_callbacks)
)
async def _deliver(
callback: litellm.CALLBACK_TYPES, payload: StandardBatchLineItemsPayload, redact_bodies: bool
) -> None:
try:
logger: Final = _resolve_logger(callback)
if logger is None:
verbose_logger.error(
"batch_line_item_callbacks: %r cannot store batch line items, part %d of batch %s was not stored",
callback,
payload["part"],
payload["batch_id"],
)
return
await logger.async_log_batch_line_items(
without_request_and_response(payload)
if redact_bodies or _logger_turned_off_message_logging(logger)
else payload
)
except Exception as error: # noqa: BLE001 # one failing sink must not stop the other sinks or the cost poller
verbose_logger.exception(
"batch_line_item_callbacks: %r failed to store the %d lines of part %d of batch %s: %s",
callback,
len(payload["lines"]),
payload["part"],
payload["batch_id"],
error,
)
def _resolve_logger(callback: litellm.CALLBACK_TYPES) -> CustomLogger | None:
match callback:
case CustomLogger():
return callback
case "s3_v2" | "gcs_bucket":
return init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None)
case _:
return None
def _logger_turned_off_message_logging(logger: CustomLogger) -> bool:
return (
getattr(logger, "turn_off_message_logging", False) is True
or getattr(logger, "message_logging", True) is not True
)

47
litellm/batches/batch_utils.py
View file

@ -1,10 +1,12 @@
import json
from collections.abc import Iterable, Iterator, Mapping
from collections.abc import Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from dataclasses import replace as dataclasses_replace
from enum import Enum
from typing import Any, Final, Literal, cast
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
@ -17,7 +19,7 @@ from litellm.llms.vertex_ai.batches.transformation import (
)
from litellm.types.llms.openai import Batch
from litellm.types.utils import ModelInfo, Usage
from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS
from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS
from litellm.utils import token_counter
@ -35,6 +37,7 @@ class BatchCostUsageResult:
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
_DECODED_BATCH_LINE: Final = TypeAdapter(object)
def _uses_native_vertex_output(
@ -70,7 +73,17 @@ def batch_cost_is_final(batch: Batch) -> bool:
async def calculate_batch_cost_and_usage(
file_content_dictionary: list[dict[str, object]],
file_content_dictionary: Sequence[dict[str, object]],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> BatchCostUsageResult:
"""Calculate the cost and usage of a batch, see batch_cost_and_usage."""
return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info)
def batch_cost_and_usage(
file_content_dictionary: Sequence[dict[str, object]], # mutable-ok: the rows calculate_batch_cost_and_usage takes
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
@ -436,7 +449,7 @@ def _provider_output_file_id(output_file_id: str) -> str:
return extracted
async def _fetch_batch_managed_file_content(
async def fetch_batch_managed_file_content(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
litellm_params: Mapping[str, object] | None = None,
@ -467,6 +480,9 @@ async def _fetch_batch_managed_file_content(
return _file_content.content
_fetch_batch_managed_file_content: Final = fetch_batch_managed_file_content
async def _fetch_batch_output_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai",
@ -553,6 +569,7 @@ def extract_file_access_credentials(
# A federated deployment holds no api_key, so without these the fetch that reads a
# finished batch's output has nothing to authenticate with and its cost is never billed.
*sorted(ANTHROPIC_WIF_KWARGS_KEYS),
*sorted(OPENAI_WIF_KWARGS_KEYS),
)
for key in credential_keys:
if key in litellm_params:
@ -604,14 +621,14 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict[str, object
whole batch's cost accounting.
"""
for line in iter_batch_input_lines(file_content):
entry = _parse_batch_output_line(line)
entry = parse_batch_output_line(line)
if entry is not None:
yield entry
def _parse_batch_output_line(line: bytes) -> dict[str, object] | None:
def parse_batch_output_line(line: bytes) -> dict[str, object] | None:
try:
parsed: Final[object] = json.loads(line)
parsed: Final = _DECODED_BATCH_LINE.validate_python(json.loads(line))
except ValueError as e:
verbose_logger.warning("skipping malformed batch output line: %s", str(e))
return None
@ -777,11 +794,17 @@ def _batch_response_was_successful(
message batch results lines report ``result.type == "succeeded"``; Bedrock
batch output lines report ``modelOutput`` (and no ``error``).
"""
if custom_llm_provider == "anthropic":
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
if custom_llm_provider == "bedrock":
return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
_response: Final[dict[str, object]] = cast( # cast-ok: batch response bodies come from provider JSON
dict[str, object], batch_job_output_file.get("response", None) or {}
envelope_key: Final = "result" if custom_llm_provider == "anthropic" else "response"
envelope: Final = batch_job_output_file.get(envelope_key)
if not isinstance(envelope, Mapping):
if envelope is not None:
verbose_logger.warning("batch output %s must be an object, treating the line as unsuccessful", envelope_key)
return False
response: Final = cast(Mapping[object, object], envelope) # cast-ok: mapping shape validated above
return (
response.get("type") == "succeeded"
if custom_llm_provider == "anthropic"
else response.get("status_code") == 200
)
return _response.get("status_code", None) == 200

16
litellm/batches/main.py
View file

@ -127,7 +127,7 @@ async def acreate_batch(
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
output_expires_after: dict[str, Any] | None = None,
output_expires_after: FileExpiresAfter | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -191,9 +191,9 @@ def create_batch(
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
output_expires_after: dict[str, Any] | None = None,
output_expires_after: FileExpiresAfter | None = None,
**kwargs,
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
"""
Creates and executes a batch from an uploaded file of request
@ -246,7 +246,7 @@ def create_batch(
extra_body=extra_body,
)
if output_expires_after is not None:
_create_batch_request["output_expires_after"] = cast(FileExpiresAfter, output_expires_after)
_create_batch_request["output_expires_after"] = output_expires_after
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
@ -309,6 +309,7 @@ def create_batch(
timeout=timeout,
max_retries=optional_params.max_retries,
_is_async=_is_async,
litellm_params=litellm_params,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@ -471,6 +472,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
organization=organization,
timeout=timeout,
max_retries=optional_params.max_retries,
litellm_params=litellm_params,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@ -568,7 +570,7 @@ def retrieve_batch(
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
"""
Retrieves a batch.
@ -819,6 +821,7 @@ def list_batches(
organization=organization,
timeout=timeout,
max_retries=optional_params.max_retries,
litellm_params=litellm_params,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
@ -941,7 +944,7 @@ def cancel_batch(
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
) -> LiteLLMBatch | Coroutine[object, None, LiteLLMBatch]:
"""
Cancels a batch.
@ -1016,6 +1019,7 @@ def cancel_batch(
organization=organization,
timeout=timeout,
max_retries=optional_params.max_retries,
litellm_params=litellm_params,
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")

9
litellm/budget_manager.py
View file

@ -13,6 +13,8 @@ import threading
import time
from typing import Final, Literal
from pydantic import TypeAdapter
import litellm
from litellm.constants import (
DAYS_IN_A_MONTH,
@ -22,6 +24,8 @@ from litellm.constants import (
)
from litellm.utils import ModelResponse
_BUDGET_SERVER_REPLY: Final = TypeAdapter(object)
class BudgetManager:
def __init__(
@ -211,6 +215,5 @@ class BudgetManager:
elif self.client_type == "hosted":
url: Final = self.api_base + "/set_budget"
data: Final = {"project_name": self.project_name, "user_dict": self.user_dict}
response = litellm.module_level_client.post(url, headers=self.headers, json=data)
response = response.json()
return response
response: Final = litellm.module_level_client.post(url, headers=self.headers, json=data)
return _BUDGET_SERVER_REPLY.validate_python(response.json())

4
litellm/caching/azure_blob_cache.py
View file

@ -65,7 +65,7 @@ class AzureBlobCache(BaseCache):
print_verbose(f"Get Azure Blob Cache: key: {key}")
as_bytes: Final = self.container_client.download_blob(key).readall()
as_str: Final = as_bytes.decode("utf-8")
cached_response: Final = json.loads(as_str)
cached_response: Final[object] = json.loads(as_str)
verbose_logger.debug(
"Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
@ -86,7 +86,7 @@ class AzureBlobCache(BaseCache):
blob: Final = await self.async_container_client.download_blob(key)
as_bytes: Final = await blob.readall()
as_str: Final = as_bytes.decode("utf-8")
cached_response: Final = json.loads(as_str)
cached_response: Final[object] = json.loads(as_str)
verbose_logger.debug(
"Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
key,

27
litellm/caching/caching.py
View file

@ -18,14 +18,14 @@ from contextlib import contextmanager
from enum import Enum
from typing import Any, Final, Literal, cast
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
import litellm
from litellm._internal_context import current_service_target, service_target
from litellm._logging import verbose_logger
from litellm.constants import CACHED_STREAMING_CHUNK_DELAY
from litellm.integrations.otel.runtime import phase_span
from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
from litellm.litellm_core_utils.model_param_helper import ModelParamHelper, peek_reasoning_summary_aliases
from litellm.types.caching import *
from litellm.types.utils import EmbeddingResponse, is_litellm_owned_kwarg
@ -42,6 +42,8 @@ from .redis_semantic_cache import RedisSemanticCache
from .s3_cache import S3Cache
RESPONSE_CACHE_TARGET: Final = "llm_response"
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
_DECODED_LITERAL: Final = TypeAdapter(object)
@contextmanager
@ -374,7 +376,7 @@ class Cache:
)
return "".join(f"{field}: {value}" for field, value in scope_values if value is not None)
def get_cache_key(self, **kwargs) -> str:
def get_cache_key(self, **kwargs: object) -> str:
"""
Get the cache key for the given arguments.
@ -398,21 +400,26 @@ class Cache:
for param in kwargs:
if param in scope_excluded_params:
continue
if param in combined_kwargs:
if param in combined_kwargs or param in _ALWAYS_KEYED_PROVIDER_PARAMS:
param_value: str | None = self._get_param_value(param, kwargs)
if param in _ALWAYS_KEYED_PROVIDER_PARAMS and not param_value:
continue
if param_value is not None:
cache_key += f"{param}: {param_value}"
elif not is_litellm_owned_kwarg(param):
if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now
if kwargs[param] is None:
continue # ignore None params
param_value = kwargs[param]
cache_key += f"{param}: {param_value}"
cache_key += f"{param}: {kwargs[param]}"
if is_semantic_cache:
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
reasoning_summary: Final = peek_reasoning_summary_aliases(kwargs)
reasoning_summary_material: Final = (
"" if reasoning_summary is None else f"reasoning_summary: {reasoning_summary}"
)
hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material)
hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
verbose_logger.debug(
"\nCreated cache key: %s (source material length: %d)",
@ -599,9 +606,11 @@ class Cache:
if isinstance(cached_response, dict):
pass
else:
cached_response = json.loads(cached_response) # Convert string to dictionary
cached_response = _DECODED_LITERAL.validate_python(
json.loads(cached_response)
) # Convert string to dictionary
except Exception:
cached_response = ast.literal_eval(cached_response)
cached_response = _DECODED_LITERAL.validate_python(ast.literal_eval(cached_response))
return cached_response
return cached_result

21
litellm/caching/caching_handler.py
View file

@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
import litellm
from litellm._internal_context import post_response_phase
from litellm._internal_context import is_internal_call, post_response_phase
from litellm._logging import print_verbose, verbose_logger
from litellm.caching import InMemoryCache
from litellm.caching.caching import S3Cache, response_cache_phase
@ -346,6 +346,7 @@ class LLMCachingHandler:
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
verbose_logger.debug("Checking Async Cache")
logging_obj.reset_caching_details()
cached_result = await self._retrieve_from_cache(
call_type=call_type,
kwargs=kwargs,
@ -357,7 +358,7 @@ class LLMCachingHandler:
verbose_logger.debug("Cache Hit!")
cache_hit: Final = True
end_time: Final = datetime.datetime.now()
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
provider_model, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=model,
custom_llm_provider=kwargs.get("custom_llm_provider", None),
api_base=kwargs.get("api_base", None),
@ -366,7 +367,7 @@ class LLMCachingHandler:
cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000
self._update_litellm_logging_obj_environment(
logging_obj=logging_obj,
model=model,
model=provider_model,
kwargs=kwargs,
cached_result=cached_result,
is_async=True,
@ -381,7 +382,7 @@ class LLMCachingHandler:
call_type=call_type,
kwargs=kwargs,
logging_obj=logging_obj,
model=model,
model=provider_model,
custom_llm_provider=kwargs.get("custom_llm_provider", None),
args=args,
)
@ -467,6 +468,7 @@ class LLMCachingHandler:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
print_verbose("Checking Sync Cache")
logging_obj.reset_caching_details()
with response_cache_phase("get"):
cached_result = litellm.cache.get_cache(**new_kwargs)
if is_response_without_output(cached_result):
@ -511,8 +513,10 @@ class LLMCachingHandler:
custom_llm_provider=custom_llm_provider,
)
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
logging_obj.handle_sync_success_callbacks_for_async_calls(
if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks(
cached_result=cached_result
):
logging_obj.handle_sync_cache_hit_success_callbacks(
result=cached_result,
start_time=start_time,
end_time=end_time,
@ -820,7 +824,12 @@ class LLMCachingHandler:
start_time (datetime): The start time of the operation.
end_time (datetime): The end time of the operation.
cache_hit (bool): Whether it was a cache hit.
An internal sub-call's hit fires nothing, since the parent call that folds it in logs it.
"""
if is_internal_call.get():
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(

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