scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge branch 'litellm_internal_staging' into litellm_/virtual-keys-e2e-tests-f40710

Browse source
This commit is contained in:
yuneng-jiang 2026-09-08 13:55:49 -07:00 • committed by GitHub
commit 4812fc77e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
323 changed files with 12815 additions and 4544 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.circleci/config.yml
View file

@ -1440,6 +1440,7 @@ jobs:
TEST_FILES=$(printf "%s\n" \
tests/local_testing/test_dual_cache.py \
tests/local_testing/test_redis_batch_optimizations.py \
tests/local_testing/test_redis_increment_with_floor.py \
tests/local_testing/test_router_utils.py)
echo "$TEST_FILES" | circleci tests run \
--verbose \

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

@ -117,6 +117,9 @@ jobs:
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- name: Run pytest tests/test_litellm_rust with the compiled extension
run: make test-rust-extension
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"

6
CLAUDE.md
View file

@ -29,7 +29,7 @@ Never test structure of code only function of it
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on `litellm_internal_staging` in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. If your branch already carries a budget edit, drop it before opening the PR
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
@ -70,7 +70,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch

4
CONTRIBUTING.md
View file

@ -315,10 +315,12 @@ Ensure the UI builds successfully before submitting your PR:
npm run build
```
Local lint and budget checks follow origin's current default branch. They refresh it from the remote instead of trusting cached `origin/HEAD`. For an intentional comparison against another branch or commit, use `make check BASE_REF=<ref>` or the standalone gate's `--base <ref>` option. An explicit ref can also be used offline once it has been fetched locally. Without an override, unavailable remote metadata stops the check
## Submitting Your PR
1. **Push your branch**: `git push origin your-feature-branch`
2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
2. **Create a PR**: Go to GitHub and open a pull request against the repository's current default branch. Run `python3 scripts/default_branch.py --branch` to check its name
3. **Fill out the PR template**: Provide clear description of changes
4. **Wait for review**: Maintainers will review and provide feedback
5. **Address feedback**: Make requested changes and push updates

73
Makefile
View file

@ -4,6 +4,7 @@
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
test-rust-extension \
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
@ -34,7 +35,7 @@ help:
@echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed"
@echo " make lint-format - Check ruff format formatting (matches CI)"
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)"
@echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
@echo " make lint-test-quality - Gate the test suite against test-quality-budget.json"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
@ -54,12 +55,16 @@ help:
@echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)"
@echo " make test-integration - Run integration tests"
@echo " make test-unit-helm - Run helm unit tests"
@echo " make test-rust-extension - Build the Rust extension and run its public Python tests"
@echo ""
@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
UV := uv
UV_RUN := $(UV) run --no-sync
BASE_REF ?=
export BASE_REF
RESOLVE_BASE = python3 scripts/default_branch.py --base "$(BASE_REF)"
# Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so
# it runs before any venv exists. See scripts/gate_slot_lock.py.
@ -67,7 +72,7 @@ GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py
LINT_DEP_INSTALL ?= install-dev
LINT_E2E_DEP_INSTALL ?= lint-install
LINT_DEP_BASE ?= lint-fetch-base
LINT_DEP_BASE ?=
LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4)
LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,)
@ -130,10 +135,8 @@ format: install-dev
format-check: install-dev
cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd ..
# Single fetch of the PR base so the delta-based gates below share one network round
# trip instead of each re-fetching when chained from `lint`.
lint-fetch-base:
git fetch origin litellm_internal_staging
@$(RESOLVE_BASE)
# Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated
# Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The
@ -150,7 +153,9 @@ lint-install:
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
@base_ref=$$($(RESOLVE_BASE)) && \
changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \
files=$$(printf '%s\n' "$$changed" | grep -v '^litellm/enterprise/' || true) || exit $$?; \
if [ -z "$$files" ]; then \
echo "No changed litellm Python files to format-check."; \
else \
@ -167,7 +172,9 @@ lint-ruff: $(LINT_DEP_INSTALL)
# https://github.com/astral-sh/ruff/discussions/10977
# https://github.com/astral-sh/ruff/discussions/4049
lint-format-changed: install-dev
@git diff origin/main --unified=0 --no-color -- '*.py' | \
@base_ref=$$($(RESOLVE_BASE)) && \
diff=$$(git diff "$$base_ref" --unified=0 --no-color -- '*.py') && \
printf '%s\n' "$$diff" | \
perl -ne '\
if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \
if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \
@ -182,20 +189,22 @@ lint-format-changed: install-dev
done
lint-ruff-dev: install-dev
@tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
@base_ref=$$($(RESOLVE_BASE)) || exit $$?; \
tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
cd litellm && \
($(UV_RUN) ruff check . --output-format=pylint || true) > "$$tmpfile" && \
$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \
$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch="$$base_ref" && \
cd .. ; \
rm -f "$$tmpfile"
lint-ruff-FULL-dev: install-dev
@files=$$(git diff --name-only origin/main -- '*.py'); \
@base_ref=$$($(RESOLVE_BASE)) && \
files=$$(git diff --name-only "$$base_ref" -- '*.py') || exit $$?; \
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
else echo "No changed .py files to check."; fi
lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)"
lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
$(UV_RUN) basedpyright tests/e2e
@ -203,37 +212,37 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
# Type-discipline budget (mutable collections / casts / type guards / kwargs /
# unexplained suppressions), the test-linting.yml step `make lint` used to omit.
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_discipline_gate.py --base "$(BASE_REF)"
# Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
# litellm module-global mutation, credential-gated skips, conftest snapshot
# inventory), counted across tests/ the same delta-vs-base way.
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)"
# --update lowers each limit by what this branch fixed since its branch point, so
# it needs the base ref fetched to resolve the merge-base.
lint-basedpyright-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_check_gate.py --update
lint-basedpyright-budget-update: install-dev
$(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)"
lint-format: format-check
lint-ruff-budget: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py
$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
# Strict gate, invoked the same way CI does in test-linting.yml so a local pass
# means the CI check will pass too.
lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
lint-ruff-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/ruff_strict_gate.py --update
lint-ruff-budget-update: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)"
lint-type-discipline-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_discipline_gate.py --update
lint-type-discipline-budget-update: install-dev
$(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)"
lint-test-quality-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/test_quality_gate.py --update
lint-test-quality-budget-update: install-dev
$(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)"
# Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
@ -249,14 +258,15 @@ check-import-safety: $(LINT_DEP_INSTALL)
# runs the diff-scoped ruff format check, whole-tree ruff check, the strict-rule /
# type-discipline / basedpyright budgets as a delta vs the base, then the circular-import
# and import-safety checks. Steps that compare against the base resolve it the same way CI
# does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client,
# does (merge-base with origin's current default branch). Setup (env sync, Prisma client,
# base fetch) runs once up front; the checks themselves are independent, so a sub-make
# fans them out with -j and the fast ones finish under basedpyright's shadow.
lint:
@$(GATE_SLOT_LOCK) $(MAKE) lint-inner
lint-inner: lint-install lint-fetch-base
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-inner: lint-install
@base_ref=$$($(RESOLVE_BASE)) && \
$(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
@ -281,6 +291,17 @@ pre-commit:
@$(MAKE) check
# Testing targets
test-rust-extension:
@temporary=$$(mktemp -d) && \
trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
set -- "$$temporary"/wheels/*.whl && \
[ "$$#" -eq 1 ] && \
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
test: install-test-deps
$(UV_RUN) pytest tests/

62
ci_cd/run_migration.py
View file

@ -6,12 +6,9 @@ import subprocess
import sys
from datetime import datetime
from pathlib import Path
import testing.postgresql
from typing import Final
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:
@ -94,31 +91,57 @@ def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
print(banner, file=out)
def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
def _default_base_branch(root_dir: Path) -> str:
try:
result: Final = subprocess.run(
[
sys.executable,
str(Path(__file__).resolve().parents[1] / "scripts" / "default_branch.py"),
"--repo-root",
str(root_dir),
"--branch",
],
check=True,
capture_output=True,
text=True,
timeout=90,
)
except (OSError, subprocess.SubprocessError) as exc:
_print_freshness_failure(
"default branch",
"Could not discover origin's default branch. Pass --base-branch <name> to choose one.",
exc.stderr if isinstance(exc, subprocess.CalledProcessError) else str(exc),
)
sys.exit(3)
return result.stdout.strip()
def _check_branch_freshness(root_dir: Path, base_branch: str | None = None) -> None:
"""Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
resolved_branch: Final = base_branch or _default_base_branch(root_dir)
cwd = str(root_dir)
try:
subprocess.run(
["git", "fetch", "origin", base_branch],
["git", "fetch", "origin", f"+refs/heads/{resolved_branch}:refs/remotes/origin/{resolved_branch}"],
check=True,
capture_output=True,
text=True,
cwd=cwd,
)
except FileNotFoundError:
_print_freshness_failure(base_branch, "git executable not found on PATH")
_print_freshness_failure(resolved_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",
resolved_branch,
f"`git fetch origin {resolved_branch}` failed",
e.stderr or "",
)
sys.exit(3)
try:
result = subprocess.run(
["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
["git", "rev-list", "--count", f"HEAD..origin/{resolved_branch}"],
check=True,
capture_output=True,
text=True,
@ -127,23 +150,23 @@ def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
behind = int(result.stdout.strip())
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git rev-list HEAD..origin/{base_branch}` failed",
resolved_branch,
f"`git rev-list HEAD..origin/{resolved_branch}` failed",
e.stderr or "",
)
sys.exit(3)
except ValueError:
_print_freshness_failure(
base_branch,
resolved_branch,
"could not parse commit count from `git rev-list`",
)
sys.exit(3)
if behind > 0:
_print_stale_branch_refusal(base_branch, behind)
_print_stale_branch_refusal(resolved_branch, behind)
sys.exit(3)
print(f"Branch freshness OK: up to date with origin/{base_branch}.")
print(f"Branch freshness OK: up to date with origin/{resolved_branch}.")
def _print_destructive_refusal(destructive_lines: list) -> None:
@ -198,7 +221,7 @@ def _print_destructive_refusal(destructive_lines: list) -> None:
def create_migration(
migration_name: str = None,
allow_destructive: bool = False,
base_branch: str = DEFAULT_BASE_BRANCH,
base_branch: str | None = None,
skip_freshness_check: bool = False,
):
"""
@ -211,7 +234,7 @@ def create_migration(
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").
(default: origin's current default branch).
skip_freshness_check (bool): Skip the "branch is up to date" check.
Only for intentional migrations against an older base.
"""
@ -225,6 +248,8 @@ def create_migration(
else:
_check_branch_freshness(root_dir, base_branch)
import testing.postgresql
try:
migrations_dir = (
root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
@ -342,9 +367,8 @@ if __name__ == "__main__":
)
parser.add_argument(
"--base-branch",
default=DEFAULT_BASE_BRANCH,
help=(
f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
"Branch to check freshness against (default: origin's current default branch). "
"The script fetches origin/<base-branch> and refuses to run if HEAD "
"is behind it."
),

6
docker/component_entrypoint.sh
View file

@ -1,5 +1,11 @@
#!/bin/sh
# stale samples from a previous container incarnation would be summed into the aggregate
if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then
mkdir -p "$PROMETHEUS_MULTIPROC_DIR"
rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"

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

@ -152,6 +152,13 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
@ -189,6 +196,11 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.metricsServer.enabled }}
- name: metrics
containerPort: {{ .Values.metricsServer.port }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}

17
helm/litellm-helm/templates/service-metrics.yaml Normal file
View file

