scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #26283 from BerriAI/litellm_internal_staging

Browse source 
Sync litellm_staging_03_22_2026 with litellm_internal_staging
This commit is contained in:
Cesar Garcia 2026-04-22 19:55:27 -03:00 • committed by GitHub
commit 25c0aa8bfd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
178 changed files with 18032 additions and 2036 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1193
.circleci/config.yml
View file

File diff suppressed because it is too large Load diff

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

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

@ -31,8 +31,15 @@ jobs:
test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py"
workers: 8
timeout: 20
# test_proxy_utils.py is large (168+ parametrized tests) — run it on its
# own matrix so --dist=loadscope doesn't pin all of it to a single xdist
# worker and push the "remaining" group past the job timeout.
- test-group: proxy-utils
test-path: "tests/proxy_unit_tests/test_proxy_utils.py"
workers: 8
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"
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 --ignore=tests/proxy_unit_tests/test_proxy_utils.py"
workers: 8
timeout: 30
uses: ./.github/workflows/_test-unit-services-base.yml

13
Dockerfile
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

299
ci_cd/run_migration.py
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,
)

13
docker/Dockerfile.database
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

117
docker/Dockerfile.non_root
View file

@ -15,29 +15,21 @@ COPY --from=uvbin /uv /usr/local/bin/uv
COPY --from=uvbin /uvx /usr/local/bin/uvx
RUN for i in 1 2 3; do \
apk add --no-cache \
python3 \
python3-dev \
clang \
llvm \
lld \
gcc \
linux-headers \
build-base \
bash \
coreutils \
curl \
openssl \
openssl-dev \
nodejs \
npm \
libsndfile && break || sleep 5; \
apk add --no-cache \
python3 \
python3-dev \
gcc \
bash \
coreutils \
curl \
openssl \
libsndfile \
nodejs && break || sleep 5; \
done
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
NVM_DIR=/root/.nvm \
PATH="/root/.nvm/versions/node/v20.20.2/bin:/app/.venv/bin:${PATH}" \
PATH="/app/.venv/bin:${PATH}" \
LITELLM_NON_ROOT=true \
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
PRISMA_CLI_BINARY_TARGETS="debian-openssl-3.0.x" \
@ -49,7 +41,8 @@ COPY enterprise/pyproject.toml enterprise/
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
@ -62,38 +55,12 @@ COPY . .
# Set non-root flag for build time consistency
ENV LITELLM_NON_ROOT=true
# Build Admin UI once and stage the static output for the runtime image.
# NOTE: .npmrc files (which may set ignore-scripts=true and min-release-age=3d)
# are temporarily renamed during npm install/ci so they don't block lifecycle
# scripts needed by the build. This is safe because npm ci installs from
# package-lock.json with pinned versions + integrity hashes.
# Stage the pre-built Admin UI from the checked-in Next.js static export.
# _experimental/out/ is regenerated as part of the release runbook.
# Restructure extensionless routes (foo.html -> foo/index.html) to match the layout
# proxy_server.py expects, and drop a readiness marker.
RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
([ -f /app/.npmrc ] && mv /app/.npmrc /app/.npmrc.bak || true) && \
NVM_VERSION="v0.40.4" && \
NVM_CHECKSUM="4b7412c49960c7d31e8df72da90c1fb5b8cccb419ac99537b737028d497aba4f" && \
NODE_VERSION="v20.20.2" && \
NVM_SCRIPT="/tmp/install-nvm.sh" && \
curl -fsSL "https://raw.githubusercontent.com/nvm-sh/nvm/${NVM_VERSION}/install.sh" -o "$NVM_SCRIPT" && \
echo "${NVM_CHECKSUM} ${NVM_SCRIPT}" | sha256sum -c - && \
bash "$NVM_SCRIPT" && \
export NVM_DIR="$HOME/.nvm" && \
. "$NVM_DIR/nvm.sh" && \
nvm install "${NODE_VERSION}" && \
nvm use "${NODE_VERSION}" && \
npm install -g npm@11.12.1 && \
npm install -g node-gyp@12.2.0 && \
ln -sf "$(npm root -g)/node-gyp" "$(npm root -g)/npm/node_modules/node-gyp" && \
npm cache clean --force && \
cd /app/ui/litellm-dashboard && \
if [ -f "/app/enterprise/enterprise_ui/enterprise_colors.json" ]; then \
cp /app/enterprise/enterprise_ui/enterprise_colors.json ./ui_colors.json; \
fi && \
([ -f .npmrc ] && mv .npmrc .npmrc.bak || true) && \
npm ci --no-audit --no-fund && \
([ -f .npmrc.bak ] && mv .npmrc.bak .npmrc || true) && \
([ -f /app/.npmrc.bak ] && mv /app/.npmrc.bak /app/.npmrc || true) && \
npm run build && \
cp -r /app/ui/litellm-dashboard/out/* /var/lib/litellm/ui/ && \
cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \
cp /app/litellm/proxy/logo.jpg /var/lib/litellm/assets/logo.jpg && \
( cd /var/lib/litellm/ui && \
for html_file in *.html; do \
@ -103,10 +70,10 @@ RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
mv "$html_file" "$folder_name/index.html"; \
fi; \
done && \
touch .litellm_ui_ready ) && \
cd /app/ui/litellm-dashboard && rm -rf ./out
touch .litellm_ui_ready )
RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
@ -123,10 +90,7 @@ RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
--python python3; \
fi
RUN mkdir -p /app/.cache/npm && \
prisma generate --schema=./schema.prisma && \
prisma --version && \
prisma migrate diff --from-empty --to-schema-datamodel ./schema.prisma --script > /dev/null 2>&1 || true
RUN prisma generate --schema=./schema.prisma
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
@ -137,33 +101,11 @@ WORKDIR /app
USER root
RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python3 bash openssl tzdata nodejs npm supervisor libsndfile && break || sleep 5; \
done && \
apk upgrade --no-cache nodejs && \
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
GLOBAL="$(npm root -g)" && \
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
done && \
find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \
sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \
npm cache clean --force && \
{ apk del --no-cache npm 2>/dev/null || true; }
apk add --no-cache python3 bash openssl tzdata supervisor libsndfile nodejs && break || sleep 5; \
done
COPY --from=builder /app /app
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
@ -179,15 +121,10 @@ ENV PATH="/app/.venv/bin:${PATH}" \
PRISMA_SKIP_POSTINSTALL_GENERATE=1 \
PRISMA_HIDE_UPDATE_MESSAGE=1 \
PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \
NPM_CONFIG_CACHE=/app/.cache/npm \
NPM_CONFIG_PREFER_OFFLINE=true \
PRISMA_OFFLINE_MODE=true
RUN sed -i 's/\r$//' docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && \
chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh && \
mkdir -p /nonexistent /.npm /var/lib/litellm/assets /var/lib/litellm/ui /tmp/.npm && \
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent /.npm /tmp/.npm && \
RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
chown -R nobody:nogroup "$PRISMA_PATH" && \
LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \

155
docs/my-website/docs/adaptive_router.md Normal file
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

1
docs/my-website/docs/completion/prompt_caching.md
View file

@ -10,6 +10,7 @@ Supported Providers:
- Vertex AI (`vertex_ai/`, `vertex_ai_beta/`)
- Bedrock (`bedrock/`, `bedrock/invoke/`, `bedrock/converse`) ([All models bedrock supports prompt caching on](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html))
- Deepseek API (`deepseek/`)
- xAI (`xai/`)
For the supported providers, LiteLLM follows the OpenAI prompt caching usage object format:

25
docs/my-website/docs/completion/prompt_compression.md
View file

@ -8,6 +8,7 @@ The function keeps high-relevance and recent context, replaces low-relevance con
```python
import litellm
from litellm.types.utils import CallTypes
messages = [
{"role": "system", "content": "You are a coding assistant."},
@ -19,6 +20,7 @@ messages = [
compressed = litellm.compress(
messages=messages,
model="gpt-4o",
call_type=CallTypes.completion,
compression_trigger=1000,
compression_target=500,
)
@ -45,6 +47,7 @@ response = litellm.completion(
- `messages` (`List[dict]`, required): input conversation messages
- `model` (`str`, required): model name used for token counting
- `call_type` (`CallTypes`, default `CallTypes.completion`): the LiteLLM call type whose message schema these messages follow. Supported values: `CallTypes.completion` / `CallTypes.acompletion` (OpenAI chat-completions shape) and `CallTypes.anthropic_messages` (Anthropic Messages shape)
- `compression_trigger` (`int`, default `200000`): compress only if input token count exceeds this
- `compression_target` (`Optional[int]`, default `70% of compression_trigger`): desired post-compression token budget
- `embedding_model` (`Optional[str]`): if set, combines BM25 + embedding relevance scoring
@ -70,6 +73,28 @@ args = json.loads(tool_call.function.arguments)
full_content = compressed["cache"][args["key"]]
```
## Server-side Callback Loop (`/v1/messages`)
You can enable callback-based compression interception to make retrieval loops
transparent for Anthropic Messages calls:
```yaml
litellm_settings:
callbacks: ["compression_interception"]
compression_interception_params:
enabled: true
compression_trigger: 10000
compression_target: 7000
```
With this enabled, LiteLLM runs the following server-side flow:
1. Compresses inbound messages before the first provider call.
2. Injects the `litellm_content_retrieve` tool.
3. Detects retrieval `tool_use` blocks in the model response.
4. Resolves retrieval keys from the compression cache.
5. Reruns the model via agentic loop and returns the final answer.
## Performance
Benchmarked on [SWE-bench Lite](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_bm25_27K) (real GitHub issues with ~27k tokens of BM25-retrieved repo context per problem).

41
docs/my-website/docs/providers/scaleway.md
View file

@ -60,3 +60,44 @@ curl http://localhost:4000/chat/completions \
## Supported features
Scaleway provider supports all features in [Generative APIs reference documentation ↗](https://www.scaleway.com/en/developers/api/generative-apis/), such as streaming, structured outputs and tool calling.
## Audio transcription
Scaleway's `/audio/transcriptions` endpoint is OpenAI-compatible and works with Whisper models.
### Python SDK
```python
import os
from litellm import transcription
os.environ["SCW_SECRET_KEY"] = "your-scaleway-secret-key"
with open("speech.mp3", "rb") as audio_file:
response = transcription(
model="scaleway/whisper-large-v3",
file=audio_file,
)
print(response.text)
```
### Proxy config
```yaml
model_list:
- model_name: scaleway-whisper
litellm_params:
model: scaleway/whisper-large-v3
api_key: "os.environ/SCW_SECRET_KEY"
```
### Proxy request
```bash
curl http://localhost:4000/v1/audio/transcriptions \
-H "Authorization: Bearer YOUR_LITELLM_MASTER_KEY" \
-F model="scaleway-whisper" \
-F file="@speech.mp3"
```
Supported optional params: `language`, `prompt`, `response_format`, `temperature`, `timestamp_granularities`.

95
docs/my-website/docs/proxy/agentic_loop_hook.md Normal file
View file

@ -0,0 +1,95 @@
# Agentic Loop Hook
Build a `CustomLogger` callback that intercepts a model response, fulfills tool calls server-side, and reruns the model — transparently to the caller.
:::info Supported call types
- `async` only (sync calls do not trigger the hook)
- Non-streaming only (streaming responses cannot be inspected for tool calls)
- Works on both `/v1/messages` and `/v1/chat/completions`
:::
## Implement the callback
Override two methods on `CustomLogger`:
```python
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
MY_TOOL = "my_tool"
class MyToolCallback(CustomLogger):
async def async_should_run_agentic_loop(
self, response, model, messages, tools, stream, custom_llm_provider, kwargs
):
# Return (True, context_dict) if there are tool calls to handle
content = getattr(response, "content", None) or []
calls = [b for b in content if isinstance(b, dict)
and b.get("type") == "tool_use" and b.get("name") == MY_TOOL]
if not calls:
return False, {}
return True, {"tool_calls": calls}
async def async_build_agentic_loop_plan(
self, tools, model, messages, response,
anthropic_messages_provider_config,
anthropic_messages_optional_request_params,
logging_obj, stream, kwargs,
):
calls = tools["tool_calls"]
results = [f"result for {c['input']}" for c in calls] # your logic here
follow_up = messages + [
{"role": "assistant", "content": [
{"type": "tool_use", "id": c["id"], "name": c["name"], "input": c["input"]}
for c in calls
]},
{"role": "user", "content": [
{"type": "tool_result", "tool_use_id": c["id"], "content": results[i]}
for i, c in enumerate(calls)
]},
]
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=AgenticLoopRequestPatch(messages=follow_up),
)
```
For `/v1/chat/completions`, override `async_build_chat_completion_agentic_loop_plan` instead — same idea, `optional_params` replaces `anthropic_messages_optional_request_params`.
## Register it
```python
import litellm
litellm.callbacks = [MyToolCallback()]
```
Or in `config.yaml`:
```yaml
litellm_settings:
callbacks: ["my_module.MyToolCallback"]
```
## `AgenticLoopPlan` fields
| Field | Effect |
|---|---|
| `run_agentic_loop=True` + `request_patch` | Reruns the model with the patched request |
| `response_override` | Returns this value directly to the caller (no rerun) |
| `terminate=True` | Stops the loop, returns the current response |
| `run_agentic_loop=False` (default) | Skips; next callback is checked |
`AgenticLoopRequestPatch` accepts: `model`, `messages`, `tools`, `max_tokens`, `optional_params`, `kwargs`.
## Loop safety
- Default max reruns: `3` — override per-request with `kwargs["max_agentic_loops"]`
- Identical tool-call fingerprints abort the loop automatically
- Current depth is in `kwargs["_agentic_loop_depth"]`
## Examples in this repo
- `litellm/integrations/compression_interception/handler.py`
- `litellm/integrations/websearch_interception/handler.py`

16
docs/my-website/docs/tutorials/prompt_caching.md
View file

@ -8,6 +8,22 @@ Reduce costs by up to 90% by using LiteLLM to auto-inject prompt caching checkpo
<Image img={require('../../img/auto_prompt_caching.png')} style={{ width: '800px', height: 'auto' }} />
Supported Providers (`cache_control` marker):
- Anthropic API (`anthropic/`)
- AWS Bedrock - Claude (`bedrock/`)
- Vertex AI - Claude and Gemini (`vertex_ai/`)
- Google AI Studio - Gemini (`gemini/`)
- Azure AI - Claude (`azure_ai/`)
- OpenRouter - Claude, Gemini, MiniMax, GLM, z-ai routes (`openrouter/`)
- Databricks - Claude (`databricks/`)
- DashScope / Qwen (`dashscope/`)
- MiniMax (`minimax/`)
- Z.ai / GLM (`zai/`)
Provider Managed (automatic, no marker needed):
- OpenAI (`openai/`)
- DeepSeek (`deepseek/`)
- xAI (`xai/`)
## How it works

1
docs/my-website/package-lock.json generated
View file

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

1
docs/my-website/package.json
View file

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

2
docs/my-website/sidebars.js
View file

@ -536,6 +536,7 @@ const sidebars = {
description: "Modify requests, responses, and more",
items: [
"proxy/call_hooks",
"proxy/agentic_loop_hook",
"proxy/rules",
]
},
@ -1059,6 +1060,7 @@ const sidebars = {
},
items: [
"routing",
"adaptive_router",
"scheduler",
"proxy/auto_routing",
"proxy/load_balancing",

39
litellm-proxy-extras/litellm_proxy_extras/migrations/20260418000000_add_adaptive_router_tables/migration.sql Normal file
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);

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

@ -1223,3 +1223,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")
}

309
litellm-proxy-extras/litellm_proxy_extras/utils.py
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()

50
litellm-proxy-extras/migration_runbook.md
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

242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py Normal file
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"

4
litellm/__init__.py
View file

@ -148,6 +148,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"vantage",
"posthog",
"levo",
"compression_interception",
]
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
@ -1501,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 (

5
litellm/_lazy_imports_registry.py
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": (

412
litellm/compression/compress.py
View file

@ -1,9 +1,9 @@
"""
Main compress() function — orchestrates BM25/embedding scoring, message stubbing,
and retrieval tool injection.
Main compress() function — normalizes input messages, orchestrates BM25/embedding
scoring, message stubbing, and retrieval tool injection.
"""
from typing import Any, Dict, List, Optional, Set, Union, cast
from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
from litellm.caching.dual_cache import DualCache
from litellm.compression.message_stubbing import (
@ -15,27 +15,196 @@ from litellm.compression.retrieval_tool import build_retrieval_tool
from litellm.compression.scoring.bm25 import bm25_score_messages
from litellm.litellm_core_utils.token_counter import token_counter
from litellm.types.compression import CompressedResult
from litellm.types.utils import AllMessageValues, Message
from litellm.types.utils import CallTypes
# CallTypes that produce Anthropic-shaped messages (structured content blocks).
# Everything else is treated as OpenAI chat-completions shape.
_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value})
# CallTypes that are valid targets for compression. Compression operates on
# message-shaped inputs, so we only accept call types whose payload is a list
# of role/content messages.
_SUPPORTED_CALL_TYPES = frozenset(
{
CallTypes.completion.value,
CallTypes.acompletion.value,
CallTypes.anthropic_messages.value,
}
)
def _normalize_call_type(call_type: Union[CallTypes, str]) -> str:
"""Return the string value for a ``CallTypes`` enum or a raw string."""
if isinstance(call_type, CallTypes):
return call_type.value
return call_type
def _is_anthropic_call_type(call_type: str) -> bool:
return call_type in _ANTHROPIC_CALL_TYPES
def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]:
"""
Build retrieval tool definitions in the target request schema.
- Chat-completions call types: keep OpenAI function-tool schema.
- Anthropic messages call type: remap to Anthropic's custom tool schema.
"""
if not keys:
return []
openai_tools = [build_retrieval_tool(keys)]
if not _is_anthropic_call_type(call_type):
return openai_tools
# Lazy import to avoid introducing provider transformation imports during
# module import for non-Anthropic call paths.
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools)
return cast(List[dict], anthropic_tools)
def _content_to_text(content: Any) -> str:
"""
Convert OpenAI/Anthropic message content blocks to plain text.
Text extraction policy:
- Include text-bearing fields only (`text` blocks + string values).
- For `tool_result`, expand into nested `content` items.
- Ignore non-textual blocks (images/documents/tool metadata/thinking metadata).
Implemented iteratively (stack-based) to avoid unbounded recursion.
"""
parts: List[str] = []
stack: List[Any] = [content]
while stack:
item = stack.pop()
if isinstance(item, str):
parts.append(item)
elif isinstance(item, list):
# Push list items in reverse order so they are processed left-to-right.
for element in reversed(item):
stack.append(element)
elif isinstance(item, dict):
item_type = item.get("type")
if item_type == "text":
parts.append(str(item.get("text", "")))
elif item_type == "tool_result":
stack.append(item.get("content", ""))
return " ".join(parts)
def _normalize_messages_for_compression(
messages: List[dict],
call_type: str,
) -> Tuple[List[dict], List[dict]]:
"""
Normalize each original message to a text-surrogate content for scoring.
Returns:
(normalized_messages, original_messages_copy)
"""
if call_type not in _SUPPORTED_CALL_TYPES:
raise ValueError(
f"Unsupported call_type={call_type!r} for compression. "
f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}."
)
original_messages: List[Dict[str, Any]] = [dict(m) for m in messages]
normalized_messages: List[dict] = []
for msg in original_messages:
normalized_messages.append(
{
**msg,
"content": _content_to_text(msg.get("content", "")),
}
)
return normalized_messages, original_messages
def _extract_last_user_message(messages: List[dict]) -> str:
"""Return the text content of the last user message."""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
elif isinstance(part, str):
parts.append(part)
return " ".join(parts)
return _content_to_text(msg.get("content", ""))
return ""
def _extract_tool_use_ids(content: Any) -> List[str]:
if not isinstance(content, list):
return []
tool_use_ids: List[str] = []
for part in content:
if not isinstance(part, dict):
continue
if part.get("type") != "tool_use":
continue
tool_use_id = part.get("id")
if isinstance(tool_use_id, str) and tool_use_id:
tool_use_ids.append(tool_use_id)
return tool_use_ids
def _extract_tool_result_ids(content: Any) -> Set[str]:
if not isinstance(content, list):
return set()
tool_result_ids: Set[str] = set()
for part in content:
if not isinstance(part, dict):
continue
if part.get("type") != "tool_result":
continue
tool_use_id = part.get("tool_use_id")
if isinstance(tool_use_id, str) and tool_use_id:
tool_result_ids.add(tool_use_id)
return tool_result_ids
def _extract_anthropic_tool_exchange_spans(
messages: List[dict],
) -> Tuple[List[Set[int]], Optional[str]]:
"""
Return atomic 2-message spans for Anthropic tool exchanges.
Each assistant message containing `tool_use` must be immediately followed by a
user message containing matching `tool_result` blocks for all tool_use ids.
"""
spans: List[Set[int]] = []
i = 0
while i < len(messages):
current = messages[i]
if current.get("role") != "assistant":
i += 1
continue
tool_use_ids = _extract_tool_use_ids(current.get("content"))
if not tool_use_ids:
i += 1
continue
if i + 1 >= len(messages):
return [], "invalid_anthropic_tool_sequence"
next_msg = messages[i + 1]
if next_msg.get("role") != "user":
return [], "invalid_anthropic_tool_sequence"
tool_result_ids = _extract_tool_result_ids(next_msg.get("content"))
if not tool_result_ids:
return [], "invalid_anthropic_tool_sequence"
for tool_use_id in tool_use_ids:
if tool_use_id not in tool_result_ids:
return [], "invalid_anthropic_tool_sequence"
spans.append({i, i + 1})
i += 2
return spans, None
def _get_protected_indices(messages: List[dict]) -> List[int]:
"""
Return indices of messages that must never be compressed:
@ -87,9 +256,98 @@ def _combine_scores(
return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)]
def _select_kept_indices_for_budget(
normalized_messages: List[dict],
original_messages: List[dict],
combined_scores: List[float],
compression_target: int,
model: str,
initial_kept_indices: Set[int],
tool_exchange_spans: List[Set[int]],
) -> Tuple[Set[int], Dict[int, dict]]:
kept_indices = set(initial_kept_indices)
current_tokens = 0
for i in kept_indices:
current_tokens += token_counter(
model=model,
text=cast(str, normalized_messages[i].get("content", "") or ""),
)
# Fill token budget from highest-scoring units.
# A unit is either:
# 1) a single message index, or
# 2) an Anthropic tool-exchange span that must be kept/dropped atomically.
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
span_id_by_index: Dict[int, int] = {}
for span_id, span in enumerate(tool_exchange_spans):
for idx in span:
span_id_by_index[idx] = span_id
# Build single-message candidate units (non-span messages).
candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = []
for idx in range(len(normalized_messages)):
if idx in span_id_by_index or idx in kept_indices:
continue
candidate_units.append((combined_scores[idx], (idx,), True))
# Build span candidate units (atomic keep/drop for tool exchanges).
for span in tool_exchange_spans:
span_indices = tuple(sorted(span))
if any(idx in kept_indices for idx in span_indices):
continue
span_score = max(combined_scores[idx] for idx in span_indices)
candidate_units.append((span_score, span_indices, False))
# Sort by descending relevance score.
candidate_units.sort(key=lambda item: item[0], reverse=True)
for _score, indices, can_truncate in candidate_units:
if any(idx in kept_indices for idx in indices):
continue
msg_tokens = 0
for idx in indices:
msg_tokens += token_counter(
model=model,
text=cast(str, normalized_messages[idx].get("content", "") or ""),
)
remaining = compression_target - current_tokens
if remaining <= 0:
break # budget exhausted
if current_tokens + msg_tokens <= compression_target:
# Fits entirely
kept_indices.update(indices)
current_tokens += msg_tokens
elif can_truncate and len(indices) == 1 and remaining >= 100:
# Too large to fit whole single message, but we have budget — truncate it.
idx = indices[0]
truncated = truncate_message(original_messages[idx], remaining)
truncated_tokens = token_counter(
model=model,
text=truncated.get("content", "") or "",
)
truncated_overrides[idx] = truncated
kept_indices.add(idx)
current_tokens += truncated_tokens
return kept_indices, truncated_overrides
def _get_dropped_tool_span_indices(
kept_indices: Set[int], tool_exchange_spans: List[Set[int]]
) -> Set[int]:
dropped_tool_span_indices: Set[int] = set()
for span in tool_exchange_spans:
if not any(idx in kept_indices for idx in span):
dropped_tool_span_indices.update(span)
return dropped_tool_span_indices
def compress(
messages: List[dict],
model: str,
call_type: Union[CallTypes, str] = CallTypes.completion,
compression_trigger: int = 200_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
@ -108,6 +366,12 @@ def compress(
Parameters:
messages: The conversation messages to (potentially) compress.
model: The LLM model name — used for token counting.
call_type: The LiteLLM call type whose message schema these messages
follow. Supported values:
- ``CallTypes.completion`` / ``CallTypes.acompletion`` — OpenAI
chat-completions shape (default)
- ``CallTypes.anthropic_messages`` — Anthropic Messages shape
(structured content blocks + atomic tool exchanges)
compression_trigger: Only compress if input exceeds this token count.
compression_target: Target token count after compression.
Defaults to ``compression_trigger // 2``.
@ -122,29 +386,37 @@ def compress(
A ``CompressedResult`` dict containing compressed messages, token
counts, a cache of original content, and the retrieval tool definition.
"""
call_type_str = _normalize_call_type(call_type)
normalized_messages, original_messages = _normalize_messages_for_compression(
messages=messages,
call_type=call_type_str,
)
if compression_target is None:
compression_target = compression_trigger * 7 // 10
original_tokens = token_counter(
model=model, messages=cast(List[Union[AllMessageValues, Message]], messages)
model=model,
messages=cast(List[Any], original_messages),
)
# Pass through if below trigger
if original_tokens <= compression_trigger:
return CompressedResult(
messages=messages,
messages=original_messages,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
compression_ratio=0.0,
cache={},
tools=[],
compression_skipped_reason="below_trigger",
)
# Extract query for relevance scoring
query = _extract_last_user_message(messages)
query = _extract_last_user_message(normalized_messages)
# Score each message
bm25_scores = bm25_score_messages(query, messages)
bm25_scores = bm25_score_messages(query, normalized_messages)
if embedding_model:
from litellm.compression.scoring.embedding_scorer import (
@ -153,7 +425,7 @@ def compress(
emb_scores = embedding_score_messages(
query,
messages,
normalized_messages,
model=embedding_model,
cache=compression_cache,
embedding_model_params=embedding_model_params,
@ -162,85 +434,69 @@ def compress(
else:
combined_scores = bm25_scores
# Sort message indices by score descending
ranked_indices = sorted(
range(len(messages)),
key=lambda i: combined_scores[i],
reverse=True,
)
# Protected messages are never compressed
protected_indices = _get_protected_indices(messages)
protected_indices = _get_protected_indices(normalized_messages)
kept_indices: Set[int] = set(protected_indices)
# Count tokens for protected messages
current_tokens = 0
for i in kept_indices:
current_tokens += token_counter(
model=model, text=messages[i].get("content", "") or ""
tool_exchange_spans: List[Set[int]] = []
if _is_anthropic_call_type(call_type_str):
tool_exchange_spans, tool_sequence_error = (
_extract_anthropic_tool_exchange_spans(original_messages)
)
# Fill token budget from highest-scoring messages.
# For each candidate (ranked by relevance):
# - If it fits entirely → keep it as-is.
# - If it doesn't fit but there's meaningful remaining budget → truncate it
# to fill as much of the budget as possible.
# - Otherwise → stub it (pointer only, content goes to cache).
# Multiple messages may be truncated so we preserve partial content from
# several high-scoring messages rather than fully stubbing all but one.
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
for idx in ranked_indices:
if idx in kept_indices:
continue
msg_content = messages[idx].get("content", "") or ""
msg_tokens = token_counter(model=model, text=msg_content)
remaining = compression_target - current_tokens
if remaining <= 0:
break # budget exhausted
if current_tokens + msg_tokens <= compression_target:
# Fits entirely
kept_indices.add(idx)
current_tokens += msg_tokens
elif remaining >= 100:
# Too large to fit whole, but we have budget — truncate it.
truncated = truncate_message(messages[idx], remaining)
truncated_tokens = token_counter(
model=model,
text=truncated.get("content", "") or "",
if tool_sequence_error is not None:
return CompressedResult(
messages=original_messages,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
compression_ratio=0.0,
cache={},
tools=[],
compression_skipped_reason=tool_sequence_error,
)
truncated_overrides[idx] = truncated
kept_indices.add(idx)
current_tokens += truncated_tokens
for span in tool_exchange_spans:
# If any message in the span is protected, keep the whole span.
if any(idx in kept_indices for idx in span):
kept_indices.update(span)
kept_indices, truncated_overrides = _select_kept_indices_for_budget(
normalized_messages=normalized_messages,
original_messages=original_messages,
combined_scores=combined_scores,
compression_target=compression_target,
model=model,
initial_kept_indices=kept_indices,
tool_exchange_spans=tool_exchange_spans,
)
# Build compressed messages and cache
compressed_messages: List[dict] = []
cache: Dict[str, str] = {}
used_keys: Set[str] = set()
dropped_tool_span_indices = _get_dropped_tool_span_indices(
kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans
)
for i, msg in enumerate(messages):
for i, msg in enumerate(original_messages):
if i in dropped_tool_span_indices:
continue
if i in kept_indices:
# Use the truncated version if we made one, otherwise the original
compressed_messages.append(truncated_overrides.get(i, msg))
else:
key = extract_key(msg, fallback_index=i, used_keys=used_keys)
content = msg.get("content", "")
if isinstance(content, list):
content = " ".join(
p.get("text", "") if isinstance(p, dict) else str(p)
for p in content
)
key = extract_key(
normalized_messages[i], fallback_index=i, used_keys=used_keys
)
content = _content_to_text(msg.get("content", ""))
cache[key] = content
compressed_messages.append(stub_message(msg, key))
# Build retrieval tool
tools = [build_retrieval_tool(list(cache.keys()))] if cache else []
# Build retrieval tool in the target request schema
tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str)
compressed_tokens = token_counter(
model=model,
messages=cast(List[Union[AllMessageValues, Message]], compressed_messages),
messages=cast(List[Any], compressed_messages),
)
return CompressedResult(

1
litellm/constants.py
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

14
litellm/integrations/compression_interception/__init__.py Normal file
View file

@ -0,0 +1,14 @@
"""
Compression Interception Module
Provides server-side prompt compression + retrieval tool fulfillment for
Anthropic Messages agentic loops.
"""
from litellm.integrations.compression_interception.handler import (
CompressionInterceptionLogger,
)
__all__ = [
"CompressionInterceptionLogger",
]

399
litellm/integrations/compression_interception/handler.py Normal file
View file

@ -0,0 +1,399 @@
"""
Compression Interception Handler
CustomLogger that compresses inbound Anthropic Messages requests and fulfills
litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
"""
import time
import uuid
from typing import Any, Dict, List, Optional, Tuple, cast
from litellm._logging import verbose_logger
from litellm.compression import compress
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.integrations.compression_interception import (
CompressionInterceptionConfig,
)
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.utils import CallTypes
LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve"
_CACHE_TTL_SECONDS = 15 * 60
class CompressionInterceptionLogger(CustomLogger):
"""
CustomLogger that implements transparent prompt compression + retrieval loops.
Flow:
1. Compress inbound /v1/messages requests in pre-call hook.
2. Inject litellm_content_retrieve tool and persist compressed cache by call_id.
3. Detect retrieval tool_use blocks in first model response.
4. Build typed rerun plan with tool_result blocks from the compressed cache.
"""
def __init__(
self,
enabled: bool = True,
compression_trigger: int = 200_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
embedding_model_params: Optional[Dict[str, Any]] = None,
):
super().__init__()
self.enabled = enabled
self.compression_trigger = compression_trigger
self.compression_target = compression_target
self.embedding_model = embedding_model
self.embedding_model_params = embedding_model_params
self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {}
@classmethod
def from_config_yaml(
cls, config: CompressionInterceptionConfig
) -> "CompressionInterceptionLogger":
return cls(
enabled=bool(config.get("enabled", True)),
compression_trigger=int(config.get("compression_trigger", 200_000)),
compression_target=config.get("compression_target"),
embedding_model=config.get("embedding_model"),
embedding_model_params=config.get("embedding_model_params"),
)
@staticmethod
def initialize_from_proxy_config(
litellm_settings: Dict[str, Any],
callback_specific_params: Dict[str, Any],
) -> "CompressionInterceptionLogger":
compression_params: CompressionInterceptionConfig = {}
if "compression_interception_params" in litellm_settings:
compression_params = litellm_settings["compression_interception_params"]
elif "compression_interception" in callback_specific_params:
compression_params = callback_specific_params["compression_interception"]
return CompressionInterceptionLogger.from_config_yaml(compression_params)
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
) -> Optional[dict]:
if not self.enabled:
return None
if call_type is not None and call_type != CallTypes.anthropic_messages:
return None
if int(kwargs.get("_agentic_loop_depth", 0) or 0) > 0:
return None
messages = kwargs.get("messages")
model = kwargs.get("model")
if not isinstance(messages, list) or not isinstance(model, str):
return None
if self._has_retrieval_tool(kwargs.get("tools")):
return None
self._prune_expired_cache()
compressed = compress( # type: ignore
messages=messages,
model=model,
call_type=CallTypes.anthropic_messages,
compression_trigger=self.compression_trigger,
compression_target=self.compression_target,
embedding_model=self.embedding_model,
embedding_model_params=self.embedding_model_params,
)
cache = cast(Dict[str, str], compressed.get("cache", {}))
skip_reason = cast(Optional[str], compressed.get("compression_skipped_reason"))
compressed_tools = cast(List[Dict[str, Any]], compressed.get("tools", []))
# Only mutate kwargs when compression actually produced a result.
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
# leave ``messages`` and ``tools`` untouched — injecting an empty
# ``tools: []`` onto a request that originally had no tools breaks
# Anthropic Messages requests.
if cache:
kwargs["messages"] = compressed["messages"]
if compressed_tools:
kwargs["tools"] = self._merge_tools(
existing_tools=cast(
Optional[List[Dict[str, Any]]], kwargs.get("tools")
),
compressed_tools=compressed_tools,
)
call_id = cast(Optional[str], kwargs.get("litellm_call_id"))
if not call_id:
call_id = str(uuid.uuid4())
kwargs["litellm_call_id"] = call_id
self._compression_cache_by_call_id[call_id] = (cache, time.time())
verbose_logger.debug(
"CompressionInterception: compressed request [call_id=%s original=%d compressed=%d cached_keys=%d]",
call_id,
compressed.get("original_tokens"),
compressed.get("compressed_tokens"),
len(cache),
)
elif skip_reason is not None:
verbose_logger.debug(
"CompressionInterception: compression skipped [reason=%s original=%d compressed=%d]",
skip_reason,
compressed.get("original_tokens"),
compressed.get("compressed_tokens"),
)
return kwargs
async def async_should_run_agentic_loop(
self,
response: Any,
model: str,
messages: List[Dict],
tools: Optional[List[Dict]],
stream: bool,
custom_llm_provider: str,
kwargs: Dict,
) -> Tuple[bool, Dict]:
if not self.enabled:
return False, {}
if not self._has_retrieval_tool(tools):
return False, {}
tool_calls, thinking_blocks = self._extract_retrieval_tool_calls(
response=response
)
if not tool_calls:
return False, {}
return True, {
"tool_calls": tool_calls,
"thinking_blocks": thinking_blocks,
"tool_type": "compression_retrieval",
}
async def async_build_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
self._prune_expired_cache()
tool_calls = cast(List[Dict[str, Any]], tools.get("tool_calls", []))
thinking_blocks = cast(List[Dict[str, Any]], tools.get("thinking_blocks", []))
call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
cache = self._get_cache(call_id=call_id)
retrieval_results = [
self._resolve_retrieval_content(tc, cache) for tc in tool_calls
]
assistant_message = {
"role": "assistant",
"content": thinking_blocks
+ [
{
"type": "tool_use",
"id": tc.get("id"),
"name": tc.get("name", LITELLM_CONTENT_RETRIEVE_TOOL_NAME),
"input": tc.get("input", {}),
}
for tc in tool_calls
],
}
user_message = {
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_calls[i].get("id"),
"content": retrieval_results[i],
}
for i in range(len(tool_calls))
],
}
follow_up_messages = messages + [assistant_message, user_message]
max_tokens = cast(
Optional[int],
anthropic_messages_optional_request_params.get("max_tokens")
or kwargs.get("max_tokens"),
)
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
full_model_name = model
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = cast(str, agentic_params.get("model", model))
request_patch = AgenticLoopRequestPatch(
model=full_model_name,
messages=follow_up_messages,
max_tokens=max_tokens,
optional_params=optional_params_without_max_tokens,
kwargs=self._prepare_followup_kwargs(kwargs=kwargs),
)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata={"tool_type": "compression_retrieval", "call_id": call_id or ""},
)
def _prune_expired_cache(self) -> None:
now = time.time()
self._compression_cache_by_call_id = {
call_id: (cache, created_at)
for call_id, (
cache,
created_at,
) in self._compression_cache_by_call_id.items()
if now - created_at <= _CACHE_TTL_SECONDS
}
def _get_cache(self, call_id: Optional[str]) -> Dict[str, str]:
if not call_id:
return {}
cache_entry = self._compression_cache_by_call_id.get(call_id)
if cache_entry is None:
return {}
return cache_entry[0]
def _resolve_call_id(
self, logging_obj: Any, kwargs: Dict[str, Any]
) -> Optional[str]:
if logging_obj is not None:
logging_call_id = getattr(logging_obj, "litellm_call_id", None)
if isinstance(logging_call_id, str) and logging_call_id:
return logging_call_id
kwargs_call_id = kwargs.get("litellm_call_id")
return cast(
Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None
)
def _resolve_retrieval_content(
self, tool_call: Dict[str, Any], cache: Dict[str, str]
) -> str:
raw_input = tool_call.get("input", {})
key = ""
if isinstance(raw_input, dict):
key = str(raw_input.get("key", "") or "")
if not key:
return "No retrieval key provided."
if key in cache:
return cache[key]
return f"[compressed content key '{key}' not found]"
def _extract_retrieval_tool_calls(
self, response: Any
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
if isinstance(response, dict):
content = response.get("content", [])
else:
content = getattr(response, "content", []) or []
if not isinstance(content, list):
return [], []
tool_calls: List[Dict[str, Any]] = []
thinking_blocks: List[Dict[str, Any]] = []
for block in content:
if isinstance(block, dict):
block_type = block.get("type")
block_name = block.get("name")
if block_type in ("thinking", "redacted_thinking"):
thinking_blocks.append(block)
if (
block_type == "tool_use"
and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
):
tool_calls.append(
{
"id": block.get("id"),
"type": "tool_use",
"name": block_name,
"input": block.get("input", {}),
}
)
else:
block_type = getattr(block, "type", None)
block_name = getattr(block, "name", None)
if block_type == "thinking":
thinking_blocks.append(
{
"type": "thinking",
"thinking": getattr(block, "thinking", ""),
"signature": getattr(block, "signature", ""),
}
)
elif block_type == "redacted_thinking":
thinking_blocks.append(
{
"type": "redacted_thinking",
"data": getattr(block, "data", ""),
}
)
if (
block_type == "tool_use"
and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
):
tool_calls.append(
{
"id": getattr(block, "id", None),
"type": "tool_use",
"name": block_name,
"input": getattr(block, "input", {}) or {},
}
)
return tool_calls, thinking_blocks
def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
internal_keys = {"litellm_logging_obj"}
return {
k: v
for k, v in kwargs.items()
if not k.startswith("_compression_interception") and k not in internal_keys
}
def _has_retrieval_tool(self, tools: Any) -> bool:
if not isinstance(tools, list):
return False
for tool in tools:
if not isinstance(tool, dict):
continue
function = tool.get("function")
if tool.get("type") == "function" and isinstance(function, dict):
if function.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME:
return True
if (
tool.get("type") == "custom"
and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
):
return True
return False
def _merge_tools(
self,
existing_tools: Optional[List[Dict[str, Any]]],
compressed_tools: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
merged = list(existing_tools or [])
if self._has_retrieval_tool(merged):
return merged
merged.extend(compressed_tools)
return merged

12
litellm/integrations/custom_guardrail.py
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,

39
litellm/integrations/custom_logger.py
View file

@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
from litellm.types.integrations.argilla import ArgillaItem
from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.integrations.custom_logger import AgenticLoopPlan
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
CallTypes,
@ -239,7 +240,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional[PreRoutingHookResponse]:
@ -676,6 +677,26 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
pass
async def async_build_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
"""
Build a typed rerun plan for Anthropic Messages agentic loops.
Override this method to separate callback decision/tool execution from
follow-up request execution (handled by BaseLLMHTTPHandler).
"""
return AgenticLoopPlan(run_agentic_loop=False)
async def async_should_run_chat_completion_agentic_loop(
self,
response: Any,
@ -707,6 +728,22 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
pass
async def async_build_chat_completion_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
optional_params: Dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
"""
Build a typed rerun plan for chat-completions agentic loops.
"""
return AgenticLoopPlan(run_agentic_loop=False)
# Useful helpers for custom logger classes
def truncate_standard_logging_payload_content(

12
litellm/integrations/opentelemetry.py
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):

13
litellm/integrations/prometheus.py
View file

@ -51,6 +51,7 @@ if TYPE_CHECKING:
else:
AsyncIOScheduler = Any
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger):
amount: float = 1.0,
) -> None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name=metric_name
),
supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name),
enum_values=enum_values,
label_context=label_context,
)
@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger):
user_api_key = hash_token(user_api_key)
label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request.
label_context = PrometheusLabelFactoryContext(
enum_values
) # amortized per request.
# increment total LLM requests and spend metric
self._increment_top_level_request_and_spend_metrics(
@ -3490,7 +3491,9 @@ def _prometheus_labels_from_context(
}
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user()
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = (
ctx.get_resolved_end_user()
)
for sk, val in ctx._custom_by_sanitized_key.items():
if sk in supported_enum_labels:

3
litellm/integrations/prometheus_helpers/__init__.py
View file

@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext:
self.enum_values = enum_values
enum_dict = enum_values.model_dump()
self._sanitized_enum: Dict[str, Optional[str]] = {
k: _sanitize_prometheus_label_value(v)
for k, v in enum_dict.items()
k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items()
}
self._custom_by_sanitized_key: Dict[str, Optional[str]] = {}
if enum_values.custom_metadata_labels is not None:

336
litellm/integrations/websearch_interception/handler.py
View file

@ -28,6 +28,10 @@ from litellm.integrations.websearch_interception.transformation import (
from litellm.types.integrations.websearch_interception import (
WebSearchInterceptionConfig,
)
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
@ -573,6 +577,35 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs=kwargs,
)
async def async_build_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
tool_calls = tools["tool_calls"]
thinking_blocks = tools.get("thinking_blocks", [])
request_patch = await self._build_anthropic_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
thinking_blocks=thinking_blocks,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs,
)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata={"tool_type": "websearch", "response_format": "anthropic"},
)
async def async_run_chat_completion_agentic_loop(
self,
tools: Dict,
@ -608,6 +641,33 @@ class WebSearchInterceptionLogger(CustomLogger):
response_format=response_format,
)
async def async_build_chat_completion_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
optional_params: Dict,
logging_obj: Any,
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
tool_calls = tools["tool_calls"]
response_format = tools.get("response_format", "openai")
request_patch = await self._build_chat_completion_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
optional_params=optional_params,
kwargs=kwargs,
response_format=response_format,
)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata={"tool_type": "websearch", "response_format": response_format},
)
@staticmethod
def _resolve_max_tokens(
optional_params: Dict,
@ -672,7 +732,48 @@ class WebSearchInterceptionLogger(CustomLogger):
stream: bool,
kwargs: Dict,
) -> Any:
"""Execute litellm.search() and make follow-up request"""
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch = await self._build_anthropic_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
thinking_blocks=thinking_blocks,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs,
)
if request_patch.messages is None:
raise ValueError("WebSearchInterception: missing follow-up messages")
optional_params = dict(anthropic_messages_optional_request_params)
optional_params.update(request_patch.optional_params)
max_tokens = request_patch.max_tokens
if max_tokens is None:
max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None))
else:
optional_params.pop("max_tokens", None)
if max_tokens is None:
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
return await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=request_patch.messages,
model=request_patch.model or model,
**optional_params,
**request_patch.kwargs,
)
async def _build_anthropic_request_patch(
self,
model: str,
messages: List[Dict],
tool_calls: List[Dict],
thinking_blocks: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
kwargs: Dict,
) -> AgenticLoopRequestPatch:
"""Execute litellm.search() and build follow-up request patch."""
# Extract search queries from tool_use blocks
search_tasks = []
@ -721,20 +822,8 @@ class WebSearchInterceptionLogger(CustomLogger):
thinking_blocks=thinking_blocks,
)
# Make follow-up request with search results
# Type cast: user_message is a Dict for Anthropic format (default response_format)
follow_up_messages = messages + [assistant_message, cast(Dict, user_message)]
verbose_logger.debug(
"WebSearchInterception: Making follow-up request with search results"
)
verbose_logger.debug(
f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}"
)
verbose_logger.debug(
f"WebSearchInterception: Last message (tool_result): {user_message}"
)
# Correlation context for structured logging
_call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get(
"litellm_call_id", "unknown"
@ -742,61 +831,41 @@ class WebSearchInterceptionLogger(CustomLogger):
full_model_name = model # safe default before try block
# Use anthropic_messages.acreate for follow-up request
try:
max_tokens = self._resolve_max_tokens(
anthropic_messages_optional_request_params, kwargs
)
max_tokens = self._resolve_max_tokens(
anthropic_messages_optional_request_params, kwargs
)
verbose_logger.debug(
f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request"
)
verbose_logger.debug(
f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request"
)
# Create a copy of optional params without max_tokens (since we pass it explicitly)
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
kwargs_for_followup = self._prepare_followup_kwargs(kwargs)
kwargs_for_followup = self._prepare_followup_kwargs(kwargs)
# Get model from logging_obj.model_call_details["agentic_loop_params"]
# This preserves the full model name with provider prefix (e.g., "bedrock/invoke/...")
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = agentic_params.get("model", model)
verbose_logger.debug(
f"WebSearchInterception: Using model name: {full_model_name}"
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
final_response = await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=follow_up_messages,
model=full_model_name,
**optional_params_without_max_tokens,
**kwargs_for_followup,
)
verbose_logger.debug(
f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}"
)
verbose_logger.debug(
f"WebSearchInterception: Final response: {final_response}"
)
return final_response
except Exception as e:
verbose_logger.exception(
"WebSearchInterception: Follow-up request failed "
"[call_id=%s model=%s messages=%d searches=%d]: %s",
_call_id,
full_model_name,
len(follow_up_messages),
len(final_search_results),
str(e),
)
raise
full_model_name = agentic_params.get("model", model)
verbose_logger.debug(
"WebSearchInterception: Built anthropic request patch "
"[call_id=%s model=%s messages=%d searches=%d]",
_call_id,
full_model_name,
len(follow_up_messages),
len(final_search_results),
)
return AgenticLoopRequestPatch(
model=full_model_name,
messages=follow_up_messages,
max_tokens=max_tokens,
optional_params=optional_params_without_max_tokens,
kwargs=kwargs_for_followup,
)
async def _execute_search(self, query: str) -> str:
"""Execute a single web search using router's search tools"""
@ -883,7 +952,36 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs: Dict,
response_format: str = "openai",
) -> Any:
"""Execute litellm.search() and make follow-up chat completion request"""
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch = await self._build_chat_completion_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
optional_params=optional_params,
kwargs=kwargs,
response_format=response_format,
)
if request_patch.messages is None:
raise ValueError("WebSearchInterception: missing follow-up messages")
params = dict(optional_params)
params.update(request_patch.optional_params)
return await litellm.acompletion(
model=request_patch.model or model,
messages=request_patch.messages,
**params,
**request_patch.kwargs,
)
async def _build_chat_completion_request_patch( # noqa: PLR0915
self,
model: str,
messages: List[Dict],
tool_calls: List[Dict],
optional_params: Dict,
kwargs: Dict,
response_format: str = "openai",
) -> AgenticLoopRequestPatch:
"""Execute litellm.search() and build chat-completion rerun patch."""
# Extract search queries from tool_calls
search_tasks = []
@ -963,74 +1061,56 @@ class WebSearchInterceptionLogger(CustomLogger):
f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}"
)
# Use litellm.acompletion for follow-up request
try:
# Remove internal parameters that shouldn't be passed to follow-up request
internal_params = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
# Remove internal parameters that shouldn't be passed to follow-up request
internal_params = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception") and k not in internal_params
}
full_model_name = model
if "custom_llm_provider" in kwargs:
custom_llm_provider = kwargs["custom_llm_provider"]
if not model.startswith(custom_llm_provider) and "/" not in model:
full_model_name = f"{custom_llm_provider}/{model}"
verbose_logger.debug(
"WebSearchInterception: Built chat completion request patch model=%s messages=%d",
full_model_name,
len(follow_up_messages),
)
tools_param = optional_params.get("tools")
optional_params_clean = {
k: v
for k, v in optional_params.items()
if k
not in {
"tools",
"extra_body",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception")
and k not in internal_params
}
}
if tools_param is not None:
optional_params_clean["tools"] = tools_param
# Get full model name from kwargs
full_model_name = model
if "custom_llm_provider" in kwargs:
custom_llm_provider = kwargs["custom_llm_provider"]
# Reconstruct full model name with provider prefix if needed
if not model.startswith(custom_llm_provider):
# Check if model already has a provider prefix
if "/" not in model:
full_model_name = f"{custom_llm_provider}/{model}"
verbose_logger.debug(
f"WebSearchInterception: Using model name: {full_model_name}"
)
# Prepare tools for follow-up request (same as original)
tools_param = optional_params.get("tools")
# Remove tools and extra_body from optional_params to avoid issues
# extra_body often contains internal LiteLLM params that shouldn't be forwarded
optional_params_clean = {
k: v
for k, v in optional_params.items()
if k
not in {
"tools",
"extra_body",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
}
final_response = await litellm.acompletion(
model=full_model_name,
messages=follow_up_messages,
tools=tools_param,
**optional_params_clean,
**kwargs_for_followup,
)
verbose_logger.debug(
f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}"
)
return final_response
except Exception as e:
verbose_logger.exception(
f"WebSearchInterception: Follow-up request failed: {str(e)}"
)
raise
return AgenticLoopRequestPatch(
model=full_model_name,
messages=follow_up_messages,
optional_params=optional_params_clean,
kwargs=kwargs_for_followup,
)
async def _create_empty_search_result(self) -> str:
"""Create an empty search result for tool calls without queries"""

40
litellm/litellm_core_utils/core_helpers.py
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
@ -438,3 +439,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

9
litellm/litellm_core_utils/get_supported_openai_params.py
View file

@ -296,6 +296,15 @@ def get_supported_openai_params( # noqa: PLR0915
return OVHCloudAudioTranscriptionConfig().get_supported_openai_params(
model=model
)
elif custom_llm_provider == "scaleway":
if request_type == "transcription":
from litellm.llms.scaleway.audio_transcription.transformation import (
ScalewayAudioTranscriptionConfig,
)
return ScalewayAudioTranscriptionConfig().get_supported_openai_params(
model=model
)
elif custom_llm_provider == "elevenlabs":
if request_type == "transcription":
from litellm.llms.elevenlabs.audio_transcription.transformation import (

2
litellm/litellm_core_utils/litellm_logging.py
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,

2
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -684,7 +684,7 @@ def generic_cost_per_token( # noqa: PLR0915
- cache_creation
- image_tokens
)
# Clamp to zero: inconsistent streaming usage
# Clamp to zero: inconsistent streaming usage
if text_tokens < 0:
text_tokens = 0
prompt_tokens_details["text_tokens"] = text_tokens

37
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -34,6 +34,7 @@ from litellm.types.llms.anthropic import (
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionRequest,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
)
@ -67,6 +68,32 @@ class AnthropicMessagesHandler(BaseTranslation):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
"""Translate Anthropic request to OpenAI chat completion format."""
(
chat_completion_compatible_request,
_tool_name_mapping,
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
)
return chat_completion_compatible_request
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
"""
Convert Anthropic messages request data to OpenAI-spec structured messages.
Uses the Anthropic-to-OpenAI adapter to translate message format.
"""
messages = data.get("messages")
if messages is None:
return None
chat_completion_compatible_request = self._translate_to_openai(data)
result = cast(
List[AllMessageValues],
chat_completion_compatible_request.get("messages", []),
)
return result if result else None
async def process_input_messages(
self,
data: dict,
@ -82,13 +109,7 @@ class AnthropicMessagesHandler(BaseTranslation):
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
(
chat_completion_compatible_request,
_tool_name_mapping,
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
# Use a shallow copy to avoid mutating request data (pop on litellm_metadata).
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
)
chat_completion_compatible_request = self._translate_to_openai(data)
structured_messages = cast(
List[AllMessageValues],
@ -103,8 +124,6 @@ class AnthropicMessagesHandler(BaseTranslation):
chat_completion_compatible_request.get("tools", [])
)
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
# Step 1: Extract all text content and images
for msg_idx, message in enumerate(messages):

322
litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py Normal file
View file

@ -0,0 +1,322 @@
"""
Agentic Streaming Iterator for Anthropic Messages
Wraps the raw SSE byte stream from the Anthropic pass-through endpoint,
yields every chunk to the caller (preserving real streaming), collects
all bytes, and on stream exhaustion rebuilds the full Anthropic response
to run through agentic completion hooks. If an agentic hook fires, the
follow-up response is chained as Phase 2 of the same iterator.
"""
import json
from typing import Any, AsyncIterator, Dict, List, Optional, cast
from litellm._logging import verbose_logger
# ---------------------------------------------------------------------------
# SSE parsing helpers (module-level to keep the class lean)
# ---------------------------------------------------------------------------
def _parse_sse_events(raw: bytes) -> List[tuple]:
"""Return a list of (event_type, parsed_data_dict) from raw SSE bytes."""
text = raw.decode("utf-8", errors="replace")
lines = text.split("\n")
events: List[tuple] = []
current_event_type: Optional[str] = None
for line in lines:
stripped = line.strip()
if stripped.startswith("event:"):
current_event_type = stripped[len("event:") :].strip()
continue
if not stripped.startswith("data:"):
continue
data_str = stripped[len("data:") :].strip()
try:
data = json.loads(data_str)
except (json.JSONDecodeError, ValueError):
continue
event_type = current_event_type or data.get("type", "")
current_event_type = None
events.append((event_type, data))
return events
def _handle_message_start(data: Dict, response: Dict) -> None:
msg = data.get("message", {})
response["id"] = msg.get("id", response["id"])
response["model"] = msg.get("model", response["model"])
response["role"] = msg.get("role", response["role"])
usage = msg.get("usage", {})
if usage:
response["usage"]["input_tokens"] = usage.get("input_tokens", 0)
for key in ("cache_creation_input_tokens", "cache_read_input_tokens"):
if key in usage:
response["usage"][key] = usage[key]
def _handle_content_block_start(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", len(content_blocks))
block = data.get("content_block", {})
block_type = block.get("type", "text")
_BLOCK_TEMPLATES: Dict[str, Dict] = {
"text": {"type": "text", "text": ""},
"thinking": {"type": "thinking", "thinking": "", "signature": ""},
"redacted_thinking": {
"type": "redacted_thinking",
"data": block.get("data", ""),
},
}
if block_type == "tool_use":
content_blocks[idx] = {
"type": "tool_use",
"id": block.get("id", ""),
"name": block.get("name", ""),
"input": {},
"_partial_json": "",
}
elif block_type in _BLOCK_TEMPLATES:
content_blocks[idx] = dict(_BLOCK_TEMPLATES[block_type])
else:
content_blocks[idx] = dict(block)
def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", 0)
delta = data.get("delta", {})
delta_type = delta.get("type", "")
block = content_blocks.get(idx)
if block is None:
return
if delta_type == "text_delta":
block["text"] = block.get("text", "") + delta.get("text", "")
elif delta_type == "input_json_delta":
block["_partial_json"] = block.get("_partial_json", "") + delta.get(
"partial_json", ""
)
elif delta_type == "thinking_delta":
block["thinking"] = block.get("thinking", "") + delta.get("thinking", "")
elif delta_type == "signature_delta":
block["signature"] = delta.get("signature", block.get("signature", ""))
def _handle_content_block_stop(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", 0)
block = content_blocks.get(idx)
if block and block.get("type") == "tool_use":
partial = block.pop("_partial_json", "")
if partial:
try:
block["input"] = json.loads(partial)
except (json.JSONDecodeError, ValueError):
block["input"] = {"_raw": partial}
def _handle_message_delta(data: Dict, response: Dict) -> None:
delta = data.get("delta", {})
if "stop_reason" in delta:
response["stop_reason"] = delta["stop_reason"]
if "stop_sequence" in delta:
response["stop_sequence"] = delta["stop_sequence"]
usage = data.get("usage", {})
if usage.get("output_tokens") is not None:
response["usage"]["output_tokens"] = usage["output_tokens"]
for key in (
"input_tokens",
"cache_creation_input_tokens",
"cache_read_input_tokens",
):
if key in usage:
response["usage"][key] = usage[key]
class AgenticAnthropicStreamingIterator:
"""
Two-phase async iterator that enables agentic hooks on streaming
Anthropic Messages pass-through responses.
Phase 1: Yield raw SSE bytes from the upstream response while
accumulating them. When the inner iterator is exhausted,
rebuild the full Anthropic response dict and call agentic hooks.
Phase 2: If an agentic hook fires and returns a follow-up response
(streaming or non-streaming), yield those bytes to the caller.
"""
def __init__(
self,
completion_stream: AsyncIterator,
http_handler: Any,
model: str,
messages: List[Dict],
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
custom_llm_provider: str,
kwargs: Dict,
):
self._inner = completion_stream.__aiter__()
self._http_handler = http_handler
self._model = model
self._messages = messages
self._anthropic_messages_provider_config = anthropic_messages_provider_config
self._anthropic_messages_optional_request_params = (
anthropic_messages_optional_request_params
)
self._logging_obj = logging_obj
self._custom_llm_provider = custom_llm_provider
self._kwargs = kwargs
self._collected_bytes: List[bytes] = []
self._stream_exhausted = False
self._hook_processing_done = False
self._follow_up_iterator: Optional[AsyncIterator] = None
def __aiter__(self):
return self
async def __anext__(self) -> bytes:
# Phase 1: yield from upstream, collect bytes
if not self._stream_exhausted:
try:
chunk = await self._inner.__anext__()
self._collected_bytes.append(chunk)
return chunk
except StopAsyncIteration:
self._stream_exhausted = True
await self._process_agentic_hooks()
# Fall through to Phase 2
# Phase 2: yield from follow-up stream if one was created
if self._follow_up_iterator is not None:
chunk = await self._follow_up_iterator.__anext__()
return chunk
raise StopAsyncIteration
async def _process_agentic_hooks(self) -> None:
"""Rebuild the Anthropic response from collected SSE bytes and call hooks."""
if self._hook_processing_done:
return
self._hook_processing_done = True
if not self._collected_bytes:
return
try:
rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes)
if rebuilt is None:
verbose_logger.debug(
"AgenticStreamingIterator: Could not rebuild response from SSE bytes"
)
return
[
(
f"{b.get('type')}({b.get('name', '')})"
if b.get("type") == "tool_use"
else b.get("type")
)
for b in rebuilt.get("content", [])
]
result = await self._http_handler._call_agentic_completion_hooks(
response=rebuilt,
model=self._model,
messages=self._messages,
anthropic_messages_provider_config=self._anthropic_messages_provider_config,
anthropic_messages_optional_request_params=self._anthropic_messages_optional_request_params,
logging_obj=self._logging_obj,
stream=True,
custom_llm_provider=self._custom_llm_provider,
kwargs=self._kwargs,
)
if result is None:
return
if hasattr(result, "__aiter__"):
self._follow_up_iterator = result.__aiter__()
elif isinstance(result, dict):
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
fake = FakeAnthropicMessagesStreamIterator(
response=cast(AnthropicMessagesResponse, result)
)
self._follow_up_iterator = fake.__aiter__()
else:
verbose_logger.warning(
"AgenticStreamingIterator: Unexpected result type from hooks: %s",
type(result).__name__,
)
except Exception as e:
_call_id = getattr(self._logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
"AgenticStreamingIterator: Error in agentic hook processing "
"[call_id=%s model=%s]: %s",
_call_id,
self._model,
str(e),
)
@staticmethod
def _rebuild_anthropic_response_from_sse(
raw_bytes: List[bytes],
) -> Optional[Dict[str, Any]]:
"""
Parse collected SSE bytes into an Anthropic Messages response dict.
Processes SSE events in order:
- message_start -> envelope (id, model, role, usage)
- content_block_start -> new content block
- content_block_delta -> accumulate text/json/thinking deltas
- content_block_stop -> finalize block
- message_delta -> stop_reason, output usage
- message_stop -> end
"""
events = _parse_sse_events(b"".join(raw_bytes))
response: Dict[str, Any] = {
"id": "",
"type": "message",
"role": "assistant",
"model": "",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
}
content_blocks: Dict[int, Dict[str, Any]] = {}
saw_message_start = False
for event_type, data in events:
if event_type == "message_start":
saw_message_start = True
_handle_message_start(data, response)
elif event_type == "content_block_start":
_handle_content_block_start(data, content_blocks)
elif event_type == "content_block_delta":
_handle_content_block_delta(data, content_blocks)
elif event_type == "content_block_stop":
_handle_content_block_stop(data, content_blocks)
elif event_type == "message_delta":
_handle_message_delta(data, response)
if not saw_message_start:
return None
for idx in sorted(content_blocks.keys()):
block = content_blocks[idx]
block.pop("_partial_json", None)
response["content"].append(block)
return response

11
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -5,6 +5,7 @@ if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
class BaseTranslation(ABC):
@ -101,6 +102,16 @@ class BaseTranslation(ABC):
"""
return responses_so_far
def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]:
"""
Convert request data to OpenAI-spec structured messages.
Override in subclasses for format-specific conversion.
Returns None if no convertible content is found.
"""
return None
def extract_request_tool_names(self, data: dict) -> List[str]:
"""
Extract tool names from the request body for allowlist/policy checks.

0
litellm/llms/bedrock/chat/mantle/__init__.py Normal file
View file

91
litellm/llms/bedrock/chat/mantle/transformation.py Normal file
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

26
litellm/llms/bedrock/common_utils.py
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":

46
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
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(

69
litellm/llms/bedrock/messages/mantle_transformation.py Normal file
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

439
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -78,6 +78,10 @@ from litellm.types.containers.main import (
DeleteContainerResult,
)
from litellm.types.files import TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -2047,7 +2051,23 @@ class BaseLLMHTTPHandler:
request_body=request_body,
litellm_logging_obj=logging_obj,
)
initial_response = completion_stream
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
)
initial_response = AgenticAnthropicStreamingIterator(
completion_stream=completion_stream,
http_handler=self,
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
return initial_response
else:
initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
model=model,
@ -2055,7 +2075,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
)
# Call agentic completion hooks
# Call agentic completion hooks (non-streaming path only)
final_response = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
@ -2063,7 +2083,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream or False,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
@ -4516,6 +4536,167 @@ class BaseLLMHTTPHandler:
return stream, data
return stream, data
@staticmethod
def _get_agentic_loop_settings(kwargs: Dict) -> Tuple[int, int, List[str]]:
depth = int(kwargs.get("_agentic_loop_depth", 0) or 0)
max_loops = int(kwargs.get("max_agentic_loops", 3) or 3)
fingerprints = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
return depth, max(max_loops, 1), fingerprints
@staticmethod
def _check_agentic_loop_safety(
tool_calls: Any,
fingerprints: List[str],
depth: int,
max_loops: int,
model: str,
) -> str:
"""
Evaluate agentic-loop safety guards (fingerprint cycle / max depth).
Raises ValueError on abort. Returns the current fingerprint on success.
These checks must not be swallowed by the per-callback ``except Exception``
block that wraps callback dispatch — they are bounded-loop / cycle-break
safety rails and must abort the agentic dispatch when they trip.
"""
fingerprint = BaseLLMHTTPHandler._fingerprint_agentic_tools(tool_calls)
if fingerprint in fingerprints:
raise ValueError(
"Agentic loop detected repeated tool-call fingerprint; aborting rerun"
)
if depth >= max_loops:
raise ValueError(
f"Exceeded max_agentic_loops={max_loops} for model={model}"
)
return fingerprint
@staticmethod
def _fingerprint_agentic_tools(tools: Dict) -> str:
try:
return json.dumps(tools, sort_keys=True, default=str)
except Exception:
return str(tools)
async def _execute_anthropic_agentic_plan(
self,
plan: AgenticLoopPlan,
model: str,
messages: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: "LiteLLMLoggingObj",
kwargs: Dict,
depth: int,
max_loops: int,
fingerprints: List[str],
fingerprint: str,
stream: bool = False,
) -> Any:
from litellm.anthropic_interface import messages as anthropic_messages
patch = plan.request_patch or AgenticLoopRequestPatch()
if patch.messages is None:
raise ValueError("Agentic loop plan missing patched messages")
full_model_name = model
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = cast(str, agentic_params.get("model", model))
optional_params = dict(anthropic_messages_optional_request_params)
optional_params.update(patch.optional_params)
if patch.tools is not None:
optional_params["tools"] = patch.tools
max_tokens = patch.max_tokens
if max_tokens is None:
max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None))
else:
optional_params.pop("max_tokens", None)
if max_tokens is None:
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
internal_keys = {"litellm_logging_obj"}
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception")
and not k.startswith("_compression_interception")
and k not in internal_keys
and k not in optional_params
}
kwargs_for_followup.update(patch.kwargs)
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
return await anthropic_messages.acreate(
**{
"max_tokens": max_tokens,
"messages": patch.messages,
"model": patch.model or full_model_name,
"stream": stream,
**optional_params,
**kwargs_for_followup,
}
)
async def _execute_chat_completion_agentic_plan(
self,
plan: AgenticLoopPlan,
model: str,
messages: List[Dict],
optional_params: Dict,
kwargs: Dict,
custom_llm_provider: str,
depth: int,
max_loops: int,
fingerprints: List[str],
fingerprint: str,
) -> Any:
patch = plan.request_patch or AgenticLoopRequestPatch()
if patch.messages is None:
raise ValueError("Agentic loop plan missing patched messages")
full_model_name = patch.model or model
if "/" not in full_model_name:
full_model_name = f"{custom_llm_provider}/{full_model_name}"
optional_params_for_followup = dict(optional_params)
optional_params_for_followup.update(patch.optional_params)
if patch.tools is not None:
optional_params_for_followup["tools"] = patch.tools
internal_params = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception")
and not k.startswith("_compression_interception")
and k not in internal_params
}
kwargs_for_followup.update(patch.kwargs)
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
return await litellm.acompletion(
model=full_model_name,
messages=patch.messages,
**optional_params_for_followup,
**kwargs_for_followup,
)
async def _call_agentic_completion_hooks(
self,
response: Any,
@ -4541,45 +4722,111 @@ class BaseLLMHTTPHandler:
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
tools = anthropic_messages_optional_request_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
for callback in callbacks:
if not isinstance(callback, CustomLogger):
continue
should_run: bool = False
tool_calls: Any = None
try:
if isinstance(callback, CustomLogger):
# First: Check if agentic loop should run
(
should_run,
tool_calls,
) = await callback.async_should_run_agentic_loop(
response=response,
# First: Check if agentic loop should run. Wrap in try/except
# to shield from buggy user callbacks — a callback crash should
# not abort the whole request.
(
should_run,
tool_calls,
) = await callback.async_should_run_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
except Exception as e:
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in "
"async_should_run_agentic_loop [call_id=%s model=%s]: %s",
_call_id,
model,
str(e),
)
continue
if not should_run:
continue
# Safety guards must run OUTSIDE the callback try/except — they are
# bounded-loop / cycle-break rails that must propagate to the caller.
fingerprint = self._check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
max_loops=max_loops,
model=model,
)
try:
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
build_plan_overridden = (
callback.__class__.async_build_agentic_loop_plan
is not CustomLogger.async_build_agentic_loop_plan
)
if not build_plan_overridden:
return await callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
tools=tools,
response=response,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
kwargs=kwargs_with_provider,
)
if should_run:
# Second: Execute agentic loop
# Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = (
custom_llm_provider
)
agentic_response = await callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
response=response,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
# First hook that runs agentic loop wins
return agentic_response
plan = await callback.async_build_agentic_loop_plan(
tools=tool_calls,
model=model,
messages=messages,
response=response,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
if plan.response_override is not None:
return plan.response_override
if plan.terminate:
verbose_logger.debug(
"Agentic loop terminated by callback=%s reason=%s",
callback.__class__.__name__,
plan.stop_reason,
)
return response
if not plan.run_agentic_loop:
continue
return await self._execute_anthropic_agentic_plan(
plan=plan,
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs_with_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
stream=stream,
)
except Exception as e:
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
@ -4653,52 +4900,104 @@ class BaseLLMHTTPHandler:
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
tools = optional_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
# Check if callback has the chat completion agentic loop method
if not hasattr(
callback, "async_should_run_chat_completion_agentic_loop"
):
continue
if not isinstance(callback, CustomLogger):
continue
if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"):
continue
# First: Check if agentic loop should run
(
should_run,
tool_calls,
) = await callback.async_should_run_chat_completion_agentic_loop(
response=response,
should_run: bool = False
tool_calls: Any = None
try:
(
should_run,
tool_calls,
) = await callback.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
except Exception as e:
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in "
"async_should_run_chat_completion_agentic_loop: %s",
str(e),
)
continue
if not should_run:
continue
# Safety guards must run OUTSIDE the callback try/except — they are
# bounded-loop / cycle-break rails that must propagate to the caller.
fingerprint = self._check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
max_loops=max_loops,
model=model,
)
try:
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
build_plan_overridden = (
callback.__class__.async_build_chat_completion_agentic_loop_plan
is not CustomLogger.async_build_chat_completion_agentic_loop_plan
)
if not build_plan_overridden:
return await callback.async_run_chat_completion_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
tools=tools,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
kwargs=kwargs_with_provider,
)
if should_run:
# Second: Execute agentic loop
# Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = (
custom_llm_provider
)
agentic_response = (
await callback.async_run_chat_completion_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
)
# First hook that runs agentic loop wins
return agentic_response
plan = await callback.async_build_chat_completion_agentic_loop_plan(
tools=tool_calls,
model=model,
messages=messages,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
if plan.response_override is not None:
return plan.response_override
if plan.terminate:
verbose_logger.debug(
"Agentic chat loop terminated by callback=%s reason=%s",
callback.__class__.__name__,
plan.stop_reason,
)
return response
if not plan.run_agentic_loop:
continue
return await self._execute_chat_completion_agentic_plan(
plan=plan,
model=model,
messages=messages,
optional_params=optional_params,
kwargs=kwargs_with_provider,
custom_llm_provider=custom_llm_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
except Exception as e:
verbose_logger.exception(
f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {str(e)}"

4
litellm/llms/github_copilot/authenticator.py
View file

@ -294,9 +294,7 @@ class Authenticator:
access_token_url = os.getenv(
"GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL
)
client_id = os.getenv(
"GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID
)
client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
for attempt in range(max_attempts):
try:

22
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -48,6 +48,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
"""
Convert chat completions request data to OpenAI-spec structured messages.
Messages are already in OpenAI format, so this is a simple extraction.
"""
messages = data.get("messages")
if messages is None:
return None
return cast(List[AllMessageValues], messages)
async def process_input_messages(
self,
data: dict,
@ -68,9 +79,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
tool_calls_to_check: List[ChatCompletionToolParam] = []
text_task_mappings: List[Tuple[int, Optional[int]]] = []
tool_call_task_mappings: List[Tuple[int, int]] = []
# text_task_mappings: Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
# tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call
# Step 1: Extract all text content, images, and tool calls
for msg_idx, message in enumerate(messages):
@ -92,12 +100,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check # type: ignore
if messages:
msg_list = cast(List[AllMessageValues], messages)
structured_messages = self.get_structured_messages(data)
if structured_messages:
inputs["structured_messages"] = (
openai_messages_without_system(msg_list)
openai_messages_without_system(structured_messages)
if skip_system
else msg_list
else structured_messages
)
# Pass tools (function definitions) to the guardrail
tools = data.get("tools")

26
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -43,6 +43,7 @@ from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
)
@ -70,6 +71,24 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
"""
Convert Responses API request data to OpenAI-spec structured messages.
Transforms `input` (string or ResponseInputParam) and optional
`instructions` into chat completion messages.
"""
input_data = data.get("input")
if input_data is None:
return None
messages = (
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input_data,
responses_api_request=data,
)
)
return cast(List[AllMessageValues], messages) if messages else None
async def process_input_messages(
self,
data: dict,
@ -86,12 +105,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if input_data is None:
return data
structured_messages = (
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input_data,
responses_api_request=data,
)
)
structured_messages = self.get_structured_messages(data)
# Handle simple string input
if isinstance(input_data, str):

158
litellm/llms/scaleway/audio_transcription/transformation.py Normal file
View file

@ -0,0 +1,158 @@
"""
Support for Scaleway's OpenAI-compatible `/v1/audio/transcriptions` endpoint.
API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path-audio-create-an-audio-transcription
"""
from typing import List, Optional, Union
import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
)
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
class ScalewayAudioTranscriptionException(BaseLLMException):
pass
class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
def get_supported_openai_params(
self, model: str
) -> List[OpenAIAudioTranscriptionOptionalParams]:
return [
"language",
"prompt",
"response_format",
"temperature",
"timestamp_granularities",
]
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)
for k, v in non_default_params.items():
if k in supported_params:
optional_params[k] = v
return optional_params
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:
api_base = (
"https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/")
)
return f"{api_base}/audio/transcriptions"
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return ScalewayAudioTranscriptionException(
message=error_message,
status_code=status_code,
headers=headers,
)
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:
if api_key is None:
api_key = get_secret_str("SCW_SECRET_KEY")
if not api_key:
raise ScalewayAudioTranscriptionException(
message=(
"Scaleway API key not found. Pass `api_key=...` or set the "
"SCW_SECRET_KEY environment variable."
),
status_code=401,
headers={},
)
default_headers = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
}
default_headers.update(headers or {})
return default_headers
def transform_audio_transcription_request(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
) -> AudioTranscriptionRequestData:
processed_audio = process_audio_file(audio_file)
form_fields: dict = {"model": model}
for key in self.get_supported_openai_params(model):
value = optional_params.get(key)
if value is not None:
form_fields[key] = value
files = {
"file": (
processed_audio.filename,
processed_audio.file_content,
processed_audio.content_type,
)
}
return AudioTranscriptionRequestData(data=form_fields, files=files)
def transform_audio_transcription_response(
self,
raw_response: httpx.Response,
) -> TranscriptionResponse:
content_type = (raw_response.headers.get("content-type") or "").lower()
if "application/json" not in content_type:
return TranscriptionResponse(text=raw_response.text)
try:
response_json = raw_response.json()
except Exception:
raise ScalewayAudioTranscriptionException(
message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
text = response_json.get("text") or ""
response = TranscriptionResponse(text=text)
if "segments" in response_json:
response["segments"] = response_json["segments"]
if "language" in response_json:
response["language"] = response_json["language"]
response._hidden_params = response_json
return response

16
litellm/model_prices_and_context_window_backup.json
View file

@ -23117,6 +23117,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,

4
litellm/passthrough/utils.py
View file

@ -79,7 +79,9 @@ class BasePassthroughUtils:
for header_name, header_value in request_headers.items():
if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX):
# Strip the 'x-pass-' prefix and normalize to lowercase
actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower()
actual_header_name = header_name[
len(PASS_THROUGH_HEADER_PREFIX) :
].lower()
if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any(
actual_header_name.startswith(p)
for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES

2
litellm/provider_endpoints_support_backup.json
View file

@ -1950,7 +1950,7 @@
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_transcriptions": true,
"audio_speech": false,
"moderations": false,
"batches": false,

8
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
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(

103
litellm/proxy/_new_secret_config.yaml
View file

@ -1,32 +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"
# ---- Underlying deployments the router picks from ---------------------
- model_name: fast
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-20250514"
litellm_params:
model: anthropic/claude-sonnet-4-20250514
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)
adaptive_router_preferences:
quality_tier: 2
strengths: []
- model_name: smart
litellm_params:
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

39
litellm/proxy/auth/auth_checks.py
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
@ -3126,9 +3132,7 @@ async def _virtual_key_max_budget_alert_check(
alert_email_config: Optional[Dict[str, List[str]]] = (
_merge_budget_alert_email_configs(
global_cfg=litellm.default_key_max_budget_alert_emails,
per_key_cfg=(valid_token.metadata or {}).get(
"max_budget_alert_emails"
),
per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"),
)
)
@ -3138,7 +3142,9 @@ async def _virtual_key_max_budget_alert_check(
(int(k) for k in alert_email_config if k.isdigit()),
default=None,
)
if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0):
if min_pct is None or valid_token.spend < valid_token.max_budget * (
min_pct / 100.0
):
return
call_info = CallInfo(
@ -3164,8 +3170,7 @@ async def _virtual_key_max_budget_alert_check(
else:
# Old path: existing single 80% threshold — completely unchanged
alert_threshold = (
valid_token.max_budget
* EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
if (
@ -3666,12 +3671,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,
@ -3687,9 +3700,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}",
)

17
litellm/proxy/common_utils/callback_utils.py
View file

@ -37,6 +37,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915
if isinstance(value, list):
imported_list: List[Any] = []
for callback in value: # ["presidio", <my-custom-callback>]
if isinstance(callback, str) and callback == "compression_interception":
from litellm.integrations.compression_interception.handler import (
CompressionInterceptionLogger,
)
compression_interception_obj = (
CompressionInterceptionLogger.initialize_from_proxy_config(
litellm_settings=litellm_settings,
callback_specific_params=callback_specific_params,
)
)
imported_list.append(compression_interception_obj)
continue
# check if callback is a custom logger compatible callback
if isinstance(callback, str):
callback = LoggingCallbackManager._add_custom_callback_generic_api_str(
@ -419,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)

131
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
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,

15
litellm/proxy/db/prisma_client.py
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(

52
litellm/proxy/example_config_yaml/adaptive_router_example.yaml Normal file
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

181
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
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")

9
litellm/proxy/health_check.py
View file

@ -306,7 +306,9 @@ def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool:
return "*" in _deployment_model_string_for_health_check(litellm_params)
def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]:
def _resolve_health_check_max_tokens(
model_info: dict, litellm_params: dict
) -> Optional[int]:
"""
Pick max_tokens for the health check request.
@ -341,10 +343,7 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) ->
return int(tokens_reasoning)
if not is_reasoning and tokens_non_reasoning is not None:
return int(tokens_non_reasoning)
if (
is_reasoning
and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None
):
if is_reasoning and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None:
return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING)
if BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None:

6
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -1121,14 +1121,14 @@ async def _db_health_readiness_check():
return db_health_cache
except Exception as e:
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
PrismaDBExceptionHandler.handle_db_exception(e)
if PrismaDBExceptionHandler.is_database_transport_error(e):
try:
verbose_proxy_logger.warning(
"_db_health_readiness_check: health_check failed, attempting reconnect"
)
await prisma_client.disconnect()
await prisma_client.connect()
await prisma_client.attempt_db_reconnect(
reason="health_readiness_check"
)
await prisma_client.health_check()
verbose_proxy_logger.info(
"_db_health_readiness_check: reconnect succeeded"

34
litellm/proxy/hooks/max_budget_limiter.py
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:

6
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -1570,9 +1570,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
user_api_key_project_id = standard_logging_metadata.get(
"user_api_key_project_id"
)
user_api_key_end_user_id = kwargs.get(
"user"
) or standard_logging_metadata.get("user_api_key_end_user_id")
user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get(
"user_api_key_end_user_id"
)
model_group = get_model_group_from_litellm_kwargs(kwargs)
# Get total tokens from response

1
litellm/proxy/hooks/proxy_track_cost_callback.py
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:

57
litellm/proxy/litellm_pre_call_utils.py
View file

@ -1,5 +1,6 @@
import asyncio
import copy
import re
import time
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
@ -28,6 +29,14 @@ _SPECIAL_HEADERS_CACHE = frozenset(
v.value.lower() for v in SpecialHeaders._member_map_.values()
)
# Matches any header of the form x-<something>-session-id (case-insensitive).
# Excludes the two explicit litellm headers which are handled with higher priority.
_GENERIC_SESSION_ID_HEADER_RE = re.compile(r"^x-.+-session-id$", re.IGNORECASE)
_EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session-id"})
# Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores
# (covers UUIDs and most common session-id formats).
_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
def _sanitize_for_log(value: Any) -> str:
"""
@ -115,13 +124,43 @@ def _get_metadata_variable_name(request: Request) -> str:
return "metadata"
def _extract_generic_session_id_from_headers(
normalized: Dict[str, str],
) -> Optional[str]:
"""
Scan a normalised (lower-cased keys) header dict for any header that looks
like ``x-<vendor>-session-id`` and whose value is a plausible session/trace
identifier (alphanumeric + hyphens/underscores, at least 8 chars).
The two explicit LiteLLM headers (``x-litellm-trace-id`` /
``x-litellm-session-id``) are excluded here because they are handled with
higher priority by the caller.
Example: ``x-claude-code-session-id: e96634a3-fa28-4083-b354-55542e2dca01``
"""
for key, value in normalized.items():
if (
key not in _EXPLICIT_SESSION_HEADERS
and _GENERIC_SESSION_ID_HEADER_RE.match(key)
and isinstance(value, str)
and _SESSION_ID_VALUE_RE.match(value)
):
return value
return None
def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]:
"""
Extract chain id for call chaining from request headers.
x-litellm-trace-id and x-litellm-session-id are interchangeable; when both
are present, x-litellm-trace-id takes precedence. Header keys are matched
case-insensitively so this works with raw header dicts from any transport.
Priority order:
1. ``x-litellm-trace-id`` (explicit, highest priority)
2. ``x-litellm-session-id`` (explicit)
3. Any ``x-<vendor>-session-id`` header whose value looks like a session id
(alphanumeric / UUID, at least 8 chars). E.g. ``x-claude-code-session-id``.
Header keys are matched case-insensitively so this works with raw header
dicts from any transport.
Used by MCP (and other paths that have raw_headers but no Request) to set
litellm_trace_id/litellm_session_id for spend logs and logging consistency.
@ -129,8 +168,10 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str
if not headers:
return None
normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)}
return normalized.get("x-litellm-trace-id") or normalized.get(
"x-litellm-session-id"
return (
normalized.get("x-litellm-trace-id")
or normalized.get("x-litellm-session-id")
or _extract_generic_session_id_from_headers(normalized)
)
@ -649,10 +690,8 @@ class LiteLLMProxyRequestSetup:
#########################################################################################
agent_id_from_header = headers.get("x-litellm-agent-id")
# x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining
chain_id = headers.get("x-litellm-trace-id") or headers.get(
"x-litellm-session-id"
)
# Explicit litellm headers take precedence; fall back to any x-*-session-id header.
chain_id = get_chain_id_from_headers(dict(headers))
if agent_id_from_header:
metadata_from_headers["agent_id"] = agent_id_from_header

45
litellm/proxy/management_endpoints/common_utils.py
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:

8
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -2120,9 +2120,7 @@ async def delete_user(
for m in all_target_memberships:
if not m.organization_id:
continue
target_org_ids_by_user.setdefault(m.user_id, set()).add(
m.organization_id
)
target_org_ids_by_user.setdefault(m.user_id, set()).add(m.organization_id)
# check that all teams passed exist
for user_id in data.user_ids:
@ -2141,9 +2139,7 @@ async def delete_user(
# Org-admin may only delete users whose entire org membership is
# within their admin scope. A target with ANY org outside the
# caller's scope (or no org at all) requires PROXY_ADMIN.
if not target_org_ids or not target_org_ids.issubset(
caller_admin_org_ids
):
if not target_org_ids or not target_org_ids.issubset(caller_admin_org_ids):
raise HTTPException(
status_code=403,
detail={

27
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -1336,7 +1336,9 @@ if MCP_AVAILABLE:
return _redact_mcp_credentials(temp_record)
def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer:
def _get_cached_temporary_mcp_server_or_404(
server_id: str, request: Optional[Request] = None
) -> MCPServer:
server = 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
@ -1344,10 +1346,14 @@ if MCP_AVAILABLE:
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 +1364,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 +1378,7 @@ 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 = _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 +1407,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 +1422,7 @@ 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 = _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 +1451,14 @@ 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 = _get_cached_temporary_mcp_server_or_404(server_id, request=request)
request_data = await _read_request_body(request=request)
data: dict = {**request_data}

5
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -1078,10 +1078,7 @@ async def organization_member_update(
LitellmUserRoles.PROXY_ADMIN.value,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
):
if (
user_api_key_dict.user_role
!= LitellmUserRoles.PROXY_ADMIN.value
):
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
status_code=403,
detail={

13
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -1570,8 +1570,7 @@ async def update_team( # noqa: PLR0915
current_org_id = getattr(existing_team_row, "organization_id", None)
if (
data.organization_id != current_org_id
and user_api_key_dict.user_role
!= LitellmUserRoles.PROXY_ADMIN.value
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
):
# Is the caller org_admin of the destination org?
caller_memberships = (
@ -2609,6 +2608,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(
@ -2621,6 +2629,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

70
litellm/proxy/management_helpers/utils.py
View file

@ -140,6 +140,62 @@ 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
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 +277,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 = (

49
litellm/proxy/proxy_cli.py
View file

@ -577,6 +577,22 @@ class ProxyInitializationHelpers:
help="Exit with error if database migration fails on startup.",
envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
)
@click.option(
"--use_v2_migration_resolver",
is_flag=True,
default=False,
help=(
"Opt into the v2 migration resolver. Avoids the diff-and-force recovery "
"path that can cause schema thrashing during rolling deploys where two "
"LiteLLM versions contend for the same DB. Default is the v1 resolver."
),
)
@click.option(
"--reload",
is_flag=True,
default=False,
help="Enable uvicorn hot reload (dev only). Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
)
def run_server( # noqa: PLR0915
host,
port,
@ -618,6 +634,8 @@ def run_server( # noqa: PLR0915
keepalive_timeout,
max_requests_before_restart,
enforce_prisma_migration_check: bool,
use_v2_migration_resolver: bool,
reload: bool,
):
if setup:
from litellm.setup_wizard import run_setup_wizard
@ -886,9 +904,31 @@ def run_server( # noqa: PLR0915
):
check_prisma_schema_diff(db_url=None)
else:
if not PrismaManager.setup_database(
use_migrate=not use_prisma_db_push
):
if not use_v2_migration_resolver:
print( # noqa
"\033[1;33mLiteLLM Proxy: Using default (v1) migration resolver. "
"If your deployment has seen schema thrashing during rolling "
"deploys, try --use_v2_migration_resolver (safer: avoids the "
"diff-and-force recovery that caused the thrash).\033[0m"
)
try:
setup_ok = PrismaManager.setup_database(
use_migrate=not use_prisma_db_push,
use_v2_resolver=use_v2_migration_resolver,
)
except RuntimeError as e:
# v2 resolver raises on unrecoverable migration errors
# (e.g. non-idempotent failures, permission issues).
# v1 never raises here, so this only fires when the
# operator opted into v2.
print( # noqa
"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. "
f"{e}\033[0m",
file=sys.stderr,
flush=True,
)
sys.exit(2)
if not setup_ok:
if enforce_prisma_migration_check:
print( # noqa
"\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. "
@ -954,6 +994,9 @@ def run_server( # noqa: PLR0915
if loop_type:
uvicorn_args["loop"] = loop_type
if reload:
uvicorn_args["reload"] = True
uvicorn.run(
**uvicorn_args,
workers=num_workers,

90
litellm/proxy/proxy_server.py
View file

@ -952,6 +952,17 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
_run_background_health_check()
) # start the background health check coroutine.
# Start adaptive-router queue flusher unconditionally — adaptive routers
# may be added later via `/config/reload`, and the flusher is a no-op when
# `llm_router.adaptive_routers` is empty. Per-router DB state is loaded
# lazily by the flusher on first tick (see `_state_loaded` flag) so
# hot-reloaded routers also get their persisted priors.
if llm_router is not None and getattr(llm_router, "adaptive_routers", None):
for _ar in llm_router.adaptive_routers.values():
await _ar.load_state_from_db(prisma_client)
_ar._state_loaded = True
asyncio.create_task(_adaptive_router_flusher_loop())
## [Optional] Initialize dd tracer
ProxyStartupEvent._init_dd_tracer()
@ -1795,6 +1806,7 @@ async def increment_spend_counters(
team_id: Optional[str],
user_id: Optional[str],
response_cost: Optional[float],
org_id: Optional[str] = None,
):
"""
Atomically increment spend counters for budget enforcement.
@ -1881,6 +1893,20 @@ async def increment_spend_counters(
increment=response_cost,
)
if user_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:user:{user_id}",
source_cache_key=user_id,
increment=response_cost,
)
if org_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:org:{org_id}",
source_cache_key=f"org_id:{org_id}",
increment=response_cost,
)
async def _init_and_increment_spend_counter(
counter_key: str,
@ -2427,6 +2453,38 @@ def _write_health_state_to_router_cache(
)
_ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS = 10
async def _adaptive_router_flusher_loop():
"""
Drain every AdaptiveRouter's in-memory state + session aggregators into
Postgres on a fixed cadence. Hot-path writes go to memory; this loop is
the only writer to the adaptive router DB tables.
"""
global llm_router, prisma_client
while True:
try:
await asyncio.sleep(_ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS)
adaptive_routers = getattr(llm_router, "adaptive_routers", None) or {}
if not adaptive_routers or prisma_client is None:
continue
for ar in adaptive_routers.values():
# Lazy state load: covers adaptive routers registered via
# `/config/reload` after proxy boot.
if not getattr(ar, "_state_loaded", False):
try:
await ar.load_state_from_db(prisma_client)
finally:
ar._state_loaded = True
await ar.queue.flush_state_to_db(prisma_client)
await ar.queue.flush_session_to_db(prisma_client)
except asyncio.CancelledError:
raise
except Exception:
verbose_proxy_logger.exception("adaptive_router flusher iteration failed")
async def _run_background_health_check():
"""
Periodically run health checks in the background on the endpoints.
@ -13953,6 +14011,38 @@ async def home(request: Request):
return "LiteLLM: RUNNING"
@router.get(
"/adaptive_router/state",
tags=["adaptive_router"],
dependencies=[Depends(user_api_key_auth)],
)
async def get_adaptive_router_state(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""Return live bandit posteriors + queue depth for every configured adaptive router.
Admin-only. Returns 404 if no adaptive router is configured.
Response shape: `{"routers": [<snapshot>, ...]}` — one snapshot per
adaptive-router deployment. Each snapshot's `router_name` field identifies
which deployment it came from.
"""
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
if llm_router is None or not llm_router.adaptive_routers:
raise HTTPException(
status_code=404,
detail={"error": "No adaptive_router is configured on this proxy."},
)
snapshots = [
await ar.get_state_snapshot() for ar in llm_router.adaptive_routers.values()
]
return {"routers": snapshots}
@router.get("/routes", dependencies=[Depends(user_api_key_auth)])
async def get_routes():
"""

43
litellm/proxy/schema.prisma
View file

@ -1223,3 +1223,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")
}

7
litellm/proxy/utils.py
View file

@ -940,7 +940,11 @@ class ProxyLogging:
Result from the guardrail execution
"""
# Use unified_guardrail if callback has apply_guardrail method
use_unified = "apply_guardrail" in type(callback).__dict__
has_apply_guardrail = "apply_guardrail" in type(callback).__dict__
use_unified = has_apply_guardrail and not (
hook_type == "during_call"
and getattr(callback, "use_native_during_call_hook", False)
)
if use_unified:
data["guardrail_to_apply"] = callback
@ -1540,6 +1544,7 @@ class ProxyLogging:
if (
"apply_guardrail" in type(callback).__dict__
and user_api_key_dict is not None
and not getattr(callback, "use_native_during_call_hook", False)
):
data["guardrail_to_apply"] = callback
guardrail_task = self._run_guardrail_task_with_enrichment(

310
litellm/router.py
View file

@ -200,12 +200,20 @@ if TYPE_CHECKING:
from litellm.router_strategy.complexity_router.complexity_router import (
ComplexityRouter,
)
from litellm.router_strategy.adaptive_router.adaptive_router import (
AdaptiveRouter,
)
from litellm.router_strategy.quality_router.quality_router import (
QualityRouter,
)
Span = Union[_Span, Any]
else:
Span = Any
AutoRouter = Any
ComplexityRouter = Any
AdaptiveRouter = Any
QualityRouter = Any
PreRoutingHookResponse = Any
@ -464,6 +472,8 @@ class Router:
) # {"TEAM_ID": PatternMatchRouter}
self.auto_routers: Dict[str, "AutoRouter"] = {}
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {}
self.quality_routers: Dict[str, "QualityRouter"] = {}
# Initialize model_group_alias early since it's used in set_model_list
self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = (
@ -5364,8 +5374,13 @@ class Router:
_request_team_id: Optional[str] = (kwargs.get("metadata", {}) or {}).get(
"user_api_key_team_id"
)
all_deployments = self._get_all_deployments(
model_name=original_model_group, team_id=_request_team_id
# Use wildcard-aware lookup so order-based fallback also works for model
# groups resolved via pattern routing (e.g. `openai/*` -> `openai/gpt-4.1-mini`).
all_deployments = (
self.get_model_list(
model_name=original_model_group, team_id=_request_team_id
)
or []
)
_order_set: set = {
litellm.utils._get_deployment_order(d)
@ -5884,7 +5899,7 @@ class Router:
response = await response
## PROCESS RESPONSE HEADERS
response = await self.set_response_headers(
response=response, model_group=model_group
response=response, model_group=model_group, request_kwargs=kwargs
)
return response
@ -6810,10 +6825,15 @@ class Router:
Check if the deployment is an auto-router deployment (semantic router).
Returns True if the litellm_params model starts with "auto_router/"
but NOT "auto_router/complexity_router" (which uses complexity routing).
but NOT "auto_router/complexity_router" or "auto_router/adaptive_router"
(which use the complexity-router and adaptive-router strategies).
"""
if litellm_params.model.startswith("auto_router/complexity_router"):
return False # This is handled by complexity_router
if litellm_params.model.startswith("auto_router/adaptive_router"):
return False # This is handled by adaptive_router
if litellm_params.model.startswith("auto_router/quality_router"):
return False # This is handled by quality_router
if litellm_params.model.startswith("auto_router/"):
return True
return False
@ -6920,6 +6940,196 @@ class Router:
)
self.complexity_routers[deployment.model_name] = complexity_router
def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
"""True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
return litellm_params.model.startswith("auto_router/adaptive_router")
def _finalize_adaptive_router_if_configured(self) -> None:
"""Locate every adaptive-router deployment in the finalized model_list and
build an AdaptiveRouter for each. Safe no-op when none are configured.
Idempotent: skips any deployment whose model_name is already initialized."""
# Drop any adaptive-router hooks left over from a previous Router
# instance (e.g. after `/config/reload` replaced `llm_router`). Without
# this, stale AdaptiveRouterPostCallHook callbacks from the old Router
# remain wired up in `litellm.callbacks` and double-fire signal
# recording for every request.
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
for _cb_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
litellm.logging_callback_manager.remove_callbacks_by_type(
_cb_list, AdaptiveRouterPostCallHook
)
for entry in self.model_list or []:
lp = (
entry.get("litellm_params")
if isinstance(entry, dict)
else entry.litellm_params
)
lp_model = (
(lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None
)
if not (lp_model and lp_model.startswith("auto_router/adaptive_router")):
continue
model_name = (
entry.get("model_name") if isinstance(entry, dict) else entry.model_name
)
if not model_name or not lp:
continue
if model_name in self.adaptive_routers:
continue
deployment = Deployment(
model_name=model_name,
litellm_params=(
lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)
),
model_info=(
entry.get("model_info")
if isinstance(entry, dict)
else entry.model_info
),
)
self.init_adaptive_router_deployment(deployment=deployment)
def init_adaptive_router_deployment(self, deployment: Deployment) -> None:
"""
Build an AdaptiveRouter instance for this deployment and register its
post-call hook. Multiple adaptive routers can coexist on a single Router,
keyed by `deployment.model_name`.
`model_to_prefs` and `model_to_cost` are derived from the OTHER models
already registered in `self.model_list` whose `model_name` appears in
`available_models`. Models not yet registered fall back to defaults.
"""
# Local import: AdaptiveRouter -> hooks -> classifier all import litellm
# internals which transitively import this module. (AGENTS.md exception clause.)
from litellm.router_strategy.adaptive_router.adaptive_router import (
AdaptiveRouter,
)
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
from litellm.types.router import (
AdaptiveRouterConfig,
AdaptiveRouterPreferences,
)
raw_config = deployment.litellm_params.adaptive_router_config
if raw_config is None:
raise ValueError(
"adaptive_router_config is required for adaptive-router deployments."
)
config = AdaptiveRouterConfig(**raw_config)
model_to_prefs: Dict[str, AdaptiveRouterPreferences] = {}
model_to_cost: Dict[str, float] = {}
# O(k) via the name→indices map: only touch deployments whose name
# is listed in `available_models`, instead of scanning model_list.
for name in config.available_models:
indices = self.model_name_to_deployment_indices.get(name, [])
if not indices:
continue
d = (self.model_list or [])[indices[0]]
mi = d.get("model_info") if isinstance(d, dict) else d.model_info
mi_dict: Dict[str, Any] = (
mi if isinstance(mi, dict) else (mi.model_dump() if mi else {})
)
prefs_raw = mi_dict.get("adaptive_router_preferences")
if prefs_raw is not None:
model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
# `input_cost_per_token` is a LiteLLM_Params field per types/router.py.
lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params
lp_dict: Dict[str, Any] = (
lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
)
cost = lp_dict.get("input_cost_per_token")
if cost is not None:
model_to_cost[name] = float(cost)
if deployment.model_name in self.adaptive_routers:
raise ValueError(
f"Adaptive-router deployment {deployment.model_name} already exists. "
"Please use a different model name."
)
adaptive_router = AdaptiveRouter(
router_name=deployment.model_name,
config=config,
model_to_prefs=model_to_prefs,
model_to_cost=model_to_cost,
)
self.adaptive_routers[deployment.model_name] = adaptive_router
litellm.logging_callback_manager.add_litellm_callback(
AdaptiveRouterPostCallHook(adaptive_router=adaptive_router)
)
verbose_router_logger.info(
"AdaptiveRouter[%s] initialized with %d models",
deployment.model_name,
len(config.available_models),
)
def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
"""
Check if the deployment is a quality-router deployment.
Returns True if the litellm_params model starts with "auto_router/quality_router".
"""
if litellm_params.model.startswith("auto_router/quality_router"):
return True
return False
def init_quality_router_deployment(self, deployment: Deployment):
"""
Initialize the quality-router deployment.
Resolves the default model from either `quality_router_default_model` or
`quality_router_config["default_model"]`, then instantiates the
QualityRouter and stores it in `self.quality_routers`.
"""
# Import here to mirror the AutoRouter / ComplexityRouter init pattern
# and avoid circular imports.
from litellm.router_strategy.quality_router.quality_router import (
QualityRouter,
)
quality_router_config: Optional[dict] = (
deployment.litellm_params.quality_router_config
)
default_model: Optional[str] = (
deployment.litellm_params.quality_router_default_model
)
if default_model is None and quality_router_config:
default_model = quality_router_config.get("default_model")
if default_model is None:
raise ValueError(
"quality_router_default_model is required for quality-router deployments, "
"or set default_model in quality_router_config. Please configure it in the litellm_params"
)
quality_router: QualityRouter = QualityRouter(
model_name=deployment.model_name,
default_model=default_model,
litellm_router_instance=self,
quality_router_config=quality_router_config,
)
if deployment.model_name in self.quality_routers:
raise ValueError(
f"Quality-router deployment {deployment.model_name} already exists. Please use a different model name."
)
self.quality_routers[deployment.model_name] = quality_router
def deployment_is_active_for_environment(self, deployment: Deployment) -> bool:
"""
Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments
@ -6966,6 +7176,11 @@ class Router:
self.model_id_to_deployment_index_map = {} # Reset the index
self.model_name_to_deployment_indices = {} # Reset the model_name index
self.team_model_to_deployment_indices = {} # Reset the team_model index
# Reset per-strategy router registries so hot-reload doesn't leave
# stale routers pointing at the old model_list.
self.quality_routers = {}
self.complexity_routers = {}
self.auto_routers = {}
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
@ -7013,6 +7228,10 @@ class Router:
# Note: model_name_to_deployment_indices is already built incrementally
# by _create_deployment -> _add_model_to_list_and_index_map
# Deferred: build the AdaptiveRouter strategy now that all underlying
# deployments have been registered.
self._finalize_adaptive_router_if_configured()
def _add_deployment(self, deployment: Deployment) -> Deployment:
import os
@ -7140,6 +7359,16 @@ class Router:
):
self.init_complexity_router_deployment(deployment=deployment)
# NOTE: adaptive-router deployments are deferred to the end of
# set_model_list() because their init needs visibility into the OTHER
# deployments listed in `available_models` (which may not yet have
# been processed when this one is created).
#########################################################
# Check if this is a quality-router deployment
#########################################################
if self._is_quality_router_deployment(litellm_params=deployment.litellm_params):
self.init_quality_router_deployment(deployment=deployment)
return deployment
def _initialize_deployment_for_pass_through(
@ -8143,7 +8372,10 @@ class Router:
return returned_dict
async def set_response_headers(
self, response: Any, model_group: Optional[str] = None
self,
response: Any,
model_group: Optional[str] = None,
request_kwargs: Optional[dict] = None,
) -> Any:
"""
Add the most accurate rate limit headers for a given model response.
@ -8164,6 +8396,45 @@ class Router:
additional_headers = response._hidden_params["additional_headers"] # type: ignore
# Lift QualityRouter routing decision into response headers for
# transparency. The decision is stashed in request_kwargs.metadata
# by QualityRouter.async_pre_routing_hook.
metadata = (
(request_kwargs.get("metadata") or {})
if isinstance(request_kwargs, dict)
else {}
)
decision = (
metadata.get("quality_router_decision")
if isinstance(metadata, dict)
else None
)
if isinstance(decision, dict):
# Only emit headers for fields that have a meaningful value.
# `complexity_tier` and `matched_keyword` are mutually exclusive
# (the keyword path short-circuits classification), so each
# request emits one or the other but not both.
if decision.get("routed_model") is not None:
additional_headers["x-litellm-quality-router-model"] = str(
decision["routed_model"]
)
if decision.get("quality_tier") is not None:
additional_headers["x-litellm-quality-router-tier"] = str(
decision["quality_tier"]
)
if decision.get("routed_via") is not None:
additional_headers["x-litellm-quality-router-via"] = str(
decision["routed_via"]
)
if decision.get("matched_keyword") is not None:
additional_headers["x-litellm-quality-router-keyword"] = str(
decision["matched_keyword"]
)
if decision.get("complexity_tier") is not None:
additional_headers["x-litellm-quality-router-complexity"] = str(
decision["complexity_tier"]
)
if (
"x-ratelimit-remaining-tokens" not in additional_headers
and "x-ratelimit-remaining-requests" not in additional_headers
@ -8708,8 +8979,6 @@ class Router:
and self.routing_strategy == "latency-based-routing"
):
_settings_to_return[var] = self.lowestlatency_logger.routing_args.json()
elif var == "routing_strategy_args":
_settings_to_return[var] = None
return _settings_to_return
def update_settings(self, **kwargs):
@ -9620,7 +9889,7 @@ class Router:
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional[PreRoutingHookResponse]:
@ -9653,6 +9922,31 @@ class Router:
specific_deployment=specific_deployment,
)
#########################################################
# Check if an adaptive-router should be used
#########################################################
adaptive_router = self.adaptive_routers.get(model)
if adaptive_router is not None:
return await adaptive_router.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
#########################################################
# Check if any quality-router should be used
#########################################################
if model in self.quality_routers:
return await self.quality_routers[model].async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
return None
def get_available_deployment(

95
litellm/router_strategy/adaptive_router/README.md Normal file
View file

@ -0,0 +1,95 @@
# Adaptive Router (v0)
A request-type-aware routing strategy. For each incoming request, classify the
prompt into one of seven `RequestType` buckets (code generation, writing,
analytical reasoning, …), then Thompson-sample a Beta(α, β) bandit posterior
per `(request_type, model)` cell to pick the best model. Quality estimates are
combined with a normalized cost score via a weighted linear sum.
A post-call hook reads the response and runs lightweight regex + tool-call
detectors (see `signals.py`) to award per-turn credit/blame to the model that
served the turn. Updates are batched in-memory and flushed to Postgres every
~10s by a background task in `proxy_server.py`.
## Config example
```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
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: ["general", "factual_lookup"]
- model_name: smart-router
litellm_params:
model: auto_router/adaptive_router
adaptive_router_default_model: gpt-4o-mini
adaptive_router_config:
available_models: ["gpt-4o", "gpt-4o-mini"]
weights:
quality: 0.7
cost: 0.3
```
Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key
`min_quality_tier: 3`) to force selection from tier-3-or-higher models only.
## Behavior summary
- **Cold start.** Each `(request_type, model)` cell starts with a
Beta prior whose mean = `BASE_TIER_WEIGHT[tier] (+ STRENGTH_BONUS if declared)`
and total mass = `COLD_START_MASS` (10). About ten real observations move it
meaningfully.
- **Per-request decision.** Sample once per eligible model, score with
`quality_weight·sample + cost_weight·normalized_cost`, pick the argmax.
Routing is stateless per-turn — no sticky lookup. Each call resamples.
- **Owner-cache attribution.** Post-call, the conversation's first picked
model claims an "owner slot" for `OWNER_CACHE_TTL_SECONDS` (24h). Later
turns of the same conversation only fire bandit/state updates if the
same model handled them — mismatches are dropped (no attribution) and
counted in `skipped_updates_total`. Conversation identity is the
client-supplied `litellm_session_id` if present, otherwise a sha256 over
caller identity (api key hash, team, user, end-user) + the first message.
- **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation,
disengagement, failure → +β` (each). `loop → +0.5β`. `exhaustion → 0`
(uptime, not quality). Skipped if conversation has fewer than
`SIGNAL_GATE_MIN_MESSAGES` messages.
- **Persistence.** Bandit cells: aggregated deltas, eventually consistent.
Session rows: last-write-wins snapshots.
## Known v0 limitations
- **Latency is not in the score.** Quality + cost only. A pathologically slow
model can still be picked.
- **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped.
No rescaling — drift is a v1 concern.
- **24h owner-cache TTL.** No explicit eviction below TTL. The in-memory map
can grow if traffic patterns produce many one-shot sessions.
- **Owner-recovery skew.** If model A "owns" a conversation but is then
dethroned in the bandit, later turns served by model B are dropped — so
bandit updates for that conversation flatline until A's TTL expires.
Tracked via `skipped_updates_total`.
- **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity,
no exemplar storage. Signals are best-effort and biased toward English.
- **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments
on the same `litellm.Router` raise at init.
- **Bandit-delta mapping is unvalidated.** `_compute_bandit_delta` is a v0
guess; expect to retune after the first ~1000 sessions of real traffic.
- **`request_type` is classified per turn from the latest user message.** For
non-GENERAL turns, the current-turn type is used for bandit attribution (so
genuine mid-session topic shifts update the correct cell). For GENERAL turns
("thanks!", "ok", "sounds good"), attribution falls back to the session's
original type to avoid misattributing closing pleasantries.

6
litellm/router_strategy/adaptive_router/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""Adaptive router strategy. See README.md for design overview."""
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"]

454
litellm/router_strategy/adaptive_router/adaptive_router.py Normal file
View file

@ -0,0 +1,454 @@
"""
Main adaptive router strategy. See README.md for design overview.
One AdaptiveRouter instance per router_name. Holds in-memory caches:
- _cells: Beta(alpha, beta) bandit posteriors per (request_type, model)
- _owner_cache: session_key -> (owner_model, expires_at) — the first model
picked for a conversation owns its bandit-update slot
- _session_states: (session_key, model) -> SessionState for incremental signal updates
Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist
state and session snapshots back to Postgres.
Routing is stateless per-turn (Thompson sample fresh on every call). The
owner cache is consulted only at post-call time to decide whether a turn's
signals should fire a bandit update — turns served by a different model than
the conversation's owner are skipped to avoid cross-model misattribution.
"""
from __future__ import annotations
import asyncio
import time
from dataclasses import asdict
from typing import Any, Dict, List, Optional, Tuple, Union, cast
from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_last_user_message,
)
from litellm.router_strategy.adaptive_router.bandit import (
BanditCell,
apply_delta,
initial_cell,
pick_best,
)
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
from litellm.router_strategy.adaptive_router.config import (
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
MIN_QUALITY_TIER_HEADER,
MIN_QUALITY_TIER_METADATA_KEY,
OWNER_CACHE_TTL_SECONDS,
)
from litellm.router_strategy.adaptive_router.signals import (
SessionState,
SignalDelta,
Turn,
apply_turn,
)
from litellm.router_strategy.adaptive_router.update_queue import (
AdaptiveRouterUpdateQueue,
)
# Sweep session-state cache when it exceeds this many live entries. Expired
# entries are dropped in bulk; amortizes to O(1) per insert.
_SESSION_STATE_SWEEP_THRESHOLD: int = 1024
# Same pattern for the owner cache.
_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import (
AdaptiveRouterConfig,
AdaptiveRouterPreferences,
PreRoutingHookResponse,
RequestType,
)
def _default_prefs() -> AdaptiveRouterPreferences:
"""Tier-2 prior with no declared strengths; used when a model omits prefs."""
return AdaptiveRouterPreferences(quality_tier=2, strengths=[])
class AdaptiveRouter:
"""One instance per router_name. Holds in-memory caches + the update queue."""
def __init__(
self,
router_name: str,
config: AdaptiveRouterConfig,
model_to_prefs: Dict[str, AdaptiveRouterPreferences],
model_to_cost: Dict[str, float],
) -> None:
self.router_name = router_name
self.config = config
self.model_to_prefs = model_to_prefs
self.model_to_cost = model_to_cost
self.queue = AdaptiveRouterUpdateQueue()
self._cells: Dict[Tuple[RequestType, str], BanditCell] = {}
self._owner_cache: Dict[str, Tuple[str, float]] = {}
self._session_states: Dict[Tuple[str, str], SessionState] = {}
# Parallel expiry map for _session_states, same TTL as _owner_cache.
# Evicted opportunistically in `get_or_create_session_state`.
self._session_states_expiry: Dict[Tuple[str, str], float] = {}
self._skipped_updates_total: int = 0
# Set to True once the proxy flusher has loaded persisted priors from
# Postgres. Checked to support lazy-load on hot-reloaded routers.
self._state_loaded: bool = False
self._lock = asyncio.Lock()
self._init_cold_start_cells()
# ---- Cold-start ------------------------------------------------------
def _init_cold_start_cells(self) -> None:
"""Populate _cells with cold-start priors for every (rt, model) combination."""
for rt in RequestType:
for model in self.config.available_models:
prefs = self.model_to_prefs.get(model) or _default_prefs()
self._cells[(rt, model)] = initial_cell(prefs, rt)
async def load_state_from_db(self, prisma_client: Any) -> None:
"""Override cold-start cells with persisted state. Called once at startup."""
if prisma_client is None:
return
try:
rows = await prisma_client.db.litellm_adaptiverouterstate.find_many(
where={"router_name": self.router_name}
)
loaded = 0
for row in rows:
try:
rt = RequestType(row.request_type)
except ValueError:
# Unknown taxonomy entry from an older/newer version. Skip.
continue
if row.model_name not in self.config.available_models:
continue
self._cells[(rt, row.model_name)] = BanditCell(
alpha=row.alpha, beta=row.beta
)
loaded += 1
verbose_router_logger.info(
"AdaptiveRouter[%s]: loaded %d cells from DB",
self.router_name,
loaded,
)
except Exception as e:
verbose_router_logger.exception(
"AdaptiveRouter[%s]: failed to load state from DB: %s",
self.router_name,
e,
)
# ---- Pre-routing hook ------------------------------------------------
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict[str, Any],
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional[PreRoutingHookResponse]:
"""
Plugin entry point invoked by `Router.async_pre_routing_hook` when the
inbound `model` matches this adaptive router's `router_name`.
Classifies the last user message, picks a logical model via the bandit,
and stashes the chosen model on `request_kwargs["metadata"]` so the
post-call hook can surface it as a response header.
Routing is stateless per-turn: every call Thompson-samples fresh,
regardless of any prior pick for the same session. Cross-turn
attribution is enforced post-call via the owner cache (see
`claim_or_check_owner`).
"""
user_text = (
get_last_user_message(cast(List[AllMessageValues], messages or [])) or ""
)
request_type = classify_prompt(user_text)
min_quality_tier = self._extract_min_quality_tier(request_kwargs)
chosen_model = await self.pick_model(
request_type=request_type, min_quality_tier=min_quality_tier
)
verbose_router_logger.debug(
"AdaptiveRouter[%s]: classified=%s -> chose %s",
self.router_name,
request_type.value,
chosen_model,
)
# Relay the chosen logical model to the post-call hook, which surfaces
# it as the `x-litellm-adaptive-router-model` response header. We use
# `metadata` (not a top-level kwarg) so the value doesn't leak into
# `litellm.acompletion(**input_kwargs)`.
kwargs_metadata = request_kwargs.setdefault("metadata", {})
if isinstance(kwargs_metadata, dict):
kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = chosen_model
return PreRoutingHookResponse(model=chosen_model, messages=messages)
# ---- Pick model ------------------------------------------------------
async def pick_model(
self,
request_type: RequestType,
min_quality_tier: Optional[int] = None,
) -> str:
"""Thompson-sample across eligible models. Stateless per-turn."""
eligible = self._eligible_models(min_quality_tier)
if not eligible:
raise ValueError(
f"AdaptiveRouter[{self.router_name}]: no models meet "
f"min_quality_tier={min_quality_tier}"
)
cells = {m: self._cells[(request_type, m)] for m in eligible}
costs = {m: self.model_to_cost.get(m, 0.0) for m in eligible}
return pick_best(
cells,
costs,
quality_weight=self.config.weights.quality,
cost_weight=self.config.weights.cost,
)
def claim_or_check_owner(self, session_key: str, current_model: str) -> bool:
"""Resolve attribution for a turn under stateless routing.
Returns True iff this turn should fire a bandit/state update. The
first call for a `session_key` claims ownership for `current_model`
and returns True. Subsequent calls return True only if the owner is
still live AND matches `current_model`. Mismatches (a different
model handled this turn) and expired owners both increment
`_skipped_updates_total` and return False — no attribution.
"""
now = time.time()
existing = self._owner_cache.get(session_key)
if existing is not None and existing[1] > now:
owner_model, _ = existing
if owner_model == current_model:
return True
self._skipped_updates_total += 1
return False
# Opportunistic bulk sweep — sessions that never come back would
# otherwise pile up here forever. Same threshold pattern as the
# session-state cache.
if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD:
self._evict_expired_owner_cache(now)
# No live owner -> claim for current_model.
self._owner_cache[session_key] = (
current_model,
now + OWNER_CACHE_TTL_SECONDS,
)
return True
def _evict_expired_owner_cache(self, now: float) -> None:
expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now]
for k in expired:
self._owner_cache.pop(k, None)
async def get_state_snapshot(self) -> Dict[str, Any]:
"""In-memory snapshot for the introspection endpoint. Cheap; no DB hit."""
cells = []
for (rt, model), cell in sorted(
self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])
):
total = cell.alpha + cell.beta
cells.append(
{
"request_type": rt.value,
"model": model,
"alpha": cell.alpha,
"beta": cell.beta,
# Net observations that have moved the posterior, excluding
# the cold-start prior mass. `alpha + beta` would show the
# initial COLD_START_MASS (e.g. 10) before any real traffic
# arrives, which confuses operators reading the endpoint.
"samples": cell.total_samples,
"quality_mean": cell.alpha / total if total > 0 else 0.0,
}
)
queue = await self.queue.queue_size()
now = time.time()
owner_cache_live = sum(1 for _, exp in self._owner_cache.values() if exp > now)
return {
"router_name": self.router_name,
"available_models": list(self.config.available_models),
"weights": {
"quality": self.config.weights.quality,
"cost": self.config.weights.cost,
},
"model_costs": dict(self.model_to_cost),
"cells": cells,
"owner_cache_live": owner_cache_live,
"skipped_updates_total": self._skipped_updates_total,
"queue": queue,
}
@staticmethod
def _extract_min_quality_tier(
request_kwargs: Dict[str, Any],
) -> Optional[int]:
"""Pull `min_quality_tier` from request headers or metadata.
Precedence: headers (`x-litellm-min-quality-tier`) over metadata
(`min_quality_tier`). Headers arrive lowercased from the proxy but we
lookup case-insensitively to be safe. Unparseable values are ignored
(treated as "not set") rather than raising — a bad header shouldn't
fail the request.
"""
headers = request_kwargs.get("headers") or {}
if isinstance(headers, dict):
for k, v in headers.items():
if isinstance(k, str) and k.lower() == MIN_QUALITY_TIER_HEADER:
try:
return int(v)
except (TypeError, ValueError):
return None
metadata = request_kwargs.get("metadata") or {}
if isinstance(metadata, dict):
raw = metadata.get(MIN_QUALITY_TIER_METADATA_KEY)
if raw is not None:
try:
return int(raw)
except (TypeError, ValueError):
return None
return None
def _eligible_models(self, min_quality_tier: Optional[int]) -> List[str]:
if min_quality_tier is None:
return list(self.config.available_models)
return [
m
for m in self.config.available_models
if (self.model_to_prefs.get(m) or _default_prefs()).quality_tier
>= min_quality_tier
]
# ---- Session state ---------------------------------------------------
def get_or_create_session_state(
self,
session_id: str,
model_name: str,
request_type: RequestType,
) -> SessionState:
key = (session_id, model_name)
now = time.time()
# Opportunistic bulk sweep when the cache grows past the threshold.
# Cheap relative to the alternative of a bounded LRU — conversations
# naturally become inactive within OWNER_CACHE_TTL_SECONDS.
if len(self._session_states) >= _SESSION_STATE_SWEEP_THRESHOLD:
self._evict_expired_session_states(now)
state = self._session_states.get(key)
if state is None:
state = SessionState(
session_id=session_id,
router_name=self.router_name,
model_name=model_name,
classified_type=request_type.value,
)
self._session_states[key] = state
self._session_states_expiry[key] = now + OWNER_CACHE_TTL_SECONDS
return state
def _evict_expired_session_states(self, now: float) -> None:
"""Drop session states whose TTL has passed. O(n) but amortized O(1)
per insert thanks to `_SESSION_STATE_SWEEP_THRESHOLD`."""
expired = [k for k, exp in self._session_states_expiry.items() if exp <= now]
for k in expired:
self._session_states.pop(k, None)
self._session_states_expiry.pop(k, None)
async def record_turn(
self,
session_id: str,
model_name: str,
request_type: RequestType,
turn: Turn,
) -> SignalDelta:
"""Apply one turn, push session snapshot + bandit deltas to the queue."""
state = self.get_or_create_session_state(session_id, model_name, request_type)
delta = apply_turn(state, turn)
verbose_router_logger.debug(
"AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta
)
# Strip the raw conversation content before persisting. The
# last_user/assistant_content and tool_call_history fields are only
# needed in-memory for the next turn's incremental signal detection;
# writing user prompts and tool payloads to the DB would store PII
# for every adaptive-router conversation. Counts + bookkeeping is
# all the persisted row needs.
snapshot = asdict(state)
for sensitive in (
"last_user_content",
"last_assistant_content",
"tool_call_history",
"pending_tool_calls",
):
snapshot.pop(sensitive, None)
await self.queue.add_session_state(
session_id, self.router_name, model_name, snapshot
)
d_alpha, d_beta = self._compute_bandit_delta(delta)
verbose_router_logger.debug(
"AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f",
self.router_name,
d_alpha,
d_beta,
)
if d_alpha != 0 or d_beta != 0:
# For non-GENERAL turns, attribute to the current-turn classification
# so genuine mid-session topic shifts (e.g. code → math) update the
# correct cell. For GENERAL turns ("thanks!", "ok", "sounds good"), fall
# back to the session's original type so closing pleasantries don't
# misattribute the reward.
attribution_type = (
request_type
if request_type != RequestType.GENERAL
else RequestType(state.classified_type)
)
cell_key = (attribution_type, model_name)
self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta)
await self.queue.add_state_delta(
self.router_name,
attribution_type.value,
model_name,
d_alpha,
d_beta,
)
return delta
@staticmethod
def _compute_bandit_delta(delta: SignalDelta) -> Tuple[float, float]:
"""
Translate per-turn signal deltas into bandit-cell deltas.
v0 mapping (UNVALIDATED — D6):
- satisfaction -> +1 alpha
- misalignment, stagnation,
disengagement, failure -> +1 beta each
- loop -> +0.5 beta (weak; could be model OR user)
- exhaustion -> 0 (uptime issue, tracked separately later)
"""
d_alpha = float(delta.satisfaction)
d_beta = (
float(
delta.misalignment
+ delta.stagnation
+ delta.disengagement
+ delta.failure
)
+ 0.5 * delta.loop
)
return d_alpha, d_beta

142
litellm/router_strategy/adaptive_router/bandit.py Normal file
View file

@ -0,0 +1,142 @@
"""
Thompson sampling and prior initialization for the adaptive router bandit.
Each (router, request_type, model) cell is a Beta(alpha, beta) posterior.
- alpha = pseudo-successes
- beta = pseudo-failures
- mean = alpha / (alpha + beta)
- total samples = alpha + beta - COLD_START_MASS (informative prior, not data)
Hot path: thompson_sample() — pure function, no I/O.
"""
import random
from dataclasses import dataclass
from typing import Dict, List, Optional
from litellm.router_strategy.adaptive_router.config import (
BASE_TIER_WEIGHT,
COLD_START_MASS,
DEFAULT_COST_WEIGHT,
DEFAULT_QUALITY_WEIGHT,
SAMPLE_CAP,
STRENGTH_BONUS,
)
from litellm.types.router import AdaptiveRouterPreferences, RequestType
@dataclass(frozen=True)
class BanditCell:
"""Posterior state for a single (router, request_type, model) cell."""
alpha: float
beta: float
@property
def mean(self) -> float:
total = self.alpha + self.beta
return self.alpha / total if total > 0 else 0.5
@property
def total_samples(self) -> int:
return max(0, int(self.alpha + self.beta - COLD_START_MASS))
def initial_cell(
prefs: AdaptiveRouterPreferences, request_type: RequestType
) -> BanditCell:
"""
Cold-start prior for a (model, request_type) cell.
mean = base_tier_weight[tier] + (STRENGTH_BONUS if request_type in strengths else 0)
capped at 0.95 to avoid an over-confident prior.
Total mass = COLD_START_MASS so that ~10 real observations can move it noticeably.
"""
if prefs.quality_tier not in BASE_TIER_WEIGHT:
valid = sorted(BASE_TIER_WEIGHT)
raise ValueError(
f"quality_tier={prefs.quality_tier} is not supported; "
f"valid tiers are {valid}"
)
base = BASE_TIER_WEIGHT[prefs.quality_tier]
bonus = STRENGTH_BONUS if request_type in prefs.strengths else 0.0
mean = min(0.95, base + bonus)
alpha = mean * COLD_START_MASS
beta = (1.0 - mean) * COLD_START_MASS
return BanditCell(alpha=alpha, beta=beta)
def apply_delta(cell: BanditCell, delta_alpha: float, delta_beta: float) -> BanditCell:
"""
Apply a learning update to a cell, enforcing the sample cap.
SAMPLE_CAP is a HARD cap on (alpha + beta). When the cap would be exceeded,
we drop the update. (D5: hard cap, no rescaling — keep v0 simple.)
"""
new_alpha = cell.alpha + delta_alpha
new_beta = cell.beta + delta_beta
if new_alpha + new_beta > SAMPLE_CAP:
return cell
return BanditCell(alpha=new_alpha, beta=new_beta)
def thompson_sample(cell: BanditCell, rng: Optional[random.Random] = None) -> float:
"""Draw a sample from Beta(alpha, beta). Returns a quality estimate in [0, 1]."""
r = rng if rng is not None else random
return r.betavariate(cell.alpha, cell.beta)
def normalized_cost(model_cost: float, all_costs: List[float]) -> float:
"""
Map a raw $/1k-token cost into [0, 1] where 0 = most expensive, 1 = cheapest.
Returns 0.5 when there's no spread.
"""
if not all_costs:
return 0.5
lo, hi = min(all_costs), max(all_costs)
if hi == lo:
return 0.5
return 1.0 - ((model_cost - lo) / (hi - lo))
def score(
quality_sample: float,
model_cost: float,
all_costs: List[float],
quality_weight: float = DEFAULT_QUALITY_WEIGHT,
cost_weight: float = DEFAULT_COST_WEIGHT,
) -> float:
"""
Multi-objective score. V0 is a weighted linear sum of (quality, normalized_cost).
Higher is better. Both inputs are in [0, 1].
"""
cost_score = normalized_cost(model_cost, all_costs)
return quality_weight * quality_sample + cost_weight * cost_score
def pick_best(
cells: Dict[str, BanditCell],
model_costs: Dict[str, float],
quality_weight: float = DEFAULT_QUALITY_WEIGHT,
cost_weight: float = DEFAULT_COST_WEIGHT,
rng: Optional[random.Random] = None,
) -> str:
"""
Sample once per model, score each, return the model with highest score.
cells: {model_name: BanditCell}
model_costs: {model_name: $/1k tokens}
"""
if not cells:
raise ValueError("pick_best called with no models")
all_costs = list(model_costs.values())
best_model: Optional[str] = None
best_score = float("-inf")
for model, cell in cells.items():
q = thompson_sample(cell, rng=rng)
s = score(q, model_costs[model], all_costs, quality_weight, cost_weight)
if s > best_score:
best_score = s
best_model = model
assert best_model is not None
return best_model

140
litellm/router_strategy/adaptive_router/classifier.py Normal file
View file

@ -0,0 +1,140 @@
"""
Rule-based classifier mapping a user prompt to a RequestType.
V0 design choice: deterministic regex over the FIRST user message in a session.
Result is cached per session (caller's responsibility, not ours).
Order matters: we check more specific types first, falling back to GENERAL.
"""
import re
from typing import List, Pattern, Tuple
from litellm.types.router import RequestType
_RULES: List[Tuple[Pattern[str], RequestType]] = [
(
re.compile(
r"\b(write|create|generate|implement|build)\s+(?:a |an |the |me )?(?:python|javascript|typescript|java|rust|go|c\+\+|sql|bash|shell)\b",
re.IGNORECASE,
),
RequestType.CODE_GENERATION,
),
(
re.compile(
r"\b(write|create|implement|build)\b(?:\s+\w+){0,4}?\s+(function|class|method|script|program|api|endpoint|microservice)\b",
re.IGNORECASE,
),
RequestType.CODE_GENERATION,
),
(
re.compile(
r"\b(explain|describe|understand|walk me through|what does)\b.*\b(code|function|method|class|algorithm|snippet)\b",
re.IGNORECASE,
),
RequestType.CODE_UNDERSTANDING,
),
(
re.compile(
r"\b(debug|fix|why (?:is|does|isn't)|what.s wrong|trace)\b.*\b(error|bug|exception|stacktrace|stack trace|traceback)\b",
re.IGNORECASE,
),
RequestType.CODE_UNDERSTANDING,
),
(
re.compile(
r"\b(review|critique)\s+(?:this |my |the )?(?:code|pr|pull request|diff|patch)\b",
re.IGNORECASE,
),
RequestType.CODE_UNDERSTANDING,
),
(
re.compile(
r"\b(design|architect|plan|architecture)\b.*\b(system|service|api|database|schema|module|microservice)\b",
re.IGNORECASE,
),
RequestType.TECHNICAL_DESIGN,
),
(
re.compile(
r"\b(should i (?:use|choose|pick)|tradeoffs? between|compare)\b.*\b(library|framework|language|database|protocol|postgres|postgresql|mongodb|dynamodb|mysql|redis|kafka|sql|nosql)\b",
re.IGNORECASE,
),
RequestType.TECHNICAL_DESIGN,
),
(
re.compile(
r"\bhow (?:should|do) i (?:design|structure|organize|model)\b",
re.IGNORECASE,
),
RequestType.TECHNICAL_DESIGN,
),
(
re.compile(
r"\b(solve|compute|calculate|prove|derive)\b.*\b(equation|integral|derivative|theorem|proof|problem)\b",
re.IGNORECASE,
),
RequestType.ANALYTICAL_REASONING,
),
(
re.compile(r"\b(if .+ then|given .+ find|suppose|assume)\b", re.IGNORECASE),
RequestType.ANALYTICAL_REASONING,
),
(
re.compile(
r"\b(probability|statistics|combinatorics|optimization problem)\b",
re.IGNORECASE,
),
RequestType.ANALYTICAL_REASONING,
),
(
re.compile(
r"\b(write|draft|compose|rewrite|edit|proofread|polish)\b.*\b(email|essay|blog|post|article|letter|memo|copy|paragraph|sentence)\b",
re.IGNORECASE,
),
RequestType.WRITING,
),
(
re.compile(
r"\b(make (?:this|it)|help me)\s+(?:more |less )?(?:concise|formal|casual|professional|persuasive)\b",
re.IGNORECASE,
),
RequestType.WRITING,
),
(
re.compile(
r"^\s*(who|what|when|where|which)\s+(?:is|was|were|are)\b", re.IGNORECASE
),
RequestType.FACTUAL_LOOKUP,
),
(
re.compile(r"^\s*(define|definition of|meaning of)\b", re.IGNORECASE),
RequestType.FACTUAL_LOOKUP,
),
(
re.compile(
r"^\s*how (?:do you spell|to spell|many .* are there|tall is)\b",
re.IGNORECASE,
),
RequestType.FACTUAL_LOOKUP,
),
]
def classify_prompt(text: str) -> RequestType:
"""
Classify a single user prompt.
Falls back to GENERAL when no rule matches. Empty/whitespace-only also
returns GENERAL.
"""
if not text or not text.strip():
return RequestType.GENERAL
truncated = text[:2000]
for pattern, request_type in _RULES:
if pattern.search(truncated):
return request_type
return RequestType.GENERAL

54
litellm/router_strategy/adaptive_router/config.py Normal file
View file

@ -0,0 +1,54 @@
"""
Configuration constants for the adaptive_router strategy.
All magic numbers are first-pass guesses (D3-D6 in the handoff plan).
Expect to retune after first 1000 sessions of real traffic.
"""
from typing import Dict
from litellm.types.router import RequestType # re-export for convenience # noqa: F401
# D3 — Score weights (default; user-overridable via AdaptiveRouterConfig.weights)
DEFAULT_QUALITY_WEIGHT: float = 0.7 # UNVALIDATED — calibrated against [0] sessions
DEFAULT_COST_WEIGHT: float = 0.3 # UNVALIDATED — calibrated against [0] sessions
# D4 — Cold-start prior: (alpha + beta) total mass = COLD_START_MASS
# Mean of Beta = base_tier_weight + (strength_bonus if declared)
BASE_TIER_WEIGHT: Dict[int, float] = {1: 0.3, 2: 0.5, 3: 0.7} # UNVALIDATED
STRENGTH_BONUS: float = 0.3 # UNVALIDATED
COLD_START_MASS: float = 10.0
# D5 — Sample cap. Hard cap, no rescaling (drift handling is v1).
SAMPLE_CAP: int = 200
# D6 — Clean-trace credit: minimum turns before α += 1 can fire.
MIN_TURNS_FOR_CLEAN_CREDIT: int = 3
# D2 — Owner-cache TTL (seconds). 24h.
# A conversation's first-picked model "owns" the bandit-update slot for
# this long. Subsequent turns of the same conversation only contribute a
# bandit/state update when the same model is re-sampled.
OWNER_CACHE_TTL_SECONDS: int = 24 * 3600
# Below this many messages we skip post-call signal recording. Most signals
# (misalignment, stagnation, satisfaction-in-response-to-prior-turn) need at
# least one full prior exchange to be meaningful.
SIGNAL_GATE_MIN_MESSAGES: int = 4
# Detector thresholds (from Plano/Chen 2026 paper).
MISALIGNMENT_JACCARD_THRESHOLD: float = 0.45
STAGNATION_JACCARD_NEAR_DUP: float = 0.50
LOOP_REPEAT_THRESHOLD: int = 3
TOOL_CALL_HISTORY_MAX: int = 20
# D1 — Caller filter for min quality tier.
MIN_QUALITY_TIER_HEADER: str = "x-litellm-min-quality-tier"
MIN_QUALITY_TIER_METADATA_KEY: str = "min_quality_tier"
# Pre-routing -> post-call relay: the chosen logical model is stashed on
# request_kwargs["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] by the
# pre-routing hook, then read by the post-call hook to surface as the
# ADAPTIVE_ROUTER_RESPONSE_HEADER response header.
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY: str = "adaptive_router_chosen_model"
ADAPTIVE_ROUTER_RESPONSE_HEADER: str = "x-litellm-adaptive-router-model"

278
litellm/router_strategy/adaptive_router/hooks.py Normal file
View file

@ -0,0 +1,278 @@
"""
Post-call hook for the adaptive router.
On each successful or failed completion, build a Turn from the request/response
and push it through `AdaptiveRouter.record_turn`. The router then updates the
in-memory bandit cell + session state and queues writes for the proxy flusher.
All work happens after the response has been returned to the caller. Any
exception is swallowed — signal recording must never break a request.
"""
from __future__ import annotations
import hashlib
import json
from typing import Any, Dict, List, Optional
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
from litellm.router_strategy.adaptive_router.config import (
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
ADAPTIVE_ROUTER_RESPONSE_HEADER,
SIGNAL_GATE_MIN_MESSAGES,
)
from litellm.router_strategy.adaptive_router.signals import Turn
# Identity fields hashed into a derived session key so the same conversation
# from the same caller produces a stable key, while different keys/teams/users
# stay segregated even if they happen to send identical first messages.
_IDENTITY_FIELDS = (
"user_api_key_hash",
"user_api_key_team_id",
"user_api_key_user_id",
"user_api_key_end_user_id",
)
def _resolve_session_key(kwargs: Dict[str, Any]) -> Optional[str]:
"""Pick a stable per-conversation key for owner-cache attribution.
Order:
1. Honor a client-supplied session id (`litellm_session_id` on either
`litellm_params` or `litellm_params.metadata`, or `session_id` on
metadata) — backward compat for callers already wired up.
2. Otherwise derive a sha256 over (identity fields, first
SIGNAL_GATE_MIN_MESSAGES messages) so the key is stable across turns
and only materialises once there is enough context for the bandit to
act on (matching the gate in the signal-processing path).
Returns None if the conversation is shorter than SIGNAL_GATE_MIN_MESSAGES.
"""
litellm_params = kwargs.get("litellm_params") or {}
sid = litellm_params.get("litellm_session_id")
if sid:
return str(sid)
metadata = litellm_params.get("metadata") or {}
if isinstance(metadata, dict):
sid = metadata.get("session_id") or metadata.get("litellm_session_id")
if sid:
return str(sid)
messages = kwargs.get("messages") or []
if len(messages) < SIGNAL_GATE_MIN_MESSAGES:
# Don't attribute until we have enough turns to match the signal gate —
# ensures the hash is stable (same N messages every time) and avoids
# crediting the bandit for conversations that are too short to signal.
return None
identity = ":".join(
str(metadata.get(f) or "") if isinstance(metadata, dict) else ""
for f in _IDENTITY_FIELDS
)
anchor = messages[:SIGNAL_GATE_MIN_MESSAGES]
payload = (
identity
+ "|"
+ json.dumps(
[{"role": m.get("role"), "content": m.get("content")} for m in anchor],
sort_keys=True,
default=str,
)
)
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
def _last_user_content(messages: Optional[List[Dict[str, Any]]]) -> Optional[str]:
if not messages:
return None
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
# OpenAI vision-style content: pick first text part.
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
return part.get("text")
return None
return None
def _recent_tool_results(
messages: Optional[List[Dict[str, Any]]]
) -> List[Dict[str, Any]]:
"""Extract the current turn's tool result payloads from the request messages.
Tool results are `role == "tool"` messages that sit at the tail of the
conversation — i.e. after the most recent assistant message with
`tool_calls`, waiting for the model to produce a user-facing reply. Walk
backwards from the end and collect the contiguous run of tool messages;
stop at the first non-tool message.
Each result is normalized to `{content, is_error}` — the only fields
`signals._detect_failure` / `_detect_exhaustion` actually read.
"""
if not messages:
return []
results: List[Dict[str, Any]] = []
for msg in reversed(messages):
if not isinstance(msg, dict):
break
if msg.get("role") != "tool":
break
content = msg.get("content")
# Some providers (Anthropic-style) carry an explicit error flag; OpenAI
# tool results don't, so fall back to an empty/missing content heuristic
# inside `_detect_failure`.
is_error = bool(msg.get("is_error"))
results.append({"content": content, "is_error": is_error})
results.reverse()
return results
def _assistant_content_and_tool_calls(response_obj: Any) -> tuple:
"""Return (assistant_text, tool_calls_list) extracted from a ModelResponse-ish object."""
if response_obj is None:
return None, []
try:
choices = getattr(response_obj, "choices", None) or response_obj.get("choices")
except Exception:
return None, []
if not choices:
return None, []
msg = choices[0]
msg = getattr(msg, "message", None) or (
msg.get("message") if isinstance(msg, dict) else None
)
if msg is None:
return None, []
content = getattr(msg, "content", None)
if content is None and isinstance(msg, dict):
content = msg.get("content")
raw_tool_calls = getattr(msg, "tool_calls", None)
if raw_tool_calls is None and isinstance(msg, dict):
raw_tool_calls = msg.get("tool_calls")
tool_calls: List[Dict[str, Any]] = []
for tc in raw_tool_calls or []:
if isinstance(tc, dict):
tool_calls.append(tc)
else:
try:
tool_calls.append(tc.model_dump())
except Exception:
tool_calls.append({"name": getattr(tc, "name", ""), "arguments": ""})
return content, tool_calls
class AdaptiveRouterPostCallHook(CustomLogger):
"""One hook instance per AdaptiveRouter. Registered into litellm.callbacks."""
def __init__(self, adaptive_router: AdaptiveRouter) -> None:
self.adaptive_router = adaptive_router
async def async_post_call_response_headers_hook(
self,
data: Dict[str, Any],
user_api_key_dict: Any,
response: Any,
request_headers: Optional[Dict[str, str]] = None,
litellm_call_info: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, str]]:
"""
Surface the chosen logical model as the `x-litellm-adaptive-router-model`
response header for both streaming and non-streaming responses.
`async_post_call_success_hook` fires after the stream is fully consumed,
so writing to `_hidden_params["additional_headers"]` there is too late for
streaming — the StreamingResponse headers are already frozen. This hook is
called during header construction (before StreamingResponse is built), so
the header is included for both paths.
"""
metadata = data.get("metadata") or {}
chosen = (
metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY)
if isinstance(metadata, dict)
else None
)
if not chosen:
return None
return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen}
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await self._record(kwargs, response_obj, response_status=200)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
status = kwargs.get("response_status")
if status is None:
exc = kwargs.get("exception")
status = getattr(exc, "status_code", 500) if exc is not None else 500
await self._record(kwargs, response_obj, response_status=int(status))
async def _record(
self,
kwargs: Dict[str, Any],
response_obj: Any,
response_status: int,
) -> None:
try:
messages = kwargs.get("messages") or []
if len(messages) < SIGNAL_GATE_MIN_MESSAGES:
# Too few turns for any signal to be meaningful — skip.
return
session_key = _resolve_session_key(kwargs)
if not session_key:
return
# The bandit cells are keyed by the *logical* model name from
# `available_models` (e.g. "smart"/"fast"). `kwargs["model"]` at
# post-call time is the physical upstream model
# (e.g. "anthropic/claude-opus-4-7"), so it cannot be used directly.
# The pre-routing hook stashes the logical pick under this key.
litellm_params = kwargs.get("litellm_params") or {}
metadata = litellm_params.get("metadata") or {}
current_model = (
metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY)
if isinstance(metadata, dict)
else None
)
if not current_model:
return
if not self.adaptive_router.claim_or_check_owner(
session_key, current_model
):
# A different model owns this conversation — skip attribution.
return
user_text = _last_user_content(messages)
assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj)
tool_results = _recent_tool_results(messages)
request_type = classify_prompt(user_text or "")
turn = Turn(
user_content=user_text,
assistant_content=(
assistant_text if isinstance(assistant_text, str) else None
),
tool_calls=tool_calls,
tool_results=tool_results,
response_status=response_status,
)
await self.adaptive_router.record_turn(
session_id=session_key,
model_name=current_model,
request_type=request_type,
turn=turn,
)
except Exception as e:
verbose_router_logger.exception(
"AdaptiveRouterPostCallHook: failed to record turn: %s", e
)

287
litellm/router_strategy/adaptive_router/signals.py Normal file
View file

@ -0,0 +1,287 @@
"""
Incremental signal detection for the adaptive router.
Each session maintains a SessionState. On every turn, we call apply_turn(state, turn)
which mutates the state in place and returns a SignalDelta listing which signals
fired on THIS turn. The router then queues the delta to be flushed to DB.
Design constraint: O(1) work per turn. No re-scanning the full session history.
We keep small bounded windows: last_user_content, last_assistant_content, and a
bounded list of recent tool call signatures.
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set
from litellm.router_strategy.adaptive_router.config import (
LOOP_REPEAT_THRESHOLD,
MIN_TURNS_FOR_CLEAN_CREDIT,
MISALIGNMENT_JACCARD_THRESHOLD,
STAGNATION_JACCARD_NEAR_DUP,
TOOL_CALL_HISTORY_MAX,
)
# ---- Public types ---------------------------------------------------------
@dataclass
class SignalDelta:
"""Which signals fired on a single turn. Counts are 0 or 1 (one delta per turn)."""
misalignment: int = 0
stagnation: int = 0
disengagement: int = 0
satisfaction: int = 0
failure: int = 0
loop: int = 0
exhaustion: int = 0
def any_fired(self) -> bool:
return any(
[
self.misalignment,
self.stagnation,
self.disengagement,
self.satisfaction,
self.failure,
self.loop,
self.exhaustion,
]
)
@dataclass
class SessionState:
"""In-memory rolling state for one session.
Mirrors the LiteLLM_AdaptiveRouterSession DB row (Wave 0 schema). The flusher
later persists this. We keep this as a plain dataclass — no DB coupling.
"""
session_id: str
router_name: str
model_name: str
classified_type: str
misalignment_count: int = 0
stagnation_count: int = 0
disengagement_count: int = 0
satisfaction_count: int = 0
failure_count: int = 0
loop_count: int = 0
exhaustion_count: int = 0
last_user_content: Optional[str] = None
last_assistant_content: Optional[str] = None
tool_call_history: List[str] = field(default_factory=list)
pending_tool_calls: Dict[str, str] = field(default_factory=dict)
turn_count: int = 0
last_processed_turn: int = -1
clean_credit_awarded: bool = False
terminal_status: Optional[int] = None
@dataclass
class Turn:
"""One turn of input. Caller assembles this from the request/response."""
user_content: Optional[str] = None
assistant_content: Optional[str] = None
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
tool_results: List[Dict[str, Any]] = field(default_factory=list)
response_status: Optional[int] = None
# ---- Detection helpers ----------------------------------------------------
_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
def _tokens(text: Optional[str]) -> Set[str]:
if not text:
return set()
return {t.lower() for t in _TOKEN_RE.findall(text)}
def _jaccard(a: Set[str], b: Set[str]) -> float:
union = a | b
if not union:
return 0.0
return len(a & b) / len(union)
_DISENGAGEMENT_PATTERNS = [
re.compile(
r"\b(forget it|never mind|give up|talk to (?:a )?human|cancel)\b", re.IGNORECASE
),
re.compile(r"\b(this (?:isn'?t|is not) working|stop|abort)\b", re.IGNORECASE),
re.compile(r"\bi'?ll do it (?:myself|manually)\b", re.IGNORECASE),
]
_SATISFACTION_PATTERNS = [
re.compile(
r"\b(that worked|that did it|works now|fixed it|solved it|nice)\b",
re.IGNORECASE,
),
re.compile(r"\b(thanks|thank you|thx|appreciated|appreciate it)\b", re.IGNORECASE),
re.compile(r"\b(perfect|great|excellent|exactly)\b", re.IGNORECASE),
]
def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> bool:
"""Fires when consecutive user messages share *some* topic (jaccard > 0)
but are sufficiently different (jaccard < threshold) — i.e. user is
rephrasing, not changing topic, not repeating."""
if not prev_user or not curr_user:
return False
j = _jaccard(_tokens(prev_user), _tokens(curr_user))
return 0.0 < j < MISALIGNMENT_JACCARD_THRESHOLD
def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bool:
"""Fires when consecutive assistant messages are near-duplicates."""
if not prev_asst or not curr_asst:
return False
j = _jaccard(_tokens(prev_asst), _tokens(curr_asst))
return j >= STAGNATION_JACCARD_NEAR_DUP
def _detect_disengagement(curr_user: Optional[str]) -> bool:
if not curr_user:
return False
return any(p.search(curr_user) for p in _DISENGAGEMENT_PATTERNS)
def _detect_satisfaction(curr_user: Optional[str]) -> bool:
if not curr_user:
return False
return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS)
def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool:
"""Any tool result explicitly flagged as an error.
We do NOT treat empty content as failure — many tools legitimately return
empty output (zero-result searches, silent bash commands, void writes) and
penalizing the model for those would corrupt the bandit posterior.
"""
for r in tool_results:
if r.get("is_error"):
return True
return False
def _signature(call: Dict[str, Any]) -> str:
"""Stable signature for loop detection: name + sorted JSON-ish args."""
name = call.get("name") or call.get("function", {}).get("name", "")
call_args = call.get("arguments")
if call_args is None:
call_args = call.get("function", {}).get("arguments", "")
if isinstance(call_args, dict):
call_args = ",".join(f"{k}={call_args[k]}" for k in sorted(call_args.keys()))
return f"{name}({call_args})"
def _detect_loop(history: List[str], new_calls: List[Dict[str, Any]]) -> bool:
"""Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times
in recent history (so this call would be the Nth)."""
if not new_calls:
return False
for call in new_calls:
sig = _signature(call)
recent_count = history.count(sig)
if recent_count >= LOOP_REPEAT_THRESHOLD - 1:
return True
return False
_EXHAUSTION_STATUSES = {408, 413, 429, 503, 504}
_EXHAUSTION_KEYWORDS = (
"context length",
"context window",
"token limit",
"rate limit",
"too many requests",
"timeout",
)
def _detect_exhaustion(
status: Optional[int], tool_results: List[Dict[str, Any]]
) -> bool:
if status is not None and status in _EXHAUSTION_STATUSES:
return True
for r in tool_results:
content = str(r.get("content", "")).lower()
if any(kw in content for kw in _EXHAUSTION_KEYWORDS):
return True
return False
# ---- Public entrypoint ----------------------------------------------------
def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
"""
Detect signals on this turn, mutate state, return the delta.
O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history
(which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload.
"""
delta = SignalDelta()
if _detect_misalignment(state.last_user_content, turn.user_content):
delta.misalignment = 1
if _detect_stagnation(state.last_assistant_content, turn.assistant_content):
delta.stagnation = 1
if _detect_disengagement(turn.user_content):
delta.disengagement = 1
if _detect_satisfaction(turn.user_content):
# Gate: only award satisfaction credit once per session, and only
# after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks"
# on turn 1-2 is noise, not a validated quality signal.
current_turn_index = state.turn_count + 1
if (
not state.clean_credit_awarded
and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT
):
delta.satisfaction = 1
state.clean_credit_awarded = True
if _detect_failure(turn.tool_results):
delta.failure = 1
if _detect_loop(state.tool_call_history, turn.tool_calls):
delta.loop = 1
if _detect_exhaustion(turn.response_status, turn.tool_results):
delta.exhaustion = 1
state.misalignment_count += delta.misalignment
state.stagnation_count += delta.stagnation
state.disengagement_count += delta.disengagement
state.satisfaction_count += delta.satisfaction
state.failure_count += delta.failure
state.loop_count += delta.loop
state.exhaustion_count += delta.exhaustion
if turn.user_content:
state.last_user_content = turn.user_content
if turn.assistant_content:
state.last_assistant_content = turn.assistant_content
for call in turn.tool_calls:
state.tool_call_history.append(_signature(call))
if len(state.tool_call_history) > TOOL_CALL_HISTORY_MAX:
state.tool_call_history = state.tool_call_history[-TOOL_CALL_HISTORY_MAX:]
if turn.response_status is not None:
state.terminal_status = turn.response_status
state.turn_count += 1
state.last_processed_turn = state.turn_count
return delta

213
litellm/router_strategy/adaptive_router/update_queue.py Normal file
View file

@ -0,0 +1,213 @@
"""
In-memory queues for adaptive router state and session updates.
Pattern follows DailySpendUpdateQueue: hot path is fully in-memory; a background
flusher task drains the aggregator and writes batches to Postgres.
Two logical queues (one class):
1. STATE updates: increments to (router, request_type, model) bandit cell.
Aggregator key = (router_name, request_type, model_name)
Aggregated payload = {"delta_alpha": float, "delta_beta": float, "samples_added": int}
2. SESSION updates: full snapshot of a session row (last-write-wins per session+router+model).
Aggregator key = (session_id, router_name, model_name)
Aggregated payload = the full session state dict.
Hot-path API is non-blocking and synchronous from the caller's POV (it just appends
to the in-memory aggregator). Flush is async and batched.
"""
from __future__ import annotations
import asyncio
from typing import Any, Dict, Tuple
from litellm._logging import verbose_router_logger
StateKey = Tuple[str, str, str] # (router_name, request_type, model_name)
SessionKey = Tuple[str, str, str] # (session_id, router_name, model_name)
class AdaptiveRouterUpdateQueue:
"""
Single class managing both state-update aggregation and session-snapshot aggregation.
Held by the AdaptiveRouter strategy instance and started by the proxy on boot.
"""
def __init__(self) -> None:
self._state_agg: Dict[StateKey, Dict[str, float]] = {}
self._session_agg: Dict[SessionKey, Dict[str, Any]] = {}
self._lock = asyncio.Lock()
self._max_state_size_seen = 0
self._max_session_size_seen = 0
# ---- Hot-path: state delta -------------------------------------------
async def add_state_delta(
self,
router_name: str,
request_type: str,
model_name: str,
delta_alpha: float,
delta_beta: float,
) -> None:
"""Aggregate a bandit-cell delta. Multiple deltas to the same cell sum."""
key: StateKey = (router_name, request_type, model_name)
async with self._lock:
current = self._state_agg.get(key)
if current is None:
self._state_agg[key] = {
"delta_alpha": delta_alpha,
"delta_beta": delta_beta,
"samples_added": 1,
}
else:
current["delta_alpha"] += delta_alpha
current["delta_beta"] += delta_beta
current["samples_added"] += 1
if len(self._state_agg) > self._max_state_size_seen:
self._max_state_size_seen = len(self._state_agg)
# ---- Hot-path: session snapshot --------------------------------------
async def add_session_state(
self,
session_id: str,
router_name: str,
model_name: str,
state_dict: Dict[str, Any],
) -> None:
"""
Last-write-wins per session row. The state_dict is a snapshot of the
SessionState (signals counts + bookkeeping fields). The flusher will
upsert this into LiteLLM_AdaptiveRouterSession.
"""
key: SessionKey = (session_id, router_name, model_name)
async with self._lock:
self._session_agg[key] = state_dict
if len(self._session_agg) > self._max_session_size_seen:
self._max_session_size_seen = len(self._session_agg)
# ---- Flushers (called by background task) ----------------------------
async def flush_state_to_db(self, prisma_client: Any) -> int:
"""
Drain state aggregator and apply to LiteLLM_AdaptiveRouterState.
Returns number of cells flushed.
"""
async with self._lock:
batch = self._state_agg
self._state_agg = {}
if not batch:
return 0
# Sort keys to give deterministic write order across writers and
# reduce the chance of cross-row deadlocks when other workers race us.
for key in sorted(batch.keys()):
router, rt, model = key
payload = batch[key]
try:
# Atomic increment: push the delta directly into the DB so
# concurrent flushers from multiple pods don't overwrite each
# other. The upsert creates the row with the delta as the
# initial value on first write, then increments on subsequent
# writes — no read-modify-write race.
await prisma_client.db.litellm_adaptiverouterstate.upsert(
where={
"router_name_request_type_model_name": {
"router_name": router,
"request_type": rt,
"model_name": model,
}
},
data={
"create": {
"router_name": router,
"request_type": rt,
"model_name": model,
"alpha": payload["delta_alpha"],
"beta": payload["delta_beta"],
"total_samples": int(payload["samples_added"]),
},
"update": {
"alpha": {"increment": payload["delta_alpha"]},
"beta": {"increment": payload["delta_beta"]},
"total_samples": {
"increment": int(payload["samples_added"])
},
},
},
)
except Exception as e:
verbose_router_logger.exception(
"AdaptiveRouterUpdateQueue: failed to flush state for %s: %s",
key,
e,
)
return len(batch)
async def flush_session_to_db(self, prisma_client: Any) -> int:
"""
Drain session aggregator and upsert into LiteLLM_AdaptiveRouterSession.
Returns number of session rows flushed.
"""
async with self._lock:
batch = self._session_agg
self._session_agg = {}
if not batch:
return 0
for key in sorted(batch.keys()):
session_id, router, model = key
payload = batch[key]
try:
# NOTE: Prisma client lower-cases model names, so
# `LiteLLM_AdaptiveRouterSession` -> `litellm_adaptiveroutersession`
# (single 's', not 'litellm_adaptiverouterssession').
# Strip PK fields from the update payload — Prisma rejects
# writes to fields that are part of the @@id. asdict(state)
# always carries them, so build a separate update dict.
update_payload = {
k: v
for k, v in payload.items()
if k not in ("session_id", "router_name", "model_name")
}
await prisma_client.db.litellm_adaptiveroutersession.upsert(
where={
"session_id_router_name_model_name": {
"session_id": session_id,
"router_name": router,
"model_name": model,
}
},
data={
"create": {
"session_id": session_id,
"router_name": router,
"model_name": model,
**update_payload,
},
"update": update_payload,
},
)
except Exception as e:
verbose_router_logger.exception(
"AdaptiveRouterUpdateQueue: failed to flush session for %s: %s",
key,
e,
)
return len(batch)
# ---- Observability ---------------------------------------------------
async def queue_size(self) -> Dict[str, int]:
async with self._lock:
return {
"state_pending": len(self._state_agg),
"session_pending": len(self._session_agg),
"max_state_seen": self._max_state_size_seen,
"max_session_seen": self._max_session_size_seen,
}

28
litellm/router_strategy/auto_router/auto_router.py
View file

@ -82,11 +82,34 @@ class AutoRouter(CustomLogger):
)
return auto_router_routes
@staticmethod
def _extract_text_from_messages(messages: List[Dict[str, Any]]) -> str:
"""
Extract text content from the last user message for routing.
Handles tool-call conversations (where the last message may be an
assistant or tool message with non-string content) and multimodal
messages (where content is a list of content blocks).
"""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content")
if content is None:
return ""
if isinstance(content, list):
return " ".join(
block.get("text", "")
for block in content
if isinstance(block, dict) and block.get("type") == "text"
)
return str(content)
return ""
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
@ -120,8 +143,7 @@ class AutoRouter(CustomLogger):
auto_sync=self.auto_sync_value,
)
user_message: Dict[str, str] = messages[-1]
message_content: str = user_message.get("content", "")
message_content = self._extract_text_from_messages(messages)
route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer(
text=message_content
)

128
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -332,45 +332,68 @@ class ComplexityRouter(CustomLogger):
f"No model configured for tier {tier_key} and no default_model set"
)
async def async_pre_routing_hook(
def _resolve_messages(
self,
model: str,
messages: Optional[List[Dict[str, Any]]],
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
) -> Optional[List[Dict[str, Any]]]:
"""
Pre-routing hook called before the routing decision.
Resolve messages from the request, converting from other formats if needed.
Classifies the request by complexity and returns the appropriate model.
Args:
model: The original model name requested.
request_kwargs: The request kwargs.
messages: The messages in the request.
input: Optional input for embeddings.
specific_deployment: Whether a specific deployment was requested.
Returns:
PreRoutingHookResponse with the routed model, or None if no routing needed.
Uses the guardrail translation handler dispatch to convert Responses API
``input`` (or other non-chat-completions formats) into OpenAI-spec messages.
"""
from litellm.types.router import PreRoutingHookResponse
if messages:
return messages
if messages is None or len(messages) == 0:
verbose_router_logger.debug(
"ComplexityRouter: No messages provided, skipping routing"
)
return None
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
from litellm.llms import load_guardrail_translation_mappings
from litellm.types.utils import CallTypes
# Extract the last user message and the last system prompt
mappings = load_guardrail_translation_mappings()
call_type: Optional[CallTypes] = None
# 1. Try route-based inference from proxy metadata
route = request_kwargs.get("litellm_metadata", {}).get(
"user_api_key_request_route"
)
if route:
call_types_list = get_call_types_for_route(route)
if call_types_list:
for ct in call_types_list:
if ct in mappings:
call_type = ct
break
# 2. Fallback: try each mapped handler until one produces messages
handlers_to_try: List[Any] = []
if call_type is not None and call_type in mappings:
handlers_to_try.append(mappings[call_type]())
else:
handlers_to_try.extend(handler_cls() for handler_cls in mappings.values())
for handler in handlers_to_try:
structured = handler.get_structured_messages(request_kwargs)
if structured:
return [
msg if isinstance(msg, dict) else msg.model_dump() # type: ignore
for msg in structured
]
return None
@staticmethod
def _extract_user_message_and_system_prompt(
messages: List[Dict[str, Any]],
) -> Tuple[Optional[str], Optional[str]]:
"""Extract the last user message text and last system prompt from messages."""
user_message: Optional[str] = None
system_prompt: Optional[str] = None
for msg in reversed(messages):
role = msg.get("role", "")
content = msg.get("content") or ""
# content may be a list of content parts (e.g. [{"type": "text", "text": "..."}])
if isinstance(content, list):
text_parts = [
part.get("text", "")
@ -383,6 +406,52 @@ class ComplexityRouter(CustomLogger):
user_message = content
elif role == "system" and system_prompt is None:
system_prompt = content
if user_message is not None and system_prompt is not None:
break
return user_message, system_prompt
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
"""
Pre-routing hook called before the routing decision.
Classifies the request by complexity and returns the appropriate model.
Supports chat completions (messages), Responses API (input), and other
formats via the guardrail translation handler dispatch.
Args:
model: The original model name requested.
request_kwargs: The request kwargs.
messages: The messages in the request.
input: Optional input for Responses API or embeddings.
specific_deployment: Whether a specific deployment was requested.
Returns:
PreRoutingHookResponse with the routed model, or None if no routing needed.
"""
from litellm.types.router import PreRoutingHookResponse
resolved_messages = self._resolve_messages(messages, request_kwargs)
if not resolved_messages:
verbose_router_logger.debug(
"ComplexityRouter: No messages could be resolved, skipping routing"
)
return None
# Determine whether the original request used messages directly
has_original_messages = messages is not None and len(messages) > 0
user_message, system_prompt = self._extract_user_message_and_system_prompt(
resolved_messages
)
if user_message is None:
verbose_router_logger.debug(
@ -391,13 +460,10 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=self.config.default_model
or self.get_model_for_tier(ComplexityTier.MEDIUM),
messages=messages,
messages=messages if has_original_messages else None,
)
# Classify the request
tier, score, signals = self.classify(user_message, system_prompt)
# Get the model for this tier
routed_model = self.get_model_for_tier(tier)
verbose_router_logger.info(
@ -407,5 +473,5 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
messages=messages if has_original_messages else None,
)

21
litellm/router_strategy/quality_router/__init__.py Normal file
View file

@ -0,0 +1,21 @@
"""
Quality-tier auto-router.
Re-uses the ComplexityRouter's classification to decide a request's complexity,
then maps that complexity to an admin-configured quality tier and resolves the
target model from each candidate's `model_info.litellm_routing_preferences`.
"""
from .config import (
DEFAULT_COMPLEXITY_TO_QUALITY,
QualityRouterConfig,
RoutingPreferences,
)
from .quality_router import QualityRouter
__all__ = [
"QualityRouter",
"QualityRouterConfig",
"RoutingPreferences",
"DEFAULT_COMPLEXITY_TO_QUALITY",
]

74
litellm/router_strategy/quality_router/config.py Normal file
View file

@ -0,0 +1,74 @@
"""
Configuration models for the QualityRouter.
"""
from typing import Dict, List, Optional
from pydantic import BaseModel, ConfigDict, Field
# Default mapping from ComplexityTier name (string) to quality tier (int).
# Higher tier = higher capability requirement.
DEFAULT_COMPLEXITY_TO_QUALITY: Dict[str, int] = {
"SIMPLE": 1,
"MEDIUM": 2,
"COMPLEX": 3,
"REASONING": 4,
}
class QualityRouterConfig(BaseModel):
"""Configuration for the QualityRouter."""
available_models: List[str] = Field(
default_factory=list,
description=(
"List of candidate model names this router may route to. Each model "
"must declare its quality_tier in model_info.litellm_routing_preferences."
),
)
default_model: Optional[str] = Field(
default=None,
description="Fallback model when no quality tier resolves.",
)
complexity_to_quality: Dict[str, int] = Field(
default_factory=lambda: DEFAULT_COMPLEXITY_TO_QUALITY.copy(),
description="Mapping from ComplexityTier name to quality tier (int).",
)
model_config = ConfigDict(extra="allow")
class RoutingPreferences(BaseModel):
"""Per-deployment routing preferences declared on model_info."""
quality_tier: int = Field(
...,
description="The quality tier this deployment satisfies.",
)
keywords: List[str] = Field(
default_factory=list,
description=(
"Substring keywords (case-insensitive) that, when present in the "
"user message, route the request to this deployment. See `order` "
"for explicit collision handling, otherwise ties fall through to "
"(highest quality_tier, then cheapest model_info.input_cost_per_token)."
),
)
order: Optional[int] = Field(
default=None,
description=(
"Explicit priority used to break ties between deployments at the "
"same quality tier. Lower values win. Applies both to keyword "
"collisions and to picking between multiple deployments at the "
"same quality_tier. Tiebreak order is "
"(quality_tier DESC, order ASC, input_cost_per_token ASC, "
"model_name ASC) — quality always wins first, then explicit "
"order, then price."
),
)
model_config = ConfigDict(extra="allow")

446
litellm/router_strategy/quality_router/quality_router.py Normal file
View file

@ -0,0 +1,446 @@
"""
Quality-tier Auto Router.
Routes a request to a model at a target quality tier. The quality tier is
inferred by re-using the existing ComplexityRouter's classification, then
mapped through an admin-configured `complexity_to_quality` table. Each
candidate model declares its own `quality_tier` in
`model_info.litellm_routing_preferences`.
Optional keyword override: deployments may also declare `keywords` in
`litellm_routing_preferences`. If any declared keyword appears in the user
message (case-insensitive substring match), the router short-circuits the
complexity-classification flow and routes to the matching deployment. When
multiple deployments match, ties are broken by (highest quality_tier first,
then cheapest `model_info.input_cost_per_token`).
"""
import math
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_strategy.complexity_router.complexity_router import (
ComplexityRouter,
)
from .config import QualityRouterConfig, RoutingPreferences
if TYPE_CHECKING:
from litellm.router import Router
from litellm.types.router import PreRoutingHookResponse
else:
Router = Any
PreRoutingHookResponse = Any
class QualityRouter(CustomLogger):
"""
Routes requests to a model at a target quality tier, with an optional
keyword override.
"""
def __init__(
self,
model_name: str,
litellm_router_instance: "Router",
default_model: Optional[str] = None,
quality_router_config: Optional[Dict[str, Any]] = None,
):
self.model_name = model_name
self.litellm_router_instance = litellm_router_instance
if quality_router_config:
self.config = QualityRouterConfig(**quality_router_config)
else:
self.config = QualityRouterConfig()
# Explicit default_model arg overrides anything in the config dict.
if default_model:
self.config.default_model = default_model
# Internal scorer — re-use the existing rule-based classifier.
self._scorer = ComplexityRouter(
model_name=f"{model_name}::scorer",
litellm_router_instance=litellm_router_instance,
)
# Per-model indices populated alongside the tier index. `_model_keywords`
# stores keywords lowercased so we can substring-match against the
# lowercased user message in O(total-keyword-count). `_model_quality`,
# `_model_cost`, and `_model_order` drive tiebreaking — `_model_order`
# is the explicit priority (lower wins, unset = +inf).
self._model_keywords: Dict[str, List[str]] = {}
self._model_quality: Dict[str, int] = {}
self._model_cost: Dict[str, Optional[float]] = {}
self._model_order: Dict[str, Optional[int]] = {}
# Tier → models index. Built lazily on first access so the QualityRouter
# deployment does NOT need to appear after all its referenced models in
# the config — when `_build_tier_index` runs eagerly in `__init__`, the
# router instance's `model_list` is still being assembled incrementally
# by `_create_deployment`, and any `available_models` defined AFTER the
# router entry in config.yaml would silently be reported as missing.
self._tier_to_models_cache: Optional[Dict[int, List[str]]] = None
verbose_router_logger.debug(
f"QualityRouter initialized for {model_name} with "
f"available_models={self.config.available_models}, "
f"default_model={self.config.default_model}"
)
@property
def _tier_to_models(self) -> Dict[int, List[str]]:
"""Lazy tier→models index; built on first access."""
if self._tier_to_models_cache is None:
self._tier_to_models_cache = self._build_tier_index()
return self._tier_to_models_cache
def _get_routing_preferences(self, deployment: Any) -> Optional[Dict[str, Any]]:
"""
Extract litellm_routing_preferences from a deployment, handling both
dict-shaped and Pydantic-object-shaped deployments.
"""
# Dict-shaped deployment.
if isinstance(deployment, dict):
model_info = deployment.get("model_info") or {}
if isinstance(model_info, dict):
return model_info.get("litellm_routing_preferences")
# Pydantic ModelInfo nested in a dict.
return getattr(model_info, "litellm_routing_preferences", None)
# Pydantic-object deployment.
model_info = getattr(deployment, "model_info", None)
if model_info is None:
return None
if isinstance(model_info, dict):
return model_info.get("litellm_routing_preferences")
return getattr(model_info, "litellm_routing_preferences", None)
def _get_deployment_input_cost(self, deployment: Any) -> Optional[float]:
"""
Extract `input_cost_per_token` from a deployment's model_info.
Returns None when not declared — None is treated as "infinite cost"
for the cheapest-tiebreak ordering, so unpriced models lose ties to
priced ones. (Admins who want a model to win on price must declare it.)
"""
if isinstance(deployment, dict):
model_info = deployment.get("model_info") or {}
else:
model_info = getattr(deployment, "model_info", None) or {}
if isinstance(model_info, dict):
cost = model_info.get("input_cost_per_token")
else:
cost = getattr(model_info, "input_cost_per_token", None)
if cost is None:
return None
try:
return float(cost)
except (TypeError, ValueError):
return None
def _get_deployment_model_name(self, deployment: Any) -> Optional[str]:
"""Extract `model_name` from a dict- or object-shaped deployment."""
if isinstance(deployment, dict):
return deployment.get("model_name")
return getattr(deployment, "model_name", None)
def _build_tier_index(self) -> Dict[int, List[str]]:
"""
Build {quality_tier: [model_name, ...]} for every model in
`available_models`, plus side indices `_model_keywords`,
`_model_quality`, and `_model_cost`. Raises if any listed model is
missing `litellm_routing_preferences`.
"""
model_list = getattr(self.litellm_router_instance, "model_list", None) or []
available = set(self.config.available_models)
# Track which available models we've matched so we can error on missing.
seen: Dict[str, bool] = {name: False for name in available}
tier_to_models: Dict[int, List[str]] = {}
for deployment in model_list:
name = self._get_deployment_model_name(deployment)
if name is None or name not in available:
continue
raw_prefs = self._get_routing_preferences(deployment)
if raw_prefs is None:
raise ValueError(
f"QualityRouter: model '{name}' is listed in available_models "
f"but has no model_info.litellm_routing_preferences"
)
# Validate via the Pydantic model so we get a clear error for
# missing quality_tier, wrong types, etc. This also means
# `RoutingPreferences` is the single source of truth for the
# accepted shape — readers relied on raw dicts before.
try:
if isinstance(raw_prefs, RoutingPreferences):
prefs = raw_prefs
elif isinstance(raw_prefs, dict):
prefs = RoutingPreferences(**raw_prefs)
else:
# A Pydantic object of some other shape — coerce via its dict.
prefs = RoutingPreferences(
**(
raw_prefs.model_dump()
if hasattr(raw_prefs, "model_dump")
else dict(raw_prefs)
)
)
except Exception as e:
raise ValueError(
f"QualityRouter: model '{name}' has invalid "
f"litellm_routing_preferences: {e}"
) from e
tier_int = int(prefs.quality_tier)
tier_to_models.setdefault(tier_int, []).append(name)
self._model_keywords[name] = [str(k).lower() for k in prefs.keywords if k]
self._model_quality[name] = tier_int
self._model_cost[name] = self._get_deployment_input_cost(deployment)
self._model_order[name] = prefs.order
seen[name] = True
missing = [name for name, found in seen.items() if not found]
if missing:
raise ValueError(
f"QualityRouter: the following available_models are not present in "
f"the router's model_list (or are missing routing preferences): {missing}"
)
# Sort each tier's model list so `_resolve_model_for_quality_tier`
# (which picks index [0]) honors (order ASC, cost ASC, name ASC).
# Quality is moot within a single tier; keep parity with the keyword
# tiebreak by ordering on (order, cost, name) here.
for models in tier_to_models.values():
models.sort(key=lambda n: (self._order_key(n), self._cost_key(n), n))
return tier_to_models
def _order_key(self, model_name: str) -> float:
"""`order` lookup as a float — unset becomes +inf so explicit wins."""
order = self._model_order.get(model_name)
return float(order) if order is not None else math.inf
def _cost_key(self, model_name: str) -> float:
"""`input_cost_per_token` as a float — unset becomes +inf."""
cost = self._model_cost.get(model_name)
return float(cost) if cost is not None else math.inf
def _keyword_override(self, user_message: str) -> Optional[Tuple[str, str]]:
"""
Find a deployment whose declared keywords appear in `user_message`.
Returns (model_name, matched_keyword) or None when no keyword matches.
When multiple deployments match, sorts by:
1. quality_tier DESC (best quality always wins first)
2. `order` ASC (explicit priority — unset = +inf so explicit wins
within the same tier)
3. input_cost_per_token ASC (unpriced = +inf so priced wins)
4. model_name ASC (deterministic stability)
"""
# Touch the lazy index so `_model_keywords` / `_model_quality` /
# `_model_cost` / `_model_order` are populated.
_ = self._tier_to_models
text = user_message.lower()
matches: List[Tuple[str, str]] = [] # (model_name, matched_keyword)
for model_name, keywords in self._model_keywords.items():
for kw in keywords:
if kw and kw in text:
matches.append((model_name, kw))
break # one match per model is enough
if not matches:
return None
def sort_key(match: Tuple[str, str]) -> Tuple[int, float, float, str]:
name = match[0]
quality = self._model_quality.get(name, 0)
order_val = self._order_key(name)
cost = self._model_cost.get(name)
cost_val = cost if cost is not None else math.inf
# Negate quality so higher tier sorts first under ASC sort.
return (-quality, order_val, cost_val, name)
matches.sort(key=sort_key)
return matches[0]
def _resolve_model_for_quality_tier(self, tier: int) -> str:
"""
Resolve a quality tier to a concrete model name.
Strategy:
1. Exact tier match → first model registered at that tier.
2. Round UP to the next higher tier that has a model (closer to a
request we might lack capacity for).
3. Round DOWN to the closest lower tier that has a model (degrade
gracefully instead of jumping straight to `default_model`,
which may be off-tier).
4. Fall back to `config.default_model`.
5. Otherwise raise.
"""
tier_index = self._tier_to_models
if tier in tier_index and tier_index[tier]:
return tier_index[tier][0]
# Round up.
higher_tiers = sorted(t for t in tier_index if t > tier)
for t in higher_tiers:
if tier_index[t]:
return tier_index[t][0]
# Round down — closest lower tier first.
lower_tiers = sorted((t for t in tier_index if t < tier), reverse=True)
for t in lower_tiers:
if tier_index[t]:
return tier_index[t][0]
if self.config.default_model:
return self.config.default_model
raise ValueError(
f"QualityRouter: no model available for quality tier {tier} and "
f"no default_model configured"
)
def _stash_decision(
self,
request_kwargs: Optional[Dict[str, Any]],
decision: Dict[str, Any],
) -> None:
"""
Stash the routing decision in request_kwargs.metadata so the Router can
lift it into response headers (`x-litellm-quality-router-*`). The same
dict object flows from here through to `make_call.set_response_headers`.
"""
if request_kwargs is None:
return
metadata = request_kwargs.setdefault("metadata", {})
if isinstance(metadata, dict):
metadata["quality_router_decision"] = decision
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
"""Try keyword override first; fall back to complexity-tier routing."""
from litellm.types.router import PreRoutingHookResponse
if messages is None or len(messages) == 0:
verbose_router_logger.debug(
"QualityRouter: No messages provided, skipping routing"
)
return None
# Extract last user message and last system prompt — same rules as
# ComplexityRouter.async_pre_routing_hook.
user_message: Optional[str] = None
system_prompt: Optional[str] = None
for msg in reversed(messages):
role = msg.get("role", "")
content = msg.get("content") or ""
if isinstance(content, list):
text_parts = [
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
]
content = " ".join(text_parts).strip()
if isinstance(content, str) and content:
if role == "user" and user_message is None:
user_message = content
elif role == "system" and system_prompt is None:
system_prompt = content
if user_message is None:
verbose_router_logger.debug(
"QualityRouter: No user message found, routing to default model"
)
if not self.config.default_model:
raise ValueError(
"QualityRouter: no user message and no default_model configured"
)
return PreRoutingHookResponse(
model=self.config.default_model,
messages=messages,
)
# Try keyword override first — it short-circuits complexity classification.
keyword_match = self._keyword_override(user_message)
if keyword_match is not None:
routed_model, matched_keyword = keyword_match
verbose_router_logger.info(
f"QualityRouter: keyword override matched='{matched_keyword}' "
f"routed_model={routed_model} "
f"(quality_tier={self._model_quality.get(routed_model)}, "
f"input_cost_per_token={self._model_cost.get(routed_model)})"
)
self._stash_decision(
request_kwargs,
{
"router_model_name": self.model_name,
"routed_model": routed_model,
"routed_via": "keyword",
"matched_keyword": matched_keyword,
"quality_tier": self._model_quality.get(routed_model),
"complexity_tier": None,
},
)
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
)
# No keyword match → complexity classification flow.
complexity_tier, score, signals = self._scorer.classify(
user_message, system_prompt
)
complexity_name = (
complexity_tier.value
if hasattr(complexity_tier, "value")
else str(complexity_tier)
)
quality_tier = self.config.complexity_to_quality.get(complexity_name)
if quality_tier is None:
raise ValueError(
f"QualityRouter: complexity tier '{complexity_name}' not present "
f"in complexity_to_quality mapping {self.config.complexity_to_quality}"
)
routed_model = self._resolve_model_for_quality_tier(int(quality_tier))
verbose_router_logger.info(
f"QualityRouter: complexity={complexity_name}, score={score:.3f}, "
f"signals={signals}, quality_tier={quality_tier}, "
f"routed_model={routed_model}"
)
self._stash_decision(
request_kwargs,
{
"router_model_name": self.model_name,
"routed_model": routed_model,
"routed_via": "quality_tier",
"matched_keyword": None,
"quality_tier": int(quality_tier),
"complexity_tier": complexity_name,
},
)
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
)

10
litellm/types/compression.py
View file

@ -2,7 +2,14 @@
Type definitions for litellm.compress().
"""
from typing import Dict, List, TypedDict
import sys
if sys.version_info >= (3, 11):
from typing import Dict, List, NotRequired, TypedDict
else:
from typing import Dict, List, TypedDict
from typing_extensions import NotRequired
class CompressedResult(TypedDict):
@ -12,3 +19,4 @@ class CompressedResult(TypedDict):
compression_ratio: float # fraction reduced, e.g. 0.6 means 60% reduction
cache: Dict[str, str] # key -> original content (for retrieval tool responses)
tools: List[dict] # [litellm_content_retrieve tool definition]
compression_skipped_reason: NotRequired[str]

27
litellm/types/integrations/compression_interception.py Normal file
View file

@ -0,0 +1,27 @@
"""
Type definitions for Compression Interception integration.
"""
from typing import Any, Dict, Optional, TypedDict
class CompressionInterceptionConfig(TypedDict, total=False):
"""
Configuration parameters for CompressionInterceptionLogger.
Used in proxy_config.yaml under litellm_settings:
litellm_settings:
compression_interception_params:
enabled: true
compression_trigger: 100000
compression_target: 70000
embedding_model: "text-embedding-3-small"
embedding_model_params:
dimensions: 512
"""
enabled: bool
compression_trigger: int
compression_target: Optional[int]
embedding_model: Optional[str]
embedding_model_params: Optional[Dict[str, Any]]

30
litellm/types/integrations/custom_logger.py
View file

@ -1,6 +1,6 @@
from typing import Optional
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
from pydantic import BaseModel, Field
class StandardCustomLoggerInitParams(BaseModel):
@ -9,3 +9,29 @@ class StandardCustomLoggerInitParams(BaseModel):
"""
turn_off_message_logging: Optional[bool] = False
class AgenticLoopRequestPatch(BaseModel):
"""
Patch returned by callbacks to request a follow-up LLM call.
"""
model: Optional[str] = None
messages: Optional[List[Dict[str, Any]]] = None
tools: Optional[List[Dict[str, Any]]] = None
max_tokens: Optional[int] = None
optional_params: Dict[str, Any] = Field(default_factory=dict)
kwargs: Dict[str, Any] = Field(default_factory=dict)
class AgenticLoopPlan(BaseModel):
"""
Typed callback response for agentic-loop reruns.
"""
run_agentic_loop: bool = False
request_patch: Optional[AgenticLoopRequestPatch] = None
response_override: Optional[Any] = None
terminate: bool = False
stop_reason: Optional[str] = None
metadata: Dict[str, Any] = Field(default_factory=dict)

2
litellm/types/integrations/prometheus.py
View file

@ -784,7 +784,7 @@ class UserAPIKeyLabelValues:
org_id: Optional[str] = None
org_alias: Optional[str] = None
#Added for test compatibility.
# Added for test compatibility.
def __init__(self, **kwargs: Any) -> None:
"""
Match former Pydantic behavior: unknown keys are ignored; ``api_key_hash`` maps to

44
litellm/types/llms/bedrock.py
View file

@ -997,3 +997,47 @@ class BedrockToolBlock(TypedDict, total=False):
toolSpec: Optional[ToolSpecBlock]
systemTool: Optional[SystemToolBlock] # For Nova grounding
cachePoint: Optional[CachePointBlock]
class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False):
"""
Top-level request body accepted by AWS Bedrock `InvokeModel` /
`InvokeModelWithResponseStream` when calling an Anthropic Claude model with
the Messages API format. The LiteLLM /v1/messages → Bedrock Invoke
transformation filters outgoing requests to the keys of this TypedDict; any
other field (Anthropic-only extension, internal metadata, future addition)
is dropped before signing so Bedrock doesn't 400 with
"Extra inputs are not permitted".
Reference:
https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html
https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html
Editing this type is the single source of truth — the runtime allowlist in
`AmazonAnthropicClaudeMessagesConfig.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS`
is derived from `__annotations__`, and a test asserts the resolved set
exactly, so any edit forces a conscious review.
Value types are intentionally loose (`list`, `dict`) — this type exists to
pin the allowed field names, not to validate nested structure.
"""
# Required by Bedrock
anthropic_version: str
max_tokens: int
messages: list
# Documented optional fields
anthropic_beta: List[str]
system: object # str or list[TextBlock]
stop_sequences: List[str]
temperature: float
top_p: float
top_k: int
tools: list
tool_choice: dict
# `thinking` is required for Opus 4.5 / Sonnet 4 extended thinking,
# `metadata` is part of the common Anthropic Messages API shape.
thinking: dict
metadata: dict

50
litellm/types/router.py
View file

@ -8,7 +8,7 @@ from dataclasses import dataclass
from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hints
import httpx
from pydantic import BaseModel, ConfigDict, Field, model_validator
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing_extensions import Required, TypedDict
from litellm._uuid import uuid
@ -221,6 +221,13 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
complexity_router_config: Optional[Dict] = None
complexity_router_default_model: Optional[str] = None
# adaptive-router params
adaptive_router_default_model: Optional[str] = None
adaptive_router_config: Optional[Dict] = None
# quality-router params
quality_router_config: Optional[Dict] = None
quality_router_default_model: Optional[str] = None
# Batch/File API Params
s3_bucket_name: Optional[str] = None
s3_encryption_key_id: Optional[str] = None
@ -790,3 +797,44 @@ class PreRoutingHookResponse(BaseModel):
model: str
messages: Optional[List[Dict[str, Any]]]
class RequestType(str, enum.Enum):
"""Fixed v0 taxonomy. User-extensible types come in v1."""
CODE_GENERATION = "code_generation"
CODE_UNDERSTANDING = "code_understanding"
TECHNICAL_DESIGN = "technical_design"
ANALYTICAL_REASONING = "analytical_reasoning"
WRITING = "writing"
FACTUAL_LOOKUP = "factual_lookup"
GENERAL = "general"
class AdaptiveRouterWeights(BaseModel):
quality: float = Field(default=0.7, ge=0.0, le=1.0)
cost: float = Field(default=0.3, ge=0.0, le=1.0)
@field_validator("cost")
@classmethod
def _weights_sum_to_one(cls, v, info):
q = info.data.get("quality", 0.7)
if abs(q + v - 1.0) > 0.001:
raise ValueError(
f"weights must sum to 1.0, got quality={q} + cost={v} = {q + v}"
)
return v
class AdaptiveRouterConfig(BaseModel):
available_models: List[str]
weights: AdaptiveRouterWeights = Field(default_factory=AdaptiveRouterWeights)
class AdaptiveRouterPreferences(BaseModel):
"""model_info.adaptive_router_preferences — declared by each model."""
model_config = ConfigDict(use_enum_values=False)
quality_tier: int = Field(ge=1, le=3)
strengths: List[RequestType] = Field(default_factory=list)

2
litellm/types/utils.py
View file

@ -2851,6 +2851,7 @@ class StandardAuditLogPayload(TypedDict):
class StandardLoggingPayload(TypedDict):
id: str
trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries)
litellm_call_id: Optional[str] # UUID returned in x-litellm-call-id response header
call_type: str
stream: Optional[bool]
response_cost: float
@ -3290,6 +3291,7 @@ class LlmProviders(str, Enum):
MANUS = "manus"
WANDB = "wandb"
OVHCLOUD = "ovhcloud"
SCALEWAY = "scaleway"
LEMONADE = "lemonade"
AMAZON_NOVA = "amazon_nova"
A2A_AGENT = "a2a_agent"

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