scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

fix(anthropic): retain both imports when merging main

Browse source
This commit is contained in:
林SO 2026-10-15 07:51:37 +08:00
commit 8a2cf196d4
265 changed files with 17791 additions and 1624 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
helm/litellm/README.md
View file

@ -5,3 +5,21 @@ This chart deploys separate gateway, backend and UI services. Configure Lens thr
Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for the complete values and verification flow. It also links the GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
For an ingress with separate component hostnames, set all three `ingress.componentHosts` values, leave `ingress.host` and `ingress.extraPaths` empty, configure `ui.backendUrl` as the backend's public origin and `backend.gatewayUrl` as the gateway's public origin, and set `ui.publicUrl` to the browser-facing UI origin. The UI Deployment receives `LITELLM_BACKEND_URL` from `ui.backendUrl` only when `ingress.componentHosts` is configured.
```yaml
ingress:
enabled: true
host: ""
componentHosts:
ui: ui.example.com
backend: backend.example.com
gateway: gateway.example.com
extraPaths: []
ui:
publicUrl: https://ui.example.com
backendUrl: https://backend.example.com
backend:
gatewayUrl: https://gateway.example.com
```

10
helm/litellm/templates/NOTES.txt
View file

@ -22,12 +22,14 @@ Reminders:
- database.writer.{host,port,dbname} (writer connection pieces)
- database.writer.passwordSecret.{name,usernameKey,passwordKey}
(Secret holding the writer DB username + password)
- database.writer.passwordSecret.{hostKey,urlKey} (optional, read writer host or full URL from that Secret)
- database.writer.useIAMAuth: true (optional — chart sets IAM_TOKEN_DB_AUTH=true and
omits DATABASE_PASSWORD / DATABASE_URL so the proxy
mints the URL from an IAM token at startup)
- database.reader.host (optional — enables read-replica routing; reader
.passwordSecret.name is required when set, unless
.useIAMAuth is true)
- database.reader.host / database.reader.passwordSecret.{hostKey,urlKey}
(optional, enables read-replica routing. The Secret
name is required for hostKey, urlKey, and for host
with password auth)
- database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true —
chart emits DATABASE_*_READ_REPLICA env vars and
omits DATABASE_PASSWORD_READ_REPLICA /
@ -35,6 +37,8 @@ Reminders:
reader URL from an IAM token at startup)
- redis.passwordSecret.name (optional — set when redis.host is provided and the
cache requires auth)
- redis.passwordSecret.{hostKey,urlKey} (optional, read Redis host or full URL from that Secret)
(the Secret name is required when either key is set)
- redis.cluster: true (optional — chart sets REDIS_CLUSTER_NODES from
redis.host / redis.port so the proxy's Cache()
constructs a RedisClusterCache; the cluster client

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

@ -27,6 +27,30 @@ Common naming + label helpers shared by gateway, backend, and ui templates.
{{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}}
{{- end -}}
{{- define "litellm.ui.publicOrigin" -}}
{{- $origin := .Values.ui.publicUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "ui.publicUrl must be an origin like https://ui.example.com, without a path" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.ui.backendOrigin" -}}
{{- $origin := .Values.ui.backendUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.backend.gatewayOrigin" -}}
{{- $origin := .Values.backend.gatewayUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.commonLabels" -}}
app.kubernetes.io/name: {{ include "litellm.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
@ -202,7 +226,9 @@ the chart context (needed for .Values), `component` selects which component's
Sensitive values (master key, DB username + password, Redis password) come
only from referenced Secrets; the chart never accepts inline values for them.
The chart never assembles DATABASE_URL itself. It emits only the discrete
The chart never assembles DATABASE_URL itself. When
database.writer.passwordSecret.urlKey is set, the chart passes the operator's
URL through from the Secret. Otherwise it emits only the discrete
DATABASE_HOST/PORT/USER/NAME/SCHEMA (+ DATABASE_PASSWORD for password auth)
vars; the proxy's entrypoint (DatabaseURLSettings in
litellm/proxy/db/db_url_settings.py) builds the URL from them and
@ -218,11 +244,11 @@ token instead of a static password. `database.writer.useAzureEntraAuth: true`
does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
for Azure Database for PostgreSQL. The two are mutually exclusive.
The read replica is opt-in via `database.reader.host`. The chart emits
DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
password auth. When `database.reader.useIAMAuth: true` (or
The read replica is opt-in via `database.reader.host`, `hostKey`, or `urlKey`.
The chart emits DATABASE_URL_READ_REPLICA alone for urlKey, or discrete
DATABASE_*_READ_REPLICA vars for host and hostKey configurations, plus
DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for password auth.
When `database.reader.useIAMAuth: true` (or
`database.reader.useAzureEntraAuth: true`) it omits
DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
same way. Reader token auth only takes effect when the writer uses the same
@ -242,8 +268,36 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
value: {{ $component.logLevel | quote }}
{{- end }}
{{- with $root.Values.database.writer }}
{{- $hostKey := .passwordSecret.hostKey }}
{{- $urlKey := .passwordSecret.urlKey }}
{{- if and .host $hostKey }}
{{- fail "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or .host $hostKey .dbname .schema .sslMode .sslRootCert .useIAMAuth .useAzureEntraAuth) }}
{{- fail "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" }}
{{- end }}
{{- if and (not .host) (not $hostKey) (not $urlKey) }}
{{- fail "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not .passwordSecret.name) }}
{{- fail "database.writer.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $urlKey }}
- name: DATABASE_URL
valueFrom:
secretKeyRef:
name: {{ .passwordSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else }}
- name: DATABASE_HOST
value: {{ required "database.writer.host is required" .host | quote }}
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ .passwordSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ .host | quote }}
{{- end }}
- name: DATABASE_PORT
value: {{ .port | default 5432 | quote }}
- name: DATABASE_USER
@ -282,8 +336,29 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
key: {{ .passwordSecret.passwordKey | default "password" }}
{{- end }}
{{- end }}
{{- end }}
{{- with $root.Values.database.reader }}
{{- if .host }}
{{- $readerSecret := .passwordSecret | default dict }}
{{- $hostKey := $readerSecret.hostKey }}
{{- $urlKey := $readerSecret.urlKey }}
{{- $readerEnabled := or .host $hostKey $urlKey }}
{{- if and .host $hostKey }}
{{- fail "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or .host $hostKey .dbname .schema .useIAMAuth .useAzureEntraAuth) }}
{{- fail "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not $readerSecret.name) }}
{{- fail "database.reader.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $readerEnabled }}
{{- if $urlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ $readerSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else }}
{{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
{{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
{{- end }}
@ -291,11 +366,18 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
{{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
{{- end }}
- name: DATABASE_HOST_READ_REPLICA
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ $readerSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ .host | quote }}
{{- end }}
- name: DATABASE_PORT_READ_REPLICA
value: {{ .port | default 5432 | quote }}
- name: DATABASE_NAME_READ_REPLICA
value: {{ required "database.reader.dbname is required when database.reader.host is set" .dbname | quote }}
value: {{ required "database.reader.dbname is required when database.reader.host, hostKey, or urlKey is set" .dbname | quote }}
{{- if .schema }}
- name: DATABASE_SCHEMA_READ_REPLICA
value: {{ .schema | quote }}
@ -325,6 +407,7 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
{{- end }}
{{- end }}
{{- end }}
{{- end }}
{{/*
The migrations Job (helm.sh/hook: pre-upgrade) is the single owner of
`prisma migrate deploy`. Without this, every gateway/backend pod also runs
@ -341,9 +424,34 @@ harmless no-op for the Job and authoritative for the app pods.
tracking, pod lock manager) via its REDIS_* env fallback. An explicit
`general_settings.coordination_redis` block in proxy_config takes
precedence over anything emitted here. */}}
{{- if $root.Values.redis.host }}
{{- $redisSecret := $root.Values.redis.passwordSecret | default dict }}
{{- $hostKey := $redisSecret.hostKey }}
{{- $urlKey := $redisSecret.urlKey }}
{{- if and $root.Values.redis.host $hostKey }}
{{- fail "redis.host and redis.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or $root.Values.redis.host $hostKey $root.Values.redis.cluster) }}
{{- fail "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not $redisSecret.name) }}
{{- fail "redis.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $urlKey }}
- name: REDIS_URL
valueFrom:
secretKeyRef:
name: {{ $redisSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else if or $root.Values.redis.host $hostKey }}
- name: REDIS_HOST
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ $redisSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ $root.Values.redis.host | quote }}
{{- end }}
- name: REDIS_PORT
value: {{ $root.Values.redis.port | quote }}
{{- if $root.Values.redis.passwordSecret.name }}
@ -360,7 +468,11 @@ harmless no-op for the Job and authoritative for the app pods.
configured endpoint; the cluster client discovers the remaining nodes from
CLUSTER SLOTS at startup. */}}
- name: REDIS_CLUSTER_NODES
{{- if $hostKey }}
value: {{ printf "[{\"host\":\"$(REDIS_HOST)\",\"port\":%v}]" (int $root.Values.redis.port) | quote }}
{{- else }}
value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }}
{{- end }}
{{- end }}
{{- end }}
{{- with $component.extraEnv }}

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

@ -1,4 +1,6 @@
{{- if .Values.backend.enabled }}
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
apiVersion: apps/v1
kind: Deployment
metadata:
@ -58,7 +60,17 @@ spec:
protocol: TCP
env:
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
{{- if .Values.ui.publicUrl }}
- name: LITELLM_CORS_ORIGINS
value: {{ include "litellm.ui.publicOrigin" . | quote }}
- name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
{{- end }}
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
{{- if .Values.backend.gatewayUrl }}
- name: LITELLM_GATEWAY_URL
value: {{ if $hasComponentHosts }}{{ include "litellm.backend.gatewayOrigin" . | quote }}{{ else }}{{ .Values.backend.gatewayUrl | quote }}{{ end }}
{{- end }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml

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

@ -56,6 +56,10 @@ spec:
protocol: TCP
env:
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
{{- if .Values.ui.publicUrl }}
- name: LITELLM_CORS_ORIGINS
value: {{ include "litellm.ui.publicOrigin" . | quote }}
{{- end }}
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH

94
helm/litellm/templates/ingress.yaml
View file

@ -1,3 +1,36 @@
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
{{- if $hasComponentHosts }}
{{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }}
{{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }}
{{- end }}
{{- if not .Values.gateway.enabled }}
{{- fail "ingress.componentHosts requires gateway.enabled=true" }}
{{- end }}
{{- if not .Values.backend.enabled }}
{{- fail "ingress.componentHosts requires backend.enabled=true" }}
{{- end }}
{{- if not .Values.ui.enabled }}
{{- fail "ingress.componentHosts requires ui.enabled=true" }}
{{- end }}
{{- if .Values.ingress.host }}
{{- fail "ingress.host must be empty when ingress.componentHosts is configured" }}
{{- end }}
{{- if .Values.ingress.extraPaths }}
{{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }}
{{- end }}
{{- if not (trim .Values.ui.backendUrl) }}
{{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- $_ := include "litellm.ui.backendOrigin" . }}
{{- if not (trim .Values.backend.gatewayUrl) }}
{{- fail "backend.gatewayUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- $_ := include "litellm.backend.gatewayOrigin" . }}
{{- if not (trim .Values.ui.publicUrl) }}
{{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- end }}
{{- if .Values.ingress.enabled -}}
{{- $gatewayName := include "litellm.gateway.fullname" . -}}
{{- $backendName := include "litellm.backend.fullname" . -}}
@ -90,6 +123,66 @@
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
{{- if $hasComponentHosts }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
name: {{ include "litellm.fullname" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
{{- with .Values.ingress.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
{{- with .Values.ingress.className }}
ingressClassName: {{ . | quote }}
{{- end }}
{{- with .Values.ingress.tls }}
tls:
{{- toYaml . | nindent 4 }}
{{- end }}
rules:
- host: {{ $componentHosts.ui | quote }}
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
- host: {{ $componentHosts.backend | quote }}
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: {{ $backendName }}
port:
number: {{ $backendPort }}
- host: {{ $componentHosts.gateway | quote }}
http:
paths:
{{- if or (eq (include "litellm.lens.mode" .) "bundled") (and (eq (include "litellm.lens.mode" .) "external") .Values.lensWorker.externalServiceName) }}
- path: /lens-ingest
pathType: Prefix
backend:
service:
name: {{ .Values.lensWorker.externalServiceName | default (printf "%s-lens-worker" (include "litellm.fullname" .)) }}
port:
number: {{ .Values.lensWorker.service.port }}
{{- end }}
- path: /
pathType: Prefix
backend:
service:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
{{- else }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -217,3 +310,4 @@ spec:
port:
number: {{ $backendPort }}
{{- end }}
{{- end }}

6
helm/litellm/templates/ui/deployment.yaml
View file

@ -1,3 +1,5 @@
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
{{- if .Values.ui.enabled }}
apiVersion: apps/v1
kind: Deployment
@ -56,9 +58,9 @@ spec:
- name: LITELLM_LOG
value: {{ .Values.ui.logLevel | quote }}
{{- end }}
{{- if .Values.ui.backendUrl }}
{{- if and .Values.ui.backendUrl $hasComponentHosts }}
- name: LITELLM_BACKEND_URL
value: {{ .Values.ui.backendUrl | quote }}
value: {{ include "litellm.ui.backendOrigin" . | quote }}
{{- end }}
{{- with .Values.ui.extraEnv }}
{{- toYaml . | nindent 12 }}

441
helm/litellm/tests/component_hosts_tests.yaml Normal file
View file

@ -0,0 +1,441 @@
suite: ingress component hosts
templates:
- ingress.yaml
- backend/deployment.yaml
- gateway/deployment.yaml
- gateway/configmap.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: routes each hostname to its component and injects the external URLs
template: ingress.yaml
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
lensWorker.mode: disabled
ingress.tls:
- hosts:
- ui.example.com
- backend.example.com
- gateway.example.com
secretName: component-hosts
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- equal:
path: spec.rules[0].host
value: ui.example.com
- equal:
path: spec.tls[0]
value:
hosts:
- ui.example.com
- backend.example.com
- gateway.example.com
secretName: component-hosts
- equal:
path: spec.rules[0].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- equal:
path: spec.rules[1].host
value: backend.example.com
- equal:
path: spec.rules[1].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- equal:
path: spec.rules[2].host
value: gateway.example.com
- equal:
path: spec.rules[2].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- it: keeps Lens ingestion ahead of the gateway catch-all
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
lensWorker.mode: external
lensWorker.externalServiceName: existing-lens
lensWorker.externalUrl: https://lens.example.com
lensWorker.publicUrl: https://lens.example.com
lensWorker.serviceTokenSecret.name: existing-lens-signing
lensWorker.gateway.secretName: existing-lens-signing
asserts:
- equal:
path: spec.rules[2].http.paths[0].path
value: /lens-ingest
- equal:
path: spec.rules[2].http.paths[0].backend.service.name
value: existing-lens
- equal:
path: spec.rules[2].http.paths[1].path
value: /
template: ingress.yaml
- it: rejects a partial component host configuration
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts.ui, backend, and gateway must all be set together
template: ingress.yaml
- it: rejects component hosts when the gateway is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
gateway.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires gateway.enabled=true
template: ingress.yaml
- it: rejects component hosts when the backend is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
backend.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires backend.enabled=true
template: ingress.yaml
- it: rejects component hosts when the UI is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
ui.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires ui.enabled=true
template: ingress.yaml
- it: rejects a legacy host alongside component hosts
set:
ingress.enabled: true
ingress.host: legacy.example.com
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
asserts:
- failedTemplate:
errorMessage: ingress.host must be empty when ingress.componentHosts is configured
template: ingress.yaml
- it: rejects extra paths alongside component hosts
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ingress.extraPaths:
- path: /custom
asserts:
- failedTemplate:
errorMessage: ingress.extraPaths must be empty when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the UI backend URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.backendUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the backend gateway URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the public UI origin in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.publicUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: exposes the gateway URL in the backend Deployment
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: " https://gateway.example.com/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_GATEWAY_URL
value: https://gateway.example.com
- it: preserves the backend gateway URL verbatim in single-host mode
set:
ingress.enabled: true
ingress.host: litellm.example.com
backend.gatewayUrl: " https://gateway.example.com/api/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_GATEWAY_URL
value: " https://gateway.example.com/api/ "
- it: exposes the backend URL in the UI Deployment
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: " https://backend.example.com/ "
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
template: ui/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BACKEND_URL
value: https://backend.example.com
- it: omits the backend URL from the UI Deployment in single-host mode
set:
ingress.enabled: true
ingress.host: litellm.example.com
ui.backendUrl: https://litellm.example.com/api
template: ui/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BACKEND_URL
any: true
- it: enables credentialed CORS from the public UI origin on the backend
set:
ui.publicUrl: https://ui.example.com
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
- it: lets backend extraEnv override generated CORS values
set:
ui.publicUrl: https://ui.example.com
backend.extraEnv:
- name: LITELLM_CORS_ORIGINS
value: https://backend-override.example.com
- name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "false"
template: backend/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].env[0].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[0].value
value: https://ui.example.com
- equal:
path: spec.template.spec.containers[0].env[1].name
value: LITELLM_CORS_ALLOW_CREDENTIALS
- equal:
path: spec.template.spec.containers[0].env[1].value
value: "true"
- equal:
path: spec.template.spec.containers[0].env[10].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[10].value
value: https://backend-override.example.com
- equal:
path: spec.template.spec.containers[0].env[11].name
value: LITELLM_CORS_ALLOW_CREDENTIALS
- equal:
path: spec.template.spec.containers[0].env[11].value
value: "false"
- it: exposes the public UI origin to the gateway without enabling credentials
set:
ui.publicUrl: https://ui.example.com
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
- it: trims the public UI origin before setting backend CORS
set:
ui.publicUrl: " https://ui.example.com/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- it: trims the public UI origin before setting gateway CORS
set:
ui.publicUrl: " https://ui.example.com/ "
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- it: rejects a UI public URL with a path
set:
ui.publicUrl: https://ui.example.com/app
template: gateway/deployment.yaml
asserts:
- failedTemplate:
errorMessage: ui.publicUrl must be an origin like https://ui.example.com, without a path
- it: lets gateway extraEnv override the generated CORS origin
set:
ui.publicUrl: https://ui.example.com
gateway.extraEnv:
- name: LITELLM_CORS_ORIGINS
value: https://gateway-override.example.com
template: gateway/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].env[0].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[0].value
value: https://ui.example.com
- equal:
path: spec.template.spec.containers[0].env[9].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[9].value
value: https://gateway-override.example.com
---
suite: component host UI backend origin validation
templates:
- ingress.yaml
values:
- ./values/required.yaml
tests:
- it: rejects a pathful UI backend URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com/api
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set
- it: rejects a pathful backend gateway URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com/v1
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set
- it: rejects a backend gateway URL without a scheme in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: gateway.example.com
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set

397
helm/litellm/tests/database_auth_tests.yaml
View file

@ -191,3 +191,400 @@ tests:
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true
- it: writer hostKey reads the host from the Secret and supports IAM auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.hostKey: hostname
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- it: writer urlKey supplies only DATABASE_URL to each database client container and keeps pool env
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.schema: ""
database.writer.sslMode: ""
database.writer.sslRootCert: ""
database.writer.passwordSecret.urlKey: connection-url
database.connectionPool.enabled: true
gateway.collector.enabled: true
gateway.metricsServer.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: gateway/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: gateway/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[2].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_USER
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_PASSWORD
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: backend/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: backend/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: backend/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: backend/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: migrations-job.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: migrations-job.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: migrations-job.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: migrations-job.yaml
- it: writer rejects host together with hostKey
template: gateway/deployment.yaml
set:
database.writer.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set"
- it: writer rejects urlKey with discrete connection settings
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.schema: public
database.writer.sslMode: require
database.writer.sslRootCert: /etc/ssl/ca.crt
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer rejects urlKey with IAM auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer rejects urlKey with Entra auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer requires an endpoint
template: gateway/deployment.yaml
set:
database.writer.host: ""
asserts:
- failedTemplate:
errorMessage: "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required"
- it: writer hostKey requires a Secret name
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.hostKey: hostname
database.writer.passwordSecret.name: ""
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.name is required when hostKey or urlKey is set"
- it: reader renders without replica env vars when passwordSecret is null
template: gateway/deployment.yaml
set:
database.reader.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
any: true
- it: reader hostKey enables replica routing and reads its host from the Secret
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.reader.passwordSecret.hostKey: hostname
database.reader.dbname: litellm
database.reader.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME_READ_REPLICA
value: litellm
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: reader urlKey emits only the replica URL
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: connection-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA_READ_REPLICA
any: true
- it: reader rejects host together with hostKey
template: gateway/deployment.yaml
set:
database.reader.host: reader.example.com
database.reader.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set"
- it: reader rejects urlKey with discrete connection settings
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
database.reader.dbname: litellm
database.reader.schema: public
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
- it: reader rejects urlKey with token auth
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.reader.passwordSecret.urlKey: connection-url
database.reader.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
- it: reader urlKey requires a Secret name
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
database.reader.passwordSecret.name: ""
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.name is required when hostKey or urlKey is set"

213
helm/litellm/tests/redis_env_tests.yaml
View file

@ -21,6 +21,27 @@ tests:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway renders without redis env vars when passwordSecret is null
template: gateway/deployment.yaml
set:
redis.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- it: gateway emits host, port and password when redis is configured
template: gateway/deployment.yaml
set:
@ -68,6 +89,27 @@ tests:
name: redis-secret
key: redis-password
- it: backend renders without redis env vars when passwordSecret is null
template: backend/deployment.yaml
set:
redis.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- it: gateway omits REDIS_PASSWORD for an auth-less redis
template: gateway/deployment.yaml
set:
@ -107,3 +149,174 @@ tests:
content:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway gets redis host from hostKey and keeps port and password
template: gateway/deployment.yaml
set:
redis.host: ""
redis.port: 6380
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.passwordKey: redis-password
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
valueFrom:
secretKeyRef:
name: redis-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6380"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: redis-password
- it: backend gets redis host from hostKey with the same coordination env
template: backend/deployment.yaml
set:
redis.host: ""
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
valueFrom:
secretKeyRef:
name: redis-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6379"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: password
- it: gateway expands the redis cluster seed from hostKey at runtime
template: gateway/deployment.yaml
set:
redis.host: ""
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.cluster: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
value: '[{"host":"$(REDIS_HOST)","port":6379}]'
- equal:
path: spec.template.spec.containers[0].env[8].name
value: REDIS_HOST
- equal:
path: spec.template.spec.containers[0].env[11].name
value: REDIS_CLUSTER_NODES
- it: gateway gets REDIS_URL from urlKey without discrete Redis env
template: gateway/deployment.yaml
set:
redis.passwordSecret.name: redis-secret
redis.passwordSecret.urlKey: connection-url
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
valueFrom:
secretKeyRef:
name: redis-secret
key: connection-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
any: true
- it: redis rejects host together with hostKey
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.host and redis.passwordSecret.hostKey cannot both be set"
- it: redis urlKey rejects host
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis urlKey rejects hostKey
template: gateway/deployment.yaml
set:
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis urlKey rejects cluster mode
template: gateway/deployment.yaml
set:
redis.cluster: true
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis hostKey requires a Secret name
template: gateway/deployment.yaml
set:
redis.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
- it: redis urlKey requires a Secret name
template: gateway/deployment.yaml
set:
redis.passwordSecret.urlKey: connection-url
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"

18
helm/litellm/values.yaml
View file

@ -23,7 +23,11 @@ ingress:
# wildcard pathType, so that rule could never match there.
controller: alb
annotations: {}
host: "" # optional; if set, becomes the rule's host
host: "" # optional; if set, becomes the single ingress rule's host
componentHosts:
ui: ""
backend: ""
gateway: ""
tls: []
# Extra HTTP paths appended to the ingress rule. Additive: every built-in
# UI / gateway / backend path is still rendered, these entries are placed
@ -208,6 +212,8 @@ database:
name: litellm-writer-secret
usernameKey: username
passwordKey: password
hostKey: "" # Secret key containing the writer database host
urlKey: "" # Secret key containing the complete writer database URL
# libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
# in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
# sslRootCert on its own implies sslMode verify-full
@ -229,6 +235,8 @@ database:
name: litellm-reader-secret
usernameKey: username
passwordKey: password
hostKey: "" # Secret key containing the reader database host
urlKey: "" # Secret key containing the complete reader database URL
# In-container connection pool (PgBouncer, transaction mode) shared by every
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
@ -274,6 +282,8 @@ redis:
passwordSecret:
name: "" # Leave empty for auth-less Redis
passwordKey: password
hostKey: "" # Secret key containing the Redis host
urlKey: "" # Secret key containing the complete Redis URL
# ---------- gateway (LLM data plane) ----------
gateway:
@ -493,6 +503,7 @@ gateway:
backend:
enabled: true
logLevel: INFO
gatewayUrl: ""
extraEnv: []
envConfigMaps: []
envSecrets: []
@ -575,9 +586,10 @@ ui:
service:
type: ClusterIP
port: 3000
# The dashboard expects to know where to reach the backend API. Set this to
# the externally-routable URL (typically the ingress host + /api or similar).
# The backend's public origin, passed to the UI only when ingress.componentHosts is configured.
backendUrl: ""
# Public origin of the UI, used for backend and gateway CORS.
publicUrl: ""
resources:
requests:
cpu: 500m

21
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py
View file

@ -13,6 +13,7 @@ from litellm_proxy_extras.prisma_toolchain import (
)
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
if TYPE_CHECKING:
import psycopg
@ -32,7 +33,7 @@ def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
if not database_url:
return environment
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
return {**environment, "DATABASE_URL": migration_url}
return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
@dataclass(frozen=True, slots=True)
@ -55,10 +56,15 @@ class MigrationCoordinator:
def check_connection(self) -> None:
self.connection.execute("SELECT 1")
def release_statement_snapshot(self) -> None:
# a parameterized statement leaves its portal snapshot registered on the backend until the
# next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
self.connection.execute("SELECT 1")
def acquire_prisma_lock(self) -> None:
deadline: Final = time.monotonic() + migration_lock_timeout()
while time.monotonic() < deadline:
if _try_lock(self.connection, 72707369):
if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
return
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
raise RuntimeError(
@ -95,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
)
@contextmanager
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
"""Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
pinned connection while a Prisma command runs with its session advisory lock disabled, so a
transaction pooler can never strand that lock on an idle backend after the command exits."""
with migration_lock(database_url) as coordinator:
coordinator.acquire_prisma_lock()
coordinator.release_statement_snapshot()
yield
@contextmanager
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql Normal file
View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_add_access_group_skill_names/migration.sql Normal file
View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_skill_names" TEXT[] DEFAULT ARRAY[]::TEXT[];

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

@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
mcp_tool_permissions Json?
access_agent_ids String[] @default([])
access_skill_names String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])

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

@ -9,6 +9,7 @@ import tempfile
import threading
import time
from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
@ -1021,7 +1022,11 @@ class ProxyExtrasDBManager:
def _setup_database_v2(use_migrate: bool) -> bool:
if not use_migrate:
return ProxyExtrasDBManager._run_database_v2(False)
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
from litellm_proxy_extras.migration_lock import (
migration_environment,
migration_lock,
prisma_migration_lock,
)
from litellm_proxy_extras.migration_recovery import (
baseline_current_schema,
recover_completed_migration,
@ -1055,7 +1060,9 @@ class ProxyExtrasDBManager:
migration_environment(_get_prisma_env()),
)
while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
while not ProxyExtrasDBManager._run_database_v2(
True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url)
):
continue
return True
@ -1064,6 +1071,7 @@ class ProxyExtrasDBManager:
use_migrate: bool,
recover_completed: Callable[[str], bool] = lambda name: False,
baseline_existing: "Callable[[str], None] | None" = None,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> bool:
"""
v2 migration resolver (opt-in via --use_v2_migration_resolver).
@ -1126,11 +1134,12 @@ class ProxyExtrasDBManager:
while not budget.exhausted:
attempt_started = time.monotonic()
try:
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=deploy_timeout,
env=migration_env,
)
with deploy_lock():
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=deploy_timeout,
env=migration_env,
)
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
return True
@ -1155,6 +1164,7 @@ class ProxyExtrasDBManager:
e,
budget,
time.monotonic() - attempt_started,
deploy_lock=deploy_lock,
)
if next_budget.attempts_left < budget.attempts_left:
@ -1187,15 +1197,19 @@ class ProxyExtrasDBManager:
return match.group(1) if match else None
@staticmethod
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
def _v2_roll_back_migration_best_effort(
migration_name: str,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> None:
from litellm_proxy_extras.migration_lock import migration_environment
try:
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
timeout=prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
with deploy_lock():
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
timeout=prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass
@ -1204,6 +1218,7 @@ class ProxyExtrasDBManager:
error: subprocess.CalledProcessError,
budget: "_MigrateAttemptBudget",
attempt_seconds: float = 0.0,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> "_MigrateAttemptBudget":
"""Recover from one failed `prisma migrate deploy`, and price the pass.
@ -1225,14 +1240,14 @@ class ProxyExtrasDBManager:
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
logger.info(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
logger.info(
@ -1276,7 +1291,7 @@ class ProxyExtrasDBManager:
"rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
@ -1286,7 +1301,7 @@ class ProxyExtrasDBManager:
"its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
raise RuntimeError(

7
litellm/_redis.py
View file

@ -675,7 +675,10 @@ def _get_redis_client_logic(**env_overrides):
if redis_kwargs.get("startup_nodes")
else redis.Redis
)
return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
coerced_kwargs: Final = _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
if coerced_kwargs.get("socket_timeout") is None:
coerced_kwargs["socket_timeout"] = REDIS_SOCKET_TIMEOUT
return coerced_kwargs
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
@ -722,7 +725,6 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:
@ -755,7 +757,6 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:

5
litellm/_service_logger.py
View file

@ -152,6 +152,11 @@ class ServiceLogging(CustomLogger):
except Exception as e:
verbose_logger.exception("Error dispatching service event - %s", e)
def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
prometheus_logger: Final[PrometheusServicesLogger | None] = getattr(self, "prometheusServicesLogger", None)
if prometheus_logger is not None:
prometheus_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
def service_success_hook(
self,
service: ServiceTypes,

123
litellm/caching/redis_cache.py
View file

@ -18,7 +18,9 @@ import json
import logging
import threading
import time
from collections.abc import Awaitable, Callable, Iterator, Sequence
import weakref
from collections import Counter
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import timedelta
@ -51,12 +53,12 @@ from .base_cache import BaseCache
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prometheus_client import Counter as _PromCounter
from prometheus_client import Gauge as _PromGauge
from redis.asyncio import Redis, RedisCluster
from redis.asyncio.client import Pipeline
from redis.asyncio.cluster import ClusterPipeline
from litellm._service_logger import ServiceLogging
pipeline = Pipeline
cluster_pipeline = ClusterPipeline
async_redis_client = Redis
@ -243,7 +245,9 @@ class RedisCircuitBreaker:
recovery_timeout: int,
enabled: bool = True,
timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
service_logger: "ServiceLogging | None" = None,
) -> None:
self._service_logger = service_logger
self.failure_threshold = failure_threshold
self.recovery_timeout = recovery_timeout
self.enabled = enabled
@ -255,7 +259,12 @@ class RedisCircuitBreaker:
self._opened_at: float | None = None
self._state = self.CLOSED
self._generation = 0
_breaker_metrics().record_state_change(None, self._state)
_LIVE_BREAKERS.add(self)
self._report(transitions_to=None, failure_class=None)
@property
def state(self) -> str:
return self._state
@property
def generation(self) -> int:
@ -301,7 +310,7 @@ class RedisCircuitBreaker:
self._timeout_count = 0
self._timeout_streak_started_at = None
self._opened_at = now
_breaker_metrics().record_failure("timeout" if is_timeout else "connectivity")
self._report(transitions_to=None, failure_class="timeout" if is_timeout else "connectivity")
if self._should_open(now):
if self._state != self.OPEN:
verbose_logger.warning(
@ -327,10 +336,13 @@ class RedisCircuitBreaker:
def _set_state(self, state: str) -> None:
if state == self._state:
return
_breaker_metrics().record_transition(state)
_breaker_metrics().record_state_change(self._state, state)
self._state = state
self._generation += 1
self._report(transitions_to=state, failure_class=None)
def _report(self, transitions_to: str | None, failure_class: str | None) -> None:
if self._service_logger is not None:
self._service_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
_RedisCallResult = TypeVar("_RedisCallResult")
@ -400,12 +412,19 @@ def _redis_timeout_error_types() -> tuple[type, ...]:
``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from
py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own
class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither.
``MaxConnectionsError`` is the async cluster client's pool-exhausted signal (its nodes have
no blocking pool, so it is raised instantly instead of after a wait): a busy pool, not a
dead Redis, like the standalone pool's chained wait timeout.
"""
try:
from redis.exceptions import TimeoutError as RedisTimeoutError
from redis import exceptions as redis_exceptions
except ImportError:
return (TimeoutError, asyncio.TimeoutError)
return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError)
max_connections_error: Final = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
"type[Exception]",
redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # types-redis 4.6 stubs predate it
)
return (redis_exceptions.TimeoutError, TimeoutError, asyncio.TimeoutError, max_connections_error)
_MAX_EXCEPTION_CAUSE_DEPTH: Final = 20
@ -430,60 +449,6 @@ def is_redis_timeout_failure(exc: BaseException) -> bool:
return any(isinstance(link, timeout_types) for link in _explicit_causes(exc))
class _BreakerMetrics:
"""Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent.
Registered lazily on the default registry (which /metrics serves) via the module-level
``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers.
"""
def __init__(self) -> None:
self._state_gauge: _PromGauge | None = None
self._transitions: _PromCounter | None = None
self._failures: _PromCounter | None = None
try:
from prometheus_client import Counter as PromCounter
from prometheus_client import Gauge
except ImportError:
return
self._state_gauge = Gauge(
"litellm_redis_circuit_breaker_state",
"Number of Redis circuit breakers currently in each state",
labelnames=("state",),
multiprocess_mode="livesum",
)
self._transitions = PromCounter(
"litellm_redis_circuit_breaker_transitions",
"Redis circuit breaker state transitions",
labelnames=("state",),
)
self._failures = PromCounter(
"litellm_redis_circuit_breaker_failures",
"Redis health failures counted by the circuit breaker",
labelnames=("failure_class",),
)
def record_state_change(self, old_state: str | None, new_state: str) -> None:
if self._state_gauge is None:
return
if old_state is not None:
self._state_gauge.labels(old_state).dec()
self._state_gauge.labels(new_state).inc()
def record_transition(self, state: str) -> None:
if self._transitions is not None:
self._transitions.labels(state).inc()
def record_failure(self, failure_class: str) -> None:
if self._failures is not None:
self._failures.labels(failure_class).inc()
@functools.lru_cache(maxsize=1)
def _breaker_metrics() -> _BreakerMetrics:
return _BreakerMetrics()
def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
"""Record a Redis failure that the calling method is about to swallow.
@ -499,6 +464,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet()
def redis_breaker_state_counts() -> Mapping[str, int]:
"""How many live Redis circuit breakers sit in each state, for the Prometheus services logger."""
return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS)))
class RedisCircuitBreakerOpenError(Exception):
pass
@ -638,9 +611,14 @@ def _redis_circuit_breaker_guard(method):
"""
@functools.wraps(method)
async def wrapper(self, *args, **kwargs):
async def wrapper(self: "RedisCache", *args, **kwargs):
async def guarded() -> object:
return await method(self, *args, **kwargs)
return await _run_under_circuit_breaker(
self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
self._circuit_breaker, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator
wrapper.__name__,
lambda: self._admit_async_call(guarded), # pyright: ignore[reportPrivateUsage] # same decorator
)
return wrapper
@ -685,12 +663,9 @@ class RedisCache(BaseCache):
redis_kwargs["socket_timeout"] = socket_timeout
### HEALTH MONITORING OBJECT ###
if kwargs.get("service_logger_obj", None) is not None and isinstance(
kwargs["service_logger_obj"], ServiceLogging
):
self.service_logger_obj: ServiceLogging = ServiceLogging()
if isinstance(kwargs.get("service_logger_obj"), ServiceLogging):
self.service_logger_obj = kwargs.pop("service_logger_obj")
else:
self.service_logger_obj = ServiceLogging()
redis_kwargs.update(kwargs)
self.redis_client = get_redis_client(**redis_kwargs)
@ -717,6 +692,7 @@ class RedisCache(BaseCache):
failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
enabled=REDIS_CIRCUIT_BREAKER_ENABLED,
service_logger=self.service_logger_obj,
)
self._setup_health_pings()
@ -726,6 +702,9 @@ class RedisCache(BaseCache):
else:
super().__init__() # defaults to 60s
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
return call()
def _setup_health_pings(self):
"""Setup async and sync health pings for Redis."""
# ASYNC HEALTH PING
@ -1106,7 +1085,9 @@ class RedisCache(BaseCache):
litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor)
return await executor(keys=keys, args=args, client=client)
return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute)
return await _run_under_circuit_breaker(
self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute)
)
return run_script

105
litellm/caching/redis_cluster_cache.py
View file

@ -5,9 +5,17 @@ Key differences:
- RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
"""
from typing import TYPE_CHECKING, Any, Final
import asyncio
import contextvars
import os
import weakref
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar
from litellm._logging import verbose_logger
from litellm.caching.redis_cache import RedisCache
from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -23,9 +31,100 @@ else:
Span = Any
_T: Final = TypeVar("_T")
def _connection_limit(max_connections: object) -> int | None:
try:
limit: Final = int(max_connections) # pyright: ignore[reportArgumentType] # config value, str or int
except (TypeError, ValueError):
return None
return limit if limit > 0 else None
_ClientSemaphores: TypeAlias = weakref.WeakValueDictionary[str, asyncio.Semaphore]
@dataclass(slots=True)
class _SemaphoreRegistry:
"""Weak on both layers: a semaphore binds to its loop once a waiter queues on it, so a strong value would keep a
closed loop alive. In-flight ``ClusterConnectionGate.run`` calls hold the semaphore they acquired."""
_per_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _ClientSemaphores]" = field(
default_factory=weakref.WeakKeyDictionary, repr=False
)
def get(self, client_key: str, limit: int) -> asyncio.Semaphore:
per_loop: Final = self._per_loop.setdefault(asyncio.get_running_loop(), weakref.WeakValueDictionary())
existing: Final = per_loop.get(client_key)
if existing is not None:
return existing
created: Final = asyncio.Semaphore(limit)
per_loop[client_key] = created
return created
_Admission: TypeAlias = tuple[str, "asyncio.Task[object] | None"]
_SEMAPHORES: Final = _SemaphoreRegistry()
_ADMITTED: Final[contextvars.ContextVar[frozenset[_Admission]]] = contextvars.ContextVar(
"redis_cluster_gate_admitted", default=frozenset()
)
@dataclass(frozen=True, slots=True)
class ClusterConnectionGate:
"""The async ``RedisCluster`` has no blocking pool, so past ``max_connections`` a node raises
``MaxConnectionsError`` instead of waiting and the circuit breaker opens on a healthy Redis. This gate makes
guarded commands wait for a slot the way the standalone ``BlockingConnectionPool`` does. A command that finds no
slot within ``wait_timeout`` runs anyway, so a genuinely saturated pool still reaches the client's own error
handling and the breaker."""
limit: int
wait_timeout: float = REDIS_CONNECTION_POOL_TIMEOUT
@classmethod
def from_max_connections(cls, max_connections: object) -> "ClusterConnectionGate | None":
limit: Final = _connection_limit(max_connections)
return None if limit is None else cls(limit=limit)
async def run(self, call: Callable[[], Awaitable[_T]], client_key: str) -> _T:
admission: Final[_Admission] = (client_key, asyncio.current_task())
if admission in _ADMITTED.get():
return await call()
semaphore: Final = _SEMAPHORES.get(client_key, self.limit)
try:
await asyncio.wait_for(semaphore.acquire(), timeout=self.wait_timeout)
except asyncio.TimeoutError:
verbose_logger.debug(
"Redis Cluster: no free connection slot within %ss (max_connections=%s), sending anyway",
self.wait_timeout,
self.limit,
)
return await call()
token: Final = _ADMITTED.set(_ADMITTED.get() | {admission})
try:
return await call()
finally:
_ADMITTED.reset(token)
semaphore.release()
class RedisClusterCache(RedisCache):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __init__(self, *args, max_connections: object = None, connection_gate: object = None, **kwargs):
pool_kwargs: Final = {} if max_connections is None else {"max_connections": max_connections}
super().__init__(*args, **pool_kwargs, **kwargs)
self._connection_gate: Final = (
connection_gate
if isinstance(connection_gate, ClusterConnectionGate)
else ClusterConnectionGate.from_max_connections(
os.environ.get("REDIS_MAX_CONNECTIONS") if max_connections is None else max_connections
)
)
def _admit_async_call(self, call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]:
if self._connection_gate is None:
return call()
return self._connection_gate.run(call, self._get_async_client_cache_key())
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""

6
litellm/constants.py
View file

@ -60,6 +60,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
"enable_tag_filtering",
"tag_routing_prefix",
"optional_pre_call_checks",
"fallback_on_unsupported_params",
}
)
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
@ -75,6 +76,9 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
}
)
DEFAULT_BATCH_SIZE: Final = get_env_int("DEFAULT_BATCH_SIZE", 512)
SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE: Final = get_env_int_in_range(
"SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE", 1000, 1, 10_000
)
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)
CLICKHOUSE_BATCH_SIZE: Final = get_env_int("CLICKHOUSE_BATCH_SIZE", 10_000)
CLICKHOUSE_FLUSH_INTERVAL_SECONDS: Final = get_env_float("CLICKHOUSE_FLUSH_INTERVAL_SECONDS", 1.0)
@ -547,7 +551,7 @@ DEFAULT_POLLING_INTERVAL: Final = get_env_float(
AZURE_OPERATION_POLLING_TIMEOUT: Final = get_env_int("AZURE_OPERATION_POLLING_TIMEOUT", 120)
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30"))
AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = get_env_int("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96)
REDIS_SOCKET_TIMEOUT: Final = get_env_float("REDIS_SOCKET_TIMEOUT", 0.1)
REDIS_SOCKET_TIMEOUT: Final = get_env_float("REDIS_SOCKET_TIMEOUT", 5.0)
CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0
REDIS_CONNECTION_POOL_TIMEOUT: Final = get_env_int("REDIS_CONNECTION_POOL_TIMEOUT", 5)
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = get_env_int("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)

8
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -799,12 +799,16 @@ class AnthropicCacheControlHook(CustomPromptManagement):
The hook sees neither the tools nor the request kwargs, so the target dialect
and the client's breakpoint count outside the message list ride on the points.
Builds copies because config-owned point dicts are shared across requests.
A zero count replaces an earlier layer's stamp (the chat bridge runs behind the
Responses layer, whose outside marks sit in the bridge's own message list) and is
otherwise left off. Builds copies because config-owned point dicts are shared
across requests.
"""
with_dialect: Final = AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
)
if external_breakpoints == 0:
stamped_earlier: Final = any(EXTERNAL_BREAKPOINTS_STAMP in point for point in with_dialect)
if external_breakpoints == 0 and not stamped_earlier:
return with_dialect
return AnthropicCacheControlHook._stamped(with_dialect, EXTERNAL_BREAKPOINTS_STAMP, external_breakpoints)

59
litellm/integrations/prometheus_services.py
View file

@ -3,7 +3,9 @@
# On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers)
from typing import Final
from dataclasses import dataclass
from functools import lru_cache
from typing import TYPE_CHECKING, Final
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -17,6 +19,42 @@ from litellm.types.services import (
FAILED_REQUESTS_LABELS: Final = ["error_class", "function_name"]
if TYPE_CHECKING:
from prometheus_client import Counter as PrometheusCounter
from prometheus_client import Gauge
@dataclass(frozen=True, slots=True)
class _RedisBreakerCollectors:
state: "Gauge"
transitions: "PrometheusCounter"
failures: "PrometheusCounter"
@lru_cache(maxsize=1)
def _redis_breaker_collectors() -> _RedisBreakerCollectors:
from prometheus_client import Counter as PrometheusCounter
from prometheus_client import Gauge
return _RedisBreakerCollectors(
state=Gauge(
"litellm_redis_circuit_breaker_state",
"Number of Redis circuit breakers currently in each state",
labelnames=("state",),
multiprocess_mode="livesum",
),
transitions=PrometheusCounter(
"litellm_redis_circuit_breaker_transitions",
"Redis circuit breaker state transitions",
labelnames=("state",),
),
failures=PrometheusCounter(
"litellm_redis_circuit_breaker_failures",
"Redis health failures counted by the circuit breaker",
labelnames=("failure_class",),
),
)
class PrometheusServicesLogger:
# Class variables or attributes
@ -41,6 +79,11 @@ class PrometheusServicesLogger:
self.Counter = Counter
self.Gauge = Gauge
self.REGISTRY = REGISTRY
breaker_collectors: Final = _redis_breaker_collectors()
self.redis_breaker_state = breaker_collectors.state
self.redis_breaker_transitions = breaker_collectors.transitions
self.redis_breaker_failures = breaker_collectors.failures
self._refresh_redis_breaker_gauge()
verbose_logger.debug("in init prometheus services metrics")
@ -98,6 +141,20 @@ class PrometheusServicesLogger:
return DEFAULT_METRICS
return metrics
def _refresh_redis_breaker_gauge(self) -> None:
from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts
counts: Final = redis_breaker_state_counts()
for state in (RedisCircuitBreaker.CLOSED, RedisCircuitBreaker.OPEN, RedisCircuitBreaker.HALF_OPEN):
self.redis_breaker_state.labels(state).set(counts.get(state, 0))
def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
self._refresh_redis_breaker_gauge()
if transitions_to is not None:
self.redis_breaker_transitions.labels(transitions_to).inc()
if failure_class is not None:
self.redis_breaker_failures.labels(failure_class).inc()
def is_metric_registered(self, metric_name) -> bool:
# Use _names_to_collectors (O(1)) instead of REGISTRY.collect() (O(n)) to avoid
# perf regression when a new Router is created per request (e.g. router_settings in DB).

40
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
View file

@ -6,7 +6,7 @@ import traceback
from collections.abc import Mapping, Sequence
from typing import Final, Literal, cast
from pydantic import ConfigDict, TypeAdapter, with_config
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
from typing_extensions import ReadOnly, TypedDict
import litellm
@ -222,6 +222,9 @@ async def convert_to_streaming_response_async(
role=choice["message"]["role"],
function_call=choice["message"].get("function_call", None),
tool_calls=choice["message"].get("tool_calls", None),
provider_specific_fields=_cached_provider_specific_fields(
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(choice["message"])
),
)
finish_reason = choice.get("finish_reason", None)
@ -243,6 +246,9 @@ async def convert_to_streaming_response_async(
completion_tokens=response_object["usage"].get("completion_tokens", 0),
prompt_tokens=response_object["usage"].get("prompt_tokens", 0),
total_tokens=response_object["usage"].get("total_tokens", 0),
completion_tokens_details=_cached_completion_tokens_details(
_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
),
),
)
@ -334,6 +340,9 @@ def convert_to_streaming_response(
model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0)
model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0)
model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0)
model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details(
_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
)
if "id" in response_object:
model_response_object.id = response_object["id"]
@ -390,6 +399,35 @@ class _ParallelToolUseArgs(TypedDict):
_PARALLEL_TOOL_USE_ARGS: Final = TypeAdapter(_ParallelToolUseArgs)
class _CachedMessageProviderFields(TypedDict, total=False):
provider_specific_fields: ReadOnly[object]
class _CachedUsageDetails(TypedDict, total=False):
completion_tokens_details: ReadOnly[object]
_CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields)
_CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails)
_FIELD_MAP: Final = TypeAdapter(dict[str, object])
_COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper)
def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None:
fields: Final = message.get("provider_specific_fields")
return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else None
def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionTokensDetailsWrapper | None:
details: Final = usage.get("completion_tokens_details")
if not isinstance(details, dict):
return None
try:
return _COMPLETION_TOKENS_DETAILS.validate_python(details)
except ValidationError:
return None
def handle_invalid_parallel_tool_calls(
tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall],
) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None:

6
litellm/litellm_core_utils/streaming_handler.py
View file

@ -280,6 +280,7 @@ class CustomStreamWrapper:
self._response_headers = _response_headers
self.response_id: str | None = None
self.received_model_id: str | None = None
self.logging_loop = None
self.rules = Rules()
self.stream_options = stream_options or getattr(logging_obj, "stream_options", None)
@ -838,6 +839,7 @@ class CustomStreamWrapper:
self.response_id = id
if id and isinstance(id, str) and id.strip():
self.received_model_id = id
model_response.hidden_params["received_model_id"] = id
if self.response_id is not None and isinstance(self.response_id, str):
@ -1756,7 +1758,9 @@ class CustomStreamWrapper:
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
def finish_reason_handler(self):
model_response: Final = self.model_response_creator()
model_response: Final = self.model_response_creator(
hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id}
)
_finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason
if _finish_reason is not None:
model_response.choices[0].finish_reason = _finish_reason

11
litellm/llms/anthropic/chat/handler.py
View file

@ -66,7 +66,7 @@ from litellm.types.utils import (
)
from ...base import BaseLLM
from ..common_utils import AnthropicError, process_anthropic_headers
from ..common_utils import AnthropicError, process_anthropic_headers, refusal_provider_specific_fields
from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig, JsonModeRequest
if TYPE_CHECKING:
@ -1104,7 +1104,8 @@ class ModelResponseIterator:
# These are automatically handled by Anthropic API, we just pass them through
pass
elif type_chunk == "message_delta":
finish_reason, usage, container = self._handle_message_delta(chunk)
finish_reason, usage, container, refusal_fields = self._handle_message_delta(chunk)
provider_specific_fields.update(refusal_fields)
text = self._release_answer_close() + self._drain_split_answer_blocks()
if container:
provider_specific_fields["container"] = container
@ -1299,7 +1300,9 @@ class ModelResponseIterator:
return ""
return _split_answer_text(blocks, self.response_format_schema)
def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]:
def _handle_message_delta(
self, chunk: dict
) -> tuple[str, Usage | None, dict[str, Any] | None, Mapping[str, object]]:
"""
Handle message_delta event for finish_reason, usage, and container.
@ -1321,7 +1324,7 @@ class ModelResponseIterator:
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
)
container: Final = message_delta["delta"].get("container")
return finish_reason, usage, container
return finish_reason, usage, container, refusal_provider_specific_fields(message_delta["delta"])
def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None:
"""

47
litellm/llms/anthropic/chat/transformation.py
View file

@ -104,10 +104,12 @@ from litellm.utils import (
)
from ..common_utils import (
REFUSAL_STOP_FIELDS,
AnthropicError,
AnthropicModelInfo,
eager_input_streaming_flag,
process_anthropic_headers,
refusal_provider_specific_fields,
requires_native_compaction_beta,
strip_advisor_blocks_from_messages,
)
@ -294,7 +296,6 @@ DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING: Final = (
DROP_UNSUPPORTED_SPEED_WARNING: Final = (
"Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models."
)
SERVICE_TIER_PROVIDERS: Final = frozenset({"anthropic", "azure_ai"})
OPENAI_TO_ANTHROPIC_SERVICE_TIER: Final = MappingProxyType(
@ -306,6 +307,12 @@ DROP_UNSUPPORTED_SERVICE_TIER_WARNING: Final = (
"Anthropic only accepts 'auto' and 'standard_only'."
)
MESSAGES_SERVICE_TIER_PROVIDERS: Final = SERVICE_TIER_PROVIDERS | {"bedrock_mantle"}
DROP_SERVICE_TIER_PROVIDER_WARNING: Final = (
"Dropping `service_tier` for model=%s on provider=%s. Only the Anthropic API routes accept it."
)
class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"""
@ -539,6 +546,24 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
) -> None:
return cls._maybe_drop_speed_param(model, optional_params, drop_params, custom_llm_provider)
@staticmethod
def provider_supports_service_tier(custom_llm_provider: str | None) -> bool:
"""Anthropic's ``service_tier`` is accepted by the direct API and the Foundry and Mantle gateways in
front of it. Vertex rejects the field and Bedrock Invoke only takes its own tier names."""
return custom_llm_provider is None or custom_llm_provider in MESSAGES_SERVICE_TIER_PROVIDERS
@classmethod
def maybe_drop_service_tier_param(
cls,
model: str,
optional_params: dict[str, object], # mutable-ok: mirrors override contract
custom_llm_provider: str | None = None,
) -> None:
if "service_tier" not in optional_params or cls.provider_supports_service_tier(custom_llm_provider):
return
litellm.verbose_logger.warning(DROP_SERVICE_TIER_PROVIDER_WARNING, model, custom_llm_provider)
optional_params.pop("service_tier", None)
@staticmethod
def _raise_invalid_reasoning_effort(model: str, value: object, llm_provider: str) -> NoReturn:
"""Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``.
@ -717,12 +742,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
# Already handled above
continue
if key == "properties" and isinstance(value, dict):
result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()}
if key in ("properties", "$defs") and isinstance(value, dict):
nested_schemas: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict
result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in nested_schemas.items()}
elif key == "items" and isinstance(value, dict):
result[key] = AnthropicConfig.filter_anthropic_output_schema(value)
elif key == "$defs" and isinstance(value, dict):
result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()}
nested_schema: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict
result[key] = AnthropicConfig.filter_anthropic_output_schema(nested_schema)
elif key == "anyOf" and isinstance(value, list):
result["anyOf"] = result.get("anyOf", []) + [
AnthropicConfig.filter_anthropic_output_schema(item) for item in value
@ -2700,7 +2725,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if compaction_blocks is not None:
provider_specific_fields["compaction_blocks"] = compaction_blocks
return provider_specific_fields
refusal_fields: Final = refusal_provider_specific_fields(
REFUSAL_STOP_FIELDS.validate_python(completion_response)
)
return {**provider_specific_fields, **refusal_fields}
def transform_parsed_response(
self,
@ -2783,7 +2811,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_message.provider_specific_fields = provider_specific_fields
if json_mode_message is not None:
refusal_fields: Final = refusal_provider_specific_fields(
REFUSAL_STOP_FIELDS.validate_python(completion_response)
)
completion_response["stop_reason"] = "stop"
if refusal_fields:
json_mode_message.provider_specific_fields = dict(refusal_fields)
_message = json_mode_message
model_response.choices[0].message = _message

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

@ -8,11 +8,12 @@ import re
from collections.abc import Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, ClassVar, Final, Literal, TypeVar
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
from urllib.parse import quote, urlparse
import httpx
from pydantic import ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.constants import (
@ -2150,3 +2151,35 @@ def create_anthropic_model_list_response(
"first_id": data[0]["id"] if data else None,
"last_id": data[-1]["id"] if data else None,
}
class RefusalStopFields(TypedDict, total=False):
stop_reason: ReadOnly[object]
stop_details: ReadOnly[object]
def _is_dict(value: object) -> bool:
return isinstance(value, dict)
def _no_stop_fields() -> Mapping[str, object]:
return {}
def _dict_or_empty(value: object) -> object:
return value if _is_dict(value) else _no_stop_fields()
REFUSAL_STOP_FIELDS: Final[TypeAdapter[RefusalStopFields]] = TypeAdapter(
Annotated[RefusalStopFields, BeforeValidator(_dict_or_empty)]
)
def refusal_provider_specific_fields(message: Mapping[str, object]) -> Mapping[str, object]:
stop_reason: Final = message.get("stop_reason")
stop_details: Final = message.get("stop_details")
if stop_details is not None:
return {"stop_reason": stop_reason, "stop_details": stop_details}
if stop_reason == "refusal":
return {"stop_reason": stop_reason}
return {}

1
litellm/llms/anthropic/pass_through/messages/transformation.py
View file

@ -98,6 +98,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
"output_format",
"inference_geo",
"speed",
"service_tier",
"output_config",
"reasoning_effort",
"safeguards",

5
litellm/llms/anthropic/pass_through/messages/utils.py
View file

@ -205,6 +205,11 @@ class AnthropicMessagesRequestUtils:
drop_params=drop_params,
custom_llm_provider=custom_llm_provider,
)
AnthropicConfig.maybe_drop_service_tier_param(
model=model,
optional_params=filtered_params,
custom_llm_provider=custom_llm_provider,
)
for param in ("temperature", "top_p", "top_k"):
if param in filtered_params:
AnthropicModelInfo._apply_sampling_param( # pyright: ignore[reportPrivateUsage] # same gating the /chat/completions path applies; forking it would drift

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

@ -1191,8 +1191,10 @@ def bedrock_model_accepts_cache_points(model: str | None) -> bool:
never silently degrade. ``litellm.utils.supports_prompt_caching`` is not reusable here:
it returns False for unmapped models, the opposite polarity.
"""
if model is None:
return True
return model is None or bedrock_declared_cache_point_support(model) is not False
def bedrock_declared_cache_point_support(model: str) -> bool | None:
if _OPENAI_FAMILY_MODEL_RE.search(model):
return False
map_keys: Final = (model, get_bedrock_base_model(model))
@ -1209,7 +1211,7 @@ def bedrock_model_accepts_cache_points(model: str | None) -> bool:
if explicit_marker_support is not None:
return explicit_marker_support
if not any(key in litellm.model_cost for key in map_keys):
return True
return None
return any(
entry.get("supports_prompt_caching") is True
for key in map_keys

4
litellm/llms/bedrock/messages/mantle_transformation.py
View file

@ -56,6 +56,10 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
model ID in the request body (unlike Bedrock Invoke which puts it in the URL).
"""
BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = (
AmazonAnthropicClaudeMessagesConfig.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS | {"service_tier"}
)
@property
def beta_headers_provider(self) -> str:
return "bedrock_mantle"

3
litellm/llms/chatgpt/responses/transformation.py
View file

@ -226,7 +226,8 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def _extract_error_message(self, parsed_chunk: dict[str, Any]) -> str | None:
error_obj: Final = parsed_chunk.get("error") or (parsed_chunk.get("response") or {}).get("error")
if error_obj is None:
return None
top_level_message: Final = parsed_chunk.get("message")
return top_level_message if isinstance(top_level_message, str) else None
if isinstance(error_obj, dict):
return error_obj.get("message") or str(error_obj)
return str(error_obj)

6
litellm/llms/litellm_proxy/skills/skill_search.py
View file

@ -107,7 +107,11 @@ class SkillSearchIndex:
if isinstance(scores, EmbeddingFailed):
return SkillSearchEmbeddingFailed(reason=scores.reason)
ranked: Final = sorted(
(SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)),
(
SkillSearchHit(skill=skill, score=score)
for skill, score in zip(skills, scores, strict=True)
if score is not None
),
key=lambda hit: hit.score,
reverse=True,
)

8
litellm/llms/openrouter/chat/transformation.py
View file

@ -40,6 +40,11 @@ class CacheControlSupportedModels(str, Enum):
ZAI = "z-ai"
def openrouter_model_supports_cache_control(model: str) -> bool:
model_lower: Final = model.lower()
return any(supported_model.value in model_lower for supported_model in CacheControlSupportedModels)
class OpenrouterConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list:
"""
@ -82,8 +87,7 @@ class OpenrouterConfig(OpenAIGPTConfig):
Returns:
bool: True if model supports cache_control (Claude or Gemini models)
"""
model_lower: Final = model.lower()
return any(supported_model.value in model_lower for supported_model in CacheControlSupportedModels)
return openrouter_model_supports_cache_control(model)
def remove_cache_control_flag_from_messages_and_tools(
self,

16
litellm/llms/snowflake/chat/transformation.py
View file

@ -96,7 +96,7 @@ def _decoded_messages(response: _MessagesResponse) -> _AnthropicMessagesResponse
return response.json()
def _is_claude_model(model: str) -> bool:
def is_snowflake_claude_model(model: str) -> bool:
"""Return True if model name (after stripping snowflake/ prefix) is a Claude model."""
name: Final = model.lower().removeprefix("snowflake/")
return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES)
@ -222,7 +222,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
"tools",
"tool_choice",
]
if _is_claude_model(model):
if is_snowflake_claude_model(model):
params.append("thinking")
return params
@ -236,7 +236,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
stream: bool | None = None,
) -> str:
api_base = self._get_api_base(api_base, optional_params)
if _is_claude_model(model):
if is_snowflake_claude_model(model):
return f"{api_base}/cortex/v1/messages"
return f"{api_base}/cortex/v1/chat/completions"
@ -259,7 +259,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
api_key=api_key,
api_base=api_base,
)
if _is_claude_model(model):
if is_snowflake_claude_model(model):
headers["anthropic-version"] = ANTHROPIC_VERSION
return headers
@ -406,7 +406,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
stream: Final[bool] = optional_params.pop("stream", False) or False
extra_body: Final = optional_params.pop("extra_body", {})
if _is_claude_model(model):
if is_snowflake_claude_model(model):
return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body)
return self._transform_request_openai(model, messages, optional_params, stream, extra_body)
@ -422,7 +422,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
litellm_params: dict[str, object], # mutable-ok: BaseConfig signature
headers: dict[str, object], # mutable-ok: BaseConfig signature
) -> dict[str, object]: # mutable-ok: BaseConfig signature
inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages
inlined_messages: Final = (
await async_inline_remote_media(messages) if is_snowflake_claude_model(model) else messages
)
return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers)
def _transform_request_openai(
@ -529,7 +531,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
if _is_claude_model(model):
if is_snowflake_claude_model(model):
return self._transform_response_anthropic(
model, raw_response, model_response, logging_obj, request_data, messages
)

6
litellm/model_prices_and_context_window_backup.json
View file

@ -72571,10 +72571,12 @@
"gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -72585,6 +72587,7 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
@ -72617,10 +72620,12 @@
"vertex_ai/gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -72631,6 +72636,7 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",

3
litellm/models/access_group.py
View file

@ -5,6 +5,7 @@ Canonical definition for ``litellm_accessgrouptable``. Re-exported from
``litellm.proxy._types`` for backwards compatibility.
"""
from collections.abc import Mapping, Sequence
from datetime import datetime
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@ -16,7 +17,9 @@ class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase):
description: str | None = None
access_model_names: list[str] = []
access_mcp_server_ids: list[str] = []
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str] = []
access_skill_names: tuple[str, ...] = ()
assigned_team_ids: list[str] = []
assigned_key_ids: list[str] = []
created_at: datetime | None = None

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

@ -1,3 +1,4 @@
import functools
import re
import secrets
from collections.abc import Awaitable, Callable, Mapping, Sequence
@ -42,6 +43,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
)
from litellm.proxy._types import (
UI_TEAM_ID,
LiteLLM_AccessGroupTable,
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable,
ProxyException,
@ -53,8 +55,10 @@ from litellm.proxy._types import (
user_api_key_has_admin_view,
)
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
AgentAccessGroupCeiling,
CeilingResolver,
resolve_agent_access_group_ceiling,
resolve_managed_agent_ceilings,
)
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy
@ -126,6 +130,22 @@ class UnloadableEntitlementError(Exception):
it places no ceiling; denying there would refuse MCP to every caller during a cold-cache fault."""
@dataclass(frozen=True, slots=True)
class KeyGroupTools:
"""``owned``: tools the key's owned groups grant on a server, ``None`` when they place no bound.
``owns_server``: whether an owned group grants the server at all, so an unbounded owned grant is told
apart from no owned grant. ``attached``: the bound every attached group, owned or not, places there,
applied only when the key holds no grant of its own on the server."""
owned: Sequence[str] | None
owns_server: bool
attached: Sequence[str] | None
_NO_KEY_GROUP_TOOLS: Final = KeyGroupTools(owned=None, owns_server=False, attached=None)
_UNREADABLE_KEY_GROUP_TOOLS: Final = KeyGroupTools(owned=(), owns_server=False, attached=())
def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: list[str] | None = None) -> list[str] | None:
"""Resolve the single MCP server name a cold-start passthrough bypass may
target. Delegates parsing to
@ -251,6 +271,15 @@ def _agent_capped_servers(
)
def _intersect_tool_ceiling(own: Sequence[str] | None, ceiling: Sequence[str] | None) -> Sequence[str] | None:
"""The agent's own tool grants capped by its access-group ceiling, ``None`` only when neither bounds."""
if ceiling is None:
return own
if own is None:
return ceiling
return tuple(set(own) & set(ceiling))
def is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
"""True when this auth is a keyless subject admitted by the gateway session / bridge user
path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``.
@ -2363,16 +2392,177 @@ class MCPRequestHandler:
def _union_tool_grants(
direct: Sequence[str] | None,
via_toolsets: Sequence[str] | None,
via_access_groups: Sequence[str] | None = None,
) -> Sequence[str] | None:
"""Union of one level's direct tool grants and its toolset-granted tools on one server,
``None`` when neither source restricts (allow-all from this level). A direct grant
containing ``MCP_ALL_TOOLS_WILDCARD`` makes the level unrestricted, so it returns
``None`` whatever the toolsets name."""
"""Union of one level's direct tool grants, its toolset-granted tools and its unified
access-group tools on one server, ``None`` when no source restricts (allow-all from this
level). A direct grant containing ``MCP_ALL_TOOLS_WILDCARD`` makes the level unrestricted,
so it returns ``None`` whatever the toolsets or groups name."""
if direct is not None and MCP_ALL_TOOLS_WILDCARD in direct:
return None
if direct is None and via_toolsets is None:
if direct is None and via_toolsets is None and via_access_groups is None:
return None
return tuple({*(direct or ()), *(via_toolsets or ())})
return tuple({*(direct or ()), *(via_toolsets or ()), *(via_access_groups or ())})
@staticmethod
def access_group_tools_for_server(
groups: Sequence[LiteLLM_AccessGroupTable], server_id: str
) -> Sequence[str] | None:
"""Tools the unified access groups grant on ``server_id``: the union of each granting group's
``mcp_tool_permissions`` entry for it. A granting group with no entry grants every tool, so
``None`` wins. A group that does not grant the server contributes nothing, and an explicit
``[]`` entry on a granting group grants no tools. ``None`` when no group grants the server."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
def granted_tools(group: LiteLLM_AccessGroupTable) -> Sequence[str] | None:
if not MCPRequestHandler._group_grants_server(group, server_id):
return None
permissions: Final = {
server: list(allowed) for server, allowed in (group.mcp_tool_permissions or {}).items()
}
tools: Final = global_mcp_server_manager.expand_tool_permissions(permissions)
return tools.get(server_id, (MCP_ALL_TOOLS_WILDCARD,))
grants: Final = tuple(granted_tools(group) for group in groups)
named: Final = tuple(grant for grant in grants if grant is not None)
if not named or any(MCP_ALL_TOOLS_WILDCARD in grant for grant in named):
return None
return tuple({tool for grant in named for tool in grant}) # comprehension-ok: flatten
@staticmethod
def _group_grants_server(group: LiteLLM_AccessGroupTable, server_id: str) -> bool:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
return server_id in global_mcp_server_manager.expand_permission_list(list(group.access_mcp_server_ids))
@staticmethod
def _key_tool_scope(
direct: Sequence[str] | None,
via_toolsets: Sequence[str] | None,
groups: KeyGroupTools,
grants_server: bool,
) -> Sequence[str] | None:
"""The key's own tool scope on a server when it holds any grant there (a direct server grant with
or without a tool entry, a toolset or an owned group), so an unrestricted grant stays unrestricted.
Only a key with no grant of its own takes the bound its attached groups place."""
if grants_server or direct is not None or via_toolsets is not None or groups.owns_server:
return MCPRequestHandler._union_tool_grants(direct, via_toolsets, groups.owned)
return groups.attached
@staticmethod
def _group_authorizes_key(group: LiteLLM_AccessGroupTable, team_id: str | None, token: str | None) -> bool:
"""Mirror of ``get_authorized_resources_from_key_access_groups``: a group a key merely names
binds only when the group lists the key's team or the key itself as an owner."""
return bool(team_id and team_id in group.assigned_team_ids) or bool(token and token in group.assigned_key_ids)
@staticmethod
def _loaded_group_tools_for_server(
groups: Sequence[LiteLLM_AccessGroupTable], requested: Sequence[str], server_id: str
) -> Sequence[str] | None:
"""Tool grants of the agent groups that loaded. A named agent group that did not load exists but
could not be read (a group that does not exist never reaches the ceiling), so the level grants no
tools on ``server_id`` rather than shedding a restriction while the fault lasts."""
if {group.access_group_id for group in groups} != set(requested):
return ()
return MCPRequestHandler.access_group_tools_for_server(groups, server_id)
@staticmethod
async def _key_direct_servers(
key_obj_perm: LiteLLM_ObjectPermissionTable | None, *, requires_fresh_policy: bool
) -> AbstractSet[str]:
"""Servers the key's object permission grants on its own: ``mcp_servers`` by id, name or alias and
legacy ``mcp_access_groups`` by the tag the server carries, exactly as server admission resolves
them. A tag never doubles as a server name, so a tag that collides with one grants nothing by itself."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
if key_obj_perm is None:
return frozenset()
try:
tagged: Final = (
await MCPRequestHandler.get_mcp_servers_from_access_groups(
key_obj_perm.mcp_access_groups, requires_fresh_policy=requires_fresh_policy
)
if key_obj_perm.mcp_access_groups
else ()
)
except Exception as e: # noqa: BLE001 # a key naming tags we cannot resolve must deny, not widen
raise UnloadableEntitlementError(f"key access group tags unreadable: {e}") from e
return frozenset(global_mcp_server_manager.expand_permission_list(key_obj_perm.mcp_servers or [])) | frozenset(
tagged
)
@staticmethod
async def _key_access_group_tools_for_server(user_api_key_auth: UserAPIKeyAuth, server_id: str) -> KeyGroupTools:
"""Tool bounds the key's attached groups place on ``server_id``. Attaching a group to a key admits
the group's servers (``_get_key_access_group_mcp_server_extras``), so an attached group the key
does not own still carries its own tool permissions onto the servers it admits, it just cannot
add tools to a grant the key already holds there. Only owned groups are grant sources. A named
group that does not exist grants and bounds nothing, one that exists but cannot be read denies."""
from litellm.proxy.auth.auth_checks import get_access_group_tables
if not user_api_key_auth.access_group_ids:
return _NO_KEY_GROUP_TOOLS
try:
groups: Final = await get_access_group_tables(
access_group_ids=user_api_key_auth.access_group_ids,
check_db_only=bool(user_api_key_auth.requires_fresh_policy),
unreadable="raise",
)
except Exception as e: # noqa: BLE001 # a key naming groups we cannot read must deny, not widen
if user_api_key_auth.requires_fresh_policy:
raise UnloadableEntitlementError(f"key access groups unreadable: {e}") from e
return _UNREADABLE_KEY_GROUP_TOOLS
owned: Final = tuple(
group
for group in groups
if MCPRequestHandler._group_authorizes_key(group, user_api_key_auth.team_id, user_api_key_auth.token)
)
return KeyGroupTools(
owned=MCPRequestHandler.access_group_tools_for_server(owned, server_id),
owns_server=any(MCPRequestHandler._group_grants_server(group, server_id) for group in owned),
attached=MCPRequestHandler.access_group_tools_for_server(groups, server_id),
)
@staticmethod
async def _team_access_group_tools_for_server(
user_api_key_auth: UserAPIKeyAuth, server_id: str
) -> Sequence[str] | None:
from litellm.proxy.auth.auth_checks import get_access_group_tables, get_team_object
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if not user_api_key_auth.team_id or user_api_key_auth.team_id == UI_TEAM_ID or prisma_client is None:
return None
try:
team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object(
team_id=user_api_key_auth.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_db_only=bool(user_api_key_auth.requires_fresh_policy),
)
except Exception as e: # noqa: BLE001 # a team we cannot read must deny, not widen
raise UnloadableEntitlementError(f"team unreadable: {e}") from e
if not team_obj.access_group_ids:
return None
try:
groups: Final = await get_access_group_tables(
access_group_ids=team_obj.access_group_ids,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
check_db_only=bool(user_api_key_auth.requires_fresh_policy),
unreadable="raise",
)
except Exception as e: # noqa: BLE001 # team groups we cannot read grant no tools, they never widen
if user_api_key_auth.requires_fresh_policy:
raise UnloadableEntitlementError(f"team access groups unreadable: {e}") from e
return ()
return MCPRequestHandler.access_group_tools_for_server(groups, server_id)
@staticmethod
async def _key_object_permission_hydrated(
@ -2493,7 +2683,17 @@ class MCPRequestHandler:
else None
)
key_tools: Final = _as_list(MCPRequestHandler._union_tool_grants(key_direct_tools, key_toolset_tools))
key_group_tools: Final = await MCPRequestHandler._key_access_group_tools_for_server(
user_api_key_auth, server_id
)
key_direct_servers: Final = await MCPRequestHandler._key_direct_servers(
key_obj_perm, requires_fresh_policy=bool(user_api_key_auth.requires_fresh_policy)
)
key_tools: Final = _as_list(
MCPRequestHandler._key_tool_scope(
key_direct_tools, key_toolset_tools, key_group_tools, server_id in key_direct_servers
)
)
team_direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id)
if team_obj_perm
@ -2505,11 +2705,17 @@ class MCPRequestHandler:
team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(
team_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy
)
team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools)
team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server(
user_api_key_auth, server_id
)
team_tools: Final = MCPRequestHandler._union_tool_grants(
team_direct_tools, team_toolset_tools, team_group_tools
)
# Apply same inheritance logic as get_allowed_mcp_servers
if team_tools:
if key_tools:
# Apply same inheritance logic as get_allowed_mcp_servers. ``is not None`` so an explicit
# empty allowlist stays a deny instead of reading as "no restriction".
if team_tools is not None:
if key_tools is not None:
# Both have restrictions → intersection
allowed_tools = list(set(team_tools) & set(key_tools))
else:
@ -3525,6 +3731,9 @@ class MCPRequestHandler:
team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(
team_obj_perm, server_id, requires_fresh_policy=caller_auth.requires_fresh_policy
)
team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server(
caller_auth, server_id
)
except Exception as e: # noqa: BLE001 # an unresolved caller team must deny, not widen
verbose_logger.warning(
"MCP agent caller team tool ceiling unresolvable, denying tools on %r: %s", server_id, e
@ -3535,7 +3744,9 @@ class MCPRequestHandler:
if team_obj_perm
else None
)
team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools)
team_tools: Final = MCPRequestHandler._union_tool_grants(
team_direct_tools, team_toolset_tools, team_group_tools
)
team_capped: Final = (
allowed_tools
if team_tools is None
@ -3729,12 +3940,50 @@ class MCPRequestHandler:
global_mcp_server_manager,
)
ceilings: Final = await MCPRequestHandler._agent_access_group_ceilings(user_api_key_auth, resolve_ceiling)
if not ceilings:
return None
expanded: Final[tuple[frozenset[str], ...]] = tuple(
frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids)))
for ceiling in ceilings
)
return functools.reduce(lambda acc, servers: acc & servers, expanded[1:], expanded[0])
@staticmethod
async def _get_agent_access_group_tool_ceiling(
user_api_key_auth: UserAPIKeyAuth,
server_id: str,
resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
) -> Sequence[str] | None:
"""Ceiling the agent's attached unified access groups place on ``server_id`` tools, the tools axis
twin of ``_get_agent_access_group_server_ceiling``. None when the agent has no groups attached or
no attached group bounds tools on the server. The agent's own grants are intersected with it."""
ceilings: Final = await MCPRequestHandler._agent_access_group_ceilings(user_api_key_auth, resolve_ceiling)
if not ceilings:
return None
bounds: Final = tuple(
MCPRequestHandler._loaded_group_tools_for_server(ceiling.access_groups, ceiling.access_group_ids, server_id)
for ceiling in ceilings
)
return functools.reduce(_intersect_tool_ceiling, bounds[1:], bounds[0])
@staticmethod
async def _agent_access_group_ceilings(
user_api_key_auth: UserAPIKeyAuth,
resolve_ceiling: CeilingResolver,
) -> tuple[AgentAccessGroupCeiling, ...]:
"""A managed agent is bounded by the groups on its admitted policy, the same source server admission
reads, never by the registry row that policy may have superseded. Empty means no groups attached."""
if not user_api_key_auth.agent_id:
return None
ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id)
if ceiling is None:
return None
return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids)))
return ()
managed: Final = managed_agent_policy(user_api_key_auth)
try:
if managed is not None:
return await resolve_managed_agent_ceilings(managed)
ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id)
except Exception as e: # noqa: BLE001 # agent groups we cannot read must deny, not widen
raise UnloadableEntitlementError(f"agent access groups unreadable: {e}") from e
return (ceiling,) if ceiling is not None else ()
@staticmethod
async def get_agent_tool_permissions_for_server(
@ -3764,14 +4013,17 @@ class MCPRequestHandler:
if agent_object_permission is not None
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
)
if obj_perm is None:
return None
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
group_tools: Final = await MCPRequestHandler._get_agent_access_group_tool_ceiling(
user_api_key_auth, server_id
)
if obj_perm is None:
return _as_list(group_tools)
direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
if obj_perm.mcp_tool_permissions
@ -3780,7 +4032,8 @@ class MCPRequestHandler:
toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(
obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy
)
agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
own_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
agent_tools: Final = _intersect_tool_ceiling(own_tools, group_tools)
return list(agent_tools) if agent_tools is not None else None
except Exception as e:
if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError):

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

@ -258,9 +258,11 @@ def _split_core_tools(tools: Sequence[Tool], core_tools: Sequence[str]) -> tuple
def _top_hits(
tools: Sequence[Tool], scores: Sequence[float], minimum: float, limit: int
tools: Sequence[Tool], scores: Sequence[float | None], minimum: float, limit: int
) -> tuple[tuple[float, Tool], ...]:
hits: Final = ((score, tool) for score, tool in zip(scores, tools, strict=True) if score >= minimum)
hits: Final = (
(score, tool) for score, tool in zip(scores, tools, strict=True) if score is not None and score >= minimum
)
return tuple(sorted(hits, key=lambda hit: hit[0], reverse=True)[:limit])

89
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -648,6 +648,20 @@
],
"title": "Access Model Names"
},
"access_skill_names": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"title": "Access Skill Names"
},
"assigned_key_ids": {
"anyOf": [
{
@ -686,6 +700,23 @@
}
],
"title": "Description"
},
"mcp_tool_permissions": {
"anyOf": [
{
"additionalProperties": {
"items": {
"type": "string"
},
"type": "array"
},
"type": "object"
},
{
"type": "null"
}
],
"title": "Mcp Tool Permissions"
}
},
"required": [
@ -812,6 +843,14 @@
"title": "Access Model Names",
"type": "array"
},
"access_skill_names": {
"default": [],
"items": {
"type": "string"
},
"title": "Access Skill Names",
"type": "array"
},
"assigned_key_ids": {
"items": {
"type": "string"
@ -867,6 +906,23 @@
],
"title": "Description"
},
"mcp_tool_permissions": {
"anyOf": [
{
"additionalProperties": {
"items": {
"type": "string"
},
"type": "array"
},
"type": "object"
},
{
"type": "null"
}
],
"title": "Mcp Tool Permissions"
},
"updated_at": {
"format": "date-time",
"title": "Updated At",
@ -957,6 +1013,20 @@
],
"title": "Access Model Names"
},
"access_skill_names": {
"anyOf": [
{
"items": {
"type": "string"
},
"type": "array"
},
{
"type": "null"
}
],
"title": "Access Skill Names"
},
"assigned_key_ids": {
"anyOf": [
{
@ -995,6 +1065,23 @@
}
],
"title": "Description"
},
"mcp_tool_permissions": {
"anyOf": [
{
"additionalProperties": {
"items": {
"type": "string"
},
"type": "array"
},
"type": "object"
},
{
"type": "null"
}
],
"title": "Mcp Tool Permissions"
}
},
"title": "AccessGroupUpdateRequest",
@ -7644,7 +7731,7 @@
"paths": {
"/claude-code/marketplace.json": {
"get": {
"description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add <url>\n- claude plugin install <name>@<marketplace>\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```",
"description": "Serve marketplace.json for Claude Code plugin discovery.\n\nThis endpoint is accessed by Claude Code CLI when users run:\n- claude plugin marketplace add <url>\n- claude plugin install <name>@<marketplace>\n\nWithout `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`\nthe key is authenticated and the catalog also holds the disabled plugins granted\nto it through `object_permission.skills` or an access group on the key or its team.\n\nReturns:\n Marketplace catalog with list of available plugins and their git sources.\n\nExample:\n ```bash\n claude plugin marketplace add http://localhost:4000/claude-code/marketplace.json\n claude plugin marketplace add \"http://localhost:4000/claude-code/marketplace.json?key=sk-...\"\n claude plugin install my-plugin@litellm\n ```",
"operationId": "get_marketplace_claude_code_marketplace_json_get",
"parameters": [
{

4
litellm/proxy/_types.py
View file

@ -4260,6 +4260,7 @@ class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase):
models: tuple[str, ...]
mcp_server_ids: tuple[str, ...] = ()
agent_ids: tuple[str, ...] = ()
skill_names: tuple[str, ...] = ()
class TeamInfoMember(Member):
@ -4297,6 +4298,9 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
access_group_models: list[str] | None = None
access_group_mcp_server_ids: list[str] | None = None
access_group_agent_ids: list[str] | None = None
access_group_skill_names: tuple[str, ...] | None = None
# Listed group ids with no row; skill grants are suspended while this is non-empty.
unresolved_access_group_ids: tuple[str, ...] | None = None
access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None
# Parent org's model ceiling, reported only to callers who can manage the team.
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.

6
litellm/proxy/agent_endpoints/agent_search.py
View file

@ -116,7 +116,11 @@ class AgentSearchIndex:
if isinstance(scores, EmbeddingFailed):
return AgentSearchEmbeddingFailed(reason=scores.reason)
ranked: Final = sorted(
(AgentSearchHit(agent=agent, score=score) for agent, score in zip(agents, scores, strict=True)),
(
AgentSearchHit(agent=agent, score=score)
for agent, score in zip(agents, scores, strict=True)
if score is not None
),
key=lambda hit: hit.score,
reverse=True,
)

42
litellm/proxy/agent_endpoints/auth/agent_access_groups.py
View file

@ -25,11 +25,19 @@ class AgentAccessGroupCeiling:
models: frozenset[str]
mcp_server_ids: frozenset[str]
agent_ids: frozenset[str]
access_groups: tuple[LiteLLM_AccessGroupTable, ...] = ()
CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]]
class _MissingGroup:
__slots__ = ()
_MISSING: Final = _MissingGroup()
async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
@ -38,7 +46,7 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup:
from litellm.proxy.auth.auth_checks import get_access_object
from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
@ -52,6 +60,8 @@ async def _load_access_group(access_group_id: str, *, check_db_only: bool = Fals
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
)
except AccessGroupNotFoundError:
raise
except HTTPException as e:
if check_db_only:
raise
@ -68,23 +78,35 @@ async def resolve_agent_access_group_ceiling(
*,
check_db_only: bool = False,
) -> AgentAccessGroupCeiling | None:
"""``None`` when the agent has no access groups attached, so nothing is capped."""
access_group_ids: Final = await load_access_group_ids(agent_id)
if not access_group_ids:
"""``None`` when the agent has no access groups attached, so nothing is capped. A named group that
provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group
that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once
any group is named, so an agent whose every group is missing is still capped to nothing."""
named_ids: Final = await load_access_group_ids(agent_id)
if not named_ids:
return None
loaded: Final = await asyncio.gather(
*(
_load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id)
for group_id in access_group_ids
)
from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup:
try:
return await (
_load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id)
)
except AccessGroupNotFoundError:
return _MISSING
results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids))
access_group_ids: Final = tuple(
group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup)
)
groups: Final = tuple(group for group in loaded if group is not None)
groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable))
return AgentAccessGroupCeiling(
access_group_ids=access_group_ids,
models=frozenset(model for group in groups for model in group.access_model_names),
mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids),
agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids),
access_groups=groups,
)

