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

merge: resolve MongoDB sidecar staging conflicts

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-08 13:59:56 -07:00
commit 5c55ad7db0
No known key found for this signature in database
136 changed files with 6481 additions and 1085 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 \

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

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

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

5
litellm/constants.py
View file

@ -73,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))
@ -1458,6 +1461,8 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"})
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"

40
litellm/cost_calculator.py
View file

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

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

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]):

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

24
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,
@ -9814,7 +9826,7 @@ class BaseLLMHTTPHandler:
vector_store_search_optional_params=vector_store_search_optional_params,
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params={**dict(litellm_params), "timeout": timeout},
litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)),
extra_body=extra_body,
embedding_executor=embedding_executor,
)
@ -9861,7 +9873,9 @@ class BaseLLMHTTPHandler:
)
except httpx.TimeoutException:
raise vector_store_provider_config.get_error_class(
error_message="Vector store search exceeded the caller timeout.", status_code=408, headers={}
error_message="Vector store search exceeded the caller timeout.",
status_code=408,
headers=httpx.Headers(),
) from None
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)
@ -9947,7 +9961,7 @@ class BaseLLMHTTPHandler:
vector_store_search_optional_params=vector_store_search_optional_params,
api_base=api_base,
litellm_logging_obj=logging_obj,
litellm_params={**dict(litellm_params), "timeout": timeout},
litellm_params=MappingProxyType(dict(litellm_params, timeout=timeout)),
extra_body=extra_body,
embedding_executor=embedding_executor,
)
@ -9996,7 +10010,9 @@ class BaseLLMHTTPHandler:
)
except httpx.TimeoutException:
raise vector_store_provider_config.get_error_class(
error_message="Vector store search exceeded the caller timeout.", status_code=408, headers={}
error_message="Vector store search exceeded the caller timeout.",
status_code=408,
headers=httpx.Headers(),
) from None
except Exception as e:
raise self._handle_error(e=e, provider_config=vector_store_provider_config)

51
litellm/llms/mongodb/vector_stores/transformation.py
View file

@ -58,7 +58,7 @@ class _Content(BaseModel):
class _Result(BaseModel):
model_config = ConfigDict(frozen=True, strict=True, allow_inf_nan=False)
score: float | None
content: list[_Content]
content: Sequence[_Content]
file_id: str | None
filename: str | None
@ -67,7 +67,7 @@ class _SearchResponse(BaseModel):
model_config = ConfigDict(frozen=True, strict=True)
object: Literal["vector_store.search_results.page"]
search_query: str
data: list[_Result]
data: Sequence[_Result]
class _MongoDBSearchParams(BaseModel):
@ -185,17 +185,21 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
return min(max(limit * NUM_CANDIDATES_MULTIPLIER, MIN_NUM_CANDIDATES), MAX_NUM_CANDIDATES)
def validate_environment(
self, headers: dict[str, object], litellm_params: GenericLiteLLMParams | None
) -> dict[str, object]:
self, headers: Mapping[str, object], litellm_params: GenericLiteLLMParams | None
) -> dict[str, object]: # mutable-ok: the shared HTTP handler requires writable headers
if litellm_params is None:
raise config_error("Configure api_base and api_key for the MongoDB BETA sidecar.")
self._reject_unknown_params(dict(litellm_params))
self._reject_unknown_params(MappingProxyType(dict(litellm_params)))
api_key: Final = litellm_params.api_key or get_secret_str("MONGODB_SIDECAR_API_KEY")
if not api_key:
raise config_error("MongoDB sidecar api_key is required. Set api_key or MONGODB_SIDECAR_API_KEY.")
return {**headers, "Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
return {
**headers,
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
} # mutable-ok: writable HTTP headers
def get_complete_url(self, api_base: str | None, litellm_params: dict[str, object]) -> str:
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
if not api_base:
raise config_error("MongoDB sidecar api_base is required, for example http://127.0.0.1:8080.")
try:
@ -268,7 +272,7 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
api_base: str,
embedding_response: EmbeddingResponse,
timeout: object,
) -> tuple[str, dict[str, object]]:
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
if not embedding_response.data:
raise config_error(
"The embedding model returned no embedding for the search query. Check litellm_embedding_model."
@ -277,17 +281,20 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
if not vector or any(not isinstance(value, (float, int)) or not isfinite(value) for value in vector):
raise config_error("The embedding model must return a non-empty, finite query vector.")
limit: Final = cls._limit(optional_params)
return f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search", {
"query": query_text,
"query_vector": tuple(vector),
"mongodb_database": params.require_database(),
"mongodb_collection": params.require_collection(),
"mongodb_embedding_field": params.embedding_field,
"mongodb_text_field": params.text_field,
"mongodb_num_candidates": cls._num_candidates(limit, params.mongodb_num_candidates),
"max_num_results": limit,
"timeout_ms": cls._timeout_ms(timeout),
}
return (
f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search",
{ # mutable-ok: JSON transport requires a dict
"query": query_text,
"query_vector": tuple(vector),
"mongodb_database": params.require_database(),
"mongodb_collection": params.require_collection(),
"mongodb_embedding_field": params.embedding_field,
"mongodb_text_field": params.text_field,
"mongodb_num_candidates": cls._num_candidates(limit, params.mongodb_num_candidates),
"max_num_results": limit,
"timeout_ms": cls._timeout_ms(timeout),
},
)
def transform_search_vector_store_request(
self,
@ -299,7 +306,7 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
litellm_params: Mapping[str, object],
extra_body: Mapping[str, object] | None = None,
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
) -> tuple[str, dict[str, object]]:
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body)
query_text: Final = self._query_text(query)
response: Final = (embedding_executor or self.embedding_executor).embed(
@ -325,7 +332,7 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
litellm_params: Mapping[str, object],
extra_body: Mapping[str, object] | None = None,
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
) -> tuple[str, dict[str, object]]:
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body)
query_text: Final = self._query_text(query)
response: Final = await (embedding_executor or self.embedding_executor).aembed(
@ -355,7 +362,7 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig):
) from None
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
self, error_message: str, status_code: int, headers: Mapping[str, object] | httpx.Headers
) -> BaseLLMException:
if status_code == 400:
raise config_error(error_message)

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

40
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -19886,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.",

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

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

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

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,

5
litellm/proxy/litellm_pre_call_utils.py
View file

@ -18,6 +18,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
@ -325,7 +326,9 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
# ``attempted_fallbacks`` and ``original_model_group`` are written by the router
# and read by spend logs as fact; a client value has no legitimate meaning and no
# key or team setting keeps it, so the strip is never gated.
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset({"attempted_fallbacks", "original_model_group"})
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
)
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
# Request fields whose value, when URL-valued, becomes the outbound destination

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

39
litellm/proxy/proxy_server.py
View file