@ -0,0 +1,17 @@
{{- if .Values.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.fullname" . }}-metrics
labels:
{{- include "litellm.labels" . | nindent 4 }}
spec:
type: ClusterIP
ports:
- port: {{ .Values.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.selectorLabels" . | nindent 4 }}
{{- end }}

2
helm/litellm-helm/templates/servicemonitor.yaml
View file

@ -26,7 +26,7 @@ spec:
{{- toYaml .namespaceSelector.matchNames | nindent 4 }}
{{- end }}
endpoints:
- port: http
- port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
path: /metrics/
interval: {{ .interval }}
scrapeTimeout: {{ .scrapeTimeout }}

106
helm/litellm-helm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: separate metrics server
templates:
- configmap-litellm.yaml
- deployment.yaml
- service.yaml
- service-metrics.yaml
- servicemonitor.yaml
tests:
- it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default
asserts:
- notContains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
any: true
template: deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- hasDocuments:
count: 0
template: service-metrics.yaml
- it: should scrape the proxy port when the metrics server is disabled
template: servicemonitor.yaml
set:
serviceMonitor.enabled: true
asserts:
- equal:
path: spec.endpoints[0].port
value: http
- it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor
set:
metricsServer.enabled: true
metricsServer.port: 4101
serviceMonitor.enabled: true
service.type: LoadBalancer
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "4101"
template: deployment.yaml
- contains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
containerPort: 4101
protocol: TCP
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-metrics
template: service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
template: service-metrics.yaml
- equal:
path: spec.endpoints[0].port
value: metrics
template: servicemonitor.yaml
- equal:
path: spec.endpoints[0].path
value: /metrics/
template: servicemonitor.yaml
- it: should reject a metrics port equal to the proxy port
template: deployment.yaml
set:
metricsServer.enabled: true
metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: metricsServer.port must differ from service.port

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

@ -180,6 +180,16 @@ proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT)
# so a scrape never runs on an inference worker. Adds a `metrics` port to the
# container and a dedicated ClusterIP `<release>-metrics` Service, and the
# ServiceMonitor scrapes it instead of the proxy port. The separate port has
# no virtual-key auth: keep it off public ingress. Needs the proxy image
# v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an

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

@ -441,3 +441,5 @@ ImplementationSpecific
{{- .pathType -}}
{{- end -}}
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}

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

@ -64,14 +64,25 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
{{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
{{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
{{- end }}
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -97,16 +108,54 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: metrics
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- {{ .Values.gateway.metricsServer.port | quote }}
env:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
ports:
- name: metrics
containerPort: {{ .Values.gateway.metricsServer.port }}
protocol: TCP
volumeMounts:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
readinessProbe:
tcpSocket: { port: metrics }
periodSeconds: 10
livenessProbe:
tcpSocket: { port: metrics }
periodSeconds: 15
failureThreshold: 6
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

18
helm/litellm/templates/gateway/service-metrics.yaml Normal file
View file

@ -0,0 +1,18 @@
{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.gateway.fullname" . }}-metrics
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
spec:
type: ClusterIP
ports:
- port: {{ .Values.gateway.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
{{- end }}

148
helm/litellm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,148 @@
suite: test gateway metrics sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/service.yaml
- gateway/service-metrics.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, volume, env or service port when the metrics server is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
any: true
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- hasDocuments:
count: 0
template: gateway/service-metrics.yaml
- it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4101
gateway.service.type: LoadBalancer
gateway.image.tag: v1.101.0
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.101.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- "4101"
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].env
value:
- name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].ports
value:
- name: metrics
containerPort: 4101
protocol: TCP
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].volumeMounts
value:
- name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 50m
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
emptyDir: {}
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: gateway/service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway-metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: gateway/service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
template: gateway/service-metrics.yaml
- it: rejects a metrics port equal to the gateway port
template: gateway/deployment.yaml
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: gateway.metricsServer.port must differ from the gateway port 4000

16
helm/litellm/values.yaml
View file

@ -268,6 +268,22 @@ gateway:
config:
create: true
proxy_config: {}
# Serve Prometheus /metrics from a `metrics` sidecar container (same image,
# `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the
# workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a
# scrape never runs on an inference worker. Adds a `metrics` port to the pod
# and a dedicated ClusterIP `<gateway>-metrics` Service; point your scrape
# config at it. The port has no virtual-key auth: keep it off public ingress.
# Needs the gateway image v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
requests:
cpu: 50m
memory: 128Mi
limits:
memory: 512Mi
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql Normal file
View file

@ -0,0 +1,3 @@
ALTER TABLE "LiteLLM_AutoRouterSession"
ADD COLUMN IF NOT EXISTS "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0;

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

@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])

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

@ -8,14 +8,10 @@ import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
from typing import TYPE_CHECKING, Final, Optional
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
prisma_command_timeout,
prisma_migrate_deploy_timeout,
)
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
if TYPE_CHECKING:
import psycopg
import psycopg.sql
def str_to_bool(value: Optional[str]) -> bool:
@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
_INVALID_LITELLM_INDEXES_SQL: Final = (
"SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
"FROM pg_index i "
"JOIN pg_class c ON c.oid = i.indexrelid "
"JOIN pg_class t ON t.oid = i.indrelid "
"JOIN pg_namespace n ON n.oid = t.relnamespace "
"WHERE NOT i.indisvalid "
" AND c.relkind = 'i' "
" AND n.nspname = %s "
" AND t.relname LIKE %s "
" AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
"ORDER BY c.relname"
)
@dataclass(frozen=True, slots=True)
class _InvalidIndex:
schema: str
name: str
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
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
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
parsed = urlparse(url)
if not parsed.query:
@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
"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)))
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
)
@staticmethod
def _invalid_litellm_indexes(
conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
) -> tuple[_InvalidIndex, ...]:
rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
return tuple(_InvalidIndex(*row) for row in rows)
@staticmethod
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
from psycopg import sql
target: Final = sql.Identifier(index.schema, index.name)
if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
@staticmethod
def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
import psycopg
statement, action = ProxyExtrasDBManager._index_repair(index)
try:
conn.execute(statement)
except psycopg.Error as e:
logger.warning(
"Could not repair invalid index %s.%s, will retry on the next startup. "
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
index.schema,
index.name,
statement.as_string(conn),
e,
)
return
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
@staticmethod
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
INVALID (a migration deadlock between replicas is the usual cause; the
retried migration skips them because of IF NOT EXISTS). Never raises:
returns True when no invalid index remains, False when the repair was
skipped or failed and will be retried on the next startup. Looks in the
schema DATABASE_URL names, the only URL Prisma migrates through, but
connects over DIRECT_URL when set: the session settings, the advisory
lock and REINDEX CONCURRENTLY all need one server session, which a
transaction pooler does not give."""
prisma_url: Final = os.getenv("DATABASE_URL")
if not prisma_url:
return False
try:
import psycopg
from psycopg import sql
except ImportError:
logger.warning(
"psycopg is not installed; skipping the invalid index check. "
"Install the litellm[extra_proxy] extra, which includes psycopg."
)
return False
schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public"
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url)
try:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
conn.execute("SET statement_timeout = 0")
conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
if not found:
return True
logger.warning(
"Found %d invalid index(es) left by an interrupted CREATE INDEX "
"CONCURRENTLY, rebuilding: %s",
len(found),
", ".join(f"{index.name} (table size {index.table_size})" for index in found),
)
lock_row: Final = conn.execute(
"SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
).fetchone()
if lock_row is None or not lock_row[0]:
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
return False
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
ProxyExtrasDBManager._repair_index(conn, index)
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
except psycopg.Error as e:
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
return False
return not remaining
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
"""
@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated

11
litellm-proxy-extras/migration_runbook.md
View file

@ -48,7 +48,7 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## What It Does
1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
1. **Verifies the current branch is up to date with origin's current default branch** (see [Branch freshness](#branch-freshness-check))
2. Creates temp PostgreSQL DB
3. Applies existing migrations
4. Compares with `schema.prisma`
@ -57,11 +57,11 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## 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.
Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. The default base is discovered from origin's advertised HEAD on each run, so an existing clone follows a default-branch change without trusting cached `origin/HEAD`. If discovery or fetching fails, migration generation stops. 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`.
- `--base-branch <name>` — check against a different base (e.g. a release branch). Defaults to origin's current default branch
- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
When the guard fires:
@ -69,8 +69,9 @@ 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
base_branch=$(python3 scripts/default_branch.py --branch) &&
git fetch origin "+refs/heads/$base_branch:refs/remotes/origin/$base_branch" &&
git rebase "origin/$base_branch"
```
2. Re-run `run_migration.py`.

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

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

3
litellm/__init__.py
View file

@ -2002,6 +2002,9 @@ if TYPE_CHECKING:
from .llms.hosted_vllm.responses.transformation import (
HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig,
)
from .llms.fireworks_ai.responses.transformation import (
FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig,
)
from .llms.github_copilot.chat.transformation import (
GithubCopilotConfig as GithubCopilotConfig,
)

5
litellm/_lazy_imports_registry.py
View file

@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = (
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
"HostedVLLMResponsesAPIConfig",
"FireworksAIResponsesAPIConfig",
"VolcEngineResponsesAPIConfig",
"PerplexityResponsesConfig",
"DatabricksResponsesAPIConfig",
@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.hosted_vllm.responses.transformation",
"HostedVLLMResponsesAPIConfig",
),
"FireworksAIResponsesAPIConfig": (
".llms.fireworks_ai.responses.transformation",
"FireworksAIResponsesAPIConfig",
),
"VolcEngineResponsesAPIConfig": (
".llms.volcengine.responses.transformation",
"VolcEngineResponsesAPIConfig",

1
litellm/batches/main.py
View file

@ -319,6 +319,7 @@ def create_batch(
timeout=timeout,
max_retries=optional_params.max_retries,
create_batch_data=_create_batch_request,
custom_endpoint=optional_params.get("custom_endpoint"),
)
else:
raise litellm.exceptions.BadRequestError(

65
litellm/caching/redis_cache.py
View file

@ -20,6 +20,8 @@ from contextvars import ContextVar
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
from pydantic import TypeAdapter
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
@ -80,11 +82,29 @@ class _AsyncRedisCommands(Protocol):
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
_BREAKER_GUARD_FRAME_NAMES: Final = frozenset(
{"<lambda>", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"}
)
_INCREMENT_WITH_FLOOR_LUA: Final = (
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]) "
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end "
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
"return count"
)
_LUA_COUNT: Final = TypeAdapter(int)
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]:
return _OPTIONAL_COUNTS.validate_python(
tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values)
)
def _get_call_stack_info(num_frames: int = 2) -> str:
"""
@ -736,6 +756,43 @@ class RedisCache(BaseCache):
)
raise e
@_redis_circuit_breaker_guard_sync
def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
"""Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call.
A counter whose key expired while a request was still in flight would otherwise be
recreated negative by that request's decrement. Clamping inside the same call is what
keeps it safe: a separate corrective write could land after another pod's increment and
erase it.
The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was
created rather than ``ttl`` after it was last touched. Refreshing it on every touch
would keep a count a dead worker never decremented alive for as long as the group
takes traffic. Returns the resulting count.
"""
namespaced_key: Final = self.check_and_fix_namespace(key=key)
count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval
_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl
)
return _LUA_COUNT.validate_python(count)
@_redis_circuit_breaker_guard_sync
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
"""Read integer counters for ``key_list``, in order, raising when Redis cannot answer.
``batch_get_cache`` swallows every failure and returns an empty dict, which the caller
cannot tell apart from "every counter is unset". A caller that has to fall back to its
own numbers when Redis is unreachable needs the failure, not a dict of zeros.
"""
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
"""Async twin of ``batch_get_counts``, raising on failure the same way."""
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_scan_iter(self, pattern: str, count: int = 100) -> list:
start_time: Final = time.time()
@ -1241,6 +1298,14 @@ class RedisCache(BaseCache):
result = result.decode()
return float(result)
@_redis_circuit_breaker_guard
async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int:
"""Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees."""
_redis_client: Final = self._async_commands()
namespaced_key: Final = self.check_and_fix_namespace(key=key)
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
return _LUA_COUNT.validate_python(count)
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

9
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -370,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
and isinstance(tool_call.get("custom"), dict)
)
for msg in messages:
leading_system_count: Final = next(
(index for index, msg in enumerate(messages) if msg.get("role") != "system"),
len(messages),
)
for index, msg in enumerate(messages):
role = msg.get("role")
content = msg.get("content", "")
tool_calls = msg.get("tool_calls")
@ -378,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if role == "system":
# Extract system message as instructions
if isinstance(content, str):
if isinstance(content, str) and index < leading_system_count:
if instructions:
# Concatenate multiple system prompts with a space
instructions = f"{instructions} {content}"

6
litellm/constants.py
View file

@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
budget_reservation_disabled_info_emitted = False
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096))
DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3))
DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1))
DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5))
DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float(
os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1")
)
DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
@ -1457,6 +1461,8 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"})
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"

40
litellm/cost_calculator.py
View file

@ -2414,6 +2414,46 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
)
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
class _ResponsesWsEventResponse(BaseModel):
usage: Mapping[str, object] | None = None
class _ResponsesWsEvent(BaseModel):
type: str = ""
response: _ResponsesWsEventResponse | None = None
class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
@staticmethod
def collect_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> tuple[Usage, ...]:
events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results)
return tuple(
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses
event.response.usage
)
for event in events
if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
and event.response is not None
and event.response.usage is not None
)
@staticmethod
def collect_and_combine_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> Usage:
collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results(
results
)
return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(
list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter
)
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"

3
litellm/exceptions.py
View file

@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError):
num_retries: int | None = None,
headers: dict | None = None,
exception_status_code: int | None = None,
response: httpx.Response | None = None,
):
request: Final = httpx.Request(
method="POST",
@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError):
self.max_retries = max_retries
self.num_retries = num_retries
self.headers = headers
if response is not None:
self.response = response
# custom function to convert to str
def __str__(self):

82
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, cast
from urllib.parse import urlparse
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_prompt_management import CustomPromptManagement
@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
)
OPENAI_API_HOST: Final = "api.openai.com"
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
def _validated_object_mapping(value: object) -> dict[object, object] | None:
try:
return _OBJECT_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return None
def _validated_object_list(value: object) -> list[object] | None:
try:
return _OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
if model_map_flag is not None:
@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_
class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _request_value(request_kwargs: object, key: str) -> object:
request_mapping: Final = _validated_object_mapping(request_kwargs)
if request_mapping is None:
return None
return request_mapping.get(key)
@staticmethod
def _request_user_agent(request_kwargs: object) -> str | None:
proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request")
proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request)
if proxy_server_request_mapping is None:
return None
headers: Final = proxy_server_request_mapping.get("headers")
headers_mapping: Final = _validated_object_mapping(headers)
if headers_mapping is None:
return None
user_agent: Final = next(
(value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"),
None,
)
return user_agent if isinstance(user_agent, str) else None
@staticmethod
def _request_system(request_kwargs: object) -> str | list[object] | None:
system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system")
if isinstance(system, str):
return system
return _validated_object_list(system)
def get_chat_completion_prompt(
self,
model: str,
@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
points: Sequence[CacheControlInjectionPoint],
messages: list[AllMessageValues],
tools: list[object] | None,
cache_control: object,
model: str,
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
) -> Sequence[Mapping[str, object]] | None:
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
return None
return AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None,
cache_control: object = None,
) -> bool:
"""Whether configured injection points must yield to client-set cache_control.
@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
"""
if all(point.get("_litellm_judged") for point in points):
return False
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
@staticmethod
def _request_has_cache_control(
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None = None,
cache_control: object = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if cache_control is not None:
return True
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
return True
if tools is not None:
@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
tools: list | None = None,
enable_prompt_caching: bool | None = None,
cache_control: object = None,
request_kwargs: object = None,
) -> list[CacheControlInjectionPoint]:
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
return []
if is_claude_code_one_shot_subagent_request(
messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs)
):
return []
control: Final = AnthropicCacheControlHook._default_control()
@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
models: Iterable[str],
tools: list[AllToolParamValues] | None = None,
enable_prompt_caching: bool | None = None,
request_kwargs: object = None,
) -> list[AllMessageValues]:
"""Return the messages auto prompt caching will send, default breakpoints included.
@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
for candidate in (
AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
custom_llm_provider=None,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
system=AnthropicCacheControlHook._request_system(request_kwargs),
cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"),
request_kwargs=request_kwargs,
)
for model in models
)
@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
non_default_params["cache_control_injection_points"],
messages,
tools,
non_default_params.get("cache_control"),
model,
custom_llm_provider,
api_base,
@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider=custom_llm_provider,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
cache_control=non_default_params.get("cache_control"),
request_kwargs=non_default_params,
)
if points:
non_default_params["cache_control_injection_points"] = points
@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy
bool | None, kwargs.pop("enable_prompt_caching", None)
)
cache_control: Final = kwargs.get("cache_control")
configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
if configured and AnthropicCacheControlHook._should_stand_down(
configured, typed_messages, system, tools, cache_control
):
return messages, system
injection_points: list[CacheControlInjectionPoint] = configured or []
if not injection_points and model is not None:
@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
model=model,
custom_llm_provider=custom_llm_provider,
enable_prompt_caching=enable_prompt_caching,
cache_control=cache_control,
request_kwargs=kwargs,
)
if not injection_points:
return messages, system

6
litellm/integrations/callback_configs.json
View file

@ -356,6 +356,12 @@
"description": "OpenTelemetry collector endpoint URL",
"required": true
},
"otel_traces_endpoint": {
"type": "text",
"ui_name": "Traces Endpoint URL",
"description": "Complete trace export URL used verbatim when the collector does not serve /v1/traces (OTel v2 only)",
"required": false
},
"otel_headers": {
"type": "text",
"ui_name": "Headers",

20
litellm/integrations/custom_guardrail.py
View file

@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger):
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None,
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
allowed_hooks: Final = frozenset(supported_event_hooks) | (
frozenset((GuardrailEventHooks.logging_only,))
if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
else frozenset()
)
def _validate_event_hook_list_is_in_supported_event_hooks(
event_hook: list[GuardrailEventHooks] | list[str],
supported_event_hooks: list[GuardrailEventHooks],
@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger):
for hook in event_hook:
if isinstance(hook, str):
hook = GuardrailEventHooks(hook)
if hook not in supported_event_hooks:
if hook not in allowed_hooks:
raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}")
if event_hook is None:
@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger):
default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default]
_validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks)
elif isinstance(event_hook, GuardrailEventHooks):
if event_hook not in supported_event_hooks:
if event_hook not in allowed_hooks:
raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}")
@staticmethod
@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger):
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
def _deployment_pre_call_target(self) -> "CustomLogger":
def _deployment_hook_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
return self
try:
@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger):
# CHECK IF GUARDRAIL REJECTS THE REQUEST
if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
target: Final = self._deployment_pre_call_target()
target: Final = self._deployment_hook_target()
if target is not self:
kwargs["guardrail_to_apply"] = self
result: Final = await target.async_pre_call_hook(
@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger):
return None
# CHECK IF GUARDRAIL REJECTS THE REQUEST
result: Final = await self.async_post_call_success_hook(
target: Final = self._deployment_hook_target()
hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger):
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=request_data,
data=hook_request_data,
response=response,
)

19
litellm/integrations/otel/model/config.py
View file

@ -72,6 +72,14 @@ class ExporterSpec(BaseModel):
description="console | in_memory | otlp_http | otlp_grpc | <factory kind>",
)
endpoint: str | None = None
traces_endpoint: str | None = Field(
default=None,
description=(
"Complete OTLP/HTTP trace URL, used verbatim. Set this when the "
"collector serves traces on a path other than ``/v1/traces``; "
"``endpoint`` is a base URL the signal path is appended to."
),
)
headers: str | None = None
owner: ExporterOwner | None = Field(
default=None,
@ -127,6 +135,14 @@ class OpenTelemetryV2Config(BaseSettings):
default=None,
validation_alias=AliasChoices("OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"),
)
traces_endpoint: str | None = Field(
default=None,
validation_alias=AliasChoices("OTEL_TRACES_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"),
description=(
"Complete OTLP/HTTP trace URL for the single-destination shorthand, "
"used verbatim instead of ``endpoint`` + ``/v1/traces``."
),
)
headers: str | None = Field(
default=None,
validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"),
@ -250,7 +266,7 @@ class OpenTelemetryV2Config(BaseSettings):
@model_validator(mode="after")
def _normalize(self) -> "OpenTelemetryV2Config":
# An endpoint with the default exporter kind implies OTLP/HTTP.
if self.endpoint and self.exporter == "console":
if (self.endpoint or self.traces_endpoint) and self.exporter == "console":
self.exporter = "otlp_http"
# When no explicit destinations are given, fold the single-destination
# shorthand into one spec so the provider always has a destination.
@ -259,6 +275,7 @@ class OpenTelemetryV2Config(BaseSettings):
ExporterSpec(
kind=self.exporter,
endpoint=self.endpoint,
traces_endpoint=self.traces_endpoint,
headers=self.headers,
)
]

11
litellm/integrations/otel/plumbing/providers.py
View file

@ -170,7 +170,7 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
)
return HTTPExporter(
endpoint=_otlp_traces_endpoint(spec.endpoint),
endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint),
headers=parse_headers(spec.headers),
)
if kind in _OTLP_GRPC_KINDS:
@ -201,7 +201,14 @@ def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple
exporters, populate ``config.exporters`` directly.
"""
return _exporter_from_spec(ExporterSpec(kind=config.exporter, endpoint=config.endpoint, headers=config.headers))
return _exporter_from_spec(
ExporterSpec(
kind=config.exporter,
endpoint=config.endpoint,
traces_endpoint=config.traces_endpoint,
headers=config.headers,
)
)
def _otlp_metrics_endpoint(endpoint: str | None) -> str | None:

15
litellm/litellm_core_utils/credential_accessor.py
View file

@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem
class CredentialAccessor:
@staticmethod
def find_credential(credential_name: str) -> CredentialItem | None:
return next(
(credential for credential in litellm.credential_list if credential.credential_name == credential_name),
None,
)
@staticmethod
def get_credential_values(credential_name: str) -> dict:
"""Safe accessor for credentials."""
if not litellm.credential_list:
return {}
for credential in litellm.credential_list:
if credential.credential_name == credential_name:
return credential.credential_values.copy()
return {}
credential: Final = CredentialAccessor.find_credential(credential_name)
return {} if credential is None else credential.credential_values.copy()
@staticmethod
def upsert_credentials(credentials: list[CredentialItem]):

11
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -860,6 +860,7 @@ def _map_bedrock_exception(
message=mantle_context_window_message,
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
)
if (
"too many tokens" in error_str
@ -873,6 +874,7 @@ def _map_bedrock_exception(
message=f"BedrockException: Context Window Error - {error_str}",
model=model,
llm_provider="bedrock",
response=getattr(original_exception, "response", None),
)
elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str:
raise BadRequestError(
@ -924,12 +926,14 @@ def _map_bedrock_exception(
message=f"BedrockException: Timeout Error - {error_str}",
model=model,
llm_provider="bedrock",
response=getattr(original_exception, "response", None),
)
elif "Could not process image" in error_str:
raise litellm.InternalServerError(
message=f"BedrockException - {error_str}",
model=model,
llm_provider="bedrock",
response=getattr(original_exception, "response", None),
)
elif hasattr(original_exception, "status_code"):
if original_exception.status_code == 500:
@ -937,10 +941,7 @@ def _map_bedrock_exception(
message=f"BedrockException - {original_exception.message}",
llm_provider="bedrock",
model=model,
response=httpx.Response(
status_code=500,
request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
),
response=getattr(original_exception, "response", None),
)
elif original_exception.status_code == 401:
raise AuthenticationError(
@ -969,6 +970,7 @@ def _map_bedrock_exception(
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
)
elif original_exception.status_code == 422:
raise BadRequestError(
@ -1001,6 +1003,7 @@ def _map_bedrock_exception(
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
exception_status_code=original_exception.status_code,
response=getattr(original_exception, "response", None),
)

12
litellm/litellm_core_utils/litellm_logging.py
View file

@ -48,6 +48,7 @@ from litellm.constants import (
)
from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
_select_model_name_for_cost_calc,
)
from litellm.exceptions import (
@ -2028,6 +2029,17 @@ class Logging(LiteLLMLoggingBaseClass):
results=result,
)
elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
combined_ws_usage: Final = (
ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
)
logging_result = LiteLLMRealtimeStreamLoggingObject(
usage=combined_ws_usage,
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
elif (
self.call_type == CallTypes.llm_passthrough_route.value
or self.call_type == CallTypes.allm_passthrough_route.value

5
litellm/litellm_core_utils/llm_response_utils/get_headers.py
View file

@ -1,7 +1,8 @@
from collections.abc import Mapping
from typing import Final
def get_response_headers(_response_headers: dict | None = None) -> dict:
def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict:
"""
Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header}
@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict:
return {**llm_provider_headers, **openai_headers}
def _get_llm_provider_headers(response_headers: dict) -> dict:
def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict:
"""
Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider

90
litellm/llms/anthropic/common_utils.py
View file

@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
_INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
_DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
_DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
def is_claude_code_user_agent(user_agent: str) -> bool:
return user_agent.startswith("claude-cli/")
def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
try:
return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return None
def _validated_claude_code_list(value: object) -> list[object] | None:
try:
return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None:
stripped: Final = text.strip()
if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX):
return None
fields: Final = tuple(
field
for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";")
if (field := raw_field.strip())
)
if not fields or any("=" not in field for field in fields):
return None
parsed_fields: Final = tuple(
(parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),)
)
if any(not key or not value for key, value in parsed_fields):
return None
return parsed_fields
def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
if isinstance(system, str):
return (system,)
blocks: Final = _validated_claude_code_list(system)
if blocks is None:
return None
block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
if any(block is None for block in block_mappings):
return None
text_values: Final = tuple(
block.get("text") for block in block_mappings if block is not None and block.get("type") == "text"
)
if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values):
return None
meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip())
return meaningful_text or None
def _is_claude_code_subagent_billing_system(system: object) -> bool:
billing_texts: Final = _claude_code_billing_texts(system)
if billing_texts is None:
return False
billing_fields: Final = tuple(
fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None
)
if len(billing_fields) != len(billing_texts):
return False
subagent_values: Final = tuple(
value for fields in billing_fields for key, value in fields if key == "cc_is_subagent"
)
return subagent_values == ("true",)
def is_claude_code_one_shot_subagent_request(
messages: list[AllMessageValues],
system: object,
tools: object,
user_agent: str | None,
) -> bool:
only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None
return (
user_agent is not None
and is_claude_code_user_agent(user_agent)
and not tools
and only_message is not None
and only_message.get("role") == "user"
and _is_claude_code_subagent_billing_system(system)
)
def _strip_bedrock_id_suffixes(model: str) -> str:

2
litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
View file

@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp(
LiteLLM_Proxy_MCP_Handler,
)
mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
if not mcp_references:
return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(

19
litellm/llms/bedrock/chat/agentcore/transformation.py
View file

@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=str(response.read()))
raise BedrockError(
status_code=response.status_code,
message=str(response.read()),
headers=response.headers,
response=response,
)
# LOGGING
logging_obj.post_call(
@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
status_code=response.status_code,
message=f"AgentCore: Failed to read/parse JSON response body: {e}",
headers=response.headers,
)
parsed: Final = self._parse_json_response(response_json)
@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=str(await response.aread()))
raise BedrockError(
status_code=response.status_code,
message=str(await response.aread()),
headers=response.headers,
response=response,
)
# LOGGING
logging_obj.post_call(
@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
status_code=response.status_code,
message=f"AgentCore: Failed to read/parse JSON response body: {e}",
headers=response.headers,
)
parsed: Final = self._parse_json_response(response_json)
@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error processing response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
def validate_environment(
@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
return headers
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def should_fake_stream(
self,

23
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
from ..common_utils import BedrockError, _get_all_bedrock_regions
from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
@ -66,7 +66,12 @@ def make_sync_call(
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=str(response.read()))
raise BedrockError(
status_code=response.status_code,
message=str(response.read()),
headers=response.headers,
response=response,
)
if fake_stream:
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@ -247,7 +252,12 @@ class BedrockConverseLLM(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -594,7 +604,12 @@ class BedrockConverseLLM(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

1
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -2255,6 +2255,7 @@ class AmazonConverseConfig(BaseConfig):
raise BedrockError(
message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues",
status_code=422,
headers=response.headers,
)
"""

3
litellm/llms/bedrock/chat/invoke_agent/transformation.py
View file

@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error processing response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
def validate_environment(
@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
return headers
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def should_fake_stream(
self,

33
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk
from ..common_utils import (
BedrockError,
build_bedrock_stream_error,
error_response_text,
get_bedrock_response_stream_shape,
get_bedrock_tool_name,
)
@ -184,7 +185,12 @@ async def make_call(
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=response.text)
raise BedrockError(
status_code=response.status_code,
message=error_response_text(response),
headers=response.headers,
response=response,
)
if fake_stream:
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@ -228,9 +234,16 @@ async def make_call(
)
return completion_stream, response.headers
except BedrockError:
raise
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
except Exception as e:
@ -270,7 +283,12 @@ def make_sync_call(
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=response.text)
raise BedrockError(
status_code=response.status_code,
message=error_response_text(response),
headers=response.headers,
response=response,
)
if fake_stream:
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@ -314,9 +332,16 @@ def make_sync_call(
)
return completion_stream, response.headers
except BedrockError:
raise
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
except Exception as e:

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
View file

@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
View file

@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py
View file

@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
raise BedrockError(
message=f"Error parsing response: {raw_response.text}, error: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
verbose_logger.debug(
@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
raise BedrockError(
message=f"Error setting response content: {e}. Response: {completion_response}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
# Calculate usage from headers

10
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
try:
completion_response: Final = raw_response.json()
except Exception:
raise BedrockError(message=raw_response.text, status_code=raw_response.status_code)
raise BedrockError(
message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
verbose_logger.debug(
"bedrock invoke response % s",
json.dumps(completion_response, indent=4, default=str),
@ -363,6 +367,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error processing={raw_response.text}, Received error={e}",
status_code=422,
headers=raw_response.headers,
)
try:
@ -384,6 +389,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error parsing received text={outputText}.\nError-{e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
## CALCULATING USAGE - bedrock returns usage in the headers
@ -431,7 +437,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)
@track_llm_api_timing()
async def get_async_custom_stream_wrapper(

11
litellm/llms/bedrock/claude_platform/common_utils.py
View file

@ -1,7 +1,10 @@
from typing import Final
import httpx
import litellm
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.secret_managers.main import get_secret_str
CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic"
@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str:
class BedrockClaudePlatformMixin(BaseAWSLLM):
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
@staticmethod
def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None:
workspace_id = (

47
litellm/llms/bedrock/common_utils.py
View file

@ -33,8 +33,53 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"
def error_response_text(response: httpx.Response) -> str:
try:
return response.text
except httpx.ResponseNotRead:
return response.reason_phrase
def _synthesize_error_response(
*, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None
) -> tuple[httpx.Request, httpx.Response]:
error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL)
safe_headers: Final = (
headers
if isinstance(headers, httpx.Headers)
else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes)))
)
return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request)
class BedrockError(BaseLLMException):
pass
def __init__(
self,
status_code: int,
message: str,
headers: dict[str, object] | httpx.Headers | None = None,
request: httpx.Request | None = None,
response: httpx.Response | None = None,
body: dict[str, object] | None = None,
status_code_is_synthesized: bool = False,
) -> None:
error_request, error_response = (
_synthesize_error_response(status_code=status_code, headers=headers, request=request)
if response is None and headers
else (request, response)
)
super().__init__(
status_code=status_code,
message=message,
headers=headers,
request=error_request,
response=error_response,
body=body,
status_code_is_synthesized=status_code_is_synthesized,
)
_BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (

4
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -102,6 +102,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise BedrockError(
status_code=response.status_code,
message=error_text,
headers=response.headers,
response=response,
)
bedrock_response: Final = response.json()
@ -124,6 +126,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,
headers=e.response.headers,
response=e.response,
)
except Exception as e:
verbose_logger.error("Error in CountTokens handler: %s", e)

14
litellm/llms/bedrock/embed/embedding.py
View file

@ -132,7 +132,12 @@ class BedrockEmbedding(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -161,7 +166,12 @@ class BedrockEmbedding(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

9
litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py
View file

@ -20,6 +20,7 @@ import httpx
from litellm._logging import verbose_logger
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
"""
return _supports_nova_canvas_image_edit_from_model_cost(model or "")
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_supported_openai_params(self, model: str) -> list:
return [
"n",

14
litellm/llms/bedrock/image_edit/handler.py
View file

@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

9
litellm/llms/bedrock/image_edit/stability_transformation.py
View file

@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.llms.stability import (
OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO,
@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
return True
return False
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_supported_openai_params(self, model: str) -> list:
"""
Return list of OpenAI params supported by Bedrock Stability.

14
litellm/llms/bedrock/image_generation/image_handler.py
View file

@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
### FORMAT RESPONSE TO OPENAI FORMAT ###
@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

9
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
AmazonInvokeConfig,
)
from litellm.llms.bedrock.common_utils import (
BedrockError,
apply_bedrock_invoke_structured_output,
ensure_bedrock_anthropic_messages_tool_names,
get_anthropic_beta_from_headers,
@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig(
BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys())
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def __init__(self, **kwargs):
BaseAnthropicMessagesConfig.__init__(self, **kwargs)
AmazonInvokeConfig.__init__(self, **kwargs)

11
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -2,13 +2,14 @@ import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, cast
import httpx
from httpx import Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo
from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo
if TYPE_CHECKING:
from httpx import URL
@ -18,6 +19,14 @@ if TYPE_CHECKING:
class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig):
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
return "stream" in endpoint

10
litellm/llms/bedrock/realtime/transformation.py
View file

@ -9,12 +9,14 @@ import json
import uuid as uuid_lib
from typing import Final, cast
import httpx
from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm
from litellm.types.llms.openai import (
OpenAIRealtimeContentPartDone,
@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self._cumulative_usage = BedrockUsageEvent()
self._reported_usage = BedrockUsageEvent()
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict:
"""Validate environment - no special validation needed for Bedrock."""
return headers

14
litellm/llms/bedrock/rerank/handler.py
View file

@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

6
litellm/llms/bedrock/search/transformation.py
View file

@ -39,7 +39,6 @@ from typing import Final
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.search.transformation import (
BaseSearchConfig,
SearchResponse,
@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
raise BedrockError(
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
message=f"AgentCore gateway MCP error: {error}",
headers=raw_response.headers,
)
# A failed tools/call is reported in-band, as HTTP 200 with result.isError
@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
raise BedrockError(
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}",
headers=raw_response.headers,
)
text_items: Final = tuple(
@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
raise BedrockError(
status_code=502,
message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}",
headers=raw_response.headers,
)
def get_error_class(
@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
status_code: int,
headers: dict, # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict
) -> Exception:
return BaseLLMException(
return BedrockError(
status_code=status_code,
message=error_message,
headers=headers,

9
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -8,6 +8,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.integrations.rag.bedrock_knowledgebase import (
BedrockKBContent,
BedrockKBResponse,
@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
BaseVectorStoreConfig.__init__(self)
BaseAWSLLM.__init__(self)
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials:
return {}

9
litellm/llms/bedrock_mantle/chat/transformation.py
View file

@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the
from collections.abc import AsyncIterator, Iterator
from typing import Any, Final
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
from ...base_llm.chat.transformation import BaseLLMException
from ...bedrock.common_utils import BedrockError
from ...openai_like.chat.transformation import OpenAILikeChatConfig
from ..common_utils import mantle_base_segment
@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
def get_config(cls):
return super().get_config()
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def _get_openai_compatible_provider_info(
self,
api_base: str | None,

8
litellm/llms/bedrock_mantle/responses/transformation.py
View file

@ -19,11 +19,14 @@ import json
from collections.abc import Mapping
from typing import Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock_mantle.common_utils import (
MANTLE_HOST_RE,
BedrockMantleAuthMixin,
@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK_MANTLE
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_complete_url(
self,
api_base: str | None,

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

@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
return self._transform_ocr_response(
provider_config=provider_config,
model=model,
@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
# Use async response transform for async operations
return await provider_config.async_transform_ocr_response(
model=model,

44
litellm/llms/fireworks_ai/common_utils.py
View file

@ -1,3 +1,5 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from httpx import Headers
@ -13,7 +15,7 @@ class FireworksAIException(BaseLLMException):
pass
def get_fireworks_session_id(litellm_params: dict) -> str | None:
def get_fireworks_session_id(litellm_params: Mapping[str, object]) -> str | None:
"""
Session id to send as `x-session-affinity`, or None when the caller gave none.
@ -23,19 +25,39 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
"""
params: Final = litellm_params
metadata: Final = params.get("metadata")
if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
if isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return None
for key in ("litellm_session_id", "session_id"):
value = params.get(key)
if value:
return str(value)
if isinstance(metadata, dict):
if isinstance(metadata, Mapping):
value = metadata.get("session_id")
if value:
return str(value)
return None
def with_fireworks_session_affinity(
headers: Mapping[str, str], litellm_params: Mapping[str, object]
) -> Mapping[str, str]:
if any(key.lower() == "x-session-affinity" for key in headers):
return headers
session_id: Final = get_fireworks_session_id(litellm_params)
if not session_id:
return headers
return MappingProxyType({**headers, "x-session-affinity": session_id})
def resolve_fireworks_api_key(api_key: str | None) -> str | None:
return api_key or (
get_secret_str("FIREWORKS_API_KEY")
or get_secret_str("FIREWORKS_AI_API_KEY")
or get_secret_str("FIREWORKSAI_API_KEY")
or get_secret_str("FIREWORKS_AI_TOKEN")
)
AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-"
@ -63,13 +85,7 @@ class FireworksAIMixin:
)
def _get_api_key(self, api_key: str | None) -> str | None:
dynamic_api_key: Final = api_key or (
get_secret_str("FIREWORKS_API_KEY")
or get_secret_str("FIREWORKS_AI_API_KEY")
or get_secret_str("FIREWORKSAI_API_KEY")
or get_secret_str("FIREWORKS_AI_TOKEN")
)
return dynamic_api_key
return resolve_fireworks_api_key(api_key)
def validate_environment(
self,
@ -92,9 +108,5 @@ class FireworksAIMixin:
return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params)
def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict:
if any(key.lower() == "x-session-affinity" for key in headers):
return headers
session_id: Final = get_fireworks_session_id(litellm_params)
if not session_id:
return headers
return {**headers, "x-session-affinity": session_id}
pinned: Final = with_fireworks_session_affinity(headers, litellm_params)
return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place

102
litellm/llms/fireworks_ai/responses/transformation.py Normal file
View file

@ -0,0 +1,102 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from urllib.parse import unquote
import httpx
from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam
from litellm.llms.fireworks_ai.common_utils import (
resolve_fireworks_api_key,
resolve_fireworks_resource_name,
with_fireworks_session_affinity,
)
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import ResponseInputParam
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1"
def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]:
extras: Final[Mapping[str, object]] = litellm_params.model_extra or MappingProxyType({})
return MappingProxyType(
{"litellm_session_id": extras.get("litellm_session_id"), "metadata": extras.get("litellm_metadata")}
)
def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam:
if "role" not in item or item["role"] != "developer":
return item
return EasyInputMessageParam(role="system", content=item["content"], type="message")
def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam:
if isinstance(input, str):
return input
return [_developer_item_as_system(item) for item in input]
class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.FIREWORKS_AI
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
litellm_params: GenericLiteLLMParams | None,
) -> dict: # mutable-ok: overrides the base class signature
params: Final = litellm_params or GenericLiteLLMParams()
api_key: Final = resolve_fireworks_api_key(params.api_key)
if api_key is None:
raise ValueError("FIREWORKS_API_KEY is not set")
authorized: Final = MappingProxyType(
{"Content-Type": "application/json", **headers, "Authorization": f"Bearer {api_key}"}
)
pinned: Final = with_fireworks_session_affinity(authorized, _session_params(params))
return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/")
return f"{base}/responses"
def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
return _developer_items_as_system(super()._validate_input_param(input))
def transform_responses_api_request(
self,
model: str,
input: str | ResponseInputParam,
response_api_optional_request_params: dict, # mutable-ok: overrides the base class signature
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: overrides the base class signature
) -> dict: # mutable-ok: overrides the base class signature
return super().transform_responses_api_request(
model=resolve_fireworks_resource_name(model),
input=input,
response_api_optional_request_params=response_api_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
def transform_delete_response_api_response(
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> DeleteResponseResult:
deleted_id: Final = unquote(raw_response.request.url.path.rsplit("/", 1)[-1])
return DeleteResponseResult(id=deleted_id, object="response", deleted=True)
def supports_native_websocket(self) -> bool:
return False
def supports_native_file_search(self) -> bool:
return False

131
litellm/llms/vertex_ai/batches/handler.py
View file

@ -1,6 +1,7 @@
import json
from collections.abc import Coroutine
from collections.abc import Coroutine, Sequence
from typing import TYPE_CHECKING, Final, Protocol
from urllib.parse import urlparse
import httpx
from typing_extensions import ReadOnly, TypedDict
@ -12,11 +13,13 @@ from litellm.litellm_core_utils.url_utils import (
safe_get,
)
from litellm.llms.custom_httpx.http_handler import (
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path
from litellm.types.llms.openai import CreateBatchRequest
from litellm.types.llms.vertex_ai import (
VERTEX_CREDENTIALS_TYPES,
@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict):
response: ReadOnly[httpx.Response]
class _VertexEndpointDeployedModel(TypedDict, total=False):
model: ReadOnly[str]
class _VertexEndpointResponse(TypedDict, total=False):
deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]]
class _VertexEndpointPayloadView(TypedDict):
"""Holds one decoded GET endpoints/<id> response so the payload reads back typed."""
payload: ReadOnly[_VertexEndpointResponse]
def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse:
return response.json()
@ -78,7 +95,17 @@ class VertexAIBatchPrediction(VertexLLM):
vertex_location: str | None,
timeout: float | httpx.Timeout,
max_retries: int | None,
custom_endpoint: bool | None = None,
) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]:
if custom_endpoint:
raise VertexAIError(
status_code=400,
message=(
"Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
"The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
"use a publisher model or fine-tuned Gemini endpoint deployment instead."
),
)
sync_handler: Final = _get_httpx_client()
access_token, project_id = self._ensure_access_token(
@ -87,6 +114,26 @@ class VertexAIBatchPrediction(VertexLLM):
custom_llm_provider="vertex_ai",
)
headers: Final = {
"Content-Type": "application/json; charset=utf-8",
"Authorization": f"Bearer {access_token}",
}
transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
request=create_batch_data,
vertex_project=vertex_project or project_id,
vertex_location=vertex_location or "us-central1",
)
)
vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
vertex_batch_request=transformed_batch_request,
headers=headers,
sync_handler=sync_handler,
api_base=api_base,
vertex_location=vertex_location or "us-central1",
)
default_api_base: Final = self.create_vertex_batch_url(
vertex_location=vertex_location or "us-central1",
vertex_project=vertex_project or project_id,
@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM):
vertex_api_version="v1",
)
headers: Final = {
"Content-Type": "application/json; charset=utf-8",
"Authorization": f"Bearer {access_token}",
}
vertex_batch_request: Final[VertexAIBatchPredictionJob] = (
VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
request=create_batch_data
)
)
if _is_async is True:
return self._async_create_batch(
vertex_batch_request=vertex_batch_request,
@ -142,6 +178,77 @@ class VertexAIBatchPrediction(VertexLLM):
)
return vertex_batch_response
@staticmethod
def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str:
"""
Builds the GET url for resolving an endpoint resource (`projects/../endpoints/<id>`).
A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the
version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a
mount prefix in front of the full default path. The `:operation` suffix convention from
`_check_custom_proxy` does not apply to a plain resource GET.
"""
default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
if not api_base:
return default_endpoint_url
api_base_path: Final = urlparse(api_base).path.rstrip("/")
if api_base_path in ("/v1", "/v1beta1"):
return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url)
return api_base.rstrip("/") + urlparse(default_endpoint_url).path
def _resolve_fine_tuned_endpoint_model(
self,
vertex_batch_request: VertexAIBatchPredictionJob,
headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers
sync_handler: HTTPHandler,
api_base: str | None,
vertex_location: str,
) -> VertexAIBatchPredictionJob:
"""
A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only
accepts Model resources, so swap the endpoint resource for its deployed tuned model
(`projects/../locations/../models/<id>`) read from GET endpoints/<id>.
"""
model: Final = vertex_batch_request.get("model", "")
if "/endpoints/" not in model:
return vertex_batch_request
endpoint_url: Final = self._build_endpoint_resolution_url(
api_base=api_base,
model=model,
vertex_location=vertex_location,
)
# ``api_base`` can come from caller-supplied request kwargs, so wrap the
# fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata
# targets before the bearer token leaves the process (mirrors retrieve_batch).
fetched: Final[_FetchedResponseView] = {
"response": safe_get(
sync_handler,
endpoint_url,
headers=headers,
)
}
response: Final = fetched["response"]
if response.status_code != 200:
raise VertexAIError(
status_code=response.status_code,
message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}",
)
payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()}
deployed_models: Final = payload_view["payload"].get("deployedModels") or ()
deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else ""
if not deployed_model:
raise VertexAIError(
status_code=400,
message=(
f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model "
"resource to run batch predictions against"
),
)
resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model}
return resolved_request
async def _async_create_batch(
self,
vertex_batch_request: VertexAIBatchPredictionJob,

70
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -22,6 +22,8 @@ class VertexAIBatchTransformation:
def transform_openai_batch_request_to_vertex_ai_batch_request(
cls,
request: CreateBatchRequest,
vertex_project: str | None = None,
vertex_location: str | None = None,
) -> VertexAIBatchPredictionJob:
"""
Transforms OpenAI Batch requests to Vertex AI Batch requests
@ -31,7 +33,11 @@ class VertexAIBatchTransformation:
if input_file_id is None:
raise ValueError("input_file_id is required, but not provided")
input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl")
model: Final[str] = cls._get_model_from_gcs_file(input_file_id)
model: Final[str] = cls._get_batch_job_model(
input_file_id=input_file_id,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
output_config: Final[OutputConfig] = OutputConfig(
predictionsFormat="jsonl",
gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)),
@ -188,6 +194,33 @@ class VertexAIBatchTransformation:
path_parts: Final = input_file_id.rsplit("/", 1)
return path_parts[0]
@classmethod
def _get_batch_job_model(
cls,
input_file_id: str,
vertex_project: str | None,
vertex_location: str | None,
) -> str:
"""
Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or
the full `projects/../locations/../endpoints/<id>` resource name for a fine-tuned endpoint.
The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource
to its deployed tuned model (`projects/../locations/../models/<id>`) before sending the job.
"""
parsed_model: Final = cls._get_model_from_gcs_file(input_file_id)
if not parsed_model.startswith("endpoints/"):
return parsed_model
if not vertex_project:
raise VertexAIError(
status_code=400,
message=(
f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require "
"`vertex_project` to build the endpoint resource name"
),
)
return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}"
@classmethod
def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str:
"""
@ -202,6 +235,9 @@ class VertexAIBatchTransformation:
gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8
returns: "publishers/google/models/gemini-1.5-flash-001"
Fine-tuned Gemini endpoints are stored as `endpoints/<numeric id>` in the uri and returned
in that form.
Raises a 400 `VertexAIError` when the uri carries no parseable model path.
"""
model: Final = cls._parse_model_from_gcs_file(gcs_file_uri)
@ -210,11 +246,13 @@ class VertexAIBatchTransformation:
status_code=400,
message=(
"Vertex AI batch creation requires the model to be part of `input_file_id`, but "
f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' path segment. "
f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' or "
"'endpoints/<numeric endpoint id>' path segment. "
"Either upload the input file through LiteLLM (POST /v1/files with "
"custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or "
"pass a uri of the form "
"gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file>"
"gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file> "
"(or gs://<bucket>/<prefix>/endpoints/<numeric endpoint id>/<file> for fine-tuned models)"
),
)
return model
@ -222,18 +260,26 @@ class VertexAIBatchTransformation:
@classmethod
def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None:
"""
Returns the `publishers/<publisher>/models/<model>` path from a gcs uri, or None if the uri
does not contain one.
Returns the `publishers/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
gcs uri, or None if the uri does not contain one.
A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is
used, so a user-configured bucket prefix that happens to contain `endpoints/<digits>` cannot
override the model path LiteLLM appended after it.
"""
_, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
if not separator:
return None
unquoted_uri: Final = unquote(gcs_file_uri)
_, separator, model_path = unquoted_uri.partition("publishers/")
if separator:
parts: Final = model_path.split("/")
if len(parts) >= 3 and parts[1] == "models" and parts[2]:
return f"publishers/{'/'.join(parts[:3])}"
parts: Final = model_path.split("/")
if len(parts) < 3 or parts[1] != "models" or not parts[2]:
return None
_, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/")
endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else ""
if endpoint_id.isdigit():
return f"endpoints/{endpoint_id}"
return f"publishers/{'/'.join(parts[:3])}"
return None
@classmethod
def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:

13
litellm/llms/vertex_ai/common_utils.py
View file

@ -370,6 +370,19 @@ def get_vertex_base_model_name(model: str) -> str:
return model
def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None:
"""
Fine-tuned Gemini deployments are addressed by a numeric endpoint id,
configured as `vertex_ai/<id>` or `vertex_ai/gemini/<id>`.
Returns the endpoint id, or None when `model` is a regular publisher model.
Mirrors the online chat path in `_get_vertex_url`, which sends numeric
models to `endpoints/{id}` instead of `publishers/google/models/{model}`.
"""
candidate: Final = model.split("/")[-1] if "gemini/" in model else model
return candidate if candidate.isdigit() else None
def validate_vertex_location(vertex_location: str | None) -> str:
"""
Validate a Vertex AI location before interpolating it into a request host or

51
litellm/llms/vertex_ai/files/transformation.py
View file

@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import (
)
from litellm.llms.vertex_ai.common_utils import (
_convert_vertex_datetime_to_openai_datetime,
get_vertex_ai_fine_tuned_endpoint_id,
)
from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@ -707,20 +708,39 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _get_gcs_object_name_from_batch_jsonl(
self,
openai_jsonl_content: list[dict[str, Any]],
deployment_model: str | None = None,
) -> str:
"""
Gets a unique GCS object name for the VertexAI batch prediction job
named as: litellm-vertex-{model}-{uuid}
The stored model path decides which Vertex model the batch job later executes against, so
`deployment_model` (the deployment's own configured model) wins over the user-supplied
JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no
deployment config.
Fine-tuned Gemini deployments (numeric endpoint ids) are stored under
`endpoints/<id>` so the batch transformation can round-trip them into a
`projects/../locations/../endpoints/<id>` batch job model instead of a
nonexistent publisher model.
"""
_model = openai_jsonl_content[0].get("body", {}).get("model", "")
if "publishers/google/models" not in _model:
_model = f"publishers/google/models/{_model}"
safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model")
raw_model: Final = (
deployment_model.removeprefix("vertex_ai/")
if deployment_model
else openai_jsonl_content[0].get("body", {}).get("model", "")
)
endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
model_path: Final = (
f"endpoints/{endpoint_id}"
if endpoint_id is not None
else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}")
)
safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model")
object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}"
return object_name
def get_object_name(self, file_data: FileTypes, purpose: str) -> str:
def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str:
"""
Get the object name for the request.
@ -728,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
upload is never materialized just to derive the GCS object name.
"""
if purpose == "batch":
## 1. If jsonl, derive the object name from the first entry's model
## 1. If jsonl, derive the object name from the deployment model (or the first entry's)
first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None)
if first_entry is not None:
return self._get_gcs_object_name_from_batch_jsonl([first_entry])
return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model)
## 2. If not jsonl, store under a server-generated managed object name
filename, _ = extract_file_metadata(file_data)
@ -761,6 +781,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
Get the complete url for the request
"""
if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"):
raise VertexAIError(
status_code=400,
message=(
"Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
"The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
"remove this deployment from the batch request (e.g. `target_model_names`) or "
"use a publisher model / fine-tuned Gemini endpoint instead."
),
)
bucket_name = self._get_configured_bucket_name(litellm_params)
bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name)
file_data: Final = data.get("file")
@ -769,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
raise ValueError("file is required")
if purpose is None:
raise ValueError("purpose is required")
object_name = self.get_object_name(file_data, purpose)
configured_model: Final = litellm_params.get("model")
object_name = self.get_object_name(
file_data,
purpose,
deployment_model=configured_model if isinstance(configured_model, str) else None,
)
if object_prefix:
object_name = f"{object_prefix}/{object_name}"
encoded_object_name: Final = encode_gcs_object_name_for_url(object_name)

13
litellm/models/mcp_server.py
View file

@ -6,10 +6,12 @@ Canonical definition for ``litellm_mcpservertable``. Re-exported from
"""
import enum
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Literal
from pydantic import Field
from pydantic import Field, ValidationInfo, field_validator
from litellm.types.llms.base import LiteLLMPydanticObjectBase
from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType
@ -115,3 +117,12 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
submitted_at: datetime | None = None
reviewed_at: datetime | None = None
review_notes: str | None = None
@field_validator("static_headers", "env", mode="before")
@classmethod
def decode_stored_secret_map(cls, value: object, info: ValidationInfo) -> Mapping[str, str] | None:
from litellm.proxy.common_utils.encrypt_decrypt_utils import decode_secret_map
if value is None and info.field_name == "env":
return MappingProxyType({})
return decode_secret_map(value, key=info.field_name or "secret map")