10
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py
View file

@ -33,7 +33,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, ProxyException, UserAPIKeyAuth
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import (
SkillVisibility,
skill_visibility,
resolve_skill_visibility,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
@ -122,7 +122,7 @@ async def get_marketplace(request: Request, key: str | None = None):
Without `key` the catalog holds the enabled (public) plugins. With `?key=sk-...`
the key is authenticated and the catalog also holds the disabled plugins granted
to it through `object_permission.skills` on the key or its team.
to it through `object_permission.skills` or an access group on the key or its team.
Returns:
Marketplace catalog with list of available plugins and their git sources.
@ -140,7 +140,7 @@ async def get_marketplace(request: Request, key: str | None = None):
caller: Final[UserAPIKeyAuth | None] = (
await user_api_key_auth(request=request, api_key=f"Bearer {key}") if key else None
)
visibility: Final[SkillVisibility] = skill_visibility(caller)
visibility: Final[SkillVisibility] = await resolve_skill_visibility(caller)
plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many(
where=visibility.where()
)
@ -464,7 +464,7 @@ async def list_plugins(
try:
prisma_client: Final = await get_prisma_client()
visibility: Final[SkillVisibility] = skill_visibility(user_api_key_dict)
visibility: Final[SkillVisibility] = await resolve_skill_visibility(user_api_key_dict)
plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many(
where={"enabled": True} if enabled_only else visibility.where()
)
@ -545,7 +545,7 @@ async def get_plugin(
detail={"error": f"Plugin '{plugin_name}' not found"},
)
if not skill_visibility(user_api_key_dict).allows(plugin):
if not (await resolve_skill_visibility(user_api_key_dict)).allows(plugin):
raise HTTPException(
status_code=403,
detail={"error": f"Plugin '{plugin_name}' is not granted to this key"},

116
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_skill_access.py
View file

@ -1,18 +1,22 @@
"""
Claude Code marketplace visibility: enabled plugins are public, disabled plugins
are private and resolve only for proxy admins or keys granted them via
``object_permission.skills``.
``object_permission.skills`` or via ``access_skill_names`` on an access group
attached to the key or its team.
"""
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
from litellm.proxy._types import UI_TEAM_ID, LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
if TYPE_CHECKING:
from prisma.types import LiteLLM_ClaudeCodePluginTableWhereInput
from litellm.proxy._types import LiteLLM_TeamTable
class _SkillRecord(Protocol):
name: str
@ -23,14 +27,28 @@ def _skills_of(permission: LiteLLM_ObjectPermissionTable | None) -> frozenset[st
return frozenset(permission.skills or ()) if permission is not None else frozenset()
def granted_skills(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]:
@dataclass(frozen=True, slots=True)
class AccessGroupSkills:
"""Plugin names inherited from access groups, split by the scope the group is attached to."""
key: frozenset[str] = frozenset()
team: frozenset[str] = frozenset()
unresolved: bool = False
"""A listed group returned no row, so the policy is incomplete and only public plugins show."""
NO_ACCESS_GROUP_SKILLS: Final = AccessGroupSkills()
def granted_skills(
user_api_key_dict: UserAPIKeyAuth, inherited: AccessGroupSkills = NO_ACCESS_GROUP_SKILLS
) -> frozenset[str]:
"""Key grant intersected with the team grant when both are non-empty; either alone applies as is.
An empty list is the Prisma column default for every object-permission row, so it means
"no private grants configured here" and defers to the other scope, same as the agents check.
An empty list is the Prisma column default, so it means "no private grants configured here".
"""
key_skills: Final = _skills_of(user_api_key_dict.object_permission)
team_skills: Final = _skills_of(user_api_key_dict.team_object_permission)
key_skills: Final = _skills_of(user_api_key_dict.object_permission) | inherited.key
team_skills: Final = _skills_of(user_api_key_dict.team_object_permission) | inherited.team
match (bool(key_skills), bool(team_skills)):
case (True, True):
return key_skills & team_skills
@ -59,9 +77,87 @@ class SkillVisibility:
PUBLIC_ONLY: Final = SkillVisibility(granted=frozenset(), sees_private=False)
def skill_visibility(user_api_key_dict: UserAPIKeyAuth | None) -> SkillVisibility:
def skill_visibility(
user_api_key_dict: UserAPIKeyAuth | None, inherited: AccessGroupSkills = NO_ACCESS_GROUP_SKILLS
) -> SkillVisibility:
if user_api_key_dict is None:
return PUBLIC_ONLY
if is_proxy_admin(user_api_key_dict):
return SkillVisibility(granted=frozenset(), sees_private=True)
return SkillVisibility(granted=granted_skills(user_api_key_dict), sees_private=False)
if inherited.unresolved:
return PUBLIC_ONLY
return SkillVisibility(granted=granted_skills(user_api_key_dict, inherited), sees_private=False)
async def _team_object(user_api_key_dict: UserAPIKeyAuth, check_db_only: bool) -> "LiteLLM_TeamTable | None":
from litellm.proxy.auth.auth_checks import TeamNotFoundError, get_team_object
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
team_id: Final = user_api_key_dict.team_id
if not team_id or team_id == UI_TEAM_ID or prisma_client is None:
return None
try:
team_obj: Final = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
)
except TeamNotFoundError:
return None
return None if team_obj.blocked else team_obj
async def access_group_skills(user_api_key_dict: UserAPIKeyAuth) -> AccessGroupSkills:
"""Resolve ``access_skill_names`` of the groups attached to the key and to its team.
Key-side groups count only when the group lists the key or its team as an owner, same as
the model and MCP checks, so a key cannot claim a group it was never assigned to.
"""
from litellm.proxy.auth.auth_checks import (
AccessGroupNotFoundError,
get_authorized_resources_from_key_access_groups,
get_skill_names_from_access_groups,
)
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
return NO_ACCESS_GROUP_SKILLS
check_db_only: Final = bool(user_api_key_dict.requires_fresh_policy)
team_obj: Final = await _team_object(user_api_key_dict, check_db_only)
try:
key_skills: Final = frozenset(
await get_authorized_resources_from_key_access_groups(
valid_token=user_api_key_dict,
team_object=team_obj,
resource_field="access_skill_names",
check_db_only=check_db_only,
strict_errors=True,
)
)
team_skills: Final = await get_skill_names_from_access_groups(
access_group_ids=tuple(team_obj.access_group_ids or ()) if team_obj is not None else (),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
strict_errors=True,
)
return AccessGroupSkills(key=key_skills, team=team_skills)
except AccessGroupNotFoundError:
return AccessGroupSkills(unresolved=True)
AccessGroupSkillsLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[AccessGroupSkills]]
async def resolve_skill_visibility(
user_api_key_dict: UserAPIKeyAuth | None,
load_access_group_skills: AccessGroupSkillsLoader = access_group_skills,
) -> SkillVisibility:
"""``skill_visibility`` with the access-group grants of the key and its team folded in."""
if user_api_key_dict is None or is_proxy_admin(user_api_key_dict):
return skill_visibility(user_api_key_dict)
return skill_visibility(user_api_key_dict, await load_access_group_skills(user_api_key_dict))

159
litellm/proxy/auth/auth_checks.py
View file

@ -3411,6 +3411,19 @@ class TeamNotFoundError(HTTPException):
super().__init__(status_code=404, detail=detail)
class AccessGroupNotFoundError(HTTPException):
"""The access group row is provably absent, as opposed to merely unreadable. ``get_access_object``
reports every failure as a 404, so a group that was never created and a database that would not
answer look the same to its callers. Policy code that must not treat a degraded read as a definitive
answer keys on this subclass. It stays a 404 with the same detail, so every other caller is unaffected."""
def __init__(self, access_group_id: str) -> None:
super().__init__(
status_code=404,
detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."},
)
@log_db_metrics
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False
@ -3667,10 +3680,7 @@ async def get_access_object(
).find_unique(where={"access_group_id": access_group_id})
if response is None:
raise HTTPException(
status_code=404,
detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."},
)
raise AccessGroupNotFoundError(access_group_id=access_group_id)
_response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict())
@ -4532,67 +4542,99 @@ async def get_org_object_for_request(
return org
async def _get_resources_from_access_groups(
async def get_access_group_tables(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
prisma_client: DatabaseClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
check_db_only: bool = False,
) -> list[str]:
"""
Fetch access groups by their IDs (from cache or DB) and collect
the specified resource field across all of them.
Args:
access_group_ids: List of access group IDs to fetch
resource_field: Which resource list to extract from each access group
- "access_model_names": model names (for model access checks)
- "access_mcp_server_ids": MCP server IDs (for MCP access checks)
- "access_agent_ids": agent IDs (for agent access checks)
prisma_client: Optional PrismaClient (lazy-imported from proxy_server if None)
user_api_key_cache: Optional DualCache (lazy-imported from proxy_server if None)
proxy_logging_obj: Optional ProxyLogging (lazy-imported from proxy_server if None)
Returns:
Deduplicated list of resource identifiers from all resolved access groups.
"""
*,
unreadable: Literal["skip", "raise"] = "skip",
missing: Literal["skip", "raise"] = "skip",
) -> tuple[LiteLLM_AccessGroupTable, ...]:
"""Load unified access groups by id (cache first, then DB). A group that provably does not exist is
skipped, it grants nothing, unless ``missing`` is ``"raise"`` and the caller wants to see
``AccessGroupNotFoundError`` itself. A group that exists but cannot be read is skipped too unless
``check_db_only`` is set or ``unreadable`` is ``"raise"``, in which case the failure propagates so a
caller enforcing policy can deny instead of shedding a restriction it could not read."""
if not access_group_ids:
return []
return ()
# Lazy import to avoid circular imports
if prisma_client is None or user_api_key_cache is None:
from litellm.proxy.proxy_server import prisma_client as _prisma_client
from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
prisma_client = prisma_client or _prisma_client
user_api_key_cache = user_api_key_cache or _user_api_key_cache
proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj
prisma_client = prisma_client or _prisma_client # rebind-ok: lazy proxy_server fallback
user_api_key_cache = user_api_key_cache or _user_api_key_cache # rebind-ok: lazy proxy_server fallback
proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj # rebind-ok: lazy proxy_server fallback
if user_api_key_cache is None:
return []
return ()
resources: Final[list[str]] = []
for ag_id in access_group_ids:
async def load(ag_id: str) -> LiteLLM_AccessGroupTable | None:
try:
ag = await get_access_object(
return await get_access_object(
access_group_id=ag_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
)
resources.extend(getattr(ag, resource_field, []))
except Exception:
if check_db_only:
except AccessGroupNotFoundError:
if missing == "raise":
raise
verbose_proxy_logger.debug(
"Could not fetch access group %s for resource field %s",
ag_id,
resource_field,
)
return list(set(resources))
verbose_proxy_logger.debug("One of the requested access groups does not exist, it grants nothing")
return None
except Exception:
if check_db_only or unreadable == "raise":
raise
verbose_proxy_logger.debug("Could not fetch one of the requested access groups")
return None
loaded: Final = [await load(ag_id) for ag_id in access_group_ids]
return tuple(ag for ag in loaded if ag is not None)
async def _get_resources_from_access_groups(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names"],
prisma_client: DatabaseClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
check_db_only: bool = False,
strict_errors: bool = False,
) -> list[str]:
"""Deduplicated ``resource_field`` values across the resolved access groups. With ``strict_errors`` a
missing or unreadable group raises instead of being skipped, so policy code never loses a ceiling."""
strictness: Final[Literal["skip", "raise"]] = "raise" if strict_errors else "skip"
groups: Final = await get_access_group_tables(
access_group_ids=access_group_ids,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
unreadable=strictness,
missing=strictness,
)
return list(
{resource for ag in groups for resource in _access_group_field(ag, resource_field)} # comprehension-ok: flatten
)
def _access_group_field(
group: LiteLLM_AccessGroupTable,
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names"],
) -> list[str]:
match resource_field:
case "access_model_names":
return group.access_model_names
case "access_mcp_server_ids":
return group.access_mcp_server_ids
case "access_agent_ids":
return group.access_agent_ids
case "access_skill_names":
return list(group.access_skill_names or ())
async def get_models_from_access_groups(
@ -4665,6 +4707,28 @@ async def get_agent_ids_from_access_groups(
_get_agent_ids_from_access_groups: Final = get_agent_ids_from_access_groups
async def get_skill_names_from_access_groups(
access_group_ids: Sequence[str],
prisma_client: PrismaClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
check_db_only: bool = False,
strict_errors: bool = False,
) -> frozenset[str]:
"""Claude Code plugin names granted by unified access groups, matched by plugin name."""
return frozenset(
await _get_resources_from_access_groups(
access_group_ids=access_group_ids,
resource_field="access_skill_names",
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
strict_errors=strict_errors,
)
)
def _resolve_all_team_model_sentinel_for_auth_check(
models: list[str],
llm_router: Router | None,
@ -5584,7 +5648,9 @@ async def can_team_access_model(
async def get_authorized_resources_from_key_access_groups(
valid_token: UserAPIKeyAuth | None,
team_object: LiteLLM_TeamTable | None,
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids", "access_skill_names"],
check_db_only: bool = False,
strict_errors: bool = False,
) -> list[str]:
"""
For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row
@ -5618,8 +5684,15 @@ async def get_authorized_resources_from_key_access_groups(
prisma_client=_prisma_client,
user_api_key_cache=_user_api_key_cache,
proxy_logging_obj=_proxy_logging_obj,
check_db_only=check_db_only,
)
except AccessGroupNotFoundError:
if check_db_only or strict_errors:
raise
continue
except Exception:
if check_db_only or strict_errors:
raise
continue
team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or []))
key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or []))

167
litellm/proxy/common_utils/semantic_text_index.py
View file

@ -3,7 +3,8 @@
from __future__ import annotations
import math
from collections.abc import Awaitable, Mapping, Sequence
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain, islice
from types import MappingProxyType
@ -13,7 +14,7 @@ from fastapi import HTTPException
from openai import OpenAIError
from pydantic import ConfigDict
from litellm.exceptions import BudgetExceededError
from litellm.exceptions import BudgetExceededError, ContextWindowExceededError
from litellm.types.llms.base import LiteLLMBaseModel
if TYPE_CHECKING:
@ -25,6 +26,7 @@ Vector: TypeAlias = tuple[float, ...]
DEFAULT_MAX_CACHED_VECTORS: Final = 5000
"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first."""
DEFAULT_TOO_LONG_TTL_SECONDS: Final = 3600.0
class Embedder(Protocol):
@ -36,6 +38,18 @@ class EmbeddingFailed:
reason: str
@dataclass(frozen=True, slots=True)
class _InputTooLong:
reason: str
@dataclass(frozen=True, slots=True)
class _TextVectors:
vectors: Mapping[str, Vector]
too_long: frozenset[str]
failure: EmbeddingFailed | None = None
class _EmbeddingItem(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
@ -102,11 +116,13 @@ def router_embedder(
_CacheKey: TypeAlias = tuple[str, str]
async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed:
async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed | _InputTooLong:
try:
vectors: Final = tuple(await embed(texts))
except HTTPException:
raise
except ContextWindowExceededError as exc:
return _InputTooLong(reason=f"embedding the search query failed: {exc}")
except (OpenAIError, ValueError, BudgetExceededError) as exc:
return EmbeddingFailed(reason=f"embedding the search query failed: {exc}")
if len(vectors) != len(texts):
@ -114,32 +130,97 @@ async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...
return vectors
async def _embed_texts(embed: Embedder, texts: Sequence[str]) -> _TextVectors:
if not texts:
return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset())
embedded: Final = await _embed_all(embed, texts)
if isinstance(embedded, EmbeddingFailed):
return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset(), failure=embedded)
if not isinstance(embedded, _InputTooLong):
return _TextVectors(vectors=MappingProxyType(dict(zip(texts, embedded, strict=True))), too_long=frozenset())
if len(texts) == 1:
return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset({texts[0]}))
middle: Final = len(texts) // 2
left: Final = await _embed_texts(embed, texts[:middle])
if left.failure is not None:
return left
right: Final = await _embed_texts(embed, texts[middle:])
return _TextVectors(
vectors=MappingProxyType({**left.vectors, **right.vectors}),
too_long=left.too_long | right.too_long,
failure=right.failure,
)
async def _embed_query_with(
embed: Embedder, query: str, texts: Sequence[str]
) -> tuple[Vector, _TextVectors] | EmbeddingFailed:
if not texts:
embedded_query: Final = await _embed_all(embed, (query,))
if isinstance(embedded_query, EmbeddingFailed):
return embedded_query
if isinstance(embedded_query, _InputTooLong):
return EmbeddingFailed(reason=embedded_query.reason)
return embedded_query[0], _TextVectors(vectors=MappingProxyType({}), too_long=frozenset())
embedded: Final = await _embed_all(embed, (query, *texts))
if isinstance(embedded, EmbeddingFailed):
return embedded
if not isinstance(embedded, _InputTooLong):
return embedded[0], _TextVectors(
vectors=MappingProxyType(dict(zip(texts, embedded[1:], strict=True))), too_long=frozenset()
)
query_embedding: Final = await _embed_all(embed, (query,))
if isinstance(query_embedding, EmbeddingFailed):
return query_embedding
if isinstance(query_embedding, _InputTooLong):
return EmbeddingFailed(reason=query_embedding.reason)
vectors: Final = await _embed_texts(embed, texts)
return query_embedding[0], vectors
@dataclass(frozen=True, slots=True)
class _Embedded:
query_vector: Vector
vectors: Mapping[str, Vector]
too_long: frozenset[str]
failure: EmbeddingFailed | None
def _same_dimension(query_vector: Vector, vectors: Mapping[str, Vector], texts: Sequence[str]) -> bool:
return all(len(vectors[text]) == len(query_vector) for text in texts)
return all(len(vectors[text]) == len(query_vector) for text in texts if text in vectors)
async def _embed_query_and_texts(
embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector]
embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector], too_long: frozenset[str]
) -> _Embedded | EmbeddingFailed:
missing: Final = tuple(dict.fromkeys(text for text in texts if text not in cached))
embedded: Final = await _embed_all(embed, (query, *missing))
if isinstance(embedded, EmbeddingFailed):
pending: Final = (text for text in texts if text not in cached and text not in too_long)
missing: Final = tuple(dict.fromkeys(pending))
query_embedding: Final = await _embed_query_with(embed, query, missing)
if isinstance(query_embedding, EmbeddingFailed):
return query_embedding
query_vector, new_text_vectors = query_embedding
vectors: Final = MappingProxyType(dict(chain(cached.items(), new_text_vectors.vectors.items())))
embedded: Final = _Embedded(
query_vector=query_vector,
vectors=vectors,
too_long=new_text_vectors.too_long,
failure=new_text_vectors.failure,
)
if embedded.failure is not None:
return embedded
vectors: Final = MappingProxyType(dict(chain(cached.items(), zip(missing, embedded[1:], strict=True))))
if _same_dimension(embedded[0], vectors, texts):
return _Embedded(query_vector=embedded[0], vectors=vectors)
unique: Final = tuple(dict.fromkeys(texts))
reembedded: Final = await _embed_all(embed, (query, *unique))
if _same_dimension(query_vector, vectors, texts):
return embedded
unique: Final = tuple(dict.fromkeys(text for text in texts if text in vectors))
reembedded: Final = await _embed_query_with(embed, query, unique)
if isinstance(reembedded, EmbeddingFailed):
return reembedded
reembedded_query, reembedded_text_vectors = reembedded
return _Embedded(
query_vector=reembedded[0], vectors=MappingProxyType(dict(zip(unique, reembedded[1:], strict=True)))
query_vector=reembedded_query,
vectors=reembedded_text_vectors.vectors,
too_long=new_text_vectors.too_long | reembedded_text_vectors.too_long,
failure=reembedded_text_vectors.failure,
)
@ -148,18 +229,49 @@ class SemanticTextIndex:
Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first."""
def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None:
def __init__(
self,
max_entries: int = DEFAULT_MAX_CACHED_VECTORS,
too_long_ttl_seconds: float = DEFAULT_TOO_LONG_TTL_SECONDS,
clock: Callable[[], float] = time.monotonic,
) -> None:
self._max_entries: Final = max_entries
self._too_long_ttl_seconds: Final = too_long_ttl_seconds
self._clock: Final = clock
self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({})
self._too_long: Mapping[_CacheKey, float] = MappingProxyType({})
def _cached(self, embedding_model: str) -> Mapping[str, Vector]:
return MappingProxyType(
{text: vector for (model, text), vector in self._vectors.items() if model == embedding_model}
)
def _too_long_for(self, embedding_model: str, now: float) -> frozenset[str]:
return frozenset(
text
for (model, text), seen_at in self._too_long.items()
if model == embedding_model and now - seen_at < self._too_long_ttl_seconds
)
def _merged_too_long(
self, embedding_model: str, embedded: _Embedded, texts: Sequence[str], now: float
) -> Mapping[_CacheKey, float]:
fresh: Final = MappingProxyType(
{key: seen_at for key, seen_at in self._too_long.items() if now - seen_at < self._too_long_ttl_seconds}
)
newly_too_long: Final = (
(embedding_model, text)
for text in texts
if text in embedded.too_long and (embedding_model, text) not in fresh
)
ordered: Final = MappingProxyType({**fresh, **dict.fromkeys(newly_too_long, now)})
return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None)))
def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]:
dimension: Final = len(embedded.query_vector)
touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts})
touched: Final = MappingProxyType(
{(embedding_model, text): embedded.vectors[text] for text in texts if text in embedded.vectors}
)
untouched: Final = MappingProxyType(
{
key: vector
@ -175,14 +287,27 @@ class SemanticTextIndex:
async def scores(
self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str
) -> tuple[float, ...] | EmbeddingFailed:
"""Cosine similarity of `query` to each entry of `texts`, in the same order."""
) -> tuple[float | None, ...] | EmbeddingFailed:
"""Cosine similarity for embeddable entries of `texts`, with None for omitted entries."""
if not texts:
return ()
embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model))
now: Final = self._clock()
embedded: Final = await _embed_query_and_texts(
embed, query, texts, self._cached(embedding_model), self._too_long_for(embedding_model, now)
)
if isinstance(embedded, EmbeddingFailed):
return embedded
if not _same_dimension(embedded.query_vector, embedded.vectors, texts):
same_dimension: Final = _same_dimension(embedded.query_vector, embedded.vectors, texts)
if embedded.failure is not None:
if same_dimension:
self._vectors = self._merged(embedding_model, embedded, texts)
self._too_long = self._merged_too_long(embedding_model, embedded, texts, now)
return embedded.failure
if not same_dimension:
return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions")
self._vectors = self._merged(embedding_model, embedded, texts)
return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts)
self._too_long = self._merged_too_long(embedding_model, embedded, texts, now)
return tuple(
cosine_similarity(embedded.query_vector, embedded.vectors[text]) if text in embedded.vectors else None
for text in texts
)

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

@ -11,9 +11,9 @@ import subprocess
import time
import urllib
import urllib.parse
from collections.abc import Callable
from collections.abc import Awaitable, Callable, Coroutine
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.db_span import db_span
@ -96,10 +96,20 @@ class _PrismaDrainTracker:
await self._drained.wait()
_EngineResponseT: Final = TypeVar("_EngineResponseT")
class _TrackedPrismaEngine:
"""Engine wrapper that counts in-flight work and never lets a cancelled caller drop an engine request.
A dropped request leaves the engine holding a busy pool connection and can crash it once that statement
finishes, so every request runs to completion in its own task and only the caller is cancelled.
"""
def __init__(self, engine: _PrismaEngine, tracker: _PrismaDrainTracker) -> None:
self._engine = engine
self.tracker = tracker
self._detached: frozenset[asyncio.Task[object]] = frozenset()
@property
def process(self) -> _PrismaProcess:
@ -109,7 +119,60 @@ class _TrackedPrismaEngine:
return getattr(self._engine, name)
async def query(self, content: str, *, tx_id: str | None) -> object:
return await self._finish_even_if_cancelled(self._tracked_query(content, tx_id))
async def start_transaction(self, *, content: str) -> str:
return await self._finish_even_if_cancelled(
self._tracked_start_transaction(content), abandoned=self.rollback_transaction
)
async def commit_transaction(self, tx_id: str) -> None:
await self._finish_even_if_cancelled(self._tracked_commit_transaction(tx_id))
async def rollback_transaction(self, tx_id: str) -> None:
await self._finish_even_if_cancelled(self._tracked_rollback_transaction(tx_id))
async def _finish_even_if_cancelled(
self,
request: Coroutine[object, object, _EngineResponseT],
abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None = None,
) -> _EngineResponseT:
self.tracker.begin_operation()
task: Final = asyncio.ensure_future(request)
try:
return await asyncio.shield(task)
except asyncio.CancelledError:
self._detached = self._detached.union((task,))
task.add_done_callback(lambda done: self._detached_request_finished(done, abandoned))
raise
def _detached_request_finished(
self,
done: "asyncio.Task[_EngineResponseT]",
abandoned: Callable[[_EngineResponseT], Awaitable[None]] | None,
) -> None:
self._detached = self._detached.difference((done,))
if done.cancelled():
return
failure: Final = done.exception()
if failure is not None:
verbose_proxy_logger.debug("Prisma request abandoned by a cancelled caller failed: %s", failure)
return
if abandoned is None:
return
settle: Final = asyncio.ensure_future(abandoned(done.result()))
self._detached = self._detached.union((settle,))
settle.add_done_callback(self._abandoned_transaction_settled)
def _abandoned_transaction_settled(self, settle: "asyncio.Task[object]") -> None:
self._detached = self._detached.difference((settle,))
if settle.cancelled() or settle.exception() is None:
return
verbose_proxy_logger.debug(
"Rolling back a transaction abandoned by a cancelled caller failed: %s", settle.exception()
)
async def _tracked_query(self, content: str, tx_id: str | None) -> object:
try:
if db_io_claimed():
record_db_io()
@ -121,9 +184,8 @@ class _TrackedPrismaEngine:
finally:
self.tracker.end_operation()
async def start_transaction(self, *, content: str) -> str:
async def _tracked_start_transaction(self, content: str) -> str:
record_db_io()
self.tracker.begin_operation()
try:
transaction_id: Final = await self._engine.start_transaction(content=content)
except (Exception, asyncio.CancelledError):
@ -132,18 +194,16 @@ class _TrackedPrismaEngine:
self.tracker.transaction_started(transaction_id)
return transaction_id
async def commit_transaction(self, tx_id: str) -> None:
async def _tracked_commit_transaction(self, tx_id: str) -> None:
record_db_io()
self.tracker.begin_operation()
try:
await self._engine.commit_transaction(tx_id)
finally:
self.tracker.end_operation()
self.tracker.transaction_finished(tx_id)
async def rollback_transaction(self, tx_id: str) -> None:
async def _tracked_rollback_transaction(self, tx_id: str) -> None:
record_db_io()
self.tracker.begin_operation()
try:
await self._engine.rollback_transaction(tx_id)
finally:

6
litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
View file

@ -33,9 +33,15 @@ async def get_ui_config():
is_control_plane: Final = len(proxy_config.worker_registry) > 0
raw_gateway_base_url: Final[str | None] = os.getenv("LITELLM_GATEWAY_URL")
gateway_base_url: Final[str | None] = (
raw_gateway_base_url.strip().rstrip("/") or None if raw_gateway_base_url is not None else None
)
return UiDiscoveryEndpoints(
server_root_path=get_server_root_path(),
proxy_base_url=get_proxy_base_url(),
gateway_base_url=gateway_base_url,
auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso,
admin_ui_disabled=admin_ui_disabled,
sso_configured=sso_configured,

32
litellm/proxy/guardrails/guardrail_hooks/airia/__init__.py Normal file
View file

@ -0,0 +1,32 @@
from typing import TYPE_CHECKING, Final
from litellm.proxy.guardrails.guardrail_hooks.airia.airia import AiriaGuardrail
from litellm.types.guardrails import SupportedGuardrailIntegrations
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> AiriaGuardrail:
import litellm
_airia_callback: Final = AiriaGuardrail(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
timeout=litellm_params.timeout,
guardrail_name=guardrail["guardrail_name"],
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)
litellm.logging_callback_manager.add_litellm_callback(_airia_callback)
return _airia_callback
guardrail_initializer_registry: Final = {
SupportedGuardrailIntegrations.AIRIA.value: initialize_guardrail,
}
guardrail_class_registry: Final = {
SupportedGuardrailIntegrations.AIRIA.value: AiriaGuardrail,
}

246
litellm/proxy/guardrails/guardrail_hooks/airia/airia.py Normal file
View file

@ -0,0 +1,246 @@
import os
import uuid
from collections.abc import Mapping
from typing import (
Any, # noqa: TID251 # the only type CustomGuardrail.__init__ accepts for its open-ended kwargs
Final,
Literal,
)
import httpx
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import GuardrailRaisedException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # the decorator is untyped upstream
)
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # the factory's params are untyped upstream
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.airia import AiriaGuardrailConfigModel
from litellm.types.utils import GenericGuardrailAPIInputs
GUARDRAIL_PATH: Final = "/v1/guardrails/litellm"
ACTION_NONE: Final = "NONE"
ACTION_BLOCKED: Final = "BLOCKED"
ACTION_INTERVENED: Final = "GUARDRAIL_INTERVENED"
DEFAULT_BLOCKED_MESSAGE: Final = "Blocked by your organization's content policy."
SUPPORTED_EVENT_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call)
def _stream_cannot_carry_a_rewrite(request_data: Mapping[str, object] | None) -> bool:
"""Whether this is a streamed response on a route the framework can only block, not rewrite.
``incremental_diff``, the mode that writes a rewrite into the stream, exists only for OpenAI chat
completions: ``UnifiedLLMGuardrails._resolve_transform_call_type`` sends every other streamed route
to ``block_only``, which drops the rewrite and replays the buffered original. On those routes a
REDACT verdict can only be honoured as a block. The route check mirrors that resolver so the two
cannot disagree, and a route it cannot resolve fails closed the same way the resolver does.
"""
if request_data is None or not _is_streamed(request_data):
return False
route: Final = _request_route(request_data)
call_types: Final = get_call_types_for_route(route) if route else None
if not call_types:
return True
# Deferred like the resolver's own import: the handler package imports the proxy.
from litellm.llms import load_guardrail_translation_mappings # noqa: PLC0415 # circular import
from litellm.llms.openai.chat.guardrail_translation.handler import ( # noqa: PLC0415 # circular import
OpenAIChatCompletionsHandler,
)
handler: Final[object] = load_guardrail_translation_mappings().get(call_types[0])
return not (isinstance(handler, type) and issubclass(handler, OpenAIChatCompletionsHandler))
def _is_streamed(request_data: Mapping[str, object]) -> bool:
if request_data.get("stream") is True:
return True
proxy_request: Final = _as_object(request_data.get("proxy_server_request"))
body: Final = _as_object(proxy_request.get("body")) if proxy_request is not None else None
return body is not None and body.get("stream") is True
def _request_route(request_data: Mapping[str, object]) -> str | None:
"""The route the proxy stamps as ``user_api_key_request_route``, in whichever metadata bucket the
route uses: ``metadata`` for chat completions, ``litellm_metadata`` for Responses and Anthropic."""
for bucket_name in ("metadata", "litellm_metadata"):
bucket = _as_object(request_data.get(bucket_name))
route = bucket.get("user_api_key_request_route") if bucket is not None else None
if isinstance(route, str) and route:
return route
return None
def _as_object(value: object) -> Mapping[str, object] | None:
"""`value` as a JSON object, or None. States the key type once so the checker does not see
`dict[Unknown, Unknown]` at every read."""
return value if isinstance(value, dict) else None # pyright: ignore[reportUnknownVariableType] # isinstance narrows to dict[Unknown, Unknown]; JSON keys are str
class AiriaGuardrail(CustomGuardrail):
"""Evaluates prompts and responses against your Airia guardrail policy."""
def __init__(
self,
api_base: str | None = None,
api_key: str | None = None,
timeout: float | None = None,
supported_event_hooks: list[GuardrailEventHooks] | None = None, # mutable-ok: matches CustomGuardrail.__init__
async_handler: AsyncHTTPHandler | None = None, # dependency-injected in tests; defaults to a real client
**kwargs: Any, # kwargs-ok: passed straight through to CustomGuardrail.__init__
) -> None:
resolved_timeout: Final = timeout or float(os.getenv("AIRIA_TIMEOUT", "10"))
self.async_handler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
params={"timeout": httpx.Timeout(timeout=resolved_timeout, connect=5.0)},
)
self.api_base = (api_base or os.getenv("AIRIA_GATEWAY_URL", "")).rstrip("/")
self.api_key = api_key or os.getenv("AIRIA_API_KEY")
self.streaming_transform_mode: Final[Literal["block_only", "incremental_diff"]] = "incremental_diff"
self.streaming_end_of_stream_only: Final = True
# incremental_diff exists only for streamed chat completions. Every other streamed route falls
# back to block_only, where the framework releases the original text live unless told to hold
# it; holding it keeps the text behind the end-of-stream verdict, and apply_guardrail turns a
# REDACT that block_only cannot deliver into a block. Never read on the incremental_diff path.
self.streaming_buffer_until_moderated: Final = True
if not self.api_base:
raise ValueError("AiriaGuardrail requires api_base, or the AIRIA_GATEWAY_URL environment variable.")
if not self.api_key:
raise ValueError("AiriaGuardrail requires api_key, or the AIRIA_API_KEY environment variable.")
self.optional_params = kwargs
super().__init__( # pyright: ignore[reportUnknownMemberType] # parent takes untyped **kwargs
supported_event_hooks=supported_event_hooks or [*SUPPORTED_EVENT_HOOKS],
**kwargs,
)
verbose_proxy_logger.info("AiriaGuardrail initialized with gateway: %s", self.api_base)
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object], # mutable-ok: the base class declares this parameter as a dict
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
call_id: Final = (
logging_obj.litellm_call_id
if logging_obj
else (request_data.get("litellm_call_id") if request_data else None)
) or str(uuid.uuid4())
try:
response: Final = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # the handler's `data` param is an untyped dict upstream
f"{self.api_base}{GUARDRAIL_PATH}",
json={
"input_type": input_type,
"texts": inputs.get("texts") or (),
"images": inputs.get("images") or (),
"structured_messages": inputs.get("structured_messages") or (),
"tools": inputs.get("tools") or (),
"tool_calls": inputs.get("tool_calls") or (),
"model": inputs.get("model"),
"litellm_call_id": call_id,
},
headers={"Authorization": f"Bearer {self.api_key}"},
)
response.raise_for_status()
raw_body: Final[object] = response.json()
body: Final = _as_object(raw_body)
if body is None:
raise TypeError(f"expected a JSON object in the response body, got {type(raw_body).__name__}")
except Exception as error:
verbose_proxy_logger.error(
"Airia guardrail could not evaluate the request (litellm_call_id=%s, input_type=%s): %s",
call_id,
input_type,
error,
)
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=f"Airia guardrail could not evaluate the request: {error}",
blocked_content=False,
) from error
action: Final = body.get("action")
if action == ACTION_BLOCKED:
reason: Final = body.get("blocked_reason")
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=reason if isinstance(reason, str) and reason else DEFAULT_BLOCKED_MESSAGE,
blocked_content=True,
)
if action == ACTION_INTERVENED:
if input_type == "response" and _stream_cannot_carry_a_rewrite(request_data):
raise self._blocked()
return self._rewritten(body, inputs)
if action != ACTION_NONE:
raise self._blocked()
return inputs
def _rewritten(
self,
body: Mapping[str, object],
inputs: GenericGuardrailAPIInputs,
) -> GenericGuardrailAPIInputs:
texts: Final = body.get("texts")
structured_messages: Final = body.get("structured_messages")
if not self._is_applicable_rewrite(texts, inputs.get("texts")):
raise self._blocked()
if not self._is_applicable_rewrite(structured_messages, inputs.get("structured_messages")):
raise self._blocked()
# An intervention that rewrites nothing would hand the original content through under a
# verdict that said it must not pass.
if texts is None and structured_messages is None:
raise self._blocked()
rewritten: Final[GenericGuardrailAPIInputs] = {**inputs}
if isinstance(texts, list):
rewritten["texts"] = texts
if isinstance(structured_messages, list):
rewritten["structured_messages"] = structured_messages
return rewritten
@staticmethod
def _is_applicable_rewrite(rewrite: object, original: object) -> bool:
"""A rewrite is safe to apply positionally only when it is absent (this field was not
touched) or a list the same length as what was sent. A different count would misalign
the positional write-back downstream, dropping content or leaving other fields unchanged.
A rewrite for a field that was never sent is a mismatch too, not a no-op: post_call sends
only `texts`, so `structured_messages: []` would otherwise count as an applied rewrite
while the `texts` that carried the content passed through untouched.
"""
if rewrite is None:
return True
if not isinstance(original, list) or not isinstance(rewrite, list):
return False
return len(rewrite) == len(original) # pyright: ignore[reportUnknownArgumentType] # only the counts matter; the items are the gateway's
def _blocked(self) -> GuardrailRaisedException:
return GuardrailRaisedException(
guardrail_name=self.guardrail_name,
message=DEFAULT_BLOCKED_MESSAGE,
blocked_content=True,
)
@staticmethod
def get_config_model() -> type[AiriaGuardrailConfigModel] | None:
return AiriaGuardrailConfigModel

39
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -2,7 +2,7 @@ import asyncio
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Protocol
from typing import TYPE_CHECKING, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, status
from typing_extensions import ReadOnly, TypedDict
@ -45,6 +45,9 @@ from litellm.types.access_group import (
AccessGroupUpdateRequest,
)
if TYPE_CHECKING:
from prisma import Json
router: Final = APIRouter(
tags=["access group management"],
)
@ -57,6 +60,9 @@ class _AccessGroupRecord(Protocol):
@property
def access_mcp_server_ids(self) -> Sequence[str] | None: ...
@property
def mcp_tool_permissions(self) -> "Mapping[str, Sequence[str]] | Json | None": ...
@property
def access_agent_ids(self) -> Sequence[str] | None: ...
@ -179,6 +185,12 @@ def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None:
)
def _tool_permissions_json(value: Mapping[str, Sequence[str]] | None) -> "Json":
from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools
return Json(None) if value is None else Json({server: list(tools) for server, tools in value.items()})
@dataclass(frozen=True, slots=True)
class _ResourceNames:
mcp_servers: Mapping[str, str]
@ -544,7 +556,9 @@ async def create_access_group(
"description": data.description,
"access_model_names": data.access_model_names or [],
"access_mcp_server_ids": data.access_mcp_server_ids or [],
"mcp_tool_permissions": _tool_permissions_json(data.mcp_tool_permissions),
"access_agent_ids": data.access_agent_ids or [],
"access_skill_names": list(data.access_skill_names or ()),
"assigned_team_ids": data.assigned_team_ids or [],
"assigned_key_ids": data.assigned_key_ids or [],
"created_by": user_api_key_dict.user_id,
@ -636,19 +650,18 @@ async def update_access_group(
update_fields: Final = data.model_dump(exclude_unset=True)
update_data: Final[dict] = {"updated_by": user_api_key_dict.user_id}
for field, value in update_fields.items():
if (
field
in (
"assigned_team_ids",
"assigned_key_ids",
"access_model_names",
"access_mcp_server_ids",
"access_agent_ids",
)
and value is None
if field in (
"assigned_team_ids",
"assigned_key_ids",
"access_model_names",
"access_mcp_server_ids",
"access_agent_ids",
"access_skill_names",
):
value = []
update_data[field] = value
value = list(value or ())
update_data[field] = (
_tool_permissions_json(data.mcp_tool_permissions) if field == "mcp_tool_permissions" else value
)
# Initialize delta lists before the try block so they remain accessible
# for cache updates after the transaction, even if an error path is added later.

198
litellm/proxy/management_endpoints/scim/scim_v2.py
View file

@ -9,9 +9,9 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from functools import partial
from itertools import chain
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, overload
from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypeAlias, overload
from fastapi import (
APIRouter,
@ -29,6 +29,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.models.user import SCIMPlaceholder
from litellm.proxy._types import (
@ -70,6 +71,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
handle_exception_on_proxy,
premium_user_check,
)
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.table_repositories import (
InvitationLinkRepository,
OrganizationMembershipRepository,
@ -565,6 +567,111 @@ class _CaseInsensitiveMatch(TypedDict):
mode: ReadOnly[str]
_InFilter: Final = TypedDict("_InFilter", {"in": ReadOnly[Sequence[str]]})
_CaseInsensitiveInFilter: Final = TypedDict(
"_CaseInsensitiveInFilter", {"in": ReadOnly[Sequence[str]], "mode": ReadOnly[str]}
)
class _UserIdIn(TypedDict):
user_id: ReadOnly[_InFilter]
class _SsoUserIdIn(TypedDict):
sso_user_id: ReadOnly[_InFilter]
class _UserEmailIn(TypedDict):
user_email: ReadOnly[_CaseInsensitiveInFilter]
_IdentityRowsClause: TypeAlias = _UserIdIn | _SsoUserIdIn | _UserEmailIn
class _IdentityRowsWhere(TypedDict):
OR: ReadOnly[Sequence[_IdentityRowsClause]]
@dataclass(frozen=True, slots=True)
class _IdentityIndex:
user_ids: frozenset[str]
by_sso_user_id: Mapping[str, tuple[str, ...]]
by_email: Mapping[str, tuple[str, ...]]
async def _identity_rows_for_chunk(chunk: Sequence[str], prisma_client: PrismaClient) -> tuple[LiteLLM_UserTable, ...]:
subjects: Final = list(dict.fromkeys(value.strip() for value in chunk))
where: Final[_IdentityRowsWhere] = {
"OR": [
{"user_id": {"in": list(chunk)}}, # bounded-ok: chunk cap is 10,000
{"sso_user_id": {"in": subjects}}, # bounded-ok: subjects are chunk-bounded
{"user_email": {"in": subjects, "mode": "insensitive"}}, # bounded-ok: subjects are chunk-bounded
]
}
return tuple(await _table(UserRepository(prisma_client)).find_many(where=where))
async def _identity_rows_for_values(
values: Sequence[str], prisma_client: PrismaClient
) -> tuple[LiteLLM_UserTable, ...]:
unique_values: Final = tuple(dict.fromkeys(values))
if not unique_values:
return ()
batch_size: Final = SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE
chunks: Final = tuple(
unique_values[start : start + batch_size] for start in range(0, len(unique_values), batch_size)
)
chunk_rows: Final = tuple([await _identity_rows_for_chunk(chunk, prisma_client) for chunk in chunks])
rows: Final = tuple(chain.from_iterable(chunk_rows))
return tuple({row.user_id: row for row in rows}.values())
def _identity_index(rows: Sequence[LiteLLM_UserTable]) -> _IdentityIndex:
sso_entries: Final = tuple(sorted((row.sso_user_id, row.user_id) for row in rows if row.sso_user_id is not None))
email_entries: Final = tuple(
sorted((row.user_email.lower(), row.user_id) for row in rows if row.user_email is not None)
)
by_sso_user_id: Final = MappingProxyType(
{
subject: tuple(sorted({user_id for _, user_id in matches}))
for subject, matches in groupby(sso_entries, key=lambda entry: entry[0])
}
)
by_email: Final = MappingProxyType(
{
email: tuple(sorted({user_id for _, user_id in matches}))
for email, matches in groupby(email_entries, key=lambda entry: entry[0])
}
)
return _IdentityIndex(
user_ids=frozenset(row.user_id for row in rows),
by_sso_user_id=by_sso_user_id,
by_email=by_email,
)
def _users_named_by(value: str, index: _IdentityIndex) -> tuple[str, ...]:
subject: Final = value.strip()
named: Final = frozenset(
chain(
index.by_sso_user_id.get(subject, ()),
index.by_email.get(subject.lower(), ()),
)
)
return tuple(sorted(named))
def _accounts_named_by(value: str, index: _IdentityIndex) -> tuple[str, ...]:
named: Final = frozenset(
chain(
_users_named_by(value, index),
(value,) if value in index.user_ids else (),
)
)
return tuple(sorted(named))
async def _users_named_by_member_value(
value: str, prisma_client: PrismaClient, *, take: int | None = 2
) -> tuple[str, ...]:
@ -597,50 +704,35 @@ async def _users_named_by_member_value(
return tuple(dict.fromkeys(row.user_id for row in rows))
async def _accounts_named_by_member_value(value: str, prisma_client: PrismaClient) -> tuple[str, ...]:
"""Every user id this member value names, by user id, SSO identity or email.
async def _scim_team_ids_among(values: Sequence[str], prisma_client: PrismaClient) -> frozenset[str]:
unique_values: Final = tuple(dict.fromkeys(values))
if not unique_values:
return frozenset()
Classification needs to know whether the value is one account's ``user_id`` and
whether it names any other account, so all three fields are read in one pass. The
id is compared exactly and unstripped, as a primary key lookup would; the
identities compare as ``_users_named_by_member_value`` describes. Two rows are
enough to tell one account from several, so the read stops there. Only a full
read that lacks the row keyed by the value leaves that row's existence open, and
only then is the id read on its own.
"""
subject: Final = value.strip()
email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"}
users: Final = _table(UserRepository(prisma_client))
rows: Final = await users.find_many(
where={
"OR": [
{"user_id": value},
{"sso_user_id": subject},
{"user_email": email},
],
},
take=2,
teams: Final = await find_many_in(
_table(TeamRepository(prisma_client)),
"team_id",
unique_values,
chunk_size=SCIM_GROUP_MEMBER_LOOKUP_BATCH_SIZE,
)
named: Final = tuple(dict.fromkeys(row.user_id for row in rows))
if len(named) < 2 or value in named:
return named
keyed: Final = await users.find_unique(where={"user_id": value})
return named if keyed is None else (value, *named)
return frozenset(team.team_id for team in teams if _team_metadata_has_scim_provenance(team.metadata))
async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember:
def _classify_group_member(
member: SCIMMember, named: tuple[str, ...], scim_team_ids: frozenset[str]
) -> _ClassifiedGroupMember:
"""
Decide what a single SCIM group member refers to.
A LiteLLM team only holds users, so a member is dropped when it declares a type
other than ``User`` or when its id names an existing team. Both of those checks
are placed around the user lookup rather than before it, because the id of a
are placed around identity resolution rather than before it, because the id of a
real user is the one thing that outranks them:
- ``"type": "Group"`` (what Entra sends for a nested group) is dropped without
a lookup. This bug provisioned nested group GUIDs as users, so those rows
identity resolution. This bug provisioned nested group GUIDs as users, so those rows
exist in the wild and would otherwise resolve as members all over again.
- any other unrecognized type is dropped only after the user lookup misses.
- any other unrecognized type is dropped only after identity resolution misses.
Clients do send non-canonical types on real members (RFC 7643 defines
``direct`` for ``User.groups``), and dropping a live user over one would
revoke that user's team access on the next full sync.
@ -650,8 +742,8 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
even if some team happens to share the id, and a team created here rather
than through SCIM is not evidence of anything about the member.
When those checks miss on an otherwise user-shaped member, its value is looked
up as an SSO identity or an email, and a match resolves to that user's
When those checks miss on an otherwise user-shaped member, its value is matched
against SSO identities or emails, and a match resolves to that user's
``user_id``. A value that names more than one account is ambiguous rather than
unknown: it names a real person we cannot identify, so it is neither guessed at
nor provisioned.
@ -670,7 +762,6 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
if member_type == "group":
return _SkippedGroupMember(value=value, reason="nested_group")
named: Final = await _accounts_named_by_member_value(value, prisma_client)
if value in named:
shared_with: Final = tuple(other for other in named if other != value)
if shared_with:
@ -687,10 +778,8 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient
if member_type is not None and member_type != "user":
return _SkippedGroupMember(value=value, reason="non_user_type")
if member_type is None:
team: Final = await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": value})
if team is not None and _team_metadata_has_scim_provenance(team.metadata):
return _SkippedGroupMember(value=value, reason="existing_team")
if member_type is None and value in scim_team_ids:
return _SkippedGroupMember(value=value, reason="existing_team")
if len(named) == 1:
verbose_proxy_logger.info(
@ -837,16 +926,10 @@ async def _member_ids_to_drop(
HTTPException: 400 when a member id names more than one current member.
"""
written: Final = frozenset(_member_value(member) for member in members)
rows: Final = await _identity_rows_for_values(sorted(written), prisma_client)
index: Final = _identity_index(rows)
matched: Final = tuple(
[
(
value,
_roster_entries_named_by(
value, roster, await _users_named_by_member_value(value, prisma_client, take=None)
),
)
for value in sorted(written)
]
(value, _roster_entries_named_by(value, roster, _users_named_by(value, index))) for value in sorted(written)
)
undecidable: Final = tuple(value for value, entries in matched if len(entries) > 1)
if undecidable:
@ -881,7 +964,22 @@ async def _resolve_group_member_ids(
other than a user. 500 when a member's user row can neither be created nor
found.
"""
classified: Final = tuple([await _classify_group_member(member, prisma_client) for member in members])
values: Final = tuple(_member_value(member) for member in members)
identity_values: Final = tuple(
value for member, value in zip(members, values) if _normalized_member_type(member) != "group"
)
rows: Final = await _identity_rows_for_values(identity_values, prisma_client)
index: Final = _identity_index(rows)
named_by: Final = MappingProxyType({value: _accounts_named_by(value, index) for value in dict.fromkeys(values)})
team_candidates: Final = tuple(
value
for member, value in zip(members, values)
if _normalized_member_type(member) is None and value not in named_by[value]
)
scim_teams: Final = await _scim_team_ids_among(team_candidates, prisma_client)
classified: Final = tuple(
_classify_group_member(member, named_by[value], scim_teams) for member, value in zip(members, values)
)
partition: Final = _partition_classified_members(classified)
for skipped in partition.skipped:

88
litellm/proxy/management_endpoints/team_admin_field_permissions.py
View file

@ -17,6 +17,7 @@ from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm.models.budget import LiteLLM_BudgetTable
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.verification_token import LiteLLM_VerificationToken
from litellm.proxy._types import (
@ -32,7 +33,16 @@ TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_field
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset(
{"tpm_limit", "rpm_limit", "max_budget", "soft_budget"}
{
"tpm_limit",
"rpm_limit",
"max_budget",
"soft_budget",
"team_member_budget",
"team_member_budget_duration",
"team_member_tpm_limit",
"team_member_rpm_limit",
}
)
TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION: Final = "raise_max_budget"
TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects"
@ -66,6 +76,19 @@ _NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"})
_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields"
@dataclass(frozen=True, slots=True)
class TeamMemberDefaults:
team_member_budget: float | None = None
team_member_budget_duration: str | None = None
team_member_tpm_limit: int | None = None
team_member_rpm_limit: int | None = None
TEAM_MEMBER_DEFAULT_FIELDS: Final[frozenset[str]] = frozenset(
{"team_member_budget", "team_member_budget_duration", "team_member_tpm_limit", "team_member_rpm_limit"}
)
@dataclass(frozen=True, slots=True)
class TeamAdminEditAllowed:
request: UpdateTeamRequest
@ -250,34 +273,76 @@ def _metadata_changes(
)
def stored_team_member_budget_id(team: LiteLLM_TeamTable) -> str | None:
budget_id: Final = _stored_metadata(_JSON_OBJECT.validate_json(team.model_dump_json())).get("team_member_budget_id")
return budget_id if isinstance(budget_id, str) else None
def team_member_defaults(budget: LiteLLM_BudgetTable | None) -> TeamMemberDefaults | None:
if budget is None:
return None
return TeamMemberDefaults(
team_member_budget=budget.max_budget,
team_member_budget_duration=budget.budget_duration,
team_member_tpm_limit=budget.tpm_limit,
team_member_rpm_limit=budget.rpm_limit,
)
def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]:
table: Final = existing_row.litellm_model_table
return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY
def _column_changed(
field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable
field: str,
submitted: Mapping[str, object],
existing: Mapping[str, object],
existing_row: LiteLLM_TeamTable,
member_defaults: Mapping[str, object],
) -> bool:
if field == "model_aliases":
return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row)
if field in member_defaults:
return submitted.get(field) != member_defaults[field]
if field in LiteLLM_TeamTable.model_fields:
return submitted.get(field) != existing.get(field)
return True
def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]:
def _member_defaults_by_field(member_defaults: TeamMemberDefaults | None) -> Mapping[str, object]:
stored: Final = member_defaults or TeamMemberDefaults()
return MappingProxyType(
{
"team_member_budget": stored.team_member_budget,
"team_member_budget_duration": stored.team_member_budget_duration,
"team_member_tpm_limit": stored.team_member_tpm_limit,
"team_member_rpm_limit": stored.team_member_rpm_limit,
}
)
def changed_team_fields(
data: UpdateTeamRequest,
existing_row: LiteLLM_TeamTable,
member_defaults: TeamMemberDefaults | None = None,
) -> frozenset[str]:
"""Logical field names whose stored value the request would change.
Request and stored row are compared as JSON values so both sides share one representation. Fields the
server folds into metadata are attributed to their own name whether they arrive top-level or inside
``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored
counterpart on the team row count as changed whenever they are sent.
``metadata``; anything else in ``metadata`` is attributed to ``metadata``. The ``team_member_*`` defaults
compare against the team's member budget row (``None`` when it has none, read as all null). Fields with no
stored counterpart count as changed whenever they are sent.
"""
submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True))
existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json())
column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS
stored_member_defaults: Final = _member_defaults_by_field(member_defaults)
column_changes: Final = frozenset(
field for field in column_fields if _column_changed(field, submitted, existing, existing_row)
field
for field in column_fields
if _column_changed(field, submitted, existing, existing_row, stored_member_defaults)
)
return column_changes | _metadata_changes(data, submitted, existing)
@ -295,14 +360,21 @@ def team_admin_edit_verdict(
data: UpdateTeamRequest,
existing: LiteLLM_TeamTable,
permitted: frozenset[str],
member_defaults: TeamMemberDefaults | None = None,
) -> TeamAdminEditVerdict:
if not permitted:
return TeamAdminEditingDisabled()
changed: Final = changed_team_fields(data, existing)
changed: Final = changed_team_fields(data, existing, member_defaults)
blocked: Final = sorted(changed - permitted)
if blocked:
return TeamAdminFieldNotPermitted(field=blocked[0])
return TeamAdminEditAllowed(request=_only_changes(data, changed))
creates_member_budget: Final = member_defaults is None and bool(changed & TEAM_MEMBER_DEFAULT_FIELDS)
sent_for_new_member_budget: Final = (
frozenset(data.model_fields_set) & TEAM_MEMBER_DEFAULT_FIELDS & permitted
if creates_member_budget
else frozenset[str]()
)
return TeamAdminEditAllowed(request=_only_changes(data, changed | sent_for_new_member_budget))
def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest:

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

@ -18,6 +18,7 @@ from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -164,10 +165,14 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
TEAM_MEMBER_DEFAULT_FIELDS,
TeamMemberDefaults,
resolve_team_admin_editable_fields,
stored_team_member_budget_id,
team_admin_edit_verdict,
team_admin_may_raise_max_budget,
team_admin_request_or_raise,
team_member_defaults,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
@ -2365,12 +2370,19 @@ async def update_team(
if access_role is None:
team_access_denied()
if access_role == "team_admin":
permitted: Final = resolve_team_admin_editable_fields(
_general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
)
reads_member_defaults: Final = bool(permitted) and bool(data.model_fields_set & TEAM_MEMBER_DEFAULT_FIELDS)
data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below
team_admin_edit_verdict(
data=data,
existing=existing_team,
permitted=resolve_team_admin_editable_fields(
_general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
permitted=permitted,
member_defaults=(
await _stored_team_member_defaults(prisma_client, existing_team)
if reads_member_defaults
else None
),
)
)
@ -4899,6 +4911,15 @@ async def validate_membership(user_api_key_dict: UserAPIKeyAuth, team_table: Lit
)
async def _stored_team_member_defaults(
prisma_client: PrismaClient, team: LiteLLM_TeamTable
) -> TeamMemberDefaults | None:
budget_id: Final = stored_team_member_budget_id(team)
if budget_id is None:
return None
return team_member_defaults(await BudgetRepository(prisma_client).find_by_id(budget_id))
async def _add_team_member_budget_table(
team_member_budget_id: str,
prisma_client: PrismaClient,
@ -4952,6 +4973,16 @@ def _parent_organization_models(team_row: BaseModel) -> list[str] | None:
return organization.models if organization is not None else None
def _effective_access_group_skill_names(
resolved_groups: "Sequence[prisma_models.LiteLLM_AccessGroupTable]", unresolved_ids: Sequence[str]
) -> tuple[str, ...]:
"""Skill names the team's keys can actually use: none while a listed group has no row,
because the skill check then clamps the team to public plugins (see skill_visibility)."""
if unresolved_ids:
return ()
return tuple(frozenset(chain.from_iterable(group.access_skill_names or () for group in resolved_groups)))
async def _resolve_team_access_group_resources(
_team_info: TeamInfoResponseObjectTeamTable,
) -> TeamInfoResponseObjectTeamTable:
@ -4963,6 +4994,9 @@ async def _resolve_team_access_group_resources(
resolved_groups: Final = tuple(
ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup
)
unresolved_ids: Final = tuple(
ag_id for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id not in ag_lookup
)
return _team_info.model_copy(
update={
"access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}),
@ -4970,6 +5004,8 @@ async def _resolve_team_access_group_resources(
{s for group in resolved_groups for s in (group.access_mcp_server_ids or [])}
),
"access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}),
"access_group_skill_names": _effective_access_group_skill_names(resolved_groups, unresolved_ids),
"unresolved_access_group_ids": unresolved_ids,
"access_group_details": tuple(
TeamAccessGroupModelGrant(
access_group_id=group.access_group_id,
@ -4977,6 +5013,7 @@ async def _resolve_team_access_group_resources(
models=tuple(group.access_model_names or ()),
mcp_server_ids=tuple(group.access_mcp_server_ids or ()),
agent_ids=tuple(group.access_agent_ids or ()),
skill_names=tuple(group.access_skill_names or ()),
)
for group in resolved_groups
),
@ -5629,6 +5666,7 @@ def _convert_teams_to_response_models(
**team_dict,
members_count=members_count,
keys_count=keys_count,
access_group_skill_names=None,
)
)
return team_list
@ -5916,6 +5954,9 @@ async def list_team_v2(
team_item.access_group_agent_ids = list(
{a for group in team_groups for a in (group.access_agent_ids or [])}
)
team_item.access_group_skill_names = _effective_access_group_skill_names(
team_groups, tuple(ag_id for ag_id in (team_item.access_group_ids or []) if ag_id not in ag_lookup)
)
return {
"teams": team_list,

2
litellm/proxy/schema.prisma
View file

@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
mcp_tool_permissions Json?
access_agent_ids String[] @default([])
access_skill_names String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])

5
litellm/responses/litellm_completion_transformation/handler.py
View file

@ -5,6 +5,8 @@ Handler for transforming responses api requests to litellm.completion requests
from collections.abc import Coroutine, Mapping
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.responses.additional_tools import hoist_additional_tools
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
@ -21,6 +23,8 @@ from litellm.types.llms.openai import (
)
from litellm.types.utils import ModelResponse
_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object])
class LiteLLMCompletionTransformationHandler:
def response_api_handler(
@ -107,6 +111,7 @@ class LiteLLMCompletionTransformationHandler:
previous_response_id=previous_response_id,
litellm_completion_request=litellm_completion_request,
instructions=responses_api_request.get("instructions"),
request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs),
)
acompletion_args: Final = {}

4
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -51,6 +51,7 @@ class ResponsesSessionHandler:
@staticmethod
async def get_chat_completion_message_history_for_previous_response_id(
previous_response_id: str,
item_cache_control: bool = False,
) -> ChatCompletionSession:
"""
Return the chat completion message history for a previous response id
@ -86,6 +87,7 @@ class ResponsesSessionHandler:
spend_log=spend_log,
proxy_server_request_dict=proxy_server_request_dict,
chat_completion_message_history=chat_completion_message_history,
item_cache_control=item_cache_control,
)
)
@ -117,6 +119,7 @@ class ResponsesSessionHandler:
| ChatCompletionResponseMessage
| Message
],
item_cache_control: bool = False,
):
"""
Extend the chat completion message history with the spend log payload
@ -148,6 +151,7 @@ class ResponsesSessionHandler:
input=response_input_param,
responses_api_request={},
replay_reasoning=True,
item_cache_control=item_cache_control,
)
)

444
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -5,8 +5,10 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion
import json
import re
import uuid
from collections.abc import Iterable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from itertools import accumulate, chain, islice
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -15,6 +17,7 @@ from typing import (
Literal,
Protocol,
TypeAlias,
TypeVar,
cast,
runtime_checkable,
)
@ -33,9 +36,12 @@ from openai.types.responses.tool_param import FunctionToolParam
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.caching import InMemoryCache
from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER
from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
)
@ -45,6 +51,11 @@ from litellm.litellm_core_utils.hidden_params import (
set_hidden_params,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.common_utils import supports_anthropic_cache_control
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import bedrock_declared_cache_point_support
from litellm.llms.openrouter.chat.transformation import openrouter_model_supports_cache_control
from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model
from litellm.responses.litellm_completion_transformation.session_handler import (
ResponsesSessionHandler,
)
@ -52,6 +63,7 @@ from litellm.types.llms.base import CachedTokensDetails
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionCachedContent,
ChatCompletionImageObject,
ChatCompletionImageUrlObject,
ChatCompletionRedactedThinkingBlock,
@ -186,6 +198,7 @@ TOOL_CALLS_CACHE: Final = InMemoryCache()
_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent)
_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]])
_TEXT_ADAPTER: Final = TypeAdapter(str)
_RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice)
@ -213,6 +226,14 @@ class _ResponsesToolCallItem(Protocol):
def get(self, key: str, /) -> object: ...
class _MarkedToolCallChunk(ChatCompletionToolCallChunk):
cache_control: ReadOnly[ChatCompletionCachedContent]
class _MarkedToolMessage(ChatCompletionToolMessage):
cache_control: ReadOnly[ChatCompletionCachedContent]
class _ToolFunctionDefinition(TypedDict, total=False):
name: ReadOnly[str]
description: ReadOnly[str]
@ -230,6 +251,258 @@ def _input_item_role(input_item: Mapping[str, object]) -> str:
return cast(str, input_item.get("role") or "user") # cast-ok: client-supplied role forwarded verbatim, unvalidated
def _input_item_cache_control(input_item: Mapping[str, object]) -> ChatCompletionCachedContent | None:
cache_control: Final = input_item.get("cache_control")
if not cache_control:
return None
return cast( # cast-ok: forwarded as the client sent it, like a message mark on chat completions
ChatCompletionCachedContent, cache_control
)
def _validated_cache_control(value: object) -> ChatCompletionCachedContent | None:
try:
return _CACHED_CONTENT_ADAPTER.validate_python(value)
except ValidationError:
return None
def _dropped_cache_control(_: object) -> None:
return None
def drops_unsupported_params(drop_params: object) -> bool:
return litellm.drop_params is True or normalize_drop_params(drop_params) is True
def _without_cache_control_key(item: object) -> object:
if not isinstance(item, dict):
return item
return {key: value for key, value in _STR_KEY_DICT_ADAPTER.validate_python(item).items() if key != "cache_control"}
def without_item_cache_control(input: str | ResponseInputParam) -> str | ResponseInputParam:
if not isinstance(input, list):
return input
return cast( # cast-ok: the same input items, each minus its own cache_control key
ResponseInputParam, [_without_cache_control_key(item) for item in input]
)
_PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai"))
def _bedrock_target_keeps_item_marks(model: str, drop_params: bool) -> bool:
declared: Final = bedrock_declared_cache_point_support(model)
return declared is True or (declared is None and not drop_params)
def _target_keeps_item_marks(model: str, custom_llm_provider: str | None, drop_params: bool) -> bool:
if supports_anthropic_cache_control(model, custom_llm_provider):
return True
if custom_llm_provider in _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL:
return True
if custom_llm_provider == "openrouter":
return openrouter_model_supports_cache_control(model)
if custom_llm_provider == "snowflake":
return is_snowflake_claude_model(model)
if custom_llm_provider == "databricks":
return "claude" in model.lower()
return custom_llm_provider == "bedrock" and _bedrock_target_keeps_item_marks(model, drop_params)
def _refuses_marks_on_tool_blocks(model: str, custom_llm_provider: str | None) -> bool:
return custom_llm_provider == "bedrock" and BaseAWSLLM.get_bedrock_invoke_provider(model) == "nova"
def _with_cache_control_on_last_block(
blocks: Sequence[object], cache_control: ChatCompletionCachedContent | None
) -> Sequence[object]:
if cache_control is None or not blocks or not isinstance(blocks[-1], dict):
return blocks
last_block: Final = _STR_KEY_DICT_ADAPTER.validate_python(blocks[-1])
if last_block.get("cache_control") is not None:
return blocks
return [*blocks[:-1], {**last_block, "cache_control": cache_control}]
def _content_blocks_carrying(content: object, cache_control: ChatCompletionCachedContent | None) -> Sequence[object]:
if isinstance(content, str):
return (
ChatCompletionTextObject(type="text", text=content)
if cache_control is None
else ChatCompletionTextObject(type="text", text=content, cache_control=cache_control),
)
return _with_cache_control_on_last_block(_OBJECT_LIST_ADAPTER.validate_python(content), cache_control)
_MARK_CARRYING_MESSAGE_FIELDS: Final = ("content", "tool_calls")
def _mark_carrying_blocks(fields: Mapping[str, object]) -> Sequence[object]:
lists: Final = tuple(fields.get(key) for key in _MARK_CARRYING_MESSAGE_FIELDS)
return tuple(
chain.from_iterable(_OBJECT_LIST_ADAPTER.validate_python(value) for value in lists if isinstance(value, list))
)
def _block_carries_cache_control(block: object) -> bool:
return isinstance(block, dict) and _STR_KEY_DICT_ADAPTER.validate_python(block).get("cache_control") is not None
def _cache_control_count(message: object) -> int:
if not isinstance(message, dict):
return 0
fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
block_marks: Final = sum(1 for block in _mark_carrying_blocks(fields) if _block_carries_cache_control(block))
return int(fields.get("cache_control") is not None) + block_marks
_CacheControlRewrite: TypeAlias = Callable[[object], ChatCompletionCachedContent | None]
def _fields_with_cache_control(fields: Mapping[str, object], rewrite: _CacheControlRewrite) -> Mapping[str, object]:
others: Final = {key: value for key, value in fields.items() if key != "cache_control"}
if "cache_control" not in fields:
return others
cache_control: Final = rewrite(fields["cache_control"])
return others if cache_control is None else {**others, "cache_control": cache_control}
def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> object:
if not isinstance(block, dict):
return block
return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite)
_MessageT = TypeVar("_MessageT")
def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT:
if not isinstance(message, dict):
return message
fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
blocks_rewritten: Final = {
key: (
[_block_with_cache_control(block, rewrite) for block in _OBJECT_LIST_ADAPTER.validate_python(value)]
if key in _MARK_CARRYING_MESSAGE_FIELDS and isinstance(value, list)
else value
)
for key, value in fields.items()
}
return cast( # cast-ok: the same message shape with each cache marker rewritten or dropped
_MessageT, _fields_with_cache_control(blocks_rewritten, rewrite)
)
def _kept_within_budget(state: tuple[int, bool], marks: int) -> tuple[int, bool]:
remaining, _ = state
if marks > remaining:
return remaining, False
return remaining - marks, True
def _messages_within_cache_control_limit(messages: Sequence[_MessageT], reserved_marks: int) -> Sequence[_MessageT]:
budget: Final = max(MAX_CACHE_CONTROL_BLOCKS - reserved_marks, 0)
newest_first: Final = tuple(reversed(messages))
budget_after_each: Final = accumulate(
(_cache_control_count(message) for message in newest_first), _kept_within_budget, initial=(budget, True)
)
kept_flags: Final = tuple(kept for _, kept in islice(budget_after_each, 1, None))
kept_newest_first: Final = tuple(
message if kept else _with_cache_control(message, _dropped_cache_control)
for message, kept in zip(newest_first, kept_flags)
)
return tuple(reversed(kept_newest_first))
_SYSTEM_MESSAGE_FIELDS: Final = frozenset(("content", "cache_control"))
def _marked_developer_message_fields(message: object) -> Mapping[str, object] | None:
if not isinstance(message, dict):
return None
fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
return fields if fields.get("role") == "developer" and "cache_control" in fields else None
def _developer_message_as_system_message(message: _MessageT) -> _MessageT:
fields: Final = _marked_developer_message_fields(message)
if fields is None:
return message
return cast( # cast-ok: the system message the chat route maps a developer message to, mark included
_MessageT, {"role": "system", **{key: value for key, value in fields.items() if key in _SYSTEM_MESSAGE_FIELDS}}
)
def _message_fields(message: object) -> Mapping[str, object] | None:
return _STR_KEY_DICT_ADAPTER.validate_python(message) if isinstance(message, dict) else None
def _without_tool_block_marks(message: _MessageT) -> _MessageT:
fields: Final = _message_fields(message)
if fields is None:
return message
if fields.get("role") == "tool":
return _with_cache_control(message, _dropped_cache_control)
tool_calls: Final = fields.get("tool_calls")
if not isinstance(tool_calls, list):
return message
return cast( # cast-ok: the same assistant message with its tool calls unmarked
_MessageT,
{
**fields,
"tool_calls": [
_block_with_cache_control(tool_call, _dropped_cache_control)
for tool_call in _OBJECT_LIST_ADAPTER.validate_python(tool_calls)
],
},
)
@dataclass(frozen=True, slots=True)
class _ItemMarkRules:
drop_params: bool
refuses_tool_block_marks: bool
external_marks: int
def sanitized(self, message: _MessageT) -> _MessageT:
valid: Final = _with_cache_control(message, _validated_cache_control)
return _without_tool_block_marks(valid) if self.refuses_tool_block_marks else valid
def checked(self, message: _MessageT) -> _MessageT:
return self.sanitized(message) if self.drop_params else message
def _item_mark_rules(
model: str, custom_llm_provider: str | None, tools: Iterable[object], request_kwargs: Mapping[str, object]
) -> _ItemMarkRules:
return _ItemMarkRules(
drop_params=drops_unsupported_params(request_kwargs.get("drop_params")),
refuses_tool_block_marks=_refuses_marks_on_tool_blocks(model, custom_llm_provider),
external_marks=AnthropicCacheControlHook.count_external_cache_breakpoints(
tools, request_kwargs.get("cache_control"), request_kwargs
),
)
def _request_messages_for_cache_points(messages: Sequence[_MessageT], rules: _ItemMarkRules) -> Sequence[_MessageT]:
as_system: Final = tuple(_developer_message_as_system_message(message) for message in messages)
if not rules.drop_params:
return as_system
return _messages_within_cache_control_limit(tuple(map(rules.checked, as_system)), rules.external_marks)
def _replayed_messages_for_cache_points(
replayed_messages: Sequence[_MessageT], current_messages: Sequence[object], rules: _ItemMarkRules
) -> Sequence[_MessageT]:
as_system: Final = tuple(_developer_message_as_system_message(message) for message in replayed_messages)
current_marks: Final = sum(_cache_control_count(message) for message in current_messages)
return _messages_within_cache_control_limit(
tuple(map(rules.sanitized, as_system)), current_marks + rules.external_marks
)
class ChatCompletionSession(TypedDict, total=False):
messages: list[
AllMessageValues
@ -482,11 +755,27 @@ class LiteLLMCompletionResponsesConfig:
api_base=kwargs.get("api_base"),
)
request_kwargs: Final = _STR_KEY_DICT_ADAPTER.validate_python(kwargs)
keeps_item_marks: Final = _target_keeps_item_marks(
model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params"))
)
bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input,
responses_api_request=responses_api_request,
replay_reasoning=True,
item_cache_control=keeps_item_marks,
)
litellm_completion_request: dict = {
"messages": LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input,
responses_api_request=responses_api_request,
replay_reasoning=True,
"messages": (
list(
_request_messages_for_cache_points(
bridged_messages,
_item_mark_rules(model, custom_llm_provider, tools, request_kwargs),
)
)
if keeps_item_marks
else bridged_messages
),
"model": model,
"tool_choice": LiteLLMCompletionResponsesConfig._transform_tool_choice(
@ -535,6 +824,7 @@ class LiteLLMCompletionResponsesConfig:
input: str | ResponseInputParam,
responses_api_request: ResponsesAPIOptionalRequestParams | dict,
replay_reasoning: bool = False,
item_cache_control: bool = False,
) -> list[
AllMessageValues
| GenericChatCompletionMessage
@ -549,11 +839,13 @@ class LiteLLMCompletionResponsesConfig:
sent to a model: prior-turn ``reasoning`` items are then rebuilt as
assistant ``reasoning_content`` and signed ``thinking_blocks`` so the
provider gets its own chain-of-thought back instead of reading it as
visible text.
visible text. ``item_cache_control`` belongs to the same callers, when
the target keeps an input item's ``cache_control`` on the message.
Callers that only inspect the messages (token counting, rate limiting,
guardrail scanning) leave it off, because they need every piece of text
in the request to stay readable as message ``content``.
guardrail scanning) leave both off, because they need every piece of text
in the request to stay readable as message ``content``, with no item mark
a strict message schema could refuse.
"""
messages: list[
AllMessageValues
@ -571,7 +863,7 @@ class LiteLLMCompletionResponsesConfig:
messages.extend(
LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message(
input=input,
input=input if item_cache_control else without_item_cache_control(input),
replay_reasoning=replay_reasoning,
)
)
@ -583,19 +875,44 @@ class LiteLLMCompletionResponsesConfig:
previous_response_id: str,
litellm_completion_request: dict,
instructions: str | None,
session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None,
request_kwargs: Mapping[str, object] = MappingProxyType({}),
) -> dict:
"""
Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages
"""
chat_completion_session = ChatCompletionSession(messages=[], litellm_session_id=None)
if previous_response_id:
chat_completion_session = (
await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
previous_response_id=previous_response_id,
)
)
request_fields: Final = _STR_KEY_DICT_ADAPTER.validate_python(litellm_completion_request)
requested_provider: Final = request_fields.get("custom_llm_provider")
model: Final = str(request_fields.get("model") or "")
provider: Final = requested_provider if isinstance(requested_provider, str) else None
keeps_item_marks: Final = _target_keeps_item_marks(
model, provider, drops_unsupported_params(request_kwargs.get("drop_params"))
)
lookup: Final = session_lookup or partial(
ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id,
item_cache_control=keeps_item_marks,
)
chat_completion_session: Final = (
await lookup(previous_response_id)
if previous_response_id
else ChatCompletionSession(messages=[], litellm_session_id=None)
)
_messages: Final = litellm_completion_request.get("messages") or []
session_messages: Final = chat_completion_session.get("messages") or []
replayed_messages: Final = chat_completion_session.get("messages") or []
session_messages: Final = (
_replayed_messages_for_cache_points(
replayed_messages,
_OBJECT_LIST_ADAPTER.validate_python(_messages),
_item_mark_rules(
model,
provider,
_OBJECT_LIST_ADAPTER.validate_python(request_fields.get("tools") or []),
request_kwargs,
),
)
if keeps_item_marks
else replayed_messages
)
instructions_end: Final = 1 if instructions else 0
carried_instructions: Final = None if instructions else chat_completion_session.get("instructions")
leading_system_messages: Final = (
@ -609,7 +926,7 @@ class LiteLLMCompletionResponsesConfig:
# Store original _messages before combining for safety check
original_new_messages: Final = _messages.copy() if _messages else []
combined_messages = leading_system_messages + session_messages + _messages[instructions_end:]
combined_messages = [*leading_system_messages, *session_messages, *_messages[instructions_end:]]
# Fix: Ensure tool_results have corresponding tool_calls in previous assistant message
# Pass tools parameter to help reconstruct tool_calls if not in cache
@ -631,7 +948,7 @@ class LiteLLMCompletionResponsesConfig:
elif session_messages:
# If we had session messages but they got filtered out,
# restore them
combined_messages = session_messages
combined_messages = list(session_messages)
else:
# Both are empty - this likely means function_call_output had empty/invalid call_id
# Provide a helpful error message
@ -955,22 +1272,23 @@ class LiteLLMCompletionResponsesConfig:
if new_content is None:
return None
previous_content: Final = last_message.get("content")
previous_cache_control: Final = _input_item_cache_control(last_message)
new_cache_control: Final = _input_item_cache_control(new_message)
content: Final = (
new_content
if not previous_content
if not previous_content and isinstance(new_content, str)
else [
block
for value in (previous_content, new_content)
for block in (
(ChatCompletionTextObject(type="text", text=value),)
if isinstance(value, str)
else _OBJECT_LIST_ADAPTER.validate_python(value)
)
*(_content_blocks_carrying(previous_content, previous_cache_control) if previous_content else ()),
*_content_blocks_carrying(new_content, new_cache_control),
]
)
message_cache_control: Final = (
None if isinstance(content, list) else new_cache_control or previous_cache_control
)
merged: Final = {
**last_message,
**{key: value for key, value in last_message.items() if key != "cache_control"},
"content": content,
**({} if message_cache_control is None else {"cache_control": message_cache_control}),
}
return cast( # cast-ok: preserves the assistant fields and content blocks
ChatCompletionAssistantMessage, merged
@ -1468,14 +1786,24 @@ class LiteLLMCompletionResponsesConfig:
# Since guardrails skip None content anyway, we return empty list to exclude it from structured messages
if content is None:
return []
return [
GenericChatCompletionMessage(
role=_input_item_role(input_item),
content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(
content
),
)
]
return [LiteLLMCompletionResponsesConfig._transform_responses_api_message_item(input_item, content)]
@staticmethod
def _transform_responses_api_message_item(
input_item: Mapping[str, object], content: object
) -> GenericChatCompletionMessage:
role: Final = _input_item_role(input_item)
content_value: Final = (
LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content(content)
)
cache_control: Final = _input_item_cache_control(input_item)
if cache_control is None:
return GenericChatCompletionMessage(role=role, content=content_value)
if isinstance(content_value, str):
return GenericChatCompletionMessage(role=role, content=content_value, cache_control=cache_control)
return GenericChatCompletionMessage(
role=role, content=list(_with_cache_control_on_last_block(content_value, cache_control))
)
@staticmethod
def _reasoning_text_from_content(input_item: Mapping[str, object]) -> str | None:
@ -1665,10 +1993,14 @@ class LiteLLMCompletionResponsesConfig:
except Exception:
return str(output)
tool_output_message: Final = ChatCompletionToolMessage(
role="tool",
content=_normalize_function_call_output_to_tool_content(tool_call_output.get("output")),
tool_call_id=str(call_id),
tool_content: Final = _normalize_function_call_output_to_tool_content(tool_call_output.get("output"))
output_cache_control: Final = _input_item_cache_control(tool_call_output)
tool_output_message: Final = (
ChatCompletionToolMessage(role="tool", content=tool_content, tool_call_id=str(call_id))
if output_cache_control is None
else _MarkedToolMessage(
role="tool", content=tool_content, tool_call_id=str(call_id), cache_control=output_cache_control
)
)
_tool_use_definition: Final = TOOL_CALLS_CACHE.get_cache(
@ -1750,24 +2082,26 @@ class LiteLLMCompletionResponsesConfig:
raw_name: Final = function_call.get("name") or ""
namespace: Final = function_call.get("namespace") or ""
qualify: Final = bool(namespace) and function_call.get("type") != "custom_tool_call"
tool_call: Final = ChatCompletionToolCallChunk(
id=function_call.get("call_id") or function_call.get("id") or "",
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=f"{namespace}__{raw_name}" if qualify else raw_name,
arguments=serialize_tool_call_arguments(raw_arguments),
),
index=0,
tool_call_id: Final = function_call.get("call_id") or function_call.get("id") or ""
tool_call_function: Final = ChatCompletionToolCallFunctionChunk(
name=f"{namespace}__{raw_name}" if qualify else raw_name,
arguments=serialize_tool_call_arguments(raw_arguments),
)
function_call_cache_control: Final = _input_item_cache_control(function_call)
tool_call: Final = (
ChatCompletionToolCallChunk(id=tool_call_id, type="function", function=tool_call_function, index=0)
if function_call_cache_control is None
else _MarkedToolCallChunk(
id=tool_call_id,
type="function",
function=tool_call_function,
index=0,
cache_control=function_call_cache_control,
)
)
# Create an assistant message with the tool call
chat_completion_response_message: Final = ChatCompletionResponseMessage(
tool_calls=[tool_call],
role="assistant",
content=None, # Function calls don't have content
)
return [chat_completion_response_message]
return [ChatCompletionResponseMessage(tool_calls=[tool_call], role="assistant", content=None)]
@staticmethod
def _resolve_file_id(item: Mapping[str, object]) -> object:

60
litellm/responses/main.py
View file

@ -10,7 +10,10 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeA
import httpx
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import assert_never
from typing_extensions import (
TypeIs, # noqa: TID251 # narrows a client input item before its mark is read
assert_never,
)
import litellm
from litellm._logging import verbose_logger
@ -568,12 +571,54 @@ def _will_bridge_to_chat_completions(
)
def _stamp_injection_points_with_dialect(
def _is_a_client_input_item(
item: object,
) -> TypeIs[Mapping[str, object]]: # guard-ok: bounds a client input item before its mark is read
return isinstance(item, Mapping)
_ITEM_TYPES_THE_BRIDGE_FORWARDS_MARKS_FROM: Final = frozenset(("function_call", "function_call_output"))
def _item_carries_a_mark_outside_the_provisional_messages(item: object) -> bool:
return (
_is_a_client_input_item(item)
and item.get("type") in _ITEM_TYPES_THE_BRIDGE_FORWARDS_MARKS_FROM
and bool(item.get("cache_control"))
)
def _function_tool_carries_a_mark(tool: Mapping[str, object]) -> bool:
return isinstance(tool, dict) and tool.get("type") == "function" and bool(tool.get("cache_control"))
def _cache_control_marks_outside_the_provisional_messages(
input: str | ResponseInputParam | None, tools: Iterable[Mapping[str, object]] | None
) -> int:
"""Client marks on the request parts the provisional message list leaves out.
A ``function_call`` or ``function_call_output`` item has no role, so the hook's pass here
never sees its mark, while the bridge forwards it onto the tool call or tool message it
builds, where the provider counts it against the cap. A ``function`` tool's mark reaches
the provider the same way, on the chat tool the bridge builds from it.
"""
item_marks: Final = (
0
if input is None or isinstance(input, str)
else sum(1 for item in input if _item_carries_a_mark_outside_the_provisional_messages(item))
)
tool_marks: Final = sum(1 for tool in tools or () if _function_tool_carries_a_mark(tool))
return item_marks + tool_marks
def _stamp_injection_points_for_the_hook(
kwargs: dict[str, object], # mutable-ok: the points are rewritten in the caller's own kwargs for the hook to read
model: str,
custom_llm_provider: str | None,
input: str | ResponseInputParam | None,
tools: Iterable[Mapping[str, object]] | None,
) -> None:
"""Carry the provider this layer resolved onto the points.
"""Carry onto the points the provider this layer resolved and the marks its pass cannot see.
The hook reads ``custom_llm_provider`` from the request kwargs, which never hold the one
resolved here, and resolving the model name alone reads a Foundry deployment of an OpenAI
@ -582,8 +627,9 @@ def _stamp_injection_points_with_dialect(
points: Final = configured_injection_points(kwargs.get("cache_control_injection_points"))
if not points:
return
kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_with_dialect(
kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_for_prompt_hook(
points,
_cache_control_marks_outside_the_provisional_messages(input, tools),
model,
custom_llm_provider,
kwargs.get("api_base") or kwargs.get("base_url"),
@ -704,7 +750,7 @@ async def aresponses(
_api_base_kwarg(kwargs),
),
):
_stamp_injection_points_with_dialect(kwargs, model, custom_llm_provider)
_stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools)
(
model,
merged_input,
@ -841,6 +887,7 @@ def _resolve_prompt_swapped_provider(
def _apply_prompt_management_to_responses_call(
input: str | ResponseInputParam,
tools: Iterable[ToolParam] | None,
model: str,
custom_llm_provider: str | None,
litellm_logging_obj: LiteLLMLoggingObj | None,
@ -875,7 +922,7 @@ def _apply_prompt_management_to_responses_call(
_api_base_kwarg(kwargs),
),
):
_stamp_injection_points_with_dialect(kwargs, model, custom_llm_provider)
_stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools)
(
model,
merged_input,
@ -1274,6 +1321,7 @@ def responses(
input, model, custom_llm_provider, deployment_model_info = _apply_prompt_management_to_responses_call(
input=input,
tools=tools,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_logging_obj=litellm_logging_obj,

19
litellm/responses/streaming_iterator.py
View file

@ -282,6 +282,21 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro
return mapped_exception
def _event_field(event: object, name: str) -> object:
return event.get(name) if _is_json_object(event) else getattr(event, name, None)
def _is_error_field_value(value: object) -> bool:
return isinstance(value, (str, int)) and not isinstance(value, bool)
def _error_event_error_object(event: object) -> object:
nested_error: Final = _event_field(event, "error")
if nested_error is not None:
return nested_error
return {field: value for field in ("message", "code") if _is_error_field_value(value := _event_field(event, field))}
def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool:
if isinstance(mapped_exception, litellm.ContentPolicyViolationError):
return True
@ -659,7 +674,7 @@ class BaseResponsesAPIStreamingIterator:
error_obj: Final[object] = (
getattr(getattr(result, "response", None), "error", None)
if chunk_type == "response.failed"
else getattr(result, "error", None)
else _error_event_error_object(result)
)
mapped_exception: Final = self._map_error_event_exception(error_obj)
@ -1837,7 +1852,7 @@ _RESPONSES_WS_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.
def _ws_event_error(event: Mapping[str, object]) -> object:
if event.get("type") == "error":
return event.get("error")
return _error_event_error_object(event)
response: Final = event.get("response")
return response.get("error") if _is_json_object(response) else None

101
litellm/router.py
View file

@ -23,6 +23,7 @@ import weakref
from collections import defaultdict
from collections.abc import (
AsyncGenerator,
AsyncIterable,
AsyncIterator,
Awaitable,
Callable,
@ -218,6 +219,7 @@ from litellm.router_utils.cooldown_handlers import (
from litellm.router_utils.fallback_event_handlers import (
MID_STREAM_FALLBACK_CONTROLS_KEY,
AttemptedFallbackTargets,
attempted_fallback_targets_from,
attempted_retries_for_request,
carry_over_pre_routing_selection,
carry_over_routed_deployment,
@ -229,6 +231,7 @@ from litellm.router_utils.fallback_event_handlers import (
get_fallback_model_group_for_lookup_groups,
get_pre_routing_selection,
has_unattempted_fallback_target,
is_bad_request_without_fallback,
mid_stream_fallback_hop_kwargs,
mid_stream_fallback_snapshot_kwargs,
mid_stream_retry_kwargs,
@ -691,6 +694,67 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception:
return error.original_exception if error.original_exception is not None else error
@dataclass(frozen=True, slots=True)
class _FirstChunk:
"""The chunk Router._acompletion / _completion already read off a stream; None means the provider closed it empty."""
chunk: ModelResponseStream | None
async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk:
"""
Read the first chunk while Router._acompletion still holds the request, so a
stream that fails before any token raises there, inside async_function_with_retries
and the fallback chain, instead of surfacing to the caller mid-iteration.
"""
from litellm.exceptions import MidStreamFallbackError
try:
return _FirstChunk(chunk=await stream.__anext__())
except StopAsyncIteration:
return _FirstChunk(chunk=None)
except MidStreamFallbackError as e:
raise _mid_stream_retry_trigger(e) from e
def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk:
"""Sync twin of _aread_first_chunk for Router._completion."""
from litellm.exceptions import MidStreamFallbackError
try:
return _FirstChunk(chunk=next(stream))
except StopIteration:
return _FirstChunk(chunk=None)
except MidStreamFallbackError as e:
raise _mid_stream_retry_trigger(e) from e
def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk | None) -> Iterator[object]:
"""The stream with the chunk _read_first_chunk already read put back in front (None: the stream as is)."""
if taken is None:
yield from model_response
return
if taken.chunk is None:
return
yield taken.chunk
yield from model_response
async def _achunks_with_first(
model_response: AsyncIterable[object], taken: _FirstChunk | None
) -> AsyncIterator[object]:
"""Async twin of _chunks_with_first."""
if taken is None:
async for item in model_response:
yield item
return
if taken.chunk is None:
return
yield taken.chunk
async for item in model_response:
yield item
def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
"""
Whether `chunk` should make Router._aanthropic_messages_streaming_iterator
@ -960,6 +1024,7 @@ class Router:
allowed_fails_policy: AllowedFailsPolicy | None = None, # set custom allowed fails policy
cooldown_time: float | None = None, # (seconds) time to cooldown a deployment after failure
disable_cooldowns: bool | None = None,
fallback_on_unsupported_params: bool = True,
routing_strategy: RoutingStrategyName = "simple-shuffle",
optional_pre_call_checks: OptionalPreCallChecks | None = None,
routing_strategy_args: dict = {}, # just for latency-based
@ -1202,6 +1267,7 @@ class Router:
self.cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS
self.cooldown_cache = CooldownCache(cache=self.cache, default_cooldown_time=self.cooldown_time)
self.disable_cooldowns = disable_cooldowns
self.fallback_on_unsupported_params = fallback_on_unsupported_params
self.enable_health_check_routing = enable_health_check_routing
self.enable_weighted_failover = enable_weighted_failover
self.health_check_ignore_transient_errors = health_check_ignore_transient_errors
@ -2791,6 +2857,7 @@ class Router:
model_response=response,
messages=messages,
initial_kwargs=input_kwargs_for_streaming_fallback,
first_chunk=_read_first_chunk(response),
)
return response
@ -3050,6 +3117,7 @@ class Router:
messages: list[dict[str, str]],
initial_kwargs: dict,
deployment_slot: contextlib.AsyncExitStack | None = None,
first_chunk: _FirstChunk | None = None,
) -> CustomStreamWrapper:
"""
Helper to iterate over a streaming response.
@ -3058,6 +3126,8 @@ class Router:
`deployment_slot` holds the deployment's max_parallel_requests semaphore; it is
released when the stream is exhausted, closed, or falls back to another deployment
`first_chunk` carries the chunk _acompletion already read (None: iterate from the start)
"""
from litellm.exceptions import MidStreamFallbackError
@ -3102,7 +3172,7 @@ class Router:
async def stream_with_fallbacks():
fallback_response = None # Track for cleanup in finally
try:
async for item in model_response:
async for item in _achunks_with_first(model_response, first_chunk):
yield item
except MidStreamFallbackError as e:
with anyio.CancelScope(shield=True):
@ -3204,7 +3274,7 @@ class Router:
wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
# weak, so the generator closing over it does not keep the wrapper out of
# refcount teardown and delay the `finally` that releases the deployment slot
wrapper_ref: Final = weakref.ref(wrapped_response)
wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response)
return wrapped_response
@staticmethod
@ -3677,6 +3747,7 @@ class Router:
model_response: CustomStreamWrapper,
messages: list[dict[str, str]],
initial_kwargs: dict,
first_chunk: _FirstChunk | None = None,
) -> CustomStreamWrapper:
"""
Sync equivalent of _acompletion_streaming_iterator.
@ -3712,7 +3783,7 @@ class Router:
def stream_with_fallbacks():
fallback_response = None
try:
for item in model_response:
for item in _chunks_with_first(model_response, first_chunk):
yield item
except MidStreamFallbackError as e:
if fallbacks_disabled_for_request(initial_kwargs) or (
@ -3808,7 +3879,7 @@ class Router:
wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
# weak, for the same reason as the async twin
wrapper_ref: Final = weakref.ref(wrapped_response)
wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response)
return wrapped_response
async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
@ -3959,6 +4030,9 @@ class Router:
and response.make_call is not None
):
await response.fetch_stream()
first_chunk: Final = (
await _aread_first_chunk(response) if isinstance(response, CustomStreamWrapper) else None
)
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -3975,6 +4049,7 @@ class Router:
messages=messages,
initial_kwargs=input_kwargs_for_streaming_fallback,
deployment_slot=deployment_slot.pop_all(),
first_chunk=first_chunk,
)
return response
@ -7649,6 +7724,7 @@ class Router:
original_exception: Final = e
fallback_model_group = None
original_model_group: Final[str | None] = kwargs.get("model")
fallback_walk: Final = attempted_fallback_targets_from(kwargs) # pyright: ignore[reportUnknownArgumentType] # keep legacy kwargs untyped
# A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier
# behind the router name, and fallbacks are configured per tier, not per router.
lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group)
@ -7657,13 +7733,22 @@ class Router:
hop_depth: Final = kwargs.get("fallback_depth")
nested_fallback_hop: Final = isinstance(hop_depth, int) and hop_depth > 0
if disable_fallbacks is True or original_model_group is None or is_guardrail_intervention(e):
unsupported_params_without_fallback: Final = is_bad_request_without_fallback(
e, self.fallback_on_unsupported_params
)
if (
disable_fallbacks is True
or original_model_group is None
or is_guardrail_intervention(e)
or unsupported_params_without_fallback
):
raise e
input_kwargs: Final = {
"litellm_router": self,
"original_exception": original_exception,
**kwargs,
"attempted_targets": fallback_walk,
}
if "max_fallbacks" not in input_kwargs:
@ -7893,7 +7978,10 @@ class Router:
and not nested_fallback_hop
):
original_exception.message += format_fallback_outcome_message(
model_group, attempted_fallback_group, fallback_failure_exception_str
model_group,
attempted_fallback_group,
fallback_failure_exception_str,
skipped=fallback_walk.skipped,
)
raise original_exception
@ -12731,6 +12819,7 @@ class Router:
"enable_weighted_failover",
"enable_tag_filtering",
"tag_routing_prefix",
"fallback_on_unsupported_params",
]
for var in vars_to_include:

54
litellm/router_utils/common_utils.py
View file

@ -15,7 +15,7 @@ from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
from litellm.types.router import CredentialLiteLLMParams
from litellm.types.router import CredentialLiteLLMParams, FallbackSkipReason, SkippedFallbackTarget
from litellm.types.utils import LlmProviders
_V = TypeVar("_V")
@ -134,15 +134,61 @@ def format_fallback_outcome_message(
model_group: str | None,
fallback_model_group: Sequence[object] | None,
fallback_failure_detail: str,
skipped: Sequence[SkippedFallbackTarget] = (),
) -> str:
"""User-facing explanation appended when the fallback orchestrator gives up and re-raises the primary error."""
lead: Final = f"\n\nLiteLLM: model group '{model_group}' failed with the error above."
if not fallback_model_group:
return f"{lead} No fallback was attempted."
def fallback_target_model_group(target: object) -> str | None:
if isinstance(target, str):
return target
if not isinstance(target, Mapping):
return None
model: Final = target.get("model")
return model if isinstance(model, str) else None
skipped_targets: Final = frozenset(target.target for target in skipped)
attempted_targets: Final = tuple(
target for target in fallback_model_group if fallback_target_model_group(target) not in skipped_targets
)
not_authorized_targets: Final = tuple(
dict.fromkeys(target.target for target in skipped if target.reason is FallbackSkipReason.NOT_AUTHORIZED)
)
over_budget_targets: Final = tuple(
dict.fromkeys(target.target for target in skipped if target.reason is FallbackSkipReason.OVER_BUDGET)
)
attempted_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in attempted_targets)
not_authorized_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in not_authorized_targets)
over_budget_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in over_budget_targets)
attempted_failure: Final = (
(f"Fallback to {attempted_text} also failed: {fallback_failure_detail}",)
if fallback_failure_detail and attempted_targets
else ()
)
not_authorized_sentence: Final = (
(
f"Fallback to {not_authorized_text} was skipped because the caller is not authorized to call "
f"{'it' if len(not_authorized_targets) == 1 else 'them'}.",
)
if not_authorized_targets
else ()
)
over_budget_sentence: Final = (
(
f"Fallback to {over_budget_text} was skipped because the caller is over budget for "
f"{'it' if len(over_budget_targets) == 1 else 'them'}.",
)
if over_budget_targets
else ()
)
sentences: Final = attempted_failure + not_authorized_sentence + over_budget_sentence
if sentences:
return f"{lead} {' '.join(sentences)}"
targets: Final = ", ".join(str(mask_sensitive_structure(target)) for target in fallback_model_group)
if not fallback_failure_detail:
return f"{lead} Fallback model group(s) configured: {targets}."
return f"{lead} Fallback to {targets} also failed: {fallback_failure_detail}"
return f"{lead} Fallback model group(s) configured: {targets}."
def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str:

65
litellm/router_utils/fallback_event_handlers.py
View file

@ -1,7 +1,7 @@
import hashlib
import json
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from types import MappingProxyType
@ -29,7 +29,12 @@ from litellm.router_utils.cooldown_handlers import (
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
increment_deployment_failures_for_current_minute,
)
from litellm.types.router import LiteLLMParamsTypedDict, reject_server_owned_wif_params
from litellm.types.router import (
FallbackSkipReason,
LiteLLMParamsTypedDict,
SkippedFallbackTarget,
reject_server_owned_wif_params,
)
if TYPE_CHECKING:
from litellm.router import Router as _Router
@ -202,6 +207,7 @@ class AttemptedFallbackTargets:
"""
keys: frozenset[str] = frozenset()
_skip_reasons: dict[str, FallbackSkipReason] = field(default_factory=dict) # mutable-ok: O(1) per caller-sized list
def __contains__(self, key: str) -> bool:
return key in self.keys
@ -209,6 +215,20 @@ class AttemptedFallbackTargets:
def record(self, key: str) -> None:
self.keys = self.keys | frozenset((key,))
@property
def skipped(self) -> tuple[SkippedFallbackTarget, ...]:
return tuple(
SkippedFallbackTarget(target=target, reason=reason) for target, reason in self._skip_reasons.items()
)
def record_skip(self, target: str, reason: FallbackSkipReason) -> None:
self._skip_reasons.setdefault(target, reason)
def attempted_fallback_targets_from(kwargs: Mapping[str, object]) -> AttemptedFallbackTargets:
carried_targets: Final[object] = kwargs.get("attempted_targets")
return carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets()
def has_unattempted_fallback_target(
fallback_model_group: Sequence[object] | None, kwargs: Mapping[str, object]
@ -222,6 +242,16 @@ def has_unattempted_fallback_target(
return any((key := fallback_attempt_key(target)) is None or key not in attempted for target in fallback_model_group)
def is_bad_request_without_fallback(exception: Exception, fallback_on_unsupported_params: bool) -> bool:
if fallback_on_unsupported_params or not isinstance(exception, litellm.BadRequestError):
return False
if isinstance(exception, litellm.UnsupportedParamsError):
return True
if isinstance(exception, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError)):
return False
return exception.status_code == 400 and bool(getattr(exception, "llm_provider", ""))
def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool:
"""
Handles wildcard routing scenario
@ -514,7 +544,7 @@ def fallback_lookup_groups(kwargs: Mapping[str, object], model_group: str | None
def _resolved_a_specific_chain(
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape
result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract
) -> bool:
resolved, generic_idx = result
@ -524,7 +554,7 @@ def _resolved_a_specific_chain(
def get_fallback_model_group_for_lookup_groups(
fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape
lookup_groups: tuple[str, ...],
) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
"""
@ -538,7 +568,10 @@ def get_fallback_model_group_for_lookup_groups(
return next((result for result in results if result[0] is not None), (None, None))
def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]:
def get_fallback_model_group(
fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape
model_group: str,
) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
"""
Returns:
- fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"]
@ -607,6 +640,8 @@ async def _is_fallback_target_authorized(
fallback_entry: str | Mapping[str, object],
original_model_group: str,
kwargs: Mapping[str, object],
*,
attempted: AttemptedFallbackTargets,
) -> bool:
access_check: Final = litellm_router.fallback_access_check
target: Final = _get_fallback_target_model_group(fallback_entry)
@ -623,6 +658,7 @@ async def _is_fallback_target_authorized(
"Skipping fallback to model_group = %s: caller is not authorized to call it",
mask_sensitive_structure(fallback_entry),
)
attempted.record_skip(target, FallbackSkipReason.NOT_AUTHORIZED)
return False
@ -631,6 +667,8 @@ async def _is_fallback_target_within_budget(
fallback_entry: str | Mapping[str, object],
original_model_group: str,
kwargs: Mapping[str, object],
*,
attempted: AttemptedFallbackTargets,
) -> bool:
budget_check: Final = litellm_router.fallback_budget_check
target: Final = _get_fallback_target_model_group(fallback_entry)
@ -642,6 +680,7 @@ async def _is_fallback_target_within_budget(
"Skipping fallback to model_group = %s: caller is over budget",
mask_sensitive_structure(fallback_entry),
)
attempted.record_skip(target, FallbackSkipReason.OVER_BUDGET)
return False
@ -730,13 +769,7 @@ async def run_async_fallback(
same_model_group_only: Final = references_provider_scoped_resource(kwargs) or creates_provider_scoped_resource(
kwargs
)
# Read out of kwargs and narrowed here rather than declared as a parameter: every caller
# reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter
# would carry an annotation that no call site can actually be checked against.
carried_targets: Final = kwargs.get("attempted_targets")
attempted: Final = (
carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets()
)
attempted: Final = attempted_fallback_targets_from(kwargs) # pyright: ignore[reportUnknownArgumentType] # keep legacy kwargs untyped
failed_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group
attempted.record(failed_model_group)
@ -758,9 +791,13 @@ async def run_async_fallback(
original_model_group,
)
continue
if not await _is_fallback_target_authorized(litellm_router, mg, original_model_group, kwargs):
if not await _is_fallback_target_authorized(
litellm_router, mg, original_model_group, kwargs, attempted=attempted
):
continue
if not await _is_fallback_target_within_budget(litellm_router, mg, original_model_group, kwargs):
if not await _is_fallback_target_within_budget(
litellm_router, mg, original_model_group, kwargs, attempted=attempted
):
continue
attempt_key = fallback_attempt_key(mg)
if attempt_key is not None:

7
litellm/types/access_group.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping, Sequence
from datetime import datetime
from litellm.types.llms.base import LiteLLMBaseModel
@ -8,7 +9,9 @@ class AccessGroupCreateRequest(LiteLLMBaseModel):
description: str | None = None
access_model_names: list[str] | None = None
access_mcp_server_ids: list[str] | None = None
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str] | None = None
access_skill_names: tuple[str, ...] | None = None
assigned_team_ids: list[str] | None = None
assigned_key_ids: list[str] | None = None
@ -18,7 +21,9 @@ class AccessGroupUpdateRequest(LiteLLMBaseModel):
description: str | None = None
access_model_names: list[str] | None = None
access_mcp_server_ids: list[str] | None = None
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str] | None = None
access_skill_names: tuple[str, ...] | None = None
assigned_team_ids: list[str] | None = None
assigned_key_ids: list[str] | None = None
@ -36,7 +41,9 @@ class AccessGroupResponse(LiteLLMBaseModel):
description: str | None = None
access_model_names: list[str]
access_mcp_server_ids: list[str]
mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str]
access_skill_names: tuple[str, ...] = ()
assigned_team_ids: list[str]
assigned_key_ids: list[str]
access_mcp_servers: tuple[AccessGroupResource, ...]

1
litellm/types/guardrails.py
View file

@ -156,6 +156,7 @@ class SupportedGuardrailIntegrations(Enum):
AGENT_365 = "agent_365"
LLM_SHIELD_PROXY = "llm_shield_proxy"
CONDUCT = "conduct"
AIRIA = "airia"
WINGBACK = "wingback"
DECISION_MODEL = "decision_model"

1
litellm/types/llms/anthropic.py
View file

@ -430,6 +430,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False):
cache_control: dict[str, Any] | None # Automatic prompt caching
reasoning_effort: str | None
safeguards: ReadOnly[list[dict[str, object]] | None]
service_tier: ReadOnly[str | None]
class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False):

6
litellm/types/llms/openai.py
View file

@ -947,6 +947,7 @@ class ChatCompletionDeveloperMessage(OpenAIChatCompletionDeveloperMessage, total
class GenericChatCompletionMessage(TypedDict, total=False):
role: Required[str]
content: Required[str | list]
cache_control: ReadOnly[ChatCompletionCachedContent]
ValidUserMessageContentTypes = [
@ -1931,7 +1932,10 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject):
class ErrorEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.ERROR]
sequence_number: int
error: ErrorEventError
error: ErrorEventError | None = None
code: str | int | None = None
message: str | None = None
param: str | Mapping[str, object] | None = None
class GenericEvent(BaseLiteLLMOpenAIResponseObject):

1
litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py
View file

@ -7,6 +7,7 @@ from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
class UiDiscoveryEndpoints(LiteLLMBaseModel):
server_root_path: str
proxy_base_url: str | None
gateway_base_url: str | None = None
auto_redirect_to_sso: bool
admin_ui_disabled: bool
sso_configured: bool

24
litellm/types/proxy/guardrails/guardrail_hooks/airia.py Normal file
View file

@ -0,0 +1,24 @@
from pydantic import Field
from .base import GuardrailConfigModel
class AiriaGuardrailConfigModel(GuardrailConfigModel): # pyright: ignore[reportMissingTypeArgument] # parameterizing adds an optional_params field to the UI form
api_base: str | None = Field(
default=None,
description="Base URL of the Airia AI Gateway, e.g. https://gateway.airia.ai. If not "
"provided, the `AIRIA_GATEWAY_URL` environment variable is checked.",
)
api_key: str | None = Field(
default=None,
description="An Airia API key. If not provided, the `AIRIA_API_KEY` environment variable is checked.",
)
timeout: float | None = Field(
default=None,
description="Request timeout in seconds. If not provided, the `AIRIA_TIMEOUT` environment "
"variable is checked, defaulting to 10.",
)
@staticmethod
def ui_friendly_name() -> str:
return "Airia Guardrail"

1
litellm/types/proxy/management_endpoints/team_endpoints.py
View file

@ -92,6 +92,7 @@ class TeamListItem(LiteLLM_TeamTable):
access_group_models: list[str] | None = None
access_group_mcp_server_ids: list[str] | None = None
access_group_agent_ids: list[str] | None = None
access_group_skill_names: tuple[str, ...] | None = None
class TeamListResponse(LiteLLMBaseModel):

12
litellm/types/router.py
View file

@ -56,6 +56,17 @@ from .utils import (
AllowedModelRegion: TypeAlias = Literal["eu", "us"]
class FallbackSkipReason(str, enum.Enum):
NOT_AUTHORIZED = "not_authorized"
OVER_BUDGET = "over_budget"
@dataclass(frozen=True, slots=True)
class SkippedFallbackTarget:
target: str
reason: FallbackSkipReason
class ConfigurableClientsideParamsCustomAuth(TypedDict):
api_base: str
@ -188,6 +199,7 @@ class UpdateRouterConfig(LiteLLMBaseModel):
weights: RouterWeights | None = None
tag_routing_prefix: str | None = None
optional_pre_call_checks: OptionalPreCallChecks | None = None
fallback_on_unsupported_params: bool | None = None
model_config = ConfigDict(protected_namespaces=())

6
model_prices_and_context_window.json
View file

@ -72571,10 +72571,12 @@
"gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -72585,6 +72587,7 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
@ -72617,10 +72620,12 @@
"vertex_ai/gemini-nano-banana-2.1": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_batches": 7.5e-08,
"cache_read_input_token_cost_flex": 7.5e-08,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_image": 0.00168,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"input_cost_per_token_priority": 2.7e-06,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 131072,
@ -72631,6 +72636,7 @@
"output_cost_per_image_token": 3e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
"search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014},
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",

1
ruff-strict.toml
View file

@ -30,6 +30,7 @@ external = [
# grows over time; typing it concretely (`object`) broke that forwarding call outright —
# basedpyright turned every named param into a reportArgumentType error. Any is correct here.
"litellm/proxy/guardrails/guardrail_hooks/alice/alice.py" = ["ANN401"]
"litellm/proxy/guardrails/guardrail_hooks/airia/airia.py" = ["ANN401"]
# Same reason: `**kwargs` forwards verbatim to CustomGuardrail.__init__, and the lifecycle
# hook signatures inherit `Any` for `response` from CustomLogger, so narrowing them here
# would break the override rather than describe it.

2
schema.prisma
View file

@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
mcp_tool_permissions Json?
access_agent_ids String[] @default([])
access_skill_names String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])

45
scripts/benchmark_ocr_callbacks.py
View file

@ -23,7 +23,7 @@ from collections.abc import Sequence
from dataclasses import asdict, dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final, cast
from typing import Final
SIZES: Final = (
1024,
@ -50,15 +50,12 @@ class Server(ThreadingHTTPServer):
def __init__(self) -> None:
super().__init__(("127.0.0.1", 0), Handler)
self.user_agents: set[str] = set()
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
server: Final = cast(Server, self.server)
server.user_agents.add(self.headers.get("User-Agent", ""))
length: Final = int(self.headers["Content-Length"])
body: Final = self.rfile.read(length)
request: Final = json.loads(body)
@ -105,10 +102,16 @@ def percentile(values: Sequence[float], quantile: float) -> float:
return ordered[index]
def verify(response: object) -> None:
def verify(response: object, expected_transport: str) -> None:
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
pages: Final = getattr(response, "pages", ())
if len(pages) != 1 or getattr(pages[0], "markdown", None) != EXPECTED_MARKDOWN:
raise RuntimeError(f"unexpected OCR response: {response!r}")
additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
rust_served: Final = isinstance(additional_headers, dict) and additional_headers.get("x-litellm-rust") == "true"
if rust_served != (expected_transport == "rust"):
raise RuntimeError(f"unexpected transport for {expected_transport}: rust_marker={rust_served}")
def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Result:
@ -125,7 +128,9 @@ def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Res
)
def sync_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]:
def sync_samples(
litellm: object, url: str, request_document: dict[str, str], count: int, expected_transport: str
) -> tuple[float, ...]:
samples: list[float] = []
for _ in range(count):
started: Final = time.perf_counter()
@ -133,11 +138,13 @@ def sync_samples(litellm: object, url: str, request_document: dict[str, str], co
model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30
)
samples.append(time.perf_counter() - started)
verify(response)
verify(response, expected_transport)
return tuple(samples)
async def async_samples(litellm: object, url: str, request_document: dict[str, str], count: int) -> tuple[float, ...]:
async def async_samples(
litellm: object, url: str, request_document: dict[str, str], count: int, expected_transport: str
) -> tuple[float, ...]:
samples: list[float] = []
for _ in range(count):
started: Final = time.perf_counter()
@ -145,7 +152,7 @@ async def async_samples(litellm: object, url: str, request_document: dict[str, s
model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30
)
samples.append(time.perf_counter() - started)
verify(response)
verify(response, expected_transport)
return tuple(samples)
@ -218,17 +225,20 @@ async def main() -> int:
try:
for size in args.sizes:
request_document: Final = document(size)
sync_samples(litellm, url, request_document, args.warmup)
sync_samples(litellm, url, request_document, args.warmup, args.expected_transport)
sync_result: Final = summarize(
args.label, "sync", size, sync_samples(litellm, url, request_document, args.iterations)
args.label,
"sync",
size,
sync_samples(litellm, url, request_document, args.iterations, args.expected_transport),
)
results.append(sync_result)
await async_samples(litellm, url, request_document, args.warmup)
await async_samples(litellm, url, request_document, args.warmup, args.expected_transport)
async_result: Final = summarize(
args.label,
"async",
size,
await async_samples(litellm, url, request_document, args.iterations),
await async_samples(litellm, url, request_document, args.iterations, args.expected_transport),
)
results.append(async_result)
sys.stdout.write(json.dumps(asdict(sync_result)) + "\n")
@ -256,12 +266,6 @@ async def main() -> int:
raise RuntimeError("callback registrations appeared in the no-callback case")
await GLOBAL_LOGGING_WORKER.stop()
user_agents: Final = tuple(sorted(server.user_agents))
python_transport: Final = any(
value.startswith("python-httpx") or value.startswith("litellm/") for value in user_agents
)
if (args.expected_transport == "python") != python_transport:
raise RuntimeError(f"unexpected transport for {args.label}: user_agents={user_agents}")
artifact: Final = {
"label": args.label,
"callbacks": args.callbacks,
@ -275,11 +279,10 @@ async def main() -> int:
"package": str(package),
"native": native_path,
"native_sha256": native_sha256,
"user_agents": user_agents,
"results": tuple(asdict(result) for result in results),
}
args.output.write_text(json.dumps(artifact, indent=2) + "\n")
sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package", "user_agents")}) + "\n")
sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package")}) + "\n")
sys.stdout.write(f"results={args.output}\n")
sys.stdout.flush()
return 0

1
terraform/provider/CHANGELOG.md
View file

@ -16,6 +16,7 @@ longer signal it.
### Added
- **team**: Optional `team_member_model_max_budget` argument on `litellm_team`, a JSON string of per-model budgets applied to every team member and sent as `team_member_model_max_budget` on `/team/new` and `/team/update`; read back from `team_member_budget_table.model_max_budget` on `/team/info` so out-of-band changes show as drift, with `budget_limit`/`time_period` treated as equal to the `max_budget`/`budget_duration` spelling the proxy stores; removing it sends `{}`, which clears every per-model member budget; `tpm_limit`/`rpm_limit` are rejected at plan time because the proxy rejects them on team budgets. Needs a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887): v1.105.0 and v1.106.0-rc.1 accept the field but silently drop it, so every plan shows a diff
- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are persisted through `/model/{id}/update` since `/model/update` ignores `model_info`; also exported by the `litellm_model` and `litellm_models` data sources
- **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones
- **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them

20
terraform/provider/docs/resources/team.md
View file

@ -98,6 +98,24 @@ resource "litellm_team" "model_dependent_team" {
}
```
### Team with Per-Member Model Budgets
```hcl
resource "litellm_team" "engineering" {
team_alias = "engineering"
models = ["claude-opus", "claude-sonnet"]
# Each member may spend $500 a day in total
team_member_budget = 500
team_member_budget_duration = "1d"
# ...of which at most $400 on Opus. Past that, Sonnet keeps working.
team_member_model_max_budget = jsonencode({
"claude-opus" = { budget_limit = 400, time_period = "1d" }
})
}
```
## Argument Reference
The following arguments are supported:
@ -150,6 +168,8 @@ The following arguments are supported:
* `team_member_key_duration` - (Optional) Lifetime for keys created by team members (e.g. `1d`, `1w`).
* `team_member_model_max_budget` - (Optional) JSON string of per-model budgets applied to each individual team member, e.g. `jsonencode({"claude-opus" = {budget_limit = 400, time_period = "1d"}})`. Each model maps to an object with `budget_limit` (or `max_budget`) and `time_period` (or `budget_duration`). The proxy rejects `tpm_limit` and `rpm_limit` on team budgets, so the provider rejects them at plan time. A member who reaches a model's cap is blocked on that model only and can keep calling the team's other models until `team_member_budget` runs out. The key is matched against the model name on the request, so for a router alias the cap applies to the alias, not the model the router picks. Remove the attribute to clear every per-model member budget. Requires a LiteLLM Enterprise license on the proxy, and a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887). Older proxies, including v1.105.0 and v1.106.0-rc.1, accept the field but silently drop it, so every plan shows a diff.
* `model_rpm_limit` - (Optional) Map of model name to requests per minute limit for that model.
* `model_tpm_limit` - (Optional) Map of model name to tokens per minute limit for that model.

90
terraform/provider/litellm/resource_team.go
View file

@ -6,6 +6,7 @@ import (
"io"
"log"
"net/http"
"reflect"
"github.com/google/uuid"
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@ -135,6 +136,13 @@ func ResourceLiteLLMTeam() *schema.Resource {
Type: schema.TypeString,
Optional: true,
},
"team_member_model_max_budget": {
Type: schema.TypeString,
Optional: true,
ValidateFunc: validateTeamMemberModelMaxBudget,
DiffSuppressFunc: teamMemberModelMaxBudgetEquivalent,
Description: "JSON string of per-model budgets applied to every team member (e.g. '{\"claude-opus\": {\"budget_limit\": 400, \"time_period\": \"1d\"}}'). Each entry takes budget_limit (or max_budget) and time_period (or budget_duration). Caps match the model name on the request, so behind a router alias the cap must be keyed on the alias, not the model the router picks. Requires a LiteLLM Enterprise license and a proxy that includes https://github.com/BerriAI/litellm/pull/43887; older proxies drop the value, so every plan shows a diff",
},
"model_rpm_limit": {
Type: schema.TypeMap,
Optional: true,
@ -280,6 +288,11 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error {
if teamResp.AllowedPassthroughRoutes != nil {
d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes)
}
var memberModelMaxBudget map[string]interface{}
if teamResp.TeamMemberBudgetTable != nil {
memberModelMaxBudget = teamResp.TeamMemberBudgetTable.ModelMaxBudget
}
d.Set("team_member_model_max_budget", keyModelMaxBudgetJSON(memberModelMaxBudget))
// Explicitly fetch the current permissions from the API
permResp, err := getTeamPermissions(client, d.Id())
@ -362,6 +375,75 @@ func resourceLiteLLMTeamDelete(d *schema.ResourceData, m interface{}) error {
return nil
}
// The proxy rejects tpm_limit/rpm_limit on team budgets, so only spend fields are allowed here.
var teamMemberModelBudgetFields = map[string]bool{
"budget_limit": true,
"max_budget": true,
"time_period": true,
"budget_duration": true,
}
func validateTeamMemberModelMaxBudget(v interface{}, k string) ([]string, []error) {
var parsed map[string]json.RawMessage
if err := json.Unmarshal([]byte(v.(string)), &parsed); err != nil || parsed == nil {
return nil, []error{fmt.Errorf("%q must be a JSON object keyed by model name, got %s", k, v)}
}
for model, cfg := range parsed {
var budget map[string]json.RawMessage
if err := json.Unmarshal(cfg, &budget); err != nil || len(budget) == 0 {
return nil, []error{fmt.Errorf("%q[%q] must be a budget object such as {\"budget_limit\": 400, \"time_period\": \"1d\"}, got %s", k, model, cfg)}
}
for field := range budget {
if !teamMemberModelBudgetFields[field] {
return nil, []error{fmt.Errorf("%q[%q] has unsupported field %q; team member budgets accept budget_limit, max_budget, time_period and budget_duration only", k, model, field)}
}
}
}
return nil, nil
}
// The proxy stores budget_limit/time_period as max_budget/budget_duration, so both spellings must compare equal.
var modelBudgetFieldAliases = map[string]string{
"budget_limit": "max_budget",
"time_period": "budget_duration",
}
func normalizeModelMaxBudget(raw string) (map[string]interface{}, bool) {
if raw == "" {
return map[string]interface{}{}, true
}
var parsed map[string]interface{}
if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
return nil, false
}
normalized := make(map[string]interface{}, len(parsed))
for model, cfg := range parsed {
budget, ok := cfg.(map[string]interface{})
if !ok {
normalized[model] = cfg
continue
}
entry := make(map[string]interface{}, len(budget))
for field, value := range budget {
if value == nil {
continue
}
if canonical, isAlias := modelBudgetFieldAliases[field]; isAlias {
field = canonical
}
entry[field] = value
}
normalized[model] = entry
}
return normalized, true
}
func teamMemberModelMaxBudgetEquivalent(k, oldValue, newValue string, d *schema.ResourceData) bool {
oldParsed, okOld := normalizeModelMaxBudget(oldValue)
newParsed, okNew := normalizeModelMaxBudget(newValue)
return okOld && okNew && reflect.DeepEqual(oldParsed, newParsed)
}
func resolveTeamID(d *schema.ResourceData) string {
if v, ok := d.GetOk("team_id"); ok {
return v.(string)
@ -402,6 +484,14 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{}
teamData["tags"] = v
}
// The proxy treats a missing field as "leave unchanged", so an empty map is sent only to clear a value that was set.
memberModelMaxBudget := parseKeyModelMaxBudget(d.Get("team_member_model_max_budget").(string))
if len(memberModelMaxBudget) > 0 {
teamData["team_member_model_max_budget"] = memberModelMaxBudget
} else if old, _ := d.GetChange("team_member_model_max_budget"); len(parseKeyModelMaxBudget(old.(string))) > 0 {
teamData["team_member_model_max_budget"] = memberModelMaxBudget
}
if metadata := buildTeamMetadata(d); metadata != nil {
teamData["metadata"] = metadata
}

235
terraform/provider/litellm/resource_team_test.go
View file

@ -472,3 +472,238 @@ func TestHandleResponseAcceptsFullSuccessRange(t *testing.T) {
})
}
}
func TestTeamCreateSendsTeamMemberModelMaxBudget(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
d := newTeamResourceData(t, map[string]interface{}{
"team_alias": "eng",
"team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`,
})
if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("create failed: %v", err)
}
want := map[string]interface{}{"claude-opus": map[string]interface{}{"budget_limit": 400.0, "time_period": "1d"}}
if got := captured["team_member_model_max_budget"]; !reflect.DeepEqual(got, want) {
t.Fatalf("payload team_member_model_max_budget = %v, want %v", got, want)
}
}
func TestTeamReadMapsTeamMemberModelMaxBudgetFromBudgetTable(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{
"team_id": "team-1",
"team_info": {
"team_id": "team-1",
"team_alias": "eng",
"team_member_budget_table": {
"max_budget": 500.0,
"model_max_budget": {"claude-opus": {"max_budget": 400.0, "budget_duration": "1d"}}
}
}
}`)
defer srv.Close()
d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"})
d.SetId("team-1")
if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("read returned error: %v", err)
}
var got map[string]interface{}
if err := json.Unmarshal([]byte(d.Get("team_member_model_max_budget").(string)), &got); err != nil {
t.Fatalf("team_member_model_max_budget in state is not JSON: %v", err)
}
want := map[string]interface{}{"claude-opus": map[string]interface{}{"max_budget": 400.0, "budget_duration": "1d"}}
if !reflect.DeepEqual(got, want) {
t.Errorf("team_member_model_max_budget = %v, want %v", got, want)
}
}
func TestTeamReadClearsTeamMemberModelMaxBudgetWhenAbsent(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng","team_member_budget_table":{"model_max_budget":null}}}`)
defer srv.Close()
d := newTeamResourceData(t, map[string]interface{}{
"team_alias": "eng",
"team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`,
})
d.SetId("team-1")
if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("read returned error: %v", err)
}
if got := d.Get("team_member_model_max_budget").(string); got != "" {
t.Errorf("team_member_model_max_budget = %q, want empty so drift shows up", got)
}
}
func TestTeamUpdateClearsRemovedTeamMemberModelMaxBudget(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
res := ResourceLiteLLMTeam()
priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{
"team_alias": "eng",
"team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`,
})
priorData.SetId("team-1")
prior := priorData.State()
config := terraform.NewResourceConfigRaw(map[string]interface{}{"team_alias": "eng"})
diff, err := res.Diff(context.Background(), prior, config, nil)
if err != nil {
t.Fatalf("diff failed: %v", err)
}
d, err := schema.InternalMap(res.Schema).Data(prior, diff)
if err != nil {
t.Fatalf("data failed: %v", err)
}
if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("update failed: %v", err)
}
if got, ok := captured["team_member_model_max_budget"]; !ok || !reflect.DeepEqual(got, map[string]interface{}{}) {
t.Errorf("payload team_member_model_max_budget = %v (present=%v), want explicit empty object", got, ok)
}
}
func TestTeamMemberModelMaxBudgetAliasesAreEquivalent(t *testing.T) {
configured := `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`
stored := `{"claude-opus":{"max_budget":400,"budget_duration":"1d"}}`
if !teamMemberModelMaxBudgetEquivalent("", stored, configured, nil) {
t.Fatalf("budget_limit/time_period should match the stored max_budget/budget_duration spelling")
}
if teamMemberModelMaxBudgetEquivalent("", stored, `{"claude-opus": {"budget_limit": 300, "time_period": "1d"}}`, nil) {
t.Fatalf("a changed limit must not be suppressed")
}
}
func TestTeamMemberModelMaxBudgetRejectsUnknownFields(t *testing.T) {
schema := ResourceLiteLLMTeam().Schema["team_member_model_max_budget"]
for _, bad := range []string{
`{"claude-opus": {"limit": 400}}`,
`{"claude-opus": {"budget_limit": 400, "time_period": "1d", "tpm_limit": 1000}}`,
`{"claude-opus": {"budget_limit": 400, "time_period": "1d", "rpm_limit": 10}}`,
} {
if _, errs := schema.ValidateFunc(bad, "team_member_model_max_budget"); len(errs) == 0 {
t.Errorf("expected %s to be rejected", bad)
}
}
if _, errs := schema.ValidateFunc(`{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`, "team_member_model_max_budget"); len(errs) != 0 {
t.Fatalf("valid config rejected: %v", errs)
}
}
func TestTeamCreateOmitsTeamMemberModelMaxBudgetWhenUnset(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"})
if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("create failed: %v", err)
}
if got, ok := captured["team_member_model_max_budget"]; ok {
t.Fatalf("payload team_member_model_max_budget = %v, want key absent", got)
}
}
func TestTeamCreateOmitsEmptyTeamMemberModelMaxBudget(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
d := newTeamResourceData(t, map[string]interface{}{
"team_alias": "eng",
"team_member_model_max_budget": `{}`,
})
if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("create failed: %v", err)
}
if got, ok := captured["team_member_model_max_budget"]; ok {
t.Fatalf("payload team_member_model_max_budget = %v, want key absent so no empty member budget row is created", got)
}
}
func teamUpdateData(t *testing.T, prior, config map[string]interface{}) *schema.ResourceData {
t.Helper()
res := ResourceLiteLLMTeam()
priorData := schema.TestResourceDataRaw(t, res.Schema, prior)
priorData.SetId("team-1")
state := priorData.State()
diff, err := res.Diff(context.Background(), state, terraform.NewResourceConfigRaw(config), nil)
if err != nil {
t.Fatalf("diff failed: %v", err)
}
d, err := schema.InternalMap(res.Schema).Data(state, diff)
if err != nil {
t.Fatalf("data failed: %v", err)
}
return d
}
func TestTeamUpdateSendsChangedTeamMemberModelMaxBudget(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
d := teamUpdateData(t,
map[string]interface{}{"team_alias": "eng", "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`},
map[string]interface{}{"team_alias": "eng", "team_member_model_max_budget": `{"claude-opus": {"budget_limit": 300, "time_period": "1d"}}`},
)
if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("update failed: %v", err)
}
want := map[string]interface{}{"claude-opus": map[string]interface{}{"budget_limit": 300.0, "time_period": "1d"}}
if got := captured["team_member_model_max_budget"]; !reflect.DeepEqual(got, want) {
t.Fatalf("payload team_member_model_max_budget = %v, want %v", got, want)
}
}
func TestTeamUpdateSendsMemberBudgetWithPerModelBudget(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
d := teamUpdateData(t,
map[string]interface{}{"team_alias": "eng"},
map[string]interface{}{
"team_alias": "eng",
"team_member_budget": 500.0,
"team_member_budget_duration": "1d",
"team_member_model_max_budget": `{"claude-opus": {"budget_limit": 400, "time_period": "1d"}}`,
},
)
if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("update failed: %v", err)
}
if got := captured["team_member_budget"]; got != 500.0 {
t.Errorf("payload team_member_budget = %v, want 500", got)
}
if got := captured["team_member_budget_duration"]; got != "1d" {
t.Errorf("payload team_member_budget_duration = %v, want 1d", got)
}
want := map[string]interface{}{"claude-opus": map[string]interface{}{"budget_limit": 400.0, "time_period": "1d"}}
if got := captured["team_member_model_max_budget"]; !reflect.DeepEqual(got, want) {
t.Errorf("payload team_member_model_max_budget = %v, want %v", got, want)
}
}
func TestTeamReadWithoutMemberBudgetTableLeavesFieldEmpty(t *testing.T) {
var captured map[string]interface{}
srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"eng"}}`)
defer srv.Close()
d := newTeamResourceData(t, map[string]interface{}{"team_alias": "eng"})
d.SetId("team-1")
if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil {
t.Fatalf("read returned error: %v", err)
}
if got := d.Get("team_member_model_max_budget").(string); got != "" {
t.Errorf("team_member_model_max_budget = %q, want empty", got)
}
}

6
terraform/provider/litellm/types.go
View file

@ -73,6 +73,12 @@ type TeamResponse struct {
ModelRPMLimit map[string]interface{} `json:"model_rpm_limit,omitempty"`
ModelTPMLimit map[string]interface{} `json:"model_tpm_limit,omitempty"`
AllowedPassthroughRoutes []string `json:"allowed_passthrough_routes,omitempty"`
TeamMemberBudgetTable *TeamMemberBudgetTable `json:"team_member_budget_table,omitempty"`
}
// TeamMemberBudgetTable is the team's default member budget row, the only place /team/info returns team_member_model_max_budget.
type TeamMemberBudgetTable struct {
ModelMaxBudget map[string]interface{} `json:"model_max_budget,omitempty"`
}
// OrganizationResponse represents a response from the API containing organization information.

8
tests/e2e/coverage_registry/mcp.yaml
View file

@ -39,6 +39,14 @@
assertions: [denied_without_permission]
source: "rest_endpoints.py:305-386"
rationale: Tool-level permission guard; multi-tenant safety
- id: mcp.tool_search.api_key.oversized_description_omitted
module: mcp
tier: P0
operation: tool_search
auth_family: api_key
assertions: [oversized_description_omitted]
source: "test_mcp_tool_search_e2e.py"
rationale: Native semantic search omits tools whose descriptions exceed the embedding input limit
- id: mcp.list_tools.bearer.succeeds
module: mcp
tier: P1

4
tests/e2e/mcp/datadog_mcp.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
import os
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Protocol
from e2e_config import datadog_mcp_url, unique_marker
@ -44,6 +44,7 @@ def register_datadog_mcp(
*,
mcp_access_groups: list[str] | None = None,
allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,),
tool_name_to_description: Mapping[str, str] | None = None,
) -> str:
"""Register the core Datadog toolset with its credentials from the env. By default
the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose
@ -61,6 +62,7 @@ def register_datadog_mcp(
},
allowed_tools=None if allowed_tools is None else list(allowed_tools),
mcp_access_groups=mcp_access_groups,
tool_name_to_description=tool_name_to_description,
)
resources.defer(lambda: client.delete_server(server_id))
return server_id

