scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit_4116_drop_params_string_coerce

Browse source 
# Conflicts:
#	tests/test_litellm/test_utils.py
This commit is contained in:
mateo-berri 2026-09-08 15:08:38 -07:00
commit 2400f1befe
297 changed files with 11368 additions and 4356 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"

13
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 \
@ -54,6 +55,7 @@ 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."
@ -289,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/

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

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

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

7
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,7 +1461,10 @@ 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"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"

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,

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)

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

21
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
@ -125,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
@ -1465,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,
@ -1727,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

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

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.

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

@ -10,8 +10,10 @@ 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,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
@ -507,6 +509,8 @@ 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,
ROUTING_REQUEST_TAGS_METADATA_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

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,

15
litellm/proxy/litellm_pre_call_utils.py
View file

@ -18,11 +18,13 @@ 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,
OTEL_SERVICE_NAME_METADATA_KEYS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
SESSION_ID_OMITTED_METADATA_KEY,
@ -289,6 +291,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
"standard_logging_object",
"proxy_server_request",
@ -325,7 +328,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 +794,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 +810,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),
)

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

5
litellm/proxy/prometheus_metrics_server.py
View file

@ -30,6 +30,7 @@ from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_a
from litellm.llms.custom_httpx.http_handler import HTTPHandler
METRICS_PATH: Final = "/metrics"
HEALTH_PATH: Final = "/health"
PID_HEADER: Final = "x-litellm-metrics-pid"
_PARENT_POLL_INTERVAL_SECONDS: Final = 1.0
_STARTUP_TIMEOUT_SECONDS: Final = 30.0
@ -77,6 +78,10 @@ def build_metrics_app(multiproc_dir: str) -> FastAPI:
app: Final = FastAPI(title="LiteLLM Prometheus metrics", docs_url=None, redoc_url=None, openapi_url=None)
app.mount(METRICS_PATH, _add_pid_header(make_metrics_asgi_app(registry)))
@app.get(HEALTH_PATH)
def health() -> dict[str, str]:
return {"status": "healthy", "multiproc_dir": multiproc_dir}
return app

44
litellm/proxy/proxy_server.py
View file