@ -631,6 +631,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 +1060,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 +1370,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}")
@ -7109,11 +7117,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
)
@ -7153,12 +7160,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()
@ -8013,7 +8017,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
@ -9597,19 +9601,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(
@ -9623,7 +9614,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(

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:

160
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,14 @@ 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,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
@ -648,6 +659,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
@ -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:
@ -4214,10 +4237,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)
@ -12604,7 +12629,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:
@ -12614,7 +12715,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)
@ -12655,6 +12771,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,
@ -12666,12 +12783,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
@ -12771,6 +12890,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,
@ -12782,12 +12902,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(

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

@ -207,17 +207,27 @@ custom_dimensions:
- 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 and their weights: 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 or the model API; this change adds no dashboard editor
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

95
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,
)
@ -82,6 +83,7 @@ from .config import (
ClassificationRubric,
ComplexityRouterConfig,
ComplexityTier,
CustomDimension,
TierDefinition,
)
from .stall_detector import detect_stalled_task
@ -879,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."""
@ -1121,7 +1132,12 @@ class ComplexityRouter(CustomLogger):
)
self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
self._custom_dimensions = tuple(
(dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns))
_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:
@ -1325,15 +1341,26 @@ 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(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight)
for dimension, patterns in self._custom_dimensions
if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords)
or any(pattern.search(scanned) is not None for pattern in patterns)
(
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:
@ -2084,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:
@ -2423,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")
@ -2445,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
@ -3505,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,
@ -3513,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 ""
@ -3540,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,
@ -3551,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)
@ -3644,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,
@ -3654,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

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

@ -667,6 +667,14 @@ class CustomDimension(BaseModel):
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":
@ -794,8 +802,9 @@ class ComplexityRouterConfig(BaseModel):
default=(),
max_length=16,
description=(
"Named binary 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. "
"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. "
@ -1080,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=(

378
litellm/router_strategy/least_busy.py
View file

@ -1,17 +1,103 @@
#### What this does ####
# identifies least busy deployment
# How is this achieved?
# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"}
# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic}
# - use litellm.success + failure callbacks to log when a request completed
# - in get_available_deployment, for a given model group name -> pick based on traffic
import random
from collections.abc import Mapping, Sequence
from typing import Final
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
class _ModelInfo(TypedDict, total=False):
id: ReadOnly[str | int | None]
class _Metadata(TypedDict, total=False):
model_group: ReadOnly[str | None]
class _LitellmParams(TypedDict, total=False):
metadata: ReadOnly[_Metadata | None]
model_info: ReadOnly[_ModelInfo | None]
class _CallKwargs(TypedDict, total=False):
litellm_params: ReadOnly[_LitellmParams | None]
class _DeploymentModelInfo(TypedDict):
id: ReadOnly[str | int]
class _Deployment(TypedDict):
model_info: ReadOnly[_DeploymentModelInfo]
_CALL_KWARGS: Final = TypeAdapter(_CallKwargs)
_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None)
def _request_count_key(model_group: str, deployment_id: str) -> str:
return f"{model_group}_request_count:{deployment_id}"
def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
try:
call: Final = _CALL_KWARGS.validate_python(kwargs)
except ValidationError:
return None
litellm_params: Final = call.get("litellm_params")
metadata: Final = litellm_params.get("metadata") if litellm_params else None
model_info: Final = litellm_params.get("model_info") if litellm_params else None
model_group: Final = metadata.get("model_group") if metadata else None
deployment_id: Final = model_info.get("id") if model_info else None
if model_group is None or deployment_id is None:
return None
return model_group, str(deployment_id)
def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
return tuple(
_request_count_key(model_group, str(deployment["model_info"]["id"]))
for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
)
def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]:
return tuple(0 if value is None else int(value) for value in values)
def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]:
values: Final = _MEMORY_COUNTS.validate_python(raw)
if values is None or len(values) != len(keys):
return (0,) * len(keys)
return _as_counts(values)
def _least_busy(
healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...]
) -> Mapping[str, object] | None:
if not healthy_deployments:
return None
return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])]
def _warn_unreadable(model_group: str, error: Exception) -> None:
verbose_router_logger.warning(
"least-busy routing could not read the shared in-flight counts for %s, "
"falling back to this worker's own counts: %s",
model_group,
error,
)
def _warn_unwritable(key: str, error: Exception) -> None:
verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error)
class LeastBusyLoggingHandler(CustomLogger):
test_flag: bool = False
@ -20,195 +106,101 @@ class LeastBusyLoggingHandler(CustomLogger):
def __init__(self, router_cache: DualCache):
self.router_cache = router_cache
self.router_cache_id = str(id(router_cache))
def log_pre_api_call(self, model, messages, kwargs):
"""
Log when a model is being used.
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
self._increment(kwargs, 1)
Caching based on model group.
"""
try:
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
def log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._increment(kwargs, -1)
if self.test_flag:
self.logged_success += 1
request_count_api_key: Final = f"{model_group}_request_count"
# update cache
request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
request_count_dict[id] = request_count_dict.get(id, 0) + 1
def log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._increment(kwargs, -1)
if self.test_flag:
self.logged_failure += 1
self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict)
except Exception:
pass
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
await self._async_increment(kwargs, -1)
if self.test_flag:
self.logged_success += 1
def log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
request_count_api_key: Final = f"{model_group}_request_count"
# decrement count in cache
request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
request_count_value: Final[int | None] = request_count_dict.get(id, 0)
if request_count_value is None:
return
request_count_dict[id] = request_count_value - 1
self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict)
### TESTING ###
if self.test_flag:
self.logged_success += 1
except Exception:
pass
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
request_count_api_key: Final = f"{model_group}_request_count"
# decrement count in cache
request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
request_count_value: Final[int | None] = request_count_dict.get(id, 0)
if request_count_value is None:
return
request_count_dict[id] = request_count_value - 1
self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict)
### TESTING ###
if self.test_flag:
self.logged_failure += 1
except Exception:
pass
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
request_count_api_key: Final = f"{model_group}_request_count"
# decrement count in cache
request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {}
request_count_value: Final[int | None] = request_count_dict.get(id, 0)
if request_count_value is None:
return
request_count_dict[id] = request_count_value - 1
await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict)
### TESTING ###
if self.test_flag:
self.logged_success += 1
except Exception:
pass
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
request_count_api_key: Final = f"{model_group}_request_count"
# decrement count in cache
request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {}
request_count_value: Final[int | None] = request_count_dict.get(id, 0)
if request_count_value is None:
return
request_count_dict[id] = request_count_value - 1
await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict)
### TESTING ###
if self.test_flag:
self.logged_failure += 1
except Exception:
pass
def _get_available_deployments(
self,
healthy_deployments: list,
all_deployments: dict,
):
"""
Helper to get deployments using least busy strategy
"""
for d in healthy_deployments:
## if healthy deployment not yet used
if d["model_info"]["id"] not in all_deployments:
all_deployments[d["model_info"]["id"]] = 0
# map deployment to id
# pick least busy deployment
min_traffic = float("inf")
min_deployment = None
for k, v in all_deployments.items():
if v < min_traffic:
min_traffic = v
min_deployment = k
if min_deployment is not None:
## check if min deployment is a string, if so, cast it to int
for m in healthy_deployments:
if m["model_info"]["id"] == min_deployment:
return m
min_deployment = random.choice(healthy_deployments)
else:
min_deployment = random.choice(healthy_deployments)
return min_deployment
async def async_log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
await self._async_increment(kwargs, -1)
if self.test_flag:
self.logged_failure += 1
def get_available_deployments(
self,
model_group: str,
healthy_deployments: list,
):
"""
Sync helper to get deployments using least busy strategy
"""
request_count_api_key: Final = f"{model_group}_request_count"
all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {}
return self._get_available_deployments(
healthy_deployments=healthy_deployments,
all_deployments=all_deployments,
)
self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
) -> Mapping[str, object] | None:
keys: Final = _request_count_keys(model_group, healthy_deployments)
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is not None:
try:
shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys)))
except Exception as e:
_warn_unreadable(model_group, e)
else:
return _least_busy(healthy_deployments, shared)
local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
return _least_busy(healthy_deployments, local)
async def async_get_available_deployments(self, model_group: str, healthy_deployments: list):
"""
Async helper to get deployments using least busy strategy
"""
request_count_api_key: Final = f"{model_group}_request_count"
all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {}
return self._get_available_deployments(
healthy_deployments=healthy_deployments,
all_deployments=all_deployments,
)
async def async_get_available_deployments(
self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
) -> Mapping[str, object] | None:
keys: Final = _request_count_keys(model_group, healthy_deployments)
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is not None:
try:
shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys)))
except Exception as e:
_warn_unreadable(model_group, e)
else:
return _least_busy(healthy_deployments, shared)
local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys)
return _least_busy(healthy_deployments, local)
def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None:
return
key: Final = _request_count_key(*ref)
redis_cache: Final = self.router_cache.redis_cache
try:
local: Final = self.router_cache.increment_cache(
key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
)
if local < 0:
self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
if redis_cache is None:
return
redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
except Exception as e:
_warn_unwritable(key, e)
async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None:
return
key: Final = _request_count_key(*ref)
redis_cache: Final = self.router_cache.redis_cache
try:
local: Final = await self.router_cache.async_increment_cache(
key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
)
if local is not None and local < 0:
await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
if redis_cache is None:
return
await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
except Exception as e:
_warn_unwritable(key, e)

54
litellm/router_strategy/lowest_latency.py
View file

@ -32,6 +32,12 @@ def _average_latency(samples: Sequence[float]) -> float:
return sum(samples) / len(samples)
def _ttft_seconds(elapsed: timedelta | float) -> float:
if isinstance(elapsed, timedelta):
return elapsed.total_seconds()
return float(elapsed)
class LowestLatencyLoggingHandler(CustomLogger):
test_flag: bool = False
logged_success: int = 0
@ -86,14 +92,13 @@ class LowestLatencyLoggingHandler(CustomLogger):
# breaks JSON serialization when the router cache syncs to
# Redis (issue #33169)
response_ms = response_ms.total_seconds()
time_to_first_token_response_time = None
time_to_first_token: float | None = None
if kwargs.get("stream", None) is not None and kwargs["stream"] is True:
# only log ttft for streaming request
time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time
time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time)
final_value: float = response_ms
time_to_first_token: float | None = None
total_tokens = 0
if isinstance(response_obj, ModelResponse):
@ -111,13 +116,6 @@ class LowestLatencyLoggingHandler(CustomLogger):
else:
final_value = response_seconds
if time_to_first_token_response_time is not None:
if isinstance(time_to_first_token_response_time, timedelta):
ttft_seconds = time_to_first_token_response_time.total_seconds()
else:
ttft_seconds = time_to_first_token_response_time
time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens)
# ------------
# Update usage
# ------------
@ -138,14 +136,14 @@ class LowestLatencyLoggingHandler(CustomLogger):
## Time to first token
if time_to_first_token is not None:
if (
len(request_count_dict[id].get("time_to_first_token", []))
len(request_count_dict[id].get("time_to_first_token_seconds", []))
< self.routing_args.max_latency_list_size
):
request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token)
request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token)
else:
request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][
1:
] + [time_to_first_token]
request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][
"time_to_first_token_seconds"
][1:] + [time_to_first_token]
if precise_minute not in request_count_dict[id]:
request_count_dict[id][precise_minute] = {}
@ -252,7 +250,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
{model_group}_map: {
id: {
"latency": [..]
"time_to_first_token": [..]
"time_to_first_token_seconds": [..]
f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
}
}
@ -273,14 +271,13 @@ class LowestLatencyLoggingHandler(CustomLogger):
# breaks JSON serialization when the router cache syncs to
# Redis (issue #33169)
response_ms = response_ms.total_seconds()
time_to_first_token_response_time = None
time_to_first_token: float | None = None
if kwargs.get("stream", None) is not None and kwargs["stream"] is True:
# only log ttft for streaming request
time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time
time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time)
final_value: float = response_ms
total_tokens = 0
time_to_first_token: float | None = None
if isinstance(response_obj, ModelResponse):
_usage: Final = getattr(response_obj, "usage", None)
@ -296,13 +293,6 @@ class LowestLatencyLoggingHandler(CustomLogger):
final_value = float(normalized_value)
else:
final_value = response_seconds
if time_to_first_token_response_time is not None:
if isinstance(time_to_first_token_response_time, timedelta):
ttft_seconds = time_to_first_token_response_time.total_seconds()
else:
ttft_seconds = time_to_first_token_response_time
time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens)
# ------------
# Update usage
# ------------
@ -328,14 +318,14 @@ class LowestLatencyLoggingHandler(CustomLogger):
## Time to first token
if time_to_first_token is not None:
if (
len(request_count_dict[id].get("time_to_first_token", []))
len(request_count_dict[id].get("time_to_first_token_seconds", []))
< self.routing_args.max_latency_list_size
):
request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token)
request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token)
else:
request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][
1:
] + [time_to_first_token]
request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][
"time_to_first_token_seconds"
][1:] + [time_to_first_token]
if precise_minute not in request_count_dict[id]:
request_count_dict[id][precise_minute] = {}
@ -433,7 +423,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
or float("inf")
)
item_latency = item_map.get("latency", [])
item_ttft_latency = item_map.get("time_to_first_token", [])
item_ttft_latency = item_map.get("time_to_first_token_seconds", [])
item_rpm = item_map.get(precise_minute, {}).get("rpm", 0)
item_tpm = item_map.get(precise_minute, {}).get("tpm", 0)

12
litellm/router_utils/auto_router_tuning_baseline.py
View file

@ -52,7 +52,17 @@ def tuning_fingerprint(complexity_router_config: object) -> str | None:
supplied: Final = ((_TUNING_FIELD_SET - frozenset(("tier_model_configs",))) & frozenset(raw)) | (
frozenset(("tier_model_configs",)) if validated.tier_model_configs else frozenset()
)
payload: Final = validated.model_dump(mode="json", include=supplied)
payload: Final = validated.model_dump(
mode="json",
include=supplied,
exclude={
"custom_dimensions": {
index: {"scoring_mode"}
for index, dimension in enumerate(validated.custom_dimensions)
if dimension.scoring_mode == "binary"
}
},
)
return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest()

48
litellm/router_utils/cooldown_cache.py
View file

@ -12,6 +12,7 @@ from typing_extensions import TypedDict
from litellm import verbose_logger
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
if TYPE_CHECKING:
@ -36,10 +37,19 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0
class CooldownCache:
def __init__(self, cache: DualCache, default_cooldown_time: float):
def __init__(
self,
cache: DualCache,
default_cooldown_time: float,
redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS,
):
self.cache = cache
self.default_cooldown_time = default_cooldown_time
self.in_memory_cache = InMemoryCache()
self._cooldown_store = DualCache(
in_memory_cache=self.in_memory_cache,
default_redis_batch_cache_expiry=redis_read_interval_seconds,
)
# Initialize the masker with custom settings for exception strings
self.exception_masker = SensitiveDataMasker(
visible_prefix=50, # Show first 50 characters
@ -48,6 +58,21 @@ class CooldownCache:
mask_short_values=False, # Truncate long messages only; keep short ones readable
)
@property
def cooldown_store(self) -> DualCache:
"""
The cache cooldown entries live in, with the router's Redis attached on first use.
It is kept separate from the router-wide cache so that a key missing from memory is
re-read from Redis every `redis_read_interval_seconds` rather than on the router
cache's much longer batch interval, which is what lets a sibling replica see a
cooldown another replica wrote, and so that unrelated router keys cannot evict a
cooldown from the in-memory tier before it expires. Redis is attached lazily because
the router builds its cooldown cache before it wires up the shared Redis client.
"""
self._cooldown_store.attach_redis_cache(self.cache.redis_cache)
return self._cooldown_store
def _common_add_cooldown_logic(
self, model_id: str, original_exception, exception_status, cooldown_time: float
) -> tuple[str, CooldownCacheValue]:
@ -93,7 +118,7 @@ class CooldownCache:
)
# Set the cache with a TTL equal to the cooldown time
self.cache.set_cache(
self.cooldown_store.set_cache(
value=cooldown_data,
key=cooldown_key,
ttl=_cooldown_time,
@ -122,13 +147,13 @@ class CooldownCache:
cooldown_cache_value: Final = CooldownCacheValue(**result) # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code
remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time
if remaining <= 0:
self.cache.in_memory_cache.delete_cache(key)
self.in_memory_cache.delete_cache(key)
return None
current_expiry: Final = self.cache.in_memory_cache.ttl_dict.get(key)
current_expiry: Final = self.in_memory_cache.ttl_dict.get(key)
if current_expiry is not None and current_expiry > current_time + remaining + 5:
corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS)
self.cache.in_memory_cache.delete_cache(key)
self.cache.in_memory_cache.set_cache(key, result, ttl=corrected_ttl)
self.in_memory_cache.delete_cache(key)
self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl)
return cooldown_cache_value
async def async_get_active_cooldowns(
@ -137,12 +162,7 @@ class CooldownCache:
# Generate the keys for the deployments
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
# Retrieve the values for the keys using mget
## more likely to be none if no models ratelimited. So just check redis every 1s
## each redis call adds ~100ms latency.
## check in memory cache first
results: Final = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span)
results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span)
active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = []
if results is None or all(v is None for v in results):
@ -164,7 +184,7 @@ class CooldownCache:
# Generate the keys for the deployments
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
# Retrieve the values for the keys using mget
results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
active_cooldowns: Final = []
current_time: Final = time.time()
@ -184,7 +204,7 @@ class CooldownCache:
keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
# Retrieve the values for the keys using mget
results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
min_cooldown_time: float | None = None
# Process the results

1
litellm/types/utils.py
View file

@ -580,6 +580,7 @@ CallTypesLiteral = Literal[
"search",
"asearch",
"_arealtime",
"_aresponses_websocket",
"create_batch",
"acreate_batch",
"create_file",

18
litellm/utils.py
View file

@ -672,11 +672,19 @@ def load_credentials_from_list(kwargs: dict):
CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor")
credential_name: Final = kwargs.get("litellm_credential_name")
if credential_name and litellm.credential_list:
credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name)
for key, value in credential_accessor.items():
if key not in kwargs:
kwargs[key] = value
if not credential_name:
return
credential: Final = CredentialAccessor.find_credential(credential_name)
if credential is None:
verbose_logger.warning(
"litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it",
credential_name,
len(litellm.credential_list),
)
return
for key, value in credential.credential_values.items():
if key not in kwargs:
kwargs[key] = value
def get_dynamic_callbacks(

6
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.101.0"
version = "1.102.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -67,7 +67,7 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.94",
"litellm-proxy-extras==0.4.95",
"litellm-enterprise==0.1.65",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
@ -327,7 +327,7 @@ members = ["enterprise", "litellm-proxy-extras"]
profile = "black"
[tool.commitizen]
version = "1.101.0"
version = "1.102.0"
version_files = [
"pyproject.toml:^version",
]

10
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 2956
"limit": 2918
},
"ANN002": {
"limit": 71
@ -9,10 +9,10 @@
"limit": 806
},
"ANN201": {
"limit": 1979
"limit": 1965
},
"ANN202": {
"limit": 831
"limit": 829
},
"ANN204": {
"limit": 683
@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
"limit": 2916
"limit": 2914
},
"C401": {
"limit": 8
@ -189,7 +189,7 @@
"limit": 0
},
"S110": {
"limit": 217
"limit": 207
},
"S112": {
"limit": 22

16
terraform/litellm/aws/README.md
View file

@ -242,6 +242,22 @@ this with `litellm_license`. To tune the export cadence, set
`LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` /
`backend_extra_env`
### Prometheus metrics sidecar
`gateway_metrics_port` adds a `metrics` sidecar
(`python -m litellm.proxy.prometheus_metrics_server`) to the gateway task that
aggregates the workers' samples over a shared task volume, so a scrape never
runs on an inference worker. The ALB never routes to that port and the tasks
security group only opens it to `gateway_metrics_scrape_cidrs`. Needs
`gateway_image` v1.101.0 or newer. See
[Prometheus metrics](https://docs.litellm.ai/docs/proxy/prometheus) for the
metrics themselves.
```hcl
gateway_metrics_port = 4001
gateway_metrics_scrape_cidrs = ["10.0.0.0/16"]
```
## Tenant deployment
Every resource the stack creates is named `${tenant}-litellm-${env}` (or

54
terraform/litellm/aws/ecs.tf
View file

@ -212,6 +212,45 @@ locals {
# pull the config from S3 first, so the command goes through `sh -c`;
# otherwise we keep the image's ENTRYPOINT and only override `command`.
gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}"
metrics_enabled = var.gateway_metrics_port != null
metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc"
metrics_volume = "prometheus-multiproc"
metrics_env = local.metrics_enabled ? [{ name = "PROMETHEUS_MULTIPROC_DIR", value = local.metrics_multiproc_dir }] : []
metrics_mount_points = local.metrics_enabled ? [{ sourceVolume = local.metrics_volume, containerPath = local.metrics_multiproc_dir }] : []
metrics_health_cmd = "import socket; socket.create_connection(('127.0.0.1', ${coalesce(var.gateway_metrics_port, 0)}), timeout=2).close()"
gateway_metrics_container = local.metrics_enabled ? [
{
name = "metrics"
image = var.gateway_image
essential = false
entryPoint = ["python", "-m", "litellm.proxy.prometheus_metrics_server"]
command = ["--port", tostring(var.gateway_metrics_port)]
portMappings = [{ containerPort = var.gateway_metrics_port, protocol = "tcp" }]
environment = local.metrics_env
mountPoints = local.metrics_mount_points
healthCheck = {
command = ["CMD", "python", "-c", local.metrics_health_cmd]
interval = 30
timeout = 5
retries = 3
startPeriod = 30
}
logConfiguration = {
logDriver = "awslogs"
options = {
awslogs-group = aws_cloudwatch_log_group.gateway.name
awslogs-region = var.region
awslogs-stream-prefix = "metrics"
}
}
}
] : []
backend_uvicorn_args = "--host 0.0.0.0 --port 4001"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac"
@ -269,7 +308,7 @@ resource "aws_ecs_task_definition" "gateway" {
execution_role_arn = aws_iam_role.task_execution.arn
task_role_arn = aws_iam_role.task.arn
container_definitions = jsonencode([
container_definitions = jsonencode(concat([
merge(
{
name = "gateway"
@ -283,8 +322,10 @@ resource "aws_ecs_task_definition" "gateway" {
local.billing_metrics_env,
local.gateway_extra_env_list,
local.proxy_config_env,
local.metrics_env,
)
secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list)
secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list)
mountPoints = local.metrics_mount_points
# Container-level healthCheck intentionally omitted — the wolfi
# runtime image doesn't ship curl/wget. The ALB target group polls
@ -301,7 +342,14 @@ resource "aws_ecs_task_definition" "gateway" {
},
local.gateway_proxy_overrides,
)
])
], local.gateway_metrics_container))
dynamic "volume" {
for_each = local.metrics_enabled ? [1] : []
content {
name = local.metrics_volume
}
}
tags = local.tags
}

3
terraform/litellm/aws/examples/default/main.tf
View file

@ -48,4 +48,7 @@ module "litellm" {
backend_extra_env = var.backend_extra_env
gateway_extra_secrets = var.gateway_extra_secrets
backend_extra_secrets = var.backend_extra_secrets
gateway_metrics_port = var.gateway_metrics_port
gateway_metrics_scrape_cidrs = var.gateway_metrics_scrape_cidrs
}

7
terraform/litellm/aws/examples/default/terraform.tfvars.example
View file

@ -102,6 +102,13 @@ env = "stage"
# }
# }
# ---------- Prometheus metrics sidecar ----------
# Serve /metrics from a sidecar in the gateway task instead of the inference
# workers. The port is not behind the ALB and has no auth: open it only to
# your Prometheus subnets.
# gateway_metrics_port = 4001
# gateway_metrics_scrape_cidrs = ["10.0.0.0/16"]
# ---------- Extra env / secrets ----------
# Plain-text env vars (non-sensitive). Land directly in the ECS task def.
# gateway_extra_env = {

12
terraform/litellm/aws/examples/default/variables.tf
View file

@ -158,3 +158,15 @@ variable "backend_extra_secrets" {
type = map(string)
default = {}
}
variable "gateway_metrics_port" {
description = "Port for the Prometheus metrics sidecar in the gateway task. Null keeps /metrics on the gateway port only."
type = number
default = null
}
variable "gateway_metrics_scrape_cidrs" {
description = "CIDRs allowed to scrape gateway_metrics_port."
type = list(string)
default = []
}

11
terraform/litellm/aws/network.tf
View file

@ -156,6 +156,17 @@ resource "aws_security_group" "tasks" {
security_groups = [aws_security_group.alb.id]
}
dynamic "ingress" {
for_each = local.metrics_enabled && length(var.gateway_metrics_scrape_cidrs) > 0 ? [1] : []
content {
description = "Prometheus scrapers to the gateway metrics sidecar"
from_port = var.gateway_metrics_port
to_port = var.gateway_metrics_port
protocol = "tcp"
cidr_blocks = var.gateway_metrics_scrape_cidrs
}
}
egress {
description = "All egress (LLM providers, RDS, Redis)"
from_port = 0

108
terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl Normal file
View file

@ -0,0 +1,108 @@
# Plan-only coverage for the Prometheus metrics sidecar wiring. Offline via
# mock_provider, same as byo_infrastructure.tftest.hcl.
mock_provider "aws" {
mock_data "aws_iam_policy_document" {
defaults = {
json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}"
}
}
}
mock_provider "random" {}
variables {
region = "us-east-1"
tenant = "acme"
env = "test"
allow_plaintext_alb = true
azs = ["us-east-1a", "us-east-1b"]
}
run "defaults_change_nothing" {
command = plan
assert {
condition = alltrue([
length(local.gateway_metrics_container) == 0,
length(local.metrics_env) == 0,
length(local.metrics_mount_points) == 0,
length([for r in aws_security_group.tasks.ingress : r if r.description == "Prometheus scrapers to the gateway metrics sidecar"]) == 0,
])
error_message = "The metrics sidecar, its env, its volume, and its security-group rule must all be absent by default."
}
}
run "metrics_port_adds_a_sidecar_volume_and_scrape_rule" {
command = plan
variables {
gateway_metrics_port = 9464
gateway_metrics_scrape_cidrs = ["10.20.0.0/16"]
}
assert {
condition = length(local.metrics_env) == 1 && local.metrics_env[0].name == "PROMETHEUS_MULTIPROC_DIR" && local.metrics_env[0].value == "/tmp/litellm_prometheus_multiproc"
error_message = "The gateway workers must write multiprocess samples to the shared dir."
}
assert {
condition = length(local.metrics_mount_points) == 1 && local.metrics_mount_points[0].sourceVolume == "prometheus-multiproc" && local.metrics_mount_points[0].containerPath == "/tmp/litellm_prometheus_multiproc"
error_message = "Gateway and sidecar must mount the same task volume at the multiproc dir."
}
assert {
condition = alltrue([
length(local.gateway_metrics_container) == 1,
local.gateway_metrics_container[0].name == "metrics",
local.gateway_metrics_container[0].essential == false,
join(" ", local.gateway_metrics_container[0].entryPoint) == "python -m litellm.proxy.prometheus_metrics_server",
join(" ", local.gateway_metrics_container[0].command) == "--port 9464",
one(local.gateway_metrics_container[0].portMappings).containerPort == 9464,
one(local.gateway_metrics_container[0].environment).value == "/tmp/litellm_prometheus_multiproc",
one(local.gateway_metrics_container[0].mountPoints).sourceVolume == "prometheus-multiproc",
strcontains(local.gateway_metrics_container[0].healthCheck.command[3], "9464"),
])
error_message = "The metrics sidecar must run prometheus_metrics_server on the configured port, share the multiproc volume, and health-check that port."
}
assert {
condition = length(aws_ecs_task_definition.gateway.volume) == 1 && one(aws_ecs_task_definition.gateway.volume).name == "prometheus-multiproc"
error_message = "The gateway task must declare the multiproc volume."
}
assert {
condition = length([
for r in aws_security_group.tasks.ingress : r
if r.from_port == 9464 && r.to_port == 9464 && r.protocol == "tcp" && r.cidr_blocks == tolist(["10.20.0.0/16"])
]) == 1
error_message = "The scrape CIDRs must be allowed to reach the metrics port on the tasks security group."
}
assert {
condition = aws_lb_target_group.gateway.port == 4000 && one(aws_ecs_service.gateway.load_balancer).container_port == 4000
error_message = "The ALB must keep targeting the gateway port only; the metrics port is never load balanced."
}
}
run "metrics_port_without_scrape_cidrs_opens_nothing" {
command = plan
variables {
gateway_metrics_port = 9464
}
assert {
condition = length(local.gateway_metrics_container) == 1 && length([for r in aws_security_group.tasks.ingress : r if r.from_port == 9464]) == 0
error_message = "Without scrape CIDRs the sidecar runs but the metrics port stays closed to everything but the ALB group."
}
}
run "metrics_port_may_not_reuse_the_gateway_port" {
command = plan
variables {
gateway_metrics_port = 4000
}
expect_failures = [var.gateway_metrics_port]
}

38
terraform/litellm/aws/variables.tf
View file

@ -549,6 +549,44 @@ variable "proxy_config" {
default = {}
}
# ---------- Prometheus metrics sidecar ----------
variable "gateway_metrics_port" {
description = <<-EOT
Serve Prometheus /metrics from a `metrics` sidecar container in the
gateway task on this port (1-65535, not 4000), so a scrape never runs on
an inference worker. The sidecar runs the gateway image with
`python -m litellm.proxy.prometheus_metrics_server` and aggregates the
workers' PROMETHEUS_MULTIPROC_DIR samples over a task volume. Null (the
default) leaves /metrics on the gateway port only. The sidecar port has
no virtual-key auth and is not routed through the ALB; open it to your
scrapers with gateway_metrics_scrape_cidrs. Needs gateway_image v1.101.0
or newer.
EOT
type = number
default = null
validation {
condition = var.gateway_metrics_port == null || (var.gateway_metrics_port >= 1 && var.gateway_metrics_port <= 65535 && var.gateway_metrics_port != 4000)
error_message = "gateway_metrics_port must be between 1 and 65535 and must not be 4000 (the gateway port)."
}
}
variable "gateway_metrics_scrape_cidrs" {
description = <<-EOT
CIDR blocks allowed to reach gateway_metrics_port on the gateway tasks
(your Prometheus or collector subnets). Empty by default, so only the
ALB can reach the tasks. Ignored when gateway_metrics_port is null.
EOT
type = list(string)
default = []
validation {
condition = alltrue([for c in var.gateway_metrics_scrape_cidrs : can(cidrnetmask(c))])
error_message = "gateway_metrics_scrape_cidrs must contain valid IPv4 CIDR blocks."
}
}
variable "log_retention_days" {
description = "CloudWatch log retention for the three services."
type = number

46
tests/local_testing/test_least_busy_routing.py
View file

@ -33,8 +33,8 @@ def test_model_added():
}
}
least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs)
request_count_api_key = f"gpt-3.5-turbo_request_count"
assert test_cache.get_cache(key=request_count_api_key) is not None
request_count_api_key = "gpt-3.5-turbo_request_count:1234"
assert test_cache.get_cache(key=request_count_api_key) == 1
def test_get_available_deployments():
@ -52,8 +52,8 @@ def test_get_available_deployments():
}
}
least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs)
request_count_api_key = f"{model_group}_request_count"
assert test_cache.get_cache(key=request_count_api_key) is not None
request_count_api_key = f"{model_group}_request_count:1234"
assert test_cache.get_cache(key=request_count_api_key) == 1
# test_get_available_deployments()
@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test):
router.leastbusy_logger.test_flag = True
model_group = "azure-model"
request_count_dict = {1: 10, 2: 54, 3: 100}
cache_key = f"{model_group}_request_count"
request_count_dict = {"1": 10, "2": 54, "3": 100}
cache_keys = {
deployment_id: f"{model_group}_request_count:{deployment_id}"
for deployment_id in request_count_dict
}
if async_test is True:
await router.cache.async_set_cache(key=cache_key, value=request_count_dict)
for deployment_id, count in request_count_dict.items():
await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count)
deployment = await router.async_get_available_deployment(
model=model_group, messages=None, request_kwargs={}
)
else:
router.cache.set_cache(key=cache_key, value=request_count_dict)
for deployment_id, count in request_count_dict.items():
router.cache.set_cache(key=cache_keys[deployment_id], value=count)
deployment = router.get_available_deployment(model=model_group, messages=None)
print(f"deployment: {deployment}")
assert deployment["model_info"]["id"] == "1"
@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test):
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
return_dict = router.cache.get_cache(key=cache_key)
# wait 2 seconds
time.sleep(2)
return_dict = {
deployment_id: router.cache.get_cache(key=cache_key)
for deployment_id, cache_key in cache_keys.items()
}
assert router.leastbusy_logger.logged_success == 1
assert return_dict[1] == 10
assert return_dict[2] == 54
assert return_dict[3] == 100
assert return_dict["1"] == 10
assert return_dict["2"] == 54
assert return_dict["3"] == 100
## Test with Real calls ##
@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming():
await asyncio.sleep(random.uniform(0, 2))
await router.atext_completion(model=model, prompt=prompt, stream=True)
cache_key = f"{model}_request_count"
## check if calls equally distributed
cache_dict = router.cache.get_cache(key=cache_key)
cache_dict = {
deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}")
for deployment_id in ("1", "2", "3")
}
for k, v in cache_dict.items():
assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}"
@ -259,8 +269,10 @@ async def test_router_completion_streaming():
await asyncio.sleep(random.uniform(0, 2))
await router.acompletion(model=model, messages=messages, stream=True)
cache_key = f"{model}_request_count"
## check if calls equally distributed
cache_dict = router.cache.get_cache(key=cache_key)
cache_dict = {
deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}")
for deployment_id in ("1", "2", "3")
}
for k, v in cache_dict.items():
assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}"

136
tests/local_testing/test_lowest_latency_routing.py
View file

@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async():
), f"Oldest latency {oldest_latency} should have been discarded"
def test_ttft_list_trimming_discards_oldest_entry():
"""
The time_to_first_token list trims the oldest entry when full, matching
the behavior of the latency list.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
ttft_values = []
for i in range(max_size + 1):
start_time = time.time()
expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4
completion_start_time = start_time + expected_ttft
end_time = start_time + float(i + 1)
ttft_values.append(expected_ttft)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
},
"stream": True,
"completion_start_time": completion_start_time,
}
# TTFT is only recorded when response_obj is a ModelResponse.
response_obj = litellm.ModelResponse(
usage=litellm.Usage(completion_tokens=1, total_tokens=1)
)
lowest_latency_logger.log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = test_cache.get_cache(key=latency_key)
ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
assert (
len(ttft_list) == max_size
), f"Expected {max_size} entries, got {len(ttft_list)}"
newest_ttft = ttft_values[-1]
oldest_ttft = ttft_values[0]
tolerance = 0.05
assert (
abs(ttft_list[-1] - newest_ttft) < tolerance
), f"Newest TTFT {newest_ttft} should be at end of list"
for ttft in ttft_list:
assert (
abs(ttft - oldest_ttft) > tolerance
), f"Oldest TTFT {oldest_ttft} should have been discarded"
@pytest.mark.asyncio
async def test_timeout_penalty_discards_oldest_entry():
"""
@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims():
assert (
abs(latency_list[i] - expected) < tolerance
), f"At index {i}, expected ~{expected}, got {latency_list[i]}"
@pytest.mark.asyncio
async def test_ttft_list_trimming_discards_oldest_entry_async():
"""
Async counterpart: the time_to_first_token list trims the oldest entry
when full. Exercises the async_log_success_event TTFT path, which only
runs when response_obj is a ModelResponse and the call is marked as
streaming with a completion_start_time.
"""
max_size = 3
test_cache = DualCache()
lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, routing_args={"max_latency_list_size": max_size}
)
model_group = "gpt-3.5-turbo"
deployment_id = "test-deployment"
ttft_values = []
for i in range(max_size + 1):
start_time = time.time()
expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4
completion_start_time = start_time + expected_ttft
end_time = start_time + float(i + 1)
ttft_values.append(expected_ttft)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": model_group,
"deployment": "azure/gpt-4.1-mini",
},
"model_info": {"id": deployment_id},
},
"stream": True,
"completion_start_time": completion_start_time,
}
response_obj = litellm.ModelResponse(
usage=litellm.Usage(completion_tokens=1, total_tokens=1)
)
await lowest_latency_logger.async_log_success_event(
response_obj=response_obj,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
latency_key = f"{model_group}_map"
cached_data = await test_cache.async_get_cache(key=latency_key)
ttft_list = cached_data[deployment_id].get("time_to_first_token", [])
assert (
len(ttft_list) == max_size
), f"Expected {max_size} entries, got {len(ttft_list)}"
newest_ttft = ttft_values[-1]
oldest_ttft = ttft_values[0]
tolerance = 0.05
assert (
abs(ttft_list[-1] - newest_ttft) < tolerance
), f"Newest TTFT {newest_ttft} should be at end of list"
for ttft in ttft_list:
assert (
abs(ttft - oldest_ttft) > tolerance
), f"Oldest TTFT {oldest_ttft} should have been discarded"

80
tests/local_testing/test_redis_increment_with_floor.py Normal file
View file

@ -0,0 +1,80 @@
"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the
create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so
these are the only tests that fail when the script itself is wrong."""
import os
import uuid
from typing import Final
import pytest
from dotenv import load_dotenv
load_dotenv()
from litellm.caching.redis_cache import RedisCache
TTL: Final = 600
@pytest.fixture
def counter():
cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT"))
key: Final = f"lit7039-{uuid.uuid4()}"
yield cache, key, cache.check_and_fix_namespace(key=key)
cache.delete_cache(key)
def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter):
cache, key, _ = counter
assert cache.increment_with_floor(key, 3, TTL) == 3
assert cache.increment_with_floor(key, 2, TTL) == 5
assert cache.batch_get_counts([key]) == (5,)
def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter):
"""A worker whose counter expired mid-request decrements a key that is no longer there.
Without the clamp that deployment reads negative, and least-busy pins every later request
on it until the count climbs back to zero."""
cache, key, _ = counter
assert cache.increment_with_floor(key, 1, TTL) == 1
assert cache.increment_with_floor(key, -5, TTL) == 0
assert cache.batch_get_counts([key]) == (0,)
def test_traffic_never_pushes_a_counters_expiry_back_out(counter):
"""The TTL is what releases a count whose worker died mid-request. Rewriting it on every
touch would keep that stuck count alive for as long as the group takes traffic."""
cache, key, namespaced_key = counter
cache.increment_with_floor(key, 1, TTL)
assert cache.redis_client.ttl(namespaced_key) > TTL - 60
cache.redis_client.expire(namespaced_key, 30)
cache.increment_with_floor(key, 1, TTL)
assert cache.redis_client.ttl(namespaced_key) <= 30
def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter):
cache, key, namespaced_key = counter
cache.increment_with_floor(key, 1, TTL)
cache.redis_client.expire(namespaced_key, 30)
assert cache.increment_with_floor(key, -5, TTL) == 0
assert cache.redis_client.ttl(namespaced_key) <= 30
@pytest.mark.asyncio
async def test_the_async_counter_behaves_the_same_way(counter):
cache, key, namespaced_key = counter
assert await cache.async_increment_with_floor(key, 2, TTL) == 2
assert await cache.async_batch_get_counts([key]) == (2,)
cache.redis_client.expire(namespaced_key, 30)
assert await cache.async_increment_with_floor(key, -9, TTL) == 0
assert cache.redis_client.ttl(namespaced_key) <= 30