79
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -3108,15 +3108,17 @@ class MCPRequestHandler:
@staticmethod
async def _get_allowed_mcp_servers_for_agent(
user_api_key_auth: UserAPIKeyAuth | None = None,
agent_object_permission=None,
agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
) -> list[str]:
"""
Get allowed MCP servers for an agent (from the agent's object_permission).
Returns the MCP servers from the agent's object_permission.
If agent has no object_permission, returns [] (no extra restriction). An entitlement the
agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the
resolver denies.
Returns the agent's direct servers, the servers in its access groups, and the servers reached
through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no
object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that
cannot be read, or a declared toolset that resolves to no grants, raises
``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the
agent as unrestricted.
Args:
user_api_key_auth: User auth with agent_id
@ -3126,31 +3128,30 @@ class MCPRequestHandler:
if not user_api_key_auth or not user_api_key_auth.agent_id:
return []
obj_perm = agent_object_permission
if obj_perm is None:
obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
obj_perm: Final = (
agent_object_permission
if agent_object_permission is not None
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
)
if obj_perm is None:
return []
try:
direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or []
if isinstance(direct_mcp_servers, str):
direct_mcp_servers = []
mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or []
if isinstance(mcp_access_groups, str):
mcp_access_groups = []
# Permission entries may be server_ids OR names/aliases — expand to ids.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers))
access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups)
all_servers: Final = expanded_direct_servers + access_group_servers
return list(set(all_servers))
expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(
obj_perm.mcp_servers or []
)
access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(
obj_perm.mcp_access_groups or []
)
toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm)
return list({*expanded_direct_servers, *access_group_servers, *toolset_grants})
except Exception as e:
if isinstance(e, UnloadableEntitlementError):
raise
verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
return []
@ -3158,13 +3159,15 @@ class MCPRequestHandler:
async def _get_agent_tool_permissions_for_server(
server_id: str,
user_api_key_auth: UserAPIKeyAuth | None = None,
agent_object_permission=None,
agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
) -> list[str] | None:
"""
Get allowed tool names for a server from the agent's object_permission.
Returns None if agent has no tool restrictions for this server. An entitlement the agent
LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the
tool resolver turns into deny-all for the server rather than an unrestricted tool list.
Get allowed tool names for a server from the agent's object_permission: the union of its
direct tool permissions and the tools its toolsets grant on that server, mirroring the key and
team levels. Returns None if agent has no tool restrictions for this server. An entitlement the
agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises
``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the
server rather than an unrestricted tool list.
Args:
server_id: Server ID to check permissions for
@ -3175,24 +3178,30 @@ class MCPRequestHandler:
if not user_api_key_auth or not user_api_key_auth.agent_id:
return None
obj_perm = agent_object_permission
if obj_perm is None:
obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
obj_perm: Final = (
agent_object_permission
if agent_object_permission is not None
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
)
if obj_perm is None:
return None
try:
mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None)
if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict):
return None
# Dict keys may be server_ids OR names/aliases; normalize before lookup.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id)
return list(tools) if tools else None
direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
if obj_perm.mcp_tool_permissions
else None
)
toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id)
agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
return list(agent_tools) if agent_tools else None
except Exception as e:
if isinstance(e, UnloadableEntitlementError):
raise
verbose_logger.warning("Failed to get agent tool permissions for server: %s", e)
return None