@ -307,6 +307,7 @@ from litellm.proxy.auth.auth_checks import (
from litellm.proxy.auth.auth_utils import (
check_response_size_is_safe,
is_request_body_safe,
log_once_if_budget_reservation_disabled,
warn_once_if_custom_auth_skips_common_checks,
)
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
@ -631,6 +632,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
router as pass_through_router,
)
from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit
from litellm.proxy.public_endpoints import router as public_endpoints_router
from litellm.proxy.public_endpoints.public_v1 import router as public_v1_router
from litellm.proxy.rag_endpoints.endpoints import router as rag_router
@ -1059,6 +1061,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
init_verbose_loggers()
prometheus_multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR")
if prometheus_multiproc_dir:
mark_dead_workers(prometheus_multiproc_dir)
## RUN WORKER STARTUP HOOKS (e.g., gflags initialization) ##
_startup_hooks_env: Final = os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "")
if _startup_hooks_env:
@ -1365,6 +1371,9 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
await proxy_shutdown_event(worker_heartbeat=worker_heartbeat)
if prometheus_multiproc_dir:
mark_worker_exit(os.getpid())
def _generate_stable_operation_id(route: "APIRoute") -> str:
operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}")
@ -5656,6 +5665,10 @@ class ProxyConfig:
run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)),
)
log_once_if_budget_reservation_disabled(
disabled=general_settings.get("disable_budget_reservation") is True,
)
custom_key_generate: Final = general_settings.get("custom_key_generate", None)
if custom_key_generate is not None:
user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path)
@ -7107,11 +7120,10 @@ class ProxyConfig:
],
)
# Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set)
if self._should_load_db_object(object_type="models"):
new_models: Final = await self._get_models_from_db(prisma_client=prisma_client)
# update llm router
load_models: Final = self._should_load_db_object(object_type="models")
new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None
await self.get_credentials(prisma_client=prisma_client)
if load_models:
still_desired_ids = await self._update_llm_router(
new_models=new_models, proxy_logging_obj=proxy_logging_obj
)
@ -7151,12 +7163,9 @@ class ProxyConfig:
async def _resync_config_from_db() -> None:
await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
async def _resync_credentials_from_db() -> None:
await self.get_credentials(prisma_client=prisma_client)
subscriber: Final = ConfigSyncSubscriber(
redis_cache=redis_cache,
resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db),
resync_callbacks=(_resync_config_from_db,),
)
self.config_sync_subscriber = subscriber
subscriber.start()
@ -8011,7 +8020,7 @@ class ProxyConfig:
async def get_credentials(self, prisma_client: PrismaClient):
try:
credentials = await CredentialsRepository(prisma_client).find_all()
credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all()
credentials = [self.decrypt_credentials(cred) for cred in credentials]
await self.delete_credentials(credentials) # delete credentials that are not in the all-up list
CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list
@ -9595,19 +9604,6 @@ class ProxyStartupEvent:
)
if store_model_in_db is True:
### GET STORED CREDENTIALS ###
scheduler.add_job(
proxy_config.get_credentials,
"interval",
seconds=config_reload_interval_seconds,
# REMOVED jitter parameter - major cause of memory leak
args=[prisma_client],
id="get_credentials_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
await proxy_config.get_credentials(prisma_client=prisma_client)
# MEMORY LEAK FIX: Increase interval from 10s to 30s minimum
# Frequent polling was causing excessive memory allocations
scheduler.add_job(
@ -9621,7 +9617,7 @@ class ProxyStartupEvent:
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
# this will load all existing models on proxy startup
# this will load all existing credentials and models on proxy startup
await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
proxy_config.start_config_sync_subscriber(

2
litellm/proxy/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])

15
litellm/proxy/spend_tracking/savings.py
View file

@ -299,6 +299,8 @@ def compute_autorouter_savings(
selected_info: ModelInfo | None = None,
baseline_info: ModelInfo | None = None,
cost_breakdown: Mapping[str, object] | None = None,
baseline_deployment_id: str | None = None,
selected_deployment_id: str | None = None,
) -> float:
"""Net dollars the router saved, or cost, by serving this request on ``selected_model``.
@ -334,11 +336,12 @@ def compute_autorouter_savings(
selected: Final = _resolve_model(selected_model, selected_provider)
if baseline is None or selected is None:
return 0.0
# Same model is only the same cost when it is also the same deployment. Two
# deployments of one model can carry different negotiated rates, and routing from
# the dear one to the cheap one is a real saving that short-circuiting on the model
# name alone reports as zero.
if baseline == selected:
same_target: Final = (
baseline_deployment_id == selected_deployment_id
if baseline_deployment_id and selected_deployment_id
else baseline == selected
)
if same_target:
return 0.0
basis: Final = _pricing_basis(cost_breakdown)
effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline)
@ -517,6 +520,8 @@ def autorouter_savings_for_request(
selected_info=_effective_model_info(router_instance, model_id, model or ""),
baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""),
cost_breakdown=cost_breakdown,
baseline_deployment_id=baseline_id,
selected_deployment_id=model_id,
)
classifier_cost: Final = classifier_cost_from_decision(decision)
return gross if classifier_cost is None else gross - classifier_cost

11
litellm/proxy/utils.py
View file

@ -177,6 +177,9 @@ from litellm.types.mcp import (
)
from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
from litellm.utils import (
_add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper
)
if TYPE_CHECKING:
from mcp.types import CallToolResult
@ -857,6 +860,14 @@ class ProxyLogging:
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
# Runs after load_config applied every litellm_settings key: logger __init__s read e.g. s3_callback_params
success_callbacks: Final = tuple(cb for cb in litellm.success_callback if isinstance(cb, str))
failure_callbacks: Final = tuple(cb for cb in litellm.failure_callback if isinstance(cb, str))
for callback in success_callbacks:
_add_custom_logger_callback_to_specific_event(callback, "success")
for callback in failure_callbacks:
_add_custom_logger_callback_to_specific_event(callback, "failure")
async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]):
# only use this if slack alerting is being used
if self.alerting is None:

8
litellm/responses/main.py
View file

@ -178,7 +178,7 @@ async def aresponses_api_with_mcp(
(
mcp_tools_with_litellm_proxy,
other_tools,
) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
# Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform)
# Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata)
@ -237,6 +237,7 @@ async def aresponses_api_with_mcp(
"timeout": timeout,
"custom_llm_provider": custom_llm_provider,
**kwargs,
"_skip_mcp_handler": True,
}
# Handle MCP streaming if requested
@ -899,13 +900,14 @@ def _responses_try_dispatch_mcp_gateway(
custom_llm_provider: str | None,
kwargs: dict[str, object],
_is_async: bool,
skip_mcp_handler: bool,
) -> Any | None:
"""Return a response when MCP gateway handles the call; otherwise None."""
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools):
return None
mcp_call_kwargs: Final = {
"input": input,
@ -1075,6 +1077,7 @@ def responses(
litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None)
_is_async: Final = kwargs.pop("aresponses", False) is True
skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False)
use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs)
client_headers: Final = kwargs.get("headers")
@ -1169,6 +1172,7 @@ def responses(
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
_is_async=_is_async,
skip_mcp_handler=skip_mcp_handler,
)
if _mcp_dispatch is not None:
return _mcp_dispatch

3
litellm/responses/mcp/chat_completions_handler.py
View file

@ -106,7 +106,7 @@ async def acompletion_with_mcp(
(
mcp_tools_with_litellm_proxy,
other_tools,
) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
if not mcp_tools_with_litellm_proxy:
# No MCP tools, proceed with regular completion
@ -114,6 +114,7 @@ async def acompletion_with_mcp(
model=model,
messages=messages,
tools=tools,
_skip_mcp_handler=True,
**kwargs,
)

153
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -1,6 +1,6 @@
import re
import traceback
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload
@ -11,6 +11,7 @@ from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._experimental.mcp_server.utils import (
iter_known_server_prefixes,
logging_safe_mcp_headers,
split_server_prefix_from_name,
strip_known_server_prefix,
@ -23,6 +24,7 @@ from litellm.types.llms.openai import (
ResponsesAPIStreamingResponse,
)
from litellm.types.llms.openai import ToolParam as ResponsesToolParam
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.utils import (
CallTypes,
ChatCompletionMessageCustomToolCall,
@ -45,6 +47,7 @@ else:
# NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling
ToolParam: TypeAlias = Mapping[str, object]
SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]]
class MCPToolResult(TypedDict):
@ -56,14 +59,74 @@ class MCPToolResult(TypedDict):
LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy"
LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/"
# Matches any URL whose path ends with /mcp/<server_name> — covers both root-path
# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments.
# A false-positive match (e.g. an external URL that happens to end with /mcp/<name>) results
# in a "server not found" error from the internal gateway, not a silent failure or data leak,
# so this broad pattern is intentional and preferred over anchoring to localhost only.
_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$")
def _mcp_server_url(tool: ToolParam) -> str | None:
if not isinstance(tool, dict) or tool.get("type") != "mcp":
return None
server_url: Final = tool.get("server_url")
return server_url if isinstance(server_url, str) else None
def _names_gateway_explicitly(tool: ToolParam) -> bool:
return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL)
def _proxy_path_mcp_name(tool: ToolParam) -> str | None:
server_url: Final = _mcp_server_url(tool)
match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url)
return None if match is None else match.group(1)
def _registered_mcp_servers() -> Collection[MCPServer]:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
return global_mcp_server_manager.get_registry().values()
def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool:
requested: Final = name.lower()
return any(
requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name))
or name in (server.access_groups or ())
for server in servers
)
async def _toolset_exists(name: str) -> bool:
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return False
return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None
except Exception as e:
verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e)
return False
async def _gateway_served_names(
names: Collection[str],
servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers,
toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists,
) -> frozenset[str]:
registered: Final = tuple(servers()) if names else ()
return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)])
async def _served_mcp_path_names(
tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]]
) -> frozenset[str]:
names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None)
return await served_names(names) if names else frozenset[str]()
class LiteLLM_Proxy_MCP_Handler:
"""
Helper class with static methods for MCP integration with Responses API.
@ -87,57 +150,41 @@ class LiteLLM_Proxy_MCP_Handler:
@staticmethod
def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool:
"""
Returns True if any MCP tool should be handled via the litellm proxy MCP gateway.
This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/<name>.
"""
if tools:
for tool in tools:
if isinstance(tool, dict) and tool.get("type") == "mcp":
server_url = tool.get("server_url", "")
if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
return True
if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url):
return True
return False
"""True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in
/mcp/<name>. `_split_mcp_tools` then settles which of the latter the gateway actually serves."""
return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ())
@staticmethod
def _parse_mcp_tools(
def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools:
items: Final = tuple(tools or ())
gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)]
other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)]
return gateway_tools, other_tools
@staticmethod
async def _split_mcp_tools(
tools: Iterable[Mapping[str, object]] | None,
) -> tuple[list[ToolParam], list[Any]]:
"""
Parse tools and separate MCP tools with litellm_proxy from other tools.
served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
) -> SplitTools:
items: Final = tuple(tools or ())
served: Final = await _served_mcp_path_names(items, served_names)
return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
[
{**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"}
if (name := _proxy_path_mcp_name(tool)) in served
else tool
for tool in items
]
)
Returns:
Tuple of (mcp_tools_with_litellm_proxy, other_tools)
"""
mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = []
other_tools: Final[list[Any]] = []
if tools:
for tool in tools:
if isinstance(tool, dict) and tool.get("type") == "mcp":
server_url = tool.get("server_url", "")
if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL):
mcp_tools_with_litellm_proxy.append(tool)
elif isinstance(server_url, str):
# Also intercept URLs like http://localhost:4000/mcp/atlassian_test
# by rewriting them to the internal litellm_proxy format.
m = _PROXY_MCP_PATH_RE.match(server_url)
if m:
rewritten = {
**tool,
"server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}",
}
mcp_tools_with_litellm_proxy.append(rewritten)
else:
other_tools.append(tool)
else:
other_tools.append(tool)
else:
other_tools.append(tool)
return mcp_tools_with_litellm_proxy, other_tools
@staticmethod
async def routes_through_gateway(
tools: Iterable[Mapping[str, object]] | None,
served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names,
) -> tuple[bool, ...]:
items: Final = tuple(tools or ())
served: Final = await _served_mcp_path_names(items, served_names)
return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items)
@staticmethod
async def _apply_toolset_permissions(

218
litellm/router.py
View file

@ -21,7 +21,16 @@ import time
import traceback
import weakref
from collections import defaultdict
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence
from collections.abc import (
AsyncGenerator,
AsyncIterator,
Callable,
Generator,
Iterator,
Mapping,
MutableMapping,
Sequence,
)
from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
@ -45,12 +54,15 @@ from litellm.caching.caching import (
RedisClusterCache,
)
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS,
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
DEFAULT_MAX_LRU_CACHE_SIZE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
OUTPUT_TOKEN_CEILING_PARAMS,
ROUTING_REQUEST_TAGS_METADATA_KEY,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
@ -648,6 +660,18 @@ class FallbackAwareStreamWrapper(CustomStreamWrapper):
self.fallback_headers_adopted = True
def as_output_cap(value: object) -> int | None:
"""A client-sent output cap coerced to an int: ints, floats and numeric strings, never bools
or negatives."""
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
return None
try:
cap: Final = int(float(value))
except (ValueError, OverflowError):
return None
return cap if cap >= 0 else None
class Router:
model_names: set = set()
cache_responses: bool | None = False
@ -955,7 +979,6 @@ class Router:
DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER
)
self.health_state_cache = DeploymentHealthCache(cache=self.cache, staleness_threshold=float(_staleness))
self.failed_calls = InMemoryCache() # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown
if num_retries is not None:
self.num_retries = num_retries
@ -1248,7 +1271,7 @@ class Router:
selector = LeastBusyLoggingHandler(router_cache=self.cache)
if register_callbacks:
if isinstance(litellm.input_callback, list):
litellm.input_callback.append(selector)
litellm.logging_callback_manager.add_litellm_input_callback(selector)
else:
litellm.input_callback = [selector]
case RoutingStrategy.USAGE_BASED_ROUTING.value:
@ -1523,9 +1546,33 @@ class Router:
self._override_selectors[strategy] = self._build_strategy_selector(
strategy=strategy,
routing_strategy_args={},
register_callbacks=False,
)
return self._override_selectors[strategy]
def _override_selector_pre_call_check(
self, strategy: str | None, selector: RouterStrategySelector | None, deployment: dict
) -> None:
"""
Override selectors are not in `litellm.callbacks`, so the pre-call check that
`routing_strategy_pre_call_checks` runs for the router's own selectors (rpm
accounting for `usage-based-routing-v2`) runs here, for the overriding request only.
"""
if selector is None or strategy is None or selector is not self._override_selectors.get(strategy):
return
selector.pre_call_check(deployment)
async def _async_override_selector_pre_call_check(
self,
strategy: str | None,
selector: RouterStrategySelector | None,
deployment: dict,
parent_otel_span: Span | None,
) -> None:
if selector is None or strategy is None or selector is not self._override_selectors.get(strategy):
return
await selector.async_pre_call_check(deployment, parent_otel_span)
def _get_routing_context(
self, model: str, request_kwargs: dict | None = None
) -> tuple[str | None, RouterStrategySelector | None]:
@ -3726,6 +3773,11 @@ class Router:
refund_stale_reservation_before_retry(self.cache, kwargs)
set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment))
kwargs[metadata_variable_name].setdefault(
ROUTING_REQUEST_TAGS_METADATA_KEY,
tuple(_get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)),
)
## DEPLOYMENT-LEVEL TAGS
deployment_tags: Final = deployment.get("litellm_params", {}).get("tags")
if deployment_tags:
@ -4214,10 +4266,12 @@ class Router:
}
)
litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object)
prompt_management_deployment: Final = self.get_available_deployment(
specific_deployment: Final = kwargs.pop("specific_deployment", None)
prompt_management_deployment: Final = await self.async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally
specific_deployment=specific_deployment,
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs)
@ -12610,7 +12664,83 @@ class Router:
request_kwargs.pop(carrier, None)
@staticmethod
def _drop_client_effort_carriers_a_tier_pin_supersedes(
def _tier_ceiling_under_the_surface_name(
tier_litellm_params: Mapping[str, object], responses_call: bool
) -> Mapping[str, object]:
"""``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are one
ceiling under three names, and each surface reads exactly one of them: the
Responses bridge builds its internal ``max_tokens`` from ``max_output_tokens``
and would overwrite the tier's, chat and /v1/messages never read
``max_output_tokens``, and litellm already renames ``max_tokens`` to
``max_completion_tokens`` for the OpenAI models that require it. Collapse
whatever the tier carries onto the surface's own name, preferring a value the
operator already wrote under that name."""
surface_key: Final = "max_output_tokens" if responses_call else "max_tokens"
carried: Final = tuple(
key
for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens")
if key in tier_litellm_params
)
if not carried:
return tier_litellm_params
return MappingProxyType(
{
**{k: v for k, v in tier_litellm_params.items() if k not in OUTPUT_TOKEN_CEILING_PARAMS},
surface_key: tier_litellm_params[carried[0]],
}
)
def _pin_tier_params_onto_request(
self,
model: str,
tier_litellm_params: Mapping[str, object] | None,
request_kwargs: dict,
responses_call: bool,
) -> bool:
"""Apply a routing strategy's per-tier litellm_params on top of the request and report
whether they pinned an output ceiling, so the caller can hand the request its own ceiling
back on a routing pass that pins none."""
if not tier_litellm_params:
return False
accepted_tier_params: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs)
surface_tier_params: Final = self._tier_ceiling_under_the_surface_name(
accepted_tier_params, responses_call=responses_call
)
self._drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface_tier_params)
request_kwargs.update(surface_tier_params)
return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params)
@staticmethod
def _restore_client_ceiling_no_tier_pins(request_kwargs: MutableMapping[str, object]) -> None:
"""A model-group fallback re-enters routing with the kwargs an earlier auto-router pass
already rewrote, so a ceiling sized for that pass's tier would ride onto a group no tier
chose. When this pass pins none, hand the request back exactly the carriers the caller
sent, which the first pinning pass stamped. The stamp lives in a metadata bucket a
caller can also write, so the proxy strips the key at ingestion and this read takes
nothing but the three ceiling carriers as integers: no other key ever reaches kwargs."""
stamped: Final = next(
(
bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY)
for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata"))
if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket
),
None,
)
if not isinstance(stamped, dict):
return
callers_ceiling: Final = MappingProxyType(
{
carrier: cap
for carrier, value in stamped.items()
if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None
}
)
for carrier in OUTPUT_TOKEN_CEILING_PARAMS:
request_kwargs.pop(carrier, None)
request_kwargs.update(callers_ceiling)
@staticmethod
def _drop_client_carriers_a_tier_pin_supersedes(
request_kwargs: dict[str, object],
tier_litellm_params: Mapping[str, object],
) -> None:
@ -12620,7 +12750,22 @@ class Router:
the ``reasoning_effort`` alias, so a pinned effort only reaches the wire
if the client's other encodings are removed before the merge. Non-effort
fields a carrier also holds (``output_config.format``,
``reasoning.summary``) are kept."""
``reasoning.summary``) are kept. An output ceiling has the same shape:
``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are
one setting under three names, and a provider handed two of them either
rejects the request or picks one by iteration order."""
if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params):
_, metadata_bucket = get_or_create_metadata_bucket(request_kwargs)
metadata_bucket.setdefault(
CLIENT_OUTPUT_CEILING_METADATA_KEY,
{
carrier: request_kwargs[carrier]
for carrier in OUTPUT_TOKEN_CEILING_PARAMS
if carrier in request_kwargs
},
)
for carrier in OUTPUT_TOKEN_CEILING_PARAMS:
request_kwargs.pop(carrier, None)
if "reasoning_effort" not in tier_litellm_params:
return
request_kwargs.pop("thinking", None)
@ -12661,6 +12806,7 @@ class Router:
# Execute Pre-Routing Hooks
# this hook can modify the model, messages before the routing decision is made
#########################################################
responses_call: Final = input is not None and messages is None
pre_routing_hook_response: Final = await self.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
@ -12672,12 +12818,14 @@ class Router:
model = pre_routing_hook_response.model
messages = pre_routing_hook_response.messages
record_pre_routing_selection(request_kwargs, model)
if pre_routing_hook_response.litellm_params:
accepted_tier_params: Final = self._tier_params_the_target_accepts(
model, pre_routing_hook_response.litellm_params, request_kwargs
)
self._drop_client_effort_carriers_a_tier_pin_supersedes(request_kwargs, accepted_tier_params)
request_kwargs.update(accepted_tier_params)
tier_pins_ceiling: Final = self._pin_tier_params_onto_request(
model=model,
tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None,
request_kwargs=request_kwargs,
responses_call=responses_call,
)
if not tier_pins_ceiling:
self._restore_client_ceiling_no_tier_pins(request_kwargs)
#########################################################
# Resolve the strategy and logger AFTER the pre-routing hook, since
@ -12694,10 +12842,16 @@ class Router:
parent_otel_span=parent_otel_span,
)
if isinstance(healthy_deployments, dict):
await self._async_override_selector_pre_call_check(
strategy, strategy_selector, healthy_deployments, parent_otel_span
)
return healthy_deployments
# When encrypted content affinity pins to a specific deployment,
if request_kwargs.get("_encrypted_content_affinity_pinned") and len(healthy_deployments) == 1:
await self._async_override_selector_pre_call_check(
strategy, strategy_selector, healthy_deployments[0], parent_otel_span
)
return healthy_deployments[0]
start_time: Final = time.time()
@ -12723,6 +12877,9 @@ class Router:
parent_otel_span=parent_otel_span,
)
raise exception
await self._async_override_selector_pre_call_check(
strategy, strategy_selector, deployment, parent_otel_span
)
verbose_router_logger.info(
"get_available_deployment for model: %s, Selected deployment: %s for model: %s",
model,
@ -12777,6 +12934,7 @@ class Router:
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs)
# 1. Execute pre-routing hook
responses_call: Final = input is not None and messages is None
pre_routing_hook_response: Final = await self.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
@ -12788,12 +12946,14 @@ class Router:
model = pre_routing_hook_response.model
messages = pre_routing_hook_response.messages
record_pre_routing_selection(request_kwargs, model)
if pre_routing_hook_response.litellm_params:
accepted_tier_params: Final = self._tier_params_the_target_accepts(
model, pre_routing_hook_response.litellm_params, request_kwargs
)
self._drop_client_effort_carriers_a_tier_pin_supersedes(request_kwargs, accepted_tier_params)
request_kwargs.update(accepted_tier_params)
tier_pins_ceiling: Final = self._pin_tier_params_onto_request(
model=model,
tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None,
request_kwargs=request_kwargs,
responses_call=responses_call,
)
if not tier_pins_ceiling:
self._restore_client_ceiling_no_tier_pins(request_kwargs)
# 2. Get healthy deployments
healthy_deployments: Final = await self.async_get_healthy_deployments(
@ -12805,6 +12965,8 @@ class Router:
parent_otel_span=parent_otel_span,
)
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
# 3. If specific deployment returned, verify if it supports pass-through
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
@ -12815,6 +12977,9 @@ class Router:
)
litellm_params: Final = healthy_deployments.get("litellm_params", {})
if litellm_params.get("use_in_pass_through"):
await self._async_override_selector_pre_call_check(
strategy, strategy_selector, healthy_deployments, parent_otel_span
)
return healthy_deployments
else:
raise litellm.BadRequestError(
@ -12835,7 +13000,6 @@ class Router:
# 5. Apply load balancing strategy
start_time: Final = time.perf_counter()
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
@ -12859,6 +13023,9 @@ class Router:
parent_otel_span=parent_otel_span,
)
raise exception
await self._async_override_selector_pre_call_check(
strategy, strategy_selector, deployment, parent_otel_span
)
verbose_router_logger.info(
"async_get_available_deployment_for_pass_through model: %s, selected deployment: %s",
@ -13434,6 +13601,7 @@ class Router:
specific_deployment=specific_deployment,
request_kwargs=request_kwargs,
)
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
@ -13442,6 +13610,7 @@ class Router:
model=model,
llm_provider="",
)
self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments)
return healthy_deployments
parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs)
@ -13517,7 +13686,6 @@ class Router:
cooldown_list=_cooldown_list,
)
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
if strategy == "simple-shuffle":
# if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm
############## Check 'weight' param set for weighted pick #################
@ -13549,6 +13717,7 @@ class Router:
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
)
self._override_selector_pre_call_check(strategy, strategy_selector, deployment)
verbose_router_logger.info(
"get_available_deployment for model: %s, Selected deployment: %s for model: %s",
model,
@ -13592,6 +13761,8 @@ class Router:
specific_deployment=specific_deployment,
)
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
# 2. If the returned is a specific deployment (Dict), verify and return directly
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
@ -13602,6 +13773,7 @@ class Router:
)
litellm_params: Final = healthy_deployments.get("litellm_params", {})
if litellm_params.get("use_in_pass_through"):
self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments)
return healthy_deployments
else:
# Specific deployment does not support pass-through
@ -13661,7 +13833,6 @@ class Router:
)
# 6. Apply load balancing strategy
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
@ -13693,6 +13864,7 @@ class Router:
enable_pre_call_checks=self.enable_pre_call_checks,
cooldown_list=_cooldown_list,
)
self._override_selector_pre_call_check(strategy, strategy_selector, deployment)
verbose_router_logger.info(
"get_available_deployment_for_pass_through model: %s, selected deployment: %s",

34
litellm/router_strategy/complexity_router/README.md
View file

@ -195,6 +195,40 @@ model_list:
session_affinity_ttl_seconds: 300
```
## Custom dimensions
Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal
```yaml
custom_dimensions:
- name: internalFrameworks
weight: 0.9
keywords: [orbitmesh, fluxgate]
- name: sqlMigration
weight: 0.7
patterns: ['\b(create|alter|drop)\s{1,4}table\b']
- name: dataPipeline
weight: 0.4
scoring_mode: match_count
keywords: [airflow, dbt, snowflake]
```
Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request
`scoring_mode` is optional and defaults to `binary`, the behavior above. `match_count` grades the dimension by how many distinct matchers hit: none scores 0 and emits no signal, one scores half the weight, two or more score the full weight. Repeated occurrences of one matcher never raise the count, keywords are distinct case-insensitively, patterns are distinct by source, and a keyword and a pattern are always distinct from each other. Matching stops as soon as the selected mode's maximum is reached, so a binary dimension still stops at its first hit. Existing configurations without the field keep binary scoring and the same tuning fingerprint, so the field only counts as a tuning change when set to `match_count`
### Weights through the API versus the dashboard
The API and YAML store exactly the weights written. A `dimension_weights` map and inline custom weights are read literally, missing recognized built-in names score zero, and nothing renormalizes the vector, so a total other than 1 is legal and scores accordingly. The dashboard's heuristic scoring editor is the one place that rebalances: editing one weight there holds it and redistributes the remainder across the other active dimensions in the draft, then Save sends the resulting explicit values, which the backend stores and scores as written. Opening a router, applying a preset, editing matchers, changing `scoring_mode`, or saving unrelated fields never normalizes existing weights
Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one
Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke
Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules
The existing heuristic-v1 tuning quota covers custom dimensions, their weights and their scoring mode: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML, the model API, or the dashboard's heuristic scoring editor
## Usage
Once configured, use the model name like any other:

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

@ -20,7 +20,7 @@ import random
import re
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import accumulate, islice, takewhile
from itertools import accumulate, chain, islice, takewhile
from threading import Lock
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
@ -31,6 +31,7 @@ from litellm._logging import verbose_router_logger
from litellm.constants import (
EMPTY_MAPPING,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
OUTPUT_TOKEN_CEILING_PARAMS,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
@ -67,6 +68,7 @@ from litellm.types.utils import (
from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section
from .config import (
CALIBRATION_EXAMPLES_HEADING,
CUSTOM_PATTERN_SCAN_CHARS,
DEFAULT_CLASSIFICATION_RUBRIC,
DEFAULT_CODE_KEYWORDS,
DEFAULT_ESCALATION_KEYWORDS,
@ -81,6 +83,7 @@ from .config import (
ClassificationRubric,
ComplexityRouterConfig,
ComplexityTier,
CustomDimension,
TierDefinition,
)
from .stall_detector import detect_stalled_task
@ -878,6 +881,15 @@ class DimensionScore:
self.signal = signal
class _CustomDimensionMatchers(NamedTuple):
"""One custom dimension's distinct matchers and the number of hits that saturates its score."""
dimension: CustomDimension
keywords: tuple[str, ...]
patterns: tuple[re.Pattern[str], ...]
saturation: int
class KeywordOverride(NamedTuple):
"""A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired."""
@ -1119,6 +1131,15 @@ class ComplexityRouter(CustomLogger):
self.config.custom_technical_keywords,
)
self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
self._custom_dimensions = tuple(
_CustomDimensionMatchers(
dimension,
tuple(dict.fromkeys(keyword.lower() for keyword in dimension.keywords)),
tuple(re.compile(pattern, re.IGNORECASE) for pattern in dict.fromkeys(dimension.patterns)),
2 if dimension.scoring_mode == "match_count" else 1,
)
for dimension in self.config.custom_dimensions
)
if self.config.has_custom_tiers:
self.escalation_keywords: tuple[str, ...] = ()
elif self.config.escalation_keywords is not None:
@ -1320,6 +1341,28 @@ class ComplexityRouter(CustomLogger):
score: Final = score_high if match_count >= high_threshold else score_low
return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
def _count_custom_hits(self, matchers: _CustomDimensionMatchers, user_text: str, scanned: str) -> int:
hits: Final = chain(
(self._keyword_matches(user_text, keyword) for keyword in matchers.keywords),
(pattern.search(scanned) is not None for pattern in matchers.patterns),
)
return sum(islice((1 for hit in hits if hit), matchers.saturation))
def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]:
if not self._custom_dimensions:
return ()
scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS]
return tuple(
(
DimensionScore(
matchers.dimension.name, hits / matchers.saturation, f"custom ({matchers.dimension.name})"
),
matchers.dimension.weight,
)
for matchers in self._custom_dimensions
if (hits := self._count_custom_hits(matchers, user_text, scanned))
)
def _score_multi_step(self, text: str) -> DimensionScore:
"""Score based on multi-step patterns."""
hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text))
@ -1415,12 +1458,13 @@ class ComplexityRouter(CustomLogger):
self._score_question_complexity(prompt),
]
# Collect signals
signals: Final = [d.signal for d in dimensions if d.signal is not None]
custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text)
signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None]
# Compute weighted score
weights: Final = self.config.dimension_weights
weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions)
weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum(
dimension.score * weight for dimension, weight in custom_dimensions
)
boundaries: Final = self._effective_tier_boundaries()
clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score()
@ -2067,11 +2111,15 @@ class ComplexityRouter(CustomLogger):
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]:
if tier is None:
return MappingProxyType({})
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ())
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else ()
entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None)
return entry.litellm_params if entry is not None else MappingProxyType({})
explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({})
if not self.config.max_tokens_from_tier_model or not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(explicit):
return explicit
ceiling: Final = self._group_output_ceiling(model)
if ceiling is None:
return explicit
return MappingProxyType({**explicit, "max_tokens": ceiling})
@staticmethod
def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str:
@ -2406,12 +2454,15 @@ class ComplexityRouter(CustomLogger):
return name if self.config.has_custom_tiers else ComplexityTier(name)
def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None:
return self._deployment_limit(group, deployment, "max_input_tokens")
def _deployment_limit(
self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"]
) -> int | None:
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
deployment_model_info: Final = deployment.get("model_info")
declared: Final = (
deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None
)
declared: Final = deployment_model_info.get(key) if isinstance(deployment_model_info, Mapping) else None
if isinstance(declared, int):
return declared
litellm_params: Final = deployment.get("litellm_params")
@ -2428,18 +2479,34 @@ class ComplexityRouter(CustomLogger):
deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts
received_model_name=group,
)
window: Final = model_info.get("max_input_tokens")
limit: Final = model_info.get(key)
except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others
return None
return window if isinstance(window, int) else None
return limit if isinstance(limit, int) else None
def _group_deployments(self, group: str) -> Sequence[Mapping[str, object]]:
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
deployments: Final = list_models(model_name=group) if callable(list_models) else None
return tuple(deployments) if isinstance(deployments, list) else ()
def _group_output_ceiling(self, group: str) -> int | None:
"""Smallest max_output_tokens across the group's deployments, or None when any deployment
declares none: the core router picks within the group without a fit check, and a ceiling
above an unmapped member's real limit is a provider 400 on that member."""
deployments: Final = self._group_deployments(group)
ceilings: Final = tuple(
ceiling
for deployment in deployments
if (ceiling := self._deployment_limit(group, deployment, "max_output_tokens")) is not None
)
return min(ceilings) if ceilings and len(ceilings) == len(deployments) else None
def _group_window_facts(self, group: str) -> tuple[int | None, bool]:
"""(smallest declared context window across the group's deployments, whether any deployment
declares none). The core router picks a deployment within the group without a fit check, so
the group is only as safe as its smallest member."""
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
deployments: Final = list_models(model_name=group) if callable(list_models) else None
if not isinstance(deployments, list) or not deployments:
deployments: Final = self._group_deployments(group)
if not deployments:
return (None, True)
windows: Final = tuple(
window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None
@ -3488,6 +3555,7 @@ class ComplexityRouter(CustomLogger):
ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
)
fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM
default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model)
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
@ -3496,7 +3564,9 @@ class ComplexityRouter(CustomLogger):
cause="default_fallback",
tier=fallback_tier,
conversation_continuing=conversation_continuing,
tier_litellm_params=default_tier_params,
),
litellm_params=default_tier_params,
)
ask: Final = user_message or ""
@ -3523,6 +3593,7 @@ class ComplexityRouter(CustomLogger):
_tier_name(plan_floor),
routed_model,
)
plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model)
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
@ -3534,7 +3605,9 @@ class ComplexityRouter(CustomLogger):
matched_keyword=plan_mode_sentinel,
escalation_keyword=escalation_keyword,
escalated=False,
tier_litellm_params=plan_tier_params,
),
litellm_params=plan_tier_params,
)
override: Final = await self._resolve_keyword_tier_override(ask, request_kwargs)
@ -3627,6 +3700,7 @@ class ComplexityRouter(CustomLogger):
outcome.signals,
fallback_model,
)
fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model)
return PreRoutingHookResponse(
model=fallback_model,
messages=messages if has_original_messages else None,
@ -3637,7 +3711,9 @@ class ComplexityRouter(CustomLogger):
signals=outcome.signals,
escalation_keyword=escalation_keyword,
escalated=False,
tier_litellm_params=fallback_tier_params,
),
litellm_params=fallback_tier_params,
)
if self.config.adaptive:
# hard_floor rather than a hard pick, and passed whenever the sentinel is present

180
litellm/router_strategy/complexity_router/config.py
View file

@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas
All values are configurable via proxy config.yaml.
"""
from collections.abc import Mapping
import math
import re
import warnings
from collections.abc import Iterable, Mapping
from enum import Enum
from types import MappingProxyType
from typing import Annotated, Final, Literal
from typing import Annotated, Final, Literal, NamedTuple
from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
import sre_constants
import sre_parse
from litellm.types.llms.openai import REASONING_EFFORT
from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin
@ -569,6 +577,125 @@ class ClassifierLLMConfig(BaseModel):
return self
MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64
MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048
MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192
MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16
CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048
_ATOM_OPCODES: Final = frozenset(
{sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY}
)
_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT})
class _PatternCost(NamedTuple):
paths: int
steps: int
def _atom_steps(node: object) -> int:
if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN:
return 1 + len(node[1])
return 1
def _repeat_cost(argument: object) -> _PatternCost | str:
if not isinstance(argument, tuple) or len(argument) != 3:
return "unsupported repeat structure"
low, high, body = argument
if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES:
return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}"
choices: Final = high - low + 1
return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices)
def _node_cost(node: object, depth: int) -> _PatternCost | str:
if not isinstance(node, tuple) or len(node) != 2:
return "unsupported regex structure"
opcode, argument = node
if opcode in _ATOM_OPCODES or opcode is sre_constants.AT:
return _PatternCost(1, _atom_steps(node))
if opcode is sre_constants.SUBPATTERN:
return _sequence_cost(argument[-1], depth + 1)
if opcode is sre_constants.BRANCH:
costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1])
refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
if refused is not None:
return refused
return _PatternCost(
sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)),
len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)),
)
if opcode in _REPEAT_OPCODES:
return _repeat_cost(argument)
return "contains an unsupported regex construct"
def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str:
if depth > MAX_CUSTOM_PATTERN_DEPTH:
return "nests deeper than 16 levels"
costs: Final = tuple(_node_cost(node, depth) for node in nodes)
refused: Final = next((cost for cost in costs if isinstance(cost, str)), None)
if refused is not None:
return refused
valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost))
# Choices multiply across a sequence; every continuation can execute once per preceding path.
total: Final = _PatternCost(
math.prod(cost.paths for cost in valid),
1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)),
)
if total.steps > MAX_CUSTOM_PATTERN_WORK:
return "exceeds the per-pattern regex work budget"
return total
def custom_pattern_work(pattern: str) -> int | str:
try:
re.compile(pattern, re.IGNORECASE)
parsed: Final = sre_parse.parse(pattern, re.IGNORECASE)
except (re.error, RecursionError, OverflowError):
return "is not a valid regex"
cost: Final = _sequence_cost(tuple(parsed), 0)
return cost if isinstance(cost, str) else cost.steps
class CustomDimension(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$")
weight: float = Field(gt=0, le=1, allow_inf_nan=False)
keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
scoring_mode: Literal["binary", "match_count"] = Field(
default="binary",
description=(
"'binary' scores 1 when any matcher hits. 'match_count' scores 0.5 when one distinct matcher hits and 1 "
"when two or more do; repeated occurrences of one matcher never raise it. Keywords are distinct "
"case-insensitively, patterns by source, and a keyword and a pattern are always distinct from each other."
),
)
@model_validator(mode="after")
def _validate_matchers(self) -> "CustomDimension":
matchers: Final = (*self.keywords, *self.patterns)
if not matchers or any(not matcher.strip() for matcher in matchers):
raise ValueError("custom dimensions require nonblank keywords and/or patterns")
if len(matchers) > 32 or sum(map(len, matchers)) > 4096:
raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each")
costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns)
rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str))
if rejected:
raise ValueError("custom dimension " + "; ".join(rejected))
return self
def pattern_work(self) -> int:
"""Combined work estimate of the validated patterns."""
return sum(
work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int)
)
class ComplexityRouterConfig(BaseModel):
"""Configuration for the ComplexityRouter."""
@ -671,6 +798,20 @@ class ComplexityRouterConfig(BaseModel):
description="Weights for each scoring dimension",
)
custom_dimensions: tuple[CustomDimension, ...] = Field(
default=(),
max_length=16,
description=(
"Named dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
"when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters; "
"scoring_mode 'match_count' instead grades half weight for one distinct matcher and full for two or more. "
"Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, "
"backreferences and lookarounds are rejected. Conservative work limits include alternation paths, "
"repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. "
"Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota."
),
)
# Keyword lists (overridable)
code_keywords: list[str] | None = Field(
default=None,
@ -948,6 +1089,20 @@ class ComplexityRouterConfig(BaseModel):
"wording the built-ins don't cover, or after a client release changes its strings."
),
)
max_tokens_from_tier_model: bool = Field(
default=True,
description=(
"Set max_tokens on every routed request to the output ceiling of the tier model it "
"lands on, replacing whatever the caller sent. A caller behind an auto-router cannot "
"pick one value that fits every tier: the smallest tier's ceiling starves a bigger "
"tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The "
"ceiling is the smallest max_output_tokens across the tier model's deployments, read "
"from each deployment's model_info and then the model cost map; a tier model with a "
"deployment whose ceiling is unknown keeps the caller's value. A max_tokens, "
"max_completion_tokens or max_output_tokens in the tier's own litellm_params still "
"wins. Set false to forward the caller's value unchanged."
),
)
route_housekeeping_to_cheapest_tier: bool = Field(
default=True,
description=(
@ -1245,6 +1400,27 @@ class ComplexityRouterConfig(BaseModel):
)
return self
@model_validator(mode="after")
def _validate_custom_dimensions(self) -> "ComplexityRouterConfig":
if not self.custom_dimensions:
return self
if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"):
raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid")
names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions)
reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS)
weighted: Final = frozenset(name.casefold() for name in self.dimension_weights)
if len(frozenset(names)) != len(names) or frozenset(names) & reserved:
raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions")
if frozenset(names) & weighted:
raise ValueError("custom dimension weights must be inline, not in dimension_weights")
work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions)
if work > MAX_CUSTOM_DIMENSIONS_WORK:
raise ValueError(
f"custom_dimensions regex work estimate is {work}; the limit across the router is "
f"{MAX_CUSTOM_DIMENSIONS_WORK}"
)
return self
@field_validator("heuristic_first_max_tier", mode="before")
@classmethod
def _coerce_heuristic_first_max_tier(cls, value: object) -> object:

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