Merge pull request #26430 from BerriAI/litellm_internal_staging
Some checks failed
Unit Tests: Security / security (push) Has been cancelled
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled

merge main
This commit is contained in:
Sameer Kankute 2026-04-24 19:59:18 +05:30 committed by GitHub
commit 7734770bd3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
204 changed files with 20050 additions and 2278 deletions

File diff suppressed because it is too large Load diff

View file

@ -32,41 +32,39 @@ on:
required: false
type: boolean
default: false
dist:
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
required: false
type: string
default: "loadscope"
artifact-name:
description: "Unique name for the coverage artifact (must be unique per run)"
required: false
type: string
default: "run"
secrets:
DATABASE_URL:
required: false
POSTGRES_USER:
required: false
POSTGRES_PASSWORD:
required: false
permissions:
contents: read
# The postgres service container below is spawned per-job on localhost and
# destroyed with the job. Nothing outside the runner can reach it. The
# user/password/database here are not secrets — they're bootstrap values
# for a throwaway container — so we hardcode them instead of attaching
# every matrix shard to a GHA environment just to read three "secrets"
# (which also produces a "temporarily deployed to …" notification on the
# PR timeline per shard per push).
jobs:
run:
name: Run tests
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.timeout-minutes }}
# Environment is derived from the enable-* flags, not caller-controllable.
# This prevents callers from passing arbitrary environment names to bypass secret scoping.
environment: >-
${{
inputs.enable-postgres && 'integration-postgres' ||
''
}}
services:
postgres:
image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14
env:
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_DB: litellm_test
ports:
- 5432:5432
@ -114,7 +112,7 @@ jobs:
- name: Run Prisma migrations
if: ${{ inputs.enable-postgres }}
env:
DATABASE_URL: ${{ secrets.DATABASE_URL }}
DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test"
run: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
@ -124,7 +122,8 @@ jobs:
MAX_FAILURES: ${{ inputs.max-failures }}
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DATABASE_URL: ${{ inputs.enable-postgres && secrets.DATABASE_URL || '' }}
DIST: ${{ inputs.dist }}
DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }}
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
@ -143,7 +142,7 @@ jobs:
-n "${WORKERS}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--dist=loadscope \
--dist="${DIST}" \
--durations=20 \
--cov=litellm \
--cov-report=xml:coverage.xml \

View file

@ -39,7 +39,7 @@ jobs:
if: github.event.action == 'opened'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.11"
python-version: "3.12"
- name: Auto-close if high-confidence duplicate
if: github.event.action == 'opened'

View file

@ -0,0 +1,65 @@
name: Create Release Branch
on:
workflow_dispatch:
inputs:
tag:
description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/<tag>"
required: true
type: string
commit_hash:
description: "Full 40-char commit SHA the branch should point to"
required: true
type: string
workflow_call:
inputs:
tag:
description: "Release tag"
required: true
type: string
commit_hash:
description: "Full 40-char commit SHA the branch should point to"
required: true
type: string
permissions: {}
jobs:
create-branch:
name: Create Release Branch
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Validate inputs
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
run: |
if ! echo "${COMMIT_HASH}" | grep -qE '^[0-9a-f]{40}$'; then
echo "::error::commit_hash must be a full 40-character commit SHA"
exit 1
fi
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
echo "::error::tag must start with vX.Y.Z"
exit 1
fi
- name: Create release branch
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
with:
script: |
const tag = process.env.TAG;
const commitHash = process.env.COMMIT_HASH;
const branchName = `release/${tag}`;
await github.rest.git.createRef({
owner: context.repo.owner,
repo: context.repo.repo,
ref: `refs/heads/${branchName}`,
sha: commitHash,
});
core.info(`Created branch ${branchName} at ${commitHash}`);

View file

@ -102,6 +102,17 @@ jobs:
body: updatedBody,
draft: false,
});
} catch (error) {
core.setFailed(error.message);
}
create-branch:
name: Create Release Branch
needs: release
permissions:
contents: write
uses: ./.github/workflows/create-release-branch.yml
with:
tag: ${{ inputs.tag }}
commit_hash: ${{ inputs.commit_hash }}

View file

@ -29,7 +29,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.11"
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7

View file

@ -29,7 +29,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.13"
python-version: "3.12"
- name: Scan for duplicate issues
env:

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