77
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -4,7 +4,7 @@ import hashlib
import json
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -23,8 +23,11 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
SecretMapDecodeError,
_get_salt_key,
decode_secret_map,
decrypt_value_helper,
encrypt_secret_map,
encrypt_value_helper,
)
from litellm.proxy.utils import PrismaClient
@ -122,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False):
server_id: str
OAuthGrantState = Literal["valid", "refreshable", "absent"]
class _OAuthTokenRefreshResponse(TypedDict, total=False):
access_token: str
refresh_token: str
@ -360,7 +366,7 @@ def _prepare_mcp_server_data(
# exclude_unset filter is respected. Reading back from ``data`` would
# reintroduce defaults (e.g. ``env={}``) for fields the caller never set.
if data_dict.get("static_headers") is not None:
data_dict["static_headers"] = safe_dumps(data_dict["static_headers"])
data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"])
# env_vars is read from ``data_dict`` (not ``data``) like every other JSON
# column so the exclude_unset filter is respected: a partial update that
@ -376,7 +382,7 @@ def _prepare_mcp_server_data(
data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"])
if data_dict.get("env") is not None:
data_dict["env"] = safe_dumps(data_dict["env"])
data_dict["env"] = encrypt_secret_map(data_dict["env"])
if "tool_name_to_display_name" in data_dict:
data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {})
@ -589,6 +595,19 @@ def decrypt_credentials(
return credentials
def _readable_mcp_servers(
rows: Iterable["prisma_db_models.LiteLLM_MCPServerTable"],
) -> Iterable[LiteLLM_MCPServerTable]:
for row in rows:
try:
table = LiteLLM_MCPServerTable.model_validate(row.model_dump())
except SecretMapDecodeError:
verbose_proxy_logger.warning("Skipping MCP server %s: cannot decrypt secret map", row.server_id)
continue
decrypt_global_env_var_values(table.env_vars)
yield table
async def get_all_mcp_servers(
prisma_client: PrismaClient,
approval_status: str | None = None,
@ -609,10 +628,7 @@ async def get_all_mcp_servers(
)
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
for table in tables:
decrypt_global_env_var_values(table.env_vars)
return tables
return list(_readable_mcp_servers(mcp_servers))
async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None:
@ -638,13 +654,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
"server_id": {"in": server_ids},
}
)
final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = []
for _mcp_server in _mcp_servers:
table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump())
decrypt_global_env_var_values(table.env_vars)
final_mcp_servers.append(table)
return final_mcp_servers
return list(_readable_mcp_servers(_mcp_servers))
async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]:
@ -852,12 +862,10 @@ async def create_mcp_server(
data_dict["created_by"] = touched_by
data_dict["updated_by"] = touched_by
new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
)
new_mcp_server: Final = await MCPServerRepository(prisma_client).table.create(data=data_dict)
_decrypt_env_vars_on_returned_row(new_mcp_server)
return new_mcp_server
return LiteLLM_MCPServerTable.model_validate(new_mcp_server.model_dump())
async def create_draft_mcp_server(
@ -1066,13 +1074,13 @@ async def update_mcp_server(
data_dict["credentials"] = Json(None)
updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
updated_mcp_server: Final = await MCPServerRepository(prisma_client).table.update(
where={"server_id": data.server_id},
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
data=data_dict,
)
_decrypt_env_vars_on_returned_row(updated_mcp_server)
return updated_mcp_server
return LiteLLM_MCPServerTable.model_validate(updated_mcp_server.model_dump()) if updated_mcp_server else None
async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None:
@ -1144,6 +1152,13 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
if rotated_env_vars is not None:
update_data["env_vars"] = safe_dumps(rotated_env_vars)
for field in ("static_headers", "env"):
try:
if secret_map := decode_secret_map(getattr(mcp_server, field, None), key=field):
update_data[field] = encrypt_secret_map(secret_map, new_encryption_key=new_master_key)
except SecretMapDecodeError:
verbose_proxy_logger.warning("Cannot rotate MCP %s for server %s", field, mcp_server.server_id)
if not update_data:
continue
@ -1453,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in
return False
def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState:
"""Classify local grant readiness without attempting a refresh or checking upstream revocation."""
if not cred or not cred.get("access_token"):
return "absent"
if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
return "valid"
return "refreshable" if cred.get("refresh_token") else "absent"
async def get_user_oauth_credential(
prisma_client: PrismaClient,
user_id: str,
@ -1715,12 +1739,11 @@ async def resolve_valid_user_oauth_token(
dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh
actually happens, so the valid-token path never requires a DB handle.
"""
if not cred or not cred.get("access_token"):
grant: Final = oauth_grant_state(cred)
if cred is None or grant == "absent":
return None
if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
if grant == "valid":
return cred
if not cred.get("refresh_token"):
return None
if prisma_client is None:
from litellm.proxy.utils import get_prisma_client_or_throw
@ -1894,9 +1917,7 @@ async def get_mcp_submissions(
order={"submitted_at": "desc"},
take=500, # safety cap; paginate if needed in a future iteration
)
items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows]
for item in items:
decrypt_global_env_var_values(item.env_vars)
items: Final = list(_readable_mcp_servers(rows))
pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review)
active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active)

69
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import (
render_token_fault,
)
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
VendorCredentialState,
aggregate_authorize,
aggregate_token,
complete_connect_flow,
describe_connect_flow,
introspect_gateway_token,
is_gateway_dcr_client_id,
is_proxy_api_resource,
@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC
return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302)
async def _bridge_authorize_access_denial(
litellm_user_id: str,
mcp_server: MCPServer,
redirect_uri: str,
state: str,
) -> RedirectResponse | None:
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.
Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the
same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting
session can actually list and call the server's tools. Without this gate the flow completes, the
client shows connected, and every tool request fail-closes to an empty list with nothing telling
the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or
deactivated user denies like a missing grant, fail closed.
"""
async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial(
)
try:
admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id)
admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id)
except HTTPException as exc:
if exc.status_code >= 500:
raise
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
if mcp_server.server_id in allowed_server_ids:
return False
return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
async def _bridge_authorize_access_denial(
litellm_user_id: str,
mcp_server: MCPServer,
redirect_uri: str,
state: str,
) -> RedirectResponse | None:
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
return None
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
@ -1910,6 +1907,38 @@ async def token_endpoint(
)
async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState:
"""Whether the gateway itself can see a live vendor credential for this user and server.
The one reading of "authorized" the connect page displays and the finish step enforces, so
the button a user sees and the grant they get cannot disagree. A read fault is neither, and
fails the scoped grant closed."""
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # circular import at module load
get_user_oauth_credential,
oauth_grant_state,
)
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load
if prisma_client is None:
return "unavailable"
try:
credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed
return "unavailable"
return "absent" if oauth_grant_state(credential) == "absent" else "present"
@router.get("/authorize/flow")
async def authorize_flow(request: Request, flow: str) -> Response:
return await describe_connect_flow(
request=request,
flow_handle=flow,
session_user_id=_session_cookie_user_id(request),
lookup_vendor_credential=_vendor_credential_state,
lookup_server_reachability=_user_can_reach_mcp_server,
)
@router.post("/authorize/complete")
async def authorize_complete(
request: Request,
@ -1934,6 +1963,8 @@ async def authorize_complete(
delivery=delivery,
team_id=team_id,
decision=decision,
lookup_vendor_credential=_vendor_credential_state,
lookup_server_reachability=_user_can_reach_mcp_server,
)

187
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code"
ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"]
ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]]
"""Injected live-user revalidation (the token endpoint's mirror of admission):
``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is
a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else
fails the grant closed."""
VendorCredentialState = Literal["present", "absent", "unavailable"]
"""The per-user vendor credential read has three outcomes: present, absent, or unavailable."""
_DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry"
_DB_FAULTED_DESCRIPTION: Final = (
@ -195,6 +193,16 @@ class ConsentTeam(BaseModel):
team_alias: str | None = None
class LookupVendorCredential(Protocol):
"""Injected read of a user's vendor credential for one server."""
def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ...
class LookupServerReachability(Protocol):
def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ...
class LookupConsentTeams(Protocol):
"""Injected lookup of the teams a signed-in user may bind a proxy-API credential to."""
@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
return "unresolvable"
async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
return "unavailable"
async def _unreachable_server(user_id: str, server_id: str) -> bool:
return False
class GatewayDcrClient(BaseModel):
"""The registration record sealed into a gateway DCR ``client_id``.
@ -449,7 +465,10 @@ def aggregate_authorize(
A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the
flow to that one server: the scope is sealed into the flow, carried into the code, and
bound into the session token, while the connect page interlude runs exactly as before.
bound into the session token. The connect URL carries only the flow handle; the page
learns the client origin, the scoped server, and whether its vendor OAuth is done from
:func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry
steers which server the page authorizes or names on the confirmation.
Validation failures respond directly with 400 and never redirect: per RFC 6749
section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and
@ -474,10 +493,7 @@ def aggregate_authorize(
resource_server_id=scoped_server.server_id if scoped_server is not None else None,
audience=None,
)
connect_url: Final = _append_query_params(
f"{base_url}/ui/connect",
(("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))),
)
connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
response: Final = RedirectResponse(connect_url, status_code=303)
_set_flow_cookie(response, request, handle, flow)
return response
@ -684,6 +700,99 @@ def _origin_only(url: str) -> str:
return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else ""
def _open_flow_for(
request: Request, flow_handle: str, session_user_id: str | None, now: datetime
) -> _ConnectFlow | Response:
sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
if sealed_flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
if flow is None or now.timestamp() >= flow.exp:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
if session_user_id is None:
return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
if session_user_id != flow.user_id:
return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
return flow
async def _flow_target(
flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability
) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]:
if flow.resource_server_id is None:
return "unscoped", None
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle
MCPServerManager,
global_mcp_server_manager,
)
server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id)
if (
server is None
or not server.is_gateway_managed_oauth2
or not await lookup_server_reachability(flow.user_id, server.server_id)
):
return "stale", None
state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive"
return state, server
class ConnectFlowDescription(TypedDict):
"""What the connect page is allowed to know about one in-flight flow."""
state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]]
client_origin: ReadOnly[str]
server_id: ReadOnly[str | None]
server_name: ReadOnly[str | None]
connected: ReadOnly[bool | None]
async def _describe_opened_flow(
flow: _ConnectFlow,
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> ConnectFlowDescription | Response:
state, server = await _flow_target(flow, lookup_server_reachability)
if state == "interactive" and server is not None:
credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id)
if credential == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
interactive_description: Final[ConnectFlowDescription] = {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
"server_id": server.server_id,
"server_name": server.server_name or server.alias or server.name,
"connected": credential == "present",
}
return interactive_description
described: Final[ConnectFlowDescription] = {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
"server_id": None if server is None else server.server_id,
"server_name": None if server is None else (server.server_name or server.alias or server.name),
"connected": state == "m2m" or None,
}
return described
async def describe_connect_flow(
request: Request,
flow_handle: str,
session_user_id: str | None,
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> Response:
opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc))
if isinstance(opened, Response):
return opened
described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
return (
described
if isinstance(described, Response)
else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS)
)
async def complete_connect_flow(
request: Request,
flow_handle: str,
@ -692,56 +801,34 @@ async def complete_connect_flow(
delivery: str | None = None,
team_id: str | None = None,
decision: str | None = None,
lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential,
lookup_server_reachability: LookupServerReachability = _unreachable_server,
) -> Response:
"""The deliberate finish step of the connect flow: mint the gateway authorization
code and send the browser back to the client.
"""Mint the code only after a deliberate POST by the sealed user.
Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly
per-flow cookie plus an exact match between the signed-in user and the user sealed
into the flow: a link crafted by another party dies here with ``access_denied``
instead of minting a code for the victim's identity. The flow is single-use (an atomic
claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
``delivery`` chooses how the code reaches the client. Default (absent or
``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"``
renders the callback URL on a page instead, for a client whose redirect URI is a
loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box,
container): the 303 would dereference the browser machine's loopback and the code
would never arrive, so the user carries it over by pasting the URL into the client or
fetching it from the client machine's terminal. Manual delivery is honored only for
loopback redirect URIs; a routable redirect URI works from any browser by
construction, so those flows always redirect. The user who sees the page is exactly
the user the 303 would have carried the code to, and the same user already sees the
code today in the dead redirect's address bar, so the page exposes the code to no new
party. Unknown ``delivery`` values are rejected rather than defaulted: a client that
asked for manual delivery and got a dead redirect instead would silently lose its
code.
``decision`` and ``team_id`` come from the native-client consent page. ``"deny"``
burns the flow and sends the client ``error=access_denied`` so it stops waiting;
``team_id`` is sealed into the code only for proxy-API flows, where it picks which of
the user's teams the minted credential is attributed to.
A scoped flow additionally requires its sealed server to have a live vendor credential
before a code can be minted. The check happens before the single-use claim, so a
premature submit can be retried after authorization; denial deliberately bypasses it.
"""
if delivery not in (None, "redirect", "manual"):
return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'")
if decision not in (None, "approve", "deny"):
return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'")
sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
if sealed_flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
if flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
now: Final = datetime.now(timezone.utc)
if now.timestamp() >= flow.exp:
return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection")
if session_user_id is None:
return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
if session_user_id != flow.user_id:
return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
opened: Final = _open_flow_for(request, flow_handle, session_user_id, now)
if isinstance(opened, Response):
return opened
if decision != "deny":
described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
if isinstance(described, Response):
return described
if described["state"] == "stale":
return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available")
if described["connected"] is False:
return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing")
flow_refusal: Final = _claim_refusal(
await _SingleUseGuard(cache).claim(
f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
),
replayed=_oauth_error(
400, "invalid_request", "this connect flow was already completed; restart the connection"
@ -750,7 +837,7 @@ async def complete_connect_flow(
if flow_refusal is not None:
return flow_refusal
response: Final = (
_denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now)
_denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now)
)
path, secure = _cookie_path_and_secure(request)
response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")

10
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -6272,8 +6272,7 @@ class MCPServerManager:
]
}
)
db_mcp_servers: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows]
verbose_logger.info("Found %s MCP servers in database", len(db_mcp_servers))
verbose_logger.info("Found %s MCP servers in database", len(raw_rows))
previous_registry: Final = self.registry
new_registry: Final[dict[str, MCPServer]] = {}
@ -6281,8 +6280,9 @@ class MCPServerManager:
# Stage one: build every server. Stage two assigns short prefixes
# against the *full* set so dedup is deterministic regardless of
# iteration order.
for server in db_mcp_servers:
for row in raw_rows:
try:
server = LiteLLM_MCPServerTable.model_validate(row.model_dump())
existing_server = previous_registry.get(server.server_id)
if (
@ -6320,8 +6320,8 @@ class MCPServerManager:
except Exception as e:
verbose_logger.exception(
"Skipping MCP server %s (%s) during DB reload: %s",
server.server_id,
getattr(server, "alias", None),
getattr(row, "server_id", None),
getattr(row, "alias", None),
e,
)

105
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict
from starlette.requests import Request as StarletteRequest
from starlette.responses import JSONResponse
from starlette.types import Message, Receive, Scope, Send
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
@ -816,6 +817,11 @@ if MCP_AVAILABLE:
}
}
return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta})
except HTTPException as e:
from mcp.shared.exceptions import McpError
from mcp.types import INVALID_REQUEST, ErrorData
raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e
except Exception as e:
verbose_logger.exception("Error in list_tools endpoint: %s", e)
# Return empty list instead of failing completely
@ -1095,6 +1101,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_client_ip,
host_progress_callback=host_progress_callback,
**data, # for logging
)
@ -1128,7 +1135,7 @@ if MCP_AVAILABLE:
except HTTPException as e:
verbose_logger.error("HTTPException in MCP tool call: %s", e)
return CallToolResult(
content=[TextContent(text=f"Error: {e.detail}", type="text")],
content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")],
isError=True,
)
except MCPUpstreamAuthError as e:
@ -1392,7 +1399,7 @@ if MCP_AVAILABLE:
########################################################
async def _get_allowed_mcp_servers_from_mcp_server_names(
mcp_servers: list[str] | None,
mcp_servers: Sequence[str] | None,
allowed_mcp_servers: list[MCPServer],
) -> list[MCPServer]:
"""
@ -1413,13 +1420,10 @@ if MCP_AVAILABLE:
server_name_matched = False
for server in allowed_mcp_servers:
if server:
match_list = [s.lower() for s in iter_known_server_prefixes(server) if s]
if server_or_group.lower() in match_list:
filtered_server[server.server_id] = server
server_name_matched = True
break
if server and _server_answers_to(server, server_or_group):
filtered_server[server.server_id] = server
server_name_matched = True
break
if not server_name_matched:
try:
@ -1449,6 +1453,72 @@ if MCP_AVAILABLE:
return allowed_mcp_servers
def _http_detail_message(detail: object) -> str:
return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
def _server_answers_to(server: MCPServer, name: str) -> bool:
requested: Final = name.lower()
return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
class _McpDeniedDetail(TypedDict):
error: ReadOnly[str]
async def raise_denied_scoped_mcp_access(
requested_names: Sequence[str],
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None = None,
) -> None:
"""A scoped request (``/mcp/<name>`` path or ``x-mcp-servers`` header) resolved to zero
allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy
server with no tools. Unknown, unauthorized, and access-group names all share one generic
error so scoping cannot probe which servers exist; the agent variant fires only when the
same request resolves once the agent binding is stripped, proving the binding caused the veto."""
agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None
if user_api_key_auth is not None and agent_id:
resolved_without_agent: Final = await _get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})),
mcp_servers=requested_names,
client_ip=client_ip,
)
def _resolved_to_server(name: str) -> bool:
return any(_server_answers_to(server, name) for server in resolved_without_agent)
vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None)
if vetoed_server is not None:
agent_denial: Final[_McpDeniedDetail] = {
"error": (
f"MCP server '{vetoed_server}' is not available to this key: the key is bound to "
f"agent '{agent_id}', whose MCP grants do not include this server. Add the server "
f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or "
f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
)
}
raise HTTPException(status_code=403, detail=agent_denial)
vetoed_group: Final = next(
(
name
for name in requested_names
if not _resolved_to_server(name)
and any(name in (server.access_groups or ()) for server in resolved_without_agent)
),
None,
)
if vetoed_group is not None:
group_denial: Final[_McpDeniedDetail] = {
"error": (
f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to "
f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the "
f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or "
f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
)
}
raise HTTPException(status_code=403, detail=group_denial)
generic_denial: Final[_McpDeniedDetail] = {
"error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}"
}
raise HTTPException(status_code=403, detail=generic_denial)
def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
"""
Check if a tool name matches any name in the filter list.
@ -1541,7 +1611,7 @@ if MCP_AVAILABLE:
async def _get_allowed_mcp_servers(
user_api_key_auth: UserAPIKeyAuth | None,
mcp_servers: list[str] | None,
mcp_servers: Sequence[str] | None,
client_ip: str | None = None,
) -> list[MCPServer]:
"""Return allowed MCP servers for a request after applying filters.
@ -1977,6 +2047,12 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
client_ip=client_ip,
)
if mcp_servers and not allowed_mcp_servers:
await raise_denied_scoped_mcp_access(
requested_names=mcp_servers,
user_api_key_auth=user_api_key_auth,
client_ip=client_ip,
)
# Pre-fetch OAuth credentials only when at least one server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
@ -2404,6 +2480,8 @@ if MCP_AVAILABLE:
)
verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
return listing
except HTTPException:
raise
except Exception as e:
verbose_logger.exception("Error getting tools from managed MCP servers: %s", e)
# Continue with an empty listing instead of failing completely
@ -3086,6 +3164,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
**kwargs: Any,
) -> CallToolResult:
"""
@ -3116,6 +3195,12 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
allowed_mcp_servers=allowed_mcp_servers,
)
if mcp_servers and not allowed_mcp_servers:
await raise_denied_scoped_mcp_access(
requested_names=mcp_servers,
user_api_key_auth=user_api_key_auth,
client_ip=client_ip,
)
if not allowed_mcp_servers:
raise HTTPException(
status_code=403,

7
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -366,6 +366,7 @@ async def handle_mcp_tool_call(
from litellm.proxy._experimental.mcp_server.server import (
_get_allowed_mcp_servers,
execute_mcp_tool,
raise_denied_scoped_mcp_access,
)
allowed_mcp_servers: Final = await _get_allowed_mcp_servers(
@ -373,6 +374,12 @@ async def handle_mcp_tool_call(
mcp_servers=mcp_servers,
client_ip=client_ip,
)
if mcp_servers and not allowed_mcp_servers:
await raise_denied_scoped_mcp_access(
requested_names=mcp_servers,
user_api_key_auth=user_api_key_dict,
client_ip=client_ip,
)
# Reject before dispatch when the key has no accessible servers; otherwise an
# unprefixed local tool name would fall through to the local registry in

54
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -2634,6 +2634,20 @@
],
"title": "Mcp Tool Permissions"
},
"mcp_toolsets": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"title": "Mcp Toolsets"
},
"models": {
"anyOf": [
{
@ -19872,6 +19886,46 @@
]
}
},
"/authorize/flow": {
"get": {
"operationId": "authorize_flow_authorize_flow_get",
"parameters": [
{
"in": "query",
"name": "flow",
"required": true,
"schema": {
"title": "Flow",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"summary": "Authorize Flow",
"tags": [
"mcp_discoverable"
]
}
},
"/callback": {
"get": {
"description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.",

3
litellm/proxy/_types.py
View file

@ -2833,7 +2833,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"Enable only if your deployment is experiencing phantom "
"BudgetExceededError responses caused by leaked reservations "
"(see GitHub issue #27639). "
"A proxy-level WARNING is logged on every request while this flag "
"An INFO notice is logged once per worker at config load while this flag "
"is active as a reminder that hard enforcement is relaxed."
),
)
@ -3585,6 +3585,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
litellm_callback_params=[
"OTEL_EXPORTER",
"OTEL_ENDPOINT",
"OTEL_TRACES_ENDPOINT",
"OTEL_HEADERS",
],
)

20
litellm/proxy/auth/auth_utils.py
View file

@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status
from pydantic import PositiveInt, TypeAdapter, ValidationError
import litellm
from litellm import Router, provider_list
from litellm import Router, constants, provider_list
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY,
@ -1390,6 +1390,24 @@ def warn_once_if_custom_auth_skips_common_checks(
_custom_auth_common_checks_warning_emitted = True
def log_once_if_budget_reservation_disabled(
*,
disabled: bool,
logger: Logger = verbose_proxy_logger,
) -> None:
if constants.budget_reservation_disabled_info_emitted or not disabled:
return
logger.info(
"disable_budget_reservation is enabled: skipping optimistic budget "
"reservation. Budget enforcement is read-time only. Concurrent "
"requests can each pass the spend check before their cost is recorded, "
"so a configured budget may be briefly exceeded under high concurrency. "
"Set disable_budget_reservation to False or remove it to restore "
"hard per-request budget enforcement."
)
constants.budget_reservation_disabled_info_emitted = True # rebind-ok: process-wide one-shot sentinel
def is_pass_through_provider_route(route: str) -> bool:
PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [
"vertex-ai",

8
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2706,14 +2706,6 @@ async def _reserve_budget_after_common_checks(
if skip_budget_checks:
return
if general_settings.get("disable_budget_reservation") is True:
verbose_proxy_logger.warning(
"disable_budget_reservation is enabled: skipping optimistic budget "
"reservation. Budget enforcement is read-time only — concurrent "
"requests can each pass the spend check before their cost is recorded, "
"so a configured budget may be briefly exceeded under high concurrency. "
"Set disable_budget_reservation to False or remove it to restore "
"hard per-request budget enforcement."
)
return
from litellm.proxy.spend_tracking.budget_reservation import (

4
litellm/proxy/common_request_processing.py
View file

@ -3508,9 +3508,13 @@ class ProxyBaseLLMRequestProcessing:
error_body: Final = await http_status_error.response.aread()
error_text: Final = error_body.decode("utf-8")
error_headers: Final = { # mutable-ok: HTTPException takes a plain header dict
k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items()
}
raise HTTPException(
status_code=http_status_error.response.status_code,
detail={"error": error_text},
headers=error_headers,
)
error_msg: Final = f"{e}"
# Check for AttributeError in the exception chain.

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

@ -10,6 +10,7 @@ import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@ -507,6 +508,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
CLIENT_OUTPUT_CEILING_METADATA_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

40
litellm/proxy/common_utils/encrypt_decrypt_utils.py
View file

@ -1,7 +1,10 @@
import base64
import os
from collections.abc import Mapping
from typing import Final, Literal, cast
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
# Versioned ciphertext marker for AES-256-GCM values.
@ -203,3 +206,40 @@ def decrypt_value(value: bytes, signing_key: str) -> str:
return plaintext
except Exception as e:
raise e
class SecretMapDecodeError(RuntimeError):
pass
_SECRET_MAP: Final = TypeAdapter(Mapping[str, str])
_STORED_SECRET_MAP: Final = TypeAdapter(Mapping[str, str] | str)
_SECRET_STRING: Final = TypeAdapter(str)
def encrypt_secret_map(value: Mapping[str, str], new_encryption_key: str | None = None) -> str:
if not value:
return "{}"
ciphertext: Final = _SECRET_STRING.validate_python(
encrypt_value_helper(_SECRET_MAP.dump_json(value).decode(), new_encryption_key=new_encryption_key), strict=True
)
return _SECRET_STRING.dump_json(ciphertext).decode()
def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None:
if value is None:
return None
try:
stored: Final = (
_STORED_SECRET_MAP.validate_json(value, strict=True)
if isinstance(value, str) and value.lstrip().startswith(("{", '"'))
else _STORED_SECRET_MAP.validate_python(value, strict=True)
)
if not isinstance(stored, str):
return stored
decrypted: Final = decrypt_value_helper(
value=stored, key=key, exception_type="debug", return_original_value=False
)
return _SECRET_MAP.validate_json(decrypted, strict=True)
except ValidationError:
raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None

1
litellm/proxy/common_utils/registry_read_through.py
View file

@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool:
)
return proxy_server.llm_router is not None
async with proxy_server.MODEL_RECONCILE_LOCK:
await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client)
proxy_server.proxy_config._add_deployment(db_models=rows)
proxy_server.llm_model_list = router.get_model_list()
return True

10
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -75,6 +75,8 @@ SELECT
COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
COALESCE(SUM(spend), 0)::float8 AS spend,
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost,
COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns,
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
FROM windowed
GROUP BY router_name, router_type
@ -95,6 +97,7 @@ class AutoRouterTurnTransaction:
total_tokens: int
spend: float
saved_spend: float
classifier_cost: float
covered: bool
cache_hit: bool
cache_ttl_seconds: int | None
@ -225,6 +228,7 @@ def build_autorouter_turn_transaction(
total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0),
saved_spend=saved_spend,
classifier_cost=classifier_cost or 0.0,
covered=cache.covered,
cache_hit=cache.read_tokens > 0,
cache_ttl_seconds=cache.write_ttl_seconds,
@ -266,7 +270,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
last_model, models, turns, unordered_turns, covered_turns, cache_hits,
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns
)
VALUES (
{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
@ -277,13 +281,15 @@ VALUES (
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
{_TIER_DELTA}
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}
)
ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
turns = t.turns + 1,
total_tokens = t.total_tokens + EXCLUDED.total_tokens,
spend = t.spend + EXCLUDED.spend,
saved_spend = t.saved_spend + EXCLUDED.saved_spend,
classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost,
classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1,
covered_turns = t.covered_turns + EXCLUDED.covered_turns,
cache_hits = t.cache_hits + EXCLUDED.cache_hits,
ttl_5m_turns = t.ttl_5m_turns + EXCLUDED.ttl_5m_turns,

39
litellm/proxy/db/check_migration.py
View file

@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]:
def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
"""Checks for differences between current database and Prisma schema.
Never raises: a diff that cannot be produced, because the runner is missing,
because the command failed, or because it outlived its budget, is reported as
"no diff" so boot continues.
Returns:
A tuple containing:
- A boolean indicating if differences were found (True) or not (False).
- A string with the diff output or error message.
Raises:
subprocess.CalledProcessError: If the Prisma command fails.
Exception: For any other errors during execution.
- The SQL commands that would close the diff, empty when there is none.
"""
verbose_logger.debug("Checking for Prisma schema diff...")
try:
result: Final = subprocess.run(
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
prisma_command_timeout,
run_prisma,
)
except ImportError as e:
print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages
f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}"
)
return False, []
verbose_logger.debug("Checking for Prisma schema diff...")
timeout: Final = prisma_command_timeout()
try:
result: Final = run_prisma(
[
"prisma",
"migrate",
@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
"./schema.prisma",
"--script",
],
capture_output=True,
text=True,
check=True,
timeout=timeout,
env=os.environ.copy(),
)
# return True, "Migration diff generated successfully."
sql_commands: Final = extract_sql_commands(result.stdout)
if sql_commands:
@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]:
return True, sql_commands
else:
return False, []
except subprocess.TimeoutExpired:
print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages
f"Timed out after {timeout}s generating the migration diff. "
f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer."
)
return False, []
except subprocess.CalledProcessError as e:
error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}"
print(error_message) # noqa: T201