264
tests/proxy_migration_tests/test_invalid_index_repair.py Normal file
View file

@ -0,0 +1,264 @@
import os
import threading
import uuid
from collections.abc import Iterator, Mapping
from types import MappingProxyType
from typing import Final
import pytest
from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager
psycopg = pytest.importorskip("psycopg")
pytestmark = pytest.mark.timeout(120)
requires_db: Final = pytest.mark.skipif(
"DATABASE_URL" not in os.environ,
reason="requires a postgres database (DATABASE_URL)",
)
HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable"
HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx"
HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC'
LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable"
LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx"
PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable"
PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx"
def _base_url() -> str:
return os.environ["DATABASE_URL"].split("?")[0]
def _index_validity(schema: str) -> Mapping[str, bool]:
with psycopg.connect(_base_url(), autocommit=True) as conn:
rows = conn.execute(
"SELECT c.relname, i.indisvalid FROM pg_index i "
"JOIN pg_class c ON c.oid = i.indexrelid "
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"WHERE n.nspname = %s",
(schema,),
).fetchall()
return MappingProxyType(dict(rows))
def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None:
"""Abort a CONCURRENTLY build while it waits on an older snapshot, the same
spot the deadlock loser dies at, so it leaves its index INVALID."""
with psycopg.connect(_base_url()) as pin:
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
with psycopg.connect(_base_url(), autocommit=True) as builder:
builder.execute("SET statement_timeout = '1s'")
with pytest.raises(psycopg.errors.QueryCanceled):
builder.execute(statement)
def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None:
_interrupt_concurrent_build(
schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})'
)
def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None:
_interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"')
@pytest.fixture
def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'CREATE SCHEMA "{schema}"')
conn.execute(
f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)'
)
conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)')
monkeypatch.delenv("DIRECT_URL", raising=False)
monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}")
yield schema
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'DROP SCHEMA "{schema}" CASCADE')
@pytest.fixture
def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
"""A brand-new database, what a first deploy sees. A scratch schema would
not do: the migrations guard on pg_constraint by name across every schema,
so a LiteLLM schema already pushed into public makes them skip and then
fail, which is exactly what CI's database looks like."""
admin_url: Final = _base_url()
name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
with psycopg.connect(admin_url, autocommit=True) as conn:
conn.execute(f'CREATE DATABASE "{name}"')
monkeypatch.delenv("DIRECT_URL", raising=False)
monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}")
yield "public"
with psycopg.connect(admin_url, autocommit=True) as conn:
conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)')
@requires_db
def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None:
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
_leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id")
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False}
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False}
@requires_db
def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None:
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
_leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX)
_leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"')
_leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"')
before: Final = _index_validity(scratch_schema)
assert len(before) == 4
assert set(before.values()) == {False}
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
@requires_db
def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None:
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})')
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
@requires_db
def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None:
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)')
conn.execute(
f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" '
f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)'
)
conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)')
assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
@requires_db
def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None:
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
with psycopg.connect(_base_url(), autocommit=True) as other_replica:
other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,))
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False}
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
@requires_db
def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None:
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
with psycopg.connect(_base_url()) as pin:
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"')
assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False
blocked: Final = _index_validity(scratch_schema)
assert blocked[HEALTH_INDEX] is False
assert [name for name in blocked if name.endswith("_ccnew")]
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None:
with psycopg.connect(_base_url()) as pin:
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
pinned.set()
pin.execute("SELECT pg_sleep(%s)", (seconds,))
@requires_db
def test_repair_outlives_a_statement_timeout_passed_through_database_url_options(
scratch_schema: str, monkeypatch: pytest.MonkeyPatch
) -> None:
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000")
pinned: Final = threading.Event()
holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0))
holder.start()
pinned.wait()
try:
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
finally:
holder.join()
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
@requires_db
def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None:
table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}"
index: Final = f"{table}_id_idx"
monkeypatch.setenv("DATABASE_URL", _base_url())
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'CREATE TABLE public."{table}" (id TEXT)')
try:
_leave_invalid_index("public", table, index, "id")
assert _index_validity("public")[index] is False
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity("public")[index] is True
finally:
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'DROP TABLE public."{table}"')
def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("DIRECT_URL", raising=False)
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever")
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
@requires_db
def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None:
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
with pytest.MonkeyPatch.context() as env:
env.setenv("DIRECT_URL", f"{_base_url()}?schema=public")
env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}")
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
def _invalidate_deployed_index(schema: str) -> None:
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"')
_leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
@requires_db
@pytest.mark.timeout(300)
@pytest.mark.parametrize("use_v2_resolver", [True, False])
def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None:
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
_invalidate_deployed_index(fresh_database)
assert _index_validity(fresh_database)[HEALTH_INDEX] is False
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
assert _index_validity(fresh_database)[HEALTH_INDEX] is True

16
tests/router_unit_tests/test_router_cooldown_per_deployment.py
View file

@ -246,12 +246,12 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time() - 120.0,
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600)
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
assert active == [], "Expired cooldown entry must not appear in active cooldowns"
assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
def test_active_entry_is_returned(self):
"""
@ -267,7 +267,7 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time(),
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60)
cc.in_memory_cache.set_cache(key, active_value, ttl=60)
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
@ -290,14 +290,14 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time() - (60.0 - remaining),
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, value, ttl=600)
cc.in_memory_cache.set_cache(key, value, ttl=600)
before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
before_expiry = cc.in_memory_cache.ttl_dict.get(key)
assert before_expiry is not None
cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
after_expiry = cc.in_memory_cache.ttl_dict.get(key)
assert after_expiry is not None
corrected_remaining = after_expiry - time.time()
assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
@ -318,12 +318,12 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time() - 120.0,
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600)
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
assert active == [], "Expired entry must not appear in async active cooldowns"
assert cc.cache.in_memory_cache.get_cache(key) is None
assert cc.in_memory_cache.get_cache(key) is None
class TestFallbackDeploymentCooldown:

44
tests/test_litellm/caching/test_redis_cache.py
View file

@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re
assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {}
def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache):
"""A caller that must fall back when Redis is unreachable needs the failure, not zeros.
The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell
apart from "every counter is unset". Least-busy routing read that as an idle deployment and
kept sending traffic to it instead of falling back to this worker's own in-flight counts.
"""
assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {}
with pytest.raises(OSError, match="redis unavailable"):
sync_batch_redis_cache.batch_get_counts(["lit7039"])
@pytest.mark.asyncio
async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None):
"""Async twin: the async batch read hides the same failure behind an empty dict."""
failing_client = AsyncMock()
failing_client.mget.side_effect = OSError("redis unavailable")
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
"litellm._redis.get_redis_client", return_value=MagicMock()
):
cache = RedisCache(host="127.0.0.1", port=6379)
with patch.object(cache, "init_async_client", return_value=failing_client):
assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {}
with pytest.raises(OSError, match="redis unavailable"):
await cache.async_batch_get_counts(["lit7039"])
@pytest.mark.parametrize("stored", [b"3", "3"])
def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None):
"""Counters come back positionally, so an unset key has to stay a hole rather than shift the
rest of the row onto the wrong deployments, and a count has to survive whether the client
hands it back as bytes or as text."""
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
"litellm._redis.get_redis_client", return_value=MagicMock()
):
cache = RedisCache(host="127.0.0.1", port=6379)
cache.redis_client.mget.return_value = [stored, None, b"0"]
assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0)
@pytest.fixture
def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]:
service_logger = ServiceLogging(mock_testing=True)

140
tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
View file

@ -4147,3 +4147,143 @@ def test_streaming_final_chunk_carries_provider_metadata():
assert chunks[-1]["content_filters"] == content_filters
assert "background" not in chunks[-1]
assert all("service_tier" not in chunk for chunk in chunks[:-1])
def _system_input_item(text: str) -> dict[str, object]:
return {"type": "message", "role": "system", "content": [{"type": "input_text", "text": text}]}
def test_mid_conversation_system_string_stays_in_input_after_a_user_turn():
handler: Final = LiteLLMResponsesTransformationHandler()
input_items, instructions = handler.convert_chat_completion_messages_to_responses_api(
[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Read the file."},
{"role": "system", "content": "<total_tokens>14982391 tokens left</total_tokens>"},
{"role": "user", "content": "Now summarize it."},
]
)
assert instructions == "You are a helpful assistant."
assert input_items == [
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Read the file."}]},
_system_input_item("<total_tokens>14982391 tokens left</total_tokens>"),
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Now summarize it."}]},
]
def test_leading_system_strings_still_join_instructions_without_a_following_turn():
handler: Final = LiteLLMResponsesTransformationHandler()
input_items, instructions = handler.convert_chat_completion_messages_to_responses_api(
[
{"role": "system", "content": "Be brief."},
{"role": "system", "content": "Answer in French."},
]
)
assert instructions == "Be brief. Answer in French."
assert input_items == []
def test_mid_conversation_system_reminder_as_string_and_as_text_block_produce_identical_input_items():
handler: Final = LiteLLMResponsesTransformationHandler()
reminder: Final = "<total_tokens>14982391 tokens left</total_tokens>"
as_string, string_instructions = handler.convert_chat_completion_messages_to_responses_api(
[{"role": "user", "content": "Read the file."}, {"role": "system", "content": reminder}]
)
as_block, block_instructions = handler.convert_chat_completion_messages_to_responses_api(
[
{"role": "user", "content": "Read the file."},
{
"role": "system",
"content": [{"type": "text", "text": reminder, "cache_control": {"type": "ephemeral"}}],
},
]
)
assert string_instructions is None
assert block_instructions is None
assert json.dumps(as_string) == json.dumps(as_block)
assert as_string[1] == _system_input_item(reminder)
def test_claude_code_shaped_history_keeps_a_byte_stable_input_prefix_across_requests():
handler: Final = LiteLLMResponsesTransformationHandler()
top_level_system: Final = [{"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}]
first_reminder: Final = "<system-reminder>27k chars of deferred tools</system-reminder>"
second_reminder: Final = "<total_tokens>14982391 tokens left</total_tokens>"
first_request_messages: Final = [
{"role": "system", "content": top_level_system},
{"role": "user", "content": "Read inventory.py."},
{
"role": "system",
"content": [{"type": "text", "text": first_reminder, "cache_control": {"type": "ephemeral"}}],
},
]
second_request_messages: Final = [
{"role": "system", "content": top_level_system},
{"role": "user", "content": "Read inventory.py."},
{"role": "system", "content": first_reminder},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "Read", "arguments": '{"file_path": "inventory.py"}'},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "ITEMS = []"},
{
"role": "system",
"content": [{"type": "text", "text": second_reminder, "cache_control": {"type": "ephemeral"}}],
},
]
first_request: Final = handler.transform_request(
model="gpt-5.6-luna",
messages=first_request_messages,
optional_params={},
litellm_params={},
headers={},
litellm_logging_obj=Mock(),
)
second_request: Final = handler.transform_request(
model="gpt-5.6-luna",
messages=second_request_messages,
optional_params={},
litellm_params={},
headers={},
litellm_logging_obj=Mock(),
)
assert "instructions" not in first_request
assert "instructions" not in second_request
assert first_request["input"][0] == _system_input_item("You are Claude Code.")
assert json.dumps(second_request["input"][: len(first_request["input"])]) == json.dumps(first_request["input"])
assert second_request["input"][len(first_request["input"]) :] == [
{"type": "function_call", "call_id": "call_1", "name": "Read", "arguments": '{"file_path": "inventory.py"}'},
{"type": "function_call_output", "call_id": "call_1", "output": [{"type": "input_text", "text": "ITEMS = []"}]},
_system_input_item(second_reminder),
]
def test_system_string_after_a_developer_message_stays_in_input_in_client_order():
handler: Final = LiteLLMResponsesTransformationHandler()
input_items, instructions = handler.convert_chat_completion_messages_to_responses_api(
[
{"role": "developer", "content": "Always answer in French."},
{"role": "system", "content": "Be brief."},
{"role": "user", "content": "Bonjour"},
]
)
assert instructions is None
assert [item["role"] for item in input_items] == ["developer", "system", "user"]
assert input_items[1] == _system_input_item("Be brief.")

2
tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py
View file

@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks():
never sees ``data["prompt"]``."""
guardrail = _guardrail()
assert guardrail.uses_apply_guardrail_interface() is True
assert guardrail._deployment_pre_call_target() is guardrail
assert guardrail._deployment_hook_target() is guardrail
@pytest.mark.asyncio

133
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -1,5 +1,5 @@
import asyncio
from typing import TYPE_CHECKING, Literal, Optional
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
from unittest.mock import AsyncMock
import pytest
@ -10,6 +10,7 @@ from litellm.integrations.custom_guardrail import (
log_guardrail_information,
)
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail
if TYPE_CHECKING:
@ -2378,11 +2379,108 @@ def _logged_call(messages: list | str) -> tuple[dict, object]:
return kwargs, response
class _NativeApplyGuardrail(_InheritedApplyGuardrail):
use_native_lifecycle_hooks: ClassVar[bool] = True
@pytest.mark.parametrize("guardrail_type", (CustomGuardrail, _NativeApplyGuardrail, _InheritedApplyGuardrail))
@pytest.mark.parametrize(
"event_hook",
(
GuardrailEventHooks.logging_only,
"logging_only",
[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only],
["pre_call", "logging_only"],
Mode(tags={"audit": "logging_only"}, default="pre_call"),
Mode(tags={"audit": ["pre_call", "logging_only"]}),
Mode(tags={"enforce": "pre_call"}, default="logging_only"),
Mode(tags={}, default=["pre_call", "logging_only"]),
),
)
def test_logging_only_requires_framework_support_or_explicit_declaration(
guardrail_type: type[CustomGuardrail],
event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode,
) -> None:
supported: Final = [GuardrailEventHooks.pre_call]
if guardrail_type is _InheritedApplyGuardrail:
guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported)
assert guardrail.event_hook == event_hook
assert supported == [GuardrailEventHooks.pre_call]
else:
with pytest.raises(ValueError, match=r"logging_only.*not in the supported event hooks"):
guardrail_type(event_hook=event_hook, supported_event_hooks=supported)
explicitly_supported: Final = guardrail_type(
event_hook=event_hook,
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only],
)
assert explicitly_supported.event_hook == event_hook
@pytest.mark.parametrize(
"event_hook",
(
GuardrailEventHooks.post_call,
"post_call",
[GuardrailEventHooks.logging_only, GuardrailEventHooks.post_call],
["logging_only", "post_call"],
Mode(tags={"enforce": "post_call"}, default="logging_only"),
Mode(tags={"enforce": ["logging_only", "post_call"]}),
Mode(tags={"audit": "logging_only"}, default="post_call"),
Mode(tags={}, default=["logging_only", "post_call"]),
),
)
def test_framework_logging_only_does_not_allow_other_unsupported_modes(
event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode,
) -> None:
with pytest.raises(ValueError, match=r"post_call.*not in the supported event hooks"):
_InheritedApplyGuardrail(event_hook=event_hook, supported_event_hooks=[GuardrailEventHooks.pre_call])
class TestLoggingOnlyApplyGuardrail:
"""LIT-4876 regression: a guardrail in mode logging_only that implements only
apply_guardrail must still run against the logged request and response and
record guardrail_information, instead of inheriting the CustomLogger no-op."""
@pytest.mark.parametrize(
"event_hook",
(
GuardrailEventHooks.logging_only,
"logging_only",
[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only],
["pre_call", "logging_only"],
Mode(tags={"audit": "logging_only"}, default="pre_call"),
Mode(tags={"audit": ["pre_call", "logging_only"]}),
Mode(tags={"enforce": "pre_call"}, default="logging_only"),
Mode(tags={}, default=["pre_call", "logging_only"]),
),
)
@pytest.mark.asyncio
async def test_content_filter_accepts_logging_only_and_records_detection(
self, event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode
) -> None:
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
guardrail: Final = ContentFilterGuardrail(
guardrail_name="content-review",
event_hook=event_hook,
default_on=True,
blocked_words=[BlockedWord(keyword="hello", action=ContentFilterAction.BLOCK)],
)
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert out_response is response
assert out_kwargs["messages"] == kwargs["messages"]
assert (
out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"]
== "guardrail_intervened"
)
@pytest.mark.asyncio
async def test_runs_apply_guardrail_observe_only_and_records_verdict(self):
guardrail = _ApplyOnlyObserver()
@ -2610,3 +2708,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
)
assert result is replacement
@pytest.mark.asyncio
async def test_apply_guardrail_interface_modifies_deployment_response(self):
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import ModelResponse
class ReplacingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
assert input_type == "response"
return {**inputs, "texts": ["filtered response"]}
guardrail = ReplacingGuardrail(
guardrail_name="test-guardrail",
event_hook=GuardrailEventHooks.post_call,
)
response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}])
request_data = {"guardrails": ["test-guardrail"]}
result = await guardrail.async_post_call_success_deployment_hook(
request_data=request_data,
response=response,
call_type=CallTypes.acompletion,
)
assert result is response
assert response.choices[0].message.content == "filtered response"
assert request_data == {"guardrails": ["test-guardrail"]}