84
tests/e2e/mcp/mcp_client.py
View file

@ -42,6 +42,7 @@ class McpServerNewBody(BaseModel):
static_headers: dict[str, str] | None = None
allowed_tools: list[str] | None = None
mcp_access_groups: list[str] | None = None
tool_name_to_description: dict[str, str] | None = None
class McpServerNewResponse(BaseModel):
@ -131,6 +132,34 @@ class McpCallToolBody(BaseModel):
server_id: str
class McpVirtualToolCallBody(BaseModel):
name: str
arguments: dict[str, McpToolArg]
class McpToolsListParams(BaseModel):
server_id: str | None = None
class McpToolSearchSettings(BaseModel):
model_config = ConfigDict(frozen=True)
embedding_model: str | None = None
top_k: int = 5
similarity_threshold: float = 0.0
core_tools: tuple[str, ...] = ()
class McpToolSearchSettingsResponse(BaseModel):
values: McpToolSearchSettings
class McpToolSearchSettingsUpdateResponse(BaseModel):
message: str
status: str
settings: McpToolSearchSettings
class McpCallContent(BaseModel):
type: str | None = None
text: str | None = None
@ -166,6 +195,7 @@ class McpClient:
static_headers: dict[str, str] | None = None,
allowed_tools: list[str] | None = None,
mcp_access_groups: list[str] | None = None,
tool_name_to_description: Mapping[str, str] | None = None,
) -> str:
return unwrap(
self.proxy.transport.post(
@ -180,6 +210,9 @@ class McpClient:
static_headers=static_headers,
allowed_tools=allowed_tools,
mcp_access_groups=mcp_access_groups,
tool_name_to_description=(
dict(tool_name_to_description) if tool_name_to_description is not None else None
),
),
response_type=McpServerNewResponse,
)
@ -244,14 +277,21 @@ class McpClient:
mcp_access_groups: list[str] | None = None,
mcp_toolsets: list[str] | None = None,
models: list[str] | None = None,
mcp_tool_search_enabled: bool | None = None,
) -> str:
object_permission = (
ObjectPermission(
mcp_servers=mcp_servers,
mcp_access_groups=mcp_access_groups,
mcp_toolsets=mcp_toolsets,
mcp_tool_search_enabled=mcp_tool_search_enabled,
)
if (
mcp_servers is not None
or mcp_access_groups is not None
or mcp_toolsets is not None
or mcp_tool_search_enabled is not None
)
if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None
else None
)
return self.proxy.generate_key(
@ -263,14 +303,37 @@ class McpClient:
)
@step("List the MCP tools the key can see from /mcp-rest/tools/list")
def list_tools(self, key: str) -> Result[McpToolsListResponse]:
def list_tools(self, key: str, server_id: str | None = None) -> Result[McpToolsListResponse]:
return self.proxy.transport.get(
"/mcp-rest/tools/list",
headers=ApiKeyHeaders(x_litellm_api_key=key),
params=NoBody(),
params=McpToolsListParams(server_id=server_id),
response_type=McpToolsListResponse,
)
@step("Read the MCP tool search settings from /get/mcp_tool_search_settings")
def get_mcp_tool_search_settings(self) -> McpToolSearchSettings:
return unwrap(
self.proxy.transport.get(
"/get/mcp_tool_search_settings",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpToolSearchSettingsResponse,
)
).values
@step("Update the MCP tool search settings with /update/mcp_tool_search_settings")
def update_mcp_tool_search_settings(self, settings: McpToolSearchSettings) -> None:
_ = unwrap(
self.proxy.transport.patch(
"/update/mcp_tool_search_settings",
headers=self.proxy.transport.master,
json=settings,
response_type=McpToolSearchSettingsUpdateResponse,
)
)
settle_propagation(time.monotonic())
@step('Wait for /mcp-rest/tools/list to show the MCP server\'s tool matching "{needle}"')
def await_tool(self, key: str, server_id: str, needle: str) -> str:
"""Poll tools/list until `server_id` serves a tool matching `needle`, and
@ -419,6 +482,21 @@ class McpClient:
response_type=McpCallToolResponse,
)
@step("Call the MCP virtual tool {name} through /mcp-rest/tools/call")
def call_virtual_tool(
self,
key: str,
*,
name: str,
arguments: McpToolArguments,
) -> Result[McpCallToolResponse]:
return self.proxy.transport.post(
"/mcp-rest/tools/call",
headers=ApiKeyHeaders(x_litellm_api_key=key),
json=McpVirtualToolCallBody(name=name, arguments=dict(arguments)),
response_type=McpCallToolResponse,
)
def _is_mcp_not_synced(
result: Result[McpCallToolResponse],

86
tests/e2e/mcp/test_mcp_tool_search_e2e.py Normal file
View file

@ -0,0 +1,86 @@
from __future__ import annotations
from typing import Final
import pytest
from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp
from e2e_config import provider_edge_base, unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from mcp_client import McpClient, McpToolSearchSettings
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
pytestmark = pytest.mark.e2e
MCP_TOOL_SEARCH_TOOL_NAME: Final = "mcp_tool_search"
OVERSIZED_DESCRIPTION: Final = f"OVERSIZED {'weather ' * 12_000}"
class TestMcpToolSearch:
@pytest.mark.covers("mcp.tool_search.api_key.oversized_description_omitted")
def test_native_tool_search_omits_oversized_description(
self, proxy: ProxyClient, client: McpClient, resources: ResourceManager
) -> None:
"""Exercise shared ranking through the native virtual tool because proxy mode does not apply description overrides."""
model_name: Final = f"e2e-mcp-tool-search-{unique_marker()}"
embedding_api_base: Final = provider_edge_base("openai")
model_id: Final = proxy.create_model(
model_name,
LiteLLMParamsBody(
model="openai/text-embedding-3-large",
api_key="os.environ/OPENAI_API_KEY",
api_base=None if embedding_api_base is None else f"{embedding_api_base}/v1",
),
)
resources.defer(lambda: proxy.delete_model(model_id))
previous_settings: Final = client.get_mcp_tool_search_settings()
resources.defer(lambda: client.update_mcp_tool_search_settings(previous_settings))
client.update_mcp_tool_search_settings(
McpToolSearchSettings(
embedding_model=model_name,
top_k=100,
similarity_threshold=0.0,
core_tools=previous_settings.core_tools,
)
)
server_id: Final = register_datadog_mcp(
client,
resources,
allowed_tools=None,
tool_name_to_description={SEARCH_LOGS_TOOL: OVERSIZED_DESCRIPTION},
)
client.await_registered(server_id)
key: Final = client.generate_key(
user_id=f"e2e-mcp-tool-search-{unique_marker()}",
mcp_servers=[server_id],
models=[model_name],
mcp_tool_search_enabled=True,
)
resources.defer(lambda: proxy.delete_key(key))
tools: Final = unwrap(client.list_tools(key, server_id=server_id))
assert len(tools.tool_names_for_server(server_id)) >= 2
oversized_tool_name: Final = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
assert oversized_tool_name is not None
oversized_tool: Final = next((tool for tool in tools.tools if tool.name == oversized_tool_name), None)
assert oversized_tool is not None
assert oversized_tool.description == OVERSIZED_DESCRIPTION
normal_tool_name: Final = next(
(name for name in tools.tool_names_for_server(server_id) if name != oversized_tool_name), None
)
assert normal_tool_name is not None
result: Final = unwrap(
client.call_virtual_tool(
key,
name=MCP_TOOL_SEARCH_TOOL_NAME,
arguments={"query": normal_tool_name, "top_k": 100},
)
)
assert result.is_error is not True, result.all_text
assert normal_tool_name in result.all_text
assert oversized_tool_name not in result.all_text

1
tests/e2e/models.py
View file

@ -67,6 +67,7 @@ class ObjectPermission(BaseModel):
mcp_servers: list[str] | None = None
mcp_access_groups: list[str] | None = None
mcp_toolsets: list[str] | None = None
mcp_tool_search_enabled: bool | None = None
class KeyGenerateBody(BaseModel):

533
tests/integration/_support/anthropic_refusal.py Normal file
View file

@ -0,0 +1,533 @@
import hashlib
import json
import os
import re
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from types import MappingProxyType
from typing import Final, TypeAlias
import yaml
from integration._support.anthropic_thinking import (
JSON_LIST,
JSON_OBJECT,
Event,
marker_of,
stream_reply,
streams,
)
from integration._support.client import eventually
from integration._support.wire import Reply, Request, Wire
from pydantic import JsonValue, TypeAdapter
from redis import Redis
DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response"
DOCS_STREAMING_REFUSAL_URL: Final = (
"https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals"
)
# Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12.
DOCS_REFUSAL: Final[Event] = {
"id": "msg_01XFUDYJgAACzvnptvVoYEL",
"type": "message",
"role": "assistant",
"model": "claude-fable-5",
"content": [],
"stop_reason": "refusal",
"stop_details": {
"type": "refusal",
"category": "cyber",
"explanation": "This request was declined because it could enable cyber harm.",
},
"usage": {"input_tokens": 412, "output_tokens": 0},
}
DOCS_PARTIAL_REFUSAL: Final[Event] = {
"role": "assistant",
"content": [{"type": "text", "text": "Hello.."}],
"stop_reason": "refusal",
"stop_details": {
"type": "refusal",
"category": "cyber",
"explanation": "This request was declined because it could enable cyber harm.",
},
}
# Captured from api.anthropic.com (claude-sonnet-5-5, prompt "Reply with the single word: ready") on 2026-10-13.
CAPTURED_TURN: Final[Event] = {
"model": "claude-sonnet-5-5",
"id": "msg_011CfyXFGze1cDkk8o7gDRRS",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "ready"}],
"container": None,
"stop_reason": "end_turn",
"stop_sequence": None,
"stop_details": None,
"usage": {
"input_tokens": 18,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
"output_tokens": 4,
"output_tokens_details": {"thinking_tokens": 0},
"service_tier": "standard",
"inference_geo": "global",
},
"diagnostics": None,
}
CAPTURED_EVENTS: Final[tuple[Event, ...]] = (
{
"type": "message_start",
"message": {
"model": "claude-sonnet-5-5",
"id": "msg_011CfyXFLwXUxsp7nNAm3gcu",
"type": "message",
"role": "assistant",
"content": [],
"container": None,
"stop_reason": None,
"stop_sequence": None,
"stop_details": None,
"usage": {
"input_tokens": 18,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
"output_tokens": 4,
"service_tier": "standard",
"inference_geo": "global",
},
"diagnostics": None,
},
},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "ping"},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ready"}},
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None},
"usage": {
"input_tokens": 18,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"output_tokens": 4,
"output_tokens_details": {"thinking_tokens": 0},
},
},
{"type": "message_stop"},
)
DOCS_STOP_DETAILS: Final = DOCS_REFUSAL["stop_details"]
REFUSAL_FIELDS: Final[Event] = {"stop_reason": "refusal", "stop_details": DOCS_STOP_DETAILS}
STOP_FIELDS: Final = ("stop_reason", "stop_details")
LONG_STOP_DETAILS: Final = "d" * 5120
MALFORMED_STOP_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"int": 7, "list": ["refusal"], "empty": "", "long": LONG_STOP_DETAILS}
)
_FIXTURE: Final = re.compile(r"refusal-fixture-([a-z]+)")
Location: TypeAlias = tuple[str | int, ...]
@dataclass(frozen=True, slots=True)
class Fixture:
status: int
body: bytes
events: tuple[Event, ...] = ()
def _text_block(index: int, block: JsonValue) -> tuple[Event, ...]:
text: Final = JSON_OBJECT.validate_python(block)["text"]
return (
{"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}},
{"type": "content_block_stop", "index": index},
)
def refusal_events(message: Mapping[str, JsonValue]) -> tuple[Event, ...]:
usage: Final = JSON_OBJECT.validate_python(message["usage"])
start: Final[Event] = {
"type": "message_start",
"message": {
"id": message["id"],
"type": "message",
"role": "assistant",
"model": message["model"],
"content": [],
"stop_reason": None,
"stop_sequence": None,
"stop_details": None,
"usage": usage,
},
}
blocks: Final = chain.from_iterable(
_text_block(index, block) for index, block in enumerate(JSON_LIST.validate_python(message["content"]))
)
delta: Final[Event] = {
"type": "message_delta",
"delta": {
"stop_reason": message["stop_reason"],
"stop_sequence": None,
"stop_details": message["stop_details"],
},
"usage": {name: usage[name] for name in ("output_tokens", "output_tokens_details") if name in usage},
}
return (start, *blocks, delta, {"type": "message_stop"})
def _without_stop_details(value: Mapping[str, JsonValue]) -> Event:
return {name: item for name, item in value.items() if name != "stop_details"}
def _captured_event(event: Event, stop_details: JsonValue, *, keep: bool) -> Event:
if event["type"] == "message_start":
message: Final = JSON_OBJECT.validate_python(event["message"])
return {**event, "message": message if keep else _without_stop_details(message)}
if event["type"] != "message_delta":
return event
delta: Final = JSON_OBJECT.validate_python(event["delta"])
replaced: Final = {**delta, "stop_details": stop_details} if keep else _without_stop_details(delta)
return {**event, "delta": replaced}
def _captured(stop_details: JsonValue, *, keep: bool = True) -> Fixture:
body: Final = {**CAPTURED_TURN, "stop_details": stop_details} if keep else _without_stop_details(CAPTURED_TURN)
return Fixture(
200,
json.dumps(body).encode(),
tuple(_captured_event(event, stop_details, keep=keep) for event in CAPTURED_EVENTS),
)
def _duplicated_stop_details() -> bytes:
pair: Final = f'"stop_details": {json.dumps(DOCS_STOP_DETAILS)}'
return json.dumps(DOCS_REFUSAL).replace(pair, f"{pair}, {pair}").encode()
def _anthropic_error(kind: str, message: str) -> bytes:
return json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode()
_PARTIAL: Final[Event] = {**DOCS_REFUSAL, **DOCS_PARTIAL_REFUSAL}
_REASONED: Final[Event] = {
**_PARTIAL,
"usage": {"input_tokens": 412, "output_tokens": 10, "output_tokens_details": {"thinking_tokens": 4}},
}
JSON_TOOL_NAME: Final = "json_tool_call"
JSON_ANSWER_SCHEMA: Final[Event] = {
"type": "object",
"properties": {"answer": {"type": "string"}},
"required": ["answer"],
"additionalProperties": False,
}
JSON_ANSWER_FORMAT: Final[Event] = {
"type": "json_schema",
"json_schema": {"name": "answer", "schema": JSON_ANSWER_SCHEMA, "strict": True},
}
DOCS_PARTIAL_TEXT: Final = "Hello.."
_DECLARED_INPUT: Final[Event] = {"answer": DOCS_PARTIAL_TEXT}
_ENVELOPED_INPUT: Final[Event] = {"parameters": _DECLARED_INPUT}
DECLARED_PARTIAL_JSON: Final = '{"answer": "Hello..'
ENVELOPED_PARTIAL_JSON: Final = '{"parameters": {"answer": "Hello..'
_JSON_TOOL_USE_ID: Final = "toolu_01RefusalFixtureJsonTool"
def _json_tool_refusal(tool_input: Event, partial_json: str) -> Fixture:
message: Final[Event] = {
**DOCS_REFUSAL,
"content": [{"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": tool_input}],
}
start, delta, stop = refusal_events(DOCS_REFUSAL)
tool_block: Final[tuple[Event, ...]] = (
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": {}},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": partial_json},
},
{"type": "content_block_stop", "index": 0},
)
return Fixture(200, json.dumps(message).encode(), (start, *tool_block, delta, stop))
FIXTURES: Final[Mapping[str, Fixture]] = MappingProxyType(
{
"refusal": Fixture(200, json.dumps(DOCS_REFUSAL).encode(), refusal_events(DOCS_REFUSAL)),
"partial": Fixture(200, json.dumps(_PARTIAL).encode(), refusal_events(_PARTIAL)),
"reasoned": Fixture(200, json.dumps(_REASONED).encode(), refusal_events(_REASONED)),
"jsontool": _json_tool_refusal(_DECLARED_INPUT, DECLARED_PARTIAL_JSON),
"jsonheld": _json_tool_refusal(_ENVELOPED_INPUT, ENVELOPED_PARTIAL_JSON),
"duplicate": Fixture(200, _duplicated_stop_details()),
"null": _captured(None),
"missing": _captured(None, keep=False),
**{name: _captured(value) for name, value in MALFORMED_STOP_DETAILS.items()},
"nondict": Fixture(200, b"[]"),
"unauthorized": Fixture(401, _anthropic_error("authentication_error", "invalid x-api-key")),
"invalid": Fixture(400, _anthropic_error("invalid_request_error", "max_tokens: Field required")),
}
)
def refusal_prompt(fixture: str, marker: str) -> str:
return f"refusal-fixture-{fixture} marker-{marker}"
CITED_FIELDS: Final[Event] = {"citations": ["https://example.com/source"]}
_PLAIN_DELTA: Final[Event] = {"role": "assistant", "content": "ready"}
_CITED_DELTA: Final[Event] = {**_PLAIN_DELTA, "provider_specific_fields": CITED_FIELDS}
_OPENAI_USAGE: Final[Event] = {
"prompt_tokens": 5,
"completion_tokens": 7,
"total_tokens": 12,
"completion_tokens_details": {"reasoning_tokens": 4, "text_tokens": 3},
}
UNVALIDATED_USAGE_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"text": {"reasoning_tokens": "abc"},
"fraction": {"reasoning_tokens": 1.5},
"object": {"reasoning_tokens": {"a": 1}},
}
)
@dataclass(frozen=True, slots=True)
class OpenAIFixture:
first_delta: Event
usage: tuple[Event, ...] = ()
_OPENAI_FIXTURES: Final[Mapping[str, OpenAIFixture]] = MappingProxyType(
{
"plain": OpenAIFixture(_PLAIN_DELTA),
"cited": OpenAIFixture(_CITED_DELTA),
"plainusage": OpenAIFixture(_PLAIN_DELTA, (_OPENAI_USAGE,)),
"citedusage": OpenAIFixture(_CITED_DELTA, (_OPENAI_USAGE,)),
}
)
def openai_id(marker: str) -> str:
return f"chatcmpl-{marker}"
def _openai_frame(marker: str, body: Event) -> bytes:
chunk: Final[Event] = {
"id": openai_id(marker),
"object": "chat.completion.chunk",
"created": 1760000000,
"model": "gpt-4o-mini",
**body,
}
return b"data: " + json.dumps(chunk).encode() + b"\n\n"
def _openai_choice(delta: Event, finish_reason: str | None) -> Event:
return {"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]}
def _openai_stream(name: str, request: Request) -> Reply:
assert streams(request), request.body
marker: Final = marker_of(request)
fixture: Final = _OPENAI_FIXTURES[name]
return Reply(
content_type="text/event-stream",
chunks=(
_openai_frame(marker, _openai_choice(fixture.first_delta, None)),
_openai_frame(marker, _openai_choice({"content": " now"}, "stop")),
*(_openai_frame(marker, {"choices": [], "usage": usage}) for usage in fixture.usage),
b"data: [DONE]\n\n",
),
)
def respond(request: Request) -> Reply:
found: Final[Sequence[str]] = _FIXTURE.findall(request.body.decode())
assert found, request.body
if request.target.endswith("/chat/completions"):
return _openai_stream(found[-1], request)
fixture: Final = FIXTURES[found[-1]]
if fixture.status != 200 or not fixture.events or not streams(request):
return Reply(status=fixture.status, body=fixture.body)
return stream_reply(request, fixture.events)
def _stop_fields(fields: JsonValue) -> Event:
if not isinstance(fields, dict):
return {}
return {name: fields[name] for name in STOP_FIELDS if name in fields}
def relayed(value: JsonValue, location: Location = ()) -> tuple[tuple[Location, Event], ...]:
if isinstance(value, list):
return tuple(chain.from_iterable(relayed(item, (*location, index)) for index, item in enumerate(value)))
if not isinstance(value, dict):
return ()
found: Final = _stop_fields(value.get("provider_specific_fields"))
here: Final = (((*location, "provider_specific_fields"), found),) if found else ()
return (*here, *chain.from_iterable(relayed(item, (*location, name)) for name, item in value.items()))
_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}")
_STORED: Final = TypeAdapter(bytes | None)
STALE_ID_SUFFIX: Final = "-stale"
def _stored(store: Redis, key: str) -> bytes:
if not _CACHE_KEY.fullmatch(key) or store.type(key) != b"string":
return b""
return _STORED.validate_python(store.get(key)) or b""
def _entry_key(store: Redis, needle: str) -> str | None:
keys: Final = (raw.decode() for raw in store.scan_iter(count=1000))
return next((key for key in keys if needle.encode() in _stored(store, key)), None)
def _with_message_fields(response: Event, fields: JsonValue) -> Event:
choice: Final = JSON_OBJECT.validate_python(JSON_LIST.validate_python(response["choices"])[0])
message: Final = JSON_OBJECT.validate_python(choice["message"])
return {**response, "choices": [{**choice, "message": {**message, "provider_specific_fields": fields}}]}
def _with_usage_details(response: Event, details: JsonValue) -> Event:
usage: Final = JSON_OBJECT.validate_python(response["usage"])
return {**response, "usage": {**usage, "completion_tokens_details": details}}
def _rewrite_cached_response(needle: str, rewrite: Callable[[Event], Event]) -> str:
store: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
key: Final = eventually(lambda: _entry_key(store, needle), lambda found: found is not None)
assert key is not None, needle
entry: Final = JSON_OBJECT.validate_json(_stored(store, key))
stored: Final = entry["response"]
encoded: Final = isinstance(stored, str)
response: Final = (
JSON_OBJECT.validate_json(stored) if isinstance(stored, str) else JSON_OBJECT.validate_python(stored)
)
rewritten: Final = {**rewrite(response), "id": f"{response['id']}{STALE_ID_SUFFIX}"}
store.set(key, json.dumps({**entry, "response": json.dumps(rewritten) if encoded else rewritten}), ex=60)
return str(rewritten["id"])
def rewrite_cached_message_fields(needle: str, fields: JsonValue) -> str:
return _rewrite_cached_response(needle, lambda response: _with_message_fields(response, fields))
def rewrite_cached_usage_details(needle: str, details: JsonValue) -> str:
return _rewrite_cached_response(needle, lambda response: _with_usage_details(response, details))
def key_hash(key: str) -> str:
return hashlib.sha256(key.encode()).hexdigest()
def _key_hash_of(event: Mapping[str, JsonValue]) -> JsonValue:
return JSON_OBJECT.validate_python(event.get("metadata") or {}).get("user_api_key_hash")
class Ledger:
def __init__(self, wire: Wire) -> None:
self.wire: Final = wire
self.batches: tuple[Request, ...] = ()
def events(self) -> tuple[Event, ...]:
self.batches = (*self.batches, *self.wire.drain())
decoded: Final = chain.from_iterable(JSON_LIST.validate_json(batch.body) for batch in self.batches)
return tuple(JSON_OBJECT.validate_python(event) for event in decoded)
def for_key(self, key: str) -> tuple[Event, ...]:
hashed: Final = key_hash(key)
return tuple(event for event in self.events() if _key_hash_of(event) == hashed)
def await_key(self, key: str, count: int = 1, seconds: float = 30) -> tuple[Event, ...]:
return eventually(lambda: self.for_key(key), lambda found: len(found) >= count, seconds=seconds)
_HIDDEN_PARAMS_PROBE: Final = """
import httpx
from litellm.integrations.custom_logger import CustomLogger
class HiddenParamsProbe(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
payload = kwargs.get("standard_logging_object") or {{}}
hidden = getattr(response_obj, "_hidden_params", None) or {{}}
event = {{
"metadata": {{"user_api_key_hash": (payload.get("metadata") or {{}}).get("user_api_key_hash")}},
"response_id": getattr(response_obj, "id", None),
"received_model_id": hidden.get("received_model_id"),
}}
async with httpx.AsyncClient(trust_env=False) as client:
await client.post({url!r}, json=[event])
instance = HiddenParamsProbe()
"""
def hidden_params_probe(directory: Path, url: str) -> str:
(directory / "hidden_params_probe.py").write_text(_HIDDEN_PARAMS_PROBE.format(url=url))
return "hidden_params_probe.instance"
def sink_config(
directory: Path,
sinks: Mapping[str, str],
*,
retries: int = 0,
extra_callbacks: Sequence[str] = (),
store_prompts: bool = False,
) -> Path:
shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
settings: Final = {
**JSON_OBJECT.validate_python(shared["litellm_settings"]),
"callbacks": [*sinks, *extra_callbacks],
"DEFAULT_FLUSH_INTERVAL_SECONDS": 1,
}
callback_settings: Final = {
name: {
"callback_type": "generic_api",
"endpoint": url,
"headers": {"x-refusal-sink": name},
"max_retries": retries,
"retry_delay": 0.5,
}
for name, url in sinks.items()
}
general_settings: Final = {
**JSON_OBJECT.validate_python(shared["general_settings"]),
"store_prompts_in_spend_logs": store_prompts,
}
path: Final = directory / "anthropic-refusal-sinks.yaml"
path.write_text(
yaml.safe_dump(
{
**shared,
"general_settings": general_settings,
"litellm_settings": settings,
"callback_settings": callback_settings,
}
)
)
return path
def frames(text: str) -> tuple[Event, ...]:
return tuple(
JSON_OBJECT.validate_json(line.removeprefix("data: "))
for line in text.splitlines()
if line.startswith("data: {")
)
def frame_types(events: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]:
return tuple(event.get("type") for event in events)

18
tests/integration/_support/redis_breaker_worker.py
View file

@ -11,10 +11,12 @@ import sys
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from typing import Final, Literal, TypeAlias
from pydantic import TypeAdapter
Scenario: TypeAlias = Literal["callbacks_before_cache", "callbacks_after_cache", "no_callbacks"]
STATE_GAUGE: Final = "litellm_redis_circuit_breaker_state"
WORKER_TIMEOUT_SECONDS: Final = 120
@ -29,9 +31,11 @@ class WorkerReport:
_REPORT: Final = TypeAdapter(WorkerReport)
def run_breaker_worker(api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ) -> WorkerReport:
def run_breaker_worker(
scenario: Scenario, api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ
) -> WorkerReport:
completed: Final = subprocess.run(
(sys.executable, "-I", __file__, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]),
(sys.executable, "-I", __file__, scenario, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]),
env={**environment, "PROMETHEUS_MULTIPROC_DIR": str(prom_dir)},
capture_output=True,
text=True,
@ -56,11 +60,15 @@ def _series(families: Iterator[object]) -> Iterator[tuple[tuple[tuple[str, str],
yield tuple(sorted(sample.labels.items())), float(sample.value)
def main(api_base: str, redis_host: str, redis_port: str) -> None:
def main(scenario: str, api_base: str, redis_host: str, redis_port: str) -> None:
import litellm
from litellm.caching.caching import Cache
if scenario == "callbacks_before_cache":
litellm.service_callback = ["prometheus_system"] # test-quality-ok: this worker process exits right after the report
litellm.cache = Cache(type="redis", host=redis_host, port=redis_port) # test-quality-ok: throwaway worker process
if scenario == "callbacks_after_cache":
litellm.service_callback = ["prometheus_system"] # test-quality-ok: throwaway worker process
response: Final = litellm.completion(
model="openai/gpt-5.6",
api_base=api_base,
@ -76,4 +84,4 @@ def main(api_base: str, redis_host: str, redis_port: str) -> None:
if __name__ == "__main__":
main(*sys.argv[1:4])
main(*sys.argv[1:5])

89
tests/integration/_support/responses_stream.py
View file

@ -2,13 +2,17 @@ import json
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Final
from integration._support.client import object_value, string_value
import httpx
from integration._support.client import JSON_OBJECT, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.wire import Reply, Request
from pydantic import JsonValue
AZURE_TARGET: Final = "/openai/v1/responses?api-version="
OPENAI_TARGET: Final = "/responses"
RATE_LIMIT_MESSAGE: Final = "Your requests to gpt-6 have exceeded token rate limit."
RATE_LIMIT_PREFIX: Final = "litellm.RateLimitError: "
SENTINEL_PREFIX: Final = "litellm.MidStreamFallbackError: "
def frame(event: Mapping[str, JsonValue]) -> bytes:
@ -36,6 +40,14 @@ def error_event(error: Mapping[str, JsonValue] | None) -> Mapping[str, JsonValue
return {"type": "error", "sequence_number": 1, **({} if error is None else {"error": dict(error)})}
def openai_error_event(code: str, message: str) -> Mapping[str, JsonValue]:
return {"type": "error", "sequence_number": 1, "code": code, "message": message, "param": None}
def openai_error_stream(identity: str, code: str, message: str) -> tuple[bytes, ...]:
return (frame(created(identity)), frame(openai_error_event(code, message)))
def azure_rate_limit() -> Mapping[str, JsonValue]:
return {
"type": "too_many_requests",
@ -126,3 +138,78 @@ def chat_content(frames: Sequence[Mapping[str, JsonValue]]) -> str:
yield string_value(object_value(object_value(choice)["delta"]).get("content") or "")
return "".join(deltas())
def chat_body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]:
return {
"model": model,
"stream": True,
"messages": [{"role": "user", "content": marker}],
"tools": function_tools(),
"num_retries": 0,
"cache": {"no-cache": True},
**extra,
}
def messages_body(model: str, marker: str) -> dict[str, JsonValue]:
return {
"model": model,
"max_tokens": 64,
"stream": True,
"messages": [{"role": "user", "content": marker}],
"tools": [
{
"name": "get_weather",
"description": "Weather for a city",
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
}
],
"num_retries": 0,
"cache": {"no-cache": True},
}
def error_body(response: httpx.Response) -> Mapping[str, JsonValue]:
return object_value(JSON_OBJECT.validate_json(response.content)["error"])
def data_frames(text: str) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(
JSON_OBJECT.validate_json(line.removeprefix("data: "))
for line in text.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
def sse_events(text: str) -> tuple[tuple[str, Mapping[str, JsonValue]], ...]:
def parse(block: str) -> tuple[str, Mapping[str, JsonValue]]:
lines: Final = block.splitlines()
event: Final = next(line.removeprefix("event: ") for line in lines if line.startswith("event: "))
data: Final = next(line.removeprefix("data: ") for line in lines if line.startswith("data: "))
return event, JSON_OBJECT.validate_json(data)
return tuple(parse(block) for block in text.strip().split("\n\n") if "event: " in block)
def spend_row(call_id: str) -> Mapping[str, JsonValue]:
rows: Final = eventually(
lambda: read_rows(
'SELECT litellm_call_id, status, model_group, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s',
(call_id,),
),
lambda found: len(found) >= 1,
seconds=70,
)
assert len(rows) == 1, rows
return rows[0]
def assert_provider_typed_rate_limit(
error: Mapping[str, JsonValue], provider_message: str = RATE_LIMIT_MESSAGE
) -> None:
assert error["type"] == "throttling_error", error
assert str(error["code"]) == "429", error
message: Final = string_value(error["message"])
assert message.startswith(RATE_LIMIT_PREFIX) and provider_message in message, message
assert SENTINEL_PREFIX not in message, message

23
tests/integration/caching/response_cache_case.py
View file

@ -336,6 +336,29 @@ def flip_entry_to_empty(store: Redis, key: str, field: str) -> str:
return stale_id
def rewrite_usage_details(store: Redis, key: str, details: JsonValue) -> str:
"""Rewrite the stored usage's ``completion_tokens_details`` to ``details`` and stamp the id ``-stale``, keeping the
encoding, and return the stale id."""
raw: Final = store.get(key)
assert raw is not None, key
entry, is_json = _decode_entry(raw)
response: Final = entry["response"]
stored: Final = object_value(json.loads(response) if isinstance(response, str) else response)
stale_id: Final = f"{string_value(stored['id'])}{STALE_ID_SUFFIX}"
usage: Final = object_value(stored["usage"])
rewritten_response: Final[dict[str, JsonValue]] = {
**stored,
"usage": {**usage, "completion_tokens_details": details},
"id": stale_id,
}
rewritten: Final[dict[str, JsonValue]] = {
**entry,
"response": json.dumps(rewritten_response) if isinstance(response, str) else rewritten_response,
}
store.set(key, json.dumps(rewritten) if is_json else str(rewritten), ex=STALE_ENTRY_TTL_SECONDS)
return stale_id
def chat_body(model: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]:
return {"model": model, "messages": [{"role": "user", "content": text}], **extra}

58
tests/integration/database/test_prisma_migration_lock.py Normal file
View file

@ -0,0 +1,58 @@
import os
import secrets
from typing import Final
import psycopg
import pytest
from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, PRISMA_ADVISORY_LOCK_KEY, prisma_migration_lock
pytestmark = pytest.mark.skipif("DATABASE_URL" not in os.environ, reason="needs a live Postgres")
def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
row: Final = connection.execute("SELECT pg_try_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)).fetchone()
taken: Final = row is not None and row[0] is True
if taken:
connection.execute("SELECT pg_advisory_unlock(%s)", (PRISMA_ADVISORY_LOCK_KEY,))
return taken
def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]:
rows: Final = connection.execute(
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted"
).fetchall()
return frozenset(int(row[0]) for row in rows)
def _index_builds_concurrently_now(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
table: Final = f'"lit_cic_{secrets.token_hex(4)}"'
connection.execute(f"CREATE TABLE {table} (a int)")
try:
connection.execute("SET lock_timeout = '5s'")
connection.execute(f"CREATE INDEX CONCURRENTLY ON {table} (a)")
except psycopg.errors.LockNotAvailable:
return False
finally:
connection.execute("RESET lock_timeout")
connection.execute(f"DROP TABLE {table}")
return True
def test_prisma_migration_lock_does_not_stall_a_concurrent_index_build():
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
with prisma_migration_lock(os.environ["DATABASE_URL"]):
assert _index_builds_concurrently_now(peer)
def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leaves_nothing_behind():
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
assert _can_take_prismas_session_lock(peer)
with prisma_migration_lock(os.environ["DATABASE_URL"]):
held: Final = _granted_advisory_keys(peer)
assert PRISMA_ADVISORY_LOCK_KEY in held
assert MIGRATION_LOCK_KEY in held
assert not _can_take_prismas_session_lock(peer)
assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY})
assert _can_take_prismas_session_lock(peer)

81
tests/integration/management/test_team_update.py
View file

@ -162,6 +162,87 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
_MEMBER_DEFAULTS = [
pytest.param("team_member_budget", 25.0, "max_budget", id="team_member_budget"),
pytest.param("team_member_budget_duration", "7d", "budget_duration", id="team_member_budget_duration"),
pytest.param("team_member_tpm_limit", 200, "tpm_limit", id="team_member_tpm_limit"),
pytest.param("team_member_rpm_limit", 9, "rpm_limit", id="team_member_rpm_limit"),
]
@pytest.mark.parametrize("field,value,budget_column", _MEMBER_DEFAULTS)
async def test_team_update_team_admin_member_default_follows_the_proxy_admin_grant(
field: str, value: str | float, budget_column: str, gateway_client, db, resources, actors
):
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "alpha", resources.prefix)
headers = {"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"}
body = {"team_id": resources.prefix, field: value}
with team_admin_permissions(gateway_client.gateway, ["tpm_limit"]):
denied = await gateway_client.post("/team/update", headers=headers, json=body)
assert denied.status_code == 403, denied.text
assert f"'{field}'" in denied.text and "Team admin editable fields" in denied.text, denied.text
with team_admin_permissions(gateway_client.gateway, [field]):
allowed = await gateway_client.post("/team/update", headers=headers, json=body)
assert allowed.status_code == 200, allowed.text
info = await gateway_client.get("/team/info", params={"team_id": resources.prefix})
assert info.status_code == 200, info.text
assert info.json()["team_info"]["team_member_budget_table"][budget_column] == value
async def test_team_update_team_admin_resending_member_defaults_leaves_the_budget_row_alone(
gateway_client, db, resources, actors
):
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "alpha", resources.prefix)
stored = {
"team_member_budget": 10.0,
"team_member_budget_duration": "30d",
"team_member_tpm_limit": 100,
"team_member_rpm_limit": 5,
}
seeded = await gateway_client.post("/team/update", json={"team_id": resources.prefix, **stored})
assert seeded.status_code == 200, seeded.text
team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert team is not None
budget_id = team.metadata["team_member_budget_id"]
before = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id})
assert before is not None
with team_admin_permissions(gateway_client.gateway, ["tpm_limit"]):
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "tpm_limit": 1000, **stored},
)
assert resp.status_code == 200, resp.text
after = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id})
assert after is not None and after.updated_at == before.updated_at
team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert team is not None and team.tpm_limit == 1000
async def test_team_update_team_admin_cannot_change_member_defaults_of_another_team(
gateway_client, db, resources, actors
):
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "beta", resources.prefix)
with team_admin_permissions(gateway_client.gateway, ["team_member_budget"]):
resp = await gateway_client.post(
"/team/update",
headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": resources.prefix, "team_member_budget": 25.0},
)
assert resp.status_code == 403, resp.text
team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
assert team is not None and "team_member_budget_id" not in (team.metadata or {}), "denied but team mutated"
async def test_team_update_team_admin_soft_budget_follows_the_proxy_admin_grant(gateway_client, db, resources, actors):
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "alpha", resources.prefix)

18
tests/integration/providers/_cache_control_marks_support.py
View file

@ -20,6 +20,7 @@ BEDROCK_CACHE_CHECKPOINT_CAP: Final = 4
ANTHROPIC_MODEL: Final = "claude-opus-5-5"
BEDROCK_MODEL: Final = "anthropic.claude-opus-5-5"
UNMAPPED_BEDROCK_ARN: Final = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/scripted"
PROVIDER_KEY: Final = "synthetic-provider-key"
SYSTEM: Final = "Use the provided tool results to answer the user."
ASK: Final = "Look up the weather in London, Paris, and Tokyo."
@ -409,6 +410,23 @@ def bedrock_peer(request: Request) -> Reply:
)
def _without_cache_points(value: JsonValue) -> JsonValue:
if isinstance(value, dict):
return {key: _without_cache_points(inner) for key, inner in value.items()}
if isinstance(value, list):
return [
_without_cache_points(inner) for inner in value if not (isinstance(inner, dict) and "cachePoint" in inner)
]
return value
def sent_without_cache_points(requests: Sequence[Request], marker: str) -> list[JsonValue]:
return [
_without_cache_points(_JSON_OBJECT.validate_json(request.body.replace(marker.encode(), b"marker")))
for request in requests
]
def gateway_injected(response_id: str) -> bool:
rows: Final = eventually(
lambda: read_rows(

337
tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py Normal file
View file

@ -0,0 +1,337 @@
import asyncio
import json
import re
import threading
import time
import uuid
from collections import Counter
from collections.abc import Callable, Mapping, Sequence
from contextlib import ExitStack
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from urllib.parse import urlsplit
import httpx
import psutil
import pytest
import yaml
from integration._support.anthropic_refusal import (
REFUSAL_FIELDS,
Ledger,
Location,
refusal_prompt,
relayed,
respond,
sink_config,
)
from integration._support.anthropic_thinking import JSON_OBJECT, MODEL, Event, chunks_of
from integration._support.client import Gateway, eventually, held
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.timing import assert_not_delayed_by
from integration._support.wire import Reply, Request, Wire, wire_server
pytestmark: Final = [
pytest.mark.timeout(2 * graceful_stop_seconds() + 120),
pytest.mark.xdist_group("anthropic_refusal_stop_details_chaos"),
]
_MODEL_NAME: Final = "anthropic-refusal-chaos"
_SLOW_SINK_SECONDS: Final = 4
_SINK: Final = "refusal_chaos_sink"
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields")
_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields")
_PATHS: Final[Mapping[str, str]] = MappingProxyType(
{"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"}
)
Surface: TypeAlias = Literal["chat", "messages", "responses"]
_RecordProperty: TypeAlias = Callable[[str, object], None]
@dataclass(frozen=True, slots=True)
class _Call:
surface: Surface
stream: bool
marker: str
@dataclass(frozen=True, slots=True)
class _Served:
call: _Call
status: int
text: str
@dataclass(frozen=True, slots=True)
class _Held:
release: threading.Event
markers: SimpleQueue[str]
def respond(self, request: Request) -> Reply:
self.markers.put(_markers_in(request.body.decode())[0])
assert self.release.wait(timeout=90), "The burst was never released"
return respond(request)
def _markers_in(text: str) -> tuple[str, ...]:
return tuple(dict.fromkeys(_MARKER.findall(text)))
def _calls(count: int) -> tuple[_Call, ...]:
surfaces: Final[tuple[Surface, ...]] = ("chat", "messages", "responses")
return tuple(
_Call(surfaces[index % 3], stream=(index // 3) % 2 == 0, marker=uuid.uuid4().hex) for index in range(count)
)
def _body(call: _Call) -> Event:
text: Final = refusal_prompt("refusal", call.marker)
shared: Final[Event] = {"model": _MODEL_NAME, "stream": call.stream, "cache": {"no-cache": True}}
match call.surface:
case "chat":
return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64}
case "messages":
return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64}
case "responses":
return {**shared, "input": text, "max_output_tokens": 64}
async def _send(client: httpx.AsyncClient, key: str, call: _Call) -> _Served:
try:
response: Final = await client.post(
_PATHS[call.surface], json=_body(call), headers={"Authorization": f"Bearer {key}"}
)
except httpx.TransportError as error:
return _Served(call, 0, repr(error))
return _Served(call, response.status_code, response.text)
async def _burst(gateway: Gateway, calls: Sequence[_Call]) -> tuple[_Served, ...]:
async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client:
return tuple(await asyncio.gather(*(_send(client, gateway.key, call) for call in calls)))
def _chaos_config(directory: Path, upstream: Wire, sink: str, *, retries: int) -> Path:
base: Final = sink_config(directory, {_SINK: sink}, retries=retries)
config: Final = JSON_OBJECT.validate_python(yaml.safe_load(base.read_text()))
model_list: Final = [
{
"model_name": _MODEL_NAME,
"litellm_params": {"model": f"anthropic/{MODEL}", "api_base": upstream.url, "api_key": "scripted-key"},
}
]
path: Final = directory / "anthropic-refusal-chaos.yaml"
path.write_text(yaml.safe_dump({**config, "model_list": model_list}))
return path
def _ledger_events(ledgers: Sequence[Ledger]) -> tuple[Event, ...]:
return tuple(chain.from_iterable(ledger.events() for ledger in ledgers))
def _marker_counts(ledgers: Sequence[Ledger], markers: frozenset[str]) -> Mapping[str, int]:
found: Final = chain.from_iterable(_markers_in(json.dumps(event)) for event in _ledger_events(ledgers))
return MappingProxyType(Counter(marker for marker in found if marker in markers))
def _events_for(ledgers: Sequence[Ledger], marker: str) -> tuple[Event, ...]:
return tuple(event for event in _ledger_events(ledgers) if marker in json.dumps(event))
def _assert_answered(item: _Served) -> None:
assert item.status == 200, (item.call, item.text)
if item.call.surface != "chat":
return
if item.call.stream:
found: Final = tuple(chain.from_iterable(relayed(chunk) for chunk in chunks_of(item.text)))
assert found == ((_DELTA_FIELDS, REFUSAL_FIELDS),), item.text
return
assert relayed(JSON_OBJECT.validate_json(item.text)) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), item.text
def _assert_logged_chat_relays(ledgers: Sequence[Ledger], calls: Sequence[_Call]) -> None:
for call in calls:
if call.surface != "chat":
continue
(event,) = _events_for(ledgers, call.marker)
assert relayed(event["response"]) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), event
def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]:
return eventually(
lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
lambda pids: len(pids) == 2,
seconds=30,
)
def _open_upstream_connections(pid: int, upstream: str) -> int:
port: Final = urlsplit(upstream).port
return sum(
1
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
)
def _port(wire: Wire) -> int:
port: Final = urlsplit(wire.url).port
assert port is not None, wire.url
return port
def _accept(_request: Request) -> Reply:
return Reply()
async def _timed_chat(gateway: Gateway) -> float:
started: Final = time.monotonic()
(served,) = await _burst(gateway, (_Call("chat", stream=False, marker=uuid.uuid4().hex),))
_assert_answered(served)
return time.monotonic() - started
async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: Gateway, tmp_path: Path) -> None:
calls: Final = _calls(30)
markers: Final = frozenset(call.marker for call in calls)
with wire_server(respond) as upstream, ExitStack() as first_sink:
sink: Final = first_sink.enter_context(wire_server(_accept))
first: Final = Ledger(sink)
config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=6)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
_worker_pids(owned)
served_first: Final = await _burst(owned.gateway, calls[:10])
await asyncio.to_thread(
eventually, lambda: sum(_marker_counts((first,), markers).values()), lambda total: total >= 1, 30
)
first.events()
first_sink.close()
served_during: Final = await _burst(owned.gateway, calls[10:20])
with wire_server(_accept, port=_port(sink)) as restarted:
second: Final = Ledger(restarted)
served_after: Final = await _burst(owned.gateway, calls[20:])
for item in (*served_first, *served_during, *served_after):
_assert_answered(item)
counts: Final = await asyncio.to_thread(
held,
lambda: _marker_counts((first, second), markers),
lambda found: set(found) == markers and set(found.values()) == {1},
holding=3,
seconds=60,
)
assert sum(counts.values()) == 30, counts
_assert_logged_chat_relays((first, second), calls)
assert len(upstream.drain()) == 30
async def test_slow_sink_neither_deadlocks_nor_uploads_a_refusal_event_twice(gateway: Gateway, tmp_path: Path) -> None:
calls: Final = _calls(30)
markers: Final = frozenset(call.marker for call in calls)
def slow(_request: Request) -> Reply:
time.sleep(_SLOW_SINK_SECONDS)
return Reply()
with wire_server(respond) as upstream, wire_server(slow) as sink:
ledger: Final = Ledger(sink)
config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
_worker_pids(owned)
control: Final = await _timed_chat(owned.gateway)
served: Final = await _burst(owned.gateway, calls)
for item in served:
_assert_answered(item)
follow_up: Final = await _timed_chat(owned.gateway)
assert_not_delayed_by((follow_up,), control=(control,), injected_delay=_SLOW_SINK_SECONDS)
counts: Final = await asyncio.to_thread(
held,
lambda: _marker_counts((ledger,), markers),
lambda found: set(found) == markers and set(found.values()) == {1},
holding=5,
seconds=90,
)
assert sum(counts.values()) == 30, counts
_assert_logged_chat_relays((ledger,), calls)
async def test_worker_sigkill_mid_burst_leaves_the_sibling_relaying_and_logging_once(
gateway: Gateway, tmp_path: Path
) -> None:
calls: Final = _calls(20)
hold: Final = _Held(threading.Event(), SimpleQueue())
with wire_server(hold.respond) as upstream, wire_server(_accept) as sink:
ledger: Final = Ledger(sink)
config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
workers: Final = _worker_pids(owned)
burst: Final = asyncio.create_task(_burst(owned.gateway, calls))
await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60)
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, upstream.url) for pid in workers})
assert sum(held_by.values()) == 20, held_by
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
with owned.replacing_worker(victim_pid) as kill:
kill.freeze()
await asyncio.to_thread(kill)
hold.release.set()
served: Final = await burst
assert held_by[survivor_pid] >= 10, held_by
completed: Final = tuple(item for item in served if item.status == 200)
assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served])
for item in completed:
_assert_answered(item)
follow_up: Final = _Call("chat", stream=True, marker=uuid.uuid4().hex)
(answered,) = await _burst(owned.gateway, (follow_up,))
_assert_answered(answered)
survivors: Final = frozenset((*(item.call.marker for item in completed), follow_up.marker))
counts: Final = await asyncio.to_thread(
held,
lambda: _marker_counts((ledger,), frozenset(call.marker for call in (*calls, follow_up))),
lambda found: set(found) == survivors and set(found.values()) == {1},
holding=3,
seconds=60,
)
assert sum(counts.values()) == len(survivors), counts
_assert_logged_chat_relays((ledger,), tuple(item.call for item in completed))
async def test_proxy_restart_mid_burst_never_logs_a_refusal_twice(
gateway: Gateway, tmp_path: Path, record_property: _RecordProperty
) -> None:
calls: Final = _calls(20)
markers: Final = frozenset(call.marker for call in calls)
hold: Final = _Held(threading.Event(), SimpleQueue())
with wire_server(hold.respond) as upstream, wire_server(_accept) as sink:
ledger: Final = Ledger(sink)
config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first:
_worker_pids(first)
burst: Final = asyncio.create_task(_burst(first.gateway, calls))
await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60)
first.begin_shutdown()
hold.release.set()
served: Final = await burst
for item in served:
_assert_answered(item)
second_directory: Final = tmp_path / "second"
second_directory.mkdir()
with owned_proxy_process(gateway, second_directory, {}, config=config, workers=2) as second:
follow_up: Final = _Call("chat", stream=False, marker=uuid.uuid4().hex)
(answered,) = await _burst(second.gateway, (follow_up,))
_assert_answered(answered)
await asyncio.to_thread(
eventually,
lambda: _marker_counts((ledger,), frozenset({follow_up.marker})),
lambda found: found[follow_up.marker] == 1,
30,
)
counts: Final = _marker_counts((ledger,), markers)
assert set(counts.values()) <= {1}, counts
_assert_logged_chat_relays((ledger,), tuple(call for call in calls if call.marker in counts))
record_property("served", len(served))
record_property("landed", sum(counts.values()))

1356
tests/integration/providers/test_anthropic_refusal_stop_details_wire.py Normal file
View file

File diff suppressed because it is too large Load diff

12
tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py
View file

@ -15,10 +15,9 @@ import psutil
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.process import owned_proxy_process
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from integration.providers._cache_control_marks_support import (
ASK_LABEL,
CITIES,
POINTS,
SYSTEM_LABEL,
@ -40,6 +39,8 @@ from integration.providers._cache_control_marks_support import (
)
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
_SURFACES: Final = ("chat", "chat-stream", "chat-unmarked", "messages", "responses")
_MODEL: Final = "capped-claude"
@ -76,8 +77,6 @@ def _surface_request(surface: str, marker: str) -> tuple[str, dict[str, JsonValu
def _expected_labels(item: _Sent) -> list[str]:
if item.surface == "responses":
return [SYSTEM_LABEL, ASK_LABEL]
if item.surface == "chat-unmarked":
return [SYSTEM_LABEL, final_label(item.marker)]
return client_marked()
@ -148,7 +147,6 @@ def _single_spend_row(item: _Sent) -> None:
assert len(rows) == 1, (item.surface, item.call_id)
@pytest.mark.timeout(240)
async def test_capped_burst_rides_out_a_provider_outage_and_logs_every_request_once(
gateway: Gateway, tmp_path: Path
) -> None:
@ -183,7 +181,6 @@ async def test_capped_burst_rides_out_a_provider_outage_and_logs_every_request_o
_single_spend_row(item)
@pytest.mark.timeout(240)
async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_capped_requests(
gateway: Gateway, tmp_path: Path
) -> None:
@ -227,7 +224,6 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_capped_reques
_single_spend_row(item)
@pytest.mark.timeout(240)
def test_yaml_auto_caching_stands_down_for_tool_call_marks_and_outranks_a_key_opt_out(
gateway: Gateway, tmp_path: Path
) -> None:
@ -262,7 +258,7 @@ def test_yaml_auto_caching_stands_down_for_tool_call_marks_and_outranks_a_key_op
]
@pytest.mark.timeout(360)
@pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
async def test_proxy_restart_mid_burst_serves_capped_requests_after_the_reboot(
gateway: Gateway, tmp_path: Path
) -> None:

28
tests/integration/providers/test_cache_control_tool_call_marks_wire.py
View file

@ -20,7 +20,7 @@ from integration._support.client import (
string_value,
)
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.providers._cache_control_marks_support import (
ANTHROPIC_MODEL,
@ -61,6 +61,7 @@ from pydantic import JsonValue, TypeAdapter
from litellm.utils import get_prompt_cache_min_tokens
_OWNED_PROXY_TIMEOUT: Final = 2 * graceful_stop_seconds() + 120
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_LANE_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
_CLIENT_MARKED: Final = client_marked()
@ -505,17 +506,17 @@ def test_messages_endpoint_points_skip_injection_when_tool_use_marks_fill_the_ca
assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED
def test_responses_bridge_keeps_system_and_user_marks_within_the_cap(gateway: Gateway) -> None:
def test_responses_bridge_points_skip_injection_when_function_call_marks_fill_the_cap(gateway: Gateway) -> None:
marker: Final = new_marker()
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = _anthropic_deployment(scenario, wire, cache_control_injection_points=POINTS)
response: Final = gateway.request("POST", "/v1/responses", responses_body(model, marker))
assert response.status_code == 200, response.text
assert _JSON_OBJECT.validate_json(response.content)["status"] == "completed", response.text
assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL, ASK_LABEL]
assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED
@pytest.mark.timeout(240)
@pytest.mark.timeout(_OWNED_PROXY_TIMEOUT)
def test_response_cache_serves_the_capped_request_once(uncapped_gateway: Gateway) -> None:
marker: Final = new_marker()
with wire_server(anthropic_peer) as wire, uncapped_gateway.scenario() as scenario:
@ -567,14 +568,16 @@ def test_assistant_without_tool_calls_keeps_configured_points(gateway: Gateway,
assert labels == [SYSTEM_LABEL, ASK_LABEL, final_label(marker)]
def test_responses_stream_bridge_keeps_system_and_user_marks_within_the_cap(gateway: Gateway) -> None:
def test_responses_stream_bridge_points_skip_injection_when_function_call_marks_fill_the_cap(
gateway: Gateway,
) -> None:
marker: Final = new_marker()
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = _anthropic_deployment(scenario, wire, cache_control_injection_points=POINTS)
status, text = _stream(gateway, "/v1/responses", {**responses_body(model, marker), "stream": True})
assert status == 200, text
assert '"type":"response.completed"' in text, text
assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL, ASK_LABEL]
assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED
def _openai_client(gateway: Gateway) -> openai.OpenAI:
@ -698,7 +701,9 @@ async def test_anthropic_sdk_async_messages_stream_keeps_the_capped_request_with
assert gateway_injected(f"msg_{marker}") is False
async def test_openai_sdk_async_responses_keeps_system_and_user_marks_within_the_cap(gateway: Gateway) -> None:
async def test_openai_sdk_async_responses_points_skip_injection_when_function_call_marks_fill_the_cap(
gateway: Gateway,
) -> None:
marker: Final = new_marker()
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
model: Final = _anthropic_deployment(scenario, wire, cache_control_injection_points=POINTS)
@ -712,7 +717,7 @@ async def test_openai_sdk_async_responses_keeps_system_and_user_marks_within_the
input=cast(ResponseInputParam, body["input"]),
)
assert response.status == "completed", response.model_dump()
assert anthropic_labels(_only_request(wire)) == [SYSTEM_LABEL, ASK_LABEL]
assert anthropic_labels(_only_request(wire)) == _CLIENT_MARKED
def test_request_level_points_skip_injection_when_client_marks_fill_the_cap(gateway: Gateway) -> None:
@ -859,7 +864,7 @@ def _cache_hit_rows(response_id: str) -> list[dict[str, JsonValue]]:
)
@pytest.mark.timeout(240)
@pytest.mark.timeout(_OWNED_PROXY_TIMEOUT)
def test_response_cache_hit_records_the_injection_on_the_first_row_only(uncapped_gateway: Gateway) -> None:
marker: Final = new_marker()
unmarked: Final = [tool_call(city) for city in CITIES]
@ -877,7 +882,7 @@ def test_response_cache_hit_records_the_injection_on_the_first_row_only(uncapped
assert hit_rows[0]["injected"] is None, hit_rows
@pytest.mark.timeout(240)
@pytest.mark.timeout(_OWNED_PROXY_TIMEOUT)
@pytest.mark.parametrize("path", ("/v1/messages", "/v1/responses"))
def test_response_cache_twins_on_messages_and_responses_stay_within_the_cap(
uncapped_gateway: Gateway, path: str
@ -891,8 +896,7 @@ def test_response_cache_twins_on_messages_and_responses_stay_within_the_cap(
responses: Final = tuple(uncapped_gateway.request("POST", path, body) for _ in range(2))
received: Final = wire.drain()
assert [response.status_code for response in responses] == [200, 200], [response.text for response in responses]
expected: Final = _CLIENT_MARKED if path == "/v1/messages" else [SYSTEM_LABEL, ASK_LABEL]
assert [anthropic_labels(request) for request in received] == [expected] * len(received)
assert [anthropic_labels(request) for request in received] == [_CLIENT_MARKED] * len(received)
assert len(received) == 1, [request.target for request in received]

85
tests/integration/providers/test_chatgpt_responses_top_level_error_wire.py Normal file
View file

@ -0,0 +1,85 @@
import uuid
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from typing import Final
import httpx
import pytest
from integration._support import codex_vendor as cv
from integration._support.client import Gateway, gateway_from_environment, object_value, string_value
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.responses_stream import (
RATE_LIMIT_MESSAGE,
RATE_LIMIT_PREFIX,
data_frames,
error_body,
openai_error_stream,
)
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_MODEL: Final = "chatgpt/gpt-6.1-sol"
_TOP_LEVEL_RATE_LIMIT: Final = "codex-top-level-rate-limit"
_VENDOR: Final = cv.CodexVendor()
@dataclass(frozen=True, slots=True)
class _Rig:
gateway: Gateway
wire: Wire
def _respond(request: Request) -> Reply:
if _TOP_LEVEL_RATE_LIMIT.encode() not in request.body:
return _VENDOR.respond(request)
stream: Final = openai_error_stream("resp_" + uuid.uuid4().hex, "rate_limit_exceeded", RATE_LIMIT_MESSAGE)
return Reply(content_type="text/event-stream", chunks=stream)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
directory: Final = tmp_path_factory.mktemp("chatgpt-top-level-error-rig")
with gateway_from_environment() as gateway, wire_server(_respond) as wire:
overrides: Final = {"CHATGPT_TOKEN_DIR": str(cv.login(directory)), "CHATGPT_API_BASE": wire.url}
config: Final = cv.proxy_config(directory, probe=False)
with owned_proxy_process(gateway, directory, overrides, config=config, workers=2) as owned:
yield _Rig(owned.gateway, wire)
def _rate_limited(rig: _Rig, *, stream: bool) -> httpx.Response:
rig.wire.drain()
marker: Final = uuid.uuid4().hex
with rig.gateway.scenario() as scenario:
model: Final = scenario.model(model=_MODEL, api_base=rig.wire.url, api_key=None)
body: Final[Mapping[str, JsonValue]] = {
"model": model,
"input": [{"role": "user", "content": f"Say marker-{marker} {_TOP_LEVEL_RATE_LIMIT}"}],
"stream": stream,
"cache": {"no-cache": True},
}
response: Final = rig.gateway.request("POST", "/v1/responses", body)
(request,) = rig.wire.drain()
cv.forwarded(request, marker)
return response
def test_chatgpt_stream_top_level_rate_limit_fails_with_the_provider_code_and_message(rig: _Rig) -> None:
response: Final = _rate_limited(rig, stream=True)
assert response.status_code == 200, response.text
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
frames: Final = data_frames(response.text)
assert [frame["type"] for frame in frames] == ["response.created", "response.failed"], response.text
error: Final = object_value(object_value(frames[-1]["response"])["error"])
assert error == {"code": "rate_limit_exceeded", "message": RATE_LIMIT_MESSAGE}, error
def test_chatgpt_aggregated_top_level_rate_limit_answers_429_with_the_provider_message(rig: _Rig) -> None:
response: Final = _rate_limited(rig, stream=False)
assert response.status_code == 429, (response.status_code, response.text)
assert response.headers["content-type"].startswith("application/json"), response.headers
message: Final = string_value(error_body(response)["message"])
assert message.startswith(RATE_LIMIT_PREFIX) and RATE_LIMIT_MESSAGE in message, message
assert "event:" not in message, message
assert rig.gateway.request("GET", "/health/liveliness", None).status_code == 200

66
tests/integration/providers/test_openai_responses_websocket_wire.py
View file

@ -1,30 +1,35 @@
import asyncio
import hashlib
import itertools
import json
import ssl
import threading
import uuid
from collections.abc import Iterator
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from queue import SimpleQueue
from typing import Final
from typing import Final, TypeAlias
import pytest
import websockets
from integration._support.client import Gateway, gateway_from_environment
from integration._support.process import owned_proxy
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.responses_stream import RATE_LIMIT_MESSAGE, openai_error_event
from integration._support.tls import server_context, write_self_signed_cert
from pydantic import JsonValue
from websockets.asyncio.server import ServerConnection, serve
pytestmark: Final = pytest.mark.timeout(180)
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
PROVIDER_MODEL: Final = "ws-peer-model"
USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}
TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"})
_Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]]
@dataclass(frozen=True, slots=True)
class Peer:
@ -58,8 +63,20 @@ def _events(response_id: str, text: str) -> tuple[dict[str, JsonValue], ...]:
)
def _seven(response_id: str) -> tuple[dict[str, JsonValue], ...]:
return _events(response_id, "seven")
def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]:
created: Final = _events(response_id, "")[0]
return (created, dict(openai_error_event("rate_limit_exceeded", RATE_LIMIT_MESSAGE)))
async def _answer(
connection: ServerConnection, paths: SimpleQueue[str], frames: SimpleQueue[dict[str, JsonValue]]
connection: ServerConnection,
paths: SimpleQueue[str],
frames: SimpleQueue[dict[str, JsonValue]],
script: _Script,
) -> None:
paths.put(connection.request.path if connection.request is not None else "")
turns: Final = itertools.count(1)
@ -68,7 +85,7 @@ async def _answer(
frames.put(frame)
if frame.get("type") != "response.create":
continue
for event in _events(f"resp_peer_{next(turns)}", "seven"):
for event in script(f"resp_peer_{next(turns)}"):
await connection.send(json.dumps(event))
@ -78,21 +95,24 @@ async def _serve(
frames: SimpleQueue[dict[str, JsonValue]],
ports: SimpleQueue[int],
stop: asyncio.Event,
script: _Script,
) -> None:
async with serve(lambda connection: _answer(connection, paths, frames), "127.0.0.1", 0, ssl=tls) as server:
async with serve(lambda connection: _answer(connection, paths, frames, script), "127.0.0.1", 0, ssl=tls) as server:
ports.put(next(iter(server.sockets)).getsockname()[1])
await stop.wait()
@contextmanager
def responses_peer(cert: tuple[Path, Path]) -> Iterator[Peer]:
def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterator[Peer]:
loop: Final = asyncio.new_event_loop()
stop: Final = asyncio.Event()
paths: Final = SimpleQueue[str]()
frames: Final = SimpleQueue[dict[str, JsonValue]]()
ports: Final = SimpleQueue[int]()
thread: Final = threading.Thread(
target=loop.run_until_complete, args=(_serve(server_context(*cert), paths, frames, ports, stop),), daemon=True
target=loop.run_until_complete,
args=(_serve(server_context(*cert), paths, frames, ports, stop, script),),
daemon=True,
)
thread.start()
try:
@ -207,3 +227,29 @@ def test_previous_response_id_from_the_first_turn_reaches_the_provider_as_its_ow
assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == list(texts)
assert "previous_response_id" not in forwarded[0]
assert forwarded[1]["previous_response_id"] == "resp_peer_1"
def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_and_message(
candidate: Gateway, cert: tuple[Path, Path]
) -> None:
with responses_peer(cert, _top_level_rate_limit) as peer, candidate.scenario() as scenario:
model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url)
key: Final = scenario.key(models=[model])
text: Final = f"say seven {uuid.uuid4().hex}"
(events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (text,)))
assert [event["type"] for event in events] == ["response.created", "error"], events
assert (events[-1]["code"], events[-1]["message"]) == ("rate_limit_exceeded", RATE_LIMIT_MESSAGE), events
(forwarded,) = _drain(peer.frames)
assert forwarded["input"][0]["content"][0]["text"] == text
(row,) = eventually(
lambda: read_rows(
'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
(hashlib.sha256(key.encode()).hexdigest(),),
),
lambda rows: len(rows) >= 1,
seconds=70,
)
assert row["status"] == "failure", row
error: Final = object_value(object_value(row["metadata"])["error_information"])
assert str(error["error_code"]) == "429", error
assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error

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