22
litellm/proxy/db/prisma_client.py
View file

@ -937,9 +937,17 @@ class PrismaManager:
use_v2_resolver=use_v2_resolver,
)
else:
try:
from litellm_proxy_extras.prisma_toolchain import (
prisma_command_timeout,
run_prisma,
)
except ImportError as e:
verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
return False
PrismaManager._raise_if_partitioned_spend_logs()
# Use prisma db push with increased timeout
subprocess.run(
run_prisma(
[
"prisma",
"db",
@ -947,13 +955,15 @@ class PrismaManager:
"--accept-data-loss",
"--skip-generate",
],
timeout=60,
check=True,
timeout=prisma_command_timeout(),
env=os.environ.copy(),
stdout=None,
stderr=None,
)
PrismaManager._apply_replica_identity_full_if_requested()
return True
except subprocess.TimeoutExpired:
verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1)
except subprocess.TimeoutExpired as e:
verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout)
time.sleep(random.randrange(5, 15))
except subprocess.CalledProcessError as e:
attempts_left = 3 - attempt

6
litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
View file

@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
for _tool_call, is_allowed, _rule_id, message in checked:
if not is_allowed and message is not None:
verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
verbose_proxy_logger.info("Tool Permission Guardrail: %s", message)
if self.on_disallowed_action == "block":
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name, message=message, blocked_content=True
@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
new_tools: Final = self._collect_request_tools(data)
if not new_tools:
verbose_proxy_logger.warning(
verbose_proxy_logger.debug(
"Tool Permission Guardrail: not running guardrail. No tools or functions in data"
)
return data
@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
is_allowed, _, message = self._check_tool_permission(tool_name, tool_type)
if not is_allowed and message is not None:
verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message)
verbose_proxy_logger.info("Tool Permission Guardrail: %s", message)
if self.on_disallowed_action == "block":
raise HTTPException(
status_code=400,

25
litellm/proxy/hooks/mcp_semantic_filter/hook.py
View file

@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference.
Reduces context window size and improves tool selection accuracy.
"""
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Optional
from fastapi import HTTPException
@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger):
)
# Parse to separate MCP tools from other tools
mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
if not mcp_tools:
return []
@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger):
return [name for name in names if name]
@staticmethod
def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]:
async def _narrow_mcp_references(
tools: Sequence[Mapping[str, object]],
selected_tool_names: list[str],
served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None,
) -> list[object]:
"""
Restrict each litellm_proxy MCP reference to the semantically selected tools.
@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger):
LiteLLM_Proxy_MCP_Handler,
)
via_gateway: Final = await (
LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, served_names)
if served_names is not None
else LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools)
)
return [
(
{**tool, "allowed_tools": selected_tool_names}
if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool])
else tool
)
for tool in tools
{**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool
for tool, routed in zip(tools, via_gateway, strict=True)
]
def _is_mcp_tool(self, tool: object) -> bool:
@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger):
filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools)
selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools)
narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names)
narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names)
data["tools"] = narrowed_tools
self._emit_filter_metadata_safe(
data=data,

13
litellm/proxy/litellm_pre_call_utils.py
View file

@ -18,6 +18,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
@ -325,7 +326,9 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
# ``attempted_fallbacks`` and ``original_model_group`` are written by the router
# and read by spend logs as fact; a client value has no legitimate meaning and no
# key or team setting keeps it, so the strip is never gated.
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset({"attempted_fallbacks", "original_model_group"})
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
)
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
# Request fields whose value, when URL-valued, becomes the outbound destination
@ -789,12 +792,6 @@ def apply_missing_session_id_policy(
)
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
extensions and the Agent SDK run through the same CLI and share that prefix."""
return user_agent.startswith("claude-cli/")
def is_codex_user_agent(user_agent: str) -> bool:
"""Codex builds its user agent as ``<originator>/<version> ...`` and ships
several first-party originators: ``codex-tui``, ``codex_cli_rs``,
@ -811,6 +808,8 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c
requests routed to providers that reject them. An explicit drop_params
from the caller or in the operator's ``litellm_settings`` always wins
over this default."""
from litellm.llms.anthropic.common_utils import is_claude_code_user_agent
if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)):
return False
if "drop_params" in data:

6
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -484,6 +484,8 @@ class _SessionAggRow(BaseModel):
total_tokens: int
spend: float
saved_spend: float
classifier_cost: float
classifier_cost_recorded_turns: int
session_seconds: float
@ -520,6 +522,7 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0,
spend=row.spend,
saved_spend=row.saved_spend,
classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None,
baseline_spend=baseline_spend,
saved_pct=_pct(row.saved_spend, baseline_spend),
saved_per_session=row.saved_spend / sessions if sessions else 0.0,
@ -552,6 +555,7 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
avg_tokens_per_session=totals.avg_tokens_per_session,
spend=totals.spend,
saved_spend=totals.saved_spend,
classifier_cost=totals.classifier_cost,
baseline_spend=totals.baseline_spend,
saved_pct=totals.saved_pct,
saved_per_session=totals.saved_per_session,
@ -582,6 +586,8 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
total_tokens=sum(row.total_tokens for row in rows),
spend=sum(row.spend for row in rows),
saved_spend=sum(row.saved_spend for row in rows),
classifier_cost=sum(row.classifier_cost for row in rows),
classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows),
session_seconds=sum(row.session_seconds for row in rows),
)

41
litellm/proxy/management_endpoints/credential_migration.py
View file

@ -43,7 +43,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
_ALGO_AES_GCM,
_ENCRYPTION_ALGORITHM_SETTING,
_V2_GCM_PREFIX,
SecretMapDecodeError,
_get_salt_key,
decode_secret_map,
decrypt_value_helper,
encrypt_value_helper,
)
@ -65,6 +67,20 @@ class LocationReport:
# Used by --check (read-only classification):
legacy: int = 0 # nacl ciphertext still awaiting migration
def count(self, classification: ValueClass | None) -> None:
if classification is None:
return
self.scanned += 1
match classification:
case "migrated":
self.already_v2 += 1
case "legacy":
self.legacy += 1
case "undecryptable":
self.undecryptable += 1
case _:
self.plaintext += 1
def as_dict(self) -> dict[str, int]:
return {
"scanned": self.scanned,
@ -441,7 +457,7 @@ def _classify_callback_value(value: object) -> ValueClass:
_COVERED_TABLE_SPECS: Final = [
("model_table", "litellm_proxymodeltable", ("litellm_params",), ()),
("credentials", "litellm_credentialstable", ("credential_values",), ()),
("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars"), ()),
("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()),
("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)),
("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)),
]
@ -472,14 +488,18 @@ def _classify_into_report(report: LocationReport, value: str) -> None:
names, base URLs, …) do not decrypt and fall through to ``plaintext``, so
over-scanning a column is harmless to the residual count.
"""
report.scanned += 1
cls: Final = classify_value(value, key="scan")
if cls == "migrated":
report.already_v2 += 1
elif cls == "legacy":
report.legacy += 1
else: # plaintext / not-a-string
report.plaintext += 1
report.count(classify_value(value, key="scan"))
def _classify_secret_map(value: object, key: str) -> ValueClass | None:
try:
decoded: Final = decode_secret_map(value, key=key)
except SecretMapDecodeError:
return "undecryptable"
if not decoded:
return None
ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value
return "migrated" if is_migrated(ciphertext) else "legacy"
async def _scan_one_table(
@ -503,6 +523,9 @@ async def _scan_one_table(
raw = getattr(row, col, None)
if raw is None:
continue
if db_attr == "litellm_mcpservertable" and col in ("static_headers", "env"):
report.count(_classify_secret_map(raw, col))
continue
if isinstance(raw, str):
try:
raw = json.loads(raw)

9
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -946,7 +946,14 @@ async def handle_bedrock_count_tokens(
except BedrockError as e:
# Convert BedrockError to HTTPException for FastAPI
verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e)
raise HTTPException(status_code=e.status_code, detail={"error": e.message})
from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers
provider_headers: Final = getattr(getattr(e, "response", None), "headers", None)
raise HTTPException(
status_code=e.status_code,
detail={"error": e.message},
headers=get_response_headers(provider_headers) if provider_headers else None,
)
except HTTPException:
# Re-raise HTTP exceptions as-is
raise

104
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import time
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal
import litellm
@ -16,7 +16,11 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import independent_snapshot
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
independent_snapshot,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -25,6 +29,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineStep,
PipelineStepResult,
)
from litellm.types.utils import StandardLoggingGuardrailInformation
try:
from fastapi.exceptions import HTTPException
@ -118,6 +123,7 @@ class PipelineExecutor:
return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
if action == "block":
_carry_working_guardrail_information(working_data=working_data, request_data=data)
return PipelineExecutionResult(
terminal_action="block",
step_results=step_results,
@ -126,6 +132,7 @@ class PipelineExecutor:
)
if action == "modify_response":
_carry_working_guardrail_information(working_data=working_data, request_data=data)
return PipelineExecutionResult(
terminal_action="modify_response",
step_results=step_results,
@ -168,34 +175,33 @@ class PipelineExecutor:
verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail)
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input))
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
if use_unified:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
try:
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = (
"apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
)
if use_unified:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
if mode == "pre_call":
response = await target.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@ -233,6 +239,12 @@ class PipelineExecutor:
else:
verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e)
return ("error", None, str(e), e)
finally:
if hook_input is not data:
_append_guardrail_information(
request_data=data,
entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:],
)
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
@ -283,6 +295,40 @@ def _restore_request_guardrails(
return {**working_data, "metadata": stripped} # mutable-ok: request dict
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]:
bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source))
recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None
return recorded if isinstance(recorded, list) else []
def _append_guardrail_information(
request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data
entries: Sequence[StandardLoggingGuardrailInformation],
) -> None:
if not entries:
return
_, request_bucket = get_or_create_metadata_bucket(request_data)
existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY)
if isinstance(existing, list):
existing.extend(entries)
return
request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries)
def _carry_working_guardrail_information(
working_data: Mapping[str, object],
request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data
) -> None:
recorded: Final = _recorded_guardrail_information(working_data)
existing: Final = _recorded_guardrail_information(request_data)
if recorded is existing:
return
_append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing])
def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str:
"""
Map pipeline step outcome to the configured action.

35
litellm/proxy/prometheus_cleanup.py
View file

@ -8,10 +8,13 @@ from __future__ import annotations
import glob
import os
import re
from typing import Final
from litellm._logging import verbose_proxy_logger
_LIVE_GAUGE_PID: Final = re.compile(r"gauge_live[a-z]*_(\d+)\.db$")
def wipe_directory(directory: str) -> None:
"""Delete all .db files in the directory. Called once before workers fork."""
@ -38,3 +41,35 @@ def mark_worker_exit(worker_pid: int) -> None:
verbose_proxy_logger.info("Prometheus cleanup: marked worker %s as dead", worker_pid)
except Exception as e:
verbose_proxy_logger.warning("Failed to mark prometheus worker %s as dead: %s", worker_pid, e)
def _is_running(pid: int) -> bool:
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
except PermissionError:
return True
return True
def mark_dead_workers(directory: str) -> tuple[int, ...]:
"""Drop the live-gauge files of workers that no longer exist and return their pids.
Uvicorn's multi-worker supervisor has no exit hook, so a replacement worker calls this at startup; without it
a crashed worker's in-flight gauges stay in the aggregate forever.
"""
owners: Final = frozenset(
int(match.group(1))
for match in map(_LIVE_GAUGE_PID.search, glob.glob(os.path.join(directory, "gauge_live*_*.db")))
if match is not None
)
dead: Final = tuple(sorted(pid for pid in owners if pid != os.getpid() and not _is_running(pid)))
if not dead:
return dead
from prometheus_client import multiprocess
for pid in dead:
multiprocess.mark_process_dead(pid, path=directory)
verbose_proxy_logger.info("Prometheus cleanup: marked dead workers %s in %s", dead, directory)
return dead

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