92
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -22,7 +22,13 @@ from litellm.litellm_core_utils.litellm_logging import (
_get_status_fields,
set_callbacks,
)
from litellm.types.utils import ModelResponse, TextCompletionResponse
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.utils import (
CallTypes,
LiteLLMRealtimeStreamLoggingObject,
ModelResponse,
TextCompletionResponse,
)
@pytest.fixture
@ -6393,6 +6399,90 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o
assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == ""
def _responses_ws_logging_obj() -> LitellmLogging:
return LitellmLogging(
model="gpt-4o",
messages=[],
stream=False,
call_type=CallTypes.aresponses_websocket.value,
start_time=time.time(),
litellm_call_id="responses-ws-usage-test",
function_id="responses-ws-usage-test",
)
def test_normalize_logging_result_extracts_usage_for_responses_websocket(monkeypatch):
"""LIT-6512: native /v1/responses WebSocket sessions logged $0 spend because the usage
carried by stored response.completed events was never extracted. The session must cost
exactly what the same usage costs over HTTP /v1/responses, discounts included."""
monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.5})
logging_obj = _responses_ws_logging_obj()
events = [
{"type": "response.created", "response": {}},
{
"type": "response.completed",
"response": {"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}},
},
{
"type": "response.completed",
"response": {"usage": {"input_tokens": 60, "output_tokens": 10, "total_tokens": 70}},
},
]
normalized = logging_obj.normalize_logging_result(result=events)
assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject)
assert normalized.usage.prompt_tokens == 160
assert normalized.usage.completion_tokens == 50
ws_cost = litellm.completion_cost(
completion_response=normalized,
model="gpt-4o",
call_type=CallTypes.aresponses_websocket.value,
custom_llm_provider="openai",
)
http_cost = litellm.completion_cost(
completion_response=ResponsesAPIResponse(
id="resp-6512",
created_at=1700000000,
output=[],
usage=ResponseAPIUsage(input_tokens=160, output_tokens=50, total_tokens=210),
),
model="gpt-4o",
call_type=CallTypes.aresponses.value,
custom_llm_provider="openai",
)
assert ws_cost > 0
assert ws_cost == http_cost
def test_normalize_logging_result_bills_incomplete_responses_websocket_turns():
"""LIT-6512: a turn cut short by max_output_tokens ends in response.incomplete, which
OpenAI bills, so its usage counts toward the session like a completed turn."""
events = [
{
"type": "response.created",
"response": {"usage": {"input_tokens": 999, "output_tokens": 999, "total_tokens": 1998}},
},
{
"type": "response.incomplete",
"response": {"usage": {"input_tokens": 15, "output_tokens": 16, "total_tokens": 31}},
},
{
"type": "response.completed",
"response": {"usage": {"input_tokens": 40, "output_tokens": 4, "total_tokens": 44}},
},
{"type": "response.failed", "response": {"usage": None}},
]
normalized = _responses_ws_logging_obj().normalize_logging_result(result=events)
assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject)
assert normalized.usage.prompt_tokens == 55
assert normalized.usage.completion_tokens == 20
assert normalized.usage.total_tokens == 75
def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj):
"""LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead
recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs)."""

64
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.llm_http_handler import (
_rust_responses_websocket_enabled,
)
from litellm.llms.azure.videos.transformation import AzureVideoConfig
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
@ -37,6 +38,69 @@ from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, Trans
_ACTIVE_KEY = "_code_interpreter_interception_active"
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
OCR_RESPONSE = {
"pages": [{"index": 0, "markdown": "OCR output", "images": []}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1},
}
def _ocr_sync_client() -> HTTPHandler:
client = HTTPHandler()
client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)))
return client
def _ocr_async_client() -> AsyncHTTPHandler:
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(
transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))
)
return client
def test_ocr_calls_post_call_with_raw_provider_response():
logging_obj = Mock()
response = BaseLLMHTTPHandler().ocr(
model="mistral-ocr-latest",
document={"type": "document_url", "document_url": "https://example.com/document.pdf"},
optional_params={},
timeout=5,
logging_obj=logging_obj,
api_key="test-key",
api_base="https://api.mistral.ai/v1/ocr",
custom_llm_provider="mistral",
client=_ocr_sync_client(),
provider_config=MistralOCRConfig(),
)
assert response.pages[0].markdown == "OCR output"
logging_obj.post_call.assert_called_once()
assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE
@pytest.mark.asyncio
async def test_async_ocr_calls_post_call_with_raw_provider_response():
logging_obj = Mock()
response = await BaseLLMHTTPHandler().async_ocr(
model="mistral-ocr-latest",
document={"type": "document_url", "document_url": "https://example.com/document.pdf"},
optional_params={},
timeout=5,
logging_obj=logging_obj,
api_key="test-key",
api_base="https://api.mistral.ai/v1/ocr",
custom_llm_provider="mistral",
client=_ocr_async_client(),
provider_config=MistralOCRConfig(),
)
assert response.pages[0].markdown == "OCR output"
logging_obj.post_call.assert_called_once()
assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE
def test_prepare_fake_stream_request():
# Initialize the BaseLLMHTTPHandler

52
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -4,6 +4,7 @@ import hashlib
import json
import time
from base64 import urlsafe_b64encode
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING
from unittest.mock import AsyncMock, MagicMock, patch
@ -18,6 +19,57 @@ if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None):
credential = {"type": "oauth2", "access_token": access_token}
if refresh_token is not None:
credential["refresh_token"] = refresh_token
if expires_in_seconds is not None:
credential["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat()
if expires_at is not None:
credential["expires_at"] = expires_at
return credential
@pytest.mark.asyncio
@pytest.mark.parametrize(
("fields", "egress_has_token"),
[
(None, False),
({"access_token": "", "refresh_token": "refresh-token"}, False),
({}, True),
({"expires_at": "never"}, True),
({"expires_in_seconds": 600}, True),
({"expires_in_seconds": 30}, False),
({"expires_in_seconds": -300}, False),
({"expires_in_seconds": -300, "refresh_token": ""}, False),
({"expires_in_seconds": 30, "refresh_token": "refresh-token"}, True),
({"expires_in_seconds": -300, "refresh_token": "refresh-token"}, True),
],
)
async def test_vendor_credential_state_agrees_with_egress_token_resolution(monkeypatch, fields, egress_has_token):
from litellm.proxy._experimental.mcp_server import db as mcp_db
from litellm.proxy._experimental.mcp_server import discoverable_endpoints
monkeypatch.setattr(mcp_db, "MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60)
credential = _stored_grant(**fields) if fields is not None else None
read = AsyncMock(return_value=credential)
refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600))
prisma = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read)
monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh)
connect = await discoverable_endpoints._vendor_credential_state("user-1", "server-1")
read.assert_awaited_once_with(prisma, "user-1", "server-1")
refresh.assert_not_awaited()
egress = await mcp_db.resolve_valid_user_oauth_token(
user_id="user-1", server=MagicMock(), cred=credential, prisma_client=prisma
)
assert (egress is not None) is egress_has_token
assert connect == ("present" if egress_has_token else "absent")
# Fixture to mock IP address check for all MCP tests
# This prevents tests from failing due to IP-based access control
@pytest.fixture(autouse=True)

254
tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
aggregate_authorize,
aggregate_token,
complete_connect_flow,
describe_connect_flow,
introspect_gateway_token,
is_gateway_dcr_client_id,
is_proxy_api_resource,
@ -230,7 +231,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co
assert location.path == "/ui/connect"
params = parse_qs(location.query)
handle = params["connect_flow"][0]
assert params["connect_client"] == ["https://claude.ai"]
assert set(params) == {"connect_flow"}
set_cookie = response.headers["set-cookie"]
assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie
assert "HttpOnly" in set_cookie
@ -889,14 +890,60 @@ def _opened_principal(payload):
return admitted.principal
async def _finish_connect_page(response):
class _VendorCredential:
def __init__(self, state="present"):
self.calls = []
self.state = state
async def __call__(self, user_id, server_id):
self.calls.append((user_id, server_id))
return self.state
class _ServerReachability:
def __init__(self, reachable=True):
self.calls = []
self.reachable = reachable
async def __call__(self, user_id, server_id):
self.calls.append((user_id, server_id))
return self.reachable
async def _complete_page(response, scoped_server=None, vendor=None, reachable=None, cache=None, **overrides):
from unittest.mock import patch
handle, cookies = _flow_cookie_from(response)
completed = await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id="u1",
cache=DualCache(),
)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.return_value = scoped_server
return await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id="u1",
cache=cache or DualCache(),
lookup_vendor_credential=vendor or _VendorCredential(),
lookup_server_reachability=reachable or _ServerReachability(),
**overrides,
)
async def _describe_page(response, scoped_server=None, vendor=None, reachable=None, session_user_id="u1", cookies=None):
from unittest.mock import patch
handle, flow_cookies = _flow_cookie_from(response)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_id.return_value = scoped_server
return await describe_connect_flow(
request=_request("/authorize/flow", cookies=flow_cookies if cookies is None else cookies),
flow_handle=handle,
session_user_id=session_user_id,
lookup_vendor_credential=vendor or _VendorCredential(),
lookup_server_reachability=reachable or _ServerReachability(),
)
async def _finish_connect_page(response, scoped_server=None):
completed = await _complete_page(response, scoped_server=scoped_server)
return parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
@ -910,23 +957,43 @@ def _sealed_wire_json(sealed, prefix, debug_key):
@pytest.mark.asyncio
async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
"""LIT-4917: a per-server RFC 8707 resource naming a gateway-managed oauth2 server
seals that server into the flow. The connect page interlude runs exactly as before
(the scope restricts, it never skips consent), and the code minted at the finish step
and the session pair it redeems for are both scoped."""
"""LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2
server seals that server into the flow. The connect URL carries only the handle; the page
learns the scoped server and its vendor state from describe_connect_flow, and the finish
step refuses to mint a scoped code until that vendor credential exists, without burning
the flow. The code minted afterwards and the session pair it redeems for are both scoped."""
from unittest.mock import patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
github = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
manager.get_mcp_server_by_name.return_value = github
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
assert response.status_code == 303
assert "/ui/connect" in response.headers["location"]
location = urlparse(response.headers["location"])
assert location.path == "/ui/connect"
assert set(parse_qs(location.query)) == {"connect_flow"}
_, cookies = _flow_cookie_from(response)
assert (
_sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id"
)
code = await _finish_connect_page(response)
described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent"))
assert json.loads(described.body) == {
"state": "interactive",
"client_origin": "https://claude.ai",
"server_id": "github-id",
"server_name": "github",
"connected": False,
}
cache = DualCache()
premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache)
assert premature.status_code == 400
assert "authorize the requested MCP server" in json.loads(premature.body)["error_description"]
present = _VendorCredential("present")
completed = await _complete_page(response, scoped_server=github, vendor=present, cache=cache)
assert completed.status_code == 303
assert present.calls == [("u1", "github-id")]
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
assert (
_sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"]
== "github-id"
@ -952,9 +1019,11 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope():
)
async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, resolves):
"""Every resource shape outside 'exactly one gateway-managed server' keeps today's flow:
connect page interlude, and NONE of the minted artifacts carry the scope key on the
wire, not the flow cookie, not the code, not the session JWT, so an unscoped flow
started on a new pod completes on a pod whose strict models predate the claim."""
the generic connect grid (describe names no server, the finish step never consults the
vendor credential), and NONE of the minted
artifacts carry the scope key on the wire, not the flow cookie, not the code, not the
session JWT, so an unscoped flow started on a new pod completes on a pod whose strict
models predate the claim."""
import base64
from unittest.mock import patch
@ -963,10 +1032,18 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource,
manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server()
response = _scoped_authorize(client_id, resource)
assert response.status_code == 303
assert "/ui/connect" in response.headers["location"]
location = urlparse(response.headers["location"])
assert location.path == "/ui/connect"
assert set(parse_qs(location.query)) == {"connect_flow"}
_, cookies = _flow_cookie_from(response)
assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")
code = await _finish_connect_page(response)
vendor = _VendorCredential("absent")
described = await _describe_page(response, vendor=vendor)
assert json.loads(described.body)["state"] == "unscoped"
assert json.loads(described.body)["server_id"] is None
completed = await _complete_page(response, vendor=vendor)
assert vendor.calls == []
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")
token_response = await _redeem(code, client_id)
payload = json.loads(token_response.body)
@ -979,19 +1056,150 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource,
@pytest.mark.asyncio
async def test_scoped_authorize_delegate_server_stays_unscoped():
"""A delegate-auth oauth2 server is outside the gateway-managed set (its keyless flow is
upstream PKCE via the relay), so a resource naming it never scopes the gateway flow."""
upstream PKCE via the relay), so a resource naming it never scopes the gateway flow and
never narrows the connect page to it."""
from unittest.mock import patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True)
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
assert "/ui/connect" in response.headers["location"]
location = urlparse(response.headers["location"])
assert location.path == "/ui/connect"
assert set(parse_qs(location.query)) == {"connect_flow"}
assert json.loads((await _describe_page(response)).body)["server_id"] is None
code = await _finish_connect_page(response)
token_response = await _redeem(code, client_id)
assert _opened_principal(json.loads(token_response.body)).resource_server_id is None
@pytest.mark.asyncio
async def test_m2m_scoped_flow_mints_without_a_user_credential():
"""A client-credentials server is already authorized by its gateway service credential, so
a resource-scoped flow finishes without consulting the per-user vault."""
from unittest.mock import patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
m2m = _scoped_mcp_server(oauth2_flow="client_credentials")
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = m2m
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
vendor = _VendorCredential("unavailable")
described = await _describe_page(response, scoped_server=m2m, vendor=vendor)
assert json.loads(described.body)["state"] == "m2m"
assert json.loads(described.body)["connected"] is True
assert vendor.calls == []
completed = await _complete_page(response, scoped_server=m2m, vendor=vendor)
assert completed.status_code == 303
assert vendor.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("oauth2_flow", ["authorization_code", "client_credentials"])
async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow):
from unittest.mock import patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
server = _scoped_mcp_server(oauth2_flow=oauth2_flow)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = server
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
reachable = _ServerReachability(False)
vendor = _VendorCredential("present")
described = await _describe_page(response, scoped_server=server, reachable=reachable, vendor=vendor)
assert json.loads(described.body) == {
"state": "stale",
"client_origin": "https://claude.ai",
"server_id": None,
"server_name": None,
"connected": None,
}
cache = DualCache()
refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache)
assert refused.status_code == 400
assert vendor.calls == []
assert reachable.calls == [("u1", "github-id"), ("u1", "github-id")]
completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache)
assert completed.status_code == 303
@pytest.mark.asyncio
async def test_stale_scoped_flow_remains_distinct_from_unscoped():
"""A server removed after authorize stays a stale scoped flow, so the page cannot offer a
broader unscoped grant or report a misleading Finish action."""
from unittest.mock import patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = _scoped_mcp_server()
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
described = await _describe_page(response, scoped_server=None)
assert json.loads(described.body)["state"] == "stale"
assert json.loads(described.body)["connected"] is None
stale = await _complete_page(response, scoped_server=None)
assert stale.status_code == 400
assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available"
@pytest.mark.asyncio
async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credential():
"""Cancel is the escape hatch: a scoped user who cannot finish the vendor step still ends
the flow with access_denied and no credential lookup. A scoped server that is no longer
gateway-managed refuses to mint (nothing could serve that code) but also burns nothing."""
from unittest.mock import patch
client_id = (await _register([REDIRECT_URI]))["client_id"]
github = _scoped_mcp_server()
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = github
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
cache = DualCache()
skipped_reachability = _ServerReachability(False)
stale = await _complete_page(response, scoped_server=None, reachable=skipped_reachability, cache=cache)
assert stale.status_code == 400
assert skipped_reachability.calls == []
assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available"
described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("unavailable"))
assert described.status_code == 503
vendor = _VendorCredential("absent")
deny_reachability = _ServerReachability(False)
denied = await _complete_page(
response,
scoped_server=github,
vendor=vendor,
reachable=deny_reachability,
cache=cache,
decision="deny",
)
assert denied.status_code == 303
assert parse_qs(urlparse(denied.headers["location"]).query)["error"] == ["access_denied"]
assert vendor.calls == []
assert deny_reachability.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
"session_user_id, cookies, expected_status, expected_error",
[
("u1", {}, 400, "invalid_request"),
("u1", {"mcp_connect_flow_wrong": "garbage"}, 400, "invalid_request"),
(None, None, 401, "login_required"),
("u2", None, 403, "access_denied"),
],
)
async def test_describe_connect_flow_refuses_exactly_like_the_finish_step(
session_user_id, cookies, expected_status, expected_error
):
"""The page's read of the flow is gated the same way minting is: the HttpOnly cookie for
that handle must open and the signed-in user must be the sealed one. A lure link with a
made-up handle therefore learns nothing and starts nothing."""
client_id = (await _register([REDIRECT_URI]))["client_id"]
response = _authorize(client_id, session_user_id="u1")
described = await _describe_page(response, session_user_id=session_user_id, cookies=cookies)
assert described.status_code == expected_status
assert json.loads(described.body)["error"] == expected_error
@pytest.mark.asyncio
async def test_token_rejects_resource_conflicting_with_sealed_scope():
"""RFC 8707 section 2.2: redeeming a scoped code (or rotating a scoped refresh token)
@ -1005,7 +1213,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope():
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = github
response = _scoped_authorize(client_id, SCOPED_RESOURCE)
code = await _finish_connect_page(response)
code = await _finish_connect_page(response, scoped_server=github)
with patch(_MANAGER_PATCH) as manager:
manager.get_mcp_server_by_name.return_value = linear

7
tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py
View file

@ -824,7 +824,7 @@ class _StopFailingSubscriber(ConfigSyncSubscriber):
raise RuntimeError("stop failed")
async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None:
async def test_proxy_config_subscriber_resyncs_deployments_only() -> None:
from litellm.proxy.proxy_server import ProxyConfig
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()]))
@ -852,10 +852,7 @@ async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() ->
await callback()
await config.stop_config_sync_subscriber()
assert calls == [
("add_deployment", prisma_client, proxy_logging_obj),
("get_credentials", prisma_client, None),
]
assert calls == [("add_deployment", prisma_client, proxy_logging_obj)]
assert config.config_sync_subscriber is None
assert subscriber._task is None

45
tests/test_litellm/proxy/common_utils/test_registry_read_through.py
View file

@ -393,6 +393,51 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc
assert not proxy_server.MODEL_RECONCILE_LOCK.locked()
@pytest.mark.asyncio
async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from unittest.mock import AsyncMock, MagicMock
import litellm
import litellm.proxy.proxy_server as proxy_server
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
from litellm.types.utils import CredentialItem
rows: Final = [MagicMock()]
prisma_client: Final = MagicMock()
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows)
router: Final = MagicMock()
router.get_model_list.return_value = []
installed: Final = MagicMock()
async def load_credentials_from_db(prisma_client: object) -> None:
CredentialAccessor.upsert_credentials(
[
CredentialItem(
credential_name="openai-cred",
credential_values={"api_key": "sk-from-db"},
credential_info={},
)
]
)
def install_models(db_models: object) -> None:
installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred"))
monkeypatch.setattr(litellm, "credential_list", [])
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", None)
monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db)
monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models)
assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True
installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"})
@pytest.mark.asyncio
async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch):
from unittest.mock import AsyncMock, MagicMock

2
tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py
View file

@ -23,7 +23,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import (
)
OPERATOR_CRON_JOB_ID = "spend_log_cleanup_job"
SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "get_credentials_job", "add_deployment_job")
SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "add_deployment_job")
async def _noop() -> None: ...

79
tests/test_litellm/proxy/db/conftest.py
View file

@ -1,5 +1,11 @@
import json
import os
import signal
import sys
import time
from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
import pytest
@ -75,3 +81,76 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]:
def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DATABASE_URL", "about-to-be-unset")
monkeypatch.delenv("DATABASE_URL")
FAKE_PRISMA_CLI = """#!{python}
import json
import os
import pathlib
import subprocess
import sys
import time
calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"])
earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else []
with calls_file.open("a") as log:
print(json.dumps(sys.argv[1:]), file=log)
if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"):
grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"])
pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid))
time.sleep(600)
sys.exit(0)
"""
@dataclass(frozen=True, slots=True)
class FakePrismaCli:
"""A stand-in `prisma` on PATH, recording every invocation.
With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree
of its own, the way the real CLI wraps Node around a Rust schema engine, so a
timeout that kills only the direct child leaves the rest of that tree running.
"""
calls_file: Path
grandchild_pidfile: Path
@property
def calls(self) -> list[list[str]]:
if not self.calls_file.exists():
return []
return [json.loads(line) for line in self.calls_file.read_text().splitlines()]
def grandchild_is_gone(self, within_seconds: float) -> bool:
deadline = time.monotonic() + within_seconds
while time.monotonic() < deadline:
try:
os.kill(int(self.grandchild_pidfile.read_text()), 0)
except ProcessLookupError:
return True
time.sleep(0.05)
return False
@pytest.fixture
def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]:
bin_dir = tmp_path / "fakebin"
bin_dir.mkdir()
script = bin_dir / "prisma"
script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable))
script.chmod(0o755)
cli = FakePrismaCli(
calls_file=tmp_path / "calls.jsonl",
grandchild_pidfile=tmp_path / "grandchild.pid",
)
monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}")
monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file))
monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile))
monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1")
monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False)
yield cli
if cli.grandchild_pidfile.exists():
try:
os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL)
except ProcessLookupError:
pass

32
tests/test_litellm/proxy/db/test_check_migration.py
View file

@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker):
check_migration.verbose_logger.exception.assert_called_once()
actual_message = check_migration.verbose_logger.exception.call_args[0][0]
assert "prisma schema out of sync with db" in actual_message
@pytest.mark.timeout(30)
def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch):
"""
`prisma migrate diff` ran unbounded, so a database that never answers hung boot
before uvicorn ever started, and interrupting the proxy orphaned the schema engine.
"""
from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])
assert fake_prisma_cli.calls == [
["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x",
"--to-schema-datamodel", "./schema.prisma", "--script"]
]
assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch):
"""
Boot calls this helper directly, so an ImportError here takes the proxy down before
uvicorn starts. An install without the runner must lose the diagnostic, not the proxy.
"""
import sys
from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])

59
tests/test_litellm/proxy/db/test_prisma_client.py
View file

@ -10,7 +10,7 @@ from fastapi.testclient import TestClient
from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema
from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema
@pytest.fixture(autouse=True)
@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client(
mock_new_prisma.connect.assert_awaited_once()
def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"]
def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url):
"""`prisma db push` bypasses litellm-proxy-extras, so it needs its own call
into the opt-in REPLICA IDENTITY FULL step."""
from litellm.proxy.db.prisma_client import PrismaManager
@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
staticmethod(lambda: applied.append(True)),
)
with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run:
assert PrismaManager.setup_database(use_migrate=False) is True
assert PrismaManager.setup_database(use_migrate=False) is True
assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
assert applied == [True]
def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
"""A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the
primary key back to ("request_id"), which Postgres rejects; the guard must
fail fast with guidance instead of running the push."""
@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch):
monkeypatch.setattr(
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True)
)
with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached
"litellm.proxy.db.prisma_client.subprocess.run"
) as mock_run:
with pytest.raises(RuntimeError) as err:
PrismaManager.setup_database(use_migrate=False)
with pytest.raises(RuntimeError) as err:
PrismaManager.setup_database(use_migrate=False)
assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR
mock_run.assert_not_called()
assert fake_prisma_cli.calls == []
def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch):
def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url):
from litellm.proxy.db.prisma_client import PrismaManager
from litellm_proxy_extras.utils import ProxyExtrasDBManager
monkeypatch.setattr(
ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False)
)
with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic
"litellm.proxy.db.prisma_client.subprocess.run"
) as mock_run:
assert PrismaManager.setup_database(use_migrate=False) is True
assert PrismaManager.setup_database(use_migrate=False) is True
assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
assert fake_prisma_cli.calls == [DB_PUSH_ARGV]
def _entra_jwt(expires_in_seconds: int) -> str:
@ -377,3 +373,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch
with pytest.raises(RuntimeError, match="DATABASE_HOST"):
wrapper.get_rds_iam_token()
@pytest.mark.timeout(45)
def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch):
"""
A timed-out `db push` used to leave Node and the schema engine writing the schema,
so the next attempt pushed into a database the abandoned one was still mutating.
"""
monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False)
monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
assert PrismaManager.setup_database(use_migrate=False) is True
assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV]
assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot(
fake_prisma_cli, unset_database_url, monkeypatch
):
"""
An ImportError out of setup_database escapes the caller's RuntimeError handler and
kills boot, bypassing the operator's enforce_prisma_migration_check choice.
"""
monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
assert PrismaManager.setup_database(use_migrate=False) is False
assert fake_prisma_cli.calls == []

75
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py
View file

@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
"""
import json
import logging
import re
from unittest.mock import patch
@ -520,6 +521,80 @@ class TestToolPermissionGuardrail:
)
assert excinfo.value.status_code == 400
@pytest.mark.asyncio
async def test_async_pre_call_hook_without_tools_logs_skip_at_debug(self, caplog):
data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
with patch.object(self.guardrail, "should_run_guardrail", return_value=True):
result = await self.guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(default_in_memory_ttl=1),
data=data,
call_type="completion",
)
assert result is data
skip_levels = [r.levelno for r in caplog.records if "No tools or functions in data" in r.getMessage()]
assert skip_levels == [logging.DEBUG]
assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == []
@pytest.mark.asyncio
async def test_async_pre_call_hook_denied_tool_logs_at_info(self, caplog):
data = {"tools": [{"type": "function", "function": {"name": "Read"}}]}
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
with patch.object(self.guardrail, "should_run_guardrail", return_value=True):
with pytest.raises(HTTPException):
await self.guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(default_in_memory_ttl=1),
data=data,
call_type="completion",
)
denied_levels = [
r.levelno
for r in caplog.records
if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'"
]
assert denied_levels == [logging.INFO]
assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == []
@pytest.mark.asyncio
async def test_async_post_call_success_hook_denied_tool_logs_at_info(self, caplog):
tool_call = {"function": {"name": "Read", "arguments": "{}"}, "type": "function"}
response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})])
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
with patch.object(self.guardrail, "should_run_guardrail", return_value=True):
with pytest.raises(GuardrailRaisedException):
await self.guardrail.async_post_call_success_hook(
data={"guardrails": ["test-tool-permission"]},
user_api_key_dict=UserAPIKeyAuth(),
response=response,
)
denied_levels = [
r.levelno
for r in caplog.records
if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'"
]
assert denied_levels == [logging.INFO]
assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == []
def test_parse_tool_call_arguments_malformed_json_logs_warning(self, caplog):
tool_call = ChatCompletionMessageToolCall(function={"name": "Bash", "arguments": "{not json"}, id="call_1")
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
parsed, error = self.guardrail._parse_tool_call_arguments(tool_call)
assert parsed is None
assert error == "arguments could not be parsed"
warning_messages = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING]
assert len(warning_messages) == 1
assert warning_messages[0].startswith("Tool Permission Guardrail: Failed to decode arguments for tool Bash")
@pytest.mark.asyncio
async def test_async_pre_call_hook_blocks_legacy_functions(self):
data = {

141
tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
View file

@ -4,20 +4,26 @@ Tests for the pipeline executor.
Uses mock guardrails to validate pipeline execution without external services.
"""
import copy
from typing import Literal
from unittest.mock import MagicMock
import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeGuardrail,
)
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
)
from litellm.types.utils import CallTypesLiteral
try:
from fastapi.exceptions import HTTPException
@ -158,11 +164,146 @@ class ContentCheckGuardrail(CustomGuardrail):
return None
class RecordingGuardrail(CustomGuardrail):
def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True):
super().__init__(
guardrail_name=guardrail_name,
event_hook="pre_call",
default_on=True,
scan_raw_request=scan_raw_request,
)
self.block = block
def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool:
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"detected": ["aws_access_key"]},
request_data=data,
guardrail_status="guardrail_intervened" if self.block else "success",
)
if self.block:
raise HTTPException(status_code=400, detail="Content policy violation")
return copy.deepcopy(data)
# ─────────────────────────────────────────────────────────────────────────────
# Tests
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
@pytest.mark.parametrize("scan_raw_request", [False, True])
@pytest.mark.parametrize("on_fail", ["block", "modify_response"])
async def test_terminal_block_carries_guardrail_information_to_request(
monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"]
):
"""
Spend logging and the Guardrails Monitor read standard_logging_guardrail_information
off the caller's request dict. A blocking step records it on the executor's
working copy (or the raw-request snapshot), so the terminal result must carry it
back onto the request or the block is never counted.
"""
guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request)
monkeypatch.setattr(litellm, "callbacks", [guard])
data = {
"messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}],
"metadata": {"user_api_key_hash": "abc"},
}
result = await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="baseline-pii-protection",
raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}},
)
assert result.terminal_action == on_fail
recorded = data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"]
assert recorded[0]["guardrail_status"] == "guardrail_intervened"
assert data["metadata"]["user_api_key_hash"] == "abc"
assert "guardrails" not in data["metadata"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch):
"""A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step
that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two
dicts at once. Those must be carried back once while every step's own entry is kept."""
first = RecordingGuardrail(guardrail_name="pii-scan", block=False)
second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True)
monkeypatch.setattr(litellm, "callbacks", [first, second])
earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"}
data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
data["metadata"]["standard_logging_guardrail_information"] = [earlier]
result = await PipelineExecutor.execute_steps(
steps=[
PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True),
PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"),
],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="baseline-pii-protection",
raw_request_snapshot={
"messages": data["messages"],
"metadata": {"standard_logging_guardrail_information": [dict(earlier)]},
},
)
assert result.terminal_action == "block"
recorded = data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch):
"""Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller,
while the entries the raw snapshot already held before the pipeline ran are not copied again."""
guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False)
monkeypatch.setattr(litellm, "callbacks", [guard])
earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"}
data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
data["metadata"]["standard_logging_guardrail_information"] = [earlier]
result = await PipelineExecutor.execute_steps(
steps=[
PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"),
PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"),
],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="raw-scan-policy",
raw_request_snapshot={
"messages": data["messages"],
"metadata": {"standard_logging_guardrail_information": [dict(earlier)]},
},
)
assert result.terminal_action == "allow"
assert result.modified_data is not None
recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_escalation_step1_fails_step2_blocks(monkeypatch):