@ -0,0 +1,136 @@
name: Code Quality Checks
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_branch
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
code-quality:
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Checkout litellm-docs (for documentation_tests)
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
repository: BerriAI/litellm-docs
path: _litellm_docs_checkout
persist-credentials: false
- name: Wire up docs path expected by documentation_tests/*
run: |
# documentation_tests scripts read from docs/my-website/docs/...
# In litellm-docs the same files live at docs/... (repo root).
# Point docs/my-website -> litellm-docs checkout so the paths resolve.
rm -rf docs/my-website
ln -s ../_litellm_docs_checkout docs/my-website
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Install dependencies
run: uv sync --frozen --all-groups --all-extras
- name: check_licenses
run: uv run --no-sync python ./tests/code_coverage_tests/check_licenses.py
- name: check_provider_folders_documented
run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
- name: test_chat_completion_imports
run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py
- name: info_log_check
run: uv run --no-sync python ./tests/code_coverage_tests/info_log_check.py
- name: check_guardrail_apply_decorator
run: uv run --no-sync python ./tests/code_coverage_tests/check_guardrail_apply_decorator.py
- name: test_ban_set_verbose
run: uv run --no-sync python ./tests/code_coverage_tests/test_ban_set_verbose.py
- name: code_qa_check_tests
run: uv run --no-sync python ./tests/code_coverage_tests/code_qa_check_tests.py
- name: check_get_model_cost_key_performance
run: uv run --no-sync python ./tests/code_coverage_tests/check_get_model_cost_key_performance.py
- name: test_proxy_types_import
run: uv run --no-sync python ./tests/code_coverage_tests/test_proxy_types_import.py
- name: callback_manager_test
run: uv run --no-sync python ./tests/code_coverage_tests/callback_manager_test.py
- name: recursive_detector
run: uv run --no-sync python ./tests/code_coverage_tests/recursive_detector.py
- name: test_router_strategy_async
run: uv run --no-sync python ./tests/code_coverage_tests/test_router_strategy_async.py
- name: litellm_logging_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/litellm_logging_code_coverage.py
- name: ensure_async_clients_test
run: uv run --no-sync python ./tests/code_coverage_tests/ensure_async_clients_test.py
- name: enforce_llms_folder_style
run: uv run --no-sync python ./tests/code_coverage_tests/enforce_llms_folder_style.py
- name: prevent_key_leaks_in_exceptions
run: uv run --no-sync python ./tests/code_coverage_tests/prevent_key_leaks_in_exceptions.py
- name: check_unsafe_enterprise_import
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
- name: ban_copy_deepcopy_kwargs
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
- name: check_fastuuid_usage
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
- name: memory_test
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py
- name: documentation_test_env_keys
run: uv run --no-sync python ./tests/documentation_tests/test_env_keys.py
- name: documentation_test_router_settings
run: uv run --no-sync python ./tests/documentation_tests/test_router_settings.py
- name: documentation_test_api_docs
run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py

39
.github/workflows/test-semgrep.yml vendored Normal file
View file

@ -0,0 +1,39 @@
name: Semgrep
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_branch
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
semgrep:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Run Semgrep (custom rules)
run: uv tool run --from 'semgrep==1.157.0' semgrep scan --config .semgrep/rules . --error

View file

@ -12,8 +12,74 @@ concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
# Semantic matrix: each shard groups tests by concern (auth, server, logging, …)
# rather than alphabetical letter ranges. Adding a new test file means adding it
# to whichever group it belongs to, not reshuffling slices.
#
# Design targets:
# * Every shard runs in <= 7 minutes of wall-clock on the default runner.
# Most of a shard's time is pytest plugin load + xdist worker imports +
# pytest-cov instrumentation, not the tests themselves. Keeping per-shard
# work low and matching worker count to runner cores is what controls it.
# * workers: 4 matches the 4-core ubuntu-latest runner. -n 8 on 4 cores
# oversubscribes 2x and workers fight for CPU during their cold-start
# imports (measured ~441% CPU for -n 8 locally, i.e. ~55% effective).
# * test_key_generate_prisma.py stays serial (workers=0) — it has event-loop
# conflicts with the logging worker when run in parallel.
# * test_proxy_utils.py runs as a single shard with --dist=worksteal so
# xdist balances its 188 parametrized cases across workers instead of
# pinning the whole file to one worker (the default --dist=loadscope
# behavior for single-file targets).
# * test_db_schema_migration.py is isolated because one test in it
# (test_aaaasschema_migration_check) takes ~170s — by itself it
# determines the shard's wall-clock floor.
jobs:
# Fast guard — fails the workflow if a test_*.py file under
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
# The semantic-shard design (no catch-all "remaining" bucket) relies on
# every test file being explicitly assigned; this guard prevents a new
# file from silently dropping out of CI.
assert-shard-coverage:
runs-on: ubuntu-latest
timeout-minutes: 2
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Assert every test_*.py is in a matrix shard
run: |
python3 - <<'PY'
import pathlib, sys, yaml
wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
referenced = set()
for entry in matrix:
for token in entry["test-path"].split():
if token.startswith("tests/proxy_unit_tests/"):
referenced.add(pathlib.PurePosixPath(token).name)
actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
and p.name != "test_configs"}
orphans = sorted(actual - referenced)
if orphans:
print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
print(" are not assigned to any shard in test-unit-proxy-db.yml:")
for o in orphans:
print(f" - {o}")
print()
print("Add each to whichever semantic shard it belongs to.")
sys.exit(1)
print(f"OK: all {len(actual)} files assigned to a shard.")
PY
proxy-db:
needs: assert-shard-coverage
# Display only the semantic shard name in the checks UI instead of GHA's
# default "proxy-db (key-generation, tests/proxy_unit_tests/…, 0, loadscope, 20)"
# which includes every matrix field and gets truncated past the test-path.
name: ${{ matrix.test-group }}
permissions:
contents: read
id-token: write
@ -22,19 +88,146 @@ jobs:
fail-fast: false
matrix:
include:
# Key generation tests must NOT run in parallel (event loop conflicts with logging worker)
# Must run serially — event-loop conflict with the logging worker.
- test-group: key-generation
test-path: "tests/proxy_unit_tests/test_key_generate_prisma.py"
workers: 0
timeout: 30
- test-group: auth-checks
test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py"
workers: 8
dist: loadscope
timeout: 20
- test-group: remaining
test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py"
workers: 8
timeout: 30
# ---- auth: split into 2 shards ----
- test-group: auth-checks
test-path: >-
tests/proxy_unit_tests/test_auth_checks.py
tests/proxy_unit_tests/test_user_api_key_auth.py
workers: 4
dist: loadscope
timeout: 15
- test-group: jwt-and-keys
test-path: >-
tests/proxy_unit_tests/test_jwt.py
tests/proxy_unit_tests/test_jwt_key_mapping.py
tests/proxy_unit_tests/test_proxy_custom_auth.py
tests/proxy_unit_tests/test_key_generate_dynamodb.py
tests/proxy_unit_tests/test_deployed_proxy_keygen.py
workers: 4
dist: loadscope
timeout: 15
# ---- test_proxy_utils.py, single shard, worksteal distribution ----
- test-group: proxy-utils
test-path: "tests/proxy_unit_tests/test_proxy_utils.py"
workers: 4
dist: worksteal
timeout: 15
# ---- proxy server: split into 2 shards ----
- test-group: proxy-server-core
test-path: >-
tests/proxy_unit_tests/test_proxy_server.py
tests/proxy_unit_tests/test_proxy_server_keys.py
tests/proxy_unit_tests/test_proxy_server_caching.py
tests/proxy_unit_tests/test_proxy_server_langfuse.py
tests/proxy_unit_tests/test_proxy_server_spend.py
tests/proxy_unit_tests/test_aproxy_startup.py
workers: 4
dist: loadscope
timeout: 15
- test-group: proxy-runtime
test-path: >-
tests/proxy_unit_tests/test_proxy_config_unit_test.py
tests/proxy_unit_tests/test_proxy_routes.py
tests/proxy_unit_tests/test_proxy_gunicorn.py
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
workers: 4
dist: loadscope
timeout: 15
# ---- logging: split into 2 shards ----
- test-group: custom-logging
test-path: >-
tests/proxy_unit_tests/test_custom_callback_input.py
tests/proxy_unit_tests/test_custom_logger_s3_gcs.py
tests/proxy_unit_tests/test_proxy_custom_logger.py
workers: 4
dist: loadscope
timeout: 15
- test-group: logging-misc
test-path: >-
tests/proxy_unit_tests/test_proxy_reject_logging.py
tests/proxy_unit_tests/test_audit_logs_proxy.py
tests/proxy_unit_tests/test_search_api_logging.py
workers: 4
dist: loadscope
timeout: 15
# ---- db-and-spend: isolate the 170s schema-migration test ----
# test_db_schema_migration.py has exactly one test, and that test
# is mostly waiting on `prisma migrate deploy` / `prisma migrate
# diff` subprocesses (~170s). It does no CPU-bound Python work
# inside the test. Running with workers=0 (serial, no xdist)
# skips the 4-worker cold-start cost we'd otherwise pay for a
# single test, saving ~4 minutes of wall-clock.
- test-group: schema-migration
test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
workers: 0
dist: loadscope
timeout: 15
- test-group: db-and-spend
test-path: >-
tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
tests/proxy_unit_tests/test_db_schema_changes.py
tests/proxy_unit_tests/test_e2e_pod_lock_manager.py
tests/proxy_unit_tests/test_skills_db.py
tests/proxy_unit_tests/test_update_daily_tag_spend.py
tests/proxy_unit_tests/test_update_spend.py
tests/proxy_unit_tests/test_project_endpoints_prisma.py
tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py
workers: 4
dist: loadscope
timeout: 15
# ---- guardrails + budget + hooks: split into 2 ----
- test-group: guardrails-hooks
test-path: >-
tests/proxy_unit_tests/test_proxy_setting_guardrails.py
tests/proxy_unit_tests/test_banned_keyword_list.py
tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
workers: 4
dist: loadscope
timeout: 15
- test-group: budgets
test-path: >-
tests/proxy_unit_tests/test_default_end_user_budget_simple.py
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py
workers: 4
dist: loadscope
timeout: 15
- test-group: endpoints-and-responses
test-path: >-
tests/proxy_unit_tests/test_blog_posts_endpoint.py
tests/proxy_unit_tests/test_models_fallback_endpoint.py
tests/proxy_unit_tests/test_google_endpoint_routing.py
tests/proxy_unit_tests/test_google_gemini_proxy_request.py
tests/proxy_unit_tests/test_get_favicon.py
tests/proxy_unit_tests/test_get_image.py
tests/proxy_unit_tests/test_ui_path_detection.py
tests/proxy_unit_tests/test_prompt_test_endpoint.py
tests/proxy_unit_tests/test_check_batch_cost.py
tests/proxy_unit_tests/test_check_responses_cost.py
tests/proxy_unit_tests/test_response_polling_handler.py
tests/proxy_unit_tests/test_response_polling_pre_call_checks.py
tests/proxy_unit_tests/test_realtime_cache.py
tests/proxy_unit_tests/test_proxy_exception_mapping.py
tests/proxy_unit_tests/test_custom_tokenizer_bug.py
tests/proxy_unit_tests/test_model_response_typing
workers: 4
dist: loadscope
timeout: 15
uses: ./.github/workflows/_test-unit-services-base.yml
with:
test-path: ${{ matrix.test-path }}
@ -42,8 +235,5 @@ jobs:
reruns: 2
timeout-minutes: ${{ matrix.timeout }}
enable-postgres: true
dist: ${{ matrix.dist }}
artifact-name: proxy-db-${{ matrix.test-group }}
secrets:
DATABASE_URL: ${{ secrets.DATABASE_URL }}
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}

View file

@ -36,6 +36,8 @@ jobs:
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
workers: 2
reruns: 2

View file

@ -1,6 +1,8 @@
name: "Unit Tests: Security"
# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
# Kept push-only (was previously required by DATABASE_URL secret scoping;
# now the postgres credentials are ephemeral localhost values but the
# push-trigger stays to match the proxy-db workflow cadence).
on:
push:
branches: [main, "litellm_**"]
@ -24,7 +26,3 @@ jobs:
timeout-minutes: 20
enable-postgres: true
artifact-name: security
secrets:
DATABASE_URL: ${{ secrets.DATABASE_URL }}
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}

View file

@ -27,10 +27,8 @@ RUN apk add --no-cache \
npm \
libsndfile
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
UV_PROJECT_ENVIRONMENT=/app/.venv \
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
# Copy dependency metadata first for layer caching
@ -94,11 +92,14 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervi
{ apk del --no-cache npm 2>/dev/null || true; }
WORKDIR /app
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
ENV PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
COPY --from=builder /root/.cache /root/.cache
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete

View file

@ -1,22 +1,231 @@
import argparse
import os
import subprocess
from pathlib import Path
from datetime import datetime
import testing.postgresql
import re
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path
import testing.postgresql
def create_migration(migration_name: str = None):
DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE)
DEFAULT_BASE_BRANCH = "litellm_internal_staging"
def _find_destructive_statements(sql: str) -> list:
"""Return SQL lines containing DROP COLUMN, DROP TABLE, or DROP INDEX."""
return [
line.strip() for line in sql.splitlines() if DESTRUCTIVE_PATTERN.search(line)
]
def _print_freshness_failure(
base_branch: str, reason: str, stderr_text: str = ""
) -> None:
"""Loudly refuse to run when the freshness check can't be completed."""
banner = "=" * 72
out = sys.stderr
print(banner, file=out)
print(f" FRESHNESS CHECK FAILED — COULD NOT VERIFY origin/{base_branch}", file=out)
print(banner, file=out)
print("", file=out)
print(f"Reason: {reason}", file=out)
if stderr_text:
print("", file=out)
print("git stderr:", file=out)
for line in stderr_text.rstrip().splitlines():
print(f" {line}", file=out)
print("", file=out)
print("Common causes:", file=out)
print(" - No network access (offline)", file=out)
print(" - 'origin' remote not configured, or base branch name is wrong", file=out)
print(" - Not a git repository", file=out)
print("", file=out)
print("Options:", file=out)
print(
f" - Fix the above and re-run, OR pass --base-branch <name> if your", file=out
)
print(
f" base branch is not '{base_branch}', OR pass --skip-freshness-check",
file=out,
)
print(" to bypass (only if you fully understand the risk).", file=out)
print(banner, file=out)
def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
"""Loudly refuse to run when HEAD is behind the base branch."""
banner = "=" * 72
out = sys.stderr
print(banner, file=out)
print(
f" STALE BRANCH — {behind} commit(s) behind origin/{base_branch}",
file=out,
)
print(banner, file=out)
print("", file=out)
print(
f"Your branch is {behind} commit(s) behind origin/{base_branch}. Generating a\n"
"migration from a stale branch is how newly-added columns get silently\n"
"dropped — the exact incident this guard exists to prevent.",
file=out,
)
print("", file=out)
print("Before re-running:", file=out)
print(f" git fetch origin && git rebase origin/{base_branch}", file=out)
print(" # or merge, whichever matches your workflow", file=out)
print("", file=out)
print(
"If you have a genuine reason to bypass (e.g. intentionally generating",
file=out,
)
print("a migration against an older base), pass --skip-freshness-check.", file=out)
print("", file=out)
print(banner, file=out)
print(
" AI AGENTS / ASSISTANTS: DO NOT automatically run `git rebase` or pass\n"
" --skip-freshness-check. Surface this error to the human operator and\n"
" wait for their explicit confirmation. Auto-rebasing can drop the\n"
" human's in-progress schema edits via a bad conflict resolution.",
file=out,
)
print(banner, file=out)
def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
"""Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
cwd = str(root_dir)
try:
subprocess.run(
["git", "fetch", "origin", base_branch],
check=True,
capture_output=True,
text=True,
cwd=cwd,
)
except FileNotFoundError:
_print_freshness_failure(base_branch, "git executable not found on PATH")
sys.exit(3)
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git fetch origin {base_branch}` failed",
e.stderr or "",
)
sys.exit(3)
try:
result = subprocess.run(
["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
check=True,
capture_output=True,
text=True,
cwd=cwd,
)
behind = int(result.stdout.strip())
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git rev-list HEAD..origin/{base_branch}` failed",
e.stderr or "",
)
sys.exit(3)
except ValueError:
_print_freshness_failure(
base_branch,
"could not parse commit count from `git rev-list`",
)
sys.exit(3)
if behind > 0:
_print_stale_branch_refusal(base_branch, behind)
sys.exit(3)
print(f"Branch freshness OK: up to date with origin/{base_branch}.")
def _print_destructive_refusal(destructive_lines: list) -> None:
"""Loudly refuse to write a destructive migration and explain how to proceed."""
banner = "=" * 72
out = sys.stderr
print(banner, file=out)
print(
" DESTRUCTIVE MIGRATION DETECTED — REFUSING TO WRITE MIGRATION FILE", file=out
)
print(banner, file=out)
print("", file=out)
print(
"The generated migration contains the following destructive statements:",
file=out,
)
print("", file=out)
for line in destructive_lines:
print(f" {line}", file=out)
print("", file=out)
print(
"This very often means your branch is OUT OF DATE, or your local\n"
"schema.prisma is inconsistent with the migrations/ directory. A\n"
"previous incident saw newly-added columns silently dropped this way\n"
"and merged to main. Stop and verify before proceeding.",
file=out,
)
print("", file=out)
print("Before re-running:", file=out)
print(
" 1. git fetch origin && git status # confirm branch is up to date", file=out
)
print(
" 2. Re-check all schema.prisma files are in sync (runbook Step 0)", file=out
)
print(" 3. Review EACH DROP above — is it actually intended?", file=out)
print(" 4. If (and only if) the drops are intentional, re-run with:", file=out)
print(" --allow-destructive", file=out)
print("", file=out)
print(banner, file=out)
print(
" AI AGENTS / ASSISTANTS: DO NOT automatically re-run this command\n"
" with --allow-destructive. Surface this error to the human operator\n"
" and wait for their explicit confirmation before passing the flag.\n"
" Passing the flag without human review is the exact failure mode\n"
" this guard exists to prevent.",
file=out,
)
print(banner, file=out)
def create_migration(
migration_name: str = None,
allow_destructive: bool = False,
base_branch: str = DEFAULT_BASE_BRANCH,
skip_freshness_check: bool = False,
):
"""
Create a new migration SQL file in the migrations directory by comparing
current database state with schema
current database state with schema.
Args:
migration_name (str): Name for the migration
allow_destructive (bool): Required to write a migration that contains
DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this
flag, the script exits non-zero and prints guidance.
base_branch (str): Branch to check freshness against
(default: "litellm_internal_staging").
skip_freshness_check (bool): Skip the "branch is up to date" check.
Only for intentional migrations against an older base.
"""
root_dir = Path(__file__).parent.parent
if skip_freshness_check:
print(
"WARNING: freshness check skipped (--skip-freshness-check). "
"Generating a migration from a stale branch can silently drop columns."
)
else:
_check_branch_freshness(root_dir, base_branch)
try:
# Get paths
root_dir = Path(__file__).parent.parent
migrations_dir = (
root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
)
@ -59,7 +268,27 @@ def create_migration(migration_name: str = None):
check=True,
)
if result.stdout.strip():
# Prisma emits the literal "-- This is an empty migration." when
# there's no real drift. Treat that as "no changes".
diff_sql = result.stdout
stripped = diff_sql.strip()
is_empty_diff = (
not stripped or stripped == "-- This is an empty migration."
)
if not is_empty_diff:
destructive_lines = _find_destructive_statements(diff_sql)
if destructive_lines and not allow_destructive:
_print_destructive_refusal(destructive_lines)
sys.exit(2)
if destructive_lines and allow_destructive:
print(
"WARNING: writing destructive migration "
"(--allow-destructive passed). Statements:"
)
for line in destructive_lines:
print(f" {line}")
# Generate timestamp and create migration directory
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
migration_name = migration_name or "unnamed_migration"
@ -68,7 +297,7 @@ def create_migration(migration_name: str = None):
# Write the SQL to migration.sql
migration_file = migration_dir / "migration.sql"
migration_file.write_text(result.stdout)
migration_file.write_text(diff_sql)
print(f"Created migration in {migration_dir}")
return True
@ -90,8 +319,48 @@ def create_migration(migration_name: str = None):
if __name__ == "__main__":
# If running directly, can optionally pass migration name as argument
import sys
migration_name = sys.argv[1] if len(sys.argv) > 1 else None
create_migration(migration_name)
parser = argparse.ArgumentParser(
description=(
"Generate a Prisma migration by diffing the temp DB "
"(existing migrations applied) against schema.prisma."
)
)
parser.add_argument(
"migration_name",
nargs="?",
default=None,
help="Name for the migration (used in the generated directory name).",
)
parser.add_argument(
"--allow-destructive",
action="store_true",
help=(
"Required to write a migration that contains DROP COLUMN, "
"DROP TABLE, or DROP INDEX. Without this flag, destructive "
"diffs are refused."
),
)
parser.add_argument(
"--base-branch",
default=DEFAULT_BASE_BRANCH,
help=(
f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
"The script fetches origin/<base-branch> and refuses to run if HEAD "
"is behind it."
),
)
parser.add_argument(
"--skip-freshness-check",
action="store_true",
help=(
"Bypass the 'branch is up to date' check. Only for intentional "
"migrations against an older base. Pairs poorly with automation."
),
)
args = parser.parse_args()
create_migration(
args.migration_name,
allow_destructive=args.allow_destructive,
base_branch=args.base_branch,
skip_freshness_check=args.skip_freshness_check,
)

View file

@ -26,10 +26,8 @@ RUN apk add --no-cache \
npm \
libsndfile
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
UV_PROJECT_ENVIRONMENT=/app/.venv \
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
# Copy dependency metadata first for layer caching
@ -92,11 +90,14 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervi
{ apk del --no-cache npm 2>/dev/null || true; }
WORKDIR /app
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
XDG_CACHE_HOME=/app/.cache \
PATH="/app/.venv/bin:${PATH}"
ENV PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
COPY --from=builder /root/.cache /root/.cache
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete

View file

@ -138,7 +138,7 @@ RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w "$LITELLM_PROXY_EXTRAS_PATH" || true && \
chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets /app/.cache
USER nobody
USER 65534
RUN prisma generate --schema=./schema.prisma

View file

@ -2,7 +2,7 @@ schemaVersion: 2.0.0
metadataTest:
entrypoint: ["docker/prod_entrypoint.sh"]
user: "nobody"
user: "65534"
workdir: "/app"
fileExistenceTests:

View file

@ -0,0 +1,155 @@
# [BETA] Adaptive Router
:::info
Beta feature. Share feedback on [Discord](https://discord.gg/wuPM9dRgDw) or [Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA).
:::
**Requirements:** LiteLLM Proxy with a Postgres database. Quality estimates are stored in Postgres and loaded on startup — without a database the router works but forgets everything learned on restart.
You have a cheap model and an expensive one. You want to use the cheap one when it's good enough, and the expensive one when it actually matters — without hardcoding rules you'll spend months tuning.
The adaptive router does this automatically. It tracks which model performs best for each type of request (code, writing, analysis, etc.) and routes accordingly, balancing quality against cost based on weights you control.
## Quick start
```yaml
model_list:
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
model_info:
input_cost_per_token: 0.0000025
adaptive_router_preferences:
quality_tier: 3 # 1=budget, 2=mid, 3=frontier
strengths: ["code_generation", "analytical_reasoning"]
- model_name: gpt-4o-mini
litellm_params:
model: openai/gpt-4o-mini
model_info:
input_cost_per_token: 0.00000015
adaptive_router_preferences:
quality_tier: 2
strengths: ["factual_lookup"]
- model_name: my-router
litellm_params:
model: auto_router/adaptive_router
adaptive_router_config:
available_models: ["gpt-4o", "gpt-4o-mini"]
weights:
quality: 0.7 # raise this if quality complaints; lower if bill too high
cost: 0.3 # must sum to 1.0 with quality
```
Route to it by setting `model` to your adaptive router's name:
```bash
curl -X POST {{baseURL}}/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $LITELLM_API_KEY" \
-d '{
"model": "my-router",
"messages": [
{"role": "user", "content": "build me a python script that parses CSV"},
{"role": "assistant", "content": "Here is a script using csv.DictReader..."},
{"role": "user", "content": "now add error handling for missing files"},
{"role": "assistant", "content": "Wrap the open() call in a try/except FileNotFoundError..."},
{"role": "user", "content": "perfect, that worked. thanks!"}
]
}'
```
The response includes a header telling you which model was actually picked:
```
x-litellm-adaptive-router-model: gpt-4o
```
The "thanks!" turn in the example above fires a satisfaction signal — that's what moves the bandit.
## Tuning cost vs. quality
The `weights` are your main lever:
| Goal | quality | cost |
|---|---|---|
| Minimize cost, quality is secondary | 0.3 | 0.7 |
| Balanced | 0.5 | 0.5 |
| Quality-first (default) | 0.7 | 0.3 |
| Quality non-negotiable | 0.9 | 0.1 |
The router learns over time. For the first ~10 requests per model, it relies on the tiers you declared. After that, real performance data takes over.
## Force a minimum quality tier per request
If a specific request needs a frontier model regardless of cost, pass this header:
```
x-litellm-min-quality-tier: 3
```
You can also pass `min_quality_tier` via request metadata instead of a header.
## What's being learned
The router classifies each request into one of 7 types and tracks how each model performs on each independently. A model that's great at factual lookup but poor at code will win factual requests and lose code requests — even if it's cheaper overall.
| Type | Example |
|---|---|
| `code_generation` | "write me a Python sort function" |
| `code_understanding` | "explain what this function does" |
| `technical_design` | "how should I design this API?" |
| `analytical_reasoning` | "calculate the probability that..." |
| `writing` | "draft an email to my team about..." |
| `factual_lookup` | "what is the capital of France?" |
| `general` | anything else |
[**See classifier code**](https://github.com/BerriAI/litellm/blob/litellm_adaptive_routing/litellm/router_strategy/adaptive_router/classifier.py)
Learning signals are inspired by [Signals: Trajectory Sampling and Triage for Agentic Interactions](https://arxiv.org/pdf/2604.00356).
## Inspect the current state
```
GET /adaptive_router/{router_name}/state
```
Returns current quality estimates per model per request type. Useful for understanding why a model is or isn't being picked.
```json
{
"routers": [
{
"router_name": "smart-cheap-router",
"available_models": ["fast", "smart"],
"weights": { "quality": 0.7, "cost": 0.3 },
"cells": [
{
"request_type": "analytical_reasoning",
"model": "fast",
"quality_mean": 0.5,
"samples": 0
},
{
"request_type": "analytical_reasoning",
"model": "smart",
"quality_mean": 0.95,
"samples": 0
}
]
}
]
}
```
`quality_mean` is the key number — it's the router's current estimate of how well that model handles that request type. `samples` counts how many real observations have moved the prior (starts at 0; the cold-start prior mass is excluded).
## Known limitations
- Latency isn't scored — a slow model can still win on quality + cost
- Signals are regex-based and English-biased — no LLM judge
- Hard cap of 200 observations per cell; no decay yet
- Once a model is picked for a session, other models' turns in that session don't contribute to learning

View file

@ -2061,7 +2061,7 @@ assert isinstance(
## Media Resolution Control (Images & Videos)
For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types.
LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter for all Gemini models. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types.
**Supported `detail` values:**
- `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos)
@ -2146,12 +2146,12 @@ response = completion(
</Tabs>
:::info
**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models.
**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types across all Gemini models.
:::
## Video Metadata Control
For Gemini 3+ models, LiteLLM supports fine-grained video processing control through the `video_metadata` field. This allows you to specify frame extraction rates and time ranges for video analysis.
LiteLLM supports fine-grained video processing control through the `video_metadata` field for all Gemini models (1.x, 2.x, 3+). This allows you to specify frame extraction rates and time ranges for video analysis.
**Supported `video_metadata` parameters:**
@ -2168,8 +2168,11 @@ For Gemini 3+ models, LiteLLM supports fine-grained video processing control thr
- `fps` remains unchanged
:::
:::tip
Video clipping (`start_offset`/`end_offset`) and frame rate control (`fps`) are supported by all Gemini models, but analysis quality is significantly higher with the **Gemini 2.5 series** (e.g., `gemini-2.5-flash`, `gemini-2.5-pro`).
:::
:::warning
- **Gemini 3+ Only:** This feature is only available for Gemini 3.0 and newer models
- **Video Files Recommended:** While `video_metadata` is designed for video files, error handling for other media types is delegated to the Vertex AI API
- **File Formats Supported:** Works with `gs://`, `https://`, and base64-encoded video files
:::

View file

@ -26,6 +26,7 @@
},
"devDependencies": {
"@docusaurus/module-type-aliases": "3.8.1",
"ajv": "^8.18.0",
"dotenv": "16.6.1"
},
"engines": {

View file

@ -32,6 +32,7 @@
},
"devDependencies": {
"@docusaurus/module-type-aliases": "3.8.1",
"ajv": "^8.18.0",
"dotenv": "16.6.1"
},
"browserslist": {

View file

@ -1060,6 +1060,7 @@ const sidebars = {
},
items: [
"routing",
"adaptive_router",
"scheduler",
"proxy/auto_routing",
"proxy/load_balancing",

View file

@ -0,0 +1,39 @@
-- One row per (router, request_type, model). Hot path on every routing decision.
CREATE TABLE "LiteLLM_AdaptiveRouterState" (
router_name TEXT NOT NULL,
request_type TEXT NOT NULL,
model_name TEXT NOT NULL,
alpha DOUBLE PRECISION NOT NULL,
beta DOUBLE PRECISION NOT NULL,
total_samples INTEGER NOT NULL DEFAULT 0,
last_updated_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (router_name, request_type, model_name)
);
-- One row per (session, router, model). Updated per turn via the queue.
CREATE TABLE "LiteLLM_AdaptiveRouterSession" (
session_id TEXT NOT NULL,
router_name TEXT NOT NULL,
model_name TEXT NOT NULL,
classified_type TEXT NOT NULL,
misalignment_count INTEGER NOT NULL DEFAULT 0,
stagnation_count INTEGER NOT NULL DEFAULT 0,
disengagement_count INTEGER NOT NULL DEFAULT 0,
satisfaction_count INTEGER NOT NULL DEFAULT 0,
failure_count INTEGER NOT NULL DEFAULT 0,
loop_count INTEGER NOT NULL DEFAULT 0,
exhaustion_count INTEGER NOT NULL DEFAULT 0,
last_user_content TEXT,
last_assistant_content TEXT,
tool_call_history JSONB NOT NULL DEFAULT '[]',
pending_tool_calls JSONB NOT NULL DEFAULT '{}',
turn_count INTEGER NOT NULL DEFAULT 0,
last_processed_turn INTEGER NOT NULL DEFAULT -1,
clean_credit_awarded BOOLEAN NOT NULL DEFAULT FALSE,
terminal_status INTEGER,
last_activity_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (session_id, router_name, model_name)
);
CREATE INDEX "idx_adaptive_router_session_activity"
ON "LiteLLM_AdaptiveRouterSession" (last_activity_at);

View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
user_id String
team_id String
spend Float @default(0.0)
total_spend Float @default(0.0)
budget_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
@@id([user_id, team_id])
@ -1223,3 +1224,46 @@ model LiteLLM_ClaudeCodePluginTable {
@@map("LiteLLM_ClaudeCodePluginTable")
}
// Per-(router, request_type, model) Beta posterior for the adaptive router.
model LiteLLM_AdaptiveRouterState {
router_name String
request_type String
model_name String
alpha Float
beta Float
total_samples Int @default(0)
last_updated_at DateTime @default(now()) @updatedAt
@@id([router_name, request_type, model_name])
}
// Per-(session, router, model) signal counters for the adaptive router.
model LiteLLM_AdaptiveRouterSession {
session_id String
router_name String
model_name String
classified_type String
misalignment_count Int @default(0)
stagnation_count Int @default(0)
disengagement_count Int @default(0)
satisfaction_count Int @default(0)
failure_count Int @default(0)
loop_count Int @default(0)
exhaustion_count Int @default(0)
last_user_content String?
last_assistant_content String?
tool_call_history Json @default("[]")
pending_tool_calls Json @default("{}")
turn_count Int @default(0)
last_processed_turn Int @default(-1)
clean_credit_awarded Boolean @default(false)
terminal_status Int?
last_activity_at DateTime @default(now()) @updatedAt
@@id([session_id, router_name, model_name])
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}

View file

@ -30,6 +30,26 @@ def _get_prisma_env() -> dict:
return prisma_env
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
def _migration_timestamp(name: str) -> int:
"""Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
Returns 0 if the name doesn't match the Prisma pattern — unexpected-format
entries sort as "oldest" and are treated as historical.
"""
m = _MIGRATION_TS_RE.match(name)
return int(m.group(1)) if m else 0
def _max_migration_timestamp(names) -> int:
"""Max timestamp in a set/list of migration names (0 if empty)."""
if not names:
return 0
return max(_migration_timestamp(n) for n in names)
def _get_prisma_command() -> str:
"""Get the Prisma command to use, bypassing Python wrapper in offline mode."""
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
@ -383,18 +403,301 @@ class ProxyExtrasDBManager:
)
@staticmethod
def setup_database(use_migrate: bool = False) -> bool:
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
schema, etc.) from DATABASE_URL so psycopg can parse it."""
from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
parsed = urlparse(url)
if not parsed.query:
return url
libpq_params = {
"sslmode",
"sslcert",
"sslkey",
"sslrootcert",
"sslpassword",
"application_name",
"connect_timeout",
"client_encoding",
"options",
"service",
"gssencmode",
"krbsrvname",
"target_session_attrs",
}
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
return urlunparse(parsed._replace(query=urlencode(kept)))
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
"""
Log a warning if _prisma_migrations contains applied migrations with
timestamps newer than every migration this build ships.
This is informational only for the v2 resolver it tells the operator
the DB was likely migrated by a newer deployment, which is usually a
signal that this (older) version shouldn't run against it. We do NOT
block startup: many users have weird _prisma_migrations state from
prior thrashing bugs, and blocking them would be a breaking change.
Safe no-op if psycopg isn't installed or DB isn't reachable.
"""
database_url = os.getenv("DATABASE_URL")
if not database_url:
return
try:
import psycopg
except ImportError:
return
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
known = set(ProxyExtrasDBManager._get_migration_names(migrations_dir))
try:
# autocommit=True keeps the SELECT outside a transaction. Without
# it, psycopg3's `with conn` calls COMMIT on clean exit — which
# fails after `UndefinedTable` (fresh DB) leaves the transaction
# in an aborted state.
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
try:
rows = conn.execute(
"SELECT migration_name FROM _prisma_migrations "
"WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL"
).fetchall()
except psycopg.errors.UndefinedTable:
return
except (psycopg.OperationalError, psycopg.DatabaseError):
# Swallow connection failures AND any other DB-layer error
# (e.g. InsufficientPrivilege if the runtime user lacks SELECT
# on _prisma_migrations). This is an informational check —
# never block startup on it.
return
applied = {r[0] for r in rows}
unknown = applied - known
if not unknown:
return
head_newest_ts = _max_migration_timestamp(known)
hostile = {
name for name in unknown if _migration_timestamp(name) > head_newest_ts
}
if not hostile:
return
sorted_hostile = sorted(hostile)
logger.warning(
"Database has %d migration(s) applied that are NEWER than any "
"migration this LiteLLM version ships. This usually means the "
"database was migrated by a newer LiteLLM deployment. Some API "
"endpoints may fail because this proxy's Prisma client does not "
"know about those schema changes. Consider upgrading this "
"deployment. Unknown: %s",
len(hostile),
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
)
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
"""
v2 migration resolver (opt-in via --use_v2_migration_resolver).
Runs `prisma migrate deploy` and handles standard recovery paths
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
NOT call `_resolve_all_migrations` the diff-and-force recovery that
caused schema thrashing when two LiteLLM versions contended for the
same DB during rolling deploys.
Ahead-of-HEAD state (DB has migrations newer than this build ships)
is logged as a warning, not a fatal error users whose DBs got into
weird shapes from the old thrashing should still be able to start.
"""
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
if not use_migrate:
# Preserve `prisma db push` path unchanged.
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=60,
check=True,
env=_get_prisma_env(),
)
return True
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as e:
# Re-raise as RuntimeError so proxy_cli.py's
# `except RuntimeError` catches it and exits cleanly.
raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e
finally:
os.chdir(original_dir)
# Informational — never blocks.
ProxyExtrasDBManager._warn_if_db_ahead_of_head(migrations_dir)
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
for attempt in range(4):
try:
result = subprocess.run(
[_get_prisma_command(), "migrate", "deploy"],
timeout=60,
check=True,
capture_output=True,
text=True,
env=_get_prisma_env(),
)
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
return True
except subprocess.TimeoutExpired:
logger.info(
f"prisma migrate deploy attempt {attempt + 1} timed out, retrying"
)
time.sleep(random.randrange(5, 15))
continue
except subprocess.CalledProcessError as e:
stderr = e.stderr or ""
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info(
"Schema exists but no migrations ledger — creating baseline"
)
ProxyExtrasDBManager._create_baseline_migration(schema_path)
continue
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're already inside the outer
# `except CalledProcessError` handler —
# re-raising CalledProcessError from here
# would escape as itself, bypassing
# proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from e
migration_match = re.search(
r"Migration name: (\d+_\S+)", stderr
)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
raise RuntimeError(
"Database migration failed after 4 attempts (retry loop "
"exhausted by timeouts or repeated idempotent-recovery "
"continues). Check database connectivity, load, and "
"_prisma_migrations ledger state."
)
finally:
os.chdir(original_dir)
@staticmethod
def setup_database(
use_migrate: bool = False, use_v2_resolver: bool = False
) -> bool:
"""
Set up the database using either prisma migrate or prisma db push
Uses migrations from litellm-proxy-extras package
Args:
schema_path (str): Path to the Prisma schema file
use_migrate (bool): Whether to use prisma migrate instead of db push
use_migrate: Whether to use prisma migrate instead of db push
use_v2_resolver: Opt into the v2 migration resolver (safer during
rolling deploys; does not run the diff-and-force recovery
that causes schema thrashing). Defaults to False for
backwards compatibility.
Returns:
bool: True if setup was successful, False otherwise
"""
if use_v2_resolver:
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
for attempt in range(4):
original_dir = os.getcwd()

View file

@ -2,6 +2,8 @@
This is a runbook for creating and running database migrations for the LiteLLM proxy. For use for litellm engineers only.
> **AI AGENTS / ASSISTANTS:** If the script refuses with either a "STALE BRANCH" or "DESTRUCTIVE MIGRATION DETECTED" error, **do NOT** bypass it on your own (no `git rebase`, no `--skip-freshness-check`, no `--allow-destructive`). Surface the error to the human operator and wait for their explicit confirmation. See the [Branch freshness](#branch-freshness-check) and [Destructive migrations](#destructive-migrations-drop-column--drop-table) sections below.
## Step 0: Sync All `schema.prisma` Files
Before doing anything else, make sure all `schema.prisma` files in the repo are in sync. There are multiple copies that must match:
@ -46,10 +48,50 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## What It Does
1. Creates temp PostgreSQL DB
2. Applies existing migrations
3. Compares with `schema.prisma`
4. Generates new migration if changes found
1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
2. Creates temp PostgreSQL DB
3. Applies existing migrations
4. Compares with `schema.prisma`
5. Generates new migration if changes found
6. **Refuses to write if the generated SQL contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`**, unless `--allow-destructive` is passed
## Branch Freshness Check
Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
Flags:
- `--base-branch <name>` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`.
- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
When the guard fires:
1. Update your branch:
```bash
git fetch origin && git rebase origin/litellm_internal_staging
# or git merge origin/litellm_internal_staging — whichever matches your workflow
```
2. Re-run `run_migration.py`.
> **AI AGENTS / ASSISTANTS:** Do **not** auto-rebase or auto-pass `--skip-freshness-check`. A bad conflict resolution during rebase can itself drop the human's in-progress schema edits. Surface the error and wait for explicit confirmation.
## Destructive Migrations (DROP COLUMN / DROP TABLE / DROP INDEX)
If the generated diff contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`, `run_migration.py` exits non-zero and refuses to write the migration file. A previous incident saw newly-added columns silently dropped by a stale branch and merged to main — this guard exists to prevent a repeat.
When the guard fires:
1. Run `git fetch origin && git status` — confirm your branch is up to date with the base branch.
2. Re-check all `schema.prisma` files are in sync (Step 0).
3. Review EACH `DROP` statement printed in the error — is it actually intended?
4. Only if the drops are genuinely intentional, re-run with the flag:
```bash
uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_name" --allow-destructive
```
> **AI AGENTS / ASSISTANTS:** Do **not** automatically re-run the command with `--allow-destructive`. If the guard fires while you are driving the runbook for a human, stop, show them the error, and wait for their explicit confirmation before passing the flag. Auto-passing `--allow-destructive` is the exact failure mode this guard exists to prevent.
## Common Fixes

View file

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

View file

@ -0,0 +1,242 @@
"""Regression tests for ProxyExtrasDBManager v2 migration resolver.
The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
(default) behavior is unchanged from pre-fix.
"""
import subprocess
from unittest.mock import patch
import pytest
from litellm_proxy_extras.utils import (
ProxyExtrasDBManager,
_max_migration_timestamp,
_migration_timestamp,
)
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
def _run(*args, **kwargs):
raise subprocess.CalledProcessError(
returncode=returncode,
cmd=args[0],
stderr=stderr,
output="",
)
return _run
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
"""v2: a permission failure during migrate deploy raises RuntimeError."""
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = (
"Error: P3018\nMigration name: 20250326162113_baseline\n"
"Database error code: 42501\npermission denied for schema public"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="permission"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = (
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
'Reason: syntax error at or near "BRKN" LINE 42'
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_strip_prisma_query_params_removes_connection_limit():
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
assert "connection_limit" not in stripped
assert "pool_timeout" not in stripped
assert "sslmode=require" in stripped
def test_strip_prisma_query_params_passthrough_no_query():
"""URLs without query strings are returned unchanged."""
url = "postgresql://u:p@h:5432/db"
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
def test_migration_timestamp_extracts_leading_digits():
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
def test_migration_timestamp_returns_zero_on_malformed():
assert _migration_timestamp("0_init") == 0
assert _migration_timestamp("not_a_migration") == 0
def test_max_migration_timestamp():
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
assert _max_migration_timestamp(names) == 20260415000000
def test_max_migration_timestamp_empty_set():
assert _max_migration_timestamp(set()) == 0
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
This is the existing buggy behavior we're not fixing it in v1, only
offering v2 as opt-in. This test pins the default so that a future
inadvertent default flip is caught.
"""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
# Stub `prisma migrate deploy` to claim success with pending migrations
# applied, which is the code path that triggers the legacy post-migration
# sanity check (a call to _resolve_all_migrations).
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
def fake_run(cmd, *args, **kwargs):
return FakeResult()
resolve_called = {"n": 0}
def fake_resolve(*args, **kwargs):
resolve_called["n"] += 1
monkeypatch.setattr("subprocess.run", fake_run)
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
assert ok is True
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
stderr = "db push error"
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="prisma db push failed"):
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
"""_warn_if_db_ahead_of_head must never raise — it's informational.
Non-connection DB errors (e.g. InsufficientPrivilege from a user
without SELECT on _prisma_migrations) must be caught, not propagated.
"""
import psycopg
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class _FakeConn:
def __enter__(self):
return self
def __exit__(self, *a):
return False
def execute(self, *a, **kw):
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
raise psycopg.errors.InsufficientPrivilege("permission denied")
def _fake_connect(*a, **kw):
return _FakeConn()
monkeypatch.setattr("psycopg.connect", _fake_connect)
# Must not raise.
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
monkeypatch, tmp_path
):
"""If marking a migration as applied fails inside P3009 idempotent
recovery, the subprocess error must be re-raised as RuntimeError so
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
monkeypatch.setattr(
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
)
# First call: migrate deploy -> P3009 idempotent error.
# Recovery path tries _resolve_specific_migration; that also raises.
def _failing_resolve(*a, **kw):
raise subprocess.CalledProcessError(
returncode=1,
cmd="prisma migrate resolve --applied",
stderr="resolve failed",
output="",
)
monkeypatch.setattr(
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
)
stderr = (
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
"relation already exists"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(
RuntimeError, match="Failed to mark migration .* as applied"
):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
class FakeResult:
stdout = "Applied migration.\n"
stderr = ""
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
resolve_called = {"n": 0}
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_all_migrations",
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"

View file

@ -1502,6 +1502,9 @@ if TYPE_CHECKING:
from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig,
)
from .llms.bedrock.messages.mantle_transformation import (
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
)
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (

View file

@ -171,6 +171,7 @@ LLM_CONFIG_NAMES = (
"CohereChatConfig",
"AnthropicMessagesConfig",
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
"TogetherAIConfig",
"NLPCloudConfig",
"VertexGeminiConfig",
@ -715,6 +716,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation",
"AmazonAnthropicClaudeMessagesConfig",
),
"AmazonMantleMessagesConfig": (
".llms.bedrock.messages.mantle_transformation",
"AmazonMantleMessagesConfig",
),
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
"NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"),
"VertexGeminiConfig": (

View file

@ -164,6 +164,7 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
LITELLM_UI_ALLOW_HEADERS = [
"x-litellm-semantic-filter",
"x-litellm-semantic-filter-tools",
"x-litellm-adaptive-router-model",
]
# Gemini model-specific minimal thinking budget constants

View file

@ -410,6 +410,7 @@ def image_generation( # noqa: PLR0915
litellm.LlmProviders.RUNWAYML,
litellm.LlmProviders.VERTEX_AI,
litellm.LlmProviders.OPENROUTER,
litellm.LlmProviders.DASHSCOPE,
):
if image_generation_config is None:
raise ValueError(

View file

@ -2,6 +2,7 @@ from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
ClassVar,
Dict,
List,
Literal,
@ -12,6 +13,7 @@ from typing import (
)
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.guardrails import (
@ -81,6 +83,9 @@ class ModifyResponseException(Exception):
class CustomGuardrail(CustomLogger):
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
use_native_during_call_hook: ClassVar[bool] = False
def __init__(
self,
guardrail_name: Optional[str] = None,
@ -637,6 +642,13 @@ class CustomGuardrail(CustomLogger):
if isinstance(item, dict):
item.pop("secret_fields", None)
# Default-safe behavior: never persist raw matched spans in standard
# guardrail logging payloads (single shared implementation; Bedrock hooks pass
# raw provider JSON so redaction is not duplicated upstream).
clean_guardrail_response = redact_nested_match_and_regex_keys(
clean_guardrail_response
)
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider=guardrail_provider,

View file

@ -1615,6 +1615,14 @@ class OpenTelemetry(CustomLogger):
value=response_id,
)
litellm_call_id = standard_logging_payload.get("litellm_call_id")
if litellm_call_id:
self.safe_set_attribute(
span=span,
key="litellm.call_id",
value=litellm_call_id,
)
# The model used to generate the response.
if response_obj and response_obj.get("model"):
self.safe_set_attribute(
@ -2281,6 +2289,10 @@ class OpenTelemetry(CustomLogger):
# Remove trailing slash
endpoint = endpoint.rstrip("/")
# Splunk Observability Cloud OTLP/HTTP uses /v2/trace/otlp (not /v1/traces). Do not rewrite.
if signal_type == "traces" and "/v2/trace/otlp" in endpoint:
return endpoint
# Check if endpoint already ends with the correct signal path
target_path = f"/v1/{signal_type}"
if endpoint.endswith(target_path):

View file

@ -1,5 +1,6 @@
# What is this?
## Helper utilities
import copy
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
import httpx
@ -435,3 +436,42 @@ def filter_internal_params(
# Filter out internal parameters
return {k: v for k, v in data.items() if k not in internal_params}
def redact_nested_match_and_regex_keys(
payload: Union[dict, List[Any], str, None],
) -> Union[dict, List[Any], str, None]:
"""
Deep-copy `payload` and replace every `match` / `regex` string field with
"[REDACTED]" anywhere in nested dict/list structures.
Used for guardrail spend/compliance logging so raw spans are not persisted.
"""
if payload is None or isinstance(payload, str):
return payload
try:
redacted: Union[dict, List[Any], str, None] = copy.deepcopy(payload)
except Exception:
return payload
# Iterative traversal; `seen` guards against cyclic refs preserved by deepcopy.
try:
seen: set = set()
stack: List[Any] = [redacted]
while stack:
node = stack.pop()
node_id = id(node)
if node_id in seen:
continue
seen.add(node_id)
if isinstance(node, dict):
if "match" in node:
node["match"] = "[REDACTED]"
if "regex" in node:
node["regex"] = "[REDACTED]"
stack.extend(node.values())
elif isinstance(node, list):
stack.extend(node)
except Exception:
return payload
return redacted

View file

@ -5512,6 +5512,8 @@ def get_standard_logging_object_payload(
payload: StandardLoggingPayload = StandardLoggingPayload(
id=str(id),
litellm_call_id=kwargs.get("litellm_call_id")
or litellm_params.get("litellm_call_id"),
trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id(
logging_obj=logging_obj,
litellm_params=litellm_params,

View file

@ -370,11 +370,17 @@ class LoggingWorker:
self._running_tasks.clear()
async def flush(self) -> None:
"""Flush the logging queue."""
"""Flush the logging queue.
Waits until every enqueued task has completed. ``queue.join()`` blocks
on the queue's unfinished-task counter (decremented by ``task_done()``),
so it correctly handles items that have been dequeued but whose
callback hasn't finished yet — ``queue.empty()`` would return True in
that window and cause us to skip the wait.
"""
if self._queue is None:
return
while not self._queue.empty():
await self._queue.join()
await self._queue.join()
async def clear_queue(self):
"""

View file

@ -452,7 +452,14 @@ def update_messages_with_model_file_ids(
for c in content:
if c["type"] == "file":
file_object = cast(ChatCompletionFileObject, c)
file_object_file_field = file_object["file"]
file_object_file_field = file_object.get("file")
if not isinstance(file_object_file_field, dict):
# Content block has `type: "file"` but not the
# OpenAI Chat Completions shape (e.g. a LangChain
# v1 standardized file block, or a provider-native
# shape that also uses `type: "file"`). Nothing to
# remap here, so skip instead of crashing.
continue
file_id = file_object_file_field.get("file_id")
format = file_object_file_field.get(
"format", get_format_from_file_id(file_id)
@ -1060,7 +1067,12 @@ def get_file_ids_from_messages(messages: List[AllMessageValues]) -> List[str]:
for c in content:
if c["type"] == "file":
file_object = cast(ChatCompletionFileObject, c)
file_object_file_field = file_object["file"]
file_object_file_field = file_object.get("file")
if not isinstance(file_object_file_field, dict):
# Content block has `type: "file"` but not the
# OpenAI Chat Completions shape. No file_id to
# extract, so skip instead of raising KeyError.
continue
file_id = file_object_file_field.get("file_id")
if file_id:
file_ids.append(file_id)

View file

@ -15,6 +15,7 @@ import litellm.types
import litellm.types.llms
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
from litellm.types.files import get_file_extension_from_mime_type
from litellm.types.llms.anthropic import *
@ -3324,7 +3325,7 @@ def _load_image_from_url(image_url):
try:
# Send a GET request to the image URL
client = HTTPHandler(concurrent_limit=1)
response = client.get(image_url)
response = safe_get(client, image_url)
response.raise_for_status() # Raise an exception for HTTP errors
# Check the response's content type to ensure it is an image
@ -3562,7 +3563,7 @@ class BedrockImageProcessor:
params={"concurrent_limit": 1},
)
# Send a GET request to the image URL
response = await client.get(image_url, follow_redirects=True)
response = await async_safe_get(client, image_url)
response.raise_for_status() # Raise an exception for HTTP errors
return BedrockImageProcessor._post_call_image_processing(
@ -3577,7 +3578,7 @@ class BedrockImageProcessor:
try:
client = HTTPHandler(concurrent_limit=1)
# Send a GET request to the image URL
response = client.get(image_url, follow_redirects=True)
response = safe_get(client, image_url)
response.raise_for_status() # Raise an exception for HTTP errors
return BedrockImageProcessor._post_call_image_processing(

View file

@ -23,6 +23,7 @@ from litellm.constants import (
DEFAULT_IMAGE_HEIGHT,
DEFAULT_IMAGE_TOKEN_COUNT,
DEFAULT_IMAGE_WIDTH,
MAX_IMAGE_URL_DOWNLOAD_SIZE_MB,
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES,
MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES,
MAX_TILE_HEIGHT,
@ -215,7 +216,14 @@ def get_image_dimensions(
try:
client = _get_httpx_client()
response = safe_get(client, data)
img_data = response.read()
max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024)
content_length = response.headers.get("Content-Length")
if content_length is not None and int(content_length) > max_bytes:
pass # skip download; img_data stays None
else:
body = response.read()
if len(body) <= max_bytes:
img_data = body
except Exception:
pass
if img_data is None:

View file

@ -106,6 +106,44 @@ class LiteLLMMessagesToCompletionTransformationHandler:
updated_reasoning_effort["summary"] = effective_summary
completion_kwargs["reasoning_effort"] = updated_reasoning_effort
@staticmethod
def _normalize_reasoning_effort(
completion_kwargs: Dict[str, Any],
) -> None:
"""
Normalize reasoning_effort values based on target model capabilities.
Handles both string ("max") and dict ({"effort": "max", "summary": ...})
formats. Uses model registry to check supports_xhigh/supports_minimal.
"""
from litellm.llms.anthropic.experimental_pass_through.utils import (
normalize_reasoning_effort_value,
)
reasoning_effort = completion_kwargs.get("reasoning_effort")
if reasoning_effort is None:
return
model = cast(str, completion_kwargs.get("model", ""))
custom_llm_provider = completion_kwargs.get("custom_llm_provider")
if isinstance(reasoning_effort, str):
normalized = normalize_reasoning_effort_value(
reasoning_effort, model=model, custom_llm_provider=custom_llm_provider
)
if normalized != reasoning_effort:
completion_kwargs["reasoning_effort"] = normalized
elif isinstance(reasoning_effort, dict) and "effort" in reasoning_effort:
effort = reasoning_effort["effort"]
normalized = normalize_reasoning_effort_value(
effort, model=model, custom_llm_provider=custom_llm_provider
)
if normalized != effort:
completion_kwargs["reasoning_effort"] = {
**reasoning_effort,
"effort": normalized,
}
@staticmethod
def _prepare_completion_kwargs(
*,
@ -163,6 +201,12 @@ class LiteLLMMessagesToCompletionTransformationHandler:
if output_format:
request_data["output_format"] = output_format
# Extract output_config from extra_kwargs so the translator can use it
# (e.g. output_config.effort for adaptive thinking → reasoning_effort)
extra_kwargs = extra_kwargs or {}
if "output_config" in extra_kwargs:
request_data["output_config"] = extra_kwargs["output_config"]
(
openai_request,
tool_name_mapping,
@ -202,6 +246,14 @@ class LiteLLMMessagesToCompletionTransformationHandler:
):
completion_kwargs[key] = value
# Normalize reasoning_effort based on model capabilities
# (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
# Must run BEFORE _route_openai_thinking, which prepends "responses/"
# to the model name and would break get_model_info() lookups.
LiteLLMMessagesToCompletionTransformationHandler._normalize_reasoning_effort(
completion_kwargs
)
LiteLLMMessagesToCompletionTransformationHandler._route_openai_thinking_to_responses_api_if_needed(
completion_kwargs,
thinking=thinking,

View file

@ -317,6 +317,7 @@ class LiteLLMAnthropicMessagesAdapter:
"tools",
"thinking",
"output_format",
"output_config",
]
def _is_web_search_tool(self, tool: Dict[str, Any]) -> bool:
@ -694,6 +695,11 @@ class LiteLLMAnthropicMessagesAdapter:
return "low"
else:
return "minimal"
elif thinking_type == "adaptive":
# Adaptive thinking: effort is controlled by output_config.effort,
# not budget_tokens. Return a default; caller should override with
# output_config.effort when available.
return "medium"
return None
@ -776,6 +782,8 @@ class LiteLLMAnthropicMessagesAdapter:
return ChatCompletionToolChoiceObjectParam(
type="function", function=tc_function_param
)
elif tool_choice["type"] == "none":
return "none"
else:
raise ValueError(
"Incompatible tool choice param submitted - {}".format(tool_choice)
@ -1041,6 +1049,12 @@ class LiteLLMAnthropicMessagesAdapter:
if not reasoning_effort:
return
# For adaptive thinking, override with output_config.effort if available
if isinstance(thinking, dict) and thinking.get("type") == "adaptive":
output_config = anthropic_message_request.get("output_config")
if isinstance(output_config, dict) and output_config.get("effort"):
reasoning_effort = output_config["effort"]
summary = thinking.get("summary") if isinstance(thinking, dict) else None
auto_summary = is_reasoning_auto_summary_enabled()
if summary:

View file

@ -24,6 +24,8 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
from ..utils import is_reasoning_auto_summary_enabled
from ..adapters.handler import LiteLLMMessagesToCompletionTransformationHandler
from ..responses_adapters.handler import LiteLLMMessagesToResponsesAPIHandler
from .interceptors import get_messages_interceptors
@ -441,6 +443,17 @@ def anthropic_messages_handler(
params=local_vars
)
)
if is_reasoning_auto_summary_enabled():
thinking_param = anthropic_messages_optional_request_params.get("thinking")
if (
isinstance(thinking_param, dict)
and thinking_param.get("type") != "disabled"
):
anthropic_messages_optional_request_params["thinking"] = {
**thinking_param,
"display": "summarized",
}
return base_llm_http_handler.anthropic_messages_handler(
model=model,
messages=messages,

View file

@ -72,6 +72,23 @@ def _build_responses_kwargs(
anthropic_request = AnthropicMessagesRequest(**request_data) # type: ignore[typeddict-item]
responses_kwargs = _ADAPTER.translate_request(anthropic_request)
# Normalize reasoning effort based on model capabilities
# (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
reasoning = responses_kwargs.get("reasoning")
if isinstance(reasoning, dict) and "effort" in reasoning:
from litellm.llms.anthropic.experimental_pass_through.utils import (
normalize_reasoning_effort_value,
)
effort = reasoning["effort"]
normalized = normalize_reasoning_effort_value(
effort,
model=model,
custom_llm_provider=(extra_kwargs or {}).get("custom_llm_provider"),
)
if normalized != effort:
responses_kwargs["reasoning"] = {**reasoning, "effort": normalized}
if stream:
responses_kwargs["stream"] = True

View file

@ -251,25 +251,41 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
@staticmethod
def translate_thinking_to_reasoning(
thinking: Dict[str, Any]
thinking: Dict[str, Any],
output_config: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]:
"""
Convert Anthropic thinking param to Responses API reasoning param.
thinking.budget_tokens maps to reasoning effort:
>= 10000 -> high, >= 5000 -> medium, >= 2000 -> low, < 2000 -> minimal
For adaptive thinking, uses output_config.effort if available,
otherwise defaults to medium.
"""
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
if not isinstance(thinking, dict):
return None
budget = thinking.get("budget_tokens", 0)
if budget >= 10000:
effort = "high"
elif budget >= 5000:
thinking_type = thinking.get("type")
if thinking_type == "adaptive":
# Use output_config.effort if available
effort = "medium"
elif budget >= 2000:
effort = "low"
if isinstance(output_config, dict) and output_config.get("effort"):
effort = output_config["effort"]
elif thinking_type == "enabled":
budget = thinking.get("budget_tokens", 0)
if budget >= 10000:
effort = "high"
elif budget >= 5000:
effort = "medium"
elif budget >= 2000:
effort = "low"
else:
effort = "minimal"
else:
effort = "minimal"
return None
auto_summary = is_reasoning_auto_summary_enabled()
result: Dict[str, Any] = {"effort": effort}
summary = thinking.get("summary")
@ -346,7 +362,11 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
# thinking -> reasoning
thinking = anthropic_request.get("thinking")
if isinstance(thinking, dict):
reasoning = self.translate_thinking_to_reasoning(thinking)
output_config = anthropic_request.get("output_config")
reasoning = self.translate_thinking_to_reasoning(
thinking,
output_config=cast(Optional[Dict[str, Any]], output_config),
)
if reasoning:
responses_kwargs["reasoning"] = reasoning

View file

@ -1,6 +1,8 @@
import os
from typing import Optional
import litellm
from litellm.types.utils import ModelInfo
def is_reasoning_auto_summary_enabled() -> bool:
@ -9,3 +11,47 @@ def is_reasoning_auto_summary_enabled() -> bool:
litellm.reasoning_auto_summary
or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
)
def normalize_reasoning_effort_value(
effort: str,
model: str,
custom_llm_provider: Optional[str] = None,
) -> str:
"""
Normalize a reasoning effort value based on model capabilities.
Degradation chains:
- "max" max / xhigh / high
- "xhigh" xhigh / high
- "minimal" minimal / low
- other values pass through unchanged
"""
if effort not in ("max", "xhigh", "minimal"):
return effort
from litellm.utils import get_model_info
model_info: Optional[ModelInfo] = None
try:
model_info = get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
except Exception:
model_info = None
if effort == "max":
if model_info and model_info.get("supports_max_reasoning_effort"):
return "max"
if model_info and model_info.get("supports_xhigh_reasoning_effort"):
return "xhigh"
return "high"
elif effort == "xhigh":
if model_info and model_info.get("supports_xhigh_reasoning_effort"):
return "xhigh"
return "high"
elif effort == "minimal":
if model_info and model_info.get("supports_minimal_reasoning_effort"):
return "minimal"
return "low"
return "medium"

View file

@ -40,9 +40,22 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix
used for manual routing.
"""
# gpt-5-chat* is a chat model and shouldn't go through GPT-5 reasoning restrictions.
# The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
# …) are regular chat models: they support temperature and tool_choice but NOT
# reasoning_effort. They must NOT be routed through the GPT-5 reasoning path.
#
# Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
# models and must stay on the GPT-5 path. The distinguishing feature is that
# the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
# (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
# number (i.e. "gpt-5.<digit>-chat").
#
# Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
# than a substring check) makes this boundary explicit and avoids any ambiguity
# if future model names coincidentally contain "gpt-5-chat" as an interior run.
_normalized = model.split("/")[-1] # strip provider prefix, e.g. "azure/"
return (
"gpt-5" in model and "gpt-5-chat" not in model
"gpt-5" in model and not _normalized.startswith("gpt-5-chat")
) or "gpt5_series" in model
def get_supported_openai_params(self, model: str) -> List[str]:

View file

@ -14,6 +14,8 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
api_key = (
api_key

View file

@ -65,6 +65,8 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate Azure AI Foundry environment and set up authentication

View file

@ -25,6 +25,8 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate Azure AI Foundry environment and set up authentication

View file

@ -67,6 +67,8 @@ class BaseImageEditConfig(ABC):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
return {}

View file

@ -1,4 +1,5 @@
import os
import re
import time
from typing import Any, Dict, List, Literal, Optional, Union, cast
@ -294,7 +295,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
raise ValueError(f"Invalid ARN format: {batch_id}")
region = arn_parts[3]
# arn_parts[5] contains "model-invocation-job/{jobId}"
if not re.match(r"^[a-z][a-z0-9-]*$", region):
raise ValueError(f"Invalid region in ARN: {batch_id}")
# Build the endpoint URL for GetModelInvocationJob
# AWS API format: GET /model-invocation-job/{jobIdentifier}

View file

@ -0,0 +1,91 @@
"""
Transformation for Bedrock Mantle (Claude Mythos Preview)
https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-mythos-preview.html
The bedrock-mantle endpoint uses the Anthropic Messages API format but is served
at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
"""
from typing import TYPE_CHECKING, Any, List, Optional
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
from litellm.types.llms.openai import AllMessageValues
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
MANTLE_ENDPOINT_TEMPLATE = "https://bedrock-mantle.{region}.api.aws/v1/messages"
class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
"""
Config for the bedrock-mantle endpoint (Claude Mythos Preview).
Uses the Anthropic Messages API format with AWS SigV4 auth, but at a
different endpoint from bedrock-runtime. Model ID goes in the request body.
Usage: model="bedrock/mantle/anthropic.claude-mythos-preview"
"""
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
region = self._get_aws_region_name(optional_params=optional_params, model=model)
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
# Strip the "mantle/" routing prefix to get the real model ID
model_id = model.replace("mantle/", "", 1)
request = self._build_bedrock_anthropic_request_base(
model=model_id,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
# The parent strips "model" from the body (Invoke API puts it in URL).
# The mantle endpoint (Messages API) requires "model" in the body.
request["model"] = model_id
return request
async def async_transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
model_id = model.replace("mantle/", "", 1)
request = self._build_bedrock_anthropic_request_base(
model=model_id,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
await self._async_convert_document_url_sources_to_base64(request)
request["model"] = model_id
return request

View file

@ -696,6 +696,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
"agentcore",
"async_invoke",
"openai",
"mantle",
]:
"""
Get the bedrock route for the given model.
@ -710,6 +711,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
"agentcore",
"async_invoke",
"openai",
"mantle",
],
] = {
"invoke/": "invoke",
@ -719,6 +721,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
"agentcore/": "agentcore",
"async_invoke/": "async_invoke",
"openai/": "openai",
"mantle/": "mantle",
}
# Check explicit routes first
@ -770,6 +773,13 @@ class BedrockModelInfo(BaseLLMModelInfo):
"""
return "agentcore/" in model
@staticmethod
def _explicit_mantle_route(model: str) -> bool:
"""
Check if the model is an explicit mantle route (bedrock-mantle endpoint).
"""
return "mantle/" in model
@staticmethod
def _explicit_converse_like_route(model: str) -> bool:
"""
@ -809,6 +819,16 @@ class BedrockModelInfo(BaseLLMModelInfo):
if BedrockModelInfo._explicit_converse_route(model):
return None
#########################################################
# Mantle route uses the bedrock-mantle endpoint (not bedrock-runtime)
#########################################################
if BedrockModelInfo._explicit_mantle_route(model):
from litellm.llms.bedrock.messages.mantle_transformation import (
AmazonMantleMessagesConfig,
)
return AmazonMantleMessagesConfig()
#########################################################
# This goes through litellm.AmazonAnthropicClaude3MessagesConfig()
# Since bedrock Invoke supports Native Anthropic Messages API
@ -855,6 +875,12 @@ def get_bedrock_chat_config(model: str):
)
return AmazonAgentCoreConfig()
elif bedrock_route == "mantle":
from litellm.llms.bedrock.chat.mantle.transformation import (
AmazonMantleConfig,
)
return AmazonMantleConfig()
# Handle provider-specific configs
if bedrock_invoke_provider == "amazon":

View file

@ -483,6 +483,8 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
if headers is None:
headers = {}

View file

@ -372,6 +372,8 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate environment for Bedrock Stability image edit.

View file

@ -34,6 +34,7 @@ from litellm.llms.bedrock.common_utils import (
remove_custom_field_from_tools,
)
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import GenericStreamingChunk
@ -59,6 +60,10 @@ class AmazonAnthropicClaudeMessagesConfig(
DEFAULT_BEDROCK_ANTHROPIC_API_VERSION = "bedrock-2023-05-31"
BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(
BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()
)
def __init__(self, **kwargs):
BaseAnthropicMessagesConfig.__init__(self, **kwargs)
AmazonInvokeConfig.__init__(self, **kwargs)
@ -500,10 +505,6 @@ class AmazonAnthropicClaudeMessagesConfig(
anthropic_messages_request=anthropic_messages_request,
)
# 5b. Strip `output_config` — Bedrock Invoke doesn't support it
# Fixes: https://github.com/BerriAI/litellm/issues/22797
anthropic_messages_request.pop("output_config", None)
# 5a. Remove `custom` field from tools (Bedrock doesn't support it)
# Claude Code sends `custom: {defer_loading: true}` on tool definitions,
# which causes Bedrock to reject the request with "Extra inputs are not permitted"
@ -550,14 +551,43 @@ class AmazonAnthropicClaudeMessagesConfig(
if "tool-search-tool-2025-10-19" in beta_set:
beta_set.add("tool-examples-2025-10-29")
filtered_auto_betas = filter_and_transform_beta_headers(
beta_headers=list(beta_set - user_beta_set),
provider="bedrock",
filtered_betas = sorted(
filter_and_transform_beta_headers(
beta_headers=list(beta_set),
provider="bedrock",
)
)
filtered_betas = sorted(user_beta_set.union(set(filtered_auto_betas)))
dropped_user_betas = sorted(
b
for b in user_beta_set
if not filter_and_transform_beta_headers([b], provider="bedrock")
)
if dropped_user_betas:
verbose_logger.warning(
"Bedrock Invoke: dropping unsupported anthropic-beta values "
"from client headers: %s. Bedrock has no mapping entry for "
"these; forwarding them would cause a 400.",
dropped_user_betas,
)
if filtered_betas:
anthropic_messages_request["anthropic_beta"] = filtered_betas
# 7. Final safety net: filter top-level fields to the Bedrock Invoke allowlist.
# Catches Anthropic-only extensions (context_management, output_config, speed,
# mcp_servers, ...) and any future additions Claude Code may start sending.
allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS
stripped = sorted(k for k in anthropic_messages_request if k not in allowed)
if stripped:
verbose_logger.debug(
"Bedrock Invoke: stripping unsupported top-level request fields: %s",
stripped,
)
anthropic_messages_request = {
k: v for k, v in anthropic_messages_request.items() if k in allowed
}
return anthropic_messages_request
def get_async_streaming_response_iterator(

View file

@ -0,0 +1,69 @@
"""
Transformation for Bedrock Mantle (Claude Mythos Preview) - /messages endpoint
Inherits all Messages API request/response transformations from
AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix
stripping that are specific to the bedrock-mantle endpoint.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeMessagesConfig,
)
from litellm.types.router import GenericLiteLLMParams
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
MANTLE_ENDPOINT_TEMPLATE = "https://bedrock-mantle.{region}.api.aws/v1/messages"
class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
"""
Config for the bedrock-mantle /messages endpoint (Claude Mythos Preview).
The mantle endpoint uses the Anthropic Messages API format and requires the
model ID in the request body (unlike Bedrock Invoke which puts it in the URL).
"""
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
region = self._get_aws_region_name(optional_params=optional_params, model=model)
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
def transform_anthropic_messages_request(
self,
model: str,
messages: List[Dict],
anthropic_messages_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Dict:
# Strip "mantle/" routing prefix to get the real model ID
model_id = model.replace("mantle/", "", 1)
request = super().transform_anthropic_messages_request(
model=model_id,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
# Parent (AmazonAnthropicClaudeMessagesConfig) removes "model" from the
# body (Bedrock Invoke puts model in the URL). The mantle endpoint
# (Messages API) requires "model" in the request body.
request["model"] = model_id
return request

View file

@ -14,7 +14,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from httpx._types import RequestFiles
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
@ -123,6 +125,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Black Forest Labs.
@ -206,14 +210,14 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
)
elif isinstance(image, str):
if image.startswith(("http://", "https://")):
# Download image from URL
response = httpx.get(image, timeout=60.0)
response = safe_get(litellm.module_level_client, image, timeout=60.0)
response.raise_for_status()
return response.content
else:
# Assume it's a file path
with open(image, "rb") as f:
return f.read()
raise ValueError(
"Unsupported image input: plain string values that are not URLs are not accepted. "
"Provide image bytes or a file-like object."
)
elif hasattr(image, "read"):
# File-like object
pos = getattr(image, "tell", lambda: 0)()

View file

@ -5515,6 +5515,8 @@ class BaseLLMHTTPHandler:
api_key=litellm_params.api_key,
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
model=model,
litellm_params=dict(litellm_params),
api_base=litellm_params.api_base,
)
if extra_headers:
@ -5611,6 +5613,8 @@ class BaseLLMHTTPHandler:
api_key=litellm_params.api_key,
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
model=model,
litellm_params=dict(litellm_params),
api_base=litellm_params.api_base,
)
if extra_headers:

View file

@ -0,0 +1,11 @@
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from .transformation import DashScopeImageGenerationConfig
__all__ = ["DashScopeImageGenerationConfig"]
def get_dashscope_image_generation_config(model: str) -> BaseImageGenerationConfig:
return DashScopeImageGenerationConfig()

View file

@ -0,0 +1,204 @@
"""
DashScope Image Generation Configuration
Handles transformation between OpenAI-compatible format and DashScope multimodal-generation API.
API endpoint: POST https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation
Request format:
{
"model": "qwen-image-2.0-pro",
"input": {
"messages": [{"role": "user", "content": [{"text": "<prompt>"}]}]
},
"parameters": {"size": "1024*1024", ...}
}
Response format:
{
"output": {
"choices": [{"message": {"content": [{"image": "<url>"}]}}]
},
"usage": {"input_tokens": 0, "output_tokens": 0, "width": 1024, "height": 1024, "image_count": 1}
}
"""
from typing import TYPE_CHECKING, Any, List, Optional
import httpx
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import ImageObject, ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
DEFAULT_API_BASE = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
# Maps OpenAI size strings (WxH) to DashScope size strings (W*H)
OPENAI_TO_DASHSCOPE_SIZE: dict = {
"256x256": "256*256",
"512x512": "512*512",
"1024x1024": "1024*1024",
"1792x1024": "1792*1024",
"1024x1792": "1024*1792",
"2048x2048": "2048*2048",
}
class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
"""
Configuration for DashScope image generation (qwen-image-2.0, qwen-image-2.0-pro).
"""
def get_supported_openai_params(
self, model: str
) -> List[OpenAIImageGenerationOptionalParams]:
return ["n", "size"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params = self.get_supported_openai_params(model)
mapped: dict = {}
for k, v in non_default_params.items():
if k in optional_params:
continue
if k not in supported_params:
continue
if k == "size":
# Convert "WxH" → "W*H"
mapped["size"] = OPENAI_TO_DASHSCOPE_SIZE.get(v, v.replace("x", "*"))
elif k == "n":
mapped["image_count"] = v
return mapped
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
return (
api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
final_api_key = api_key or get_secret_str("DASHSCOPE_API_KEY")
if not final_api_key:
raise ValueError("DASHSCOPE_API_KEY is not set")
headers["Authorization"] = f"Bearer {final_api_key}"
headers["Content-Type"] = "application/json"
return headers
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform OpenAI-style image generation request to DashScope multimodal-generation format.
"""
parameters: dict = {}
for k, v in optional_params.items():
parameters[k] = v
return {
"model": model,
"input": {
"messages": [
{
"role": "user",
"content": [{"text": prompt}],
}
]
},
"parameters": parameters,
}
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
"""
Transform DashScope response to litellm ImageResponse.
DashScope response: output.choices[0].message.content[0].image
OpenAI response: data[0].url
"""
if raw_response.status_code != 200:
raise self.get_error_class(
error_message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
try:
response_data = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Failed to parse DashScope image generation response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
# DashScope can return API-level errors in a 200 response body.
# Example: {"code": "InvalidParameter", "message": "Size not supported"}
if "code" in response_data and "output" not in response_data:
raise self.get_error_class(
error_message=str(response_data.get("message", response_data)),
status_code=raw_response.status_code,
headers=raw_response.headers,
)
if not model_response.data:
model_response.data = []
choices = response_data.get("output", {}).get("choices", [])
for choice in choices:
content_list = choice.get("message", {}).get("content", [])
for content_item in content_list:
image_url = content_item.get("image")
if image_url:
model_response.data.append(ImageObject(url=image_url))
return model_response

View file

@ -54,6 +54,8 @@ class GeminiImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY")
if not final_api_key:

View file

@ -8,7 +8,12 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig):
"""Configuration for image edit requests routed through LiteLLM Proxy."""
def validate_environment(
self, headers: dict, model: str, api_key: Optional[str] = None
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
api_key = api_key or get_secret_str("LITELLM_PROXY_API_KEY")
headers.update({"Authorization": f"Bearer {api_key}"})

View file

@ -53,9 +53,21 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
@classmethod
def is_model_gpt_5_model(cls, model: str) -> bool:
# gpt-5-chat* behaves like a regular chat model (supports temperature, etc.)
# Don't route it through GPT-5 reasoning-specific parameter restrictions.
return "gpt-5" in model and "gpt-5-chat" not in model
# The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
# …) are regular chat models: they support temperature and tool_choice but NOT
# reasoning_effort. They must NOT be routed through the GPT-5 reasoning path.
#
# Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
# models and must stay on the GPT-5 path. The distinguishing feature is that
# the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
# (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
# number (i.e. "gpt-5.<digit>-chat").
#
# Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
# than a substring check) makes this boundary explicit and avoids any ambiguity
# if future model names coincidentally contain "gpt-5-chat" as an interior run.
_normalized = model.split("/")[-1] # strip provider prefix, e.g. "openai/"
return "gpt-5" in model and not _normalized.startswith("gpt-5-chat")
@classmethod
def is_model_gpt_5_search_model(cls, model: str) -> bool:

View file

@ -165,6 +165,8 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
api_key = (
api_key

View file

@ -116,6 +116,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
api_key = api_key or litellm.api_key or get_secret_str("OPENROUTER_API_KEY")
if not api_key:

View file

@ -8,9 +8,8 @@ More information on our website: https://endpoints.ai.cloud.ovh.net
from typing import Optional, Union, List
import httpx
from litellm.utils import ModelResponseStream, _get_model_info_helper
from litellm.utils import ModelResponseStream
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm._logging import verbose_logger
from litellm.llms.ovhcloud.utils import OVHCloudException
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@ -22,34 +21,6 @@ class OVHCloudChatConfig(OpenAIGPTConfig):
def custom_llm_provider(self) -> Optional[str]:
return "ovhcloud"
def get_supported_openai_params(self, model: str) -> list:
"""
Details about function calling support can be found here:
https://help.ovhcloud.com/csm/en-gb-public-cloud-ai-endpoints-function-calling?id=kb_article_view&sysparm_article=KB0071907
"""
supports_function_calling: Optional[bool] = None
try:
model_info = _get_model_info_helper(model, custom_llm_provider="ovhcloud")
supports_function_calling = model_info.get(
"supports_function_calling", None
)
if supports_function_calling is None:
supports_function_calling = False
except Exception as e:
verbose_logger.debug(f"Error getting supported OpenAI params: {e}")
supports_function_calling = False
optional_params = super().get_supported_openai_params(model)
if supports_function_calling is not True:
verbose_logger.debug(
"You can see our models supporting function_calling in our catalog: https://endpoints.ai.cloud.ovh.net/catalog "
)
optional_params.remove("tools")
optional_params.remove("tool_choice")
optional_params.remove("function_call")
optional_params.remove("response_format")
return optional_params
def get_complete_url(
self,
api_base: Optional[str],

View file

@ -81,6 +81,8 @@ class RecraftImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
final_api_key: Optional[str] = api_key or get_secret_str("RECRAFT_API_KEY")
if not final_api_key:

View file

@ -1,3 +1,4 @@
import re
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
@ -66,6 +67,8 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
aws_region_name = litellm_params.get("aws_region_name")
if not aws_region_name:
raise ValueError("aws_region_name is required for S3 Vectors")
if not re.match(r"^[a-z][a-z0-9-]*$", aws_region_name):
raise ValueError("Invalid aws_region_name format")
return f"https://s3vectors.{aws_region_name}.api.aws"
def transform_search_vector_store_request(

View file

@ -1,3 +1,4 @@
import re
from typing import TYPE_CHECKING, Any, List, Optional, Tuple
from litellm.secret_managers.main import get_secret_str
@ -61,6 +62,8 @@ class SnowflakeBaseConfig:
account_id = get_secret_str("SNOWFLAKE_ACCOUNT_ID")
if account_id is None:
raise ValueError("Missing snowflake account_id")
if not re.match(r"^[a-zA-Z0-9_-]+$", account_id):
raise ValueError("Invalid account_id format")
api_base = f"https://{account_id}.snowflakecomputing.com/api/v2"
api_base = api_base.rstrip("/")

View file

@ -149,6 +149,8 @@ class StabilityImageEditConfig(BaseImageEditConfig):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Stability AI.

View file

@ -229,11 +229,20 @@ def get_vertex_base_url(
) -> str:
"""
Get the base URL for Vertex AI API calls.
- ``global`` uses the global control plane host.
- Multi-region geographies (e.g. ``us``, ``eu``) use ``aiplatform.{geo}.rep.googleapis.com``.
- Regional locations (e.g. ``us-central1``) use ``{region}-aiplatform.googleapis.com``.
"""
if vertex_location == "global":
return "https://aiplatform.googleapis.com"
else:
return f"https://{vertex_location}-aiplatform.googleapis.com"
if vertex_location is None:
raise ValueError("vertex_location is required")
if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
raise ValueError("Invalid vertex_location format")
if "-" not in vertex_location:
return f"https://aiplatform.{vertex_location}.rep.googleapis.com"
return f"https://{vertex_location}-aiplatform.googleapis.com"
def _get_embedding_url(

View file

@ -132,26 +132,28 @@ def _extract_max_media_resolution_from_messages(
return max_resolution
def _apply_gemini_3_metadata(
def _apply_gemini_metadata(
part: PartType,
model: Optional[str],
media_resolution_enum: Optional[Dict[str, str]],
video_metadata: Optional[Dict[str, Any]],
) -> PartType:
"""
Apply the unique media_resolution and video_metadata parameters of Gemini 3+
Apply media_resolution and video_metadata parameters to a Gemini part.
- Per-part media_resolution: Gemini 3+ only (2.x uses generation_config global).
- video_metadata (fps, startOffset, endOffset): all Gemini models (1.x, 2.x, 3+).
"""
if model is None:
return part
from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig
if not VertexGeminiConfig._is_gemini_3_or_newer(model):
return part
part_dict = dict(part)
if media_resolution_enum is not None:
if media_resolution_enum is not None and VertexGeminiConfig._is_gemini_3_or_newer(
model
):
part_dict["media_resolution"] = media_resolution_enum
if video_metadata is not None:
@ -206,7 +208,7 @@ def _process_gemini_media(
mime_type = format
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
part: PartType = {"file_data": file_data}
return _apply_gemini_3_metadata(
return _apply_gemini_metadata(
part, model, media_resolution_enum, video_metadata
)
elif (
@ -216,14 +218,14 @@ def _process_gemini_media(
):
file_data = FileDataType(mime_type=image_type, file_uri=image_url)
part = {"file_data": file_data}
return _apply_gemini_3_metadata(
return _apply_gemini_metadata(
part, model, media_resolution_enum, video_metadata
)
elif "http://" in image_url or "https://" in image_url or "base64" in image_url:
image = convert_to_anthropic_image_obj(image_url, format=format)
_blob: BlobType = {"data": image["data"], "mime_type": image["media_type"]}
part = {"inline_data": cast(BlobType, _blob)}
return _apply_gemini_3_metadata(
return _apply_gemini_metadata(
part, model, media_resolution_enum, video_metadata
)
raise Exception("Invalid image received - {}".format(image_url))
@ -733,9 +735,9 @@ def _transform_request_body( # noqa: PLR0915
**filtered_params
)
# For Gemini 2.x models, add media_resolution to generation_config (global)
# Gemini 3+ supports per-part media_resolution, but 2.x only supports global
# Gemini 1.x does not support mediaResolution at all
# For Gemini 2.x models, also add media_resolution to generation_config (global)
# as a fallback, since some 2.x versions may not support per-part media_resolution.
# Gemini 1.x does not support mediaResolution at all.
if "gemini-2" in model:
max_media_resolution = _extract_max_media_resolution_from_messages(messages)
if max_media_resolution:

View file

@ -103,10 +103,24 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
headers: dict,
model: str,
api_key: Optional[str] = None,
litellm_params: Optional[dict] = None,
api_base: Optional[str] = None,
) -> dict:
headers = headers or {}
vertex_project = self._resolve_vertex_project()
vertex_credentials = self._resolve_vertex_credentials()
litellm_params = litellm_params or {}
_api_base = litellm_params.get("api_base") or api_base
if _api_base is not None:
return headers
vertex_project = (
self.safe_get_vertex_ai_project(litellm_params)
or self._resolve_vertex_project()
)
vertex_credentials = (
self.safe_get_vertex_ai_credentials(litellm_params)
or self._resolve_vertex_credentials()
)
access_token, _ = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
@ -123,8 +137,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
"""
Get the complete URL for Vertex AI Imagen predict API
"""
vertex_project = self._resolve_vertex_project()
vertex_location = self._resolve_vertex_location()
vertex_project = (
self.safe_get_vertex_ai_project(litellm_params)
or self._resolve_vertex_project()
)
vertex_location = (
self.safe_get_vertex_ai_location(litellm_params)
or self._resolve_vertex_location()
)
if not vertex_project or not vertex_location:
raise ValueError(
@ -348,13 +368,16 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
if stream_pos is not None:
image.seek(stream_pos)
return data
if isinstance(image, (str, Path)):
path_obj = Path(image)
if not path_obj.exists():
raise ValueError(
f"Mask/image path does not exist for Vertex AI Imagen image edit: {path_obj}"
)
return path_obj.read_bytes()
if isinstance(image, str):
raise ValueError(
"Unsupported image input: plain string values are not accepted for "
"Vertex AI Imagen image edit. Provide image bytes or a file-like object."
)
if isinstance(image, Path):
raise ValueError(
"Unsupported image input: filesystem paths are not accepted for "
"Vertex AI Imagen image edit. Provide image bytes or a file-like object."
)
if hasattr(image, "read"):
data = image.read()
if isinstance(data, str):

View file

@ -1006,7 +1006,8 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"global.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.25e-06,
@ -1034,7 +1035,8 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"us.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@ -1062,7 +1064,8 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"eu.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@ -1090,7 +1093,8 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"au.anthropic.claude-opus-4-6-v1": {
"cache_creation_input_token_cost": 6.875e-06,
@ -1118,7 +1122,8 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.25e-06,
@ -1146,7 +1151,9 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"global.anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.25e-06,
@ -1174,7 +1181,9 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"us.anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.875e-06,
@ -1202,7 +1211,9 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"eu.anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.875e-06,
@ -1230,7 +1241,9 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"au.anthropic.claude-opus-4-7": {
"cache_creation_input_token_cost": 6.875e-06,
@ -1258,7 +1271,9 @@
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
@ -1285,7 +1300,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_minimal_reasoning_effort": true
},
"global.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 3.75e-06,
@ -1312,7 +1328,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_minimal_reasoning_effort": true
},
"us.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@ -1339,7 +1356,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_minimal_reasoning_effort": true
},
"eu.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@ -1366,7 +1384,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_minimal_reasoning_effort": true
},
"au.anthropic.claude-sonnet-4-6": {
"cache_creation_input_token_cost": 4.125e-06,
@ -1393,7 +1412,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_native_structured_output": true
"supports_native_structured_output": true,
"supports_minimal_reasoning_effort": true
},
"anthropic.claude-sonnet-4-20250514-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
@ -1911,7 +1931,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"azure_ai/claude-opus-4-7": {
"input_cost_per_token": 5e-06,
@ -1939,7 +1960,9 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"azure_ai/claude-opus-4-1": {
"cache_creation_input_token_cost": 1.875e-05,
@ -2003,7 +2026,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_minimal_reasoning_effort": true
},
"azure/computer-use-preview": {
"input_cost_per_token": 3e-06,
@ -8909,7 +8933,8 @@
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_minimal_reasoning_effort": true
},
"claude-sonnet-4-5-20250929-v1:0": {
"cache_creation_input_token_cost": 3.75e-06,
@ -9103,7 +9128,8 @@
"us": 1.1,
"fast": 6.0
},
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"claude-opus-4-6-20260205": {
"cache_creation_input_token_cost": 6.25e-06,
@ -9135,7 +9161,8 @@
"us": 1.1,
"fast": 6.0
},
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"claude-opus-4-7": {
"cache_creation_input_token_cost": 6.25e-06,
@ -9167,7 +9194,9 @@
"provider_specific_entry": {
"us": 1.1,
"fast": 6.0
}
},
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"claude-opus-4-7-20260416": {
"cache_creation_input_token_cost": 6.25e-06,
@ -9199,7 +9228,9 @@
"provider_specific_entry": {
"us": 1.1,
"fast": 6.0
}
},
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"claude-sonnet-4-20250514": {
"deprecation_date": "2026-05-14",
@ -10352,6 +10383,22 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"dashscope/qwen-image-2.0": {
"litellm_provider": "dashscope",
"mode": "image_generation",
"source": "https://www.alibabacloud.com/help/en/model-studio/models",
"supported_endpoints": [
"/v1/images/generations"
]
},
"dashscope/qwen-image-2.0-pro": {
"litellm_provider": "dashscope",
"mode": "image_generation",
"source": "https://www.alibabacloud.com/help/en/model-studio/models",
"supported_endpoints": [
"/v1/images/generations"
]
},
"databricks/databricks-bge-large-en": {
"input_cost_per_token": 1.0003e-07,
"input_dbu_cost_per_token": 1.429e-06,
@ -19226,6 +19273,42 @@
"supports_xhigh_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"gpt-5.5": {
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "openai",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 3e-05,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_none_reasoning_effort": true,
"supports_xhigh_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"gpt-5.4": {
"cache_read_input_token_cost": 2.5e-07,
"cache_read_input_token_cost_above_272k_tokens": 5e-07,
@ -22872,6 +22955,22 @@
"supports_video_input": true,
"supports_vision": true
},
"moonshot/kimi-k2.6": {
"cache_read_input_token_cost": 1.6e-07,
"input_cost_per_token": 9.5e-07,
"litellm_provider": "moonshot",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://platform.kimi.ai/docs/pricing/chat-k26",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"moonshot/kimi-latest": {
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 2e-06,
@ -25052,7 +25151,8 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
"tool_use_system_prompt_tokens": 159,
"supports_minimal_reasoning_effort": true
},
"openrouter/anthropic/claude-opus-4.5": {
"cache_creation_input_token_cost": 6.25e-06,
@ -25090,7 +25190,8 @@
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_minimal_reasoning_effort": true
},
"openrouter/anthropic/claude-sonnet-4.5": {
"input_cost_per_image": 0.0048,
@ -30118,7 +30219,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_minimal_reasoning_effort": true
},
"vercel_ai_gateway/anthropic/claude-sonnet-4": {
"cache_creation_input_token_cost": 3.75e-06,
@ -31345,7 +31447,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"vertex_ai/claude-opus-4-6@default": {
"cache_creation_input_token_cost": 6.25e-06,
@ -31372,7 +31475,8 @@
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 346,
"supports_max_reasoning_effort": true
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"vertex_ai/claude-opus-4-7": {
"cache_creation_input_token_cost": 6.25e-06,
@ -31399,7 +31503,9 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"vertex_ai/claude-opus-4-7@default": {
"cache_creation_input_token_cost": 6.25e-06,
@ -31426,7 +31532,9 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_xhigh_reasoning_effort": true,
"tool_use_system_prompt_tokens": 346
"tool_use_system_prompt_tokens": 346,
"supports_max_reasoning_effort": true,
"supports_minimal_reasoning_effort": true
},
"vertex_ai/claude-sonnet-4-5": {
"cache_creation_input_token_cost": 3.75e-06,
@ -31478,7 +31586,8 @@
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
}
},
"supports_minimal_reasoning_effort": true
},
"vertex_ai/claude-sonnet-4-5@20250929": {
"cache_creation_input_token_cost": 3.75e-06,
@ -38345,7 +38454,8 @@
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
}
},
"supports_minimal_reasoning_effort": true
},
"duckduckgo/search": {
"litellm_provider": "duckduckgo",

View file

@ -323,6 +323,14 @@ async def authorize_with_server(
)
parsed = urlparse(redirect_uri)
if parsed.scheme not in ("http", "https"):
raise HTTPException(
status_code=400,
detail={
"error": "invalid_redirect_uri",
"message": "redirect_uri must use http or https scheme",
},
)
base_url = urlunparse(parsed._replace(query=""))
request_base_url = get_request_base_url(request)
encoded_state = encode_state_with_base_url(

View file

@ -7,6 +7,7 @@ Filters MCP tools semantically for /chat/completions and /responses endpoints.
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR
if TYPE_CHECKING:
from semantic_router.routers import SemanticRouter
@ -214,20 +215,89 @@ class SemanticMCPToolFilter:
return []
@staticmethod
def _name_matches_canonical(client_name: str, canonical: str) -> bool:
"""
Return True if a client-side tool name refers to the given canonical
MCP tool name.
MCP clients (e.g. opencode) commonly wrap the proxy's canonical tool
name with an additive namespace prefix of their own
(``<client_alias><sep><canonical>``). The prefix can use either a
dash or an underscore as separator regardless of what
``MCP_TOOL_PREFIX_SEPARATOR`` is set to on the proxy, because the
client doesn't know the proxy's separator.
The match is anchored: ``canonical`` must form the complete suffix
of ``client_name`` and be preceded by a separator character, so
``rain_gear`` does not match canonical ``ear``.
Suffix matching is additionally gated on ``canonical`` itself
containing ``MCP_TOOL_PREFIX_SEPARATOR``. Server-registered MCP
tools are always emitted as
``<server_name><MCP_TOOL_PREFIX_SEPARATOR><tool_name>`` (see
``add_server_prefix_to_name``), so a canonical without the
separator is not a namespaced MCP tool and falling back to
suffix matching would spuriously collide with unrelated local
user functions whose names end in the same characters.
"""
if client_name == canonical:
return True
if MCP_TOOL_PREFIX_SEPARATOR not in canonical:
return False
if len(client_name) <= len(canonical):
return False
if not client_name.endswith(canonical):
return False
separator = client_name[-len(canonical) - 1]
return separator in ("_", "-")
def _get_tools_by_names(
self, tool_names: List[str], available_tools: List[Any]
) -> List[Any]:
"""Get tools from available_tools by their names, preserving order."""
# Match tools from available_tools (preserves format - dict or MCPTool)
matched_tools = []
for tool in available_tools:
tool_name, _ = self._extract_tool_info(tool)
if tool_name in tool_names:
matched_tools.append(tool)
"""
Get tools from available_tools by their names, preserving the
semantic router's ordering.
# Reorder to match semantic router's ordering
tool_map = {self._extract_tool_info(t)[0]: t for t in matched_tools}
return [tool_map[name] for name in tool_names if name in tool_map]
Matching is tolerant of client-side namespace prefixes: if an
incoming tool arrived as ``<client_alias>_<canonical>`` while the
router returned ``<canonical>`` (see
``_name_matches_canonical``), that tool is still selected. The
returned tool object is the original from ``available_tools``, so
the client-facing name is preserved for tool-call round-trips.
"""
# Build an index of incoming tools by their client-facing name.
# Exact matches win over suffix matches when both are present, and
# each incoming tool is returned at most once even if two canonical
# names happen to be tail-compatible with the same incoming name.
available_by_name: Dict[str, Any] = {}
for tool in available_tools:
client_name, _ = self._extract_tool_info(tool)
if client_name and client_name not in available_by_name:
available_by_name[client_name] = tool
matched: List[Any] = []
used_ids: set = set()
for canonical in tool_names:
tool = available_by_name.get(canonical)
if tool is None:
# Prefer the shortest qualifying name. When several
# incoming tools suffix-match the same canonical (e.g.
# "my_search" and "my_tag_search" both end in "search"),
# the one closest in length to the canonical is the
# least-wrapped and most likely the intended target.
best_name: Optional[str] = None
for client_name in available_by_name:
if not self._name_matches_canonical(client_name, canonical):
continue
if best_name is None or len(client_name) < len(best_name):
best_name = client_name
if best_name is not None:
tool = available_by_name[best_name]
if tool is not None and id(tool) not in used_ids:
matched.append(tool)
used_ids.add(id(tool))
return matched
def extract_user_query(self, messages: List[Dict[str, Any]]) -> str:
"""

View file

@ -1,42 +1,83 @@
# model_list:
# - model_name: claude-sonnet-4-6
# litellm_params: {model: anthropic/claude-sonnet-4-6}
# model_info:
# litellm_routing_preferences:
# quality_tier: 1
# keywords: [tin]
# - model_name: gpt-4o-mini
# litellm_params: {model: openai/gpt-4o-mini}
# model_info:
# litellm_routing_preferences:
# quality_tier: 1
# keywords: []
# - model_name: gpt-4o
# litellm_params: {model: openai/gpt-4o}
# model_info:
# litellm_routing_preferences:
# quality_tier: 2
# keywords: [vision, function_calling]
# - model_name: opus
# litellm_params: {model: anthropic/claude-opus-4-7}
# model_info:
# litellm_routing_preferences:
# quality_tier: 3
# keywords: ["architecture", "design"]
# - model_name: my-quality-router
# litellm_params:
# model: auto_router/adaptive_router
# adaptive_router_default_model: gpt-4o-mini
# adaptive_router_config:
# available_models: [gpt-4o-mini, gpt-4o, opus, claude-sonnet-4-6]
# Example proxy config for the adaptive router (v0).
#
# Wires one logical router ("smart-cheap-router") that adaptively picks between
# two real deployments ("fast" and "smart") based on per-session feedback signals.
#
# How to use from a client:
# POST /v1/chat/completions { "model": "smart-cheap-router", ... }
# Add { "metadata": { "litellm_session_id": "<your-session-id>" } } to enable
# sticky-session routing within a conversation.
#
# Required env vars: OPENAI_API_KEY, DATABASE_URL.
model_list:
# OpenAI model for /v1/chat/completions test — 200x custom pricing
- model_name: "gpt-4.1-mini"
# ---- The adaptive router "control" deployment -------------------------
# `model_name` is what clients call. `available_models` lists the underlying
# deployments the router is allowed to pick from (must match other model_name
# entries in this list).
- model_name: smart-cheap-router
litellm_params:
model: openai/gpt-4.1-mini
api_key: os.environ/OPENAI_API_KEY
model_info:
id: gpt-4.1-mini-custom-pricing
input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004)
output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016)
model: auto_router/adaptive_router
adaptive_router_config:
available_models: ["fast", "smart"]
weights:
quality: 0.7
cost: 0.3
# OpenAI model for /v1/responses test — 100x custom pricing
- model_name: "gpt-5"
litellm_params:
model: openai/gpt-5
api_key: os.environ/OPENAI_API_KEY
model_info:
id: gpt-5-custom-pricing
mode: "chat"
input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125)
output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001)
# Anthropic model for /v1/messages test — 100x custom pricing
- model_name: "claude-sonnet-4-6"
# ---- Underlying deployments the router picks from ---------------------
- model_name: fast
litellm_params:
model: anthropic/claude-sonnet-4-6
api_key: os.environ/ANTHROPIC_API_KEY
input_cost_per_token: 0.00000015
model_info:
id: claude-sonnet-4-custom-pricing
input_cost_per_token: 0.0003 # 100x standard ($0.000003)
output_cost_per_token: 0.0015 # 100x standard ($0.000015)
- model_name: my-auto
adaptive_router_preferences:
quality_tier: 2
strengths: []
- model_name: smart
litellm_params:
model: auto_router/complexity_router
complexity_router_config:
tiers:
SIMPLE: "gpt-4.1-mini"
COMPLEX: claude-sonnet-4-6
tier_boundaries:
simple_medium: 0.30
complexity_router_default_model: small-model
model: anthropic/claude-opus-4-7
api_key: os.environ/ANTHROPIC_API_KEY
input_cost_per_token: 0.0000050
model_info:
adaptive_router_preferences:
quality_tier: 3
strengths: ["code_generation", "technical_design", "analytical_reasoning"]
litellm_settings:
drop_params: True
general_settings:
master_key: sk-1234 # REPLACE in production

View file

@ -427,7 +427,8 @@ class LiteLLMRoutes(enum.Enum):
"/v1/skills/{skill_id}",
]
mcp_routes = [
# MCP tool-call / passthrough routes — data-plane. Gated by DISABLE_LLM_API_ENDPOINTS.
mcp_inference_routes = [
"/mcp",
"/mcp/",
"/mcp/{subpath}",
@ -436,10 +437,18 @@ class LiteLLMRoutes(enum.Enum):
"/mcp/tools/call",
"/mcp-rest/tools/list",
"/mcp-rest/tools/call",
]
# MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS.
mcp_management_routes = [
"/v1/mcp/server",
"/v1/mcp/server/{path:path}",
]
# Backwards-compat union — virtual keys may be configured with
# allowed_routes=["mcp_routes"], which should cover both halves.
mcp_routes = mcp_inference_routes + mcp_management_routes
agent_routes = [
"/v1/agents",
"/v1/agents/{agent_id}",
@ -477,7 +486,7 @@ class LiteLLMRoutes(enum.Enum):
+ mapped_pass_through_routes
+ passthrough_routes_wildcard
+ apply_guardrail_routes
+ mcp_routes
+ mcp_inference_routes
+ litellm_native_routes
+ agent_routes
)
@ -530,40 +539,44 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_ALIASES.value,
]
management_routes = [
# user
"/user/new",
"/user/update",
"/user/bulk_update",
"/user/delete",
"/user/info",
"/user/list",
"/user/daily/activity",
"/user/daily/activity/aggregated",
# team
"/team/new",
"/team/update",
"/team/delete",
"/team/list",
"/v2/team/list",
"/team/info",
"/team/block",
"/team/unblock",
"/team/available",
"/team/permissions_list",
"/team/permissions_update",
"/team/daily/activity",
# model
"/model/new",
"/model/update",
"/model/delete",
"/model/info",
"/jwt/key/mapping/new",
"/jwt/key/mapping/update",
"/jwt/key/mapping/delete",
"/jwt/key/mapping/list",
"/jwt/key/mapping/info",
] + key_management_routes
management_routes = (
[
# user
"/user/new",
"/user/update",
"/user/bulk_update",
"/user/delete",
"/user/info",
"/user/list",
"/user/daily/activity",
"/user/daily/activity/aggregated",
# team
"/team/new",
"/team/update",
"/team/delete",
"/team/list",
"/v2/team/list",
"/team/info",
"/team/block",
"/team/unblock",
"/team/available",
"/team/permissions_list",
"/team/permissions_update",
"/team/daily/activity",
# model
"/model/new",
"/model/update",
"/model/delete",
"/model/info",
"/jwt/key/mapping/new",
"/jwt/key/mapping/update",
"/jwt/key/mapping/delete",
"/jwt/key/mapping/list",
"/jwt/key/mapping/info",
]
+ key_management_routes
+ mcp_management_routes
)
spend_tracking_routes = [
# spend
@ -1997,7 +2010,12 @@ class TeamRequest(LiteLLMPydanticObjectBase):
class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
"""Represents user-controllable params for a LiteLLM_BudgetTable record"""
"""Represents user-controllable params for a LiteLLM_BudgetTable record.
Budget-write paths use `model_fields.keys()` on this class as an allowlist
for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
`LiteLLM_BudgetTableFull` so they aren't user-settable.
"""
budget_id: Optional[str] = None
soft_budget: Optional[float] = None
@ -2015,7 +2033,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
"""Represents all params for a LiteLLM_BudgetTable record"""
"""LiteLLM_BudgetTable + server-managed fields returned on API responses."""
budget_reset_at: Optional[datetime] = None
created_at: datetime
@ -3695,7 +3713,11 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
team_id: str
budget_id: Optional[str] = None
spend: Optional[float] = 0.0
litellm_budget_table: Optional[LiteLLM_BudgetTable]
total_spend: Optional[float] = 0.0
# Union so Pydantic picks Full when data has server-managed fields
# (/team/info) and Base when callers/tests construct with only
# user-settable fields.
litellm_budget_table: Optional[Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]]
def safe_get_team_member_rpm_limit(self) -> Optional[int]:
if self.litellm_budget_table is not None:
@ -3898,7 +3920,7 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse):
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
team_member_budget_table: Optional[LiteLLM_BudgetTable] = None
team_member_budget_table: Optional[LiteLLM_BudgetTableFull] = None
# Resources inherited from access groups (separate from direct assignments)
access_group_models: Optional[List[str]] = None
access_group_mcp_server_ids: Optional[List[str]] = None

View file

@ -626,11 +626,17 @@ async def common_checks( # noqa: PLR0915
and user_object.max_budget is not None
):
user_budget = user_object.max_budget
if user_budget < user_object.spend:
from litellm.proxy.proxy_server import get_current_spend
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
)
if user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_object.spend,
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
)
## 4.2 check team member budget, if team key
@ -899,6 +905,63 @@ async def get_default_end_user_budget(
return None
@log_db_metrics
async def get_team_member_default_budget(
budget_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
) -> Optional[LiteLLM_BudgetTable]:
"""
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
This budget is applied to team members whose TeamMembership row has no
linked budget. Results are cached for performance.
Args:
budget_id: The budget_id pulled from team.metadata["team_member_budget_id"]
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving budget data
Returns:
LiteLLM_BudgetTable if found, None otherwise
"""
if prisma_client is None:
return None
cache_key = f"team_member_default_budget:{budget_id}"
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
if isinstance(cached_budget, LiteLLM_BudgetTable):
return cached_budget
if isinstance(cached_budget, dict):
return LiteLLM_BudgetTable(**cached_budget)
try:
budget_record = await prisma_client.db.litellm_budgettable.find_unique(
where={"budget_id": budget_id}
)
if budget_record is None:
verbose_proxy_logger.warning(
f"Team-default member budget not found in database: {budget_id}"
)
return None
await user_api_key_cache.async_set_cache(
key=cache_key,
value=budget_record.dict(),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return LiteLLM_BudgetTable(**budget_record.dict())
except Exception:
verbose_proxy_logger.exception(
f"Error fetching team-default member budget {budget_id}"
)
return None
async def _apply_default_budget_to_end_user(
end_user_obj: LiteLLM_EndUserTable,
prisma_client: PrismaClient,
@ -3224,13 +3287,31 @@ async def _check_team_member_budget(
proxy_logging_obj=proxy_logging_obj,
)
# Per-member override wins; otherwise fall back to the team-level
# default configured via team.metadata["team_member_budget_id"].
team_member_budget: Optional[float] = None
if (
team_membership is not None
and team_membership.litellm_budget_table is not None
and team_membership.litellm_budget_table.max_budget is not None
):
team_member_budget = team_membership.litellm_budget_table.max_budget
team_member_spend = team_membership.spend or 0.0
else:
default_budget_id = (team_object.metadata or {}).get(
"team_member_budget_id"
)
if isinstance(default_budget_id, str):
default_budget = await get_team_member_default_budget(
budget_id=default_budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
if default_budget is not None:
team_member_budget = default_budget.max_budget
if team_member_budget is not None:
team_member_spend = (
team_membership.spend if team_membership is not None else 0.0
) or 0.0
# Read from cross-pod counter (Redis-first) if available
from litellm.proxy.proxy_server import get_current_spend
@ -3665,12 +3746,20 @@ async def _organization_max_budget_check(
if org_max_budget is None or org_max_budget <= 0:
return
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
from litellm.proxy.proxy_server import get_current_spend
org_spend = await get_current_spend(
counter_key=f"spend:org:{org_id}",
fallback_spend=org_table.spend or 0.0,
)
# Check if organization spend exceeds max budget
if org_table.spend >= org_max_budget:
if org_spend >= org_max_budget:
# Trigger budget alert
call_info = CallInfo(
token=valid_token.token,
spend=org_table.spend,
spend=org_spend,
max_budget=org_max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
@ -3686,9 +3775,9 @@ async def _organization_max_budget_check(
)
raise litellm.BudgetExceededError(
current_cost=org_table.spend,
current_cost=org_spend,
max_budget=org_max_budget,
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}",
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
)

View file

@ -151,7 +151,15 @@ def is_request_body_safe(
A malicious user can set the api_base to their own domain and invoke POST /chat/completions to intercept and steal the OpenAI API key.
Relevant issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997
"""
banned_params = ["api_base", "base_url", "user_config"]
banned_params = [
"api_base",
"base_url",
"user_config",
"aws_sts_endpoint",
"aws_web_identity_token",
"aws_role_name",
"vertex_credentials",
]
for param in banned_params:
if (

View file

@ -300,7 +300,7 @@ class RouteChecks:
return True
if RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value
):
return True
@ -358,7 +358,9 @@ class RouteChecks:
"""
Check if route is a management route
"""
return route in LiteLLMRoutes.management_routes.value
return RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.management_routes.value
)
@staticmethod
def is_info_route(route: str) -> bool:

View file

@ -433,7 +433,8 @@ def add_guardrail_to_applied_guardrails_header(
return
_metadata = request_data.get("metadata", None) or {}
if "applied_guardrails" in _metadata:
_metadata["applied_guardrails"].append(guardrail_name)
if guardrail_name not in _metadata["applied_guardrails"]:
_metadata["applied_guardrails"].append(guardrail_name)
else:
_metadata["applied_guardrails"] = [guardrail_name]
# Ensure metadata is set back to request_data (important when metadata didn't exist)

View file

@ -632,20 +632,27 @@ class ResetBudgetJob:
now = datetime.utcnow()
# Note on raw SQL: prisma-client-python does not support null-filtering
# on `Json?` columns (no DbNull/JsonNull sentinel — see
# RobertCraigie/prisma-client-py#714). We use `query_raw` with
# `IS NOT NULL` so we don't materialize every key/team row on each
# tick of the reset job. Writes still go through the ORM.
# --- Keys ---
try:
all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where={"budget_limits": {"not": None}} # type: ignore[arg-type]
key_rows = await self.prisma_client.db.query_raw(
'SELECT token, budget_limits FROM "LiteLLM_VerificationToken" '
"WHERE budget_limits IS NOT NULL"
)
for key in all_keys:
raw = key.budget_limits # type: ignore[attr-defined]
for row in key_rows:
raw = row["budget_limits"]
if not raw:
continue
windows: list = raw if isinstance(raw, list) else json.loads(raw)
changed = False
for window in windows:
counter_key = (
f"spend:key:{key.token}:window:{window['budget_duration']}"
f"spend:key:{row['token']}:window:{window['budget_duration']}"
)
if await ResetBudgetJob._reset_expired_window(
window, counter_key, spend_counter_cache, now
@ -653,7 +660,7 @@ class ResetBudgetJob:
changed = True
if changed:
await self.prisma_client.db.litellm_verificationtoken.update(
where={"token": key.token},
where={"token": row["token"]},
data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type]
)
except Exception as e:
@ -663,26 +670,25 @@ class ResetBudgetJob:
# --- Teams ---
try:
all_teams = await self.prisma_client.db.litellm_teamtable.find_many(
where={"budget_limits": {"not": None}} # type: ignore[arg-type]
team_rows = await self.prisma_client.db.query_raw(
'SELECT team_id, budget_limits FROM "LiteLLM_TeamTable" '
"WHERE budget_limits IS NOT NULL"
)
for team in all_teams:
raw = team.budget_limits # type: ignore[attr-defined]
for row in team_rows:
raw = row["budget_limits"]
if not raw:
continue
windows = raw if isinstance(raw, list) else json.loads(raw)
changed = False
for window in windows:
counter_key = (
f"spend:team:{team.team_id}:window:{window['budget_duration']}"
)
counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}"
if await ResetBudgetJob._reset_expired_window(
window, counter_key, spend_counter_cache, now
):
changed = True
if changed:
await self.prisma_client.db.litellm_teamtable.update(
where={"team_id": team.team_id},
where={"team_id": row["team_id"]},
data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type]
)
except Exception as e:

View file

@ -1300,7 +1300,10 @@ class DBSpendUpdateWriter:
batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id, "user_id": user_id},
data={"spend": {"increment": response_cost}},
data={
"spend": {"increment": response_cost},
"total_spend": {"increment": response_cost},
},
)
# Transaction succeeded, break out of retry loop
break

View file

@ -22,6 +22,7 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
@ -29,6 +30,8 @@ from litellm.proxy._types import (
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType,
SpendUpdateQueueItem,
)
from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
@ -259,9 +262,36 @@ class RedisUpdateBuffer:
if len(rpush_list) == 0:
return
result_lengths = await self.redis_cache.async_rpush_pipeline(
rpush_list=rpush_list,
)
try:
result_lengths = await self.redis_cache.async_rpush_pipeline(
rpush_list=rpush_list,
)
except Exception as e:
# The in-memory queues were already drained above. If we let the
# exception propagate without restoring, the aggregated spend is
# permanently lost. Re-enqueue so the next scheduler tick retries.
verbose_proxy_logger.error(
"Spend tracking - failed to push aggregated spend updates to Redis. "
"Restoring %d transaction sets to in-memory queues for retry on next tick. "
"Error: %s",
len(rpush_list),
str(e),
)
await self._restore_spend_updates_to_in_memory_queues(
db_spend_update_transactions=db_spend_update_transactions,
daily_spend_update_transactions=daily_spend_update_transactions,
daily_team_spend_update_transactions=daily_team_spend_update_transactions,
daily_org_spend_update_transactions=daily_org_spend_update_transactions,
daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions,
daily_agent_spend_update_transactions=daily_agent_spend_update_transactions,
spend_update_queue=spend_update_queue,
daily_spend_update_queue=daily_spend_update_queue,
daily_team_spend_update_queue=daily_team_spend_update_queue,
daily_org_spend_update_queue=daily_org_spend_update_queue,
daily_end_user_spend_update_queue=daily_end_user_spend_update_queue,
daily_agent_spend_update_queue=daily_agent_spend_update_queue,
)
return
# Emit gauge events for each queue
for i, queue_size in enumerate(result_lengths):
@ -271,6 +301,101 @@ class RedisUpdateBuffer:
service=service_types[i],
)
@staticmethod
async def _restore_spend_updates_to_in_memory_queues(
db_spend_update_transactions: Optional[DBSpendUpdateTransactions],
daily_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]],
daily_team_spend_update_transactions: Optional[
Dict[str, BaseDailySpendTransaction]
],
daily_org_spend_update_transactions: Optional[
Dict[str, BaseDailySpendTransaction]
],
daily_end_user_spend_update_transactions: Optional[
Dict[str, BaseDailySpendTransaction]
],
daily_agent_spend_update_transactions: Optional[
Dict[str, BaseDailySpendTransaction]
],
spend_update_queue: SpendUpdateQueue,
daily_spend_update_queue: DailySpendUpdateQueue,
daily_team_spend_update_queue: DailySpendUpdateQueue,
daily_org_spend_update_queue: DailySpendUpdateQueue,
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
daily_agent_spend_update_queue: DailySpendUpdateQueue,
) -> None:
"""
Put drained-but-unpushed transactions back into in-memory queues.
Called when the Redis rpush pipeline raises. Without this, all spend
data aggregated during the current scheduler tick is permanently lost
because the source queues were already drained before the rpush.
"""
if db_spend_update_transactions is not None:
entity_entries: List[
Tuple[Litellm_EntityType, Optional[Dict[str, float]]]
] = [
(
Litellm_EntityType.USER,
db_spend_update_transactions.get("user_list_transactions"),
),
(
Litellm_EntityType.END_USER,
db_spend_update_transactions.get("end_user_list_transactions"),
),
(
Litellm_EntityType.KEY,
db_spend_update_transactions.get("key_list_transactions"),
),
(
Litellm_EntityType.TEAM,
db_spend_update_transactions.get("team_list_transactions"),
),
(
Litellm_EntityType.TEAM_MEMBER,
db_spend_update_transactions.get("team_member_list_transactions"),
),
(
Litellm_EntityType.ORGANIZATION,
db_spend_update_transactions.get("org_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
),
(
Litellm_EntityType.AGENT,
db_spend_update_transactions.get("agent_list_transactions"),
),
]
for entity_type, entities in entity_entries:
if not entities:
continue
for entity_id, cost in entities.items():
await spend_update_queue.add_update(
SpendUpdateQueueItem(
entity_type=entity_type,
entity_id=entity_id,
response_cost=cost,
)
)
daily_pairs: List[
Tuple[Optional[Dict[str, BaseDailySpendTransaction]], DailySpendUpdateQueue]
] = [
(daily_spend_update_transactions, daily_spend_update_queue),
(daily_team_spend_update_transactions, daily_team_spend_update_queue),
(daily_org_spend_update_transactions, daily_org_spend_update_queue),
(
daily_end_user_spend_update_transactions,
daily_end_user_spend_update_queue,
),
(daily_agent_spend_update_transactions, daily_agent_spend_update_queue),
]
for daily_txns, daily_queue in daily_pairs:
if daily_txns:
await daily_queue.update_queue.put(daily_txns)
@staticmethod
def _number_of_transactions_to_store_in_redis(
db_spend_update_transactions: DBSpendUpdateTransactions,

View file

@ -403,10 +403,18 @@ class PrismaManager:
return dname
@staticmethod
def setup_database(use_migrate: bool = False) -> bool:
def setup_database(
use_migrate: bool = False, use_v2_resolver: bool = False
) -> bool:
"""
Set up the database using either prisma migrate or prisma db push
Args:
use_migrate: Use `prisma migrate deploy` instead of `db push`.
use_v2_resolver: Opt into the v2 migration resolver that avoids
the diff-and-force recovery behavior (which caused schema
thrashing during rolling deploys). Defaults to False.
Returns:
bool: True if setup was successful, False otherwise
"""
@ -427,7 +435,10 @@ class PrismaManager:
prisma_dir = PrismaManager._get_prisma_dir()
return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate)
return ProxyExtrasDBManager.setup_database(
use_migrate=use_migrate,
use_v2_resolver=use_v2_resolver,
)
else:
# Use prisma db push with increased timeout
subprocess.run(

View file

@ -0,0 +1,52 @@
# Example proxy config for the adaptive router (v0).
#
# Wires one logical router ("smart-cheap-router") that adaptively picks between
# two real deployments ("fast" and "smart") based on per-session feedback signals.
#
# How to use from a client:
# POST /v1/chat/completions { "model": "smart-cheap-router", ... }
# Add { "metadata": { "litellm_session_id": "<your-session-id>" } } to enable
# sticky-session routing within a conversation.
#
# Required env vars: OPENAI_API_KEY, DATABASE_URL.
model_list:
# ---- The adaptive router "control" deployment -------------------------
# `model_name` is what clients call. `available_models` lists the underlying
# deployments the router is allowed to pick from (must match other model_name
# entries in this list).
- model_name: smart-cheap-router
litellm_params:
model: auto_router/adaptive_router # required prefix -- triggers adaptive-router init
adaptive_router_config:
available_models: ["fast", "smart"]
weights:
quality: 0.7
cost: 0.3
# ---- Underlying deployments the router picks from ---------------------
- model_name: fast
litellm_params:
model: openai/gpt-4o-mini
api_key: os.environ/OPENAI_API_KEY
input_cost_per_token: 0.00000015
model_info:
adaptive_router_preferences:
quality_tier: 2
strengths: []
- model_name: smart
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
input_cost_per_token: 0.0000050
model_info:
adaptive_router_preferences:
quality_tier: 3
strengths: ["code_generation", "technical_design", "analytical_reasoning"]
litellm_settings:
drop_params: True
general_settings:
master_key: sk-1234 # REPLACE in production

View file

@ -5,7 +5,6 @@
# +-------------------------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
import copy
import os
import sys
@ -18,6 +17,7 @@ from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
ClassVar,
Dict,
List,
Literal,
@ -33,6 +33,7 @@ from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.caching import DualCache
from litellm.exceptions import GuardrailInterventionNormalStringError
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -62,6 +63,7 @@ from litellm.types.utils import (
CallTypesLiteral,
Choices,
GuardrailStatus,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
@ -78,56 +80,33 @@ class GuardrailMessageFilterResult(NamedTuple):
def _redact_pii_matches(response_json: dict) -> dict:
try:
# Create a deep copy to avoid modifying the original response
redacted_response = copy.deepcopy(response_json)
"""
Redact match-like fields from a Bedrock ApplyGuardrail JSON payload.
# Get assessments from the response
# NOTE: We use `.get("key") or []` instead of `.get("key", [])` because
# the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null).
# In Python, dict.get("key", []) returns None (not []) when the key exists
# with a None/null value. The `or []` ensures we always get an iterable,
# preventing "TypeError: 'NoneType' object is not iterable".
assessments = redacted_response.get("assessments") or []
if not assessments:
return redacted_response
Delegates to :func:`redact_nested_match_and_regex_keys` (same rules as spend
logging). Kept as a Bedrock-module entry point for existing unit tests.
"""
redacted = redact_nested_match_and_regex_keys(response_json)
return redacted if isinstance(redacted, dict) else response_json
for assessment in assessments:
# Redact PII entities in sensitive information policy
sensitive_info_policy = assessment.get("sensitiveInformationPolicy")
if sensitive_info_policy:
pii_entities = sensitive_info_policy.get("piiEntities") or []
for pii_entity in pii_entities:
if "match" in pii_entity:
pii_entity["match"] = "[REDACTED]"
# Redact regex matches
regexes = sensitive_info_policy.get("regexes") or []
for regex_match in regexes:
if "match" in regex_match:
regex_match["match"] = "[REDACTED]"
def _redact_assessment_match_fields(assessments: List[dict]) -> List[dict]:
"""
Redact sensitive match-like fields from blocked assessment summaries.
# Redact custom word matches in word policy
word_policy = assessment.get("wordPolicy")
if word_policy:
custom_words = word_policy.get("customWords") or []
for custom_word in custom_words:
if "match" in custom_word:
custom_word["match"] = "[REDACTED]"
managed_words = word_policy.get("managedWordLists") or []
for managed_word in managed_words:
if "match" in managed_word:
managed_word["match"] = "[REDACTED]"
return redacted_response
except Exception as e:
# We do not want to fail in any case so this is just a warning
verbose_proxy_logger.warning("Guardrail log redaction failed: %s", str(e))
return response_json
This is used for customer-visible error payloads (HTTPException.detail) where
we want to preserve policy/type/action metadata without echoing raw matched
content.
"""
redacted = redact_nested_match_and_regex_keys(assessments)
return redacted if isinstance(redacted, list) else assessments
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
use_native_during_call_hook: ClassVar[bool] = True
def __init__(
self,
guardrailIdentifier: Optional[str] = None,
@ -418,6 +397,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
messages: Optional[List[AllMessageValues]] = None,
response: Optional[Union[Any, litellm.ModelResponse]] = None,
request_data: Optional[dict] = None,
logging_event_type: Optional[GuardrailEventHooks] = None,
) -> BedrockGuardrailResponse:
from datetime import datetime
@ -455,11 +435,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prepared_request.headers,
)
event_type = (
GuardrailEventHooks.pre_call
if source == "INPUT"
else GuardrailEventHooks.post_call
)
# UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API
# body, which must not be confused with the proxy hook (pre_call / during_call /
# post_call). When omitted, keep legacy mapping for backward compatibility.
if logging_event_type is not None:
event_type = logging_event_type
else:
event_type = (
GuardrailEventHooks.pre_call
if source == "INPUT"
else GuardrailEventHooks.post_call
)
try:
httpx_response = await self.async_handler.post(
@ -514,9 +500,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
# Add guardrail information to request trace
#########################################################
_json_response = httpx_response.json()
# Raw Bedrock JSON is passed here; match/regex redaction runs once inside
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=httpx_response.json(),
guardrail_json_response=_json_response,
request_data=request_data or {},
guardrail_status=self._get_bedrock_guardrail_response_status(
response=httpx_response
@ -529,9 +518,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
if httpx_response.status_code == 200:
# check if the response was flagged
_json_response = httpx_response.json()
redacted_response = _redact_pii_matches(_json_response)
verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response)
verbose_proxy_logger.debug(
"Bedrock AI response : %s",
redact_nested_match_and_regex_keys(_json_response),
)
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
if self._should_raise_guardrail_blocked_exception(
bedrock_guardrail_response
@ -808,7 +798,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
assessments = self._extract_blocked_assessments(response)
if assessments:
detail["assessments"] = assessments
detail["assessments"] = _redact_assessment_match_fields(assessments)
return HTTPException(status_code=400, detail=detail)
@ -830,8 +820,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return False
# Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED)
# NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API.
# See _redact_pii_matches() for detailed explanation of the null safety pattern.
# NOTE: Use `.get("k") or []` not `.get("k", [])` — Bedrock can return explicit
# JSON null; dict.get("k", []) then yields None, and `for x in None` raises.
assessments = response.get("assessments") or []
if not assessments:
return False
@ -951,7 +941,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
source="INPUT",
messages=filtered_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.pre_call,
)
except GuardrailInterventionNormalStringError as e:
bedrock_guardrail_response = e.message
@ -1023,7 +1016,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
source="INPUT",
messages=filtered_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.during_call,
)
except GuardrailInterventionNormalStringError as e:
bedrock_guardrail_response = e.message
@ -1127,9 +1123,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
source="INPUT",
messages=input_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.post_call,
)
output_task = self.make_bedrock_api_request(
source="OUTPUT", response=response, request_data=data
source="OUTPUT",
response=response,
request_data=data,
logging_event_type=GuardrailEventHooks.post_call,
)
# Execute both requests in parallel
@ -1143,7 +1143,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# Only run OUTPUT validation (INPUT was already validated in pre_call or during_call)
try:
output_content_bedrock = await self.make_bedrock_api_request(
source="OUTPUT", response=response, request_data=data
source="OUTPUT",
response=response,
request_data=data,
logging_event_type=GuardrailEventHooks.post_call,
)
except GuardrailInterventionNormalStringError as e:
output_content_bedrock = e.message
@ -1270,9 +1273,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
source="INPUT",
messages=input_messages,
request_data=request_data,
logging_event_type=GuardrailEventHooks.post_call,
) # Only input messages
output_task = self.make_bedrock_api_request(
source="OUTPUT", response=assembled_model_response
source="OUTPUT",
response=assembled_model_response,
request_data=request_data,
logging_event_type=GuardrailEventHooks.post_call,
) # Only response
# Execute both requests in parallel
@ -1286,7 +1293,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# Only run OUTPUT validation (INPUT was already validated in pre_call or during_call)
try:
output_guardrail_response = await self.make_bedrock_api_request(
source="OUTPUT", response=assembled_model_response
source="OUTPUT",
response=assembled_model_response,
request_data=request_data,
logging_event_type=GuardrailEventHooks.post_call,
)
except GuardrailInterventionNormalStringError as e:
output_guardrail_response = e.message
@ -1563,11 +1573,50 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# Bedrock will throw an error if there is no text to process
if filtered_messages:
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=request_data,
_log_hook = (
GuardrailEventHooks.pre_call
if input_type == "request"
else GuardrailEventHooks.post_call
)
# Map the abstract input_type to the Bedrock source parameter.
# "request" -> INPUT (scan user-supplied content)
# "response" -> OUTPUT (scan model-generated content)
# Bedrock guardrail policies are often configured differently
# for Input vs Output (e.g. PII blocking only on Output), so
# the source MUST match where the text originated.
bedrock_source: Literal["INPUT", "OUTPUT"] = (
"OUTPUT" if input_type == "response" else "INPUT"
)
if bedrock_source == "OUTPUT":
# Build a synthetic ModelResponse whose choices carry the
# text(s) to scan, so _create_bedrock_output_content_request
# can produce the correct Bedrock OUTPUT payload.
synthetic_response = ModelResponse(
choices=[
Choices(
index=_idx,
message=Message(
role="assistant",
content=str(_msg.get("content") or ""),
),
finish_reason="stop",
)
for _idx, _msg in enumerate(filtered_messages)
]
)
bedrock_response = await self.make_bedrock_api_request(
source="OUTPUT",
response=synthetic_response,
request_data=request_data,
logging_event_type=_log_hook,
)
else:
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=request_data,
logging_event_type=_log_hook,
)
# Apply any masking that was applied by the guardrail
output_list = bedrock_response.get("output")

View file

@ -21,20 +21,30 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
):
try:
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
cache_key = f"{user_api_key_dict.user_id}_user_api_key_user_id"
user_row = await cache.async_get_cache(
cache_key, parent_otel_span=user_api_key_dict.parent_otel_span
max_budget = user_api_key_dict.user_max_budget
user_id = user_api_key_dict.user_id
if max_budget is None or user_id is None:
return
# Personal budget applies only to non-team requests, matching
# the explicit team-key exemption in common_checks section 4.1.
if user_api_key_dict.team_id is not None:
return
from litellm.proxy.proxy_server import get_current_spend
curr_spend = await get_current_spend(
counter_key=f"spend:user:{user_id}",
fallback_spend=user_api_key_dict.user_spend or 0.0,
)
if user_row is None: # value not yet cached
return
max_budget = user_row["max_budget"]
curr_spend = user_row["spend"]
if max_budget is None:
return
if curr_spend is None:
return
verbose_proxy_logger.debug(
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
user_id,
curr_spend,
max_budget,
)
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:

View file

@ -213,6 +213,7 @@ class _ProxyDBLogger(CustomLogger):
team_id=team_id,
user_id=user_id,
response_cost=response_cost,
org_id=org_id,
)
# update cache (fire-and-forget for backward compat:

View file

@ -285,6 +285,13 @@ async def image_edit_api(
if mask_files:
data["mask"] = mask_files
for _field in ("image", "mask"):
if _field in data and isinstance(data[_field], str):
raise HTTPException(
status_code=422,
detail=f"'{_field}' must be provided as a multipart file upload, not a string.",
)
# Ensure prompt exists in data (default to None for models that don't require it)
if "prompt" not in data:
data["prompt"] = None

View file

@ -355,6 +355,7 @@ async def _upsert_budget_and_membership(
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
allowed_models: Optional[List[str]] = None,
team_default_budget_id: Optional[str] = None,
):
"""
Helper function to Create/Update or Delete the budget within the team membership
@ -368,6 +369,11 @@ async def _upsert_budget_and_membership(
tpm_limit: Tokens per minute limit for the team member
rpm_limit: Requests per minute limit for the team member
allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce.
team_default_budget_id: The team's shared default member budget id (from
team metadata.team_member_budget_id), if any. When the membership's
existing_budget_id matches this, we clone-on-write so editing one
member's budget does not mutate the shared default (and therefore
every other member who still points at it).
If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership.
If any of these values exist, a budget is updated or created and linked to the team membership.
@ -385,7 +391,13 @@ async def _upsert_budget_and_membership(
)
return
if existing_budget_id is not None:
is_shared_default = (
existing_budget_id is not None
and team_default_budget_id is not None
and existing_budget_id == team_default_budget_id
)
if existing_budget_id is not None and not is_shared_default:
# Update the existing budget in-place to preserve fields not being changed.
# Only write fields that the caller explicitly provided (non-None).
update_data: Dict[str, Any] = {
@ -405,11 +417,40 @@ async def _upsert_budget_and_membership(
)
return
# No existing budget — create a new one and link it to the membership.
# Either there is no existing budget, OR the membership is still pointing
# at the team's shared default member budget. In both cases we create a
# NEW private budget for this user and (re)link the membership to it.
create_data: Dict[str, Any] = {
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
}
# If we're forking off the shared default, seed the new row with the
# default's values so fields the caller did not change carry over.
if is_shared_default:
default_budget_row = await tx.litellm_budgettable.find_unique(
where={"budget_id": existing_budget_id}
)
if default_budget_row is not None:
default_budget_dict = default_budget_row.model_dump()
for field in (
"max_budget",
"soft_budget",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
"model_max_budget",
"budget_duration",
"allowed_models",
):
value = default_budget_dict.get(field)
if value is None:
continue
if isinstance(value, list) and len(value) == 0:
continue
create_data[field] = value
# Caller-provided values take precedence over the cloned defaults.
if max_budget is not None:
create_data["max_budget"] = max_budget
if tpm_limit is not None:

View file

@ -52,12 +52,17 @@ from litellm.proxy._experimental.mcp_server.utils import (
from litellm.proxy._experimental.mcp_server.utils import (
validate_and_normalize_mcp_server_payload as _base_validate_and_normalize_mcp_server_payload,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
router = APIRouter(prefix="/v1/mcp", tags=["mcp"])
MCP_AVAILABLE: bool = True
TEMPORARY_MCP_SERVER_TTL_SECONDS = 300
TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX = "litellm:mcp:temporary_server"
def does_mcp_server_exist(
@ -329,13 +334,115 @@ if MCP_AVAILABLE:
)
return server
def get_cached_temporary_mcp_server(
async def _cache_temporary_mcp_server_in_redis(
server: MCPServer, ttl_seconds: int
) -> None:
"""
Best-effort write-through to Redis so temporary MCP OAuth sessions are
shared across proxy instances. Keep local in-memory cache as fallback.
"""
if litellm.cache is None or not hasattr(litellm.cache, "cache"):
return
cache_backend = getattr(litellm.cache, "cache", None)
if cache_backend is None or not hasattr(cache_backend, "async_set_cache"):
return
payload: Dict[str, Any] = server.model_dump(mode="json")
payload_json = json.dumps(payload)
try:
encrypted_payload = encrypt_value_helper(payload_json)
except Exception as e:
verbose_proxy_logger.debug(
f"Failed to encrypt temporary MCP server payload for Redis cache: {str(e)}"
)
return
if not isinstance(encrypted_payload, str):
verbose_proxy_logger.debug(
"Encrypted temporary MCP payload is not a string; skipping Redis cache write"
)
return
try:
await cache_backend.async_set_cache(
key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server.server_id}",
value=encrypted_payload,
ttl=max(1, ttl_seconds),
)
except Exception as e:
verbose_proxy_logger.debug(
f"Failed to write temporary MCP server to Redis cache: {str(e)}"
)
async def _get_temporary_mcp_server_from_redis(
server_id: str,
) -> Optional[MCPServer]:
"""
Best-effort read from Redis shared cache. Returns None on miss/errors.
Values must be encrypted strings (same contract as _cache_temporary_mcp_server_in_redis);
legacy plaintext dict payloads are rejected.
"""
if litellm.cache is None or not hasattr(litellm.cache, "cache"):
return None
cache_backend = getattr(litellm.cache, "cache", None)
if cache_backend is None or not hasattr(cache_backend, "async_get_cache"):
return None
try:
cached_server = await cache_backend.async_get_cache(
key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server_id}"
)
except Exception as e:
verbose_proxy_logger.debug(
f"Failed reading temporary MCP server from Redis cache: {str(e)}"
)
return None
if not isinstance(cached_server, str):
verbose_proxy_logger.debug(
"Temporary MCP Redis cache value must be an encrypted string; rejecting non-string payload"
)
return None
decrypted_json = decrypt_value_helper(
value=cached_server,
key="temporary_mcp_server",
exception_type="debug",
)
if decrypted_json is None:
return None
try:
loaded = json.loads(decrypted_json)
except Exception as e:
verbose_proxy_logger.debug(
f"Invalid decrypted temporary MCP payload in Redis cache: {str(e)}"
)
return None
if not isinstance(loaded, dict):
return None
payload_dict: Dict[str, Any] = loaded
try:
return MCPServer(**payload_dict)
except Exception as e:
verbose_proxy_logger.debug(
f"Invalid temporary MCP server payload in Redis cache: {str(e)}"
)
return None
async def get_cached_temporary_mcp_server(
server_id: str,
) -> Optional[MCPServer]:
_prune_expired_temporary_mcp_servers()
entry = _temporary_mcp_servers.get(server_id)
if entry is None:
return None
redis_server = await _get_temporary_mcp_server_from_redis(server_id)
if redis_server is None:
return None
# Intentionally avoid repopulating local cache from Redis to prevent
# extending effective lifetime beyond the remaining Redis TTL.
return redis_server
return entry.server
def _redact_mcp_credentials(
@ -1325,6 +1432,10 @@ if MCP_AVAILABLE:
temporary_server,
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
)
await _cache_temporary_mcp_server_in_redis(
temporary_server,
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
)
except Exception as e:
verbose_proxy_logger.exception(
f"Error caching temporary mcp server: {str(e)}"
@ -1336,18 +1447,24 @@ if MCP_AVAILABLE:
return _redact_mcp_credentials(temp_record)
def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer:
server = get_cached_temporary_mcp_server(server_id)
async def _get_cached_temporary_mcp_server_or_404(
server_id: str, request: Optional[Request] = None
) -> MCPServer:
server = await get_cached_temporary_mcp_server(server_id)
if server is None:
# Fall back to real DB/config server (e.g. for the user-side OAuth flow
# which calls these endpoints with a real server_id, not a temp session id).
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
client_ip = IPAddressUtils.get_mcp_client_ip(request) if request else None
server = global_mcp_server_manager.get_mcp_server_by_id(
server_id
) or global_mcp_server_manager.get_mcp_server_by_name(server_id)
) or global_mcp_server_manager.get_mcp_server_by_name(
server_id, client_ip=client_ip
)
if server is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
@ -1358,10 +1475,12 @@ if MCP_AVAILABLE:
@router.get(
"/server/oauth/{server_id}/authorize",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
)
async def mcp_authorize(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
client_id: Optional[str] = None,
redirect_uri: str = Query(...),
state: str = "",
@ -1370,7 +1489,9 @@ if MCP_AVAILABLE:
response_type: Optional[str] = None,
scope: Optional[str] = None,
):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
mcp_server = await _get_cached_temporary_mcp_server_or_404(
server_id, request=request
)
# Use the server's stored client_id when the caller doesn't supply one
resolved_client_id = mcp_server.client_id or client_id or ""
if not resolved_client_id:
@ -1399,10 +1520,12 @@ if MCP_AVAILABLE:
@router.post(
"/server/oauth/{server_id}/token",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
)
async def mcp_token(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
grant_type: str = Form(...),
code: Optional[str] = Form(None),
redirect_uri: Optional[str] = Form(None),
@ -1412,7 +1535,9 @@ if MCP_AVAILABLE:
refresh_token: Optional[str] = Form(None),
scope: Optional[str] = Form(None),
):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
mcp_server = await _get_cached_temporary_mcp_server_or_404(
server_id, request=request
)
resolved_client_id = mcp_server.client_id or client_id or ""
if not resolved_client_id:
raise HTTPException(
@ -1441,9 +1566,16 @@ if MCP_AVAILABLE:
@router.post(
"/server/oauth/{server_id}/register",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
)
async def mcp_register(request: Request, server_id: str):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
async def mcp_register(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
mcp_server = await _get_cached_temporary_mcp_server_or_404(
server_id, request=request
)
request_data = await _read_request_body(request=request)
data: dict = {**request_data}

View file

@ -302,14 +302,15 @@ class TeamMemberBudgetHandler:
prisma_client: PrismaClient,
) -> None:
"""
Create team_memberships entries for existing members that don't have one.
Ensure every team member has a TeamMembership row linked to the
team_member_budget.
Called after team_member_budget is set/updated on a team to ensure
members who joined before the budget was configured also get budget
enforcement.
Only creates missing entries does not touch existing memberships
(which may carry individual per-member budgets).
Called after team_member_budget is set/updated on a team. Creates
rows for members who don't have one, and populates budget_id on
existing rows where it is NULL. Rows with a non-NULL budget_id
are left untouched, which preserves per-member overrides but also
means rows pointing to a prior team-default budget_id are not
migrated to the new one.
"""
if not members_with_roles:
return
@ -347,6 +348,21 @@ class TeamMemberBudgetHandler:
_sanitize_for_log(team_member_budget_id),
)
# Heal existing membership rows that predate the team_member_budget
# configuration: populate budget_id where it is currently NULL.
# Rows with an explicit budget_id (per-member override) are left alone.
updated = await prisma_client.db.litellm_teammembership.update_many(
where={"team_id": team_id, "budget_id": None},
data={"budget_id": team_member_budget_id},
)
if updated:
verbose_proxy_logger.info(
"Populated budget_id on %d existing team_memberships for team %s with budget %s",
updated,
_sanitize_for_log(team_id),
_sanitize_for_log(team_member_budget_id),
)
def _get_default_team_param(field: str) -> Any:
"""
@ -2608,6 +2624,15 @@ async def team_member_update(
identified_budget_id = tm.budget_id
break
# If this membership still points at the team's shared default member
# budget, _upsert_budget_and_membership will clone-on-write so that the
# update only touches this user (not every member sharing the default).
team_default_budget_id: Optional[str] = None
if team_table.metadata is not None:
raw_default_budget_id = team_table.metadata.get("team_member_budget_id")
if isinstance(raw_default_budget_id, str):
team_default_budget_id = raw_default_budget_id
### upsert new budget
async with prisma_client.db.tx() as tx:
await _upsert_budget_and_membership(
@ -2620,6 +2645,7 @@ async def team_member_update(
tpm_limit=data.tpm_limit,
rpm_limit=data.rpm_limit,
allowed_models=data.allowed_models,
team_default_budget_id=team_default_budget_id,
)
### update team member role

View file

@ -9,6 +9,7 @@ from fastapi import HTTPException, Request
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy._types import ( # key request types; user request types; team request types; customer request types
BudgetNewRequest,
DeleteCustomerRequest,
@ -140,6 +141,69 @@ async def handle_budget_for_entity(
return existing_budget_id
# Fields on LiteLLM_BudgetTable that represent the budget's *configuration*
# (i.e. the values an admin sets). We copy these when cloning a team's
# default member-budget into an individual member-budget so that the new
# row starts with the same limits as the default.
_CLONABLE_BUDGET_FIELDS: Tuple[str, ...] = (
"max_budget",
"soft_budget",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
"model_max_budget",
"budget_duration",
"allowed_models",
)
async def _clone_team_default_budget_for_member(
prisma_client: PrismaClient,
default_team_budget_id: str,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
) -> Optional[str]:
"""
Create a new budget row that copies the values from the team's default
member budget. Returns the new budget_id, or None if the default budget
no longer exists in the DB.
Used when adding a new team member without an explicit per-member budget,
so the member starts with the team default's values but gets their own
private budget row (which can be edited independently).
"""
default_budget = await prisma_client.db.litellm_budgettable.find_unique(
where={"budget_id": default_team_budget_id}
)
if default_budget is None:
return None
default_budget_dict = default_budget.model_dump()
cloned_data: dict = {
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}
for field in _CLONABLE_BUDGET_FIELDS:
value = default_budget_dict.get(field)
if value is None:
continue
# Skip empty list defaults (e.g. allowed_models = []) so the cloned
# row matches the "no value set" shape rather than carrying a default.
if isinstance(value, list) and len(value) == 0:
continue
cloned_data[field] = value
# Start the member's budget window at clone time, not the pool's reset
# timestamp — otherwise a member joining mid-cycle inherits a stale reset.
if cloned_data.get("budget_duration"):
cloned_data["budget_reset_at"] = get_budget_reset_time(
cloned_data["budget_duration"]
)
new_budget = await prisma_client.db.litellm_budgettable.create(data=cloned_data)
return new_budget.budget_id
async def add_new_member(
new_member: Member,
max_budget_in_team: Optional[float],
@ -221,8 +285,20 @@ async def add_new_member(
response = await prisma_client.db.litellm_budgettable.create(data=budget_data)
_budget_id = response.budget_id
elif default_team_budget_id is not None:
# No per-member budget was provided, but the team has a default member
# budget. Clone the default budget into a new row for this user so that
# later edits to one member's budget do not bleed into other members.
# If the default no longer exists in the DB, fall back to no budget.
_budget_id = await _clone_team_default_budget_for_member(
prisma_client=prisma_client,
default_team_budget_id=default_team_budget_id,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
else:
_budget_id = default_team_budget_id
# No per-member budget and no team default → member gets no budget.
_budget_id = None
if _budget_id and returned_user is not None and returned_user.user_id is not None:
_returned_team_membership = (

View file

@ -8,6 +8,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
import json
import os
import re
from typing import Any, Optional, Tuple, Union, cast
import httpx
@ -1496,10 +1497,18 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
def get_vertex_base_url(vertex_location: Optional[str]) -> str:
"""
Returns the base URL for Vertex AI based on the provided location.
Base URL for Vertex AI pass-through (trailing slash for URL joining).
Keep location rules aligned with ``litellm.llms.vertex_ai.common_utils.get_vertex_base_url``.
"""
if vertex_location == "global":
return "https://aiplatform.googleapis.com/"
if vertex_location is None:
raise ValueError("vertex_location is required")
if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
raise ValueError("Invalid vertex_location format")
if "-" not in vertex_location:
return f"https://aiplatform.{vertex_location}.rep.googleapis.com/"
return f"https://{vertex_location}-aiplatform.googleapis.com/"
@ -1703,7 +1712,8 @@ async def _base_vertex_proxy_route(
Base function for Vertex AI passthrough routes.
Handles common logic for all Vertex AI services.
Default base_target_url is `https://{vertex_location}-aiplatform.googleapis.com/`
Default base_target_url is derived from ``get_vertex_base_url`` in this module
(regional, ``global``, or multi-region ``.rep.`` hosts), with a trailing slash.
Args:
endpoint: The endpoint path
@ -2275,11 +2285,7 @@ async def vertex_ai_live_websocket_passthrough(
return
host_location = resolved_location or vertex_llm_base.get_default_vertex_location()
host = (
"aiplatform.googleapis.com"
if host_location == "global"
else f"{host_location}-aiplatform.googleapis.com"
)
host = get_vertex_base_url(host_location).removeprefix("https://").rstrip("/")
service_url = (
f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
)

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