25
tests/test_litellm/proxy/proxy_server/test_lifecycle.py
View file

@ -22,6 +22,7 @@ import inspect
import json
import logging
import os
import subprocess
from collections.abc import Awaitable, Callable
from typing import List, Optional, Union
from unittest.mock import AsyncMock, MagicMock, patch
@ -749,6 +750,30 @@ async def test_proxy_startup_event_invalid_missing_app_arg_raises():
pass
@pytest.mark.asyncio
async def test_proxy_startup_event_prunes_dead_workers_live_gauges(tmp_path):
"""With PROMETHEUS_MULTIPROC_DIR set, a booting worker drops the live-gauge files of pids that no longer
exist, so a crashed worker's in-flight samples leave the aggregate as soon as its replacement starts."""
exited = subprocess.Popen(["true"])
assert exited.wait(timeout=30) == 0
stale = tmp_path / f"gauge_livesum_{exited.pid}.db"
stale.touch()
counter = tmp_path / f"counter_{exited.pid}.db"
counter.touch()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path)
with patch.dict(os.environ, clean_env, clear=True):
try:
async with proxy_startup_event(app=None):
pass
except Exception:
pass
assert not stale.exists()
assert counter.exists()
def test_otel_global_provider_published_after_callback_init():
"""The OTel V2 global-provider publish must run after callback
initialization in ``proxy_startup_event``.

123
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -2939,6 +2939,129 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch
fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing")
def _stub_add_deployment_collaborators(
monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock
) -> None:
from litellm.proxy import proxy_server
fake_router = MagicMock()
fake_router.get_model_list = MagicMock(return_value=[])
async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]:
return {}
monkeypatch.setattr(litellm, "credential_list", [])
monkeypatch.setattr(pc, "get_config", fake_get_config)
monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock())
monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock())
monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None))
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(proxy_server, "master_key", "sk-master")
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "proxy_config", pc)
monkeypatch.delenv("LITELLM_SALT_KEY", raising=False)
def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]:
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
return {
"credential_name": credential_name,
"credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")},
"credential_info": {"custom_llm_provider": "openai"},
}
def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock:
fake_prisma = MagicMock()
fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(
return_value=[_encrypted_credential_row(credential_name, api_key)]
)
return fake_prisma
@pytest.mark.asyncio
async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy import proxy_server
from litellm.utils import load_credentials_from_list
pc = ProxyConfig()
fake_prisma = MagicMock()
fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
monkeypatch.setattr(proxy_server, "general_settings", {})
installed = MagicMock()
async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]:
fake_prisma.db.litellm_credentialstable.find_many.return_value = [
_encrypted_credential_row("openai-cred", "sk-from-db")
]
return [MagicMock()]
async def install_models(new_models: object, proxy_logging_obj: object) -> None:
installed(credential=CredentialAccessor.get_credential_values("openai-cred"))
monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands)
monkeypatch.setattr(pc, "_update_llm_router", install_models)
await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
installed.assert_called_once_with(credential={"api_key": "sk-from-db"})
assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"}
request_kwargs = {"litellm_credential_name": "openai-cred"}
load_credentials_from_list(request_kwargs)
assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"}
@pytest.mark.asyncio
async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy import proxy_server
pc = ProxyConfig()
fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db")
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]})
models_fetch = AsyncMock(return_value=[])
monkeypatch.setattr(pc, "_get_models_from_db", models_fetch)
await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
models_fetch.assert_not_awaited()
assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"}
@pytest.mark.asyncio
async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
pc = ProxyConfig()
writer_inner = MagicMock(name="writer_prisma")
reader_inner = MagicMock(name="reader_prisma")
writer_inner.litellm_credentialstable.find_many = AsyncMock(
return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")]
)
reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[])
fake_prisma = MagicMock()
fake_prisma.db = RoutingPrismaWrapper(
writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False),
reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False),
)
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
await pc.get_credentials(prisma_client=fake_prisma)
assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"}
reader_inner.litellm_credentialstable.find_many.assert_not_awaited()
# ---------------------------------------------------------------------------
# ProxyConfig._add_general_settings_from_db_config
# ---------------------------------------------------------------------------

98
tests/test_litellm/proxy/spend_tracking/test_savings.py
View file

@ -1162,6 +1162,104 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one():
assert result.prompt_caching > at_public_rates.prompt_caching
@pytest.mark.parametrize(
"baseline_id, selected_id, selected_multiplier, billed_input, classifier_cost, expected",
[
("baseline", "selected", 0.1, None, 0.0, 0.0135),
("baseline", "selected", 2.0, None, 0.0, -0.015),
("baseline", "selected", 1.0, None, 0.0, 0.0),
("baseline", "selected", 0.1, 0.004, 0.001, 0.01),
("baseline", "baseline", 0.1, 0.004, 0.001, -0.001),
(None, "selected", 0.1, None, 0.0, 0.0),
("baseline", None, 0.1, None, 0.0, 0.0),
(None, None, 0.1, None, 0.0, 0.0),
("", "selected", 0.1, None, 0.0, 0.0),
("baseline", "", 0.1, None, 0.0, 0.0),
],
)
def test_autorouter_savings_distinguishes_priced_deployments(
baseline_id: str | None,
selected_id: str | None,
selected_multiplier: float,
billed_input: float | None,
classifier_cost: float,
expected: float,
) -> None:
router: Final = Router(
model_list=[
{
"model_name": name,
"litellm_params": {
"model": "anthropic/claude-opus-5",
"api_key": "test-key",
"input_cost_per_token": 1e-5 * multiplier,
"output_cost_per_token": 5e-5 * multiplier,
},
"model_info": {"id": name},
}
for name, multiplier in (("baseline", 1.0), ("selected", selected_multiplier))
]
)
result: Final = compute_savings_spend(
model="claude-opus-5",
custom_llm_provider="anthropic",
compression_saved_tokens=0,
gateway_injected_cache=False,
model_id=selected_id,
llm_router=lambda: router,
routing_decision={
"savings_baseline_model": "anthropic/claude-opus-5",
"savings_baseline_deployment_id": baseline_id,
"conversation_continuing": False,
"classifier_cost": classifier_cost,
},
usage_object={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100},
cost_breakdown=None if billed_input is None else {"input_cost": billed_input, "output_cost": 0.0},
)
assert result.autorouter == pytest.approx(expected)
@pytest.mark.parametrize("selected_model", ["azure/contract-deployment", "contract-deployment"])
def test_autorouter_savings_recognizes_one_deployment_under_its_base_model(selected_model: str) -> None:
router: Final = Router(
model_list=[
{
"model_name": "contract",
"litellm_params": {
"model": "azure/contract-deployment",
"api_key": "test-key",
"api_base": "https://example.openai.azure.com",
"input_cost_per_token": 0.0001,
"output_cost_per_token": 0.0002,
"cache_read_input_token_cost": 0.00001,
},
"model_info": {"id": "contract", "base_model": "azure/gpt-5.5"},
}
]
)
result: Final = compute_savings_spend(
model=selected_model,
custom_llm_provider="azure",
compression_saved_tokens=0,
gateway_injected_cache=False,
model_id="contract",
llm_router=lambda: router,
routing_decision={
"savings_baseline_model": "azure/gpt-5.5",
"savings_baseline_deployment_id": "contract",
"conversation_continuing": True,
},
usage_object={
"prompt_tokens": 21000,
"completion_tokens": 100,
"total_tokens": 21100,
"prompt_tokens_details": {"text_tokens": 1000, "cached_tokens": 0, "cache_creation_tokens": 20000},
},
cost_breakdown={"input_cost": 2.1, "output_cost": 0.02},
)
assert result.autorouter == 0.0
def test_a_recorded_baseline_deployment_prices_at_its_configured_rate():
"""A hardest-tier deployment with a negotiated rate is what the traffic would
really have cost; pricing its model publicly misstates the saving."""

8
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -7272,7 +7272,12 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth:
)
_PLANTED_STAMPS = {"attempted_fallbacks": 99, "original_model_group": "spoofed-group", "client_key": "client_value"}
_PLANTED_STAMPS = {
"attempted_fallbacks": 99,
"original_model_group": "spoofed-group",
"_client_output_ceiling": {"api_base": "https://attacker.example"},
"client_key": "client_value",
}
@pytest.mark.asyncio
@ -7301,6 +7306,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
assert "litellm_metadata" not in updated
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
assert "_client_output_ceiling" not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"

66
tests/test_litellm/proxy/test_prometheus_cleanup.py
View file

@ -6,13 +6,77 @@ ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir.
from __future__ import annotations
import os
import subprocess
import sys
from pathlib import Path
from typing import Final
from unittest.mock import patch
import pytest
from prometheus_client import CollectorRegistry, multiprocess
from litellm.proxy.prometheus_cleanup import mark_worker_exit, wipe_directory
from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit, wipe_directory
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
_WORKER: Final = """
import sys, time
from prometheus_client import Gauge
Gauge("litellm_in_flight", "", multiprocess_mode="livesum").set(float(sys.argv[1]))
print("ready", flush=True)
if sys.argv[2] == "stay":
time.sleep(120)
"""
def _spawn_worker(directory: Path, in_flight: str, lifetime: str) -> subprocess.Popen[str]:
env = {**os.environ, "PROMETHEUS_MULTIPROC_DIR": str(directory)}
worker = subprocess.Popen(
[sys.executable, "-c", _WORKER, in_flight, lifetime], env=env, stdout=subprocess.PIPE, text=True
)
assert worker.stdout is not None and worker.stdout.readline() == "ready\n"
return worker
def _livesum(directory: Path) -> float:
registry = CollectorRegistry()
multiprocess.MultiProcessCollector(registry, path=str(directory))
value = registry.get_sample_value("litellm_in_flight")
return 0.0 if value is None else value
class TestMarkDeadWorkers:
def test_drops_live_gauges_of_exited_workers_and_keeps_running_ones(self, tmp_path: Path) -> None:
"""A worker that died mid-request leaves its livesum file behind; the replacement worker's startup prune
must remove exactly that file so the aggregate stops counting requests nobody is serving."""
dead = _spawn_worker(tmp_path, "3", "exit")
assert dead.wait(timeout=30) == 0
alive = _spawn_worker(tmp_path, "2", "stay")
try:
assert (tmp_path / f"gauge_livesum_{dead.pid}.db").exists()
assert _livesum(tmp_path) == 5.0
assert mark_dead_workers(str(tmp_path)) == (dead.pid,)
assert not (tmp_path / f"gauge_livesum_{dead.pid}.db").exists()
assert (tmp_path / f"gauge_livesum_{alive.pid}.db").exists()
assert _livesum(tmp_path) == 2.0
assert mark_dead_workers(str(tmp_path)) == ()
finally:
alive.kill()
alive.wait(timeout=30)
def test_leaves_counters_of_exited_workers_alone(self, tmp_path: Path) -> None:
(tmp_path / "counter_424242.db").touch()
(tmp_path / "histogram_424242.db").touch()
assert mark_dead_workers(str(tmp_path)) == ()
assert sorted(p.name for p in tmp_path.glob("*.db")) == ["counter_424242.db", "histogram_424242.db"]
def test_keeps_live_gauges_of_workers_it_may_not_signal(self, tmp_path: Path) -> None:
"""Signal 0 to pid 1 raises PermissionError for an unprivileged proxy; that pid is alive, not dead."""
(tmp_path / "gauge_livesum_1.db").touch()
assert mark_dead_workers(str(tmp_path)) == ()
assert (tmp_path / "gauge_livesum_1.db").exists()
class TestWipeDirectory:
def test_deletes_all_db_files(self, tmp_path):

9
tests/test_litellm/proxy/test_prometheus_metrics_server.py
View file

@ -1,5 +1,5 @@
"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose only /metrics, and follow its
parent's lifetime.
"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose /metrics plus a probe-friendly
/health, and follow its parent's lifetime.
Everything here runs on loopback against a child of this test process; no LLM keys or external network.
"""
@ -91,7 +91,10 @@ def test_metrics_app_aggregates_multiproc_dir_and_reports_pid(tmp_path: Path, mo
assert metrics.headers[PID_HEADER] == str(os.getpid())
assert 'litellm_requests_metric_total{model="gpt-5"} 5.0' in metrics.text
assert client.get("/health").status_code == 404
health: Final = client.get("/health")
assert health.status_code == 200
assert health.json() == {"status": "healthy", "multiproc_dir": str(tmp_path)}
assert client.get("/docs").status_code == 404
empty: Final = TestClient(build_metrics_app(str(other_dir))).get("/metrics")
assert empty.status_code == 200

10
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri
assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == []
def test_deployment_pre_call_target_stays_native_when_opted_out():
def test_deployment_hook_target_stays_native_when_opted_out():
"""Model-level guardrails resolve their target here rather than through ProxyLogging."""
assert _KeepsNativeHooks()._deployment_pre_call_target() is not None
assert _KeepsNativeHooks()._deployment_hook_target() is not None
opted_out = _KeepsNativeHooks()
assert opted_out._deployment_pre_call_target() is opted_out
assert _AppliesGuardrail()._deployment_pre_call_target() is not None
assert opted_out._deployment_hook_target() is opted_out
assert _AppliesGuardrail()._deployment_hook_target() is not None
routed = _AppliesGuardrail()
assert routed._deployment_pre_call_target() is not routed
assert routed._deployment_hook_target() is not routed
@pytest.mark.asyncio

109
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -822,16 +822,16 @@ def _mock_scheduled_proxy_config() -> MagicMock:
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_credentials(monkeypatch):
"""
Test that get_credentials is only called when store_model_in_db is True
"""
async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_deployment(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
# Mock dependencies
mock_prisma_client = MagicMock()
mock_proxy_logging = MagicMock(spec=ProxyLogging)
mock_proxy_logging.slack_alerting_instance = MagicMock()
@ -841,25 +841,6 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
with (
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
patch("litellm.proxy.proxy_server.store_model_in_db", False),
): # set store_model_in_db to False
# Test when store_model_in_db is False
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
# Verify get_credentials was not called
mock_proxy_config.get_credentials.assert_not_called()
# Now test with store_model_in_db = True
with (
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True),
):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
@ -870,12 +851,31 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch):
proxy_logging_obj=mock_proxy_logging,
)
# Verify get_credentials was called both directly and scheduled
assert mock_proxy_config.get_credentials.call_count == 1 # Direct call
mock_proxy_config.get_credentials.assert_not_called()
mock_proxy_config.add_deployment.assert_not_called()
# Verify a scheduled job was added for get_credentials
mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls]
assert len(mock_scheduler_calls) > 0
scheduler = AsyncIOScheduler()
try:
with (
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
assert scheduler.get_job("get_credentials_job") is None
assert scheduler.get_job("add_deployment_job") is not None
mock_proxy_config.get_credentials.assert_not_called()
assert mock_proxy_config.add_deployment.call_count == 1
finally:
scheduler.shutdown(wait=False)
@pytest.mark.asyncio
@ -924,7 +924,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch):
"""
The DB config-reload jobs (add_deployment, get_credentials) that keep multi-pod
The DB config-reload job (add_deployment) that keeps multi-pod
deployments in sync must be scheduled at the configured
proxy_config_reload_interval_seconds, not a hardcoded value.
"""
@ -967,7 +967,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(
if "id" in job_call.kwargs
}
assert scheduled_seconds["add_deployment_job"] == configured_interval
assert scheduled_seconds["get_credentials_job"] == configured_interval
assert "get_credentials_job" not in scheduled_seconds
@pytest.mark.asyncio
@ -1011,7 +1011,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte
if "id" in job_call.kwargs
}
assert scheduled_seconds["add_deployment_job"] == 30
assert scheduled_seconds["get_credentials_job"] == 30
assert "get_credentials_job" not in scheduled_seconds
@pytest.mark.asyncio
@ -3166,6 +3166,47 @@ async def test_custom_ui_sso_sign_in_handler_config_loading():
os.unlink(config_file_path)
@pytest.mark.asyncio
async def test_startup_initializes_string_callbacks_after_all_litellm_settings_load(tmp_path, monkeypatch):
from litellm.integrations.s3_v2 import S3Logger
from litellm.litellm_core_utils import litellm_logging
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import ProxyLogging
config_file = tmp_path / "config.yaml"
config_file.write_text(
"model_list: []\n"
"litellm_settings:\n"
" success_callback:\n"
" - s3_v2\n"
" failure_callback:\n"
" - s3_v2\n"
" s3_callback_params:\n"
" s3_bucket_name: ordering-regression-bucket\n"
" s3_region_name: us-west-2\n"
)
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "s3_callback_params", None)
monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [])
await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file))
ProxyLogging(user_api_key_cache=MagicMock())._init_litellm_callbacks(llm_router=None)
success_loggers = [cb for cb in litellm._async_success_callback if isinstance(cb, S3Logger)]
failure_loggers = [cb for cb in litellm._async_failure_callback if isinstance(cb, S3Logger)]
assert len(success_loggers) == 1
assert len(failure_loggers) == 1
assert success_loggers[0].s3_bucket_name == "ordering-regression-bucket"
assert success_loggers[0].s3_region_name == "us-west-2"
assert "s3_v2" not in litellm.success_callback
assert "s3_v2" not in litellm.failure_callback
@pytest.mark.asyncio
async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch):
"""
@ -7446,10 +7487,8 @@ async def test_store_model_in_db_db_override_when_config_false():
# store_model_in_db should now be True (overridden by DB)
assert ps.store_model_in_db is True
# add_deployment and get_credentials should have been called
# since store_model_in_db is now True
assert mock_proxy_config.add_deployment.call_count == 1
assert mock_proxy_config.get_credentials.call_count == 1
mock_proxy_config.get_credentials.assert_not_called()
@pytest.mark.asyncio

540
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -16,6 +16,7 @@ from pydantic import ValidationError
import litellm
from litellm import Router
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.auto_router_model_naming import (
CUSTOMIZATION_CAPABILITY,
GATED_AUTO_ROUTER_CAPABILITIES,
@ -24,7 +25,12 @@ from litellm.router_utils.auto_router_model_naming import (
)
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.constants import (
OUTPUT_TOKEN_CEILING_PARAMS,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.router import as_output_cap
from litellm.router_strategy.complexity_router.complexity_router import (
_CLASSIFICATION_CURRENT_MESSAGE_ONLY,
_CLASSIFICATION_WITH_CONVERSATION,
@ -826,6 +832,8 @@ class TestCustomDimensions:
pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"),
pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"),
pytest.param({"unknown": True}, {}, id="extra-field"),
pytest.param({"scoring_mode": "graded"}, {}, id="unknown-scoring-mode"),
pytest.param({"scoring_mode": None}, {}, id="null-scoring-mode"),
],
)
def test_custom_dimension_invalid_configuration_rejected(
@ -878,43 +886,125 @@ class TestCustomDimensions:
)
@pytest.mark.asyncio
@pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh"))
@pytest.mark.parametrize("scoring_mode", ("binary", "match_count"))
@pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh", "orbitmesh fluxgate"))
async def test_custom_dimensions_public_hook_scores_only_current_ask(
self, mock_router_instance: MagicMock, current_ask: str
self, mock_router_instance: MagicMock, current_ask: str, scoring_mode: str
) -> None:
router: Final = ComplexityRouter(
"test-router",
mock_router_instance,
{
"tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"},
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
"dimension_weights": {},
"custom_dimensions": [
{
"name": "internalFrameworks",
"weight": 0.8,
"keywords": ["orbitmesh", "fluxgate"],
"scoring_mode": scoring_mode,
}
],
},
)
result: Final = await router.async_pre_routing_hook(
model="test-router",
request_kwargs={},
messages=[
{"role": "system", "content": "orbitmesh"},
{"role": "user", "content": "orbitmesh"},
{"role": "assistant", "content": "orbitmesh is ready"},
{"role": "system", "content": "orbitmesh fluxgate"},
{"role": "user", "content": "orbitmesh fluxgate"},
{"role": "assistant", "content": "orbitmesh fluxgate is ready"},
{"role": "user", "content": current_ask},
{"role": "tool", "tool_call_id": "previous", "content": "orbitmesh"},
{"role": "tool", "tool_call_id": "previous", "content": "orbitmesh fluxgate"},
],
)
assert result is not None
assert result.routing_decision is not None
assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (current_ask == "orbitmesh")
assert result.model == ("top" if current_ask == "orbitmesh" else "cheap")
expected_score: Final = (
0.0
if current_ask == "Hello!"
else 0.4
if scoring_mode == "match_count" and current_ask == "orbitmesh"
else 0.8
)
assert result.routing_decision["score"] == expected_score
assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (expected_score > 0)
assert result.model == ("cheap" if expected_score == 0 else "strong" if expected_score == 0.4 else "top")
assert "orbitmesh" not in " ".join(result.routing_decision["signals"])
def test_custom_patterns_scan_only_the_first_2048_characters(self, mock_router_instance: MagicMock) -> None:
@pytest.mark.parametrize("scoring_mode", ("binary", "match_count"))
def test_custom_patterns_scan_only_the_first_2048_characters(
self, mock_router_instance: MagicMock, scoring_mode: str
) -> None:
router: Final = ComplexityRouter(
"test-router",
mock_router_instance,
{"custom_dimensions": [{"name": "late", "weight": 0.7, "patterns": [r"zzz{1,3}"]}]},
{
"custom_dimensions": [
{
"name": "late",
"weight": 0.7,
"patterns": [r"zzz{1,3}", r"yyy{1,3}"],
"scoring_mode": scoring_mode,
}
]
},
)
baseline: Final = ComplexityRouter("test-router", mock_router_instance)
assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2]
assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2]
second_hit_past_the_bound: Final = "yyy " + "a" * 2044 + " zzz"
contribution: Final = (
router.classify(second_hit_past_the_bound)[1] - baseline.classify(second_hit_past_the_bound)[1]
)
assert contribution == pytest.approx(0.7 if scoring_mode == "binary" else 0.35)
@pytest.mark.parametrize(
"prompt,expected_score",
[
pytest.param("Hello!", 0.0, id="no-hit"),
pytest.param("orbitmesh orbitmesh ORBITMESH again", 0.5, id="one-keyword-repeated"),
pytest.param("create table a; CREATE TABLE b; create table c", 0.5, id="one-pattern-repeated"),
pytest.param("orbitmesh and fluxgate", 1.0, id="two-keywords"),
pytest.param("orbitmesh then create table t", 1.0, id="keyword-plus-pattern"),
pytest.param("create table a; alter table b", 1.0, id="two-patterns"),
pytest.param("orbitmesh fluxgate create table a alter table b", 1.0, id="all-matchers"),
],
)
def test_match_count_grades_distinct_matchers(
self, mock_router_instance: MagicMock, prompt: str, expected_score: float
) -> None:
dimension: Final = {
"name": "graded",
"weight": 0.6,
"keywords": ["orbitmesh", "ORBITMESH", "fluxgate"],
"patterns": [r"\bcreate\s{1,4}table\b", r"\bcreate\s{1,4}table\b", r"\balter\s{1,4}table\b"],
}
baseline: Final = ComplexityRouter("test-router", mock_router_instance)
binary: Final = ComplexityRouter("test-router", mock_router_instance, {"custom_dimensions": [dimension]})
graded: Final = ComplexityRouter(
"test-router",
mock_router_instance,
{"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]},
)
_, baseline_score, baseline_signals = baseline.classify(prompt)
_, binary_score, binary_signals = binary.classify(prompt)
_, graded_score, graded_signals = graded.classify(prompt)
assert graded_score == pytest.approx(baseline_score + 0.6 * expected_score)
assert binary_score == pytest.approx(baseline_score + (0.6 if expected_score else 0.0))
expected_signals: Final = [*baseline_signals, *(["custom (graded)"] if expected_score else [])]
assert graded_signals == expected_signals
assert binary_signals == expected_signals
def test_scoring_mode_round_trips_and_defaults_to_binary(self) -> None:
dimension: Final = {"name": "graded", "weight": 0.6, "keywords": ["orbitmesh"]}
legacy: Final = ComplexityRouterConfig.model_validate({"custom_dimensions": [dimension]})
graded: Final = ComplexityRouterConfig.model_validate(
{"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]}
)
assert legacy.custom_dimensions[0].scoring_mode == "binary"
assert graded.model_dump(mode="json")["custom_dimensions"][0]["scoring_mode"] == "match_count"
assert ComplexityRouterConfig.model_validate(graded.model_dump(mode="json")) == graded
def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None:
heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]}
@ -1511,6 +1601,7 @@ class TestRouterComplexityDeploymentMethods:
def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None:
"""Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no
prompt at all, leaves a router unmetered, so several of them register under a ceiling of one."""
def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]:
llm_config: dict[str, object] = {"model": "gpt-4o-mini"}
if preset is not None:
@ -1647,6 +1738,7 @@ class TestRouterComplexityDeploymentMethods:
def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None:
"""tier_labels renames the built-in ladder without defining one, so it stays ungated: two such
routers register under a ceiling of one."""
def labeled(model_name: str, model_id: str) -> dict[str, object]:
row = self._router_row(model_name, model_id, "heuristic")
row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"}
@ -2520,9 +2612,7 @@ class TestLLMClassifier:
assert outcome.classifier_cost == pytest.approx(1.35e-05)
@pytest.mark.asyncio
async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(
self, llm_classifier_config
):
async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(self, llm_classifier_config):
real_router = Router(
model_list=[
{
@ -2565,9 +2655,7 @@ class TestLLMClassifier:
assert real_router.total_calls["openai/mock-backup-classifier"] == 0
@pytest.mark.asyncio
async def test_aclassify_enforces_total_classifier_deadline(
self, mock_router_instance, llm_classifier_config
):
async def test_aclassify_enforces_total_classifier_deadline(self, mock_router_instance, llm_classifier_config):
cancelled = asyncio.Event()
async def slow_classifier(**_kwargs: object) -> None:
@ -3333,11 +3421,11 @@ class TestRouterPreRoutingAliasOverrides:
def test_drop_client_effort_carriers_helper_edge_shapes(self):
no_pin: Dict = {"thinking": {"type": "adaptive"}}
Router._drop_client_effort_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1})
Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1})
assert no_pin == {"thinking": {"type": "adaptive"}}
non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3}
Router._drop_client_effort_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"})
Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"})
assert non_dict_carriers == {"output_config": "max", "reasoning": 3}
effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}}
@ -3791,11 +3879,11 @@ class TestRouterPreRoutingSharedAliasName:
}
@staticmethod
async def _routed_call_kwargs(router: Router, **request_params) -> dict:
async def _routed_call_kwargs(router: Router, prompt: str = "hi", **request_params) -> dict:
mock_acompletion = AsyncMock(return_value=litellm.ModelResponse(choices=[{"message": {"content": "hi"}}]))
with patch.object(litellm, "acompletion", mock_acompletion):
await router.acompletion(
model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_params
model="smart-router", messages=[{"role": "user", "content": prompt}], **request_params
)
return mock_acompletion.call_args.kwargs
@ -12414,9 +12502,7 @@ class TestTierHealthFailover:
llm_provider="",
)
filtered = (*cooling, *blocked, *excluded)
healthy = [
{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered
]
healthy = [{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered]
if not healthy:
raise RouterRateLimitError(
model=model, cooldown_time=60.0, enable_pre_call_checks=False, cooldown_list=[]
@ -12845,9 +12931,7 @@ class TestTierHealthFailover:
assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs)
@pytest.mark.asyncio
async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(
self, mock_router_instance
):
async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(self, mock_router_instance):
"""RPM exhaustion is its own verdict from the owner (RouterRateLimitErrorBasic). A peer
in that state would be rejected downstream, so it cannot be the substitute."""
from litellm.types.router import RouterRateLimitErrorBasic
@ -12880,9 +12964,7 @@ class TestTierHealthFailover:
assert {r.model for r in results} == {"live-c"}
@pytest.mark.asyncio
async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(
self, mock_router_instance
):
async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(self, mock_router_instance):
"""The Responses API carries its prompt as `input`, never as messages. The owner only
runs its context-window pre-call check when one of them is present, so dropping `input`
would silently skip window filtering on that whole surface."""
@ -12908,9 +12990,7 @@ class TestTierHealthFailover:
), "the eligibility probe must forward `input` to the owner"
@pytest.mark.asyncio
async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(
self, mock_router_instance
):
async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance):
"""The owner answers an unconfigured group with BadRequestError. Reading that as live
would both skip failover off it and let it be chosen as a substitute."""
router = self._router(
@ -13099,9 +13179,7 @@ class TestClassifierVision:
routed as default_fallback on text the request never contained.
"""
router = self._router(mock_router_instance, vision={"enabled": True})
response = await router.async_pre_routing_hook(
model="m", request_kwargs={}, messages=self._turn(IMG_PART)
)
response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART))
assert response.routing_decision["cause"] == "llm_classifier"
assert response.model == "t-complex"
assert [block["type"] for block in self._classifier_user_content(mock_router_instance)] == [
@ -13112,9 +13190,7 @@ class TestClassifierVision:
@pytest.mark.asyncio
async def test_image_only_turn_still_falls_back_when_vision_is_off(self, mock_router_instance):
router = self._router(mock_router_instance, vision={"enabled": False})
response = await router.async_pre_routing_hook(
model="m", request_kwargs={}, messages=self._turn(IMG_PART)
)
response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART))
assert response.routing_decision["cause"] == "default_fallback"
mock_router_instance.acompletion.assert_not_awaited()
@ -13184,9 +13260,7 @@ class TestClassifierVision:
makes the image the only variable; a margin loose enough to leave the score undecided
would pass whether or not the guard exists.
"""
router = self._router(
mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra
)
router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra)
response = await router.async_pre_routing_hook(
model="m", request_kwargs={}, messages=self._turn({"type": "text", "text": "what is this"}, IMG_PART)
)
@ -13200,9 +13274,7 @@ class TestClassifierVision:
self, mock_router_instance, classifier_type, extra, short_circuit_cause
):
"""The negative class: same router, same text, no image, and the scorer still decides."""
router = self._router(
mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra
)
router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra)
response = await router.async_pre_routing_hook(
model="m", request_kwargs={}, messages=[{"role": "user", "content": "what is this"}]
)
@ -13212,3 +13284,379 @@ class TestClassifierVision:
def test_max_images_must_be_positive(self):
with pytest.raises(ValidationError):
ClassifierLLMConfig(model="clf", vision={"enabled": True, "max_images": 0})
class TestMaxTokensFromTierModel:
"""The auto-router replaces the caller's output ceiling with the tier model's own, so one
client-side value no longer starves a bigger tier or gets rejected by a smaller one."""
COMPLEX_PROMPT: Final = (
"Design a distributed rate limiter with Redis, sharding and failover. Analyze the consistency "
"tradeoffs and implement the algorithm step by step with tests."
)
SMALL: Final = {
"model_name": "small",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"},
"model_info": {"max_output_tokens": 8192},
}
@staticmethod
def _router(
tier_litellm_params: dict | None = None,
max_tokens_from_tier_model: bool | None = None,
simple_deployments: list[dict] | None = None,
extra_config: dict | None = None,
) -> Router:
simple_tier: dict = {"model_name": "small"}
if tier_litellm_params:
simple_tier["litellm_params"] = tier_litellm_params
config: dict = {
"tiers": {"SIMPLE": simple_tier, "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"},
**(extra_config or {}),
}
if max_tokens_from_tier_model is not None:
config["max_tokens_from_tier_model"] = max_tokens_from_tier_model
return Router(
model_list=[
{
"model_name": "smart-router",
"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config},
},
*(simple_deployments or [TestMaxTokensFromTierModel.SMALL]),
{
"model_name": "big",
"litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "k"},
"model_info": {"max_output_tokens": 64000},
},
]
)
@staticmethod
async def _routed(router: Router, prompt: str = "hi", **request_kwargs) -> dict:
"""Drive the real routing entry point and return the request kwargs it leaves behind."""
deployment = await router.async_get_available_deployment(
model="smart-router", request_kwargs=request_kwargs, messages=[{"role": "user", "content": prompt}]
)
return {"model": deployment["litellm_params"]["model"], **request_kwargs}
@staticmethod
async def _routed_responses(router: Router, prompt: str = "hi", **request_kwargs) -> dict:
"""The Responses surface hands the router `input` both as the prompt argument and inside the
request kwargs, so the hook sees the same shape the real call carries."""
routed: dict = {"input": prompt, **request_kwargs}
deployment = await router.async_get_available_deployment(
model="smart-router", request_kwargs=routed, input=prompt
)
return {"model": deployment["litellm_params"]["model"], **routed}
@pytest.mark.asyncio
async def test_client_ceiling_is_replaced_by_the_routed_tier_models_ceiling(self):
router = self._router()
simple = await self._routed(router, max_tokens=8192)
complex_ = await self._routed(router, self.COMPLEX_PROMPT, max_tokens=8192)
assert (simple["model"], simple["max_tokens"]) == ("anthropic/claude-haiku-4-5", 8192)
assert (complex_["model"], complex_["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
assert "max_output_tokens" not in complex_
@pytest.mark.asyncio
async def test_every_client_carrier_of_the_ceiling_is_replaced(self):
sent = await self._routed(self._router(), self.COMPLEX_PROMPT, max_completion_tokens=8192)
assert sent["max_tokens"] == 64000
assert "max_completion_tokens" not in sent
@pytest.mark.asyncio
async def test_responses_surface_gets_the_ceiling_under_its_own_name(self):
sent = await self._routed_responses(self._router(), self.COMPLEX_PROMPT, max_output_tokens=8192)
assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
assert "max_tokens" not in sent
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tier_params, responses_call",
[
({"max_tokens": 4321}, False),
({"max_tokens": 4321}, True),
({"max_completion_tokens": 4321}, False),
({"max_completion_tokens": 4321}, True),
({"max_output_tokens": 4321}, False),
],
)
async def test_operators_own_tier_ceiling_wins_under_the_surface_name(self, tier_params, responses_call):
router = self._router(tier_litellm_params=tier_params)
if responses_call:
sent = await self._routed_responses(router, max_output_tokens=8192)
else:
sent = await self._routed(router, max_tokens=8192)
surface_key = "max_output_tokens" if responses_call else "max_tokens"
assert sent[surface_key] == 4321
assert not (OUTPUT_TOKEN_CEILING_PARAMS - {surface_key}) & sent.keys()
@pytest.mark.asyncio
async def test_opting_out_forwards_the_client_value_unchanged(self):
sent = await self._routed(self._router(max_tokens_from_tier_model=False), self.COMPLEX_PROMPT, max_tokens=8192)
assert sent["max_tokens"] == 8192
@pytest.mark.asyncio
async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self):
unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}}
sent = await self._routed(self._router(simple_deployments=[self.SMALL, unmapped]), max_tokens=4000)
assert sent["max_tokens"] == 4000
@pytest.mark.asyncio
async def test_a_multi_deployment_tier_model_uses_its_smallest_ceiling(self):
smaller: dict = {
**self.SMALL,
"litellm_params": {**self.SMALL["litellm_params"], "api_key": "k2"},
"model_info": {"max_output_tokens": 4096},
}
sent = await self._routed(self._router(simple_deployments=[self.SMALL, smaller]), max_tokens=100000)
assert sent["max_tokens"] == 4096
@pytest.mark.asyncio
async def test_ceiling_falls_back_to_the_cost_map(self, monkeypatch):
monkeypatch.setitem(
litellm.model_cost,
"auto-cap-probe-model",
{"litellm_provider": "openai", "mode": "chat", "max_output_tokens": 4242, "max_input_tokens": 100000},
)
mapped_only: dict = {
"model_name": "small",
"litellm_params": {"model": "openai/auto-cap-probe-model", "api_key": "k"},
}
sent = await self._routed(self._router(simple_deployments=[mapped_only]), max_tokens=8192)
assert sent["max_tokens"] == 4242
@pytest.mark.asyncio
@pytest.mark.parametrize("client_kwargs", [{}, {"max_tokens": 0}], ids=["omitted", "zero"])
async def test_omitted_and_zero_are_replaced_like_any_other_value(self, client_kwargs):
sent = await self._routed(self._router(), self.COMPLEX_PROMPT, **client_kwargs)
assert sent["max_tokens"] == 64000
@pytest.mark.parametrize(
"tier_params, responses_call, expected",
[
({"max_tokens": 1, "temperature": 0.2}, False, {"max_tokens": 1, "temperature": 0.2}),
({"max_tokens": 1}, True, {"max_output_tokens": 1}),
({"max_completion_tokens": 2}, False, {"max_tokens": 2}),
({"max_completion_tokens": 2}, True, {"max_output_tokens": 2}),
({"max_output_tokens": 3}, False, {"max_tokens": 3}),
({"max_output_tokens": 3}, True, {"max_output_tokens": 3}),
({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 1}),
({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, True, {"max_output_tokens": 3}),
({"max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 2}),
({"reasoning_effort": "low"}, True, {"reasoning_effort": "low"}),
],
)
def test_every_tier_alias_collapses_onto_the_surface_key(self, tier_params, responses_call, expected):
assert dict(Router._tier_ceiling_under_the_surface_name(tier_params, responses_call=responses_call)) == expected
@pytest.mark.asyncio
async def test_the_default_fallback_exit_carries_the_ceiling(self):
routed: dict = {"max_tokens": 8192}
deployment = await self._router().async_get_available_deployment(
model="smart-router", request_kwargs=routed, messages=[{"role": "system", "content": "be nice"}]
)
assert routed["metadata"]["routing_decision"]["cause"] == "default_fallback"
assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
@pytest.mark.asyncio
async def test_the_plan_mode_exit_carries_the_ceiling(self):
routed: dict = {"max_tokens": 8192}
deployment = await self._router(
extra_config={"plan_mode_min_tier": "REASONING"}
).async_get_available_deployment(
model="smart-router",
request_kwargs=routed,
messages=[
{"role": "user", "content": "plan the refactor"},
{"role": "system", "content": "Plan mode is active"},
],
)
assert routed["metadata"]["routing_decision"]["cause"] == "plan_mode"
assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
@pytest.mark.asyncio
async def test_a_default_model_landing_with_no_tier_still_gets_its_ceiling(self):
strategy = ComplexityRouter(
model_name="smart-router",
litellm_router_instance=self._router(),
complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"},
)
assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000}
@pytest.mark.asyncio
async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self):
"""A model-group fallback re-enters routing with the same kwargs; a Sonnet-sized ceiling
must not ride onto the plain group the caller configured as the fallback."""
big: dict = {
"model_name": "big",
"litellm_params": {
"model": "anthropic/claude-sonnet-5",
"api_key": "k",
"mock_response": "litellm.InternalServerError",
},
"model_info": {"max_output_tokens": 64000},
}
plain: dict = {
"model_name": "plain",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k", "mock_response": "ok"},
}
router = Router(
model_list=[
{
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"tiers": {"SIMPLE": "big", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"}
},
},
},
big,
plain,
],
fallbacks=[{"smart-router": ["plain"]}],
num_retries=0,
)
recorder = _OutputCeilingRecorder()
litellm.callbacks.append(recorder)
try:
await router.acompletion(
model="smart-router", messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], max_tokens=8192
)
finally:
litellm.callbacks.remove(recorder)
assert recorder.seen == [("claude-sonnet-5", 64000), ("claude-haiku-4-5", 8192)]
@pytest.mark.asyncio
async def test_a_caller_seeded_stamp_cannot_inject_kwargs_on_a_plain_group(self):
"""The stamp sits in a metadata bucket a caller can write; a planted one must yield
nothing but integer ceiling carriers, never a redirected api_base or credential."""
planted: dict = {
"api_base": "https://attacker.example",
"api_key": "stolen",
"max_tokens": "not-an-int",
"max_completion_tokens": True,
"max_output_tokens": 321,
}
routed: dict = {"max_tokens": 8192, "metadata": {"_client_output_ceiling": planted}}
await self._router().async_get_available_deployment(
model="big", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}]
)
assert {k: v for k, v in routed.items() if k not in ("metadata", "model_info")} == {"max_output_tokens": 321}
@pytest.mark.asyncio
async def test_the_pass_through_routing_entry_point_pins_and_restores_the_same_way(self):
pass_through: dict = {**self.SMALL["litellm_params"], "use_in_pass_through": True}
small: dict = {**self.SMALL, "litellm_params": pass_through}
plain: dict = {**small, "model_name": "plain"}
router = self._router(simple_deployments=[small, plain])
for deployment in router.model_list:
deployment["litellm_params"]["use_in_pass_through"] = True
routed: dict = {"max_tokens": 8192}
deployment = await router.async_get_available_deployment_for_pass_through(
model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": self.COMPLEX_PROMPT}]
)
pinned = routed["max_tokens"]
await router.async_get_available_deployment_for_pass_through(
model="plain", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}]
)
assert (deployment["litellm_params"]["model"], pinned, routed["max_tokens"]) == (
"anthropic/claude-sonnet-5",
64000,
8192,
)
@pytest.mark.asyncio
async def test_the_classifier_fallback_exit_carries_the_ceiling(self):
router = self._router(
extra_config={
"classifier_type": "llm",
"classifier_llm_config": {"model": "no-such-classifier", "timeout_ms": 400},
"classifier_fallback": "default_model",
"default_model": "big",
}
)
routed: dict = {"max_tokens": 8192}
deployment = await router.async_get_available_deployment(
model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}]
)
assert routed["metadata"]["routing_decision"]["cause"] == "default_model_fallback"
assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
@pytest.mark.parametrize(
"value, expected",
[(8192, 8192), ("8192", 8192), (100.9, 100), (0, 0), (-1, None), (True, None), ("x", None), (None, None)],
)
def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected):
assert as_output_cap(value) == expected
def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self):
stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}}
Router._restore_client_ceiling_no_tier_pins(stamped)
assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192}
coerced: dict = {
"max_tokens": 64000,
"metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}},
}
Router._restore_client_ceiling_no_tier_pins(coerced)
assert {k: v for k, v in coerced.items() if k != "metadata"} == {
"max_tokens": 8192,
"max_completion_tokens": 100,
}
unstamped: dict = {"max_tokens": 64000, "metadata": {}}
Router._restore_client_ceiling_no_tier_pins(unstamped)
assert unstamped["max_tokens"] == 64000
@pytest.mark.asyncio
async def test_pinning_stamps_the_callers_carriers_once(self):
router = self._router()
request_kwargs: dict = {"max_completion_tokens": 8192}
first = router._pin_tier_params_onto_request(
model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False
)
second = router._pin_tier_params_onto_request(
model="big", tier_litellm_params={"max_tokens": 32000}, request_kwargs=request_kwargs, responses_call=False
)
none = router._pin_tier_params_onto_request(
model="big", tier_litellm_params=None, request_kwargs=request_kwargs, responses_call=False
)
assert (first, second, none) == (True, True, False)
assert request_kwargs["max_tokens"] == 32000
assert request_kwargs["metadata"]["_client_output_ceiling"] == {"max_completion_tokens": 8192}
class _OutputCeilingRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.seen: list[tuple[str, int | None]] = []
def log_pre_api_call(self, model, messages, kwargs):
self.seen.append((model, kwargs.get("optional_params", {}).get("max_tokens")))

187
tests/test_litellm/router_strategy/test_least_busy.py Normal file
View file

@ -0,0 +1,187 @@
from typing import Final
import pytest
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler
GROUP: Final = "least-busy-group"
DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}}
DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}}
HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B]
def _call_kwargs(deployment_id: str) -> dict[str, object]:
return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}}
class SharedRedisCounters:
"""Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when
the key is created, and ordered reads that raise rather than invent a value."""
def __init__(self) -> None:
self.counts: dict[str, int] = {}
self.ttls: dict[str, int] = {}
def count(self, key: str) -> int | None:
return self.counts.get(key)
def expire(self, key: str) -> None:
self.counts.pop(key, None)
self.ttls.pop(key, None)
def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
incremented: Final = max(0, self.counts.get(key, 0) + value)
self.counts[key] = incremented
self.ttls.setdefault(key, ttl)
return incremented
async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int:
return self.increment_with_floor(key, value, ttl)
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
return tuple(self.counts.get(key) for key in key_list)
async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
return self.batch_get_counts(key_list)
def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler:
cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double
return LeastBusyLoggingHandler(router_cache=cache)
@pytest.mark.asyncio
async def test_worker_routes_around_a_request_another_worker_started() -> None:
shared: Final = SharedRedisCounters()
streaming_worker: Final = _worker(shared)
picking_worker: Final = _worker(shared)
picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None)
streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None)
assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_sync_pick_reads_the_shared_counts() -> None:
shared: Final = SharedRedisCounters()
streaming_worker: Final = _worker(shared)
picking_worker: Final = _worker(shared)
picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None)
streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None)
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_the_handler_never_pushes_a_counters_ttl_forward() -> None:
"""A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that
stuck count an hour after the key was created, which only works while nothing writes the TTL
again: a handler that refreshed it on every touch would keep the count alive for as long as
the group takes traffic, and the deployment would read busier than it is forever."""
shared: Final = SharedRedisCounters()
worker: Final = _worker(shared)
key: Final = f"{GROUP}_request_count:dep-a"
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS}
shared.ttls[key] = 5
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
worker.log_success_event(_call_kwargs("dep-a"), None, None, None)
assert shared.count(key) == 1
assert shared.ttls == {key: 5}
@pytest.mark.asyncio
async def test_counts_stay_in_memory_without_redis() -> None:
worker: Final = _worker(None)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None)
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0
class UnavailableRedis(SharedRedisCounters):
def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
raise ConnectionError("redis is down")
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
raise ConnectionError("redis is down")
@pytest.mark.asyncio
async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None:
worker: Final = _worker(UnavailableRedis())
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None:
shared: Final = SharedRedisCounters()
worker: Final = _worker(shared)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
shared.expire(f"{GROUP}_request_count:dep-a")
worker.log_success_event(_call_kwargs("dep-a"), None, None, None)
assert shared.count(f"{GROUP}_request_count:dep-a") == 0
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert shared.count(f"{GROUP}_request_count:dep-a") == 1
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
@pytest.mark.asyncio
async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None:
worker: Final = _worker(None)
in_memory: Final = worker.router_cache.in_memory_cache
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
in_memory.delete_cache(f"{GROUP}_request_count:dep-a")
await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None)
assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
def test_calls_without_a_deployment_are_ignored() -> None:
shared: Final = SharedRedisCounters()
worker: Final = _worker(shared)
worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}})
worker.log_pre_api_call(model="m", messages=[], kwargs={})
assert shared.counts == {}

127
tests/test_litellm/router_strategy/test_lowest_latency.py
View file

@ -165,6 +165,133 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds():
json.dumps({"latency": latencies})
MODEL_GROUP = "gpt-4o-mini"
FAST_TTFT_ID = "fast-ttft-short-output"
SLOW_TTFT_ID = "slow-ttft-long-output"
STREAMING_DEPLOYMENTS = [
{"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}},
{"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}},
]
def _streaming_kwargs(deployment_id: str, start_time: datetime, ttft_seconds: float):
return {
"litellm_params": {
"metadata": {"model_group": MODEL_GROUP},
"model_info": {"id": deployment_id},
},
"stream": True,
"completion_start_time": start_time + timedelta(seconds=ttft_seconds),
}
def _recorded_ttft(cache: DualCache, deployment_id: str):
cached = cache.get_cache(key=f"{MODEL_GROUP}_map") or {}
return cached.get(deployment_id, {}).get("time_to_first_token_seconds", [])
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
async def test_streaming_ttft_ranking_ignores_completion_length(sync_mode: bool):
"""Deployment A: TTFT 1s, 50 completion tokens. Deployment B: TTFT 3s, 500
completion tokens. Dividing TTFT by completion tokens made B look faster
(3/500 = 0.006 beats 1/50 = 0.02); actual TTFT must win."""
cache = DualCache()
handler = LowestLatencyLoggingHandler(router_cache=cache)
start_time = datetime(2026, 1, 1, 12, 0, 0)
end_time = start_time + timedelta(seconds=10)
samples = (
(FAST_TTFT_ID, 1.0, 50),
(SLOW_TTFT_ID, 3.0, 500),
)
for deployment_id, ttft, completion_tokens in samples:
kwargs = _streaming_kwargs(deployment_id, start_time, ttft)
response_obj = _chat_response(completion_tokens=completion_tokens)
if sync_mode:
handler.log_success_event(
response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time
)
else:
await handler.async_log_success_event(
response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time
)
assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(1.0)]
assert _recorded_ttft(cache, SLOW_TTFT_ID) == [pytest.approx(3.0)]
request_kwargs = {"stream": True, "metadata": {}}
if sync_mode:
picked = handler.get_available_deployments(
model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs
)
else:
picked = await handler.async_get_available_deployments(
model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs
)
assert picked is not None
assert picked["model_info"]["id"] == FAST_TTFT_ID
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
async def test_ttft_window_keeps_newest_samples_when_full(sync_mode: bool):
"""Float timestamps, as the SDK passes them. Once max_latency_list_size
samples exist the oldest TTFT is dropped so the window slides."""
max_size = 3
cache = DualCache()
handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": max_size})
start_time = 1_700_000_000.0
ttfts = (0.1, 0.2, 0.3, 0.4)
for ttft in ttfts:
kwargs = {
"litellm_params": {
"metadata": {"model_group": MODEL_GROUP},
"model_info": {"id": FAST_TTFT_ID},
},
"stream": True,
"completion_start_time": start_time + ttft,
}
response_obj = _chat_response(completion_tokens=1)
if sync_mode:
handler.log_success_event(
response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0
)
else:
await handler.async_log_success_event(
response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0
)
assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(ttft) for ttft in ttfts[-max_size:]]
@pytest.mark.asyncio
async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_workers():
"""Workers on the previous release share the Redis map and keep writing
seconds-per-token under the old "time_to_first_token" key during a rolling
deploy. Those samples favor SLOW; routing must only read the seconds key."""
cache = DualCache()
handler = LowestLatencyLoggingHandler(router_cache=cache)
cache.set_cache(
key=f"{MODEL_GROUP}_map",
value={
FAST_TTFT_ID: {"time_to_first_token": [0.02], "time_to_first_token_seconds": [1.0]},
SLOW_TTFT_ID: {"time_to_first_token": [0.006], "time_to_first_token_seconds": [3.0]},
},
)
picked = await handler.async_get_available_deployments(
model_group=MODEL_GROUP,
healthy_deployments=STREAMING_DEPLOYMENTS,
request_kwargs={"stream": True, "metadata": {}},
)
assert picked is not None
assert picked["model_info"]["id"] == FAST_TTFT_ID
@pytest.mark.asyncio
@pytest.mark.parametrize(
"cached_entry",

76
tests/test_litellm/router_strategy/test_router_routing_groups.py
View file

@ -12,6 +12,7 @@ import pytest
import litellm
from litellm import Router
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.router import RoutingGroup, RoutingStrategy
@ -435,6 +436,81 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp
assert router._group_selectors == {}
def test_two_least_busy_groups_count_a_request_once(monkeypatch):
"""
Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and
back down from the success hooks on `litellm.callbacks`. The success list drops a second
selector of the same class, so a pre-call list that kept both counted every request twice
and released it once, and the deployment's in-flight count climbed until it looked pinned.
"""
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
router = _build_router(
routing_strategy="least-busy",
routing_groups=[
{
"group_name": "fast",
"models": ["filtered-model"],
"routing_strategy": "least-busy",
}
],
)
kwargs = {
"litellm_params": {
"metadata": {"model_group": "filtered-model"},
"model_info": {"id": "deploy-1"},
}
}
for callback in litellm.input_callback:
if isinstance(callback, CustomLogger):
callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs)
for callback in litellm.callbacks:
if isinstance(callback, CustomLogger):
callback.log_success_event(kwargs, None, None, None)
assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0
def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch):
"""
Least-busy hangs its counting off litellm's global callback lists, and those lists keep one
logger per class unless the instances differ in a plain attribute. Two routers in one process
(a second Router, or a per-request `user_config` one) therefore have to register separately:
a second router whose selector is dropped counts nothing, reads zero for every deployment,
and sends every request to whichever one is listed first.
"""
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "input_callback", [])
first = _build_router(routing_strategy="least-busy")
second = _build_router(routing_strategy="least-busy")
kwargs = {
"litellm_params": {
"metadata": {"model_group": "filtered-model"},
"model_info": {"id": "deploy-1"},
}
}
for callback in litellm.input_callback:
if isinstance(callback, CustomLogger):
callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs)
assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1
assert (
second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"]
== "deploy-2"
)
for callback in litellm.callbacks:
if isinstance(callback, CustomLogger):
callback.log_success_event(kwargs, None, None, None)
assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0
assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0
# ---------------------------------------------------------------------------
# Direct helper coverage
# ---------------------------------------------------------------------------

77
tests/test_litellm/router_strategy/test_router_routing_plugins.py
View file

@ -56,6 +56,18 @@ class BlockEverything:
return context
class MessageRecorder:
"""Records what each plugin pass was handed, then blocks so the request stops there."""
def __init__(self):
self.seen = []
async def run(self, context: RoutingContext) -> RoutingContext:
self.seen.append(list(context.raw_messages))
context.candidate_models = []
return context
def _smart_router_model_list():
return [
{
@ -164,6 +176,71 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu
await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}])
@pytest.mark.asyncio
async def test_prompt_management_model_still_runs_the_plugin_pipeline():
"""
A prompt-management model routes through its own factory, which picked the deployment
on the synchronous path. Plugins never run there, so the guard turned every such request
into an error message about the caller's own API choice, on an async call the caller made
correctly. It also read the in-flight counts with a blocking call inside the event loop.
"""
router = Router(
model_list=[
{
"model_name": "cached-claude",
"litellm_params": {
"model": "anthropic_cache_control_hook/claude-sonnet-5",
"prompt_id": "cache-points",
},
}
],
routing_strategy="least-busy",
plugins=[BlockEverything()],
)
with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"):
await router.acompletion(
model="cached-claude",
messages=[{"role": "user", "content": "hi"}],
litellm_call_id="lit-7039",
)
@pytest.mark.asyncio
async def test_prompt_management_plugins_see_the_callers_own_messages():
"""
The prompt-management factory picks its deployment with a placeholder message, which was
harmless while that pick ran on the synchronous path (plugins never ran there at all). Now
that the pick runs the plugin pipeline, a plugin that classifies request content would score
the placeholder instead of the conversation, and the narrowing it produces decides which
deployments the real call is allowed to use.
"""
recorder = MessageRecorder()
router = Router(
model_list=[
{
"model_name": "cached-claude",
"litellm_params": {
"model": "anthropic_cache_control_hook/claude-sonnet-5",
"prompt_id": "cache-points",
},
}
],
routing_strategy="least-busy",
plugins=[recorder],
)
messages = [{"role": "user", "content": "wire me $40,000 to account 12345"}]
with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"):
await router.acompletion(
model="cached-claude",
messages=messages,
litellm_call_id="lit-7039",
)
assert recorder.seen == [messages]
@pytest.mark.asyncio
async def test_router_without_plugins_is_unaffected():
"""Regression guard: a Router with no `plugins` configured behaves exactly as before."""

73
tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py
View file

@ -7,6 +7,7 @@ from typing import Final
import pytest
from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
from litellm.router_utils.auto_router_tuning_baseline import (
DEFAULT_TUNING_FINGERPRINT,
HEURISTIC_V1_TUNING_FIELDS,
@ -21,6 +22,33 @@ from litellm.router_utils.auto_router_tuning_baseline import (
_TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}
_ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"}
_KEYWORD_DIMENSION: Final = {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}
_HISTORICAL_FINGERPRINTS: Final = (
({}, "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"),
(
{"custom_dimensions": [_KEYWORD_DIMENSION]},
"b5c3c3f3be6341a8a16148d68d9067e03f94ed01a0bbfcde7955763042744372",
),
(
{"custom_dimensions": [{"name": "sqlDdl", "weight": 0.4, "patterns": [r"\bCREATE\s{1,4}TABLE\b"]}]},
"814ce0017fc7f60a160b262f658d910e9bdf784e6139a4ba4f1e2657aa203950",
),
(
{
"tiers": _TIERS,
"dimension_weights": {"codePresence": 0.3},
"custom_dimensions": [
{
"name": "internalFrameworks",
"weight": 0.2,
"keywords": ["orbitmesh", "fluxgate"],
"patterns": [r"\bALTER\s{1,4}TABLE\b"],
}
],
},
"38970dc9224e265ab38c89674563d8d0537822591f9239b45251db6f5ca6cc39",
),
)
def _router(
@ -77,6 +105,27 @@ class TestTuningFingerprint:
def test_explicit_empty_tier_model_configs_follow_omission(self) -> None:
assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT
@pytest.mark.parametrize(("config", "fingerprint"), _HISTORICAL_FINGERPRINTS)
def test_fingerprints_recorded_before_scoring_mode_existed_are_preserved(
self, config: Mapping[str, object], fingerprint: str
) -> None:
"""Literal hashes captured from the merged implementation at 9bc9104102, before CustomDimension.scoring_mode."""
assert tuning_fingerprint(config) == fingerprint
def test_binary_scoring_mode_hashes_like_its_absence(self) -> None:
historical: Final = tuning_fingerprint({"custom_dimensions": [_KEYWORD_DIMENSION]})
explicit: Final = tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "binary"}]})
reserialized: Final = ComplexityRouterConfig.model_validate(
{"custom_dimensions": [_KEYWORD_DIMENSION]}
).model_dump(mode="json", include={"custom_dimensions"})
assert reserialized["custom_dimensions"][0]["scoring_mode"] == "binary"
assert reserialized["custom_dimensions"][0]["patterns"] == []
assert historical == explicit == tuning_fingerprint(reserialized)
assert (
tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]})
!= historical
)
def test_tier_model_overrides_change_the_fingerprint(self) -> None:
plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}})
with_override = tuning_fingerprint(
@ -218,15 +267,18 @@ class TestQuota:
is None
)
def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self) -> None:
@pytest.mark.parametrize(
"edit",
[
pytest.param({"weight": 0.9}, id="weight"),
pytest.param({"scoring_mode": "match_count"}, id="scoring-mode"),
],
)
def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self, edit: Mapping[str, object]) -> None:
baselines: Final = snapshot_tuning_baselines(())
original: Final = _router("a", {})
config: Final = {
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}]
}
edited_config: Final = {
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.9, "keywords": ["orbitmesh"]}]
}
config: Final = {"custom_dimensions": [_KEYWORD_DIMENSION]}
edited_config: Final = {"custom_dimensions": [{**_KEYWORD_DIMENSION, **edit}]}
added: Final = _router("a", config)
edited: Final = _router("a", edited_config)
second: Final = _router("b", config)
@ -240,6 +292,13 @@ class TestQuota:
assert mutable_tuned_identities((original,), baselines) == frozenset()
assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None
def test_graded_dimension_recorded_at_snapshot_is_its_own_baseline(self) -> None:
graded: Final = _router("a", {"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]})
baselines: Final = snapshot_tuning_baselines((graded,))
assert mutable_tuned_identities((graded,), baselines) == frozenset()
reverted_to_binary: Final = _router("a", {"custom_dimensions": [_KEYWORD_DIMENSION]})
assert mutable_tuned_identities((reverted_to_binary,), baselines) == {router_identity(graded)}
def test_violation_message_names_the_limit_and_remedy(self) -> None:
message = tuning_limit_violation(held=2, limit=1)
assert message is not None

184
tests/test_litellm/router_utils/test_cooldown_cache.py
View file

@ -268,12 +268,12 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time() - 120.0,
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600)
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
assert active == [], "Expired cooldown entry must not appear in active cooldowns"
assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
def test_active_entry_is_returned(self):
"""
@ -289,7 +289,7 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time(),
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60)
cc.in_memory_cache.set_cache(key, active_value, ttl=60)
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
@ -312,14 +312,14 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time() - (60.0 - remaining),
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, value, ttl=600)
cc.in_memory_cache.set_cache(key, value, ttl=600)
before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
before_expiry = cc.in_memory_cache.ttl_dict.get(key)
assert before_expiry is not None
cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
after_expiry = cc.in_memory_cache.ttl_dict.get(key)
assert after_expiry is not None
corrected_remaining = after_expiry - time.time()
assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
@ -340,12 +340,12 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time() - 120.0,
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600)
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
assert active == [], "Expired entry must not appear in async active cooldowns"
assert cc.cache.in_memory_cache.get_cache(key) is None
assert cc.in_memory_cache.get_cache(key) is None
@pytest.mark.asyncio
async def test_async_active_entry_is_returned(self):
@ -363,7 +363,7 @@ class TestCooldownCacheTTLCorrection:
"timestamp": time.time(),
"cooldown_time": 60.0,
}
cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60)
cc.in_memory_cache.set_cache(key, active_value, ttl=60)
active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
@ -389,18 +389,18 @@ class TestCorrectedActiveCooldown:
cc = self._make_cooldown_cache()
key = "deployment:expired-dep:cooldown"
entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0)
cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600)
cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)
result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
assert result is None
assert cc.cache.in_memory_cache.get_cache(key) is None
assert cc.in_memory_cache.get_cache(key) is None
def test_active_entry_within_window_returns_value(self):
cc = self._make_cooldown_cache()
key = "deployment:active-dep:cooldown"
entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60)
cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
@ -412,12 +412,12 @@ class TestCorrectedActiveCooldown:
key = "deployment:backfilled-dep:cooldown"
remaining = 30.0
entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0)
cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600)
cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)
result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
assert result is not None
corrected_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
corrected_expiry = cc.in_memory_cache.ttl_dict.get(key)
assert corrected_expiry is not None
assert corrected_expiry - time.time() <= 60.0
@ -425,10 +425,160 @@ class TestCorrectedActiveCooldown:
cc = self._make_cooldown_cache()
key = "deployment:normal-dep:cooldown"
entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60)
original_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
original_expiry = cc.in_memory_cache.ttl_dict.get(key)
cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key)
after_expiry = cc.in_memory_cache.ttl_dict.get(key)
assert after_expiry == original_expiry
class SharedRedisDouble:
"""
In-process stand-in for RedisCache, shared by several DualCache instances so that
tests can model two proxy replicas talking to one Redis.
"""
def __init__(self) -> None:
self.store: dict = {} # mutable-ok: stands in for Redis' own mutable keyspace
def set_cache(self, key, value, **kwargs):
self.store[key] = value
async def async_set_cache(self, key, value, **kwargs):
self.store[key] = value
def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
return {key: self.store.get(key) for key in key_list}
async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
return {key: self.store.get(key) for key in key_list}
class TestCooldownPropagationBetweenReplicas:
"""
A cooldown written by one replica has to reach its siblings quickly. The router's own
DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads
get their own cache with a much shorter Redis read interval.
"""
def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache:
router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis)
if read_interval is None:
return CooldownCache(cache=router_cache, default_cooldown_time=60.0)
return CooldownCache(
cache=router_cache,
default_cooldown_time=60.0,
redis_read_interval_seconds=read_interval,
)
@pytest.mark.asyncio
async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self):
redis = SharedRedisDouble()
replica_a = self._make_replica(redis, read_interval=0.25)
replica_b = self._make_replica(redis, read_interval=0.25)
model_id = "shared-deployment"
assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []
replica_a.add_deployment_to_cooldown(
model_id=model_id,
original_exception=Exception("Internal server error"),
exception_status=500,
cooldown_time=60.0,
)
time.sleep(0.3)
active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
assert [model_id] == [entry[0] for entry in active], (
"sibling replica must pick up a cooldown written by another replica within the read interval"
)
@pytest.mark.asyncio
async def test_sibling_replica_sees_cooldown_within_default_read_interval(self):
redis = SharedRedisDouble()
replica_a = self._make_replica(redis)
replica_b = self._make_replica(redis)
model_id = "default-interval-deployment"
assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []
replica_a.add_deployment_to_cooldown(
model_id=model_id,
original_exception=Exception("Internal server error"),
exception_status=500,
cooldown_time=60.0,
)
time.sleep(1.2)
active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
assert [model_id] == [entry[0] for entry in active], (
"the shipped default read interval must let a sibling replica see a cooldown about a second later"
)
def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self):
redis = SharedRedisDouble()
replica_a = self._make_replica(redis, read_interval=0.25)
replica_b = self._make_replica(redis, read_interval=0.25)
model_id = "sync-shared-deployment"
assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == []
replica_a.add_deployment_to_cooldown(
model_id=model_id,
original_exception=Exception("Internal server error"),
exception_status=500,
cooldown_time=60.0,
)
time.sleep(0.3)
active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None)
assert [model_id] == [entry[0] for entry in active]
@pytest.mark.asyncio
async def test_redis_attached_after_construction_is_still_used(self):
redis = SharedRedisDouble()
router_cache = DualCache(in_memory_cache=InMemoryCache())
writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25)
router_cache.attach_redis_cache(redis)
reader = self._make_replica(redis, read_interval=0.25)
model_id = "late-redis-deployment"
writer.add_deployment_to_cooldown(
model_id=model_id,
original_exception=Exception("Internal server error"),
exception_status=500,
cooldown_time=60.0,
)
active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None)
assert [model_id] == [entry[0] for entry in active], (
"a router that wires Redis after building its cooldown cache must still publish cooldowns to it"
)
class TestCooldownSurvivesUnrelatedCacheTraffic:
@pytest.mark.asyncio
async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self):
router_cache = DualCache(in_memory_cache=InMemoryCache())
cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0)
model_id = "busy-router-deployment"
cc.add_deployment_to_cooldown(
model_id=model_id,
original_exception=Exception("Internal server error"),
exception_status=500,
cooldown_time=30.0,
)
for i in range(400):
router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i})
active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None)
assert [model_id] == [entry[0] for entry in active], (
"unrelated router cache traffic must not evict a cooldown that is still running"
)

53
tests/test_litellm/test_component_entrypoint.py
View file

@ -223,6 +223,59 @@ def test_gating_matches_the_monolithic_entrypoint_and_get_secret_bool(
assert monolith[1] == ("args=litellm --port 4000" if traced else "args=--port 4000")
def test_wipes_the_prometheus_multiproc_dir_before_uvicorn_forks(tmp_path: Path) -> None:
"""A restarted container inherits the emptyDir of its predecessor, whose worker pids it may reuse, so the
stale .db files must be gone before any worker opens the one carrying its own pid."""
multiproc_dir = tmp_path / "multiproc"
multiproc_dir.mkdir()
(multiproc_dir / "gauge_livesum_7.db").write_bytes(b"stale")
(multiproc_dir / "counter_7.db").write_bytes(b"stale")
(multiproc_dir / "keep.txt").write_text("not a sample")
bin_dir = tmp_path / "bin"
bin_dir.mkdir()
_write_stubs(bin_dir, ("uvicorn",))
record = tmp_path / "record.txt"
env = {
**os.environ,
"PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
"RECORD": str(record),
"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir),
}
env.pop("USE_DDTRACE", None)
result = subprocess.run(
["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"],
env=env,
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
assert sorted(p.name for p in multiproc_dir.iterdir()) == ["keep.txt"]
assert record.read_text().splitlines()[0] == "exec=uvicorn"
def test_creates_a_missing_prometheus_multiproc_dir(tmp_path: Path) -> None:
bin_dir = tmp_path / "bin"
bin_dir.mkdir()
_write_stubs(bin_dir, ("uvicorn",))
missing = tmp_path / "multiproc"
env = {
**os.environ,
"PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
"RECORD": str(tmp_path / "record.txt"),
"PROMETHEUS_MULTIPROC_DIR": str(missing),
}
env.pop("USE_DDTRACE", None)
result = subprocess.run(
["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"], env=env, capture_output=True, text=True
)
assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
assert missing.is_dir()
def _copied_script(dockerfile: Path, image_path: str) -> Path:
"""Resolve the repo file a Dockerfile `COPY`s to `image_path`, so tests run what the image ships."""
matches = _COPY_RE.findall(dockerfile.read_text())

44
tests/test_litellm/test_utils.py
View file

@ -6092,3 +6092,47 @@ class TestFinalOptionalParamsLineRedaction:
assert "'max_tokens': 17" in printed
assert "'temperature': 0.25" in printed
def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()]
def test_load_credentials_from_list_warns_when_the_named_credential_is_not_loaded(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.utils import load_credentials_from_list
monkeypatch.setattr(litellm, "credential_list", [])
request_kwargs = {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"}
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
load_credentials_from_list(request_kwargs)
assert request_kwargs == {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"}
assert _credential_warnings(caplog) == [
"litellm_credential_name=openai-cred matched none of the 0 loaded credentials; the request runs without it"
]
def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_without_warning(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
from litellm.types.utils import CredentialItem
from litellm.utils import load_credentials_from_list
loaded = CredentialItem(
credential_name="openai-cred",
credential_values={"api_key": "sk-from-db", "api_base": "https://credential.example"},
credential_info={},
)
monkeypatch.setattr(litellm, "credential_list", [loaded])
request_kwargs = {"litellm_credential_name": "openai-cred", "api_base": "https://request.example"}
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
load_credentials_from_list(request_kwargs)
assert request_kwargs == {
"litellm_credential_name": "openai-cred",
"api_base": "https://request.example",
"api_key": "sk-from-db",
}
assert _credential_warnings(caplog) == []

8
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22180
"limit": 22174
},
"LIT002": {
"limit": 26729
"limit": 26715
},
"LIT003": {
"limit": 261
@ -27,10 +27,10 @@
"limit": 0
},
"LIT010": {
"limit": 16426
"limit": 16398
},
"LIT011": {
"limit": 5506
"limit": 5504
},
"LIT012": {
"limit": 4486

12
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx
View file

@ -187,10 +187,10 @@ describe("AutoRouterBenchmarksTab", () => {
});
it.each([
{ spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18" },
{ spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00" },
{ spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004" },
])("shows total classification cost across $turns turns without a per-turn rate", ({ llm, cost, ...values }) => {
{ spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18", rate: "$2.43" },
{ spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00", rate: "$0.00" },
{ spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004", rate: "$0.0040" },
])("shows total classification cost and its rate across $turns turns", ({ llm, cost, rate, ...values }) => {
const stats = totals({ ...values, saved_spend: 10126.28, baseline_spend: values.spend + 10126.28 });
mockHook({ data: response([group(stats)], stats) });
renderTab();
@ -201,7 +201,7 @@ describe("AutoRouterBenchmarksTab", () => {
.map((node) => node.textContent)
.slice(1, 3),
).toEqual([llm, cost]);
expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument();
expect(screen.getByText(`(${rate} / 1K turns)`)).toBeInTheDocument();
expect(screen.getAllByText("$10,126.28").length).toBeGreaterThan(0);
});
@ -237,7 +237,7 @@ describe("AutoRouterBenchmarksTab", () => {
expect(terms).toEqual([
"Actual auto-router spend",
"LLM spend",
"Classification cost",
"Classification cost($2.00 / 1K turns)",
"Estimated spend at highest-tier model",
]);
expect(values).toEqual(["$359.86", "$353.71", "$6.15", "$2,534.45"]);

26
ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx
View file

@ -30,7 +30,7 @@ import {
type BenchmarkView,
type BucketRow,
} from "./autoRouterBenchmarks";
import { formatRangeLabel, usd } from "./costOptimizationUtils";
import { classificationRatePer1kTurns, formatRangeLabel, usd } from "./costOptimizationUtils";
import ShadowEvalSection from "./ShadowEvalSection";
import TierTurnsChart from "./TierTurnsChart";
import { useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks";
@ -52,9 +52,17 @@ const Metric: React.FC<{ label: string; value: string; hint?: string }> = ({ lab
</Card>
);
const SpendRow: React.FC<{ label: string; value: string; subdued?: boolean }> = ({ label, value, subdued }) => (
const SpendRow: React.FC<{ label: string; value: string; hint?: string; subdued?: boolean }> = ({
label,
value,
hint,
subdued,
}) => (
<dl className="flex flex-wrap items-baseline justify-between gap-x-6 gap-y-1 py-2">
<dt className="min-w-0 text-sm text-muted-foreground">{label}</dt>
<dt className="flex min-w-0 flex-wrap items-baseline gap-x-2 text-sm text-muted-foreground">
{label}
{hint && <span className="text-xs">{hint}</span>}
</dt>
<dd
className={`min-w-0 break-all tabular-nums ${subdued ? "text-sm font-normal text-muted-foreground" : "text-base font-semibold text-foreground"}`}
>
@ -99,6 +107,11 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => {
subdued
label="Classification cost"
value={stats.classifier_cost == null ? "Unavailable" : usd(stats.classifier_cost)}
hint={
stats.classifier_cost == null
? undefined
: classificationRatePer1kTurns(stats.classifier_cost, stats.turns)
}
/>
</div>
{stats.classifier_cost == null && (
@ -277,9 +290,10 @@ const BenchmarksBody: React.FC<BenchmarksBodyProps> = ({ isPending, error, data,
<p className="text-xs text-muted-foreground">
Compares your actual routed spend with the estimated cost of using only the most expensive model configured in
the auto-router. It accounts for both the cache savings from staying on one model and the added cache costs from
switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. The
range counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets
savings by UTC day.
switching models. Savings are net of recorded LLM classification cost, which is included in actual spend.
Classification cost per 1K turns is averaged over all auto-router turns, including those that skip
classification. The range counts whole sessions that overlap it, so totals can differ slightly from the Overall
tab, which buckets savings by UTC day.
</p>
<div className="space-y-4">

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