mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/migrate-simple-tables-ac3786
# Conflicts: # ui/litellm-dashboard/eslint-suppressions.json
This commit is contained in:
commit
ce8ecd4fcf
179 changed files with 16449 additions and 1902 deletions
1
.github/CODEOWNERS
vendored
1
.github/CODEOWNERS
vendored
|
|
@ -1,2 +1,3 @@
|
|||
/ui/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/ui/litellm-dashboard/src/lib/http/schema.d.ts
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ RUN uv venv --python python && \
|
|||
"opentelemetry-api==1.28.0" \
|
||||
"opentelemetry-sdk==1.28.0" \
|
||||
"opentelemetry-exporter-otlp==1.28.0" \
|
||||
"ddtrace==2.19.0" \
|
||||
"ddtrace==4.11.0" \
|
||||
"sentry-sdk==2.21.0" \
|
||||
"mangum==0.17.0" \
|
||||
"azure-ai-contentsafety==1.0.0" \
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@
|
|||
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
||||
## This provides an LLM Guard Integration for content moderation on the proxy
|
||||
|
||||
from typing import Literal, Optional
|
||||
import asyncio
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -18,7 +19,6 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
from litellm.utils import get_formatted_prompt
|
||||
|
||||
|
||||
class _ENTERPRISE_LLMGuard(CustomLogger):
|
||||
|
|
@ -46,45 +46,44 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
async def moderation_check(self, text: str):
|
||||
async def moderation_check(self, text: str) -> str:
|
||||
"""
|
||||
Runs the LLM Guard moderation check on ``text``.
|
||||
|
||||
Raises an HTTPException when the content violates the safety policy;
|
||||
otherwise returns the sanitized prompt from LLM Guard, falling back to
|
||||
the original text when the API does not provide one.
|
||||
|
||||
[TODO] make this more performant for high-throughput scenario
|
||||
"""
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
if self.mock_redacted_text is not None:
|
||||
redacted_text = self.mock_redacted_text
|
||||
else:
|
||||
# Make the first request to /analyze
|
||||
analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
|
||||
verbose_proxy_logger.debug("Making request to: %s", analyze_url)
|
||||
analyze_payload = {"prompt": text}
|
||||
redacted_text = None
|
||||
if self.mock_redacted_text is not None:
|
||||
redacted_text = self.mock_redacted_text
|
||||
else:
|
||||
analyze_url = f"{self.llm_guard_api_base}analyze/prompt"
|
||||
verbose_proxy_logger.debug("Making request to: %s", analyze_url)
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
analyze_url, json=analyze_payload
|
||||
analyze_url, json={"prompt": text}
|
||||
) as response:
|
||||
redacted_text = await response.json()
|
||||
verbose_proxy_logger.debug(
|
||||
f"LLM Guard: Received response - {redacted_text}"
|
||||
verbose_proxy_logger.debug(
|
||||
f"LLM Guard: Received response - {redacted_text}"
|
||||
)
|
||||
if redacted_text is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
"error": f"Invalid content moderation response: {redacted_text}"
|
||||
},
|
||||
)
|
||||
if redacted_text is not None:
|
||||
if (
|
||||
redacted_text.get("is_valid", None) is not None
|
||||
and redacted_text["is_valid"] is False
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Violated content safety policy"},
|
||||
)
|
||||
else:
|
||||
pass
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
"error": f"Invalid content moderation response: {redacted_text}"
|
||||
},
|
||||
)
|
||||
if redacted_text.get("is_valid", None) is False:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Violated content safety policy"},
|
||||
)
|
||||
sanitized_prompt = redacted_text.get("sanitized_prompt")
|
||||
return sanitized_prompt if isinstance(sanitized_prompt, str) else text
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.enterprise.enterprise_hooks.llm_guard::moderation_check - Exception occurred - {}".format(
|
||||
|
|
@ -138,23 +137,75 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
|
|||
return
|
||||
|
||||
self.print_verbose("Makes LLM Guard Check")
|
||||
try:
|
||||
assert call_type in [
|
||||
"completion",
|
||||
"embeddings",
|
||||
"image_generation",
|
||||
"moderation",
|
||||
"audio_transcription",
|
||||
]
|
||||
except Exception:
|
||||
if call_type not in [
|
||||
"completion",
|
||||
"embeddings",
|
||||
"image_generation",
|
||||
"moderation",
|
||||
"audio_transcription",
|
||||
]:
|
||||
self.print_verbose(
|
||||
f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']"
|
||||
)
|
||||
return data
|
||||
|
||||
formatted_prompt = get_formatted_prompt(data=data, call_type=call_type) # type: ignore
|
||||
self.print_verbose(f"LLM Guard, formatted_prompt: {formatted_prompt}")
|
||||
return await self.moderation_check(text=formatted_prompt)
|
||||
return await self._moderate_request(data=data)
|
||||
|
||||
async def _moderate_request(self, data: dict) -> dict:
|
||||
"""
|
||||
Sanitizes the request in place using the prompt returned by LLM Guard so
|
||||
the provider-bound request carries the redacted content, then returns it.
|
||||
"""
|
||||
messages = data.get("messages")
|
||||
if messages is not None:
|
||||
data["messages"] = list(
|
||||
await asyncio.gather(
|
||||
*(self._moderate_message(message) for message in messages)
|
||||
)
|
||||
)
|
||||
return data
|
||||
|
||||
input_ = data.get("input")
|
||||
if input_ is not None:
|
||||
data["input"] = await self._moderate_input(input_)
|
||||
return data
|
||||
|
||||
prompt = data.get("prompt")
|
||||
if isinstance(prompt, str):
|
||||
data["prompt"] = await self.moderation_check(text=prompt)
|
||||
return data
|
||||
|
||||
async def _moderate_message(self, message: dict) -> dict:
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return {**message, "content": await self.moderation_check(text=content)}
|
||||
if isinstance(content, list):
|
||||
return {
|
||||
**message,
|
||||
"content": list(
|
||||
await asyncio.gather(
|
||||
*(self._moderate_content_part(part) for part in content)
|
||||
)
|
||||
),
|
||||
}
|
||||
return message
|
||||
|
||||
async def _moderate_content_part(self, part: dict) -> dict:
|
||||
if part.get("type") == "text" and isinstance(part.get("text"), str):
|
||||
return {**part, "text": await self.moderation_check(text=part["text"])}
|
||||
return part
|
||||
|
||||
async def _moderate_input(self, input_: object) -> object:
|
||||
if isinstance(input_, str):
|
||||
return await self.moderation_check(text=input_)
|
||||
if isinstance(input_, list):
|
||||
return [
|
||||
await self.moderation_check(text=item)
|
||||
if isinstance(item, str)
|
||||
else item
|
||||
for item in input_
|
||||
]
|
||||
return input_
|
||||
|
||||
async def async_post_call_streaming_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, response: str
|
||||
|
|
|
|||
|
|
@ -54,6 +54,12 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
| `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` |
|
||||
| `pdb.labels` | Extra metadata labels to add to the PDB | `{}` |
|
||||
|
||||
| `billingMetrics.enabled` | Enable enterprise billable-request metering. Requires an enterprise license. | `false` |
|
||||
| `billingMetrics.endpoint` | Collector that the billable-request counter is pushed to. | `https://telemetry.litellm.ai` |
|
||||
| `billingMetrics.secretName` | Name of an existing Secret holding the mTLS client certificate, under the keys `tls.crt` and `tls.key`. | `litellm-billing-metrics-mtls` |
|
||||
| `billingMetrics.caSecretName` | Name of an existing Secret holding a CA bundle under the key `ca.crt`. Only needed for a private or test collector whose server certificate is not on the public web PKI. | `""` |
|
||||
| `billingMetrics.exportIntervalMs` | How often the counter is pushed, in milliseconds. The proxy defaults to `60000` when unset. | `""` |
|
||||
|
||||
#### Example `proxy_config` ConfigMap from values (default):
|
||||
|
||||
```
|
||||
|
|
@ -94,6 +100,21 @@ data:
|
|||
type: Opaque
|
||||
```
|
||||
|
||||
#### Enterprise billable-request metering
|
||||
|
||||
Enterprise licenses meter billable requests by pushing a counter to LiteLLM's collector over mutual TLS. The chart does not create the client certificate; it mounts one you already hold, read-only, so the private key is never exposed through the environment. Create the Secret under the name the chart expects, then turn the block on:
|
||||
|
||||
```
|
||||
kubectl create secret tls litellm-billing-metrics-mtls --cert=client.crt --key=client.key
|
||||
```
|
||||
|
||||
```
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
```
|
||||
|
||||
Set `billingMetrics.caSecretName` only when the collector is a private or test one whose server certificate is not on the public web PKI; the production collector needs no CA override. The chart fails the render rather than deploying a proxy that silently never exports, so a missing `secretName` or an emptied `endpoint` surfaces at `helm install` time.
|
||||
|
||||
### Database Settings
|
||||
|
||||
| Name | Description | Value |
|
||||
|
|
|
|||
|
|
@ -50,6 +50,53 @@ app.kubernetes.io/name: {{ include "litellm.name" . }}
|
|||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Enterprise billable-request metering. The client certificate identifies the
|
||||
deployment to LiteLLM's collector, so it is mounted read-only from an existing
|
||||
Secret rather than passed through the environment.
|
||||
*/}}
|
||||
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
|
||||
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsEnv" -}}
|
||||
- name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
|
||||
{{- end }}
|
||||
{{- with .Values.billingMetrics.exportIntervalMs }}
|
||||
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: {{ . | quote }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumes" -}}
|
||||
- name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: {{ .Values.billingMetrics.caSecretName }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumeMounts" -}}
|
||||
- name: billing-metrics-mtls
|
||||
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
|
||||
readOnly: true
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
|
||||
readOnly: true
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Create the name of the service account to use
|
||||
*/}}
|
||||
|
|
|
|||
|
|
@ -142,6 +142,9 @@ spec:
|
|||
{{- with .Values.extraEnvVars }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- if .Values.migrationJob.enabled }}
|
||||
# Schema updates are owned by the dedicated migrations Job; skip
|
||||
# the proxy's startup `prisma db push` so N replicas don't race
|
||||
|
|
@ -220,6 +223,9 @@ spec:
|
|||
- name: npm
|
||||
mountPath: /.npm
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.volumeMounts }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
|
|
@ -252,6 +258,9 @@ spec:
|
|||
items:
|
||||
- key: {{ .Values.proxyConfigMap.key | default "config.yaml" }}
|
||||
path: "config.yaml"
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.volumes }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
|
|
|||
297
helm/litellm-helm/tests/billing_metrics_tests.yaml
Normal file
297
helm/litellm-helm/tests/billing_metrics_tests.yaml
Normal file
|
|
@ -0,0 +1,297 @@
|
|||
suite: test billingMetrics wiring on the proxy deployment
|
||||
templates:
|
||||
- deployment.yaml
|
||||
- configmap-litellm.yaml
|
||||
- migrations-job.yaml
|
||||
tests:
|
||||
- it: is off by default, adding no env, volume, or mount
|
||||
template: deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
|
||||
- it: renders the endpoint and the mounted cert paths when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: /etc/litellm/billing-mtls/tls.crt
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: /etc/litellm/billing-mtls/tls.key
|
||||
|
||||
# The conventional Secret name is the default, so enabling the block is enough.
|
||||
- it: mounts the default cert secret read-only alongside the config volume
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
|
||||
- it: honours a secretName override
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: my-billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: my-billing-mtls
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
|
||||
- it: honours an endpoint override
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
endpoint: https://collector.internal:4318
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://collector.internal:4318
|
||||
|
||||
# The production collector presents a public web-PKI certificate, so the CA
|
||||
# override must stay absent unless a private collector is configured.
|
||||
- it: omits the CA env, volume, and mount when no caSecretName is set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
mountPath: /etc/litellm/billing-mtls-ca
|
||||
readOnly: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
|
||||
- it: mounts the CA secret when caSecretName is set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
caSecretName: billing-ca
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
mountPath: /etc/litellm/billing-mtls-ca
|
||||
readOnly: true
|
||||
|
||||
- it: passes the export interval through only when set
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
exportIntervalMs: 5000
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: "5000"
|
||||
|
||||
- it: omits the export interval when unset
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: "60000"
|
||||
|
||||
# Kubernetes resolves duplicate env names last-wins, so the chart-owned billing
|
||||
# entries must render after .Values.envVars or a user could silently redirect
|
||||
# the metering export. The three billing entries are the last ones emitted here
|
||||
# (migrationJob, which appends DISABLE_SCHEMA_UPDATE, is off for this case).
|
||||
- it: renders the billing endpoint after envVars so it cannot be shadowed
|
||||
template: deployment.yaml
|
||||
set:
|
||||
migrationJob:
|
||||
enabled: false
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
envVars:
|
||||
LITELLM_BILLING_METRICS_ENDPOINT: https://shadowed.example
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://shadowed.example
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[-3]
|
||||
value:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[-2].name
|
||||
value: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[-1].name
|
||||
value: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
|
||||
- it: keeps user-supplied volumes and mounts alongside the billing secret
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
volumes:
|
||||
- name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
volumeMounts:
|
||||
- name: custom-callbacks
|
||||
mountPath: /app/callbacks
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: custom-callbacks
|
||||
mountPath: /app/callbacks
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
|
||||
- it: still mounts the proxy config when enabled
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: litellm-config
|
||||
mountPath: /etc/litellm/config.yaml
|
||||
subPath: config.yaml
|
||||
|
||||
# Only the proxy serves billable traffic. The migrations Job must never mount
|
||||
# the client certificate, and it renders its own env and volumes, so nothing
|
||||
# stops a future edit from wiring the billing include into it by mistake.
|
||||
- it: does not touch the migrations job when enabled
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- notExists:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
- notExists:
|
||||
path: spec.template.spec.volumes
|
||||
|
||||
- it: fails loudly when enabled with an emptied secretName
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
|
||||
|
||||
- it: fails loudly when enabled without an endpoint
|
||||
template: deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
endpoint: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
|
||||
|
|
@ -139,6 +139,20 @@ masterkeySecretName: ""
|
|||
# if set, use this secret key for the master key; otherwise, use the default key
|
||||
masterkeySecretKey: ""
|
||||
|
||||
# Optional: enterprise billable-request metering. When enabled, the proxy counts
|
||||
# successful requests to inference, MCP, and A2A endpoints and pushes them to
|
||||
# LiteLLM's collector over mutual TLS. Requires an enterprise license.
|
||||
# The client certificate identifies the deployment, so it is mounted read-only
|
||||
# from an existing Secret and never passed through the environment.
|
||||
billingMetrics:
|
||||
enabled: false
|
||||
endpoint: https://telemetry.litellm.ai # collector to push the counter to
|
||||
secretName: litellm-billing-metrics-mtls # existing Secret holding tls.crt and tls.key
|
||||
# Only for private or test collectors whose server certificate is not on the
|
||||
# public web PKI. The production collector needs no CA override.
|
||||
caSecretName: "" # existing Secret holding ca.crt
|
||||
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
|
||||
|
||||
proxyConfigMap:
|
||||
# when true, creates a new configmap
|
||||
create: true
|
||||
|
|
|
|||
|
|
@ -34,6 +34,57 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
|
|||
helm.sh/chart: {{ printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Enterprise billable-request metering. Wired into gateway and backend, not the
|
||||
migrations job. The gateway serves nearly all billable traffic, but the backend
|
||||
keeps the named-server MCP transport (/{mcp_server_name}/mcp), which writes a
|
||||
SpendLogs row, so metering only the gateway would silently drop that traffic.
|
||||
The client certificate identifies the deployment to LiteLLM's collector, so it is
|
||||
mounted read-only from an existing Secret rather than passed through the
|
||||
environment.
|
||||
*/}}
|
||||
{{- define "litellm.billingMetrics.certDir" -}}/etc/litellm/billing-mtls{{- end -}}
|
||||
{{- define "litellm.billingMetrics.caDir" -}}/etc/litellm/billing-mtls-ca{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsEnv" -}}
|
||||
- name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: {{ required "billingMetrics.endpoint is required when billingMetrics.enabled is true" .Values.billingMetrics.endpoint | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: {{ printf "%s/tls.crt" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
- name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: {{ printf "%s/tls.key" (include "litellm.billingMetrics.certDir" .) | quote }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: {{ printf "%s/ca.crt" (include "litellm.billingMetrics.caDir" .) | quote }}
|
||||
{{- end }}
|
||||
{{- with .Values.billingMetrics.exportIntervalMs }}
|
||||
- name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: {{ . | quote }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumes" -}}
|
||||
- name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: {{ required "billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)" .Values.billingMetrics.secretName }}
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: {{ .Values.billingMetrics.caSecretName }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.billingMetricsVolumeMounts" -}}
|
||||
- name: billing-metrics-mtls
|
||||
mountPath: {{ include "litellm.billingMetrics.certDir" . }}
|
||||
readOnly: true
|
||||
{{- if .Values.billingMetrics.caSecretName }}
|
||||
- name: billing-metrics-mtls-ca
|
||||
mountPath: {{ include "litellm.billingMetrics.caDir" . }}
|
||||
readOnly: true
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Per-component selector labels — used in both Service selectors and Deployment matchLabels.
|
||||
*/}}
|
||||
|
|
|
|||
|
|
@ -44,14 +44,20 @@ spec:
|
|||
- name: CONFIG_FILE_PATH
|
||||
value: /app/config/config.yaml
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumeMounts .Values.billingMetrics.enabled }}
|
||||
volumeMounts:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.volumeMounts }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
|
|
@ -66,13 +72,16 @@ spec:
|
|||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.backend.resources | nindent 12 }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumes }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
configMap:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-config
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.volumes }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
|
|
|||
|
|
@ -46,14 +46,20 @@ spec:
|
|||
- name: NUM_WORKERS
|
||||
value: {{ .Values.gateway.numWorkers | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
|
||||
volumeMounts:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.volumeMounts }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
|
|
@ -68,13 +74,16 @@ spec:
|
|||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.gateway.resources | nindent 12 }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
configMap:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-config
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.volumes }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
|
|
|||
249
helm/litellm/tests/billing_metrics_tests.yaml
Normal file
249
helm/litellm/tests/billing_metrics_tests.yaml
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
suite: test billingMetrics wiring on gateway and backend
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- migrations-job.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: is off by default, adding no env, volume, or mount
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
value:
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
|
||||
- it: renders the endpoint and the mounted cert paths when enabled
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_CERT
|
||||
value: /etc/litellm/billing-mtls/tls.crt
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CLIENT_KEY
|
||||
value: /etc/litellm/billing-mtls/tls.key
|
||||
|
||||
- it: mounts the cert secret read-only alongside the config volume
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
|
||||
# The production collector presents a public web-PKI certificate, so the CA
|
||||
# override must stay absent unless a private collector is configured.
|
||||
- it: omits the CA env, volume, and mount when no caSecretName is set
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
|
||||
- it: mounts the CA secret when caSecretName is set
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
caSecretName: billing-ca
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_CA_CERT
|
||||
value: /etc/litellm/billing-mtls-ca/ca.crt
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
secret:
|
||||
secretName: billing-ca
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls-ca
|
||||
mountPath: /etc/litellm/billing-mtls-ca
|
||||
readOnly: true
|
||||
|
||||
- it: passes the export interval through only when set
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
exportIntervalMs: 5000
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS
|
||||
value: "5000"
|
||||
|
||||
- it: keeps user-supplied gateway volumes alongside the billing secret
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
gateway.volumes:
|
||||
- name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
gateway.volumeMounts:
|
||||
- name: custom-callbacks
|
||||
mountPath: /app/callbacks
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: custom-callbacks
|
||||
configMap:
|
||||
name: my-callbacks
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
|
||||
# The backend keeps the named-server MCP transport (/{mcp_server_name}/mcp),
|
||||
# which writes a SpendLogs row, so it must meter too or that traffic is lost.
|
||||
- it: meters the backend as well, since it serves the MCP transport
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
mountPath: /etc/litellm/billing-mtls
|
||||
readOnly: true
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: billing-mtls
|
||||
|
||||
- it: leaves the backend alone when metering is off
|
||||
template: backend/deployment.yaml
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
|
||||
# The migrations job runs prisma and serves no traffic; it must never receive
|
||||
# the client key.
|
||||
- it: never mounts the billing cert on the migrations job
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LITELLM_BILLING_METRICS_ENDPOINT
|
||||
value: https://telemetry.litellm.ai
|
||||
- isNull:
|
||||
path: spec.template.spec.volumes
|
||||
|
||||
# The conventional Secret name is the default, so enabling metering needs no
|
||||
# secretName at all; the guard below only fires on an explicitly blanked one.
|
||||
- it: uses the conventional secret name by default
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: billing-metrics-mtls
|
||||
secret:
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
|
||||
- it: fails loudly when the secretName is explicitly blanked
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
secretName: ""
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.secretName is required when billingMetrics.enabled is true (an existing Secret with tls.crt and tls.key)
|
||||
|
||||
- it: fails loudly when enabled without an endpoint
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
billingMetrics:
|
||||
enabled: true
|
||||
endpoint: ""
|
||||
secretName: billing-mtls
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: billingMetrics.endpoint is required when billingMetrics.enabled is true
|
||||
|
|
@ -73,6 +73,25 @@ masterKey:
|
|||
secretName: litellm-master-key-secret # name of a Secret containing the master key
|
||||
secretKey: master-key
|
||||
|
||||
# Optional: enterprise billable-request metering. When enabled, the gateway and
|
||||
# backend count successful requests to inference, MCP, and A2A endpoints and push
|
||||
# them to LiteLLM's collector over mutual TLS. Both components serve billable
|
||||
# routes: the backend keeps the named-server MCP transport. Requires an
|
||||
# enterprise license. The client certificate identifies the deployment, so it is
|
||||
# mounted read-only from an existing Secret and never passed through the env.
|
||||
billingMetrics:
|
||||
enabled: false
|
||||
endpoint: https://telemetry.litellm.ai # collector to push the counter to
|
||||
# An existing Secret holding the client certificate under tls.crt and its key
|
||||
# under tls.key, usually created from the onboarding artifact. The default is
|
||||
# the conventional name, so the common path is to create that Secret and set
|
||||
# enabled: true. Override only if yours is named differently.
|
||||
secretName: litellm-billing-metrics-mtls
|
||||
# Only for private or test collectors whose server certificate is not on the
|
||||
# public web PKI. The production collector needs no CA override.
|
||||
caSecretName: "" # existing Secret holding ca.crt
|
||||
exportIntervalMs: "" # push cadence; the proxy defaults to 60000
|
||||
|
||||
# External Postgres connection.
|
||||
database:
|
||||
writer:
|
||||
|
|
|
|||
|
|
@ -85,6 +85,22 @@ class CachingHandlerResponse(BaseModel):
|
|||
in_memory_cache_obj = InMemoryCache()
|
||||
|
||||
|
||||
def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]:
|
||||
"""
|
||||
The caching handler is stored on the Logging object
|
||||
(``logging_obj._llm_caching_handler``), so keeping ``litellm_logging_obj``
|
||||
inside ``request_kwargs`` closes a reference cycle
|
||||
(Logging -> LLMCachingHandler -> kwargs -> Logging) that keeps the full
|
||||
request payload (messages included) alive until a generational GC pass
|
||||
instead of being freed by refcount when the request ends. Nothing in the
|
||||
caching layer reads the logging object from these kwargs; cache-key
|
||||
generation ignores litellm-internal params.
|
||||
"""
|
||||
if "litellm_logging_obj" not in request_kwargs:
|
||||
return request_kwargs
|
||||
return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"}
|
||||
|
||||
|
||||
def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
|
||||
cached_id = cached_result.get("id")
|
||||
if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"):
|
||||
|
|
@ -118,7 +134,7 @@ class LLMCachingHandler:
|
|||
|
||||
self.async_streaming_chunks: List[ModelResponse] = []
|
||||
self.sync_streaming_chunks: List[ModelResponse] = []
|
||||
self.request_kwargs = request_kwargs
|
||||
self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs)
|
||||
self.preset_cache_key: Optional[str] = None
|
||||
self.original_function = original_function
|
||||
self.start_time = start_time
|
||||
|
|
@ -297,7 +313,7 @@ class LLMCachingHandler:
|
|||
new_kwargs.pop("metadata", None)
|
||||
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
|
||||
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
|
||||
self.request_kwargs = new_kwargs
|
||||
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
|
||||
print_verbose("Checking Sync Cache")
|
||||
cached_result = litellm.cache.get_cache(**new_kwargs)
|
||||
if cached_result is not None:
|
||||
|
|
@ -693,7 +709,7 @@ class LLMCachingHandler:
|
|||
new_kwargs.pop("metadata", None)
|
||||
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
|
||||
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
|
||||
self.request_kwargs = new_kwargs
|
||||
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
|
||||
cached_result: Optional[Any] = None
|
||||
if call_type == CallTypes.aembedding.value:
|
||||
if isinstance(new_kwargs["input"], str):
|
||||
|
|
|
|||
|
|
@ -1496,6 +1496,7 @@ MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20))
|
|||
MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000))
|
||||
DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7))
|
||||
LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16))
|
||||
MINIMUM_CUSTOM_KEY_LENGTH = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16))
|
||||
SECRET_MANAGER_REFRESH_INTERVAL = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400))
|
||||
LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [
|
||||
"default_internal_user_params",
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union
|
|||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ddtrace.tracer import Tracer as DD_TRACER
|
||||
from ddtrace.trace import Tracer as DD_TRACER
|
||||
else:
|
||||
DD_TRACER = Any
|
||||
|
||||
|
|
|
|||
|
|
@ -925,7 +925,6 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def pre_call(self, input, api_key, model=None, additional_args={}):
|
||||
# Log the exact input to the LLM API
|
||||
litellm.error_logs["PRE_CALL"] = locals()
|
||||
try:
|
||||
self._pre_call(
|
||||
input=input,
|
||||
|
|
@ -1135,7 +1134,6 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def post_call(self, original_response, input=None, api_key=None, additional_args={}):
|
||||
# Log the exact result from the LLM API, for streaming - log the type of response received
|
||||
litellm.error_logs["POST_CALL"] = locals()
|
||||
if isinstance(original_response, dict):
|
||||
original_response = json.dumps(original_response, default=str)
|
||||
try:
|
||||
|
|
@ -3074,7 +3072,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List:
|
||||
if dynamic_success_callbacks is None:
|
||||
return list(global_callbacks)
|
||||
return list(set(dynamic_success_callbacks + global_callbacks))
|
||||
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
|
||||
|
||||
def _remove_internal_litellm_callbacks(self, callbacks: List) -> List:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -77,6 +77,22 @@ def _redact_streaming_response(streaming_response):
|
|||
streaming_response.reasoning = None
|
||||
|
||||
|
||||
def _redact_tool_calls(tool_calls) -> None:
|
||||
"""Redact tool call arguments (assistant tool calls carry prompt-derived data)."""
|
||||
if not tool_calls:
|
||||
return
|
||||
for tool_call in tool_calls:
|
||||
function = getattr(tool_call, "function", None)
|
||||
if function is not None and hasattr(function, "arguments"):
|
||||
function.arguments = "redacted-by-litellm"
|
||||
|
||||
|
||||
def _redact_function_call(function_call) -> None:
|
||||
"""Redact legacy assistant function_call arguments."""
|
||||
if function_call is not None and hasattr(function_call, "arguments"):
|
||||
function_call.arguments = "redacted-by-litellm"
|
||||
|
||||
|
||||
def _redact_choice_content(choice):
|
||||
"""Helper to redact content in a choice (message or delta)."""
|
||||
if isinstance(choice, litellm.Choices):
|
||||
|
|
@ -85,12 +101,16 @@ def _redact_choice_content(choice):
|
|||
choice.message.reasoning_content = "redacted-by-litellm"
|
||||
if hasattr(choice.message, "thinking_blocks"):
|
||||
choice.message.thinking_blocks = None
|
||||
_redact_tool_calls(getattr(choice.message, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.message, "function_call", None))
|
||||
elif isinstance(choice, litellm.utils.StreamingChoices):
|
||||
choice.delta.content = "redacted-by-litellm"
|
||||
if hasattr(choice.delta, "reasoning_content"):
|
||||
choice.delta.reasoning_content = "redacted-by-litellm"
|
||||
if hasattr(choice.delta, "thinking_blocks"):
|
||||
choice.delta.thinking_blocks = None
|
||||
_redact_tool_calls(getattr(choice.delta, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.delta, "function_call", None))
|
||||
|
||||
|
||||
def _redact_responses_api_output(output_items):
|
||||
|
|
@ -111,6 +131,9 @@ def _redact_responses_api_output(output_items):
|
|||
if hasattr(summary_item, "text"):
|
||||
summary_item.text = "redacted-by-litellm"
|
||||
|
||||
if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"):
|
||||
output_item.arguments = "redacted-by-litellm"
|
||||
|
||||
|
||||
def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
||||
"""Helper to redact ResponsesAPIResponse output items in dict form."""
|
||||
|
|
@ -131,6 +154,9 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
|||
if isinstance(summary_item, dict) and "text" in summary_item:
|
||||
summary_item["text"] = redacted_str
|
||||
|
||||
if output_item.get("type") == "function_call" and "arguments" in output_item:
|
||||
output_item["arguments"] = redacted_str
|
||||
|
||||
|
||||
def _redact_standard_logging_object(model_call_details: dict):
|
||||
"""Redact messages and response inside standard_logging_object if present."""
|
||||
|
|
@ -162,6 +188,19 @@ def _redact_standard_logging_object(model_call_details: dict):
|
|||
standard_logging_object["response"] = {"text": redacted_str}
|
||||
|
||||
|
||||
def _redact_tool_calls_dict(message: dict, redacted_str: str) -> None:
|
||||
"""Redact tool call / function_call arguments in a dict-form message or delta."""
|
||||
tool_calls = message.get("tool_calls")
|
||||
if isinstance(tool_calls, list):
|
||||
for tool_call in tool_calls:
|
||||
if isinstance(tool_call, dict) and isinstance(tool_call.get("function"), dict):
|
||||
tool_call["function"]["arguments"] = redacted_str
|
||||
|
||||
function_call = message.get("function_call")
|
||||
if isinstance(function_call, dict) and "arguments" in function_call:
|
||||
function_call["arguments"] = redacted_str
|
||||
|
||||
|
||||
def _redact_model_response_dict_choices(choices, redacted_str: str):
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
|
|
@ -173,6 +212,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
|
|||
choice["message"]["thinking_blocks"] = None
|
||||
if "audio" in choice["message"]:
|
||||
choice["message"]["audio"] = None
|
||||
_redact_tool_calls_dict(choice["message"], redacted_str)
|
||||
elif "delta" in choice and isinstance(choice["delta"], dict):
|
||||
choice["delta"]["content"] = redacted_str
|
||||
if "reasoning_content" in choice["delta"]:
|
||||
|
|
@ -181,6 +221,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str):
|
|||
choice["delta"]["thinking_blocks"] = None
|
||||
if "audio" in choice["delta"]:
|
||||
choice["delta"]["audio"] = None
|
||||
_redact_tool_calls_dict(choice["delta"], redacted_str)
|
||||
else:
|
||||
_redact_choice_content(choice)
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ secrets from strings without depending on the logging-configuration module.
|
|||
import re
|
||||
from typing import List
|
||||
|
||||
from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH
|
||||
|
||||
_REDACTED = "REDACTED"
|
||||
|
||||
|
||||
|
|
@ -30,7 +32,7 @@ def _build_secret_patterns() -> "re.Pattern[str]":
|
|||
# Basic auth headers
|
||||
r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}",
|
||||
# OpenAI / Anthropic sk- prefixed keys
|
||||
r"sk-[A-Za-z0-9\-_]{20,}",
|
||||
rf"sk-[A-Za-z0-9\-_]{{{MINIMUM_CUSTOM_KEY_LENGTH - len('sk-')},}}",
|
||||
# Generic api_key / api-key / apikey (handles 'key': 'value' dict repr)
|
||||
r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}",
|
||||
# x-api-key / api-key header values (handles 'key': 'value' dict repr)
|
||||
|
|
|
|||
|
|
@ -1403,11 +1403,6 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
assert isinstance(thinking, str)
|
||||
assert isinstance(signature, str)
|
||||
|
||||
if thinking and signature:
|
||||
raise ValueError(
|
||||
"Both `thinking` and `signature` in a single streaming chunk isn't supported."
|
||||
)
|
||||
|
||||
return "thinking", ChatCompletionThinkingBlock(
|
||||
type="thinking", thinking=thinking, signature=signature
|
||||
)
|
||||
|
|
@ -1463,17 +1458,14 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
if choice.delta.reasoning_content is not None:
|
||||
reasoning_content += choice.delta.reasoning_content
|
||||
|
||||
if reasoning_content and reasoning_signature:
|
||||
raise ValueError("Both `reasoning` and `signature` in a single streaming chunk isn't supported.")
|
||||
|
||||
if partial_json is not None:
|
||||
return "input_json_delta", ContentJsonBlockDelta(type="input_json_delta", partial_json=partial_json)
|
||||
elif reasoning_content:
|
||||
return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content)
|
||||
elif reasoning_signature:
|
||||
return "signature_delta", ContentThinkingSignatureBlockDelta(
|
||||
type="signature_delta", signature=reasoning_signature
|
||||
)
|
||||
elif reasoning_content:
|
||||
return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content)
|
||||
else:
|
||||
return "text_delta", ContentTextBlockDelta(type="text_delta", text=text)
|
||||
|
||||
|
|
|
|||
|
|
@ -1879,6 +1879,7 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params: GenericLiteLLMParams,
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> httpx.Response:
|
||||
max_attempts = max(provider_config.max_retry_on_anthropic_messages_http_error, 1)
|
||||
litellm_params_dict = dict(litellm_params)
|
||||
|
|
@ -1891,6 +1892,7 @@ class BaseLLMHTTPHandler:
|
|||
data=signed_json_body or json.dumps(request_body),
|
||||
stream=stream or False,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
|
|
@ -1925,6 +1927,32 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
raise RuntimeError("unreachable: anthropic messages HTTP retry loop exited without return")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_anthropic_messages_timeout(
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
) -> Optional[Union[float, httpx.Timeout]]:
|
||||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
)
|
||||
from litellm.utils import supports_httpx_timeout
|
||||
|
||||
stream_timeout = litellm_params.get("stream_timeout") if stream else None
|
||||
model_timeout = stream_timeout if stream_timeout is not None else litellm_params.get("timeout")
|
||||
request_timeout = litellm_params.get("request_timeout")
|
||||
global_timeout = get_configured_request_timeout()
|
||||
if model_timeout is None and request_timeout is None and global_timeout is None:
|
||||
return None
|
||||
return CompletionTimeout.resolve(
|
||||
model_timeout,
|
||||
{"request_timeout": request_timeout},
|
||||
custom_llm_provider,
|
||||
global_timeout=global_timeout,
|
||||
supports_httpx_timeout=supports_httpx_timeout,
|
||||
)
|
||||
|
||||
async def async_anthropic_messages_handler(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -2075,6 +2103,11 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
timeout=self._resolve_anthropic_messages_timeout(
|
||||
litellm_params=litellm_params,
|
||||
stream=stream or False,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
),
|
||||
)
|
||||
|
||||
# used for logging + cost tracking
|
||||
|
|
|
|||
|
|
@ -247,6 +247,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
"""
|
||||
Merge remapped guardrailed tools with original tools that were not sent
|
||||
to the guardrail (e.g. web_search, web_search_preview), preserving order.
|
||||
Tools a guardrail appended (``remapped`` longer than ``original_tools``)
|
||||
have no original slot and are kept so an injected tool is not dropped.
|
||||
"""
|
||||
if not original_tools:
|
||||
return remapped
|
||||
|
|
@ -262,6 +264,8 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if j < len(remapped):
|
||||
result.append(remapped[j])
|
||||
j += 1
|
||||
# Keep guardrail-appended tools that matched no original slot above.
|
||||
result.extend(remapped[j:])
|
||||
return result
|
||||
|
||||
def _apply_guardrailed_tools_to_data(
|
||||
|
|
|
|||
|
|
@ -44261,6 +44261,90 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-sol": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"use_openai_responses_path": true,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-terra": {
|
||||
"input_cost_per_token": 2.75e-06,
|
||||
"cache_creation_input_token_cost": 3.4375e-06,
|
||||
"cache_read_input_token_cost": 2.75e-07,
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"use_openai_responses_path": true,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-luna": {
|
||||
"input_cost_per_token": 1.1e-06,
|
||||
"cache_creation_input_token_cost": 1.375e-06,
|
||||
"cache_read_input_token_cost": 1.1e-07,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"use_openai_responses_path": true,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.5": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
|
|
@ -44373,6 +44457,7 @@
|
|||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/xai.grok-4.3": {
|
||||
"use_openai_responses_path": true,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
|
|
|
|||
|
|
@ -186,6 +186,61 @@ _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES: tuple[MCPAuth, ...] = (
|
|||
)
|
||||
|
||||
|
||||
def _blank_to_none(value: str | None) -> str | None:
|
||||
"""Collapse an absent, empty, or whitespace-only string to ``None``.
|
||||
|
||||
OAuth endpoint fields are consumed by truthiness-based merges (``row or discovered``) and by the
|
||||
corroboration gate. A whitespace-only value is truthy to ``or`` but is not a usable endpoint, so
|
||||
without this the merge would keep the blank value for redirects while the gate treats it as
|
||||
unpinned and backfills the other fields, yielding a broken half-discovered config. Normalizing
|
||||
the pinned fields once, at each build entry point, gives every downstream consumer a single
|
||||
notion of "blank" so those code paths cannot disagree.
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
return value.strip() or None
|
||||
|
||||
|
||||
def _normalized_authorize_endpoint(url: str) -> str:
|
||||
"""Compare authorize endpoints on scheme, host, and path only. The default port is elided and
|
||||
the host is lowercased so ``https://IDP.example.com:443/authorize/`` and
|
||||
``https://idp.example.com/authorize`` are the same identity; query and trailing slash are not."""
|
||||
parsed = urlparse(url)
|
||||
scheme = parsed.scheme.lower()
|
||||
host = (parsed.hostname or "").lower()
|
||||
default_port = {"https": 443, "http": 80}.get(scheme)
|
||||
try:
|
||||
port = parsed.port
|
||||
except ValueError:
|
||||
port = None
|
||||
authority = host if port is None or port == default_port else f"{host}:{port}"
|
||||
return f"{scheme}://{authority}{parsed.path.rstrip('/')}"
|
||||
|
||||
|
||||
def _endpoints_corroborate_authorization_url(
|
||||
source_authorization_url: str | None,
|
||||
trusted_authorization_url: str | None,
|
||||
) -> bool:
|
||||
"""Whether a source's ``token_url``/``registration_url`` may be paired with a trusted authorize
|
||||
endpoint. This is the single trust rule for adopting OAuth endpoints from any non-manual source.
|
||||
|
||||
Discovery is rooted at the MCP resource (RFC 9728), so a compromised upstream can advertise an
|
||||
attacker-run authorization server. When ``authorization_url`` is admin-pinned, pairing it with a
|
||||
``token_url`` from a different source is the RFC 9700 authorization-server mix-up: the user signs
|
||||
in at the trusted authorize endpoint while the gateway redeems the code, with the stored client
|
||||
secret and PKCE verifier, at the attacker's token endpoint. Endpoints are trustworthy together
|
||||
only when they share an authorization server, so a source's endpoints are adopted only when the
|
||||
same source advertised an ``authorization_endpoint`` matching the pinned value. With no pinned
|
||||
value (``trusted_authorization_url is None``) there is nothing to protect: the authorize endpoint
|
||||
comes from the same source as the token endpoint, so they corroborate each other by construction.
|
||||
"""
|
||||
if not (trusted_authorization_url and trusted_authorization_url.strip()):
|
||||
return True
|
||||
return bool(source_authorization_url) and _normalized_authorize_endpoint(
|
||||
source_authorization_url
|
||||
) == _normalized_authorize_endpoint(trusted_authorization_url)
|
||||
|
||||
|
||||
def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_server: MCPServer | None) -> None:
|
||||
"""Keep the last known good OAuth endpoints when a rebuild's re-discovery comes back empty.
|
||||
|
||||
|
|
@ -193,26 +248,82 @@ def _carry_forward_resolved_oauth_endpoints(new_server: MCPServer, previous_serv
|
|||
during re-discovery downgrades a working server (``authorization_url`` set) to a broken one
|
||||
(``None``, /authorize 400s) with no configuration change. Mirrors the ``short_prefix``
|
||||
carry-forward. Skipped when the server's ``url`` or ``auth_type`` changed, since the previous
|
||||
endpoints may then belong to a different upstream. ``registration_url`` IS carried here even
|
||||
though ``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only
|
||||
restores the same in-memory value the previous build already ran with, while persisting it
|
||||
would flip ``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for
|
||||
dcr_bridge servers that never had one configured.
|
||||
endpoints may then belong to a different upstream. ``registration_url`` IS carried even though
|
||||
``_persist_discovered_oauth_endpoints`` refuses to write it to the row: carrying only restores
|
||||
the same in-memory value the previous build already ran with, while persisting it would flip
|
||||
``_dcr_bridge_relays_client_registration`` (which keys off the stored column) for dcr_bridge
|
||||
servers that never had one configured.
|
||||
|
||||
Carry-forward is a non-manual endpoint source, so the same trust rule as discovery applies: the
|
||||
previous ``token_url``/``registration_url``/``scopes`` are carried only when the previous
|
||||
``authorization_url`` corroborates the authorize endpoint this build will use, i.e. when the
|
||||
incoming build has no pinned authorize endpoint (``None`` -> we adopt the previous one too, a
|
||||
consistent group) or pins the same one. An admin re-pointing ``authorization_url`` to a different
|
||||
server must not keep serving the old server's token endpoint or granted scopes.
|
||||
"""
|
||||
if previous_server is None:
|
||||
return
|
||||
if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type:
|
||||
return
|
||||
may_carry = _endpoints_corroborate_authorization_url(
|
||||
previous_server.authorization_url, new_server.authorization_url
|
||||
)
|
||||
if new_server.authorization_url is None and previous_server.authorization_url:
|
||||
new_server.authorization_url = previous_server.authorization_url
|
||||
if new_server.token_url is None and previous_server.token_url:
|
||||
if may_carry and new_server.token_url is None and previous_server.token_url:
|
||||
new_server.token_url = previous_server.token_url
|
||||
if new_server.registration_url is None and previous_server.registration_url:
|
||||
if may_carry and new_server.registration_url is None and previous_server.registration_url:
|
||||
new_server.registration_url = previous_server.registration_url
|
||||
if not new_server.scopes and previous_server.scopes:
|
||||
if may_carry and not new_server.scopes and previous_server.scopes:
|
||||
new_server.scopes = previous_server.scopes
|
||||
|
||||
|
||||
def _restrict_discovery_to_corroborated_authorization_server(
|
||||
metadata: MCPOAuthMetadata | None,
|
||||
manual_authorization_url: str | None,
|
||||
server_identifier: str,
|
||||
is_dcr_bridge: bool,
|
||||
) -> MCPOAuthMetadata | None:
|
||||
"""Reject discovered token/registration endpoints a manually pinned authorize endpoint cannot
|
||||
vouch for (the RFC 9700 authorization-server mix-up).
|
||||
|
||||
Discovery is rooted at the MCP resource, so a compromised upstream can advertise an attacker
|
||||
``token_endpoint``: with ``authorization_url`` admin-pinned but ``token_url`` blank, the merge
|
||||
would pair the trusted authorize endpoint with that attacker token endpoint, and the gateway would
|
||||
post the authorization code and client secret there. So the discovered ``token_url`` and
|
||||
``registration_url`` are kept only if the document corroborates the pin (its
|
||||
``authorization_endpoint`` matches). ``scopes`` are deliberately NOT gated here: per the MCP
|
||||
authorization spec Scope Selection Strategy and RFC 9700 §2.3, the scopes a client requests are
|
||||
resource-driven (the WWW-Authenticate challenge or the RFC 9728 protected-resource
|
||||
``scopes_supported``), and scope inflation by a compromised resource is bounded by the
|
||||
authorization server and user consent (RFC 6749 §3.3), not by the client second-guessing the
|
||||
request. With no pin there is no trust anchor to protect, so discovery is returned as-is.
|
||||
"""
|
||||
if metadata is None or not (manual_authorization_url and manual_authorization_url.strip()):
|
||||
return metadata
|
||||
if _endpoints_corroborate_authorization_url(metadata.authorization_url, manual_authorization_url):
|
||||
return metadata
|
||||
if not metadata.token_url and not metadata.registration_url:
|
||||
return metadata
|
||||
bridge_note = (
|
||||
" The discovered registration_url is rejected with it, so this dcr_bridge server stays on the"
|
||||
" short-circuit registration arm."
|
||||
if is_dcr_bridge and metadata.registration_url
|
||||
else ""
|
||||
)
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery for server %s advertised authorization_endpoint %s, which does not match the "
|
||||
"manually configured authorization_url %s; rejecting the discovered token_url/registration_url so "
|
||||
"authorization codes and client credentials only follow the configured authorization server. "
|
||||
"Configure Token URL manually if the mismatch is intentional.%s",
|
||||
server_identifier,
|
||||
_normalized_authorize_endpoint(metadata.authorization_url) if metadata.authorization_url else "<absent>",
|
||||
_normalized_authorize_endpoint(manual_authorization_url),
|
||||
bridge_note,
|
||||
)
|
||||
return metadata.model_copy(update={"token_url": None, "registration_url": None})
|
||||
|
||||
|
||||
def invalidate_user_env_vars_cache(user_id: str, server_id: str) -> None:
|
||||
"""Drop a cached entry after the user stores or clears their env var values
|
||||
so the next request reads the fresh value instead of a stale one."""
|
||||
|
|
@ -1026,12 +1137,15 @@ class MCPServerManager:
|
|||
)
|
||||
|
||||
auth_type = server_config.get("auth_type", None)
|
||||
manual_authorization_url = _blank_to_none(server_config.get("authorization_url"))
|
||||
manual_token_url = _blank_to_none(server_config.get("token_url"))
|
||||
manual_registration_url = _blank_to_none(server_config.get("registration_url"))
|
||||
if server_url and (
|
||||
auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
|
||||
or self._obo_needs_endpoint_discovery(
|
||||
auth_type,
|
||||
server_config.get("token_exchange_endpoint"),
|
||||
server_config.get("token_url"),
|
||||
manual_token_url,
|
||||
)
|
||||
):
|
||||
mcp_oauth_metadata = await self._descovery_metadata(
|
||||
|
|
@ -1041,20 +1155,29 @@ class MCPServerManager:
|
|||
else:
|
||||
mcp_oauth_metadata = None
|
||||
|
||||
gated_oauth_metadata = (
|
||||
_restrict_discovery_to_corroborated_authorization_server(
|
||||
mcp_oauth_metadata,
|
||||
manual_authorization_url,
|
||||
server_name or server_id,
|
||||
bool(server_config.get("dcr_bridge")),
|
||||
)
|
||||
if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
|
||||
else mcp_oauth_metadata
|
||||
)
|
||||
|
||||
# Filter blank scopes (e.g. YAML ``scopes: [""]``) the same way the DB-build path does, so
|
||||
# an all-blank list normalizes to None rather than a ``("",)`` tuple that skips the
|
||||
# entra_obo fail-closed scope precondition and POSTs an empty scope to the IdP.
|
||||
resolved_scopes = self._extract_scopes(server_config.get("scopes")) or (
|
||||
mcp_oauth_metadata.scopes if mcp_oauth_metadata else None
|
||||
gated_oauth_metadata.scopes if gated_oauth_metadata else None
|
||||
)
|
||||
resolved_authorization_url = server_config.get("authorization_url") or (
|
||||
mcp_oauth_metadata.authorization_url if mcp_oauth_metadata else None
|
||||
resolved_authorization_url = manual_authorization_url or (
|
||||
gated_oauth_metadata.authorization_url if gated_oauth_metadata else None
|
||||
)
|
||||
resolved_token_url = server_config.get("token_url") or (
|
||||
mcp_oauth_metadata.token_url if mcp_oauth_metadata else None
|
||||
)
|
||||
resolved_registration_url = server_config.get("registration_url") or (
|
||||
mcp_oauth_metadata.registration_url if mcp_oauth_metadata else None
|
||||
resolved_token_url = manual_token_url or (gated_oauth_metadata.token_url if gated_oauth_metadata else None)
|
||||
resolved_registration_url = manual_registration_url or (
|
||||
gated_oauth_metadata.registration_url if gated_oauth_metadata else None
|
||||
)
|
||||
|
||||
config_oauth2_flow = server_config.get("oauth2_flow", None)
|
||||
|
|
@ -1447,13 +1570,17 @@ class MCPServerManager:
|
|||
|
||||
auth_type = cast(MCPAuthType, mcp_server.auth_type)
|
||||
server_url = mcp_server.url
|
||||
manual_authorization_url = _blank_to_none(mcp_server.authorization_url)
|
||||
manual_token_url = _blank_to_none(mcp_server.token_url)
|
||||
manual_registration_url = _blank_to_none(mcp_server.registration_url)
|
||||
has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes)
|
||||
needs_discovery = bool(server_url) and (
|
||||
(auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not mcp_server.authorization_url)
|
||||
(auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not has_all_upstream_oauth_fields)
|
||||
or self._obo_needs_endpoint_discovery(
|
||||
auth_type,
|
||||
mcp_server.token_exchange_endpoint
|
||||
or (credentials_dict.get("token_exchange_endpoint") if credentials_dict else None),
|
||||
mcp_server.token_url,
|
||||
manual_token_url,
|
||||
)
|
||||
)
|
||||
mcp_oauth_metadata = (
|
||||
|
|
@ -1467,12 +1594,22 @@ class MCPServerManager:
|
|||
if needs_discovery and mcp_oauth_metadata is None:
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery yielded no metadata for server %s (%s); "
|
||||
"OAuth endpoints stay unresolved until a rebuild succeeds",
|
||||
"OAuth endpoints/scopes stay unresolved until a rebuild succeeds",
|
||||
mcp_server.server_id,
|
||||
server_url,
|
||||
)
|
||||
gated_oauth_metadata = (
|
||||
_restrict_discovery_to_corroborated_authorization_server(
|
||||
mcp_oauth_metadata,
|
||||
manual_authorization_url,
|
||||
mcp_server.server_id,
|
||||
bool(getattr(mcp_server, "dcr_bridge", None)),
|
||||
)
|
||||
if auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
|
||||
else mcp_oauth_metadata
|
||||
)
|
||||
|
||||
resolved_scopes = scopes or (mcp_oauth_metadata.scopes if mcp_oauth_metadata else None)
|
||||
resolved_scopes = scopes or (gated_oauth_metadata.scopes if gated_oauth_metadata else None)
|
||||
|
||||
new_server = MCPServer(
|
||||
server_id=mcp_server.server_id,
|
||||
|
|
@ -1492,9 +1629,9 @@ class MCPServerManager:
|
|||
client_secret=client_secret_value or getattr(mcp_server, "client_secret", None),
|
||||
oauth2_flow=self._explicit_oauth2_flow(getattr(mcp_server, "oauth2_flow", None)),
|
||||
scopes=resolved_scopes,
|
||||
authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None),
|
||||
token_url=mcp_server.token_url or getattr(mcp_oauth_metadata, "token_url", None),
|
||||
registration_url=mcp_server.registration_url or getattr(mcp_oauth_metadata, "registration_url", None),
|
||||
authorization_url=manual_authorization_url or getattr(gated_oauth_metadata, "authorization_url", None),
|
||||
token_url=manual_token_url or getattr(gated_oauth_metadata, "token_url", None),
|
||||
registration_url=manual_registration_url or getattr(gated_oauth_metadata, "registration_url", None),
|
||||
token_endpoint_auth_method=(
|
||||
credentials_dict.get("token_endpoint_auth_method") if credentials_dict else None
|
||||
),
|
||||
|
|
@ -1545,16 +1682,16 @@ class MCPServerManager:
|
|||
await self._persist_discovered_obo_token_url(
|
||||
server_id=mcp_server.server_id,
|
||||
auth_type=auth_type,
|
||||
existing_token_url=mcp_server.token_url,
|
||||
existing_token_url=manual_token_url,
|
||||
discovered_token_url=new_server.token_url,
|
||||
)
|
||||
await self._persist_discovered_oauth_endpoints(
|
||||
server_id=mcp_server.server_id,
|
||||
auth_type=auth_type,
|
||||
existing_authorization_url=mcp_server.authorization_url,
|
||||
existing_token_url=mcp_server.token_url,
|
||||
existing_authorization_url=manual_authorization_url,
|
||||
existing_token_url=manual_token_url,
|
||||
existing_scopes=scopes,
|
||||
metadata=mcp_oauth_metadata,
|
||||
metadata=gated_oauth_metadata,
|
||||
)
|
||||
return new_server
|
||||
|
||||
|
|
|
|||
|
|
@ -175,16 +175,18 @@ mcp_oauth2_token_cache = MCPOAuth2TokenCache()
|
|||
def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int:
|
||||
"""Compute Redis TTL for a per-user token.
|
||||
|
||||
Uses server.token_storage_ttl_seconds when configured; otherwise derives
|
||||
TTL from expires_in minus the expiry buffer; falls back to the default TTL.
|
||||
Uses server.token_storage_ttl_seconds when configured, capped at the token's
|
||||
remaining lifetime (expires_in minus the expiry buffer) so a cached entry never
|
||||
outlives the token itself; otherwise derives TTL from expires_in minus the
|
||||
expiry buffer; falls back to the default TTL.
|
||||
"""
|
||||
lifetime_bound = expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS if expires_in is not None else None
|
||||
if server.token_storage_ttl_seconds is not None:
|
||||
return max(server.token_storage_ttl_seconds, 1)
|
||||
if expires_in is not None:
|
||||
return max(
|
||||
expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS,
|
||||
1,
|
||||
)
|
||||
if lifetime_bound is None:
|
||||
return max(server.token_storage_ttl_seconds, 1)
|
||||
return max(min(server.token_storage_ttl_seconds, lifetime_bound), 1)
|
||||
if lifetime_bound is not None:
|
||||
return max(lifetime_bound, 1)
|
||||
return MCP_PER_USER_TOKEN_DEFAULT_TTL
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -7613,6 +7613,18 @@
|
|||
],
|
||||
"title": "Messages"
|
||||
},
|
||||
"metadata": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": true,
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Metadata"
|
||||
},
|
||||
"text": {
|
||||
"title": "Text",
|
||||
"type": "string"
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from fastapi import HTTPException, Request, status
|
|||
import litellm
|
||||
from litellm import Router, provider_list
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
|
||||
from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH, STANDARD_CUSTOMER_ID_HEADERS
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
|
||||
from litellm.proxy._types import *
|
||||
|
|
@ -1533,4 +1533,6 @@ def get_model_from_request(
|
|||
|
||||
|
||||
def abbreviate_api_key(api_key: str) -> str:
|
||||
if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH:
|
||||
return "sk-..."
|
||||
return f"sk-...{api_key[-4:]}"
|
||||
|
|
|
|||
|
|
@ -471,6 +471,109 @@ The token minted by `lite login` is a short-lived, per-session agent credential,
|
|||
|
||||
The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
|
||||
|
||||
### Route Every Claude Code Session Through the Proxy
|
||||
|
||||
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
|
||||
|
||||
Two things need to already be true: you've run `lite login`, since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
|
||||
|
||||
```bash
|
||||
lite login
|
||||
litellm --config litellm/proxy/dev_config.yaml &
|
||||
lite up
|
||||
```
|
||||
|
||||
`lite up` runs in the foreground and blocks. Press Ctrl-C to stop it, which restores the original settings file and exits. If the process is ever killed uncleanly instead -- `kill -9`, a crash -- the settings file is left patched, and `lite down` is the manual recovery path: run it at any later point to restore from the same backup.
|
||||
|
||||
This is a one-time file patch and restore, not a live traffic interceptor. A Claude Code session already running before `lite up` started keeps whatever `ANTHROPIC_BASE_URL` and token it loaded at its own startup, and a session still running when `lite up` stops keeps routing through the proxy until it exits; only sessions *started* while the patch is in effect are affected, and only *new* sessions after a restore go back to Anthropic directly.
|
||||
|
||||
Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI.
|
||||
|
||||
### QA Complexity-Based Auto-Routing Against Your Real Proxy
|
||||
|
||||
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
|
||||
|
||||
#### Install the CLI
|
||||
|
||||
`lite autoroute up` builds and runs a throwaway litellm proxy locally, so unlike the rest of this CLI it needs the proxy server runtime, not just the thin `litellm[cli]` client. Install `litellm[proxy]` (which ships the `lite` command too) with a single curl command -- no existing Python tooling required, `uv` is bootstrapped automatically if missing:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh
|
||||
```
|
||||
|
||||
To QA an unreleased branch or commit instead of the latest PyPI release, set `LITELLM_CLI_REF`:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/<branch-or-commit>/scripts/install.sh | \
|
||||
LITELLM_CLI_REF=<branch-or-commit> sh
|
||||
```
|
||||
|
||||
The thin `scripts/install-cli.sh` installs only `litellm[cli]`, which is enough for `lite login`, `lite claude`, and `lite up`, but not for `lite autoroute up`; running it against a `litellm[cli]` install fails fast with a message telling you to install the proxy runtime.
|
||||
|
||||
Point the CLI at your real proxy and key before running any `lite model-groups` or `lite autoroute` command -- like every other command in this CLI, they read `LITELLM_PROXY_URL`/`LITELLM_PROXY_API_KEY` (or `--base-url`/`--api-key`), no `lite login` required:
|
||||
|
||||
```bash
|
||||
export LITELLM_PROXY_URL=http://localhost:4000
|
||||
export LITELLM_PROXY_API_KEY=sk-...
|
||||
```
|
||||
|
||||
#### List Your Accessible Model Groups
|
||||
|
||||
```bash
|
||||
lite model-groups list [--format table|json]
|
||||
```
|
||||
|
||||
Lists the model groups your key can reach on the proxy, via `/model_group/info`, along with each group's mode (`chat`, `embedding`, etc.) and per-token pricing. This is also what `lite autoroute configure` uses internally to discover what it can offer you.
|
||||
|
||||
#### Configure the Auto-Router
|
||||
|
||||
```bash
|
||||
lite autoroute configure
|
||||
```
|
||||
|
||||
An interactive wizard. It runs the same model-group discovery as above, splits the results into chat-capable and embedding-capable pools, and asks you to assign one or more models from the chat pool to each of the four complexity tiers -- SIMPLE, MEDIUM, COMPLEX, REASONING. Each tier's picker is a type-to-filter fuzzy search (fzf-style) rather than a scrollable numbered list, so it stays usable even with hundreds of model groups: type a substring to narrow the list, tab to toggle a model into the selection, enter to confirm (assigning more than one model to a tier is exactly when this matters -- complexity_router picks randomly among a tier's pool per request, and adaptive mode specifically depends on having more than one candidate to choose from). From there it optionally offers: classifying prompt complexity with an LLM (again picked from your discovered pool) instead of the free built-in heuristic scorer, semantic keyword matching for tier assignment (needs an embedding model from the pool), and adaptive (bandit-based) selection layered on top of tiering.
|
||||
|
||||
The wizard writes the result to `~/.litellm/autorouter/config.yaml` with `0600` permissions, since the file embeds your real proxy API key. Every model referenced anywhere in that config -- tier targets, the classifier model, the embedding model -- becomes its own `litellm_proxy/<model-name>` deployment whose `api_base` and `api_key` point back at your real proxy. That is the trick that keeps your real proxy's config untouched: every actual network call this generates, whether it is the routed completion, an LLM-classifier call, or an embedding call, forwards transparently through your real, already-running proxy with your real key.
|
||||
|
||||
You do not need to tell Claude Code to request `autorouter` by name yourself: `lite autoroute up` also sets `ANTHROPIC_DEFAULT_SONNET_MODEL`, `ANTHROPIC_DEFAULT_HAIKU_MODEL`, and `ANTHROPIC_DEFAULT_OPUS_MODEL` to `autorouter` in `~/.claude/settings.json`, so every one of Claude Code's own model tiers requests it directly regardless of `/model` or whatever it defaults to otherwise. (A bare `model_name: "*"` deployment looks like the obvious way to catch any request instead, but litellm's Router looks up auto-router deployments by the literal requested model string with no wildcard resolution, so a `"*"` entry would never actually match real traffic -- these env var overrides are what makes it work.)
|
||||
|
||||
You must run `configure` at least once before `up`; running `up` first fails with a clear error telling you to configure first.
|
||||
|
||||
#### Launch the Ephemeral Auto-Router Proxy
|
||||
|
||||
```bash
|
||||
lite autoroute up
|
||||
```
|
||||
|
||||
Starts a local, throwaway litellm proxy on a random free port, running the config `configure` generated, with a freshly-minted random API key baked in for this session only (your real proxy key never leaves the generated config -- it only appears there, forwarding to your real proxy). It waits for the ephemeral proxy to report healthy, then patches `~/.claude/settings.json` the same way `lite up` does, except with a static `ANTHROPIC_AUTH_TOKEN` env var instead of an `apiKeyHelper`, since this key is short-lived and self-issued rather than something needing SSO refresh. Any `claude` session started afterward, from any terminal, routes through the ephemeral proxy.
|
||||
|
||||
`lite autoroute up` runs in the foreground and streams the ephemeral proxy's own log file into your terminal, so you can watch its routing decisions -- which tier and model got picked for each request -- as you use Claude Code normally. Press Ctrl-C (or send SIGTERM) to stop it; this kills the child proxy process and restores your original Claude Code settings, in that order.
|
||||
|
||||
#### Recover From an Unclean Shutdown
|
||||
|
||||
```bash
|
||||
lite autoroute down
|
||||
```
|
||||
|
||||
If the `lite autoroute up` process dies uncleanly -- `kill -9`, a crash -- rather than being stopped with Ctrl-C, `down` is the manual recovery path: it kills any leftover ephemeral proxy process found via a recorded pid file and restores Claude Code's settings from whatever backup is on disk.
|
||||
|
||||
#### Example
|
||||
|
||||
```bash
|
||||
lite autoroute configure
|
||||
lite autoroute up
|
||||
# use Claude Code as normal in another terminal; routing decisions stream live
|
||||
lite autoroute down # only needed if `up` was killed uncleanly instead of Ctrl-C'd
|
||||
```
|
||||
|
||||
#### Caveats
|
||||
|
||||
Adaptive mode's learned state does not persist across `lite autoroute up` sessions -- there is no local database, so every session starts adaptive selection cold. A Claude Code session already running before `up` started, or still running when it stops, keeps whatever settings it loaded at its own startup; like `lite up`, this is a one-time file patch and restore, not a live traffic interceptor. Only Claude Code is supported, for the same reason as `lite up`: no other supported agent (for example Cursor) has an equivalent hot-patchable config file.
|
||||
|
||||
A session that outlives `up` (or is still running the moment you stop it) keeps sending requests, master key included, to that now-freed loopback port until you restart it. Once the ephemeral proxy process exits, nothing stops another local account on the same machine from binding that same port and receiving those requests instead -- unlike `lite up`'s `apiKeyHelper`, which is re-resolved per request, `autoroute`'s master key is a static value, so whoever receives them gets a live-looking token along with the prompt content. Restart any Claude Code session before you consider the machine clean, run `lite autoroute down` promptly rather than leaving a stopped session's settings patched, and do not run `lite autoroute up` on a shared or multi-tenant host.
|
||||
|
||||
Do not run `lite up` and `lite autoroute up` at the same time. Each patches `~/.claude/settings.json` and keeps its own separate backup, with no coordination between them: whichever one you stop or crash out of last is the one whose backup gets restored, which can silently leave the *other* mode's settings (a static master key and a now-dead loopback URL, or a stale `apiKeyHelper`) active. Run `lite down` or `lite autoroute down` (whichever applies) before switching to the other mode.
|
||||
|
||||
## Environment Variables
|
||||
|
||||
The CLI respects the following environment variables:
|
||||
|
|
|
|||
|
|
@ -212,7 +212,7 @@ def _is_interactive() -> bool:
|
|||
return sys.stdin.isatty()
|
||||
|
||||
|
||||
def _resolve_api_key(ctx: click.Context) -> str:
|
||||
def resolve_api_key(ctx: click.Context) -> str:
|
||||
base_url = ctx.obj["base_url"]
|
||||
api_key = ctx.obj.get("api_key")
|
||||
if api_key:
|
||||
|
|
@ -238,7 +238,7 @@ _SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy."
|
|||
def _launch(ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool) -> None:
|
||||
base_url = ctx.obj["base_url"]
|
||||
started_interactive = _is_interactive()
|
||||
api_key = _resolve_api_key(ctx)
|
||||
api_key = resolve_api_key(ctx)
|
||||
|
||||
display_name, _ = agent_profile(binary)
|
||||
click.echo(f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}")
|
||||
|
|
@ -288,5 +288,6 @@ __all__ = [
|
|||
"agent_launch_args",
|
||||
"verify_proxy_key",
|
||||
"agent_profile",
|
||||
"resolve_api_key",
|
||||
"AgentRunError",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -277,6 +277,28 @@ def prompt_team_selection_fallback(
|
|||
return None
|
||||
|
||||
|
||||
def _response_error_detail(response: requests.Response) -> str | None:
|
||||
try:
|
||||
body = response.json()
|
||||
except ValueError:
|
||||
return None
|
||||
detail = body.get("detail") if isinstance(body, dict) else None
|
||||
if isinstance(detail, str) and detail:
|
||||
return detail
|
||||
return None
|
||||
|
||||
|
||||
def _polling_error_message(response: requests.Response) -> str:
|
||||
detail = _response_error_detail(response)
|
||||
if detail:
|
||||
return f"Polling error: HTTP {response.status_code}: {detail}"
|
||||
return f"Polling error: HTTP {response.status_code}"
|
||||
|
||||
|
||||
def _is_permanent_polling_error(status_code: int) -> bool:
|
||||
return 400 <= status_code < 500 and status_code != 429
|
||||
|
||||
|
||||
# Polling-based authentication - no local server needed
|
||||
def _poll_for_ready_data(
|
||||
url: str,
|
||||
|
|
@ -308,8 +330,14 @@ def _poll_for_ready_data(
|
|||
click.echo(pending_message)
|
||||
elif other_status_message and other_status_log_every > 0 and attempt % other_status_log_every == 0:
|
||||
click.echo(other_status_message)
|
||||
elif _is_permanent_polling_error(response.status_code):
|
||||
detail = _response_error_detail(response)
|
||||
raise ValueError(
|
||||
f"The proxy rejected the login session with HTTP {response.status_code}"
|
||||
+ (f": {detail}" if detail else f" and no error detail (from {url})")
|
||||
)
|
||||
elif http_error_log_every > 0 and attempt % http_error_log_every == 0:
|
||||
click.echo(f"Polling error: HTTP {response.status_code}")
|
||||
click.echo(_polling_error_message(response))
|
||||
except requests.RequestException as e:
|
||||
if connection_error_log_every > 0 and attempt % connection_error_log_every == 0:
|
||||
click.echo(f"Connection error (will retry): {e}")
|
||||
|
|
@ -342,12 +370,45 @@ def _normalize_teams(teams, team_details):
|
|||
|
||||
|
||||
def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]:
|
||||
response = requests.post(f"{base_url}/sso/cli/start", timeout=10)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
start_url = f"{base_url}/sso/cli/start"
|
||||
try:
|
||||
response = requests.post(start_url, timeout=10)
|
||||
except requests.RequestException as e:
|
||||
raise ValueError(
|
||||
f"Could not reach the proxy at {start_url}: {e}. "
|
||||
"Check that the proxy is running and that --base-url points at it."
|
||||
) from e
|
||||
|
||||
if response.status_code in (404, 405):
|
||||
raise ValueError(
|
||||
f"POST {start_url} returned HTTP {response.status_code}. "
|
||||
"Either --base-url is wrong, or the proxy is older than this CLI and does not support "
|
||||
"the CLI SSO login flow; upgrade the proxy or use a CLI version that matches it."
|
||||
)
|
||||
if response.status_code != 200:
|
||||
detail = _response_error_detail(response)
|
||||
raise ValueError(
|
||||
f"Starting CLI login failed: HTTP {response.status_code} from {start_url}"
|
||||
+ (f": {detail}" if detail else "")
|
||||
)
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
except ValueError:
|
||||
content_type = response.headers.get("content-type", "unknown")
|
||||
raise ValueError(
|
||||
f"The proxy returned a non-JSON response from {start_url} (content-type: {content_type}). "
|
||||
"A proxy, load balancer, or auth gateway in front of LiteLLM may be intercepting the request. "
|
||||
f"Response starts with: {response.text[:200]!r}"
|
||||
)
|
||||
|
||||
required_fields = ("login_id", "poll_secret", "user_code")
|
||||
if not all(isinstance(data.get(field), str) for field in required_fields):
|
||||
raise ValueError("Invalid CLI SSO start response")
|
||||
missing_fields = tuple(field for field in required_fields if not isinstance(data.get(field), str))
|
||||
if missing_fields:
|
||||
raise ValueError(
|
||||
f"The response from {start_url} is missing required field(s): {', '.join(missing_fields)}. "
|
||||
"The proxy version may not match this CLI; upgrade whichever is older."
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
|
|
@ -577,6 +638,10 @@ def login(ctx: click.Context):
|
|||
return
|
||||
else:
|
||||
click.echo("❌ Authentication timed out. Please try again.")
|
||||
click.echo(
|
||||
"The proxy never reported the browser sign-in as finished. If you did complete it, "
|
||||
"check the proxy logs for /sso/callback errors and confirm SSO is configured on the proxy."
|
||||
)
|
||||
return
|
||||
|
||||
except KeyboardInterrupt:
|
||||
|
|
|
|||
0
litellm/proxy/client/cli/commands/autoroute/__init__.py
Normal file
0
litellm/proxy/client/cli/commands/autoroute/__init__.py
Normal file
207
litellm/proxy/client/cli/commands/autoroute/commands.py
Normal file
207
litellm/proxy/client/cli/commands/autoroute/commands.py
Normal file
|
|
@ -0,0 +1,207 @@
|
|||
import atexit
|
||||
import json
|
||||
import secrets
|
||||
import signal
|
||||
import threading
|
||||
from types import FrameType
|
||||
|
||||
import click
|
||||
import yaml
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup
|
||||
from ..up import BackupRecord as ClaudeBackupRecord
|
||||
from .process import (
|
||||
AUTOROUTE_DIR,
|
||||
CONFIG_PATH,
|
||||
LOG_PATH,
|
||||
PidRecord,
|
||||
ProcessLaunchError,
|
||||
allocate_free_port,
|
||||
clear_pid_record,
|
||||
is_running,
|
||||
launch_proxy,
|
||||
missing_proxy_runtime_modules,
|
||||
poll_liveliness,
|
||||
read_pid_record,
|
||||
secure_create,
|
||||
stream_log,
|
||||
terminate,
|
||||
write_pid_record,
|
||||
)
|
||||
from .settings import merge_claude_settings_static_token
|
||||
from .wizard import run_configure_wizard
|
||||
|
||||
AUTOROUTE_BACKUP_PATH = AUTOROUTE_DIR / "claude_settings_backup.json"
|
||||
|
||||
_GENERATED_CONFIG_ADAPTER = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def _mint_and_embed_master_key() -> str:
|
||||
"""Generate a fresh key for this session and write it into the generated config.yaml.
|
||||
|
||||
Must go under general_settings, not litellm_settings -- the proxy server only ever
|
||||
reads general_settings.master_key (proxy_server.py:4530) to authenticate requests. A
|
||||
key placed under litellm_settings is silently ignored, leaving the ephemeral proxy with
|
||||
no real auth: any request reaches it regardless of the token Claude Code sends.
|
||||
"""
|
||||
master_key = secrets.token_urlsafe(32)
|
||||
with open(CONFIG_PATH, "r") as f:
|
||||
try:
|
||||
generated = _GENERATED_CONFIG_ADAPTER.validate_python(yaml.safe_load(f))
|
||||
except (yaml.YAMLError, ValidationError):
|
||||
raise click.ClickException(
|
||||
f"{CONFIG_PATH} is empty or corrupt. Run `lite autoroute configure` again to regenerate it."
|
||||
)
|
||||
general_settings = generated.get("general_settings")
|
||||
updated_settings: dict[str, JsonValue] = {
|
||||
**(general_settings if isinstance(general_settings, dict) else {}),
|
||||
"master_key": master_key,
|
||||
}
|
||||
updated: dict[str, JsonValue] = {**generated, "general_settings": updated_settings}
|
||||
with secure_create(CONFIG_PATH) as f:
|
||||
yaml.safe_dump(updated, f, sort_keys=False)
|
||||
return master_key
|
||||
|
||||
|
||||
@click.group(name="autoroute")
|
||||
def autoroute_group() -> None:
|
||||
"""QA complexity-based auto-routing against models your key can already use"""
|
||||
|
||||
|
||||
@autoroute_group.command("configure")
|
||||
@click.pass_context
|
||||
def configure(ctx: click.Context) -> None:
|
||||
"""Discover accessible models and generate an ephemeral auto-router config"""
|
||||
run_configure_wizard(ctx)
|
||||
|
||||
|
||||
@autoroute_group.command("up")
|
||||
def up() -> None:
|
||||
"""Launch the ephemeral auto-router proxy and route Claude Code through it"""
|
||||
if not CONFIG_PATH.exists():
|
||||
raise click.ClickException("No config found. Run `lite autoroute configure` first.")
|
||||
|
||||
missing = missing_proxy_runtime_modules()
|
||||
if missing:
|
||||
raise click.ClickException(
|
||||
"lite autoroute up launches a local litellm proxy, which needs the proxy runtime that the "
|
||||
f"thin `litellm[cli]` install does not include (missing: {', '.join(missing)}). Install the "
|
||||
"proxy runtime with `uv tool install --force 'litellm[proxy]'`, or to QA a branch, "
|
||||
"`curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/<branch>/scripts/install.sh | "
|
||||
"LITELLM_CLI_REF=<branch> sh`."
|
||||
)
|
||||
|
||||
try:
|
||||
existing_pid = read_pid_record()
|
||||
except UpError as e:
|
||||
raise click.ClickException(str(e))
|
||||
if existing_pid is not None and is_running(existing_pid.pid):
|
||||
raise click.ClickException(
|
||||
"An ephemeral proxy is already running (lite autoroute up looks already active). "
|
||||
"Run `lite autoroute down` first."
|
||||
)
|
||||
|
||||
if AUTOROUTE_BACKUP_PATH.exists():
|
||||
raise click.ClickException(
|
||||
f"{AUTOROUTE_BACKUP_PATH} already exists -- `lite autoroute up` looks like it's already "
|
||||
"running (or crashed without cleanup). Run `lite autoroute down` first."
|
||||
)
|
||||
|
||||
master_key = _mint_and_embed_master_key()
|
||||
port = allocate_free_port()
|
||||
base_url = f"http://127.0.0.1:{port}"
|
||||
process = launch_proxy(CONFIG_PATH, port, LOG_PATH)
|
||||
write_pid_record(PidRecord(pid=process.pid, port=port, config_path=str(CONFIG_PATH), log_path=str(LOG_PATH)))
|
||||
|
||||
try:
|
||||
poll_liveliness(base_url, LOG_PATH, process)
|
||||
except ProcessLaunchError as e:
|
||||
terminate(process.pid)
|
||||
clear_pid_record()
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
try:
|
||||
original_existed = CLAUDE_SETTINGS_PATH.exists()
|
||||
original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH)
|
||||
write_backup(
|
||||
ClaudeBackupRecord(existed=original_existed, content=original_settings if original_existed else None),
|
||||
AUTOROUTE_BACKUP_PATH,
|
||||
)
|
||||
merged = merge_claude_settings_static_token(original_settings, base_url, master_key)
|
||||
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
with secure_create(CLAUDE_SETTINGS_PATH) as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except UpError as e:
|
||||
terminate(process.pid)
|
||||
clear_pid_record()
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
click.echo(f"litellm: ephemeral auto-router proxy up at {base_url} (pid {process.pid})")
|
||||
click.echo("Claude Code sessions started now will route through it. Press Ctrl-C to stop and restore.")
|
||||
|
||||
stop_event = threading.Event()
|
||||
restored = threading.Lock()
|
||||
|
||||
def _teardown() -> None:
|
||||
if not restored.acquire(blocking=False):
|
||||
return
|
||||
terminate(process.pid)
|
||||
clear_pid_record()
|
||||
try:
|
||||
restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
|
||||
except UpError as e:
|
||||
# Runs from atexit/a signal handler too, outside Click's own exception
|
||||
# handling -- raising here would only produce an unhandled-exception
|
||||
# warning on stderr, not a clean message.
|
||||
click.echo(str(e), err=True)
|
||||
return
|
||||
click.echo("\nStopped ephemeral proxy and restored Claude Code settings.")
|
||||
click.echo(
|
||||
f"Restart any Claude Code session still open from this session, or another local account could "
|
||||
f"bind the now-free port {port} and receive its requests. Do not use `lite autoroute up` on a "
|
||||
f"shared or multi-tenant host."
|
||||
)
|
||||
|
||||
def _handle_signal(_signum: int, _frame: FrameType | None) -> None:
|
||||
stop_event.set()
|
||||
|
||||
signal.signal(signal.SIGINT, _handle_signal)
|
||||
signal.signal(signal.SIGTERM, _handle_signal)
|
||||
atexit.register(_teardown)
|
||||
|
||||
log_thread = threading.Thread(target=stream_log, args=(LOG_PATH, stop_event), daemon=True)
|
||||
log_thread.start()
|
||||
|
||||
stop_event.wait()
|
||||
_teardown()
|
||||
|
||||
|
||||
@autoroute_group.command("down")
|
||||
def down() -> None:
|
||||
"""Restore Claude Code settings and stop a leftover ephemeral proxy, if any"""
|
||||
try:
|
||||
record: PidRecord | None = read_pid_record()
|
||||
except UpError as e:
|
||||
# down is the crash-recovery path -- a corrupt pid record must not block it; clear the
|
||||
# unusable record and keep going rather than leaving the user with no way to clean up.
|
||||
click.echo(f"{e} Clearing it and continuing cleanup.", err=True)
|
||||
record = None
|
||||
if record is not None and is_running(record.pid):
|
||||
terminate(record.pid)
|
||||
click.echo(f"Stopped leftover ephemeral proxy (pid {record.pid}).")
|
||||
clear_pid_record()
|
||||
|
||||
try:
|
||||
restored = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
|
||||
except UpError as e:
|
||||
raise click.ClickException(str(e))
|
||||
if restored is None:
|
||||
click.echo("Nothing to restore.")
|
||||
elif restored.existed:
|
||||
click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.")
|
||||
else:
|
||||
click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite autoroute up`).")
|
||||
|
||||
|
||||
__all__ = ["autoroute_group"]
|
||||
249
litellm/proxy/client/cli/commands/autoroute/config.py
Normal file
249
litellm/proxy/client/cli/commands/autoroute/config.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
from typing import Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
|
||||
|
||||
TIER_NAMES: tuple[str, ...] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING")
|
||||
AUTOROUTER_MODEL_NAME = "autorouter"
|
||||
|
||||
|
||||
class ConfigGenerationError(Exception):
|
||||
"""Raised when an AutorouteConfig references a model the discovery step didn't find."""
|
||||
|
||||
|
||||
class DiscoveredModel(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
name: str
|
||||
mode: str = "chat"
|
||||
input_cost_per_token: float | None = None
|
||||
output_cost_per_token: float | None = None
|
||||
|
||||
|
||||
class _RawModelGroup(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
model_group: str
|
||||
# Optional: some real deployments return an explicit `"mode": null` for models that
|
||||
# were registered without a mode (seen for embedding models like voyage-4-large).
|
||||
# ModelGroupInfo's own "chat" default (litellm/types/router.py) only applies when the
|
||||
# key is missing entirely, not when it's present as null, so this must tolerate None.
|
||||
mode: str | None = "chat"
|
||||
input_cost_per_token: float | None = None
|
||||
output_cost_per_token: float | None = None
|
||||
|
||||
|
||||
_RAW_MODEL_GROUPS_ADAPTER = TypeAdapter(list[_RawModelGroup])
|
||||
|
||||
|
||||
def parse_discovered_models(raw: list[JsonValue]) -> tuple[DiscoveredModel, ...]:
|
||||
"""Validate a raw `/model_group/info` response into typed models."""
|
||||
parsed = _RAW_MODEL_GROUPS_ADAPTER.validate_python(raw)
|
||||
return tuple(
|
||||
DiscoveredModel(
|
||||
name=group.model_group,
|
||||
# A null mode means the server genuinely doesn't know what this model does;
|
||||
# "unknown" (rather than guessing "chat") keeps it out of both chat_models()
|
||||
# and embedding_models() instead of risking a wrong-mode deployment.
|
||||
mode=group.mode or "unknown",
|
||||
input_cost_per_token=group.input_cost_per_token,
|
||||
output_cost_per_token=group.output_cost_per_token,
|
||||
)
|
||||
for group in parsed
|
||||
)
|
||||
|
||||
|
||||
def chat_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]:
|
||||
return tuple(m for m in models if m.mode == "chat")
|
||||
|
||||
|
||||
def embedding_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]:
|
||||
return tuple(m for m in models if m.mode == "embedding")
|
||||
|
||||
|
||||
class HeuristicClassifier(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
kind: Literal["heuristic"] = "heuristic"
|
||||
|
||||
|
||||
class LLMClassifier(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
kind: Literal["llm"] = "llm"
|
||||
model: str
|
||||
timeout_ms: int = 3000
|
||||
|
||||
|
||||
ClassifierChoice = Union[HeuristicClassifier, LLMClassifier]
|
||||
|
||||
|
||||
class NoSemanticMatching(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
kind: Literal["none"] = "none"
|
||||
|
||||
|
||||
class KeywordTierRule(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
keywords: tuple[str, ...]
|
||||
tier: str
|
||||
|
||||
|
||||
# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules"
|
||||
# invariant with a sane starting point; the wizard lets the user override these per tier.
|
||||
DEFAULT_KEYWORD_TIER_RULES: tuple[KeywordTierRule, ...] = (
|
||||
KeywordTierRule(keywords=("hi", "hello", "thanks"), tier="SIMPLE"),
|
||||
KeywordTierRule(keywords=("explain", "how does"), tier="MEDIUM"),
|
||||
KeywordTierRule(keywords=("refactor", "implement", "debug"), tier="COMPLEX"),
|
||||
KeywordTierRule(keywords=("step by step", "think through", "prove"), tier="REASONING"),
|
||||
)
|
||||
|
||||
|
||||
class SemanticMatching(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
kind: Literal["semantic"] = "semantic"
|
||||
embedding_model: str
|
||||
match_threshold: float = 0.5
|
||||
keyword_tier_rules: tuple[KeywordTierRule, ...] = DEFAULT_KEYWORD_TIER_RULES
|
||||
|
||||
|
||||
SemanticMatchingChoice = Union[NoSemanticMatching, SemanticMatching]
|
||||
|
||||
|
||||
class AutorouteConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
base_url: str
|
||||
api_key: str
|
||||
# Each tier maps to a pool of one or more models; complexity_router picks randomly among
|
||||
# them per request (or, in adaptive mode, learns which to prefer within the pool).
|
||||
tiers: dict[str, tuple[str, ...]]
|
||||
default_model: str
|
||||
classifier: ClassifierChoice = Field(default_factory=HeuristicClassifier)
|
||||
semantic_matching: SemanticMatchingChoice = Field(default_factory=NoSemanticMatching)
|
||||
adaptive: bool = False
|
||||
|
||||
|
||||
def validate_config(config: AutorouteConfig, discovered: tuple[DiscoveredModel, ...]) -> None:
|
||||
"""Raise ConfigGenerationError if config references a model discovery didn't return."""
|
||||
chat_names: frozenset[str] = frozenset(m.name for m in chat_models(discovered))
|
||||
embedding_names: frozenset[str] = frozenset(m.name for m in embedding_models(discovered))
|
||||
|
||||
for tier, models in config.tiers.items():
|
||||
for model in models:
|
||||
if model not in chat_names:
|
||||
raise ConfigGenerationError(f"Tier {tier} references unknown chat model '{model}'")
|
||||
|
||||
if config.default_model not in chat_names:
|
||||
raise ConfigGenerationError(f"default_model '{config.default_model}' is not a known chat model")
|
||||
|
||||
if isinstance(config.classifier, LLMClassifier) and config.classifier.model not in chat_names:
|
||||
raise ConfigGenerationError(f"classifier model '{config.classifier.model}' is not a known chat model")
|
||||
|
||||
if (
|
||||
isinstance(config.semantic_matching, SemanticMatching)
|
||||
and config.semantic_matching.embedding_model not in embedding_names
|
||||
):
|
||||
raise ConfigGenerationError(
|
||||
f"embedding model '{config.semantic_matching.embedding_model}' is not a known embedding model"
|
||||
)
|
||||
|
||||
|
||||
def _litellm_proxy_deployment(name: str, base_url: str, api_key: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": f"litellm_proxy/{name}",
|
||||
"api_base": base_url,
|
||||
"api_key": api_key,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]:
|
||||
"""Build the model_list for the ephemeral proxy's config.yaml.
|
||||
|
||||
Every real model referenced anywhere (tier targets, classifier, embedding) is deduplicated
|
||||
to exactly one `litellm_proxy/<name>` deployment forwarding to the customer's real proxy,
|
||||
plus one `auto_router/complexity_router` deployment tying the tiers together.
|
||||
"""
|
||||
referenced_names = {model for models in config.tiers.values() for model in models}
|
||||
referenced_names.add(config.default_model)
|
||||
if isinstance(config.classifier, LLMClassifier):
|
||||
referenced_names.add(config.classifier.model)
|
||||
if isinstance(config.semantic_matching, SemanticMatching):
|
||||
referenced_names.add(config.semantic_matching.embedding_model)
|
||||
|
||||
proxy_deployments = [
|
||||
_litellm_proxy_deployment(name, config.base_url, config.api_key) for name in sorted(referenced_names)
|
||||
]
|
||||
|
||||
complexity_router_config: dict[str, JsonValue] = {
|
||||
"tiers": {tier: list(models) for tier, models in config.tiers.items()},
|
||||
"default_model": config.default_model,
|
||||
}
|
||||
if isinstance(config.classifier, LLMClassifier):
|
||||
complexity_router_config["classifier_type"] = "llm"
|
||||
complexity_router_config["classifier_llm_config"] = {
|
||||
"model": config.classifier.model,
|
||||
"timeout_ms": config.classifier.timeout_ms,
|
||||
}
|
||||
if isinstance(config.semantic_matching, SemanticMatching):
|
||||
complexity_router_config["semantic_keyword_matching"] = True
|
||||
complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model
|
||||
complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold
|
||||
complexity_router_config["keyword_tier_rules"] = [
|
||||
{"keywords": list(rule.keywords), "tier": rule.tier} for rule in config.semantic_matching.keyword_tier_rules
|
||||
]
|
||||
if config.adaptive:
|
||||
complexity_router_config["adaptive"] = True
|
||||
|
||||
auto_router_litellm_params: dict[str, JsonValue] = {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": complexity_router_config,
|
||||
}
|
||||
# A bare "*" model_name looks like the obvious way to catch every request Claude Code
|
||||
# might send regardless of which model it thinks it's using, but Router's auto-router
|
||||
# registry is keyed by the literal requested model string (router.py:10711-10717), not
|
||||
# resolved through pattern/wildcard matching first -- so a "*" entry here would only ever
|
||||
# match a client that literally sends model="*", never an actual wildcard catch-all. Callers
|
||||
# instead need to make Claude Code request this "autorouter" name directly (see
|
||||
# ANTHROPIC_DEFAULT_*_MODEL in settings.py's merge_claude_settings_static_token).
|
||||
return [
|
||||
*proxy_deployments,
|
||||
{"model_name": AUTOROUTER_MODEL_NAME, "litellm_params": auto_router_litellm_params},
|
||||
]
|
||||
|
||||
|
||||
def build_generated_proxy_config(config: AutorouteConfig, master_key: str) -> dict[str, JsonValue]:
|
||||
"""Full config.yaml content for the ephemeral proxy, including its own auth key.
|
||||
|
||||
master_key must live under general_settings, not litellm_settings -- the proxy server
|
||||
only ever reads general_settings.master_key (proxy_server.py:4530) to authenticate
|
||||
requests; a key placed under litellm_settings is silently ignored, leaving the proxy
|
||||
with no real auth at all.
|
||||
"""
|
||||
return {
|
||||
"model_list": build_generated_model_list(config),
|
||||
"general_settings": {"master_key": master_key},
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AUTOROUTER_MODEL_NAME",
|
||||
"TIER_NAMES",
|
||||
"AutorouteConfig",
|
||||
"ClassifierChoice",
|
||||
"ConfigGenerationError",
|
||||
"DEFAULT_KEYWORD_TIER_RULES",
|
||||
"DiscoveredModel",
|
||||
"HeuristicClassifier",
|
||||
"KeywordTierRule",
|
||||
"LLMClassifier",
|
||||
"NoSemanticMatching",
|
||||
"SemanticMatching",
|
||||
"SemanticMatchingChoice",
|
||||
"build_generated_model_list",
|
||||
"build_generated_proxy_config",
|
||||
"chat_models",
|
||||
"embedding_models",
|
||||
"parse_discovered_models",
|
||||
"validate_config",
|
||||
]
|
||||
190
litellm/proxy/client/cli/commands/autoroute/process.py
Normal file
190
litellm/proxy/client/cli/commands/autoroute/process.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
import contextlib
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
import requests
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from ..up import UpError, secure_create
|
||||
|
||||
AUTOROUTE_DIR = Path.home() / ".litellm" / "autorouter"
|
||||
CONFIG_PATH = AUTOROUTE_DIR / "config.yaml"
|
||||
LOG_PATH = AUTOROUTE_DIR / "proxy.log"
|
||||
PID_RECORD_PATH = AUTOROUTE_DIR / "proxy.pid.json"
|
||||
|
||||
|
||||
class ProcessLaunchError(Exception):
|
||||
"""Raised when the ephemeral proxy subprocess fails to come up healthy."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PidRecord:
|
||||
pid: int
|
||||
port: int
|
||||
config_path: str
|
||||
log_path: str
|
||||
|
||||
|
||||
_PID_RECORD_ADAPTER = TypeAdapter(PidRecord)
|
||||
|
||||
|
||||
_PROXY_RUNTIME_MODULES: tuple[str, ...] = ("fastapi", "uvicorn", "backoff", "orjson", "websockets", "apscheduler")
|
||||
|
||||
|
||||
def missing_proxy_runtime_modules() -> tuple[str, ...]:
|
||||
"""Proxy-server modules that ``lite autoroute up`` needs but the thin CLI install lacks.
|
||||
|
||||
``launch_proxy`` runs the full ``litellm.proxy.proxy_cli`` server, whose dependencies live in
|
||||
the ``proxy`` extra, not the ``cli`` extra that installs the ``lite`` command. On a thin
|
||||
``litellm[cli]`` install the subprocess dies with a bare ``ModuleNotFoundError``; detecting the
|
||||
gap here lets ``up`` fail with an actionable message instead.
|
||||
"""
|
||||
return tuple(name for name in _PROXY_RUNTIME_MODULES if importlib.util.find_spec(name) is None)
|
||||
|
||||
|
||||
def allocate_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return int(sock.getsockname()[1])
|
||||
|
||||
|
||||
def launch_proxy(config_path: Path, port: int, log_path: Path) -> "subprocess.Popen[bytes]":
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(log_path, "w") as log_file:
|
||||
return subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"litellm.proxy.proxy_cli",
|
||||
"--config",
|
||||
str(config_path),
|
||||
"--port",
|
||||
str(port),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
],
|
||||
stdout=log_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
|
||||
|
||||
def _tail(log_path: Path, lines: int = 40) -> str:
|
||||
if not log_path.exists():
|
||||
return "(no log output captured)"
|
||||
return "\n".join(log_path.read_text(errors="replace").splitlines()[-lines:])
|
||||
|
||||
|
||||
def poll_liveliness(base_url: str, log_path: Path, process: "subprocess.Popen[bytes]", timeout: float = 30.0) -> None:
|
||||
"""Poll /health/liveliness until it responds, the process dies, or timeout elapses."""
|
||||
deadline = time.monotonic() + timeout
|
||||
url = base_url.rstrip("/") + "/health/liveliness"
|
||||
while time.monotonic() < deadline:
|
||||
if process.poll() is not None:
|
||||
raise ProcessLaunchError(
|
||||
f"Ephemeral proxy exited early (code {process.returncode}). Last log lines:\n{_tail(log_path)}"
|
||||
)
|
||||
with contextlib.suppress(requests.RequestException):
|
||||
if requests.get(url, timeout=2).status_code == 200:
|
||||
return
|
||||
time.sleep(0.5)
|
||||
raise ProcessLaunchError(
|
||||
f"Ephemeral proxy never became healthy within {timeout}s. Last log lines:\n{_tail(log_path)}"
|
||||
)
|
||||
|
||||
|
||||
def write_pid_record(record: PidRecord, path: Path | None = None) -> None:
|
||||
resolved_path = path if path is not None else PID_RECORD_PATH
|
||||
resolved_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(resolved_path, "w") as f:
|
||||
json.dump(
|
||||
{"pid": record.pid, "port": record.port, "config_path": record.config_path, "log_path": record.log_path},
|
||||
f,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
|
||||
def read_pid_record(path: Path | None = None) -> PidRecord | None:
|
||||
resolved_path = path if path is not None else PID_RECORD_PATH
|
||||
if not resolved_path.exists():
|
||||
return None
|
||||
with open(resolved_path, "r") as f:
|
||||
content = f.read()
|
||||
try:
|
||||
return _PID_RECORD_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise UpError(f"{resolved_path} contains invalid or unexpected JSON; cannot proceed safely.")
|
||||
|
||||
|
||||
def clear_pid_record(path: Path | None = None) -> None:
|
||||
resolved_path = path if path is not None else PID_RECORD_PATH
|
||||
resolved_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def is_running(pid: int) -> bool:
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except PermissionError:
|
||||
return True
|
||||
return True
|
||||
|
||||
|
||||
def terminate(pid: int, grace_period: float = 5.0) -> None:
|
||||
"""Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed."""
|
||||
if not is_running(pid):
|
||||
return
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
deadline = time.monotonic() + grace_period
|
||||
while time.monotonic() < deadline and is_running(pid):
|
||||
time.sleep(0.2)
|
||||
if is_running(pid):
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
|
||||
|
||||
def stream_log(log_path: Path, stop_event: threading.Event) -> None:
|
||||
"""Print new lines appended to log_path until stop_event is set. Blocks the calling thread."""
|
||||
while not log_path.exists() and not stop_event.is_set():
|
||||
time.sleep(0.1)
|
||||
if stop_event.is_set() or not log_path.exists():
|
||||
return
|
||||
with open(log_path, "r") as f:
|
||||
while not stop_event.is_set():
|
||||
line = f.readline()
|
||||
if line:
|
||||
click.echo(line, nl=False)
|
||||
else:
|
||||
time.sleep(0.2)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AUTOROUTE_DIR",
|
||||
"CONFIG_PATH",
|
||||
"LOG_PATH",
|
||||
"PID_RECORD_PATH",
|
||||
"PidRecord",
|
||||
"ProcessLaunchError",
|
||||
"allocate_free_port",
|
||||
"clear_pid_record",
|
||||
"is_running",
|
||||
"launch_proxy",
|
||||
"missing_proxy_runtime_modules",
|
||||
"poll_liveliness",
|
||||
"read_pid_record",
|
||||
"secure_create",
|
||||
"stream_log",
|
||||
"terminate",
|
||||
"write_pid_record",
|
||||
]
|
||||
46
litellm/proxy/client/cli/commands/autoroute/settings.py
Normal file
46
litellm/proxy/client/cli/commands/autoroute/settings.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
from pydantic import JsonValue
|
||||
|
||||
from .config import AUTOROUTER_MODEL_NAME
|
||||
|
||||
ENV_KEY = "env"
|
||||
API_KEY_HELPER_KEY = "apiKeyHelper"
|
||||
ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY"
|
||||
ANTHROPIC_AUTH_TOKEN_KEY = "ANTHROPIC_AUTH_TOKEN"
|
||||
ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL"
|
||||
# Force every one of Claude Code's own model tiers to request the auto-router by name.
|
||||
# Router's auto-router registry is keyed by the literal requested model string
|
||||
# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
|
||||
# model_name can never work as a catch-all -- these overrides are what actually makes
|
||||
# Claude Code send "autorouter" regardless of /model or its own version-specific defaults.
|
||||
ANTHROPIC_DEFAULT_MODEL_ENV_KEYS = (
|
||||
"ANTHROPIC_DEFAULT_SONNET_MODEL",
|
||||
"ANTHROPIC_DEFAULT_HAIKU_MODEL",
|
||||
"ANTHROPIC_DEFAULT_OPUS_MODEL",
|
||||
)
|
||||
|
||||
|
||||
def merge_claude_settings_static_token(
|
||||
settings: dict[str, JsonValue], base_url: str, auth_token: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to a local ephemeral proxy with a static token.
|
||||
|
||||
Unlike up.py's merge_claude_settings (which sets apiKeyHelper for a long-lived, real
|
||||
remote proxy needing refreshable SSO tokens), this proxy is ephemeral and its key was just
|
||||
minted for this session, so a plain env var is simpler and correct. Any existing
|
||||
apiKeyHelper is cleared so it can't fight with the static token.
|
||||
"""
|
||||
raw_env = settings.get(ENV_KEY, {})
|
||||
base_env = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: dict[str, JsonValue] = {
|
||||
**base_env,
|
||||
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
|
||||
ANTHROPIC_AUTH_TOKEN_KEY: auth_token,
|
||||
**{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS},
|
||||
}
|
||||
env.pop(ANTHROPIC_API_KEY_KEY, None)
|
||||
merged: dict[str, JsonValue] = {**settings, ENV_KEY: env}
|
||||
merged.pop(API_KEY_HELPER_KEY, None)
|
||||
return merged
|
||||
|
||||
|
||||
__all__ = ["merge_claude_settings_static_token"]
|
||||
150
litellm/proxy/client/cli/commands/autoroute/wizard.py
Normal file
150
litellm/proxy/client/cli/commands/autoroute/wizard.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
import yaml
|
||||
from InquirerPy import inquirer
|
||||
from InquirerPy.base.control import Choice
|
||||
|
||||
from .... import Client
|
||||
from .config import (
|
||||
DEFAULT_KEYWORD_TIER_RULES,
|
||||
TIER_NAMES,
|
||||
AutorouteConfig,
|
||||
ConfigGenerationError,
|
||||
DiscoveredModel,
|
||||
HeuristicClassifier,
|
||||
KeywordTierRule,
|
||||
LLMClassifier,
|
||||
NoSemanticMatching,
|
||||
SemanticMatching,
|
||||
build_generated_model_list,
|
||||
chat_models,
|
||||
embedding_models,
|
||||
parse_discovered_models,
|
||||
validate_config,
|
||||
)
|
||||
from .process import CONFIG_PATH, secure_create
|
||||
|
||||
|
||||
def _is_interactive() -> bool:
|
||||
return sys.stdin.isatty()
|
||||
|
||||
|
||||
def _fuzzy_pick(models: tuple[DiscoveredModel, ...], prompt_label: str, multiselect: bool) -> list[str]:
|
||||
"""Type-to-filter picker over a (possibly huge) model pool, using InquirerPy's fzf-style fuzzy prompt.
|
||||
|
||||
A plain numbered table + typed index does not scale past a handful of models -- proxies with
|
||||
hundreds of model groups made that interaction unusable. This lets the user narrow the pool by
|
||||
typing a substring instead of scrolling/counting.
|
||||
|
||||
Assumes the caller already checked interactivity (run_configure_wizard does, once, up front) --
|
||||
checking here too would check the wrong thing under test, where InquirerPy is driven through its
|
||||
own injected input/output rather than the real process stdin.
|
||||
"""
|
||||
choices = [Choice(value=model.name, name=model.name) for model in models]
|
||||
toggle_hint = "tab to toggle, " if multiselect else ""
|
||||
while True:
|
||||
result = inquirer.fuzzy(
|
||||
message=f"{prompt_label}: type to filter, {toggle_hint}enter to confirm",
|
||||
choices=choices,
|
||||
multiselect=multiselect,
|
||||
max_height="70%",
|
||||
).execute()
|
||||
selected = result if multiselect else [result]
|
||||
if selected:
|
||||
return selected
|
||||
click.echo("Select at least one model.")
|
||||
|
||||
|
||||
def _render_and_prompt_for_model(models: tuple[DiscoveredModel, ...], prompt_label: str) -> str:
|
||||
return _fuzzy_pick(models, prompt_label, multiselect=False)[0]
|
||||
|
||||
|
||||
def _render_and_prompt_for_models(models: tuple[DiscoveredModel, ...], prompt_label: str) -> tuple[str, ...]:
|
||||
return tuple(_fuzzy_pick(models, prompt_label, multiselect=True))
|
||||
|
||||
|
||||
def _parse_keywords(raw: str) -> tuple[str, ...]:
|
||||
return tuple(keyword.strip() for keyword in raw.split(",") if keyword.strip())
|
||||
|
||||
|
||||
def _prompt_for_keyword_tier_rules() -> tuple[KeywordTierRule, ...]:
|
||||
"""Let the user supply the semantic-matching keywords per tier, since matching those
|
||||
keywords against the request is the whole point of enabling it. Each prompt is prefilled
|
||||
with the built-in default, so pressing enter keeps it."""
|
||||
click.echo("\nEnter example keywords/phrases per tier (comma-separated); press enter to keep the default:")
|
||||
defaults = {rule.tier: rule.keywords for rule in DEFAULT_KEYWORD_TIER_RULES}
|
||||
|
||||
def _rule_for(tier: str) -> KeywordTierRule:
|
||||
default_keywords = defaults.get(tier, ())
|
||||
raw = click.prompt(f" {tier} keywords", default=", ".join(default_keywords), show_default=True)
|
||||
return KeywordTierRule(keywords=_parse_keywords(raw) or default_keywords, tier=tier)
|
||||
|
||||
return tuple(_rule_for(tier) for tier in TIER_NAMES)
|
||||
|
||||
|
||||
def run_configure_wizard(ctx: click.Context) -> Path:
|
||||
"""Discover the caller's accessible models, walk them through tier assignment, write config."""
|
||||
base_url = ctx.obj["base_url"]
|
||||
api_key = ctx.obj["api_key"]
|
||||
client = Client(base_url=base_url, api_key=api_key)
|
||||
|
||||
raw_groups = client.model_groups.info()
|
||||
if not isinstance(raw_groups, list):
|
||||
raise click.ClickException(
|
||||
f"Unexpected response from /model_group/info: expected a list, got {type(raw_groups).__name__}"
|
||||
)
|
||||
discovered = parse_discovered_models(raw_groups)
|
||||
chat_pool = chat_models(discovered)
|
||||
embedding_pool = embedding_models(discovered)
|
||||
|
||||
if not chat_pool:
|
||||
raise click.ClickException("Your key has no chat-capable models available on this proxy.")
|
||||
|
||||
if not _is_interactive():
|
||||
raise click.ClickException("`lite autoroute configure` requires an interactive terminal.")
|
||||
|
||||
click.echo("Assign model(s) to each complexity tier (from what your key can access):")
|
||||
tiers = {tier: _render_and_prompt_for_models(chat_pool, tier) for tier in TIER_NAMES}
|
||||
default_model = tiers["MEDIUM"][0]
|
||||
|
||||
classifier = HeuristicClassifier()
|
||||
if click.confirm("\nUse an LLM classifier instead of the free heuristic scorer?", default=False):
|
||||
classifier_model = _render_and_prompt_for_model(chat_pool, "LLM classifier")
|
||||
classifier = LLMClassifier(model=classifier_model)
|
||||
|
||||
semantic_matching = NoSemanticMatching()
|
||||
if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False):
|
||||
embedding_model = _render_and_prompt_for_model(embedding_pool, "semantic embeddings")
|
||||
keyword_tier_rules = _prompt_for_keyword_tier_rules()
|
||||
semantic_matching = SemanticMatching(embedding_model=embedding_model, keyword_tier_rules=keyword_tier_rules)
|
||||
|
||||
adaptive = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False)
|
||||
|
||||
config = AutorouteConfig(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
tiers=tiers,
|
||||
default_model=default_model,
|
||||
classifier=classifier,
|
||||
semantic_matching=semantic_matching,
|
||||
adaptive=adaptive,
|
||||
)
|
||||
try:
|
||||
validate_config(config, discovered)
|
||||
except ConfigGenerationError as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
model_list = build_generated_model_list(config)
|
||||
CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
with secure_create(CONFIG_PATH) as f:
|
||||
yaml.safe_dump({"model_list": model_list}, f, sort_keys=False)
|
||||
|
||||
click.echo(f"\nWrote {CONFIG_PATH}")
|
||||
for tier, models in tiers.items():
|
||||
click.echo(f" {tier}: {', '.join(models)}")
|
||||
return CONFIG_PATH
|
||||
|
||||
|
||||
__all__ = ["run_configure_wizard"]
|
||||
57
litellm/proxy/client/cli/commands/model_groups.py
Normal file
57
litellm/proxy/client/cli/commands/model_groups.py
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
from typing import Literal
|
||||
|
||||
import click
|
||||
import rich
|
||||
import rich.table
|
||||
|
||||
from ... import Client
|
||||
|
||||
|
||||
def create_client(ctx: click.Context) -> Client:
|
||||
return Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
|
||||
|
||||
|
||||
@click.group(name="model-groups")
|
||||
def model_groups() -> None:
|
||||
"""Inspect model groups your key can access on the proxy"""
|
||||
|
||||
|
||||
@model_groups.command("list")
|
||||
@click.option(
|
||||
"--format",
|
||||
"output_format",
|
||||
type=click.Choice(["table", "json"]),
|
||||
default="table",
|
||||
help="Output format (table or json)",
|
||||
)
|
||||
@click.pass_context
|
||||
def list_model_groups(ctx: click.Context, output_format: Literal["table", "json"]) -> None:
|
||||
"""List model groups accessible to your key, with mode and pricing"""
|
||||
client = create_client(ctx)
|
||||
groups = client.model_groups.info()
|
||||
if not isinstance(groups, list):
|
||||
raise click.ClickException(
|
||||
f"Unexpected response from /model_group/info: expected a list, got {type(groups).__name__}"
|
||||
)
|
||||
|
||||
if output_format == "json":
|
||||
rich.print_json(data=groups)
|
||||
return
|
||||
|
||||
table = rich.table.Table(title="Accessible Model Groups")
|
||||
table.add_column("Model", style="cyan")
|
||||
table.add_column("Mode", style="green")
|
||||
table.add_column("Input $/token", style="yellow")
|
||||
table.add_column("Output $/token", style="yellow")
|
||||
|
||||
for group in groups:
|
||||
table.add_row(
|
||||
str(group.get("model_group", "")),
|
||||
str(group.get("mode", "chat")),
|
||||
str(group.get("input_cost_per_token", "")),
|
||||
str(group.get("output_cost_per_token", "")),
|
||||
)
|
||||
rich.print(table)
|
||||
|
||||
|
||||
__all__ = ["model_groups"]
|
||||
283
litellm/proxy/client/cli/commands/up.py
Normal file
283
litellm/proxy/client/cli/commands/up.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
import atexit
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import shutil
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import FrameType
|
||||
from typing import IO, Iterator, Mapping
|
||||
|
||||
import click
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
||||
|
||||
from .agents import AgentRunError, resolve_api_key, verify_proxy_key
|
||||
from .auth import load_token, login
|
||||
|
||||
ENV_KEY = "env"
|
||||
API_KEY_HELPER_KEY = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_API_KEY_KEY = "ANTHROPIC_API_KEY"
|
||||
|
||||
CLAUDE_SETTINGS_PATH = Path.home() / ".claude" / "settings.json"
|
||||
BACKUP_PATH = Path.home() / ".litellm" / "claude_settings_backup.json"
|
||||
|
||||
|
||||
class UpError(Exception):
|
||||
"""Raised for any user-actionable failure while starting/stopping interception."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BackupRecord:
|
||||
"""Snapshot of ~/.claude/settings.json taken right before `lite up` patches it."""
|
||||
|
||||
existed: bool
|
||||
content: dict[str, JsonValue] | None
|
||||
|
||||
|
||||
_SETTINGS_ADAPTER = TypeAdapter(dict[str, JsonValue])
|
||||
_BACKUP_RECORD_ADAPTER = TypeAdapter(BackupRecord)
|
||||
|
||||
|
||||
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
with open(path, "r") as f:
|
||||
content = f.read()
|
||||
if not content.strip():
|
||||
return {}
|
||||
try:
|
||||
return _SETTINGS_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.")
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to route Claude Code through the proxy.
|
||||
|
||||
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
|
||||
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
|
||||
token (same reasoning as build_agent_env in agents.py). Every other key is
|
||||
preserved untouched.
|
||||
"""
|
||||
raw_env = settings.get(ENV_KEY, {})
|
||||
base_env = raw_env if isinstance(raw_env, dict) else {}
|
||||
env = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")}
|
||||
env.pop(ANTHROPIC_API_KEY_KEY, None)
|
||||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def secure_create(path: Path) -> Iterator[IO[str]]:
|
||||
"""Open path for writing with mode 0600 fixed up before any content is written.
|
||||
|
||||
A plain `open(path, "w")` creates a *new* file at the umask-derived default (commonly 0644)
|
||||
and leaves it world- or group-readable until a later `chmod` call catches up -- a real window
|
||||
in which a file holding a credential is readable by another local account. Passing the mode to
|
||||
`os.open` closes that window for a brand-new file, but `O_CREAT`'s mode argument is only
|
||||
applied on creation: if the file already exists its old, broader permissions carry over
|
||||
untouched. `os.fchmod` right after opening -- before a single byte of the new content is
|
||||
written -- covers both cases.
|
||||
"""
|
||||
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
os.fchmod(fd, 0o600)
|
||||
f: IO[str] = os.fdopen(fd, "w")
|
||||
try:
|
||||
yield f
|
||||
finally:
|
||||
f.close()
|
||||
|
||||
|
||||
def write_backup(record: BackupRecord, backup_path: Path | None = None) -> None:
|
||||
path = backup_path if backup_path is not None else BACKUP_PATH
|
||||
path.parent.mkdir(exist_ok=True)
|
||||
with secure_create(path) as f:
|
||||
json.dump({"existed": record.existed, "content": record.content}, f, indent=2)
|
||||
|
||||
|
||||
def read_backup(backup_path: Path | None = None) -> BackupRecord | None:
|
||||
path = backup_path if backup_path is not None else BACKUP_PATH
|
||||
if not path.exists():
|
||||
return None
|
||||
with open(path, "r") as f:
|
||||
content = f.read()
|
||||
try:
|
||||
return _BACKUP_RECORD_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise UpError(f"{path} contains invalid or unexpected JSON; cannot restore from it safely.")
|
||||
|
||||
|
||||
def restore_claude_settings(settings_path: Path | None = None, backup_path: Path | None = None) -> BackupRecord | None:
|
||||
"""Restore settings_path from the backup at backup_path, then delete the backup.
|
||||
|
||||
Returns the restored record, or None if there was nothing to restore.
|
||||
"""
|
||||
resolved_settings_path = settings_path if settings_path is not None else CLAUDE_SETTINGS_PATH
|
||||
resolved_backup_path = backup_path if backup_path is not None else BACKUP_PATH
|
||||
record = read_backup(resolved_backup_path)
|
||||
if record is None:
|
||||
return None
|
||||
if record.existed and record.content is not None:
|
||||
resolved_settings_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(resolved_settings_path, "w") as f:
|
||||
json.dump(record.content, f, indent=2)
|
||||
elif resolved_settings_path.exists():
|
||||
resolved_settings_path.unlink()
|
||||
resolved_backup_path.unlink()
|
||||
return record
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
Code would otherwise use) makes `print-token` enforce that the cached
|
||||
token was actually issued for this proxy -- without it, a token minted
|
||||
for a different, previously-logged-into proxy would be handed to
|
||||
whichever server `up` currently points at.
|
||||
"""
|
||||
lite_path = shutil.which("lite")
|
||||
if lite_path is None:
|
||||
raise UpError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs "
|
||||
"an absolute path to it, so `lite up` cannot continue."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}"
|
||||
|
||||
|
||||
def _ensure_fresh_login(ctx: click.Context) -> None:
|
||||
base_url = ctx.obj["base_url"].rstrip("/")
|
||||
token_data = load_token()
|
||||
if token_data and token_data.get("base_url") == base_url and is_cli_token_fresh(token_data):
|
||||
return
|
||||
|
||||
if not sys.stdin.isatty():
|
||||
raise UpError(
|
||||
"No fresh LiteLLM login found for this proxy. Run `lite login` first (apiKeyHelper "
|
||||
"reads this token on every Claude Code request)."
|
||||
)
|
||||
|
||||
click.echo("No fresh LiteLLM login found for this proxy; starting login...")
|
||||
ctx.invoke(login)
|
||||
token_data = load_token()
|
||||
if not token_data or token_data.get("base_url") != base_url or not is_cli_token_fresh(token_data):
|
||||
raise UpError("Login did not produce a usable token; cannot start `lite up`.")
|
||||
|
||||
|
||||
def _restore_and_report() -> None:
|
||||
record = restore_claude_settings()
|
||||
if record is None:
|
||||
click.echo("Nothing to restore.")
|
||||
return
|
||||
if record.existed:
|
||||
click.echo(f"Restored {CLAUDE_SETTINGS_PATH} to its original contents.")
|
||||
else:
|
||||
click.echo(f"Removed {CLAUDE_SETTINGS_PATH} (it did not exist before `lite up`).")
|
||||
|
||||
|
||||
@click.command(name="up")
|
||||
@click.pass_context
|
||||
def up(ctx: click.Context) -> None:
|
||||
"""Route every Claude Code session through your LiteLLM proxy until stopped.
|
||||
|
||||
Patches ~/.claude/settings.json so Claude Code picks up the proxy on its own
|
||||
next startup, from any terminal -- no need to launch it through `lite`.
|
||||
Press Ctrl-C to stop and restore your original settings. Assumes the proxy
|
||||
is already running (this does not start one for you). Cursor is not
|
||||
supported: it has no equivalent file-based config to patch.
|
||||
"""
|
||||
base_url = ctx.obj["base_url"]
|
||||
|
||||
try:
|
||||
_ensure_fresh_login(ctx)
|
||||
api_key = resolve_api_key(ctx)
|
||||
verify_proxy_key(base_url, api_key)
|
||||
|
||||
if BACKUP_PATH.exists():
|
||||
raise UpError(
|
||||
f"{BACKUP_PATH} already exists -- `lite up` looks like it's already "
|
||||
"running (or crashed without cleanup). Run `lite down` first."
|
||||
)
|
||||
|
||||
api_key_helper = resolve_api_key_helper(base_url)
|
||||
original_existed = CLAUDE_SETTINGS_PATH.exists()
|
||||
original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH)
|
||||
write_backup(
|
||||
BackupRecord(
|
||||
existed=original_existed,
|
||||
content=original_settings if original_existed else None,
|
||||
)
|
||||
)
|
||||
|
||||
CLAUDE_SETTINGS_PATH.parent.mkdir(exist_ok=True)
|
||||
merged = merge_claude_settings(original_settings, base_url, api_key_helper)
|
||||
with open(CLAUDE_SETTINGS_PATH, "w") as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except (AgentRunError, UpError) as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}")
|
||||
click.echo("Press Ctrl-C to stop and restore your original settings.")
|
||||
|
||||
stop_event = threading.Event()
|
||||
restored = threading.Lock()
|
||||
|
||||
def _handle_signal(_signum: int, _frame: FrameType | None) -> None:
|
||||
stop_event.set()
|
||||
|
||||
def _restore_once() -> None:
|
||||
if not restored.acquire(blocking=False):
|
||||
return
|
||||
try:
|
||||
_restore_and_report()
|
||||
except UpError as e:
|
||||
# Runs from atexit/a signal handler, outside Click's own exception
|
||||
# handling -- raising here would only produce an unhandled-exception
|
||||
# warning on stderr, not a clean message.
|
||||
click.echo(str(e), err=True)
|
||||
|
||||
signal.signal(signal.SIGINT, _handle_signal)
|
||||
signal.signal(signal.SIGTERM, _handle_signal)
|
||||
atexit.register(_restore_once)
|
||||
|
||||
stop_event.wait()
|
||||
_restore_once()
|
||||
|
||||
|
||||
@click.command(name="down")
|
||||
def down() -> None:
|
||||
"""Restore ~/.claude/settings.json if a `lite up` session left it patched.
|
||||
|
||||
Use this after a `lite up` process was killed uncleanly (e.g. `kill -9`)
|
||||
instead of stopped with Ctrl-C.
|
||||
"""
|
||||
try:
|
||||
_restore_and_report()
|
||||
except UpError as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BACKUP_PATH",
|
||||
"CLAUDE_SETTINGS_PATH",
|
||||
"BackupRecord",
|
||||
"UpError",
|
||||
"down",
|
||||
"load_json_or_empty",
|
||||
"merge_claude_settings",
|
||||
"read_backup",
|
||||
"resolve_api_key_helper",
|
||||
"restore_claude_settings",
|
||||
"up",
|
||||
"write_backup",
|
||||
]
|
||||
|
|
@ -9,15 +9,18 @@ from litellm.proxy.client.health import HealthManagementClient
|
|||
|
||||
from .commands.agents import agent_commands
|
||||
from .commands.auth import auth_group, get_stored_api_key, login, logout, whoami
|
||||
from .commands.autoroute.commands import autoroute_group
|
||||
from .commands.chat import chat
|
||||
from .commands.credentials import credentials
|
||||
from .commands.encryption import encryption
|
||||
from .commands.http import http
|
||||
from .commands.keys import keys
|
||||
from .commands.model_groups import model_groups
|
||||
|
||||
# local imports
|
||||
from .commands.models import models
|
||||
from .commands.teams import teams
|
||||
from .commands.up import down, up
|
||||
from .commands.users import users
|
||||
from .interface import interactive_shell
|
||||
|
||||
|
|
@ -131,6 +134,13 @@ cli.add_command(users)
|
|||
# Add a top-level command per coding agent (claude, codex, opencode, ...)
|
||||
for agent_command in agent_commands():
|
||||
cli.add_command(agent_command)
|
||||
# Add the up/down commands (route Claude Code through the local LiteLLM proxy)
|
||||
cli.add_command(up)
|
||||
cli.add_command(down)
|
||||
# Add the model-groups command group (discover models your key can access)
|
||||
cli.add_command(model_groups)
|
||||
# Add the autoroute command group (QA auto-routing against your real proxy)
|
||||
cli.add_command(autoroute_group, name="autoroute")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
0
litellm/proxy/enterprise_billing/__init__.py
Normal file
0
litellm/proxy/enterprise_billing/__init__.py
Normal file
323
litellm/proxy/enterprise_billing/billing_metrics.py
Normal file
323
litellm/proxy/enterprise_billing/billing_metrics.py
Normal file
|
|
@ -0,0 +1,323 @@
|
|||
"""
|
||||
Push-based OTLP metering for enterprise litellm deployments.
|
||||
|
||||
Owns a dedicated OpenTelemetry meter provider and an OTLP/HTTP exporter
|
||||
authenticated to our global collector with a TLS client certificate. The
|
||||
collector front end terminates mutual TLS: the client certificate presented
|
||||
here is validated against our CA at the edge, and the verified subject is
|
||||
what identifies the deployment. It is intentionally isolated from the global
|
||||
meter provider so the customer's own OTEL metrics are untouched and ours
|
||||
never leak into their backend.
|
||||
|
||||
The deployment's identity rides on the TLS client certificate, not on the
|
||||
payload; the secret license key is never sent as an attribute or header.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Optional, Union
|
||||
|
||||
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
|
||||
from opentelemetry.metrics import Counter
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.middleware.billable_request_metrics_middleware import (
|
||||
BillableCategory,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import EnterpriseLicenseData
|
||||
|
||||
ENDPOINT_ENV = "LITELLM_BILLING_METRICS_ENDPOINT"
|
||||
CLIENT_CERT_ENV = "LITELLM_BILLING_METRICS_CLIENT_CERT"
|
||||
CLIENT_KEY_ENV = "LITELLM_BILLING_METRICS_CLIENT_KEY"
|
||||
CA_CERT_ENV = "LITELLM_BILLING_METRICS_CA_CERT"
|
||||
EXPORT_INTERVAL_ENV = "LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS"
|
||||
DEFAULT_EXPORT_INTERVAL_MS = 60_000
|
||||
SHUTDOWN_FLUSH_TIMEOUT_MS = 5_000
|
||||
_METRICS_PATH = "/v1/metrics"
|
||||
|
||||
# The cert env vars take a path or the PEM itself. Secret stores that inject
|
||||
# values as env content cannot mount them as files, so inline PEM is written out.
|
||||
_PEM_PREFIX = "-----BEGIN"
|
||||
_PEM_DIR_PREFIX = "litellm-billing-mtls-"
|
||||
_PEM_FILE_MODE = 0o600
|
||||
_CLIENT_CERT_FILENAME = "client.crt"
|
||||
_CLIENT_KEY_FILENAME = "client.key"
|
||||
_CA_CERT_FILENAME = "ca.crt"
|
||||
|
||||
METRIC_NAME = "litellm.enterprise.billable_requests"
|
||||
METER_NAME = "litellm.enterprise.billing"
|
||||
|
||||
AttributeValue = Union[str, int]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BillingMetricsConfig:
|
||||
endpoint: str
|
||||
client_cert_path: str
|
||||
client_key_path: str
|
||||
ca_cert_path: Optional[str]
|
||||
export_interval_ms: int
|
||||
litellm_version: str
|
||||
license_id: Optional[str]
|
||||
|
||||
|
||||
def _metrics_endpoint(endpoint: str) -> str:
|
||||
"""The OTLP/HTTP metric exporter wants the full URL including the signal path."""
|
||||
trimmed = endpoint.rstrip("/")
|
||||
return trimmed if trimmed.endswith(_METRICS_PATH) else f"{trimmed}{_METRICS_PATH}"
|
||||
|
||||
|
||||
def _resource_attributes(config: BillingMetricsConfig) -> dict[str, AttributeValue]:
|
||||
base: dict[str, AttributeValue] = {
|
||||
"service.name": "litellm-proxy",
|
||||
"litellm.version": config.litellm_version,
|
||||
}
|
||||
license_attr: dict[str, AttributeValue] = {"litellm.license.id": config.license_id} if config.license_id else {}
|
||||
return {**base, **license_attr}
|
||||
|
||||
|
||||
def _billable_attributes(
|
||||
category: BillableCategory, route: str, status_code: int, model_id: Optional[str]
|
||||
) -> dict[str, AttributeValue]:
|
||||
base: dict[str, AttributeValue] = {
|
||||
"litellm.endpoint.category": category.value,
|
||||
"http.route": route,
|
||||
"http.response.status_code": status_code,
|
||||
}
|
||||
model_attr: dict[str, AttributeValue] = {"litellm.model_id": model_id} if model_id else {}
|
||||
return {**base, **model_attr}
|
||||
|
||||
|
||||
def build_mtls_meter_provider(config: BillingMetricsConfig) -> MeterProvider:
|
||||
"""OTLP/HTTP exporter presenting a TLS client certificate.
|
||||
|
||||
The collector's load balancer terminates mutual TLS and validates the client
|
||||
certificate against our CA. Server verification uses the system trust store
|
||||
(the collector presents a public web-PKI certificate); ca_cert_path overrides
|
||||
it only for private/test collectors.
|
||||
"""
|
||||
exporter = OTLPMetricExporter(
|
||||
endpoint=_metrics_endpoint(config.endpoint),
|
||||
# None -> exporter falls back to the system trust store.
|
||||
certificate_file=config.ca_cert_path,
|
||||
client_certificate_file=config.client_cert_path,
|
||||
client_key_file=config.client_key_path,
|
||||
)
|
||||
reader = PeriodicExportingMetricReader(exporter, export_interval_millis=config.export_interval_ms)
|
||||
return MeterProvider(metric_readers=[reader], resource=Resource.create(_resource_attributes(config)))
|
||||
|
||||
|
||||
class BillingMetricsRecorder:
|
||||
"""Increments one OTLP counter per billable request. The meter provider is injected (see the factory)."""
|
||||
|
||||
def __init__(self, provider: MeterProvider) -> None:
|
||||
self._provider = provider
|
||||
self._counter: Counter = provider.get_meter(METER_NAME).create_counter(
|
||||
name=METRIC_NAME,
|
||||
unit="{request}",
|
||||
description="Count of 2xx HTTP requests to billable LLM/MCP/A2A endpoints",
|
||||
)
|
||||
|
||||
def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None:
|
||||
self._counter.add(1, _billable_attributes(category, route, status_code, model_id))
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Final flush + exporter-thread stop. Without this, up to one export
|
||||
interval of billable counts is dropped on every proxy restart."""
|
||||
self._provider.shutdown(timeout_millis=SHUTDOWN_FLUSH_TIMEOUT_MS)
|
||||
|
||||
|
||||
def _export_interval_ms() -> int:
|
||||
raw = os.getenv(EXPORT_INTERVAL_ENV)
|
||||
if raw is None:
|
||||
return DEFAULT_EXPORT_INTERVAL_MS
|
||||
try:
|
||||
return int(raw)
|
||||
except ValueError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Invalid %s=%r, falling back to %d ms", EXPORT_INTERVAL_ENV, raw, DEFAULT_EXPORT_INTERVAL_MS
|
||||
)
|
||||
return DEFAULT_EXPORT_INTERVAL_MS
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CredentialPaths:
|
||||
client_cert_path: str
|
||||
client_key_path: str
|
||||
ca_cert_path: Optional[str]
|
||||
|
||||
|
||||
def _is_pem_content(value: str) -> bool:
|
||||
return value.lstrip().startswith(_PEM_PREFIX)
|
||||
|
||||
|
||||
def _write_pem(directory: str, filename: str, pem: str) -> str:
|
||||
path = os.path.join(directory, filename)
|
||||
with open(path, "w", encoding="utf-8") as handle:
|
||||
handle.write(pem if pem.endswith("\n") else f"{pem}\n")
|
||||
os.chmod(path, _PEM_FILE_MODE)
|
||||
return path
|
||||
|
||||
|
||||
def _resolve_credential_paths(*, client_cert: str, client_key: str, ca_cert: Optional[str]) -> _CredentialPaths:
|
||||
"""
|
||||
Accept either a filesystem path or inline PEM content for each credential.
|
||||
|
||||
Secret stores that inject values as environment content rather than mounted
|
||||
files (ECS tasks reading AWS Secrets Manager, Cloud Run reading Secret
|
||||
Manager) can only deliver the certificate as a string. The OTLP exporter
|
||||
takes paths, so inline PEM is written to a private directory once, when the
|
||||
recorder is built. Raises OSError if that write fails; the caller disables
|
||||
metering rather than propagating.
|
||||
"""
|
||||
inline = tuple(value for value in (client_cert, client_key, ca_cert) if value and _is_pem_content(value))
|
||||
if not inline:
|
||||
return _CredentialPaths(client_cert, client_key, ca_cert)
|
||||
|
||||
# mkdtemp is 0o700, so the 0o600 key file it holds is unreachable by other users.
|
||||
directory = tempfile.mkdtemp(prefix=_PEM_DIR_PREFIX)
|
||||
return _CredentialPaths(
|
||||
client_cert_path=(
|
||||
_write_pem(directory, _CLIENT_CERT_FILENAME, client_cert) if _is_pem_content(client_cert) else client_cert
|
||||
),
|
||||
client_key_path=(
|
||||
_write_pem(directory, _CLIENT_KEY_FILENAME, client_key) if _is_pem_content(client_key) else client_key
|
||||
),
|
||||
ca_cert_path=(
|
||||
_write_pem(directory, _CA_CERT_FILENAME, ca_cert) if ca_cert and _is_pem_content(ca_cert) else ca_cert
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def load_billing_metrics_config(
|
||||
*, license_data: Optional["EnterpriseLicenseData"], litellm_version: str
|
||||
) -> Optional[BillingMetricsConfig]:
|
||||
endpoint = os.getenv(ENDPOINT_ENV)
|
||||
client_cert = os.getenv(CLIENT_CERT_ENV)
|
||||
client_key = os.getenv(CLIENT_KEY_ENV)
|
||||
# Optional: only for private/test collectors whose server cert is not on the
|
||||
# public web PKI. The production collector needs no CA override.
|
||||
ca_cert = os.getenv(CA_CERT_ENV)
|
||||
|
||||
missing = [
|
||||
name
|
||||
for name, value in (
|
||||
(ENDPOINT_ENV, endpoint),
|
||||
(CLIENT_CERT_ENV, client_cert),
|
||||
(CLIENT_KEY_ENV, client_key),
|
||||
)
|
||||
if not value
|
||||
]
|
||||
if not endpoint or not client_cert or not client_key:
|
||||
verbose_proxy_logger.warning(
|
||||
"Enterprise billing metrics disabled: licensed deployment missing config (%s)",
|
||||
", ".join(missing),
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
paths = _resolve_credential_paths(client_cert=client_cert, client_key=client_key, ca_cert=ca_cert)
|
||||
except OSError as exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"Enterprise billing metrics disabled: could not write inline certificate content to disk: %s", exc
|
||||
)
|
||||
return None
|
||||
|
||||
# Report the variable names, never their values. A value that is neither a
|
||||
# readable path nor recognizable PEM is still secret material, and this
|
||||
# warning would otherwise copy a client key straight into the proxy logs.
|
||||
unreadable = [
|
||||
env_name
|
||||
for env_name, path in (
|
||||
(CLIENT_CERT_ENV, paths.client_cert_path),
|
||||
(CLIENT_KEY_ENV, paths.client_key_path),
|
||||
(CA_CERT_ENV, paths.ca_cert_path),
|
||||
)
|
||||
if path and not os.path.isfile(path)
|
||||
]
|
||||
if unreadable:
|
||||
verbose_proxy_logger.warning(
|
||||
"Enterprise billing metrics disabled: %s did not resolve to a readable certificate file. "
|
||||
"Set each to a file path, or to inline PEM content beginning with '%s'.",
|
||||
", ".join(unreadable),
|
||||
_PEM_PREFIX,
|
||||
)
|
||||
return None
|
||||
|
||||
return BillingMetricsConfig(
|
||||
endpoint=endpoint,
|
||||
client_cert_path=paths.client_cert_path,
|
||||
client_key_path=paths.client_key_path,
|
||||
ca_cert_path=paths.ca_cert_path,
|
||||
export_interval_ms=_export_interval_ms(),
|
||||
litellm_version=litellm_version,
|
||||
license_id=(license_data or {}).get("user_id"),
|
||||
)
|
||||
|
||||
|
||||
class _ActiveRecorderRegistry:
|
||||
"""One-slot registry linking the factory-built recorder to the shutdown
|
||||
hook; the middleware instance holding the recorder is not reachable from
|
||||
proxy_shutdown_event."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._recorder: Optional[BillingMetricsRecorder] = None
|
||||
|
||||
def set(self, recorder: BillingMetricsRecorder) -> None:
|
||||
self._recorder = recorder
|
||||
|
||||
def pop(self) -> Optional[BillingMetricsRecorder]:
|
||||
recorder = self._recorder
|
||||
self._recorder = None
|
||||
return recorder
|
||||
|
||||
|
||||
_ACTIVE_RECORDER = _ActiveRecorderRegistry()
|
||||
|
||||
|
||||
def build_billing_metrics_recorder(
|
||||
*, premium: bool, license_data: Optional["EnterpriseLicenseData"], litellm_version: str
|
||||
) -> Optional[BillingMetricsRecorder]:
|
||||
"""Build the recorder, or None when the deployment is not licensed or metering is unconfigured."""
|
||||
if not premium:
|
||||
# Debug, not warning: unlicensed is the common case and a warning here
|
||||
# would be noise on every OSS proxy. Every other disable path warns.
|
||||
verbose_proxy_logger.debug("Enterprise billing metrics disabled: deployment is not licensed")
|
||||
return None
|
||||
|
||||
config = load_billing_metrics_config(license_data=license_data, litellm_version=litellm_version)
|
||||
if config is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
recorder = BillingMetricsRecorder(build_mtls_meter_provider(config))
|
||||
except Exception as exc: # noqa: BLE001 -- metering must never break proxy startup
|
||||
verbose_proxy_logger.warning("Enterprise billing metrics disabled: failed to initialize exporter: %s", exc)
|
||||
return None
|
||||
_ACTIVE_RECORDER.set(recorder)
|
||||
# The only positive signal that this component meters. Without it, a silent
|
||||
# return above is indistinguishable from a working exporter in the logs, and
|
||||
# a component that carries the cert but no license would look healthy.
|
||||
verbose_proxy_logger.info(
|
||||
"Enterprise billing metrics enabled: exporting to %s every %d ms",
|
||||
config.endpoint,
|
||||
config.export_interval_ms,
|
||||
)
|
||||
return recorder
|
||||
|
||||
|
||||
def shutdown_billing_metrics_recorder() -> None:
|
||||
"""Flush and stop the active recorder, if any. Idempotent; never raises."""
|
||||
recorder = _ACTIVE_RECORDER.pop()
|
||||
if recorder is None:
|
||||
return
|
||||
try:
|
||||
recorder.shutdown()
|
||||
except Exception as exc: # noqa: BLE001 -- shutdown must never block or fail proxy exit
|
||||
verbose_proxy_logger.warning("Enterprise billing metrics: final flush failed: %s", exc)
|
||||
|
|
@ -2238,7 +2238,10 @@ async def apply_guardrail(
|
|||
if litellm_logging_obj is not None:
|
||||
_patch_logging_obj_for_guardrail(litellm_logging_obj, request)
|
||||
|
||||
request_data: dict = {"messages": request.messages} if request.messages else {}
|
||||
request_data: dict = {
|
||||
**({"messages": request.messages} if request.messages is not None else {}),
|
||||
**({"metadata": request.metadata} if request.metadata is not None else {}),
|
||||
}
|
||||
_input_type = _resolve_guardrail_input_type(active_guardrail, request.input_type)
|
||||
guardrailed_inputs = await active_guardrail.apply_guardrail(
|
||||
inputs={"texts": [request.text]},
|
||||
|
|
|
|||
|
|
@ -0,0 +1,75 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.types.guardrails import (
|
||||
GuardrailEventHooks,
|
||||
Mode,
|
||||
SupportedGuardrailIntegrations,
|
||||
)
|
||||
|
||||
from .compresr import CompresrGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def _coerce_event_hook(
|
||||
mode: str | list[str] | Mode,
|
||||
) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:
|
||||
if isinstance(mode, Mode):
|
||||
return mode
|
||||
if isinstance(mode, list):
|
||||
return [GuardrailEventHooks(item) for item in mode]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object:
|
||||
if optional_params is not None:
|
||||
value = getattr(optional_params, attribute_name, None)
|
||||
if value is not None:
|
||||
return value
|
||||
return getattr(litellm_params, attribute_name, None)
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> CompresrGuardrail:
|
||||
import litellm
|
||||
|
||||
optional_params = getattr(litellm_params, "optional_params", None)
|
||||
|
||||
_callback = CompresrGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
model=litellm_params.model,
|
||||
target_compression_ratio=_get_optional_value(litellm_params, optional_params, "target_compression_ratio"),
|
||||
coarse=_get_optional_value(litellm_params, optional_params, "coarse"),
|
||||
min_chars_to_compress=_get_optional_value(litellm_params, optional_params, "min_chars_to_compress"),
|
||||
compress_tool_outputs=_get_optional_value(litellm_params, optional_params, "compress_tool_outputs"),
|
||||
compress_system=_get_optional_value(litellm_params, optional_params, "compress_system"),
|
||||
compress_history=_get_optional_value(litellm_params, optional_params, "compress_history"),
|
||||
compress_last_user=_get_optional_value(litellm_params, optional_params, "compress_last_user"),
|
||||
enable_retrieval=_get_optional_value(litellm_params, optional_params, "enable_retrieval"),
|
||||
max_bytes_per_call=_get_optional_value(litellm_params, optional_params, "max_bytes_per_call"),
|
||||
allow_bypass_header=_get_optional_value(litellm_params, optional_params, "allow_bypass_header"),
|
||||
dynamic=_get_optional_value(litellm_params, optional_params, "dynamic"),
|
||||
dynamic_min_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_min_ratio"),
|
||||
dynamic_max_ratio=_get_optional_value(litellm_params, optional_params, "dynamic_max_ratio"),
|
||||
compression_params=_get_optional_value(litellm_params, optional_params, "compression_params"),
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
event_hook=_coerce_event_hook(litellm_params.mode),
|
||||
default_on=litellm_params.default_on or False,
|
||||
unreachable_fallback=litellm_params.unreachable_fallback,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped
|
||||
_callback
|
||||
)
|
||||
return _callback
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.COMPRESR.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.COMPRESR.value: CompresrGuardrail,
|
||||
}
|
||||
1214
litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
Normal file
1214
litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -494,6 +494,7 @@ class InMemoryGuardrailHandler:
|
|||
guardrail_id=guardrail.get("guardrail_id"),
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
litellm_params=litellm_params,
|
||||
guardrail_info=guardrail.get("guardrail_info"),
|
||||
)
|
||||
|
||||
# store references to the guardrail in memory
|
||||
|
|
@ -612,6 +613,27 @@ class InMemoryGuardrailHandler:
|
|||
"""
|
||||
return self._sources.get(guardrail_id)
|
||||
|
||||
def list_config_guardrails(self) -> List[Guardrail]:
|
||||
"""
|
||||
List in-memory guardrails owned by config.yaml.
|
||||
|
||||
DB-sourced entries are excluded: a read surface that also queries the DB
|
||||
would double-count live ones, and a DB-sourced entry that's missing from
|
||||
the DB is stale (deleted on another pod, awaiting reconciliation here).
|
||||
"""
|
||||
return [g for gid, g in self.IN_MEMORY_GUARDRAILS.items() if self._sources.get(gid) == "config"]
|
||||
|
||||
def get_config_guardrail_by_id(self, guardrail_id: str) -> Optional[Guardrail]:
|
||||
"""
|
||||
Get a config-owned in-memory guardrail by its ID, or None.
|
||||
|
||||
Mirrors the fallback in get_guardrail_info: a DB-sourced in-memory entry
|
||||
that missed the DB lookup is stale and must not be surfaced.
|
||||
"""
|
||||
if self._sources.get(guardrail_id) != "config":
|
||||
return None
|
||||
return self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
|
||||
|
||||
def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]:
|
||||
"""
|
||||
Drop in-memory entries that originated from the DB but are no longer
|
||||
|
|
|
|||
|
|
@ -137,10 +137,26 @@ def _chart_from_metrics(metrics: Any) -> List[Dict[str, Any]]:
|
|||
return [{"date": d, "passed": v["passed"], "blocked": v["blocked"]} for d, v in sorted(chart_by_date.items())]
|
||||
|
||||
|
||||
def _get_guardrail_field(g: Any, field: str) -> Any:
|
||||
"""Read `field` off a guardrail whether it's a Prisma row (attr) or a dict/TypedDict (key)."""
|
||||
if isinstance(g, dict):
|
||||
return g.get(field)
|
||||
return getattr(g, field, None)
|
||||
|
||||
|
||||
def _to_dict(value: Any) -> Dict[str, Any]:
|
||||
"""Coerce a pydantic model (e.g. LitellmParams) / dict value into a plain dict."""
|
||||
if isinstance(value, BaseModel):
|
||||
return value.model_dump(exclude_none=True)
|
||||
if isinstance(value, dict):
|
||||
return value
|
||||
return {}
|
||||
|
||||
|
||||
def _get_guardrail_attrs(g: Any) -> tuple[Any, str]:
|
||||
"""Get (guardrail_id, display_name) from guardrail - handles Prisma model or dict."""
|
||||
gid = getattr(g, "guardrail_id", None) or (g.get("guardrail_id") if isinstance(g, dict) else None)
|
||||
name = getattr(g, "guardrail_name", None) or (g.get("guardrail_name") if isinstance(g, dict) else None)
|
||||
gid = _get_guardrail_field(g, "guardrail_id")
|
||||
name = _get_guardrail_field(g, "guardrail_name")
|
||||
return gid, (name or gid or "")
|
||||
|
||||
|
||||
|
|
@ -163,9 +179,9 @@ def _guardrail_overview_rows(
|
|||
break
|
||||
req, blocked = a["requests"], a["blocked"]
|
||||
fail_rate = (100.0 * blocked / req) if req else 0.0
|
||||
litellm_params = (g.litellm_params or {}) if isinstance(g.litellm_params, dict) else {}
|
||||
litellm_params = _to_dict(_get_guardrail_field(g, "litellm_params"))
|
||||
provider = str(litellm_params.get("guardrail", "Unknown"))
|
||||
guardrail_info = (g.guardrail_info or {}) if isinstance(g.guardrail_info, dict) else {}
|
||||
guardrail_info = _to_dict(_get_guardrail_field(g, "guardrail_info"))
|
||||
gtype = str(guardrail_info.get("type", "Guardrail"))
|
||||
prev_fail = 0.0
|
||||
for k in lookup_keys:
|
||||
|
|
@ -262,9 +278,15 @@ async def guardrails_usage_overview(
|
|||
end = end_date or now.strftime("%Y-%m-%d")
|
||||
start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
|
||||
try:
|
||||
# Guardrails from DB
|
||||
guardrails = await GuardrailsRepository(prisma_client).table.find_many()
|
||||
db_guardrails = await GuardrailsRepository(prisma_client).table.find_many()
|
||||
seen_ids = {gid for g in db_guardrails if (gid := _get_guardrail_field(g, "guardrail_id")) is not None}
|
||||
config_guardrails = [
|
||||
g for g in IN_MEMORY_GUARDRAIL_HANDLER.list_config_guardrails() if g.get("guardrail_id") not in seen_ids
|
||||
]
|
||||
guardrails: List[Any] = [*db_guardrails, *config_guardrails]
|
||||
|
||||
# Daily metrics in range
|
||||
metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many(
|
||||
|
|
@ -321,16 +343,18 @@ async def guardrails_usage_detail(
|
|||
end = end_date or now.strftime("%Y-%m-%d")
|
||||
start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
|
||||
|
||||
guardrail = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id})
|
||||
if not guardrail:
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
|
||||
guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique(where={"guardrail_id": guardrail_id})
|
||||
if guardrail is None:
|
||||
guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id)
|
||||
if guardrail is None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
raise HTTPException(status_code=404, detail="Guardrail not found")
|
||||
|
||||
# Metrics are keyed by logical name (from spend log metadata), not UUID
|
||||
logical_id = getattr(guardrail, "guardrail_name", None) or (
|
||||
guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None
|
||||
)
|
||||
logical_id = _get_guardrail_field(guardrail, "guardrail_name")
|
||||
metric_ids = [i for i in (logical_id, guardrail_id) if i]
|
||||
|
||||
metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many(
|
||||
|
|
@ -367,17 +391,9 @@ async def guardrails_usage_detail(
|
|||
{"date": d, "passed": v["passed"], "blocked": v["blocked"], "score": None}
|
||||
for d, v in sorted(ts_by_date.items())
|
||||
]
|
||||
_litellm_params = getattr(guardrail, "litellm_params", None) or (
|
||||
guardrail.get("litellm_params") if isinstance(guardrail, dict) else None
|
||||
)
|
||||
litellm_params = _litellm_params if isinstance(_litellm_params, dict) else {}
|
||||
_guardrail_info = getattr(guardrail, "guardrail_info", None) or (
|
||||
guardrail.get("guardrail_info") if isinstance(guardrail, dict) else None
|
||||
)
|
||||
guardrail_info = _guardrail_info if isinstance(_guardrail_info, dict) else {}
|
||||
_guardrail_name = getattr(guardrail, "guardrail_name", None) or (
|
||||
guardrail.get("guardrail_name") if isinstance(guardrail, dict) else None
|
||||
)
|
||||
litellm_params = _to_dict(_get_guardrail_field(guardrail, "litellm_params"))
|
||||
guardrail_info = _to_dict(_get_guardrail_field(guardrail, "guardrail_info"))
|
||||
_guardrail_name = _get_guardrail_field(guardrail, "guardrail_name")
|
||||
|
||||
return UsageDetailResponse(
|
||||
guardrail_id=guardrail_id,
|
||||
|
|
@ -548,11 +564,15 @@ async def guardrails_usage_logs(
|
|||
# Query by both so we match regardless of which was written.
|
||||
effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else []
|
||||
if guardrail_id:
|
||||
guardrail = await GuardrailsRepository(prisma_client).table.find_unique(
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
|
||||
guardrail: Any = await GuardrailsRepository(prisma_client).table.find_unique(
|
||||
where={"guardrail_id": guardrail_id}
|
||||
)
|
||||
if guardrail is None:
|
||||
guardrail = IN_MEMORY_GUARDRAIL_HANDLER.get_config_guardrail_by_id(guardrail_id=guardrail_id)
|
||||
if guardrail:
|
||||
logical_name = getattr(guardrail, "guardrail_name", None)
|
||||
logical_name = _get_guardrail_field(guardrail, "guardrail_name")
|
||||
if logical_name and logical_name not in effective_guardrail_ids:
|
||||
effective_guardrail_ids.append(logical_name)
|
||||
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm._uuid import uuid
|
|||
from litellm.constants import (
|
||||
LENGTH_OF_LITELLM_GENERATED_KEY,
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
MINIMUM_CUSTOM_KEY_LENGTH,
|
||||
UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
|
|
@ -1022,6 +1023,14 @@ async def _common_key_generation_helper(
|
|||
detail={"error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {_masked}"},
|
||||
)
|
||||
|
||||
if data.key is not None and len(data.key) < MINIMUM_CUSTOM_KEY_LENGTH:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Invalid key format. LiteLLM Virtual Key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long."
|
||||
},
|
||||
)
|
||||
|
||||
# check org key limits - done here to handle inheriting org id from team
|
||||
if data.organization_id is not None:
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
|
@ -1474,7 +1483,7 @@ async def generate_key_fn(
|
|||
Parameters:
|
||||
- duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||||
- key_alias: Optional[str] - User defined key alias
|
||||
- key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you.
|
||||
- key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you.
|
||||
- team_id: Optional[str] - The team id of the key
|
||||
- user_id: Optional[str] - The user id of the key
|
||||
- agent_id: Optional[str] - The agent id associated with the key.
|
||||
|
|
@ -1688,7 +1697,7 @@ async def generate_service_account_key_fn(
|
|||
Parameters:
|
||||
- duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||||
- key_alias: Optional[str] - User defined key alias
|
||||
- key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you.
|
||||
- key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you.
|
||||
- team_id: Optional[str] - The team id of the key
|
||||
- user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key
|
||||
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
|
||||
|
|
@ -4356,7 +4365,6 @@ async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
|
|||
if data and data.new_key is not None:
|
||||
# Reject custom key values if disabled by admin
|
||||
await _check_custom_key_allowed(data.new_key)
|
||||
new_token = data.new_key
|
||||
if not data.new_key.startswith("sk-"):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
|
|
@ -4364,6 +4372,12 @@ async def get_new_token(data: Optional[RegenerateKeyRequest]) -> str:
|
|||
"error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key."
|
||||
},
|
||||
)
|
||||
if len(data.new_key) < MINIMUM_CUSTOM_KEY_LENGTH:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail={"error": f"New key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long."},
|
||||
)
|
||||
new_token = data.new_key
|
||||
else:
|
||||
new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}"
|
||||
return new_token
|
||||
|
|
@ -4470,7 +4484,7 @@ async def _execute_virtual_key_regeneration(
|
|||
|
||||
new_token = await get_new_token(data=data)
|
||||
new_token_hash = hash_token(new_token)
|
||||
new_token_key_name = f"sk-...{new_token[-4:]}"
|
||||
new_token_key_name = abbreviate_api_key(api_key=new_token)
|
||||
update_data = {"token": new_token_hash, "key_name": new_token_key_name}
|
||||
|
||||
non_default_values = {}
|
||||
|
|
@ -4550,7 +4564,7 @@ async def regenerate_key_fn(
|
|||
- data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update
|
||||
- key: Optional[str] - The key to regenerate.
|
||||
- new_master_key: Optional[str] - The new master key to use, if key is the master key.
|
||||
- new_key: Optional[str] - The new key to use, if key is not the master key. If both set, new_master_key will be used.
|
||||
- new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used.
|
||||
- key_alias: Optional[str] - User-friendly key alias
|
||||
- user_id: Optional[str] - User ID associated with key
|
||||
- team_id: Optional[str] - Team ID associated with key
|
||||
|
|
|
|||
|
|
@ -1621,6 +1621,7 @@ async def update_team(
|
|||
- allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team.
|
||||
- model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit per model for this team. Example: {"gpt-4": 100, "gpt-3.5-turbo": 200}
|
||||
- model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit per model for this team. Example: {"gpt-4": 10000, "gpt-3.5-turbo": 20000}
|
||||
- mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team.
|
||||
Example - update team TPM Limit
|
||||
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
|
||||
- secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview)
|
||||
|
|
|
|||
|
|
@ -241,13 +241,34 @@ def _check_cli_sso_start_rate_limit(
|
|||
|
||||
|
||||
def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict:
|
||||
if isinstance(login_id, str) and login_id.startswith("sk-"):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
"Your litellm CLI is out of date and uses a login flow this proxy no longer supports. "
|
||||
"Upgrade it with `pip install -U 'litellm[proxy]'` and run `litellm-proxy login` again."
|
||||
),
|
||||
)
|
||||
if not _is_valid_cli_sso_login_id(login_id):
|
||||
raise HTTPException(status_code=400, detail="Invalid CLI login session")
|
||||
raise HTTPException(status_code=400, detail="Invalid CLI login session id")
|
||||
|
||||
cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id))
|
||||
flow = cache.get_cache(key=cache_key)
|
||||
if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
|
||||
raise HTTPException(status_code=400, detail="Invalid CLI login session")
|
||||
verbose_proxy_logger.warning(
|
||||
"CLI SSO login session not found in cache for login_id=%s. If the proxy runs multiple replicas, "
|
||||
"a shared Redis cache (enable_redis_auth_cache: true) is required for CLI login to work.",
|
||||
login_id,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
"CLI login session not found or expired. Run `litellm-proxy login` again. "
|
||||
"If this happens immediately after starting a login, the proxy is likely running multiple "
|
||||
"replicas without a shared cache; configure Redis with `enable_redis_auth_cache: true` "
|
||||
"so every replica can see the login session."
|
||||
),
|
||||
)
|
||||
return flow
|
||||
|
||||
|
||||
|
|
|
|||
234
litellm/proxy/middleware/billable_request_metrics_middleware.py
Normal file
234
litellm/proxy/middleware/billable_request_metrics_middleware.py
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
"""
|
||||
Counts billable HTTP requests on enterprise deployments.
|
||||
|
||||
A billable request is an inbound request to an LLM inference, MCP, or A2A
|
||||
endpoint that returns a 2xx status. The actual export happens in an injected
|
||||
recorder (see litellm.proxy.enterprise_billing.billing_metrics); when no
|
||||
recorder is injected (non-enterprise, or metering misconfigured) this
|
||||
middleware is a transparent pass-through.
|
||||
"""
|
||||
|
||||
import re
|
||||
import threading
|
||||
from enum import Enum
|
||||
from typing import Callable, Optional, Protocol, Sequence, runtime_checkable
|
||||
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import LiteLLMRoutes
|
||||
|
||||
|
||||
class BillableCategory(str, Enum):
|
||||
LLM = "llm"
|
||||
MCP = "mcp"
|
||||
A2A = "a2a"
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class BillingRecorder(Protocol):
|
||||
def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None: ...
|
||||
|
||||
|
||||
_MODEL_ID_HEADER = b"x-litellm-model-id"
|
||||
|
||||
# Ordered: a longer suffix that shares an ending with a shorter one must come
|
||||
# first, e.g. "/chat/completions" before "/completions". This is the POST
|
||||
# inference surface that writes a SpendLogs row on success, so the exported
|
||||
# count lines up with the admin UI usage page for inference traffic. Billing
|
||||
# is a deliberate lower bound on SpendLogs rows: management writes that also
|
||||
# log (batch/file/fine-tuning creation, interaction cancel) and non-POST calls
|
||||
# that log (passthrough reads) never bill, so drift only ever undercounts.
|
||||
_LLM_ROUTE_SUFFIXES: tuple[str, ...] = (
|
||||
"/chat/completions",
|
||||
"/completions",
|
||||
"/embeddings",
|
||||
"/responses",
|
||||
"/rerank",
|
||||
"/moderations",
|
||||
"/images/generations",
|
||||
"/images/edits",
|
||||
"/images/variations",
|
||||
"/audio/transcriptions",
|
||||
"/audio/translations",
|
||||
"/audio/speech",
|
||||
"/videos", # create; GET list is excluded by the POST gate
|
||||
"/remix", # /v1/videos/{id}/remix
|
||||
"/ocr",
|
||||
"/search", # /v1/search and /v1/vector_stores/{id}/search
|
||||
"/rag/query",
|
||||
"/rag/ingest",
|
||||
":generateContent", # Gemini-native /v1beta/models/{model}:generateContent
|
||||
":streamGenerateContent",
|
||||
)
|
||||
|
||||
# Exact paths only: a suffix match would also catch non-inference resources that
|
||||
# share the ending, e.g. the OpenAI Assistants route /v1/threads/{id}/messages
|
||||
# writes no SpendLogs row and must not bill, unlike Anthropic /v1/messages.
|
||||
_LLM_ROUTE_EXACT: tuple[str, ...] = (
|
||||
"/v1/messages",
|
||||
"/interactions", # Google Interactions create; /{id} reads and /cancel do not match
|
||||
"/v1beta/interactions",
|
||||
)
|
||||
|
||||
# Provider passthrough prefixes (e.g. /bedrock/..., /vertex-ai/...) carry real
|
||||
# inference calls that write SpendLogs rows, so they bill. Anchored to the
|
||||
# routes enum so new providers are picked up without touching this module.
|
||||
# /langfuse forwards observability traffic, not inference: it writes no
|
||||
# SpendLogs row and must not bill.
|
||||
_NON_BILLABLE_PASSTHROUGH_PREFIXES = frozenset({"/langfuse"})
|
||||
_PASSTHROUGH_PREFIXES: tuple[str, ...] = tuple(
|
||||
prefix
|
||||
for prefix in LiteLLMRoutes.mapped_pass_through_routes.value
|
||||
if prefix not in _NON_BILLABLE_PASSTHROUGH_PREFIXES
|
||||
)
|
||||
|
||||
|
||||
def _classify_llm_route(path: str) -> Optional[str]:
|
||||
exact_match = next((route for route in _LLM_ROUTE_EXACT if path == route), None)
|
||||
if exact_match is not None:
|
||||
return exact_match
|
||||
suffix_match = next((suffix for suffix in _LLM_ROUTE_SUFFIXES if path == suffix or path.endswith(suffix)), None)
|
||||
if suffix_match is not None:
|
||||
return suffix_match
|
||||
# Deep passthrough paths only: the bare prefix itself is not an inference call.
|
||||
return next((prefix for prefix in _PASSTHROUGH_PREFIXES if path.startswith(f"{prefix}/")), None)
|
||||
|
||||
|
||||
_MCP_MANAGEMENT_PREFIX = "/v1/mcp"
|
||||
_MCP_DYNAMIC_TRANSPORT = re.compile(r"/(?:toolset/)?[^/]+/mcp")
|
||||
# The REST wrapper's tool-call endpoint executes a tool and fires the same MCP
|
||||
# spend logging as the /mcp transport; its list/test siblings do not bill.
|
||||
_MCP_REST_TOOL_CALL = "/mcp-rest/tools/call"
|
||||
|
||||
_A2A_INVOKE_SUFFIX = "/message/send"
|
||||
_A2A_TRANSPORT_PREFIXES: tuple[str, ...] = ("/v1/a2a/", "/a2a/")
|
||||
# Bare POST /a2a/{agent_id} carries the JSON-RPC method in the body, not the
|
||||
# path. Only message/send and message/stream write a SpendLogs row there; the
|
||||
# task RPCs (tasks/get, tasks/cancel, tasks/pushNotificationConfig/*, ...) are
|
||||
# forwarded upstream and write none. A path-only classifier cannot separate
|
||||
# them, so the bare route does not bill: counting a task RPC would overcount,
|
||||
# while missing a bare-path message/send only undercounts, and undercounting is
|
||||
# the sole direction this metric is allowed to drift. The /mcp transport is
|
||||
# method-agnostic by contrast because its list path logs a SpendLogs row too.
|
||||
|
||||
|
||||
def _classify_mcp_route(path: str) -> Optional[str]:
|
||||
if path == _MCP_MANAGEMENT_PREFIX or path.startswith(f"{_MCP_MANAGEMENT_PREFIX}/"):
|
||||
return None
|
||||
if path == "/mcp" or path.startswith("/mcp/"):
|
||||
return "/mcp"
|
||||
if path == _MCP_REST_TOOL_CALL:
|
||||
return "/mcp"
|
||||
if _MCP_DYNAMIC_TRANSPORT.fullmatch(path) is not None:
|
||||
return "/mcp"
|
||||
return None
|
||||
|
||||
|
||||
def _classify_a2a_route(path: str) -> Optional[str]:
|
||||
if path.endswith(_A2A_INVOKE_SUFFIX) and any(path.startswith(prefix) for prefix in _A2A_TRANSPORT_PREFIXES):
|
||||
return "/a2a"
|
||||
return None
|
||||
|
||||
|
||||
def classify_billable_request(path: str, method: str = "POST") -> Optional[tuple[BillableCategory, str]]:
|
||||
"""Map a request path to its (category, normalized route), or None if not billable."""
|
||||
normalized = path.rstrip("/") or "/"
|
||||
|
||||
mcp_route = _classify_mcp_route(normalized)
|
||||
if mcp_route is not None:
|
||||
return (BillableCategory.MCP, mcp_route)
|
||||
|
||||
a2a_route = _classify_a2a_route(normalized)
|
||||
if a2a_route is not None:
|
||||
return (BillableCategory.A2A, a2a_route)
|
||||
|
||||
# POST-only is a conservative gate: non-POST calls can still write a
|
||||
# SpendLogs row (passthrough reads, resource GETs) but must not bill, so
|
||||
# any classifier-vs-dashboard mismatch is an undercount, never an overcount.
|
||||
if method.upper() != "POST":
|
||||
return None
|
||||
|
||||
llm_route = _classify_llm_route(normalized)
|
||||
if llm_route is not None:
|
||||
return (BillableCategory.LLM, llm_route)
|
||||
return None
|
||||
|
||||
|
||||
def _extract_model_id(headers: Sequence[tuple[bytes, bytes]]) -> Optional[str]:
|
||||
return next(
|
||||
(value.decode("latin-1") for name, value in headers if name.lower() == _MODEL_ID_HEADER and value),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
class BillableRequestMetricsMiddleware:
|
||||
"""
|
||||
Pure ASGI middleware that records one billable request per 2xx response to a
|
||||
billable endpoint. Modeled on InFlightRequestsMiddleware: it wraps `send`,
|
||||
reads the final status and the x-litellm-model-id header off the
|
||||
`http.response.start` message, and never blocks or fails the request path.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app: ASGIApp,
|
||||
recorder: Optional[BillingRecorder] = None,
|
||||
recorder_factory: Optional[Callable[[], Optional[BillingRecorder]]] = None,
|
||||
) -> None:
|
||||
self.app = app
|
||||
self.recorder = recorder
|
||||
# The factory defers recorder construction to the first request, AFTER the
|
||||
# startup event has loaded the YAML config's environment_variables (license
|
||||
# and cert env vars). Building at import time captured recorder=None for
|
||||
# deployments configured that way. Resolved exactly once; the result
|
||||
# (including None) is cached.
|
||||
self._recorder_factory = recorder_factory
|
||||
self._resolved = recorder_factory is None
|
||||
self._resolve_lock = threading.Lock()
|
||||
|
||||
def _resolve_recorder(self) -> Optional[BillingRecorder]:
|
||||
if self._resolved:
|
||||
return self.recorder
|
||||
# The lock keeps concurrent first requests from each building their own
|
||||
# MeterProvider (and leaking its background exporter thread).
|
||||
with self._resolve_lock:
|
||||
if not self._resolved:
|
||||
factory = self._recorder_factory
|
||||
self.recorder = factory() if factory is not None else self.recorder
|
||||
self._resolved = True
|
||||
return self.recorder
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
recorder = self._resolve_recorder()
|
||||
if recorder is None:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
classification = classify_billable_request(scope.get("path", ""), scope.get("method", "POST"))
|
||||
if classification is None:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
category, route = classification
|
||||
status_code = 0
|
||||
model_id: Optional[str] = None
|
||||
|
||||
async def send_wrapper(message: Message) -> None:
|
||||
nonlocal status_code, model_id
|
||||
if message["type"] == "http.response.start":
|
||||
status_code = message["status"]
|
||||
model_id = _extract_model_id(message.get("headers", []))
|
||||
await send(message)
|
||||
|
||||
await self.app(scope, receive, send_wrapper)
|
||||
|
||||
if 200 <= status_code < 300:
|
||||
try:
|
||||
recorder.record(category=category, route=route, status_code=status_code, model_id=model_id)
|
||||
except Exception: # noqa: BLE001 -- metering must never fail a request that was already served
|
||||
verbose_proxy_logger.warning("billable request metering failed for %s", route, exc_info=True)
|
||||
|
|
@ -438,10 +438,30 @@ from litellm.proxy.management_helpers.audit_logs import (
|
|||
create_object_audit_log,
|
||||
)
|
||||
from litellm.proxy.memory.memory_endpoints import router as memory_router
|
||||
from litellm.proxy.middleware.billable_request_metrics_middleware import (
|
||||
BillableRequestMetricsMiddleware,
|
||||
BillingRecorder,
|
||||
)
|
||||
from litellm.proxy.plugin_routes import (
|
||||
router as plugin_router,
|
||||
register_plugins_from_config,
|
||||
)
|
||||
from litellm.proxy.plugin_routes import (
|
||||
router as plugin_router,
|
||||
)
|
||||
|
||||
try:
|
||||
from litellm.proxy.enterprise_billing.billing_metrics import (
|
||||
build_billing_metrics_recorder as _build_billing_metrics_recorder,
|
||||
)
|
||||
from litellm.proxy.enterprise_billing.billing_metrics import (
|
||||
shutdown_billing_metrics_recorder as _shutdown_billing_metrics_recorder,
|
||||
)
|
||||
|
||||
build_billing_metrics_recorder: Optional[Callable[..., Optional[BillingRecorder]]] = _build_billing_metrics_recorder
|
||||
shutdown_billing_metrics_recorder: Optional[Callable[[], None]] = _shutdown_billing_metrics_recorder
|
||||
except ImportError:
|
||||
build_billing_metrics_recorder = None
|
||||
shutdown_billing_metrics_recorder = None
|
||||
from litellm.proxy.middleware.in_flight_requests_middleware import (
|
||||
InFlightRequestsMiddleware,
|
||||
)
|
||||
|
|
@ -461,13 +481,11 @@ from litellm.proxy.openai_files_endpoints.files_endpoints import (
|
|||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
passthrough_endpoint_router,
|
||||
vertex_ai_live_websocket_passthrough,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
router as llm_passthrough_router,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
vertex_ai_live_websocket_passthrough,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
initialize_pass_through_endpoints,
|
||||
)
|
||||
|
|
@ -552,21 +570,19 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
|
|||
from litellm.types.realtime import RealtimeQueryParams
|
||||
from litellm.types.router import (
|
||||
DeploymentTypedDict,
|
||||
)
|
||||
from litellm.types.router import ModelInfo as RouterModelInfo
|
||||
from litellm.types.router import (
|
||||
RouterGeneralSettings,
|
||||
RoutingPlugin,
|
||||
SearchToolTypedDict,
|
||||
updateDeployment,
|
||||
)
|
||||
from litellm.types.router import ModelInfo as RouterModelInfo
|
||||
from litellm.types.scheduler import DefaultPriorities
|
||||
from litellm.types.secret_managers.main import (
|
||||
KeyManagementSettings,
|
||||
KeyManagementSystem,
|
||||
)
|
||||
from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer
|
||||
from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload
|
||||
from litellm.types.utils import ModelInfo as ModelMapInfo
|
||||
from litellm.types.utils import RawRequestTypedDict, StandardLoggingPayload
|
||||
from litellm.utils import _add_custom_logger_callback_to_specific_event
|
||||
|
||||
try:
|
||||
|
|
@ -767,6 +783,11 @@ async def proxy_shutdown_event():
|
|||
if db_writer_client is not None:
|
||||
await db_writer_client.close() # type: ignore[reportGeneralTypeIssues]
|
||||
|
||||
# final flush of billable-request counts: without it, up to one export
|
||||
# interval of enterprise billing data is dropped on every restart
|
||||
if shutdown_billing_metrics_recorder is not None:
|
||||
shutdown_billing_metrics_recorder()
|
||||
|
||||
# flush remaining langfuse logs
|
||||
if "langfuse" in litellm.success_callback:
|
||||
try:
|
||||
|
|
@ -972,11 +993,11 @@ async def proxy_startup_event(app: FastAPI):
|
|||
if is_otel_v2_enabled():
|
||||
from opentelemetry import trace as _otel_trace
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import _in_memory_loggers
|
||||
from litellm.integrations.otel.logger import (
|
||||
OpenTelemetryV2,
|
||||
publish_global_otel_v2_provider,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import _in_memory_loggers
|
||||
|
||||
registered = open_telemetry_logger if isinstance(open_telemetry_logger, OpenTelemetryV2) else None
|
||||
publish_global_otel_v2_provider(
|
||||
|
|
@ -1780,6 +1801,31 @@ app.add_middleware(
|
|||
)
|
||||
|
||||
app.add_middleware(PrometheusAuthMiddleware)
|
||||
# Added before InFlightRequestsMiddleware so it nests *inside* it: Starlette
|
||||
# makes the last-added middleware outermost. The billable count is recorded
|
||||
# after the inner app returns, so if this sat outside the in-flight tracker a
|
||||
# request could be counted as drained while its record() had not yet run, and
|
||||
# proxy_shutdown_event could flush and stop the exporter underneath it.
|
||||
app.add_middleware(
|
||||
BillableRequestMetricsMiddleware,
|
||||
# Factory, not an instance: the recorder is resolved on the first request so
|
||||
# it sees premium_user and the billing env vars AFTER proxy_startup_event has
|
||||
# loaded the YAML config's environment_variables. Building it here at import
|
||||
# time would permanently capture recorder=None for YAML-configured
|
||||
# deployments. The lambda reads the module globals at call time.
|
||||
recorder_factory=lambda: (
|
||||
build_billing_metrics_recorder(
|
||||
premium=premium_user,
|
||||
# Read from the license check, not the premium_user_data module
|
||||
# global: that global is bound once at import and goes stale when
|
||||
# the license arrives via the YAML config's environment_variables.
|
||||
license_data=_license_check.airgapped_license_data,
|
||||
litellm_version=version,
|
||||
)
|
||||
if build_billing_metrics_recorder is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
app.add_middleware(InFlightRequestsMiddleware)
|
||||
app.add_middleware(SecurityHeadersMiddleware)
|
||||
|
||||
|
|
@ -3660,6 +3706,45 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache:
|
|||
litellm_config_cache.redis_cache = redis_cache
|
||||
|
||||
|
||||
def resolve_complexity_router_plugins(
|
||||
model_name: str,
|
||||
complexity_router_config: dict,
|
||||
config_file_path: str | None,
|
||||
) -> None:
|
||||
"""
|
||||
Resolves `complexity_router_config["plugins"]` dotted-path strings to live
|
||||
instances via `get_instance_fn` (the same convention `litellm_settings.callbacks`
|
||||
uses), in place. Raises at config-load time if a path resolves to something that
|
||||
doesn't implement `RoutingPlugin`, rather than deferring to a confusing
|
||||
`AttributeError` on the first request that reaches the plugin pipeline.
|
||||
"""
|
||||
plugin_paths = complexity_router_config.get("plugins")
|
||||
if not isinstance(plugin_paths, list):
|
||||
return
|
||||
|
||||
resolved_plugins = [
|
||||
get_instance_fn(value=plugin_path, config_file_path=config_file_path)
|
||||
if isinstance(plugin_path, str)
|
||||
else plugin_path
|
||||
for plugin_path in plugin_paths
|
||||
]
|
||||
for plugin_path, resolved_plugin in zip(plugin_paths, resolved_plugins):
|
||||
# `@runtime_checkable` only checks that `run` exists as an attribute, not that
|
||||
# it's a coroutine function -- a synchronous `def run(self, context)` would pass
|
||||
# isinstance() here and only fail at request time with a confusing `TypeError:
|
||||
# object RoutingContext can't be used in 'await' expression`.
|
||||
if not isinstance(resolved_plugin, RoutingPlugin) or not inspect.iscoroutinefunction(
|
||||
getattr(resolved_plugin, "run", None)
|
||||
):
|
||||
raise ValueError(
|
||||
f"complexity_router_config.plugins entry {plugin_path!r} on model {model_name!r} "
|
||||
f"resolved to {resolved_plugin!r}, which does not implement the RoutingPlugin "
|
||||
"interface (an async `run(context)` method). Fix the referenced module before "
|
||||
"starting the proxy."
|
||||
)
|
||||
complexity_router_config["plugins"] = resolved_plugins
|
||||
|
||||
|
||||
class ProxyConfig:
|
||||
"""
|
||||
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
|
||||
|
|
@ -4720,6 +4805,13 @@ class ProxyConfig:
|
|||
for k, v in model["litellm_params"].items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
model["litellm_params"][k] = get_secret(v)
|
||||
complexity_router_config = model["litellm_params"].get("complexity_router_config")
|
||||
if isinstance(complexity_router_config, dict):
|
||||
resolve_complexity_router_plugins(
|
||||
model_name=model.get("model_name", ""),
|
||||
complexity_router_config=complexity_router_config,
|
||||
config_file_path=config_file_path,
|
||||
)
|
||||
print(f"\033[32m {model.get('model_name', '')}\033[0m") # noqa: T201
|
||||
litellm_model_name = model["litellm_params"]["model"]
|
||||
litellm_model_api_base = model["litellm_params"].get("api_base", None)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from fastapi import HTTPException, status
|
|||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router_utils.common_utils import _is_proxy_admin_request
|
||||
|
||||
# Router-internal mock_testing_* flag names — kept in sync with
|
||||
# ``litellm.types.router.MockRouterTestingParams`` by the test
|
||||
|
|
@ -363,6 +364,7 @@ async def route_request(
|
|||
|
||||
team_id = get_team_id_from_data(data)
|
||||
router_model_names = llm_router.model_names if llm_router is not None else []
|
||||
is_proxy_admin_without_team = team_id is None and _is_proxy_admin_request(data)
|
||||
|
||||
# Preprocess Google GenAI generate content requests
|
||||
if route_type in ["agenerate_content", "agenerate_content_stream"]:
|
||||
|
|
@ -517,6 +519,13 @@ async def route_request(
|
|||
data["model"] = team_model_name
|
||||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
||||
elif (
|
||||
is_proxy_admin_without_team
|
||||
and data["model"] not in router_model_names
|
||||
and data["model"] in llm_router.team_public_model_names
|
||||
):
|
||||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
||||
elif data["model"] in router_model_names or llm_router.has_model_id(data["model"]):
|
||||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
||||
|
|
|
|||
|
|
@ -833,7 +833,7 @@ class ProxyLogging:
|
|||
def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List:
|
||||
if dynamic_success_callbacks is None:
|
||||
return list(global_callbacks)
|
||||
return list(set(dynamic_success_callbacks + global_callbacks))
|
||||
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
|
||||
|
||||
def _parse_pre_mcp_call_hook_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from typing import (
|
|||
AsyncGenerator,
|
||||
Callable,
|
||||
Dict,
|
||||
FrozenSet,
|
||||
Generator,
|
||||
List,
|
||||
Literal,
|
||||
|
|
@ -108,6 +109,7 @@ from litellm.router_utils.clientside_credential_handler import (
|
|||
is_clientside_credential,
|
||||
)
|
||||
from litellm.router_utils.common_utils import (
|
||||
_is_proxy_admin_request,
|
||||
filter_team_based_models,
|
||||
filter_web_search_deployments,
|
||||
)
|
||||
|
|
@ -494,6 +496,7 @@ class Router:
|
|||
self.model_name_to_deployment_indices: Dict[str, List[int]] = {}
|
||||
# Maps (team_id, team_public_model_name) -> list of indices in model_list
|
||||
self.team_model_to_deployment_indices: Dict[Tuple[str, str], List[int]] = {}
|
||||
self.team_public_model_names: FrozenSet[str] = frozenset()
|
||||
|
||||
# Initialize cache attributes that ``_invalidate_model_group_info_cache``
|
||||
# touches *before* the first ``set_model_list`` below (which calls
|
||||
|
|
@ -2983,7 +2986,7 @@ class Router:
|
|||
# here before it's wiped below, instead of relying on that attempt's
|
||||
# (possibly still-pending) failure event to do it.
|
||||
refund_stale_reservation_before_retry(self.cache, kwargs)
|
||||
set_io_token_rate_limit_request_kwargs(kwargs)
|
||||
set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=deployment_has_io_token_limits(deployment))
|
||||
|
||||
## DEPLOYMENT-LEVEL TAGS
|
||||
deployment_tags = deployment.get("litellm_params", {}).get("tags")
|
||||
|
|
@ -7284,6 +7287,14 @@ class Router:
|
|||
raise e
|
||||
return returned_healthy_deployments
|
||||
|
||||
@staticmethod
|
||||
def _json_default_stable_id(value: object) -> str:
|
||||
"""json.dumps default= for _generate_model_id: plain str() on an arbitrary
|
||||
object (e.g. a RoutingPlugin instance) falls back to object.__repr__'s
|
||||
`<module.Class object at 0x...>`, so the hash -- and deployment id -- would
|
||||
change every restart. Use the class name instead, stable across restarts."""
|
||||
return f"{type(value).__module__}.{type(value).__qualname__}"
|
||||
|
||||
def _generate_model_id(self, model_group: str, litellm_params: dict):
|
||||
"""
|
||||
Helper function to consistently generate the same id for a deployment
|
||||
|
|
@ -7299,14 +7310,14 @@ class Router:
|
|||
if isinstance(k, str):
|
||||
parts.append(k)
|
||||
elif isinstance(k, dict):
|
||||
parts.append(json.dumps(k))
|
||||
parts.append(json.dumps(k, default=self._json_default_stable_id))
|
||||
else:
|
||||
parts.append(str(k))
|
||||
|
||||
if isinstance(v, str):
|
||||
parts.append(v)
|
||||
elif isinstance(v, dict):
|
||||
parts.append(json.dumps(v))
|
||||
parts.append(json.dumps(v, default=self._json_default_stable_id))
|
||||
else:
|
||||
parts.append(str(v))
|
||||
|
||||
|
|
@ -7792,6 +7803,7 @@ class Router:
|
|||
self.model_id_to_deployment_index_map = {} # Reset the index
|
||||
self.model_name_to_deployment_indices = {} # Reset the model_name index
|
||||
self.team_model_to_deployment_indices = {} # Reset the team_model index
|
||||
self.team_public_model_names = frozenset()
|
||||
# Reset per-strategy router registries so hot-reload doesn't leave
|
||||
# stale routers pointing at the old model_list.
|
||||
self.quality_routers = {}
|
||||
|
|
@ -8143,6 +8155,9 @@ class Router:
|
|||
self.team_model_to_deployment_indices[key] = updated_indices
|
||||
else:
|
||||
del self.team_model_to_deployment_indices[key]
|
||||
self.team_public_model_names = frozenset(
|
||||
public_model_name for _, public_model_name in self.team_model_to_deployment_indices
|
||||
)
|
||||
|
||||
def _update_team_model_index(self, model: dict, idx: int) -> None:
|
||||
"""
|
||||
|
|
@ -8156,6 +8171,7 @@ class Router:
|
|||
team_public_model_name = (model.get("model_info") or {}).get("team_public_model_name")
|
||||
if team_id and team_public_model_name:
|
||||
key = (team_id, team_public_model_name)
|
||||
self.team_public_model_names = self.team_public_model_names | frozenset({team_public_model_name})
|
||||
if key not in self.team_model_to_deployment_indices:
|
||||
self.team_model_to_deployment_indices[key] = []
|
||||
if idx not in self.team_model_to_deployment_indices[key]:
|
||||
|
|
@ -9110,6 +9126,7 @@ class Router:
|
|||
"""
|
||||
self.model_name_to_deployment_indices.clear()
|
||||
self.team_model_to_deployment_indices.clear()
|
||||
self.team_public_model_names = frozenset()
|
||||
|
||||
for idx, model in enumerate(model_list):
|
||||
model_name = model.get("model_name")
|
||||
|
|
@ -10018,7 +10035,10 @@ class Router:
|
|||
return [m for m in self.model_list if m["litellm_params"]["model"] == model]
|
||||
|
||||
def _try_early_resolve_deployments_for_model_not_in_names(
|
||||
self, model: str, request_team_id: Optional[str]
|
||||
self,
|
||||
model: str,
|
||||
request_team_id: Optional[str],
|
||||
include_team_models: bool = False,
|
||||
) -> Optional[Tuple[str, Union[List, Dict]]]:
|
||||
"""
|
||||
When ``model`` is not in ``self.model_names``, try team routes, pattern routes,
|
||||
|
|
@ -10033,6 +10053,30 @@ class Router:
|
|||
team_deployments = self._get_all_deployments(model_name=model, team_id=request_team_id)
|
||||
if team_deployments:
|
||||
return model, team_deployments
|
||||
elif include_team_models:
|
||||
team_deployments = [
|
||||
self.model_list[index]
|
||||
for (_, public_model_name), indices in self.team_model_to_deployment_indices.items()
|
||||
if public_model_name == model
|
||||
for index in indices
|
||||
]
|
||||
team_ids = {
|
||||
team_id
|
||||
for deployment in team_deployments
|
||||
for team_id in [(deployment.get("model_info") or {}).get("team_id")]
|
||||
if team_id is not None
|
||||
}
|
||||
if len(team_ids) > 1:
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
f"Model name '{model}' matches deployments from multiple teams. "
|
||||
"Specify the deployment ID directly to disambiguate."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
if team_deployments:
|
||||
return model, team_deployments
|
||||
|
||||
pattern_deployments = self.pattern_router.get_deployments_by_pattern(
|
||||
model=model,
|
||||
|
|
@ -10097,7 +10141,11 @@ class Router:
|
|||
if _model_from_alias is not None:
|
||||
model = _model_from_alias
|
||||
|
||||
early = self._try_early_resolve_deployments_for_model_not_in_names(model=model, request_team_id=request_team_id)
|
||||
early = self._try_early_resolve_deployments_for_model_not_in_names(
|
||||
model=model,
|
||||
request_team_id=request_team_id,
|
||||
include_team_models=_is_proxy_admin_request(request_kwargs),
|
||||
)
|
||||
if early is not None:
|
||||
return early
|
||||
|
||||
|
|
|
|||
|
|
@ -98,9 +98,9 @@ def _sanitize_user_api_key_auth(auth: Any) -> Any:
|
|||
return auth
|
||||
|
||||
|
||||
def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
def _classifier_call_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]:
|
||||
if not metadata:
|
||||
return metadata
|
||||
return {}
|
||||
return {
|
||||
k: _sanitize_user_api_key_auth(v) if k == "user_api_key_auth" else v
|
||||
for k, v in metadata.items()
|
||||
|
|
@ -468,6 +468,38 @@ class ComplexityRouter(CustomLogger):
|
|||
def _tier_pools(self) -> dict[str, list[str]]:
|
||||
return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()}
|
||||
|
||||
async def _pick_model_for_tier(
|
||||
self,
|
||||
tier: ComplexityTier,
|
||||
raw_messages: list[dict[str, Any]] | None,
|
||||
resolved_messages: list[dict[str, Any]] | None,
|
||||
request_kwargs: dict,
|
||||
) -> str:
|
||||
if not self.config.plugins:
|
||||
return self.get_model_for_tier(tier)
|
||||
|
||||
from litellm.types.router import RoutingContext
|
||||
|
||||
tier_key = tier.value
|
||||
metadata_key = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata"
|
||||
context = RoutingContext(
|
||||
raw_messages=raw_messages or [],
|
||||
structured_messages=resolved_messages or [],
|
||||
candidate_models=list(self._tier_pools().get(tier_key, [])),
|
||||
metadata=request_kwargs.get(metadata_key) or {},
|
||||
)
|
||||
for plugin in self.config.plugins:
|
||||
context = await plugin.run(context)
|
||||
|
||||
if not context.candidate_models:
|
||||
# A plugin narrowing a tier to zero candidates is a policy decision (e.g. no
|
||||
# model this tenant's budget allows) -- falling back to default_model here
|
||||
# (which was never checked against the plugins) would let that policy be
|
||||
# silently bypassed. Raise instead, matching the Router-level plugin
|
||||
# pipeline's own fail-closed behavior for the same situation.
|
||||
raise ValueError(f"No candidate models left for tier {tier_key} after routing-plugin filtering")
|
||||
return self._pick_from_tier_value(context.candidate_models, tier_key)
|
||||
|
||||
def _ensure_adaptive_router(self) -> Any | None:
|
||||
if not self.config.adaptive:
|
||||
return None
|
||||
|
|
@ -731,8 +763,8 @@ class ComplexityRouter(CustomLogger):
|
|||
# embedding call. Forwarding it would let the embedding's cost callback finalize the
|
||||
# reservation, so the routed completion's own callback then skips incrementing the
|
||||
# key/team budget. Key/team attribution fields are preserved for spend logging.
|
||||
metadata = _classifier_call_metadata(request_kwargs.get("metadata")) or {}
|
||||
litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata")) or {}
|
||||
metadata = _classifier_call_metadata(request_kwargs.get("metadata"))
|
||||
litellm_metadata = _classifier_call_metadata(request_kwargs.get("litellm_metadata"))
|
||||
query_vector = (
|
||||
await encoder.aencode_queries([user_message], metadata=metadata, litellm_metadata=litellm_metadata)
|
||||
)[0]
|
||||
|
|
@ -861,10 +893,16 @@ class ComplexityRouter(CustomLogger):
|
|||
When `session_affinity` is enabled and a session_id is resolvable on the request,
|
||||
pins the model chosen on the session's first turn and reuses it for every later
|
||||
turn, skipping classification entirely. Otherwise delegates to `_classify_and_route`.
|
||||
|
||||
Skipped entirely when `plugins` are configured: reusing a stale pin would bypass
|
||||
the plugin pipeline on every turn after the first, since a pinned model was never
|
||||
re-checked against a policy plugin whose decision can change between turns (e.g. a
|
||||
budget plugin, once the session's spend crosses its cap).
|
||||
"""
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
session_id = self._get_session_id_from_request_kwargs(request_kwargs) if self.config.session_affinity else None
|
||||
use_session_affinity = self.config.session_affinity and not self.config.plugins
|
||||
session_id = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None
|
||||
cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
|
||||
|
||||
if cache_key is not None:
|
||||
|
|
@ -947,14 +985,26 @@ class ComplexityRouter(CustomLogger):
|
|||
|
||||
if user_message is None:
|
||||
verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model")
|
||||
if not self.config.plugins and self.config.default_model:
|
||||
# No plugins configured: preserve the pre-existing default_model-first
|
||||
# priority exactly (changing it would be a silent behavior change for
|
||||
# every non-plugin user, not just a security fix).
|
||||
routed_model = self.config.default_model
|
||||
else:
|
||||
# Plugins configured: default_model must never bypass them, so it's not
|
||||
# checked here at all -- _pick_model_for_tier -> get_model_for_tier still
|
||||
# falls back to it (after the MEDIUM tier) once the plugin pipeline runs.
|
||||
routed_model = await self._pick_model_for_tier(
|
||||
ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
|
||||
)
|
||||
return PreRoutingHookResponse(
|
||||
model=self.config.default_model or self.get_model_for_tier(ComplexityTier.MEDIUM),
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
)
|
||||
|
||||
override_tier = await self._resolve_keyword_tier_override(user_message, request_kwargs)
|
||||
if override_tier is not None:
|
||||
routed_model = self.get_model_for_tier(override_tier)
|
||||
routed_model = await self._pick_model_for_tier(override_tier, messages, resolved_messages, request_kwargs)
|
||||
cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match"
|
||||
verbose_router_logger.info(
|
||||
f"ComplexityRouter: routing decision cause={cause}, "
|
||||
|
|
@ -980,7 +1030,7 @@ class ComplexityRouter(CustomLogger):
|
|||
f"signals={signals}, routed_model={routed_model}"
|
||||
)
|
||||
else:
|
||||
routed_model = self.get_model_for_tier(tier)
|
||||
routed_model = await self._pick_model_for_tier(tier, messages, resolved_messages, request_kwargs)
|
||||
verbose_router_logger.info(
|
||||
f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, "
|
||||
f"score={score:.3f}, signals={signals}, routed_model={routed_model}"
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from typing import Literal
|
|||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
from litellm.types.router import AdaptiveRouterWeights
|
||||
from litellm.types.router import AdaptiveRouterWeights, RoutingPlugin
|
||||
|
||||
|
||||
class ComplexityTier(str, Enum):
|
||||
|
|
@ -363,10 +363,13 @@ class ComplexityRouterConfig(BaseModel):
|
|||
|
||||
# Session affinity: pin the first turn's routed model for the rest of the session
|
||||
session_affinity: bool = Field(
|
||||
default=False,
|
||||
default=True,
|
||||
description=(
|
||||
"When True and a session_id is resolvable on the request, pin the model chosen on the "
|
||||
"session's first turn and reuse it for every later turn, skipping re-classification."
|
||||
"session's first turn and reuse it for every later turn, skipping re-classification. "
|
||||
"On by default so multi-turn sessions stay on one model, preserving provider prompt "
|
||||
"caches and avoiding cross-model conversation-history errors. Set False to reclassify "
|
||||
"every turn."
|
||||
),
|
||||
)
|
||||
session_affinity_ttl_seconds: int = Field(
|
||||
|
|
@ -375,7 +378,12 @@ class ComplexityRouterConfig(BaseModel):
|
|||
description="TTL for the session affinity pin; refreshed on every cache hit",
|
||||
)
|
||||
|
||||
model_config = ConfigDict(extra="allow") # Allow additional fields
|
||||
plugins: list[RoutingPlugin] | None = Field(
|
||||
default=None,
|
||||
description="RoutingPlugin instances that narrow the classified tier's candidate models before selection",
|
||||
)
|
||||
|
||||
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) # Allow additional fields
|
||||
|
||||
@field_validator("tiers", mode="before")
|
||||
@classmethod
|
||||
|
|
@ -421,6 +429,15 @@ class ComplexityRouterConfig(BaseModel):
|
|||
raise ValueError("keyword_tier_rules must be non-empty when semantic_keyword_matching is enabled")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_plugins_adaptive_combo(self) -> "ComplexityRouterConfig":
|
||||
if self.plugins and self.adaptive:
|
||||
raise ValueError(
|
||||
"plugins and adaptive=True cannot both be set: adaptive's bandit selection doesn't yet "
|
||||
"consume plugin-narrowed candidate pools. Disable adaptive or remove plugins."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
# Combined default config
|
||||
DEFAULT_COMPLEXITY_CONFIG = ComplexityRouterConfig()
|
||||
|
|
|
|||
|
|
@ -1,14 +1,27 @@
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.llms.openai import OpenAIFileObject
|
||||
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.types.router import CredentialLiteLLMParams
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
|
||||
def _is_proxy_admin_request(request_kwargs: Optional[Mapping[str, object]]) -> bool:
|
||||
if request_kwargs is None:
|
||||
return False
|
||||
metadata_value = request_kwargs.get("metadata")
|
||||
litellm_metadata_value = request_kwargs.get("litellm_metadata")
|
||||
metadata = metadata_value if isinstance(metadata_value, Mapping) else {}
|
||||
litellm_metadata = litellm_metadata_value if isinstance(litellm_metadata_value, Mapping) else {}
|
||||
user_api_key_auth = metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth")
|
||||
return getattr(user_api_key_auth, "user_role", None) == "proxy_admin"
|
||||
|
||||
|
||||
def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str:
|
||||
"""
|
||||
Hash of the credential params, used for mapping the file id to the right model
|
||||
|
|
@ -59,6 +72,40 @@ def filter_team_based_models(
|
|||
metadata = request_kwargs.get("metadata") or {}
|
||||
litellm_metadata = request_kwargs.get("litellm_metadata") or {}
|
||||
request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id")
|
||||
if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list):
|
||||
requested_model = (
|
||||
request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group")
|
||||
)
|
||||
candidate_deployments = tuple(
|
||||
(deployment.get("model_name"), deployment.get("model_info") or {}) for deployment in healthy_deployments
|
||||
)
|
||||
team_ids = frozenset(
|
||||
team_id
|
||||
for _, model_info in candidate_deployments
|
||||
for team_id in [model_info.get("team_id")]
|
||||
if team_id is not None
|
||||
)
|
||||
matches_requested_model = (
|
||||
isinstance(requested_model, str)
|
||||
and bool(candidate_deployments)
|
||||
and all(
|
||||
model_info.get("team_id") is not None
|
||||
and (model_name == requested_model or model_info.get("team_public_model_name") == requested_model)
|
||||
for model_name, model_info in candidate_deployments
|
||||
)
|
||||
)
|
||||
if matches_requested_model and len(team_ids) > 1:
|
||||
raise BadRequestError(
|
||||
message=(
|
||||
f"Model name '{requested_model}' matches deployments from multiple teams. "
|
||||
"Specify the deployment ID directly to disambiguate."
|
||||
),
|
||||
model=requested_model,
|
||||
llm_provider="",
|
||||
)
|
||||
if matches_requested_model:
|
||||
return healthy_deployments
|
||||
|
||||
ids_to_remove = set()
|
||||
if isinstance(healthy_deployments, dict):
|
||||
return healthy_deployments
|
||||
|
|
|
|||
|
|
@ -43,14 +43,21 @@ ITPM_CACHE_KEY = "_litellm_itpm_cache_key"
|
|||
OTPM_CACHE_KEY = "_litellm_otpm_cache_key"
|
||||
|
||||
|
||||
def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]]) -> None:
|
||||
def set_io_token_rate_limit_request_kwargs(kwargs: Optional[dict[str, Any]], store_in_context: bool = True) -> None:
|
||||
# The reservation sentinels are server-only, but `metadata` is caller
|
||||
# controlled on proxy requests. Strip any client-supplied copies here (this
|
||||
# runs before the router stashes its own reservation) so a forged
|
||||
# reservation can't drive the post-call reconcile/refund against an
|
||||
# arbitrary counter and bypass the configured limits.
|
||||
_clear_reservation_from_kwargs(kwargs)
|
||||
_io_token_rate_limit_request_kwargs.set(kwargs)
|
||||
# The context slot pins the entire request kwargs (messages included) for
|
||||
# the lifetime of the surrounding context, which outlives the request when
|
||||
# the context is captured by pooled resources (e.g. a redis connection
|
||||
# created mid-request). Only ITPM/OTPM-limited deployments read it, so for
|
||||
# every other deployment overwrite the slot with None instead of the
|
||||
# kwargs; overwriting (rather than skipping) also releases a previous
|
||||
# request's kwargs when a context is reused.
|
||||
_io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None)
|
||||
|
||||
|
||||
def get_io_token_rate_limit_request_kwargs() -> Optional[dict[str, Any]]:
|
||||
|
|
|
|||
|
|
@ -56,6 +56,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import (
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.headroom import (
|
||||
HeadroomGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.compresr import (
|
||||
CompresrGuardrailConfigModel,
|
||||
)
|
||||
|
||||
"""
|
||||
Pydantic object defining how to set guardrails on litellm proxy
|
||||
|
|
@ -123,6 +126,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
VIGIL_GUARD = "vigil_guard"
|
||||
REPELLOAI = "repelloai"
|
||||
HEADROOM = "headroom"
|
||||
COMPRESR = "compresr"
|
||||
|
||||
|
||||
class Role(Enum):
|
||||
|
|
@ -806,7 +810,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
|
|||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when a guardrail endpoint is unreachable due to network errors. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', and 'headroom'. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
|
||||
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
|
||||
),
|
||||
)
|
||||
|
|
@ -899,6 +903,7 @@ class LitellmParams(
|
|||
BedrockGuardrailConfigModel,
|
||||
LakeraV2GuardrailConfigModel,
|
||||
HeadroomGuardrailConfigModel,
|
||||
CompresrGuardrailConfigModel,
|
||||
RepelloAIGuardrailConfigModel,
|
||||
LassoGuardrailConfigModel,
|
||||
PillarGuardrailConfigModel,
|
||||
|
|
@ -1034,6 +1039,7 @@ class ApplyGuardrailRequest(BaseModel):
|
|||
entities: Optional[List[PiiEntityType]] = None
|
||||
input_type: str = "request"
|
||||
messages: Optional[List[Dict[str, Any]]] = None
|
||||
metadata: Dict[str, Any] | None = None
|
||||
|
||||
|
||||
class ApplyGuardrailResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -17,6 +17,12 @@ MCPInfo = Dict[str, Any]
|
|||
|
||||
class MCPOAuthMetadata(BaseModel):
|
||||
scopes: Optional[List[str]] = None
|
||||
"""Resource-driven scopes for the authorization request: the RFC 9728 protected-resource
|
||||
``scopes_supported``, or the ``scope`` from the WWW-Authenticate 401 challenge when the resource
|
||||
supplied one, else the authorization server's ``scopes_supported``. This is the scope value a
|
||||
client requests per the MCP authorization spec Scope Selection Strategy; scope minimization and
|
||||
inflation control are the authorization server's and user's job at consent (RFC 6749 §3.3), not
|
||||
the client's."""
|
||||
authorization_url: Optional[str] = None
|
||||
token_url: Optional[str] = None
|
||||
registration_url: Optional[str] = None
|
||||
|
|
@ -122,9 +128,10 @@ class MCPServer(BaseModel):
|
|||
# response (supports dot-notation for nested fields, e.g. "team.enterprise_id").
|
||||
# Tokens that fail validation are rejected before storage.
|
||||
token_validation: Optional[Dict[str, Any]] = None
|
||||
# Optional TTL override (seconds) for the Redis per-user token cache.
|
||||
# Defaults to the token's expires_in minus the expiry buffer, or
|
||||
# MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.
|
||||
# Optional TTL override (seconds) for the Redis per-user token cache, capped
|
||||
# at the token's expires_in minus the expiry buffer so a cached entry never
|
||||
# outlives the token. Defaults to the token's expires_in minus the expiry
|
||||
# buffer, or MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.
|
||||
token_storage_ttl_seconds: Optional[int] = None
|
||||
timeout: Optional[float] = None
|
||||
# Max concurrent outbound tool calls to this server; excess calls queue.
|
||||
|
|
|
|||
135
litellm/types/proxy/guardrails/guardrail_hooks/compresr.py
Normal file
135
litellm/types/proxy/guardrails/guardrail_hooks/compresr.py
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
from typing import Any, Dict, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class CompresrGuardrailOptionalParams(BaseModel):
|
||||
"""Optional tuning knobs for the Compresr guardrail."""
|
||||
|
||||
target_compression_ratio: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Compression strength. 0-1 is the fraction of tokens to remove "
|
||||
"(0.5 = remove ~50%, the default); a value >1 is an Nx reduction "
|
||||
"factor (e.g. 4 = ~4x smaller)."
|
||||
),
|
||||
)
|
||||
coarse: bool | None = Field(
|
||||
default=None,
|
||||
description=("Paragraph-level compression (default, faster) instead of token-level (finer-grained)."),
|
||||
)
|
||||
min_chars_to_compress: int | None = Field(
|
||||
default=None,
|
||||
description=("Skip messages whose text is shorter than this many characters. Defaults to 500."),
|
||||
)
|
||||
compress_tool_outputs: bool | None = Field(
|
||||
default=None,
|
||||
description=("Compress tool/function result messages (search hits, RAG chunks, API dumps). Defaults to True."),
|
||||
)
|
||||
compress_system: bool | None = Field(
|
||||
default=None,
|
||||
description="Also compress system messages. Defaults to False.",
|
||||
)
|
||||
compress_history: bool | None = Field(
|
||||
default=None,
|
||||
description="Also compress prior (non-last) user messages. Defaults to False.",
|
||||
)
|
||||
compress_last_user: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Also compress the last user message. The query sent to Compresr "
|
||||
"is always the original verbatim text. Defaults to False."
|
||||
),
|
||||
)
|
||||
enable_retrieval: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Make compression recoverable: inject a `compresr_retrieve` tool "
|
||||
"so the model can fetch the original content behind a compression "
|
||||
"marker via the agentic loop. Defaults to True. Set to False (or "
|
||||
"run the proxy with --workers 1) for multi-worker deployments: "
|
||||
"the recovery store is per-process, so pre-call and retrieval hooks "
|
||||
"on different workers cannot see each other's originals."
|
||||
),
|
||||
)
|
||||
max_bytes_per_call: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Cap on aggregate bytes of stored originals per litellm_call_id. "
|
||||
"When a call exceeds this, oldest entries are evicted so the "
|
||||
"in-process store cannot grow without bound. Defaults to 10 MiB."
|
||||
),
|
||||
)
|
||||
allow_bypass_header: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Honor the `x-compresr-bypass: true` request header to skip "
|
||||
"compression for a single call. Off by default because the "
|
||||
"header is caller-settable; enable only on trusted deployments."
|
||||
),
|
||||
)
|
||||
dynamic: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"latte_v2 only. Let the server choose the compression amount per input "
|
||||
"(Kneedle elbow) instead of using target_compression_ratio. Defaults to True."
|
||||
),
|
||||
)
|
||||
dynamic_min_ratio: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"latte_v2 only. Floor on the adaptive ratio when `dynamic` is on. "
|
||||
"Unset lets the server default apply (~1.5)."
|
||||
),
|
||||
)
|
||||
dynamic_max_ratio: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"latte_v2 only. Ceiling on the adaptive ratio when `dynamic` is on. "
|
||||
"Unset lets the server default apply (~10.0)."
|
||||
),
|
||||
)
|
||||
compression_params: Dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Passthrough of extra parameters forwarded verbatim in the Compresr "
|
||||
"compress payload (e.g. `heuristic_chunking`, or any newer knob), so "
|
||||
"a new Compresr feature works without a guardrail update. The named "
|
||||
"fields above take precedence on collision."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CompresrGuardrailConfigModel(GuardrailConfigModel[CompresrGuardrailOptionalParams]):
|
||||
api_key: str | None = Field(
|
||||
default=None,
|
||||
description=("Compresr API key. Falls back to the COMPRESR_API_KEY env var."),
|
||||
)
|
||||
api_base: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Base URL of the Compresr API. Falls back to the COMPRESR_API_BASE "
|
||||
"env var, then https://api.compresr.ai. Point at your internal "
|
||||
"service URL for on-prem deployments."
|
||||
),
|
||||
)
|
||||
model: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Compresr compression model (not the LLM). Defaults to 'latte_v2', the query-aware compression model."
|
||||
),
|
||||
)
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when the Compresr compression service is unreachable or errors. "
|
||||
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and "
|
||||
"forwards the request uncompressed instead of blocking it."
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Compresr (context compression)"
|
||||
|
|
@ -9,7 +9,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hi
|
|||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
from typing_extensions import Protocol, Required, TypedDict
|
||||
from typing_extensions import Protocol, Required, TypedDict, runtime_checkable
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
||||
|
|
@ -852,6 +852,7 @@ class RoutingContext(BaseModel):
|
|||
signals: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RoutingPlugin(Protocol):
|
||||
"""Interface a custom routing plugin must implement to run in `Router(plugins=[...])`."""
|
||||
|
||||
|
|
|
|||
|
|
@ -44382,6 +44382,90 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-sol": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"output_cost_per_token": 3.3e-05,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"use_openai_responses_path": true,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-terra": {
|
||||
"input_cost_per_token": 2.75e-06,
|
||||
"cache_creation_input_token_cost": 3.4375e-06,
|
||||
"cache_read_input_token_cost": 2.75e-07,
|
||||
"output_cost_per_token": 1.65e-05,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"use_openai_responses_path": true,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.6-luna": {
|
||||
"input_cost_per_token": 1.1e-06,
|
||||
"cache_creation_input_token_cost": 1.375e-06,
|
||||
"cache_read_input_token_cost": 1.1e-07,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"litellm_provider": "bedrock_mantle",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "responses",
|
||||
"use_openai_responses_path": true,
|
||||
"supported_endpoints": [
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/openai.gpt-5.5": {
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
|
|
@ -44494,6 +44578,7 @@
|
|||
"supports_vision": true
|
||||
},
|
||||
"bedrock_mantle/xai.grok-4.3": {
|
||||
"use_openai_responses_path": true,
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ name = "litellm"
|
|||
version = "1.94.0"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.14"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
license = "MIT"
|
||||
license-files = ["LICENSE"]
|
||||
authors = [
|
||||
|
|
@ -66,6 +66,7 @@ proxy = [
|
|||
"litellm-enterprise==0.1.50",
|
||||
"RestrictedPython>=8.1,<9.0",
|
||||
"rich>=13.9.4,<14.0",
|
||||
"InquirerPy>=0.3.4,<1.0",
|
||||
"polars>=1.38.1,<2.0",
|
||||
"soundfile>=0.12.1,<1.0",
|
||||
"pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'",
|
||||
|
|
@ -74,11 +75,12 @@ proxy = [
|
|||
]
|
||||
# Thin client install for the `lite` CLI on developer laptops. The CLI's heavy
|
||||
# imports (fastapi, cryptography, ...) are all guarded, so it runs on the base
|
||||
# SDK plus just these three; none of the server runtime in `proxy` is pulled in.
|
||||
# SDK plus just these four; none of the server runtime in `proxy` is pulled in.
|
||||
cli = [
|
||||
"rich>=13.9.4,<14.0",
|
||||
"pyyaml>=6.0.3,<7.0",
|
||||
"requests>=2.32.0,<3.0",
|
||||
"InquirerPy>=0.3.4,<1.0",
|
||||
]
|
||||
extra_proxy = [
|
||||
"prisma>=0.11.0,<1.0",
|
||||
|
|
@ -129,7 +131,7 @@ proxy-runtime = [
|
|||
"opentelemetry-sdk==1.28.0",
|
||||
"opentelemetry-exporter-otlp==1.28.0",
|
||||
"opentelemetry-instrumentation-fastapi==0.49b0",
|
||||
"ddtrace>=2.19.0,<3.0",
|
||||
"ddtrace>=4.8.2,<5.0",
|
||||
"sentry-sdk>=2.21.0,<3.0",
|
||||
"mangum>=0.17.0,<1.0",
|
||||
"azure-ai-contentsafety>=1.0.0,<2.0",
|
||||
|
|
|
|||
|
|
@ -11,12 +11,21 @@
|
|||
# Python itself (honouring litellm's requires-python), downloading a managed one
|
||||
# when the host has no suitable interpreter.
|
||||
#
|
||||
# To try an unreleased branch instead of the latest PyPI release (for example, to
|
||||
# QA a CLI feature before it ships), set LITELLM_CLI_REF to a branch, tag, or commit:
|
||||
# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/<branch>/scripts/install-cli.sh | \
|
||||
# LITELLM_CLI_REF=<branch> sh
|
||||
#
|
||||
# NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian
|
||||
# ignores the shebang when invoked as `sh` and does not support `pipefail`).
|
||||
set -eu
|
||||
|
||||
# NOTE: before merging, this must stay as "litellm[cli]" to install from PyPI.
|
||||
LITELLM_PACKAGE="litellm[cli]"
|
||||
# Defaults to the PyPI release; LITELLM_CLI_REF opts into installing from source instead.
|
||||
if [ -n "${LITELLM_CLI_REF:-}" ]; then
|
||||
LITELLM_PACKAGE="litellm[cli] @ git+https://github.com/BerriAI/litellm.git@${LITELLM_CLI_REF}"
|
||||
else
|
||||
LITELLM_PACKAGE="litellm[cli]"
|
||||
fi
|
||||
UV_VERSION="0.10.9"
|
||||
|
||||
# ── colours ────────────────────────────────────────────────────────────────
|
||||
|
|
@ -90,7 +99,11 @@ fi
|
|||
# otherwise download a managed one. Either way uv honours litellm's requires-python,
|
||||
# so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced.
|
||||
echo ""
|
||||
header "Installing litellm[cli]…"
|
||||
if [ -n "${LITELLM_CLI_REF:-}" ]; then
|
||||
header "Installing litellm[cli] from ${LITELLM_CLI_REF}…"
|
||||
else
|
||||
header "Installing litellm[cli]…"
|
||||
fi
|
||||
echo ""
|
||||
|
||||
"$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \
|
||||
|
|
|
|||
|
|
@ -5,12 +5,24 @@
|
|||
# Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible
|
||||
# Python itself (reusing a suitable system one, else downloading a managed build).
|
||||
#
|
||||
# To install from an unreleased branch, tag, or commit instead of the latest PyPI
|
||||
# release, set LITELLM_CLI_REF:
|
||||
# curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/<branch>/scripts/install.sh | \
|
||||
# LITELLM_CLI_REF=<branch> sh
|
||||
#
|
||||
# NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian
|
||||
# ignores the shebang when invoked as `sh` and does not support `pipefail`).
|
||||
set -eu
|
||||
|
||||
# NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI.
|
||||
LITELLM_PACKAGE="litellm[proxy]"
|
||||
# LITELLM_CLI_REF opts into installing from a branch, tag, or commit instead (for
|
||||
# example, to QA lite autoroute against an unreleased branch, which needs this proxy
|
||||
# runtime, not the thin litellm[cli] install).
|
||||
if [ -n "${LITELLM_CLI_REF:-}" ]; then
|
||||
LITELLM_PACKAGE="litellm[proxy] @ git+https://github.com/BerriAI/litellm.git@${LITELLM_CLI_REF}"
|
||||
else
|
||||
LITELLM_PACKAGE="litellm[proxy]"
|
||||
fi
|
||||
UV_VERSION="0.10.9"
|
||||
|
||||
# ── colours ────────────────────────────────────────────────────────────────
|
||||
|
|
@ -81,7 +93,11 @@ fi
|
|||
|
||||
# ── install ────────────────────────────────────────────────────────────────
|
||||
echo ""
|
||||
header "Installing litellm[proxy]…"
|
||||
if [ -n "${LITELLM_CLI_REF:-}" ]; then
|
||||
header "Installing litellm[proxy] from ${LITELLM_CLI_REF}…"
|
||||
else
|
||||
header "Installing litellm[proxy]…"
|
||||
fi
|
||||
echo ""
|
||||
|
||||
# --python-preference system: reuse a compatible system Python when present,
|
||||
|
|
|
|||
|
|
@ -158,6 +158,38 @@ AgentOps) live under `proxy_config.litellm_settings.callbacks` and are
|
|||
orthogonal to the OTLP variables above; their credentials still go in
|
||||
`*_extra_secrets`.
|
||||
|
||||
### Enterprise billing metrics
|
||||
|
||||
License-gated request metering is opt-in and gated entirely on
|
||||
`billing_metrics_endpoint`. Empty (default) and no billing env is added to
|
||||
the container, so existing deployments are unchanged. Set it and both
|
||||
gateway and backend export billable-request counts over OTLP/HTTP,
|
||||
authenticating to the collector with the mTLS client certificate issued for
|
||||
your deployment.
|
||||
|
||||
The proxy accepts the certificate, key, and CA bundle as either a file path
|
||||
or literal PEM content. This stack takes the PEM, writes each one to its own
|
||||
Secrets Manager entry, grants the task-execution role
|
||||
`secretsmanager:GetSecretValue` on them, and injects them as
|
||||
`LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and `_CA_CERT` when
|
||||
set), so no volume mount is needed on Fargate.
|
||||
|
||||
```hcl
|
||||
billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics"
|
||||
```
|
||||
|
||||
```bash
|
||||
export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)"
|
||||
export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)"
|
||||
```
|
||||
|
||||
`billing_metrics_ca_cert_pem` is only for private or test collectors whose
|
||||
CA is not in the system trust store; leave it empty against
|
||||
`telemetry.litellm.ai`. Metering requires an enterprise license, so pair
|
||||
this with `litellm_license`. To tune the export cadence, set
|
||||
`LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` /
|
||||
`backend_extra_env`
|
||||
|
||||
## Tenant deployment
|
||||
|
||||
Every resource the stack creates is named `${tenant}-litellm-${env}` (or
|
||||
|
|
|
|||
|
|
@ -76,6 +76,33 @@ locals {
|
|||
{ name = "OTEL_HEADERS", valueFrom = var.otel_headers_secret_arn },
|
||||
] : []
|
||||
|
||||
# Enterprise request metering, gated on billing_metrics_endpoint. The
|
||||
# endpoint rides in as a plain env var; the mTLS material is stored in
|
||||
# Secrets Manager (secrets.tf) and injected as PEM-valued env vars, which
|
||||
# the proxy accepts in place of file paths. Each PEM is wired only when the
|
||||
# operator supplied it, so an empty ca_cert_pem falls back to the system
|
||||
# trust store.
|
||||
billing_metrics_enabled = var.billing_metrics_endpoint != ""
|
||||
billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != ""
|
||||
billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != ""
|
||||
billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != ""
|
||||
|
||||
billing_metrics_env = local.billing_metrics_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint },
|
||||
] : []
|
||||
|
||||
billing_metrics_secrets = concat(
|
||||
local.billing_metrics_client_cert_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_CLIENT_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_client_cert[0].arn },
|
||||
] : [],
|
||||
local.billing_metrics_client_key_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_CLIENT_KEY", valueFrom = aws_secretsmanager_secret.billing_metrics_client_key[0].arn },
|
||||
] : [],
|
||||
local.billing_metrics_ca_cert_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_CA_CERT", valueFrom = aws_secretsmanager_secret.billing_metrics_ca_cert[0].arn },
|
||||
] : [],
|
||||
)
|
||||
|
||||
shared_env = [
|
||||
{ name = "IAM_TOKEN_DB_AUTH", value = "true" },
|
||||
{ name = "DATABASE_HOST", value = aws_rds_cluster.this.endpoint },
|
||||
|
|
@ -108,6 +135,7 @@ locals {
|
|||
{ name = "LITELLM_LICENSE", valueFrom = aws_secretsmanager_secret.license[0].arn },
|
||||
],
|
||||
local.otel_secrets,
|
||||
local.billing_metrics_secrets,
|
||||
)
|
||||
|
||||
# Backend-only managed secrets. UI_PASSWORD is consumed by the management
|
||||
|
|
@ -179,6 +207,30 @@ locals {
|
|||
|
||||
# ---------- Gateway ----------
|
||||
resource "aws_ecs_task_definition" "gateway" {
|
||||
# Metering needs a client certificate AND its key. Each secret is created only
|
||||
# when its own PEM is supplied, so an endpoint set with a missing key would
|
||||
# otherwise apply cleanly and leave the proxy logging "missing config" and
|
||||
# never exporting. ca_cert_pem stays optional: empty means fall back to the
|
||||
# system trust store.
|
||||
#
|
||||
# The guard lives here, on an unconditional resource, rather than on the cert
|
||||
# secret: that secret is count-gated on the cert itself, so it has zero
|
||||
# instances in exactly the case this must catch. Adding count or for_each to
|
||||
# this resource would silently stop the guard from evaluating.
|
||||
#
|
||||
# endpoint cert key -> result
|
||||
# "" any any -> metering off, no secrets created
|
||||
# set set set -> metering on
|
||||
# set any-missing -> plan fails here
|
||||
lifecycle {
|
||||
precondition {
|
||||
condition = var.billing_metrics_endpoint == "" || (
|
||||
var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
|
||||
)
|
||||
error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
|
||||
}
|
||||
}
|
||||
|
||||
family = "${local.name}-gateway"
|
||||
network_mode = "awsvpc"
|
||||
requires_compatibilities = ["FARGATE"]
|
||||
|
|
@ -198,6 +250,7 @@ resource "aws_ecs_task_definition" "gateway" {
|
|||
environment = concat(
|
||||
local.shared_env,
|
||||
local.gateway_otel_env,
|
||||
local.billing_metrics_env,
|
||||
local.gateway_extra_env_list,
|
||||
local.proxy_config_env,
|
||||
)
|
||||
|
|
@ -264,6 +317,18 @@ resource "aws_ecs_service" "gateway" {
|
|||
|
||||
# ---------- Backend ----------
|
||||
resource "aws_ecs_task_definition" "backend" {
|
||||
# Same guard as the gateway: the backend meters too (it serves the named-server
|
||||
# MCP transport), and a targeted apply of just this resource must not slip a
|
||||
# billing endpoint through without the credentials to use it.
|
||||
lifecycle {
|
||||
precondition {
|
||||
condition = var.billing_metrics_endpoint == "" || (
|
||||
var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
|
||||
)
|
||||
error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
|
||||
}
|
||||
}
|
||||
|
||||
family = "${local.name}-backend"
|
||||
network_mode = "awsvpc"
|
||||
requires_compatibilities = ["FARGATE"]
|
||||
|
|
@ -284,6 +349,7 @@ resource "aws_ecs_task_definition" "backend" {
|
|||
local.shared_env,
|
||||
local.backend_default_env,
|
||||
local.backend_otel_env,
|
||||
local.billing_metrics_env,
|
||||
local.backend_extra_env_list,
|
||||
local.proxy_config_env,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -53,6 +53,9 @@ data "aws_iam_policy_document" "secrets_access" {
|
|||
[aws_secretsmanager_secret.master_key.arn],
|
||||
aws_secretsmanager_secret.license[*].arn,
|
||||
aws_secretsmanager_secret.ui_password[*].arn,
|
||||
aws_secretsmanager_secret.billing_metrics_client_cert[*].arn,
|
||||
aws_secretsmanager_secret.billing_metrics_client_key[*].arn,
|
||||
aws_secretsmanager_secret.billing_metrics_ca_cert[*].arn,
|
||||
local.extra_secret_arns,
|
||||
var.otel_headers_secret_arn == "" ? [] : [var.otel_headers_secret_arn],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -74,6 +74,61 @@ resource "aws_secretsmanager_secret_version" "ui_password" {
|
|||
secret_string = var.ui_password
|
||||
}
|
||||
|
||||
# Billing-metrics mTLS material — only created when metering is enabled
|
||||
# (billing_metrics_endpoint non-empty) and the operator supplied the PEM.
|
||||
# The task-execution role gets GetSecretValue via iam.tf, and gateway +
|
||||
# backend pick the env vars up through shared_secrets in ecs.tf.
|
||||
resource "aws_secretsmanager_secret" "billing_metrics_client_cert" {
|
||||
count = local.billing_metrics_client_cert_enabled ? 1 : 0
|
||||
|
||||
name = "${local.name}-billing-metrics-client-cert"
|
||||
description = "LITELLM_BILLING_METRICS_CLIENT_CERT for gateway + backend."
|
||||
recovery_window_in_days = 0
|
||||
|
||||
tags = local.tags
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret_version" "billing_metrics_client_cert" {
|
||||
count = local.billing_metrics_client_cert_enabled ? 1 : 0
|
||||
|
||||
secret_id = aws_secretsmanager_secret.billing_metrics_client_cert[0].id
|
||||
secret_string = var.billing_metrics_client_cert_pem
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret" "billing_metrics_client_key" {
|
||||
count = local.billing_metrics_client_key_enabled ? 1 : 0
|
||||
|
||||
name = "${local.name}-billing-metrics-client-key"
|
||||
description = "LITELLM_BILLING_METRICS_CLIENT_KEY for gateway + backend."
|
||||
recovery_window_in_days = 0
|
||||
|
||||
tags = local.tags
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret_version" "billing_metrics_client_key" {
|
||||
count = local.billing_metrics_client_key_enabled ? 1 : 0
|
||||
|
||||
secret_id = aws_secretsmanager_secret.billing_metrics_client_key[0].id
|
||||
secret_string = var.billing_metrics_client_key_pem
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret" "billing_metrics_ca_cert" {
|
||||
count = local.billing_metrics_ca_cert_enabled ? 1 : 0
|
||||
|
||||
name = "${local.name}-billing-metrics-ca-cert"
|
||||
description = "LITELLM_BILLING_METRICS_CA_CERT for gateway + backend."
|
||||
recovery_window_in_days = 0
|
||||
|
||||
tags = local.tags
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret_version" "billing_metrics_ca_cert" {
|
||||
count = local.billing_metrics_ca_cert_enabled ? 1 : 0
|
||||
|
||||
secret_id = aws_secretsmanager_secret.billing_metrics_ca_cert[0].id
|
||||
secret_string = var.billing_metrics_ca_cert_pem
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret" "db_master_password" {
|
||||
name = "${local.name}-db-master-password"
|
||||
description = "Aurora master-user password - bootstrap only. Runtime auth is IAM-token."
|
||||
|
|
|
|||
|
|
@ -533,3 +533,65 @@ variable "otel_headers_secret_arn" {
|
|||
type = string
|
||||
default = ""
|
||||
}
|
||||
|
||||
# ---------- Enterprise billing metrics ----------
|
||||
#
|
||||
# License-gated request metering. Opt-in and gated entirely on
|
||||
# billing_metrics_endpoint: leave it empty (the default) and nothing
|
||||
# metering-related lands in the container env. Set it and gateway + backend
|
||||
# export billable-request counts over OTLP/HTTP, authenticating to the
|
||||
# collector with an mTLS client cert. The proxy accepts the cert, key, and CA
|
||||
# as either a file path or literal PEM content, so on Fargate they are
|
||||
# injected straight from Secrets Manager as env vars and no volume is needed.
|
||||
|
||||
variable "billing_metrics_endpoint" {
|
||||
description = <<-EOT
|
||||
OTLP/HTTP endpoint for enterprise billing metrics (sets
|
||||
LITELLM_BILLING_METRICS_ENDPOINT). Non-empty enables request metering;
|
||||
empty (default) disables it and adds no billing env to the container.
|
||||
Requires an enterprise license. Example:
|
||||
"https://telemetry.litellm.ai/v1/metrics"
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "billing_metrics_client_cert_pem" {
|
||||
description = <<-EOT
|
||||
PEM content of the mTLS client certificate issued for this deployment.
|
||||
When billing_metrics_endpoint is set, the stack stores this in a
|
||||
`<tenant>-litellm-<env>-billing-metrics-client-cert` Secrets Manager
|
||||
entry, grants the task-execution role GetSecretValue on it, and exposes
|
||||
it to gateway + backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required
|
||||
whenever metering is enabled.
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
sensitive = true
|
||||
}
|
||||
|
||||
variable "billing_metrics_client_key_pem" {
|
||||
description = <<-EOT
|
||||
PEM content of the private key matching
|
||||
billing_metrics_client_cert_pem. Stored in a
|
||||
`<tenant>-litellm-<env>-billing-metrics-client-key` Secrets Manager
|
||||
entry and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required
|
||||
whenever metering is enabled.
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
sensitive = true
|
||||
}
|
||||
|
||||
variable "billing_metrics_ca_cert_pem" {
|
||||
description = <<-EOT
|
||||
PEM content of the CA bundle used to verify the metering collector.
|
||||
Only needed for private or test collectors whose CA is not in the
|
||||
system trust store; telemetry.litellm.ai is publicly trusted, so leave
|
||||
this empty for production. When set, it is exposed as
|
||||
LITELLM_BILLING_METRICS_CA_CERT.
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
sensitive = true
|
||||
}
|
||||
|
|
|
|||
|
|
@ -204,6 +204,40 @@ Behavior matches the AWS stack 1:1; the only naming differences are
|
|||
`otel_headers_secret` (a Secret Manager resource ID) vs AWS's
|
||||
`otel_headers_secret_arn` (a Secrets Manager ARN).
|
||||
|
||||
### Enterprise billing metrics
|
||||
|
||||
License-gated request metering is opt-in and gated entirely on
|
||||
`billing_metrics_endpoint`. Empty (default) and no billing env is added to
|
||||
the container, so existing deployments are unchanged. Set it and both
|
||||
gateway and backend export billable-request counts over OTLP/HTTP,
|
||||
authenticating to the collector with the mTLS client certificate issued for
|
||||
your deployment.
|
||||
|
||||
The proxy accepts the certificate, key, and CA bundle as either a file path
|
||||
or literal PEM content. This stack takes the PEM, writes each one to its own
|
||||
Secret Manager entry, grants the runtime service account
|
||||
`roles/secretmanager.secretAccessor` on them, and injects them as Cloud Run
|
||||
secret env vars `LITELLM_BILLING_METRICS_CLIENT_CERT` / `_CLIENT_KEY` (and
|
||||
`_CA_CERT` when set), so no volume mount is needed.
|
||||
|
||||
```hcl
|
||||
billing_metrics_endpoint = "https://telemetry.litellm.ai/v1/metrics"
|
||||
```
|
||||
|
||||
```bash
|
||||
export TF_VAR_billing_metrics_client_cert_pem="$(cat client.crt)"
|
||||
export TF_VAR_billing_metrics_client_key_pem="$(cat client.key)"
|
||||
```
|
||||
|
||||
`billing_metrics_ca_cert_pem` is only for private or test collectors whose
|
||||
CA is not in the system trust store; leave it empty against
|
||||
`telemetry.litellm.ai`. Metering requires an enterprise license, so pair
|
||||
this with `litellm_license`. To tune the export cadence, set
|
||||
`LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` /
|
||||
`backend_extra_env`
|
||||
|
||||
Behavior matches the AWS stack 1:1; the variable names are identical
|
||||
|
||||
## Tenant deployment
|
||||
|
||||
Every resource the stack creates is named `${tenant}-litellm-${env}` (or
|
||||
|
|
|
|||
|
|
@ -59,6 +59,33 @@ locals {
|
|||
{ name = "OTEL_HEADERS", secret = var.otel_headers_secret, version = "latest" },
|
||||
] : []
|
||||
|
||||
# Enterprise request metering, gated on billing_metrics_endpoint. The
|
||||
# endpoint rides in as a plain env var; the mTLS material lives in Secret
|
||||
# Manager (secrets.tf) and is injected as PEM-valued env vars, which the
|
||||
# proxy accepts in place of file paths. Each PEM is wired only when the
|
||||
# operator supplied it, so an empty ca_cert_pem falls back to the system
|
||||
# trust store.
|
||||
billing_metrics_enabled = var.billing_metrics_endpoint != ""
|
||||
billing_metrics_client_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_client_cert_pem != ""
|
||||
billing_metrics_client_key_enabled = local.billing_metrics_enabled && var.billing_metrics_client_key_pem != ""
|
||||
billing_metrics_ca_cert_enabled = local.billing_metrics_enabled && var.billing_metrics_ca_cert_pem != ""
|
||||
|
||||
billing_metrics_env_kv = local.billing_metrics_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_ENDPOINT", value = var.billing_metrics_endpoint },
|
||||
] : []
|
||||
|
||||
billing_metrics_env_secrets = concat(
|
||||
local.billing_metrics_client_cert_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_CLIENT_CERT", secret = google_secret_manager_secret.billing_metrics_client_cert[0].id, version = "latest" },
|
||||
] : [],
|
||||
local.billing_metrics_client_key_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_CLIENT_KEY", secret = google_secret_manager_secret.billing_metrics_client_key[0].id, version = "latest" },
|
||||
] : [],
|
||||
local.billing_metrics_ca_cert_enabled ? [
|
||||
{ name = "LITELLM_BILLING_METRICS_CA_CERT", secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id, version = "latest" },
|
||||
] : [],
|
||||
)
|
||||
|
||||
# Cloud Run v2 secret env vars use value_source.secret_key_ref pointing at a
|
||||
# secret resource ID. Shared between gateway and backend (the migrations
|
||||
# job has its own narrower env list — see migrations_env_secrets below).
|
||||
|
|
@ -138,6 +165,30 @@ locals {
|
|||
|
||||
# ---------- Gateway ----------
|
||||
resource "google_cloud_run_v2_service" "gateway" {
|
||||
# Metering needs a client certificate AND its key. Each secret is created only
|
||||
# when its own PEM is supplied, so an endpoint set with a missing key would
|
||||
# otherwise apply cleanly and leave the proxy logging "missing config" and
|
||||
# never exporting. ca_cert_pem stays optional: empty means fall back to the
|
||||
# system trust store.
|
||||
#
|
||||
# The guard lives here, on an unconditional resource, rather than on the cert
|
||||
# secret: that secret is count-gated on the cert itself, so it has zero
|
||||
# instances in exactly the case this must catch. Adding count or for_each to
|
||||
# this resource would silently stop the guard from evaluating.
|
||||
#
|
||||
# endpoint cert key -> result
|
||||
# "" any any -> metering off, no secrets created
|
||||
# set set set -> metering on
|
||||
# set any-missing -> plan fails here
|
||||
lifecycle {
|
||||
precondition {
|
||||
condition = var.billing_metrics_endpoint == "" || (
|
||||
var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
|
||||
)
|
||||
error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
|
||||
}
|
||||
}
|
||||
|
||||
name = "${local.name}-gateway"
|
||||
location = var.region
|
||||
ingress = "INGRESS_TRAFFIC_INTERNAL_LOAD_BALANCER"
|
||||
|
|
@ -175,7 +226,7 @@ resource "google_cloud_run_v2_service" "gateway" {
|
|||
}
|
||||
|
||||
dynamic "env" {
|
||||
for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.gateway_extra_env_kv, local.proxy_config_env)
|
||||
for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env)
|
||||
content {
|
||||
name = env.value.name
|
||||
value = env.value.value
|
||||
|
|
@ -183,7 +234,7 @@ resource "google_cloud_run_v2_service" "gateway" {
|
|||
}
|
||||
|
||||
dynamic "env" {
|
||||
for_each = concat(local.shared_env_secrets, local.otel_env_secrets, local.gateway_extra_secret_kv)
|
||||
for_each = concat(local.shared_env_secrets, local.otel_env_secrets, local.billing_metrics_env_secrets, local.gateway_extra_secret_kv)
|
||||
content {
|
||||
name = env.value.name
|
||||
value_source {
|
||||
|
|
@ -242,6 +293,9 @@ resource "google_cloud_run_v2_service" "gateway" {
|
|||
google_secret_manager_secret_iam_member.license,
|
||||
google_secret_manager_secret_iam_member.extras,
|
||||
google_secret_manager_secret_iam_member.otel_headers,
|
||||
google_secret_manager_secret_iam_member.billing_metrics_client_cert,
|
||||
google_secret_manager_secret_iam_member.billing_metrics_client_key,
|
||||
google_secret_manager_secret_iam_member.billing_metrics_ca_cert,
|
||||
google_storage_bucket_iam_member.proxy_config_runtime,
|
||||
google_sql_user.app,
|
||||
# Don't go live until the schema is migrated; otherwise the proxy boots,
|
||||
|
|
@ -252,6 +306,18 @@ resource "google_cloud_run_v2_service" "gateway" {
|
|||
|
||||
# ---------- Backend ----------
|
||||
resource "google_cloud_run_v2_service" "backend" {
|
||||
# Same guard as the gateway: the backend meters too (it serves the named-server
|
||||
# MCP transport), and a targeted apply of just this resource must not slip a
|
||||
# billing endpoint through without the credentials to use it.
|
||||
lifecycle {
|
||||
precondition {
|
||||
condition = var.billing_metrics_endpoint == "" || (
|
||||
var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != ""
|
||||
)
|
||||
error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
|
||||
}
|
||||
}
|
||||
|
||||
name = "${local.name}-backend"
|
||||
location = var.region
|
||||
ingress = "INGRESS_TRAFFIC_INTERNAL_LOAD_BALANCER"
|
||||
|
|
@ -289,7 +355,7 @@ resource "google_cloud_run_v2_service" "backend" {
|
|||
}
|
||||
|
||||
dynamic "env" {
|
||||
for_each = concat(local.shared_env_kv, local.backend_default_env_kv, local.backend_otel_env_kv, local.backend_extra_env_kv, local.proxy_config_env)
|
||||
for_each = concat(local.shared_env_kv, local.backend_default_env_kv, local.backend_otel_env_kv, local.billing_metrics_env_kv, local.backend_extra_env_kv, local.proxy_config_env)
|
||||
content {
|
||||
name = env.value.name
|
||||
value = env.value.value
|
||||
|
|
@ -297,7 +363,7 @@ resource "google_cloud_run_v2_service" "backend" {
|
|||
}
|
||||
|
||||
dynamic "env" {
|
||||
for_each = concat(local.shared_env_secrets, local.backend_managed_env_secrets, local.otel_env_secrets, local.backend_extra_secret_kv)
|
||||
for_each = concat(local.shared_env_secrets, local.backend_managed_env_secrets, local.otel_env_secrets, local.billing_metrics_env_secrets, local.backend_extra_secret_kv)
|
||||
content {
|
||||
name = env.value.name
|
||||
value_source {
|
||||
|
|
@ -357,6 +423,9 @@ resource "google_cloud_run_v2_service" "backend" {
|
|||
google_secret_manager_secret_iam_member.ui_password,
|
||||
google_secret_manager_secret_iam_member.extras,
|
||||
google_secret_manager_secret_iam_member.otel_headers,
|
||||
google_secret_manager_secret_iam_member.billing_metrics_client_cert,
|
||||
google_secret_manager_secret_iam_member.billing_metrics_client_key,
|
||||
google_secret_manager_secret_iam_member.billing_metrics_ca_cert,
|
||||
google_storage_bucket_iam_member.proxy_config_runtime,
|
||||
google_sql_user.app,
|
||||
terraform_data.migration,
|
||||
|
|
|
|||
|
|
@ -79,3 +79,29 @@ resource "google_secret_manager_secret_iam_member" "otel_headers" {
|
|||
role = "roles/secretmanager.secretAccessor"
|
||||
member = "serviceAccount:${google_service_account.runtime.email}"
|
||||
}
|
||||
|
||||
# Billing-metrics mTLS accessors — only created when request metering is
|
||||
# enabled and the matching PEM was supplied.
|
||||
resource "google_secret_manager_secret_iam_member" "billing_metrics_client_cert" {
|
||||
count = local.billing_metrics_client_cert_enabled ? 1 : 0
|
||||
|
||||
secret_id = google_secret_manager_secret.billing_metrics_client_cert[0].id
|
||||
role = "roles/secretmanager.secretAccessor"
|
||||
member = "serviceAccount:${google_service_account.runtime.email}"
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret_iam_member" "billing_metrics_client_key" {
|
||||
count = local.billing_metrics_client_key_enabled ? 1 : 0
|
||||
|
||||
secret_id = google_secret_manager_secret.billing_metrics_client_key[0].id
|
||||
role = "roles/secretmanager.secretAccessor"
|
||||
member = "serviceAccount:${google_service_account.runtime.email}"
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret_iam_member" "billing_metrics_ca_cert" {
|
||||
count = local.billing_metrics_ca_cert_enabled ? 1 : 0
|
||||
|
||||
secret_id = google_secret_manager_secret.billing_metrics_ca_cert[0].id
|
||||
role = "roles/secretmanager.secretAccessor"
|
||||
member = "serviceAccount:${google_service_account.runtime.email}"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -63,3 +63,58 @@ resource "google_secret_manager_secret_version" "ui_password" {
|
|||
secret = google_secret_manager_secret.ui_password[0].id
|
||||
secret_data = var.ui_password
|
||||
}
|
||||
|
||||
# Billing-metrics mTLS material — only created when metering is enabled
|
||||
# (billing_metrics_endpoint non-empty) and the operator supplied the PEM.
|
||||
# The runtime SA gets accessor permission via iam.tf, and gateway + backend
|
||||
# pick the env vars up through billing_metrics_env_secrets in cloudrun.tf.
|
||||
resource "google_secret_manager_secret" "billing_metrics_client_cert" {
|
||||
count = local.billing_metrics_client_cert_enabled ? 1 : 0
|
||||
|
||||
secret_id = "${local.name}-billing-metrics-client-cert"
|
||||
labels = local.labels
|
||||
replication {
|
||||
auto {}
|
||||
}
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret_version" "billing_metrics_client_cert" {
|
||||
count = local.billing_metrics_client_cert_enabled ? 1 : 0
|
||||
|
||||
secret = google_secret_manager_secret.billing_metrics_client_cert[0].id
|
||||
secret_data = var.billing_metrics_client_cert_pem
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret" "billing_metrics_client_key" {
|
||||
count = local.billing_metrics_client_key_enabled ? 1 : 0
|
||||
|
||||
secret_id = "${local.name}-billing-metrics-client-key"
|
||||
labels = local.labels
|
||||
replication {
|
||||
auto {}
|
||||
}
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret_version" "billing_metrics_client_key" {
|
||||
count = local.billing_metrics_client_key_enabled ? 1 : 0
|
||||
|
||||
secret = google_secret_manager_secret.billing_metrics_client_key[0].id
|
||||
secret_data = var.billing_metrics_client_key_pem
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret" "billing_metrics_ca_cert" {
|
||||
count = local.billing_metrics_ca_cert_enabled ? 1 : 0
|
||||
|
||||
secret_id = "${local.name}-billing-metrics-ca-cert"
|
||||
labels = local.labels
|
||||
replication {
|
||||
auto {}
|
||||
}
|
||||
}
|
||||
|
||||
resource "google_secret_manager_secret_version" "billing_metrics_ca_cert" {
|
||||
count = local.billing_metrics_ca_cert_enabled ? 1 : 0
|
||||
|
||||
secret = google_secret_manager_secret.billing_metrics_ca_cert[0].id
|
||||
secret_data = var.billing_metrics_ca_cert_pem
|
||||
}
|
||||
|
|
|
|||
|
|
@ -490,3 +490,66 @@ variable "otel_capture_message_content" {
|
|||
error_message = "otel_capture_message_content must be one of: no_content, prompt_and_completion."
|
||||
}
|
||||
}
|
||||
|
||||
# ---------- Enterprise billing metrics ----------
|
||||
#
|
||||
# License-gated request metering. Opt-in and gated entirely on
|
||||
# billing_metrics_endpoint: leave it empty (the default) and nothing
|
||||
# metering-related is added to the container env. Set it and gateway +
|
||||
# backend export billable-request counts over OTLP/HTTP, authenticating to
|
||||
# the collector with an mTLS client cert. The proxy accepts the cert, key,
|
||||
# and CA as either a file path or literal PEM content, so on Cloud Run they
|
||||
# are injected straight from Secret Manager as env vars and no volume is
|
||||
# needed.
|
||||
|
||||
variable "billing_metrics_endpoint" {
|
||||
description = <<-EOT
|
||||
OTLP/HTTP endpoint for enterprise billing metrics (sets
|
||||
LITELLM_BILLING_METRICS_ENDPOINT). Non-empty enables request metering;
|
||||
empty (default) disables it and adds no billing env to the container.
|
||||
Requires an enterprise license. Example:
|
||||
"https://telemetry.litellm.ai/v1/metrics"
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "billing_metrics_client_cert_pem" {
|
||||
description = <<-EOT
|
||||
PEM content of the mTLS client certificate issued for this deployment.
|
||||
When billing_metrics_endpoint is set, the stack stores this in a
|
||||
`<tenant>-litellm-<env>-billing-metrics-client-cert` Secret Manager
|
||||
entry, grants the runtime SA accessor on it, and exposes it to gateway +
|
||||
backend as LITELLM_BILLING_METRICS_CLIENT_CERT. Required whenever
|
||||
metering is enabled.
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
sensitive = true
|
||||
}
|
||||
|
||||
variable "billing_metrics_client_key_pem" {
|
||||
description = <<-EOT
|
||||
PEM content of the private key matching
|
||||
billing_metrics_client_cert_pem. Stored in a
|
||||
`<tenant>-litellm-<env>-billing-metrics-client-key` Secret Manager entry
|
||||
and exposed as LITELLM_BILLING_METRICS_CLIENT_KEY. Required whenever
|
||||
metering is enabled.
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
sensitive = true
|
||||
}
|
||||
|
||||
variable "billing_metrics_ca_cert_pem" {
|
||||
description = <<-EOT
|
||||
PEM content of the CA bundle used to verify the metering collector.
|
||||
Only needed for private or test collectors whose CA is not in the
|
||||
system trust store; telemetry.litellm.ai is publicly trusted, so leave
|
||||
this empty for production. When set, it is exposed as
|
||||
LITELLM_BILLING_METRICS_CA_CERT.
|
||||
EOT
|
||||
type = string
|
||||
default = ""
|
||||
sensitive = true
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,8 +16,8 @@ from pathlib import Path
|
|||
import pytest
|
||||
import yaml
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
MANIFEST_PATH = REPO_ROOT / "manifest.yaml"
|
||||
SUITE_ROOT = Path(__file__).resolve().parents[1]
|
||||
MANIFEST_PATH = SUITE_ROOT / "manifest.yaml"
|
||||
|
||||
# The PRD's "Features in v0" section, in row order.
|
||||
EXPECTED_FEATURE_IDS = [
|
||||
|
|
@ -90,14 +90,14 @@ def test_manifest_every_feature_has_human_readable_name(manifest):
|
|||
|
||||
@pytest.mark.parametrize("feature_id", EXPECTED_FEATURE_IDS)
|
||||
def test_feature_directory_exists(feature_id):
|
||||
feature_dir = REPO_ROOT / feature_id
|
||||
feature_dir = SUITE_ROOT / feature_id
|
||||
assert feature_dir.is_dir(), f"missing feature directory: {feature_dir}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("feature_id", EXPECTED_FEATURE_IDS)
|
||||
@pytest.mark.parametrize("provider", EXPECTED_PROVIDERS)
|
||||
def test_per_provider_test_file_exists(feature_id, provider):
|
||||
test_file = REPO_ROOT / feature_id / f"test_{provider}.py"
|
||||
test_file = SUITE_ROOT / feature_id / f"test_{provider}.py"
|
||||
assert test_file.is_file(), f"missing per-provider test file: {test_file}"
|
||||
|
||||
|
||||
|
|
@ -106,7 +106,7 @@ def test_feature_directory_has_init_file(feature_id):
|
|||
"""Each feature directory needs an __init__.py so pytest collects
|
||||
the per-provider test files as a package — matches the layout
|
||||
established by `basic_messaging_non_streaming/`."""
|
||||
init_file = REPO_ROOT / feature_id / "__init__.py"
|
||||
init_file = SUITE_ROOT / feature_id / "__init__.py"
|
||||
assert init_file.is_file(), f"missing __init__.py: {init_file}"
|
||||
|
||||
|
||||
|
|
@ -117,7 +117,7 @@ def test_feature_directory_has_init_file(feature_id):
|
|||
# a broken post-v0 directory still fails CI.
|
||||
@pytest.mark.parametrize("feature_id", ALL_FEATURE_IDS)
|
||||
def test_every_manifest_feature_has_directory(feature_id):
|
||||
feature_dir = REPO_ROOT / feature_id
|
||||
feature_dir = SUITE_ROOT / feature_id
|
||||
assert feature_dir.is_dir(), (
|
||||
f"manifest declares {feature_id!r} but {feature_dir} is missing — "
|
||||
"feature_id MUST match its on-disk directory (see manifest.yaml header)."
|
||||
|
|
@ -126,7 +126,7 @@ def test_every_manifest_feature_has_directory(feature_id):
|
|||
|
||||
@pytest.mark.parametrize("feature_id", ALL_FEATURE_IDS)
|
||||
def test_every_manifest_feature_has_init_file(feature_id):
|
||||
init_file = REPO_ROOT / feature_id / "__init__.py"
|
||||
init_file = SUITE_ROOT / feature_id / "__init__.py"
|
||||
assert init_file.is_file(), f"missing __init__.py: {init_file}"
|
||||
|
||||
|
||||
|
|
@ -137,7 +137,7 @@ def test_every_manifest_feature_has_per_provider_test_file(feature_id, provider)
|
|||
backed by a per-provider test file. Without this check, a missing
|
||||
file silently becomes a `not_tested` cell in the published matrix
|
||||
rather than a CI failure surfacing the layout drift."""
|
||||
test_file = REPO_ROOT / feature_id / f"test_{provider}.py"
|
||||
test_file = SUITE_ROOT / feature_id / f"test_{provider}.py"
|
||||
assert test_file.is_file(), f"missing per-provider test file: {test_file}"
|
||||
|
||||
|
||||
|
|
@ -151,7 +151,7 @@ def test_per_provider_test_file_imports_and_parametrizes_three_models(
|
|||
use plain aliases or per-provider-suffixed aliases (e.g.
|
||||
`claude-opus-4-7-bedrock-invoke`), so we check for the tier
|
||||
substrings rather than exact alias names."""
|
||||
text = (REPO_ROOT / feature_id / f"test_{provider}.py").read_text()
|
||||
text = (SUITE_ROOT / feature_id / f"test_{provider}.py").read_text()
|
||||
for tier in ("haiku-4-5", "sonnet-4-6", "opus-4-7"):
|
||||
assert (
|
||||
tier in text
|
||||
|
|
@ -171,7 +171,7 @@ def test_azure_test_file_drives_the_proxy(feature_id):
|
|||
that wraps them — both shapes drive the proxy, and we don't want
|
||||
this layout pin to block legitimate de-duplication of test bodies.
|
||||
"""
|
||||
text = (REPO_ROOT / feature_id / "test_azure.py").read_text()
|
||||
text = (SUITE_ROOT / feature_id / "test_azure.py").read_text()
|
||||
assert "run_claude" in text or "run_basic_messaging_cell" in text, (
|
||||
f"{feature_id}/test_azure.py must drive the claude CLI via run_claude() "
|
||||
"or a shared helper that wraps it; the not_applicable stub was removed "
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from claude_code.cli_driver import (
|
|||
ClaudeCLIError,
|
||||
DriverResult,
|
||||
failure_diagnostic,
|
||||
is_rate_limit_shaped,
|
||||
run_claude,
|
||||
run_claude_models_parallel,
|
||||
)
|
||||
|
|
@ -793,3 +794,199 @@ def test_failure_diagnostic_uses_last_result_event_status():
|
|||
diag = failure_diagnostic(result)
|
||||
assert "api_status=429" in diag
|
||||
assert "500" not in diag
|
||||
|
||||
|
||||
_RATE_LIMITED_STDOUT = (
|
||||
json.dumps(
|
||||
{
|
||||
"type": "assistant",
|
||||
"message": {
|
||||
"content": [
|
||||
{"type": "text", "text": "API Error: 429 Too Many Requests"}
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
+ json.dumps({"type": "result", "api_error_status": 429})
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
_OK_STDOUT = (
|
||||
json.dumps(
|
||||
{
|
||||
"type": "assistant",
|
||||
"message": {"content": [{"type": "text", "text": "pong"}]},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
|
||||
class _FlakyRunner:
|
||||
"""Fake runner that rate-limits each model N times before succeeding.
|
||||
|
||||
Keeps a per-model call count so tests can assert exactly how many
|
||||
attempts the retry loop made — the load-bearing detail a canned
|
||||
single-response runner can't express.
|
||||
"""
|
||||
|
||||
def __init__(self, failures_before_success: dict):
|
||||
self.failures_before_success = dict(failures_before_success)
|
||||
self.calls: dict = {}
|
||||
|
||||
def __call__(self, cmd, env, capture_output, text, timeout, check, input=None):
|
||||
model = cmd[cmd.index("--model") + 1]
|
||||
self.calls[model] = self.calls.get(model, 0) + 1
|
||||
if self.calls[model] <= self.failures_before_success.get(model, 0):
|
||||
return _Completed(returncode=1, stdout=_RATE_LIMITED_STDOUT)
|
||||
return _Completed(returncode=0, stdout=_OK_STDOUT)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"outcome,expected",
|
||||
[
|
||||
(ClaudeCLIError("claude CLI timed out after 120.0s"), True),
|
||||
(ClaudeCLIError("claude CLI not found at 'claude'"), False),
|
||||
(
|
||||
DriverResult(
|
||||
text="",
|
||||
events=[{"type": "result", "api_error_status": 429}],
|
||||
exit_code=1,
|
||||
),
|
||||
True,
|
||||
),
|
||||
(DriverResult(text="Too Many Requests", exit_code=1), True),
|
||||
(DriverResult(text="", stderr="throttled by upstream", exit_code=1), True),
|
||||
(DriverResult(text="rate limit exceeded", exit_code=0), False),
|
||||
(DriverResult(text="", stderr="auth failed", exit_code=2), False),
|
||||
],
|
||||
)
|
||||
def test_is_rate_limit_shaped_classification(outcome, expected):
|
||||
"""The retry trigger must match 429/throttle/timeout markers on
|
||||
failures only — a passing result mentioning '429' in its reply text
|
||||
must never be classified as retryable."""
|
||||
assert is_rate_limit_shaped(outcome) is expected
|
||||
|
||||
|
||||
def test_run_claude_models_parallel_retries_rate_limited_model_until_success():
|
||||
"""A model that 429s once must be retried after the backoff sleep and
|
||||
end up green, while an untroubled sibling model runs exactly once."""
|
||||
runner = _FlakyRunner({"flaky": 1})
|
||||
sleeps: List[float] = []
|
||||
|
||||
outcomes = run_claude_models_parallel(
|
||||
models=["flaky", "steady"],
|
||||
prompt="hi",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
runner=runner,
|
||||
rate_limit_retries=2,
|
||||
rate_limit_backoff_seconds=0.5,
|
||||
sleep=sleeps.append,
|
||||
)
|
||||
|
||||
assert isinstance(outcomes["flaky"], DriverResult)
|
||||
assert outcomes["flaky"].exit_code == 0
|
||||
assert outcomes["flaky"].text == "pong"
|
||||
assert runner.calls == {"flaky": 2, "steady": 1}
|
||||
assert sleeps == [0.5]
|
||||
|
||||
|
||||
def test_run_claude_models_parallel_does_not_retry_non_rate_limit_failures():
|
||||
"""A deterministic failure (bad auth) must fail fast: no sleeps, one
|
||||
attempt — retrying it would just triple the matrix wall time."""
|
||||
|
||||
def runner(cmd, env, capture_output, text, timeout, check, input=None):
|
||||
return _Completed(returncode=2, stdout="", stderr="auth failed")
|
||||
|
||||
sleeps: List[float] = []
|
||||
outcomes = run_claude_models_parallel(
|
||||
models=["a"],
|
||||
prompt="hi",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
runner=runner,
|
||||
rate_limit_retries=2,
|
||||
rate_limit_backoff_seconds=0.5,
|
||||
sleep=sleeps.append,
|
||||
)
|
||||
|
||||
assert outcomes["a"].exit_code == 2
|
||||
assert sleeps == []
|
||||
|
||||
|
||||
def test_run_claude_models_parallel_returns_last_failure_when_retries_exhausted():
|
||||
"""A persistently rate-limited model exhausts its budget (initial
|
||||
attempt + N retries, each preceded by one backoff sleep) and still
|
||||
surfaces the 429 diagnostic instead of masking it."""
|
||||
runner = _FlakyRunner({"stuck": 99})
|
||||
sleeps: List[float] = []
|
||||
|
||||
outcomes = run_claude_models_parallel(
|
||||
models=["stuck"],
|
||||
prompt="hi",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
runner=runner,
|
||||
rate_limit_retries=2,
|
||||
rate_limit_backoff_seconds=0.25,
|
||||
sleep=sleeps.append,
|
||||
)
|
||||
|
||||
assert runner.calls == {"stuck": 3}
|
||||
assert sleeps == [0.25, 0.25]
|
||||
assert outcomes["stuck"].exit_code == 1
|
||||
assert "429" in failure_diagnostic(outcomes["stuck"])
|
||||
|
||||
|
||||
def test_run_claude_models_parallel_retries_timeout_shaped_cli_errors():
|
||||
"""CLI timeouts are how saturated upstreams usually present (the CLI
|
||||
retries 429s internally until the harness kills it), so a timeout
|
||||
must be retried like an explicit 429."""
|
||||
calls: List[int] = []
|
||||
|
||||
def runner(cmd, env, capture_output, text, timeout, check, input=None):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
raise subprocess.TimeoutExpired(cmd="claude", timeout=1)
|
||||
return _Completed(returncode=0, stdout=_OK_STDOUT)
|
||||
|
||||
sleeps: List[float] = []
|
||||
outcomes = run_claude_models_parallel(
|
||||
models=["a"],
|
||||
prompt="hi",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
runner=runner,
|
||||
rate_limit_retries=1,
|
||||
rate_limit_backoff_seconds=0.5,
|
||||
sleep=sleeps.append,
|
||||
)
|
||||
|
||||
assert isinstance(outcomes["a"], DriverResult)
|
||||
assert outcomes["a"].text == "pong"
|
||||
assert len(calls) == 2
|
||||
assert sleeps == [0.5]
|
||||
|
||||
|
||||
def test_run_claude_models_parallel_zero_retries_disables_backoff():
|
||||
"""`rate_limit_retries=0` must restore the old single-attempt
|
||||
behavior exactly: one call, no sleeps, failure returned as-is."""
|
||||
runner = _FlakyRunner({"stuck": 99})
|
||||
sleeps: List[float] = []
|
||||
|
||||
outcomes = run_claude_models_parallel(
|
||||
models=["stuck"],
|
||||
prompt="hi",
|
||||
base_url="http://x",
|
||||
api_key="k",
|
||||
runner=runner,
|
||||
rate_limit_retries=0,
|
||||
rate_limit_backoff_seconds=0.5,
|
||||
sleep=sleeps.append,
|
||||
)
|
||||
|
||||
assert runner.calls == {"stuck": 1}
|
||||
assert sleeps == []
|
||||
assert outcomes["stuck"].exit_code == 1
|
||||
|
|
|
|||
195
tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py
Normal file
195
tests/e2e/claude_code/_driver_unit_tests/test_passthrough.py
Normal file
|
|
@ -0,0 +1,195 @@
|
|||
"""Unit tests for the shared `run_passthrough_cell` helper.
|
||||
|
||||
These tests inject a fake `run_models` callable and an explicit `env`
|
||||
mapping (both are first-class parameters, no monkeypatching), so they
|
||||
exercise the helper's branching -- env-missing guard, base-URL
|
||||
assembly, extra-env forwarding, per-model pass/fail -- without
|
||||
spawning the real CLI.
|
||||
|
||||
The env-builder tests pin the provider-mode contract itself: the
|
||||
CLAUDE_CODE_USE_* / CLAUDE_CODE_SKIP_*_AUTH flags and the passthrough
|
||||
route each mode must target. Those values are the feature -- e.g.
|
||||
dropping the `/v1` from the vertex base URL produces a request Google
|
||||
404s on -- so a mutation to any of them must fail here before it burns
|
||||
a live matrix run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Mapping, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from claude_code._passthrough import (
|
||||
ANTHROPIC_PASSTHROUGH_BASE_PATH,
|
||||
CLIENT_SIDE_AWS_REGION,
|
||||
VERTEX_PLACEHOLDER_PROJECT,
|
||||
VERTEX_PLACEHOLDER_REGION,
|
||||
bedrock_extra_env,
|
||||
foundry_extra_env,
|
||||
run_passthrough_cell,
|
||||
vertex_extra_env,
|
||||
)
|
||||
from claude_code.cli_driver import ClaudeCLIError, DriverResult
|
||||
|
||||
PROXY_ENV = {
|
||||
"LITELLM_PROXY_BASE_URL": "http://localhost:4000",
|
||||
"LITELLM_PROXY_API_KEY": "sk-test",
|
||||
}
|
||||
|
||||
|
||||
class _FakeResult:
|
||||
def __init__(self) -> None:
|
||||
self.rows: List[Dict[str, Any]] = []
|
||||
self.single: Optional[Dict[str, Any]] = None
|
||||
|
||||
def set(self, payload: Mapping[str, Any]) -> None:
|
||||
self.single = dict(payload)
|
||||
|
||||
def add(self, payload: Mapping[str, Any]) -> None:
|
||||
self.rows.append(dict(payload))
|
||||
|
||||
|
||||
def _fake_run_models(outcomes_by_model, captured: Dict[str, Any]):
|
||||
def fake(*, models, prompt, base_url, api_key, extra_env=None, **_kwargs):
|
||||
captured["models"] = list(models)
|
||||
captured["prompt"] = prompt
|
||||
captured["base_url"] = base_url
|
||||
captured["api_key"] = api_key
|
||||
captured["extra_env"] = dict(extra_env) if extra_env is not None else None
|
||||
return {model: outcomes_by_model[model] for model in models}
|
||||
|
||||
return fake
|
||||
|
||||
|
||||
def test_env_missing_guard_reports_fail_and_aborts():
|
||||
fake_result = _FakeResult()
|
||||
with pytest.raises(pytest.fail.Exception):
|
||||
run_passthrough_cell(
|
||||
compat_result=fake_result,
|
||||
models=["claude-haiku-4-5"],
|
||||
prompt="ping",
|
||||
env={},
|
||||
)
|
||||
assert fake_result.single is not None
|
||||
assert fake_result.single["status"] == "fail"
|
||||
assert "LITELLM_PROXY_BASE_URL" in fake_result.single["error"]
|
||||
|
||||
|
||||
def test_anthropic_base_path_appended_to_normalized_proxy_url():
|
||||
fake_result = _FakeResult()
|
||||
captured: Dict[str, Any] = {}
|
||||
outcome = DriverResult(text="pong")
|
||||
|
||||
run_passthrough_cell(
|
||||
compat_result=fake_result,
|
||||
models=["claude-haiku-4-5"],
|
||||
prompt="ping",
|
||||
passthrough_base_path=ANTHROPIC_PASSTHROUGH_BASE_PATH,
|
||||
run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured),
|
||||
env={**PROXY_ENV, "LITELLM_PROXY_BASE_URL": "http://localhost:4000/"},
|
||||
)
|
||||
|
||||
assert captured["base_url"] == "http://localhost:4000/anthropic"
|
||||
assert captured["extra_env"] is None
|
||||
assert fake_result.rows == [{"status": "pass"}]
|
||||
|
||||
|
||||
def test_extra_env_builder_receives_normalized_base_and_is_forwarded():
|
||||
fake_result = _FakeResult()
|
||||
captured: Dict[str, Any] = {}
|
||||
outcome = DriverResult(text="pong")
|
||||
seen_bases: List[str] = []
|
||||
|
||||
def build(proxy_base: str) -> Dict[str, str]:
|
||||
seen_bases.append(proxy_base)
|
||||
return {"SOME_FLAG": "1"}
|
||||
|
||||
run_passthrough_cell(
|
||||
compat_result=fake_result,
|
||||
models=["claude-haiku-4-5"],
|
||||
prompt="ping",
|
||||
build_extra_env=build,
|
||||
run_models=_fake_run_models({"claude-haiku-4-5": outcome}, captured),
|
||||
env={**PROXY_ENV, "LITELLM_PROXY_BASE_URL": "http://localhost:4000/"},
|
||||
)
|
||||
|
||||
assert seen_bases == ["http://localhost:4000"]
|
||||
assert captured["extra_env"] == {"SOME_FLAG": "1"}
|
||||
assert captured["base_url"] == "http://localhost:4000"
|
||||
|
||||
|
||||
def test_per_model_failures_reported_individually():
|
||||
fake_result = _FakeResult()
|
||||
captured: Dict[str, Any] = {}
|
||||
outcomes = {
|
||||
"claude-haiku-4-5": DriverResult(text="pong"),
|
||||
"claude-sonnet-4-6": ClaudeCLIError("claude CLI timed out after 120s"),
|
||||
"claude-opus-4-7": DriverResult(text="", exit_code=1),
|
||||
}
|
||||
|
||||
with pytest.raises(pytest.fail.Exception):
|
||||
run_passthrough_cell(
|
||||
compat_result=fake_result,
|
||||
models=list(outcomes.keys()),
|
||||
prompt="ping",
|
||||
run_models=_fake_run_models(outcomes, captured),
|
||||
env=PROXY_ENV,
|
||||
)
|
||||
|
||||
statuses = [row["status"] for row in fake_result.rows]
|
||||
assert statuses == ["pass", "fail", "fail"]
|
||||
assert "timed out" in fake_result.rows[1]["error"]
|
||||
assert "claude CLI failed" in fake_result.rows[2]["error"]
|
||||
|
||||
|
||||
def test_empty_assistant_text_is_a_fail():
|
||||
fake_result = _FakeResult()
|
||||
captured: Dict[str, Any] = {}
|
||||
outcomes = {"claude-haiku-4-5": DriverResult(text=" ")}
|
||||
|
||||
with pytest.raises(pytest.fail.Exception):
|
||||
run_passthrough_cell(
|
||||
compat_result=fake_result,
|
||||
models=["claude-haiku-4-5"],
|
||||
prompt="ping",
|
||||
run_models=_fake_run_models(outcomes, captured),
|
||||
env=PROXY_ENV,
|
||||
)
|
||||
|
||||
assert fake_result.rows == [
|
||||
{
|
||||
"status": "fail",
|
||||
"error": "[claude-haiku-4-5] claude returned empty assistant text",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_bedrock_extra_env_targets_proxy_bedrock_route():
|
||||
env = bedrock_extra_env("http://localhost:4000")
|
||||
assert env == {
|
||||
"CLAUDE_CODE_USE_BEDROCK": "1",
|
||||
"CLAUDE_CODE_SKIP_BEDROCK_AUTH": "1",
|
||||
"ANTHROPIC_BEDROCK_BASE_URL": "http://localhost:4000/bedrock",
|
||||
"AWS_REGION": CLIENT_SIDE_AWS_REGION,
|
||||
}
|
||||
|
||||
|
||||
def test_vertex_extra_env_keeps_the_api_version_in_the_base_url():
|
||||
env = vertex_extra_env("http://localhost:4000")
|
||||
assert env == {
|
||||
"CLAUDE_CODE_USE_VERTEX": "1",
|
||||
"CLAUDE_CODE_SKIP_VERTEX_AUTH": "1",
|
||||
"ANTHROPIC_VERTEX_BASE_URL": "http://localhost:4000/vertex_ai/v1",
|
||||
"ANTHROPIC_VERTEX_PROJECT_ID": VERTEX_PLACEHOLDER_PROJECT,
|
||||
"CLOUD_ML_REGION": VERTEX_PLACEHOLDER_REGION,
|
||||
}
|
||||
|
||||
|
||||
def test_foundry_extra_env_targets_proxy_azure_route():
|
||||
env = foundry_extra_env("http://localhost:4000")
|
||||
assert env == {
|
||||
"CLAUDE_CODE_USE_FOUNDRY": "1",
|
||||
"CLAUDE_CODE_SKIP_FOUNDRY_AUTH": "1",
|
||||
"ANTHROPIC_FOUNDRY_BASE_URL": "http://localhost:4000/azure",
|
||||
}
|
||||
196
tests/e2e/claude_code/_passthrough.py
Normal file
196
tests/e2e/claude_code/_passthrough.py
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
"""Shared body for the `passthrough` × <provider> compat cells.
|
||||
|
||||
Every other matrix row drives the proxy's `/v1/messages` translation
|
||||
layer: Claude Code speaks the first-party Anthropic wire and LiteLLM
|
||||
transforms the request per provider. This row instead exercises
|
||||
LiteLLM's *native passthrough* routes -- the "LLM gateway"
|
||||
configuration documented at https://code.claude.com/docs/en/gateway --
|
||||
where Claude Code speaks each cloud's own wire format and the proxy
|
||||
forwards it, attaching provider credentials on the way out:
|
||||
|
||||
anthropic ANTHROPIC_BASE_URL={proxy}/anthropic. The CLI's
|
||||
first-party wire, forwarded verbatim to
|
||||
api.anthropic.com, so the model ids are real
|
||||
Anthropic ids rather than proxy aliases.
|
||||
bedrock_invoke CLAUDE_CODE_USE_BEDROCK=1 +
|
||||
ANTHROPIC_BEDROCK_BASE_URL={proxy}/bedrock. The
|
||||
CLI POSTs /model/{model}/invoke-with-response-stream;
|
||||
the proxy recognizes a router alias in the model
|
||||
segment, rewrites it to the deployment's upstream
|
||||
model id, and SigV4-signs with its own AWS creds.
|
||||
vertex_ai CLAUDE_CODE_USE_VERTEX=1 +
|
||||
ANTHROPIC_VERTEX_BASE_URL={proxy}/vertex_ai/v1.
|
||||
The CLI POSTs
|
||||
.../models/{model}:streamRawPredict; the proxy
|
||||
resolves a router alias in the model segment and
|
||||
takes project, location, and credentials from the
|
||||
deployment (which is why the deployment must set
|
||||
`use_in_pass_through: true` -- see
|
||||
test_config.yaml).
|
||||
azure CLAUDE_CODE_USE_FOUNDRY=1 +
|
||||
ANTHROPIC_FOUNDRY_BASE_URL={proxy}/azure. Foundry
|
||||
mode sends the model in the JSON body, not the
|
||||
URL, so the proxy's /azure route cannot resolve a
|
||||
router alias and falls back to the env-configured
|
||||
AZURE_API_BASE / AZURE_API_KEY target.
|
||||
bedrock_converse not applicable -- Claude Code's bedrock mode is
|
||||
InvokeModel-only; no Converse-wire client exists.
|
||||
|
||||
Auth is the same in every mode: the CLI's provider-native signing is
|
||||
disabled via CLAUDE_CODE_SKIP_<PROVIDER>_AUTH, and the LiteLLM virtual
|
||||
key travels as `Authorization: Bearer` (ANTHROPIC_AUTH_TOKEN), exactly
|
||||
like the translation rows. The proxy holds the real provider
|
||||
credentials.
|
||||
|
||||
The per-mode env vars and URL shapes above were captured from a real
|
||||
`claude` CLI (2.1.210) run against a request-logging sink, not from
|
||||
docs; if a CLI release changes them, the cells fail with the CLI's own
|
||||
diagnostic rather than silently testing the wrong wire.
|
||||
|
||||
`run_models` and `env` are injection seams for
|
||||
`_driver_unit_tests/test_passthrough.py`; production callers leave
|
||||
them unset.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any, Callable, Dict, Mapping, Optional, Sequence
|
||||
|
||||
import pytest
|
||||
|
||||
from claude_code.cli_driver import (
|
||||
ClaudeCLIError,
|
||||
failure_diagnostic,
|
||||
run_claude_models_parallel,
|
||||
)
|
||||
|
||||
PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL"
|
||||
PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY"
|
||||
|
||||
ANTHROPIC_PASSTHROUGH_BASE_PATH = "/anthropic"
|
||||
|
||||
CLIENT_SIDE_AWS_REGION = "us-east-1"
|
||||
"""Satisfies the CLI's embedded AWS SDK, which refuses to construct a
|
||||
client without a region. The value never influences routing: the proxy
|
||||
signs the upstream request with its own credentials and region."""
|
||||
|
||||
VERTEX_PLACEHOLDER_PROJECT = "proxy-resolved-project"
|
||||
VERTEX_PLACEHOLDER_REGION = "us-east5"
|
||||
"""The CLI refuses to build a Vertex URL without a project id and
|
||||
region, but the proxy replaces both path segments with the resolved
|
||||
deployment's `vertex_project` / `vertex_location` before forwarding,
|
||||
so deliberately-fake values prove the resolution actually happened."""
|
||||
|
||||
|
||||
def bedrock_extra_env(proxy_base_url: str) -> Dict[str, str]:
|
||||
return {
|
||||
"CLAUDE_CODE_USE_BEDROCK": "1",
|
||||
"CLAUDE_CODE_SKIP_BEDROCK_AUTH": "1",
|
||||
"ANTHROPIC_BEDROCK_BASE_URL": f"{proxy_base_url}/bedrock",
|
||||
"AWS_REGION": CLIENT_SIDE_AWS_REGION,
|
||||
}
|
||||
|
||||
|
||||
def vertex_extra_env(proxy_base_url: str) -> Dict[str, str]:
|
||||
"""Vertex-mode CLI env pointed at the proxy's /vertex_ai route.
|
||||
|
||||
The `/v1` suffix on ANTHROPIC_VERTEX_BASE_URL is load-bearing: the
|
||||
CLI's Vertex SDK ships its API version inside its *default* base
|
||||
URL (`https://{region}-aiplatform.googleapis.com/v1`), so
|
||||
overriding the base drops the version from the request path unless
|
||||
the override carries it. LiteLLM's /vertex_ai route reuses the
|
||||
incoming path verbatim when it contains `/projects/.../locations/...`,
|
||||
so a version-less path would reach Google as
|
||||
`aiplatform.googleapis.com/projects/...` and 404.
|
||||
"""
|
||||
return {
|
||||
"CLAUDE_CODE_USE_VERTEX": "1",
|
||||
"CLAUDE_CODE_SKIP_VERTEX_AUTH": "1",
|
||||
"ANTHROPIC_VERTEX_BASE_URL": f"{proxy_base_url}/vertex_ai/v1",
|
||||
"ANTHROPIC_VERTEX_PROJECT_ID": VERTEX_PLACEHOLDER_PROJECT,
|
||||
"CLOUD_ML_REGION": VERTEX_PLACEHOLDER_REGION,
|
||||
}
|
||||
|
||||
|
||||
def foundry_extra_env(proxy_base_url: str) -> Dict[str, str]:
|
||||
return {
|
||||
"CLAUDE_CODE_USE_FOUNDRY": "1",
|
||||
"CLAUDE_CODE_SKIP_FOUNDRY_AUTH": "1",
|
||||
"ANTHROPIC_FOUNDRY_BASE_URL": f"{proxy_base_url}/azure",
|
||||
}
|
||||
|
||||
|
||||
def run_passthrough_cell(
|
||||
*,
|
||||
compat_result,
|
||||
models: Sequence[str],
|
||||
prompt: str,
|
||||
passthrough_base_path: str = "",
|
||||
build_extra_env: Optional[Callable[[str], Mapping[str, str]]] = None,
|
||||
run_models: Callable[..., Mapping[str, Any]] = run_claude_models_parallel,
|
||||
env: Optional[Mapping[str, str]] = None,
|
||||
) -> None:
|
||||
"""Run the shared `passthrough` × <provider> cell body.
|
||||
|
||||
`passthrough_base_path` is appended to the proxy base URL and
|
||||
becomes the CLI's ANTHROPIC_BASE_URL (only the anthropic column
|
||||
uses it; the cloud columns ignore ANTHROPIC_BASE_URL entirely once
|
||||
their CLAUDE_CODE_USE_* flag is set). `build_extra_env` receives
|
||||
the trailing-slash-normalized proxy base URL and returns the
|
||||
provider-mode env for the CLI subprocess.
|
||||
"""
|
||||
environ = env if env is not None else os.environ
|
||||
base_url = environ.get(PROXY_BASE_URL_ENV)
|
||||
api_key = environ.get(PROXY_API_KEY_ENV)
|
||||
if not base_url or not api_key:
|
||||
compat_result.set(
|
||||
{
|
||||
"status": "fail",
|
||||
"error": (
|
||||
f"missing required env: set {PROXY_BASE_URL_ENV} and "
|
||||
f"{PROXY_API_KEY_ENV} to point at a running LiteLLM proxy"
|
||||
),
|
||||
}
|
||||
)
|
||||
pytest.fail(
|
||||
f"{PROXY_BASE_URL_ENV} / {PROXY_API_KEY_ENV} not configured",
|
||||
pytrace=False,
|
||||
)
|
||||
|
||||
proxy_base = base_url.rstrip("/")
|
||||
extra_env = dict(build_extra_env(proxy_base)) if build_extra_env else None
|
||||
|
||||
outcomes = run_models(
|
||||
models=models,
|
||||
prompt=prompt,
|
||||
base_url=proxy_base + passthrough_base_path,
|
||||
api_key=api_key,
|
||||
extra_env=extra_env,
|
||||
)
|
||||
|
||||
failures = []
|
||||
for model in models:
|
||||
outcome = outcomes[model]
|
||||
if isinstance(outcome, ClaudeCLIError):
|
||||
error = f"[{model}] {outcome}"
|
||||
compat_result.add({"status": "fail", "error": error})
|
||||
failures.append(error)
|
||||
continue
|
||||
|
||||
if outcome.exit_code != 0:
|
||||
error = f"[{model}] claude CLI failed: {failure_diagnostic(outcome)}"
|
||||
compat_result.add({"status": "fail", "error": error})
|
||||
failures.append(error)
|
||||
continue
|
||||
|
||||
if not outcome.text.strip():
|
||||
error = f"[{model}] claude returned empty assistant text"
|
||||
compat_result.add({"status": "fail", "error": error})
|
||||
failures.append(error)
|
||||
continue
|
||||
|
||||
compat_result.add({"status": "pass"})
|
||||
|
||||
if failures:
|
||||
pytest.fail("; ".join(failures), pytrace=False)
|
||||
|
|
@ -15,6 +15,7 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
|
|
@ -40,6 +41,29 @@ DEFAULT_TIMEOUT_SECONDS = float(
|
|||
os.environ.get("LITELLM_COMPAT_CLI_TIMEOUT_SECONDS") or 120
|
||||
)
|
||||
|
||||
RATE_LIMIT_SHAPED_RE = re.compile(
|
||||
r"(?:\b429\b|rate[\s_-]?limit|too\s+many\s+requests|throttl(?:ed|ing)|"
|
||||
r"claude\s+CLI\s+timed\s+out)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
"""Heuristic shared with the conftest rate-limit summary: 429s and
|
||||
throttle markers anywhere in the failure text, plus CLI timeouts --
|
||||
the CLI retries 429s internally until the harness timeout kills it,
|
||||
so a saturated upstream usually surfaces as a timeout rather than a
|
||||
clean 429."""
|
||||
|
||||
DEFAULT_RATE_LIMIT_RETRIES = int(
|
||||
os.environ.get("LITELLM_COMPAT_RATE_LIMIT_RETRIES") or 2
|
||||
)
|
||||
DEFAULT_RATE_LIMIT_BACKOFF_SECONDS = float(
|
||||
os.environ.get("LITELLM_COMPAT_RATE_LIMIT_BACKOFF_SECONDS") or 65
|
||||
)
|
||||
"""Rate-limit-shaped failures are retried after a backoff long enough
|
||||
for a per-minute quota window (the dominant 429 source across
|
||||
Anthropic / Bedrock / Vertex) to reset. Both knobs are env-tunable so
|
||||
a matrix run can trade wall time for resilience without code edits;
|
||||
retries=0 disables the behavior entirely."""
|
||||
|
||||
# Env vars the `claude` Node CLI legitimately needs to function:
|
||||
# locating its own binary + node, basic locale/terminal plumbing.
|
||||
# Deliberately excludes every credential-bearing var that the
|
||||
|
|
@ -265,6 +289,22 @@ def run_claude(
|
|||
ModelResult = Union[DriverResult, ClaudeCLIError]
|
||||
|
||||
|
||||
def is_rate_limit_shaped(outcome: ModelResult) -> bool:
|
||||
"""Classify an outcome as a retryable rate-limit-shaped failure.
|
||||
|
||||
A `ClaudeCLIError` matches on its message (which is where the
|
||||
driver's own timeout diagnostic lands); a failing `DriverResult`
|
||||
matches on its full `failure_diagnostic` so 429s buried in the
|
||||
CLI's stdout text or `api_error_status` are both caught. Passing
|
||||
results are never rate-limit-shaped.
|
||||
"""
|
||||
if isinstance(outcome, ClaudeCLIError):
|
||||
return bool(RATE_LIMIT_SHAPED_RE.search(str(outcome)))
|
||||
if outcome.exit_code == 0:
|
||||
return False
|
||||
return bool(RATE_LIMIT_SHAPED_RE.search(failure_diagnostic(outcome)))
|
||||
|
||||
|
||||
def run_claude_models_parallel(
|
||||
*,
|
||||
models: Sequence[str],
|
||||
|
|
@ -277,6 +317,9 @@ def run_claude_models_parallel(
|
|||
cli_path: str = CLAUDE_CLI_DEFAULT,
|
||||
timeout: float = DEFAULT_TIMEOUT_SECONDS,
|
||||
runner: Optional[Callable[..., Any]] = None,
|
||||
rate_limit_retries: Optional[int] = None,
|
||||
rate_limit_backoff_seconds: Optional[float] = None,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
) -> Dict[str, ModelResult]:
|
||||
"""Invoke `run_claude` for every `models[i]` concurrently and collect outcomes.
|
||||
|
||||
|
|
@ -290,6 +333,14 @@ def run_claude_models_parallel(
|
|||
keep the synchronous CLI driver unchanged so unit tests can keep
|
||||
injecting a fake `runner`.
|
||||
|
||||
Rate-limit-shaped failures (see `is_rate_limit_shaped`) are retried
|
||||
per model up to `rate_limit_retries` times, sleeping
|
||||
`rate_limit_backoff_seconds` before each retry so per-minute quota
|
||||
windows can reset; both default to the `LITELLM_COMPAT_RATE_LIMIT_*`
|
||||
env knobs. Each retry goes back through `run_claude`, so it
|
||||
re-acquires a token from the provider rate limiter like any other
|
||||
invocation. `sleep` is an injection seam for unit tests.
|
||||
|
||||
Returns a dict keyed by model id. Each value is either the
|
||||
`DriverResult` produced by `run_claude` or the `ClaudeCLIError`
|
||||
that aborted that model's run — callers decide how to map either
|
||||
|
|
@ -300,14 +351,20 @@ def run_claude_models_parallel(
|
|||
if not models:
|
||||
raise ValueError("models must be a non-empty sequence")
|
||||
|
||||
def _one(model: str) -> Tuple[str, ModelResult, float]:
|
||||
# Per-model wall clock: this is what the matrix run actually pays for.
|
||||
# We record it whether the run succeeded or raised so the breakdown
|
||||
# log below covers both code paths and surfaces "which model is the
|
||||
# long pole?" without requiring per-test instrumentation.
|
||||
started = time.monotonic()
|
||||
retries = (
|
||||
DEFAULT_RATE_LIMIT_RETRIES
|
||||
if rate_limit_retries is None
|
||||
else max(0, rate_limit_retries)
|
||||
)
|
||||
backoff = (
|
||||
DEFAULT_RATE_LIMIT_BACKOFF_SECONDS
|
||||
if rate_limit_backoff_seconds is None
|
||||
else max(0.0, rate_limit_backoff_seconds)
|
||||
)
|
||||
|
||||
def _run_once(model: str) -> ModelResult:
|
||||
try:
|
||||
result = run_claude(
|
||||
return run_claude(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
|
|
@ -319,14 +376,8 @@ def run_claude_models_parallel(
|
|||
timeout=timeout,
|
||||
runner=runner,
|
||||
)
|
||||
elapsed = time.monotonic() - started
|
||||
# Stamp the duration onto the DriverResult so callers (tests,
|
||||
# diagnostics) can attribute slow cells without re-timing.
|
||||
result.duration_ms = int(elapsed * 1000)
|
||||
return model, result, elapsed
|
||||
except ClaudeCLIError as exc:
|
||||
elapsed = time.monotonic() - started
|
||||
return model, exc, elapsed
|
||||
return exc
|
||||
except Exception as exc:
|
||||
# Honor the documented "errors as values" contract for any
|
||||
# exception type — not just ClaudeCLIError. The rate
|
||||
|
|
@ -334,13 +385,38 @@ def run_claude_models_parallel(
|
|||
# raise ValueError on edge-case model strings, and a future
|
||||
# bug elsewhere in the call stack must not abort the entire
|
||||
# parallel batch and lose the other models' outcomes.
|
||||
elapsed = time.monotonic() - started
|
||||
wrapped = ClaudeCLIError(
|
||||
f"unexpected error running model {model!r}: "
|
||||
f"{type(exc).__name__}: {exc}"
|
||||
)
|
||||
wrapped.__cause__ = exc
|
||||
return model, wrapped, elapsed
|
||||
return wrapped
|
||||
|
||||
def _one(model: str) -> Tuple[str, ModelResult, float]:
|
||||
# Per-model wall clock: this is what the matrix run actually pays
|
||||
# for, retries and backoff sleeps included. We record it whether
|
||||
# the run succeeded or raised so the breakdown log below covers
|
||||
# both code paths and surfaces "which model is the long pole?"
|
||||
# without requiring per-test instrumentation.
|
||||
started = time.monotonic()
|
||||
outcome = _run_once(model)
|
||||
for attempt in range(retries):
|
||||
if not is_rate_limit_shaped(outcome):
|
||||
break
|
||||
print(
|
||||
f"[retry] {model}: rate-limit-shaped failure; sleeping "
|
||||
f"{backoff:.0f}s before attempt {attempt + 2}/{retries + 1}",
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
sleep(backoff)
|
||||
outcome = _run_once(model)
|
||||
elapsed = time.monotonic() - started
|
||||
if isinstance(outcome, DriverResult):
|
||||
# Stamp the duration onto the DriverResult so callers (tests,
|
||||
# diagnostics) can attribute slow cells without re-timing.
|
||||
outcome.duration_ms = int(elapsed * 1000)
|
||||
return model, outcome, elapsed
|
||||
|
||||
outcomes: Dict[str, ModelResult] = {}
|
||||
durations: Dict[str, float] = {}
|
||||
|
|
|
|||
|
|
@ -33,7 +33,6 @@ from __future__ import annotations
|
|||
import functools
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from collections import Counter, defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
|
|
@ -43,6 +42,8 @@ from typing import Any, Dict, FrozenSet, List, Optional, Tuple
|
|||
import pytest
|
||||
import yaml
|
||||
|
||||
from claude_code.cli_driver import RATE_LIMIT_SHAPED_RE
|
||||
|
||||
VALID_STATUSES = {"pass", "fail", "not_applicable", "not_tested"}
|
||||
RESULTS_ARTIFACT_ENV = "COMPAT_RESULTS_PATH"
|
||||
DEFAULT_ARTIFACT_PATH = "compat-results.json"
|
||||
|
|
@ -62,11 +63,10 @@ DEFAULT_RATE_LIMIT_SUMMARY_PATH = "compat-rate-limit-summary.json"
|
|||
# the rate limiter is supposed to back off from. False positives on a
|
||||
# genuinely slow upstream are tolerable here because the worst case is
|
||||
# the binary search runs at a slightly lower rate than necessary.
|
||||
_RATE_LIMIT_RE = re.compile(
|
||||
r"(?:\b429\b|rate[\s_-]?limit|too\s+many\s+requests|throttl(?:ed|ing)|"
|
||||
r"claude\s+CLI\s+timed\s+out)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
#
|
||||
# The pattern lives in `cli_driver` so the driver's retry-on-rate-limit
|
||||
# logic and this summary classify failures identically.
|
||||
_RATE_LIMIT_RE = RATE_LIMIT_SHAPED_RE
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
|
|||
|
|
@ -27,6 +27,15 @@ VERTEXAI_LOCATION=global
|
|||
AZURE_FOUNDRY_API_KEY=
|
||||
AZURE_FOUNDRY_API_BASE=
|
||||
|
||||
# Azure cell of the `passthrough` row. Foundry-mode Claude Code sends
|
||||
# the model in the request body, so the proxy's /azure passthrough
|
||||
# cannot resolve a router alias and falls back to these env vars.
|
||||
# AZURE_API_BASE is the Foundry resource's Anthropic surface, i.e.
|
||||
# https://<resource>.services.ai.azure.com/anthropic ; AZURE_API_KEY
|
||||
# is the same key as AZURE_FOUNDRY_API_KEY.
|
||||
AZURE_API_BASE=
|
||||
AZURE_API_KEY=
|
||||
|
||||
# REQUIRED for publishing: PAT for the `agent-shin` user, used to push
|
||||
# the daily compat-matrix branch to its fork (agent-shin/litellm-docs)
|
||||
# and open the cross-repo PR against BerriAI/litellm-docs. Scopes:
|
||||
|
|
|
|||
|
|
@ -91,6 +91,23 @@ features:
|
|||
# Code releases. The HTTP probe hits the bug surface LiteLLM
|
||||
# has actually shipped fixes for (2.1.117, 2.1.72, 2.1.70 per
|
||||
# the Claude Code release notes).
|
||||
- id: passthrough
|
||||
name: Native API passthrough
|
||||
# Drives the CLI in each cloud's native mode against LiteLLM's
|
||||
# passthrough routes instead of the /v1/messages translation
|
||||
# layer -- the "LLM gateway" setup from
|
||||
# https://code.claude.com/docs/en/gateway. anthropic uses
|
||||
# ANTHROPIC_BASE_URL={proxy}/anthropic; bedrock_invoke uses
|
||||
# CLAUDE_CODE_USE_BEDROCK=1 against {proxy}/bedrock (InvokeModel
|
||||
# wire, alias resolved from the URL by the router); vertex_ai
|
||||
# uses CLAUDE_CODE_USE_VERTEX=1 against {proxy}/vertex_ai/v1
|
||||
# (rawPredict wire, alias + project + location resolved from the
|
||||
# deployment, which therefore needs `use_in_pass_through: true`);
|
||||
# azure uses CLAUDE_CODE_USE_FOUNDRY=1 against {proxy}/azure and
|
||||
# needs AZURE_API_BASE/AZURE_API_KEY on the proxy (see
|
||||
# passthrough/test_azure.py and the cron env example).
|
||||
# bedrock_converse is structurally not_applicable: Claude Code
|
||||
# has no Converse-wire client.
|
||||
- id: long_context_1m
|
||||
name: Long context (1M)
|
||||
# Sends a ~210k-token padded prompt with the
|
||||
|
|
|
|||
0
tests/e2e/claude_code/passthrough/__init__.py
Normal file
0
tests/e2e/claude_code/passthrough/__init__.py
Normal file
44
tests/e2e/claude_code/passthrough/test_anthropic.py
Normal file
44
tests/e2e/claude_code/passthrough/test_anthropic.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
"""passthrough x Anthropic.
|
||||
|
||||
Drive the real `claude` CLI in its default first-party mode, but with
|
||||
ANTHROPIC_BASE_URL aimed at the proxy's `/anthropic` passthrough route
|
||||
instead of the `/v1/messages` translation endpoint. The proxy forwards
|
||||
the request verbatim to api.anthropic.com, swapping the virtual-key
|
||||
bearer for its own ANTHROPIC_API_KEY.
|
||||
|
||||
The (feature, provider) for this cell is inferred from the file path by
|
||||
`tests/e2e/claude_code/conftest.py`:
|
||||
|
||||
tests/e2e/claude_code/passthrough/test_anthropic.py
|
||||
^^^^^^^^^^^ ^^^^^^^^^
|
||||
feature_id provider
|
||||
|
||||
Because nothing is translated, the model ids are the real Anthropic API
|
||||
ids (which happen to equal the proxy aliases for this column). A red
|
||||
cell here means the passthrough route broke forwarding itself -- auth
|
||||
header swap, streaming SSE relay, or beta-header propagation -- since
|
||||
no per-provider transformation is involved.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from claude_code._passthrough import (
|
||||
ANTHROPIC_PASSTHROUGH_BASE_PATH,
|
||||
run_passthrough_cell,
|
||||
)
|
||||
|
||||
ANTHROPIC_MODELS = [
|
||||
"claude-haiku-4-5",
|
||||
"claude-sonnet-4-6",
|
||||
"claude-opus-4-7",
|
||||
]
|
||||
|
||||
|
||||
def test_passthrough_anthropic(compat_result):
|
||||
"""Drive the `claude` CLI through `{proxy}/anthropic` and assert a reply."""
|
||||
run_passthrough_cell(
|
||||
compat_result=compat_result,
|
||||
models=ANTHROPIC_MODELS,
|
||||
prompt="Reply with the single word 'pong' and nothing else.",
|
||||
passthrough_base_path=ANTHROPIC_PASSTHROUGH_BASE_PATH,
|
||||
)
|
||||
59
tests/e2e/claude_code/passthrough/test_azure.py
Normal file
59
tests/e2e/claude_code/passthrough/test_azure.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
"""passthrough x Azure (Microsoft Foundry).
|
||||
|
||||
Drive the real `claude` CLI in foundry mode (CLAUDE_CODE_USE_FOUNDRY=1)
|
||||
with ANTHROPIC_FOUNDRY_BASE_URL aimed at the proxy's `/azure`
|
||||
passthrough route. The CLI POSTs `/v1/messages` with the model in the
|
||||
JSON body -- unlike the bedrock/vertex modes there is no model segment
|
||||
in the URL, so the proxy's router-alias resolution cannot engage and
|
||||
the `/azure` route falls back to its env-configured target: the proxy
|
||||
must set AZURE_API_BASE to the Foundry resource's Anthropic surface
|
||||
(`https://<resource>.services.ai.azure.com/anthropic`) and
|
||||
AZURE_API_KEY to the Foundry key (see
|
||||
cron_vm/litellm-compat-matrix.env.example). The model ids are the
|
||||
Foundry deployment names, which this matrix provisions to match the
|
||||
Anthropic ids.
|
||||
|
||||
The (feature, provider) for this cell is inferred from the file path by
|
||||
`tests/e2e/claude_code/conftest.py`:
|
||||
|
||||
tests/e2e/claude_code/passthrough/test_azure.py
|
||||
^^^^^^^^^^^ ^^^^^
|
||||
feature_id provider
|
||||
|
||||
Wiring verified live at authoring time: through `{proxy}/azure` the
|
||||
Foundry Anthropic surface accepted the `api-key` / `Authorization:
|
||||
Bearer` headers the fallback sends (a bogus key 401s, the real key
|
||||
proceeds to deployment lookup), so a red cell here means missing
|
||||
AZURE_API_BASE/AZURE_API_KEY on the proxy, missing Foundry deployments
|
||||
for the three tiers, or a genuine forwarding gap -- not an auth-scheme
|
||||
mismatch.
|
||||
|
||||
Known-red at authoring time against a healthy Foundry resource: the
|
||||
`/azure` fallback assembles only its own auth headers and drops the
|
||||
rest of the client's headers, including the `anthropic-version` header
|
||||
the CLI sends, and Foundry's Anthropic surface rejects the request
|
||||
with 400 "anthropic-version: header is required" (the same request
|
||||
sent directly to Foundry with that header succeeds). This cell stays
|
||||
red until that forwarding gap is fixed, which is precisely the class
|
||||
of bug the row exists to surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from claude_code._passthrough import foundry_extra_env, run_passthrough_cell
|
||||
|
||||
AZURE_MODELS = [
|
||||
"claude-haiku-4-5",
|
||||
"claude-sonnet-4-6",
|
||||
"claude-opus-4-7",
|
||||
]
|
||||
|
||||
|
||||
def test_passthrough_azure(compat_result):
|
||||
"""Drive the `claude` CLI through `{proxy}/azure` and assert a reply."""
|
||||
run_passthrough_cell(
|
||||
compat_result=compat_result,
|
||||
models=AZURE_MODELS,
|
||||
prompt="Reply with the single word 'pong' and nothing else.",
|
||||
build_extra_env=foundry_extra_env,
|
||||
)
|
||||
34
tests/e2e/claude_code/passthrough/test_bedrock_converse.py
Normal file
34
tests/e2e/claude_code/passthrough/test_bedrock_converse.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
"""passthrough x Bedrock (Converse).
|
||||
|
||||
Structurally not applicable. In bedrock mode the `claude` CLI speaks
|
||||
only the InvokeModel wire (`/model/{id}/invoke-with-response-stream`);
|
||||
it has no Converse-wire client, so there is no Claude Code traffic a
|
||||
Converse passthrough could serve. LiteLLM's `/bedrock` route does
|
||||
accept `/model/{id}/converse-stream`, but exercising it would test a
|
||||
wire no Claude Code user can produce, which is out of scope for this
|
||||
matrix.
|
||||
|
||||
The (feature, provider) for this cell is inferred from the file path by
|
||||
`tests/e2e/claude_code/conftest.py`:
|
||||
|
||||
tests/e2e/claude_code/passthrough/test_bedrock_converse.py
|
||||
^^^^^^^^^^^ ^^^^^^^^^^^^^^^^
|
||||
feature_id provider
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def test_passthrough_bedrock_converse(compat_result):
|
||||
"""Report not_applicable: Claude Code has no Converse-wire mode."""
|
||||
compat_result.set(
|
||||
{
|
||||
"status": "not_applicable",
|
||||
"reason": (
|
||||
"Claude Code's bedrock mode speaks only the InvokeModel wire "
|
||||
"(/model/{id}/invoke-with-response-stream); it has no "
|
||||
"Converse-wire client, so there is no Claude Code surface "
|
||||
"for Converse passthrough."
|
||||
),
|
||||
}
|
||||
)
|
||||
42
tests/e2e/claude_code/passthrough/test_bedrock_invoke.py
Normal file
42
tests/e2e/claude_code/passthrough/test_bedrock_invoke.py
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
"""passthrough x Bedrock (Invoke).
|
||||
|
||||
Drive the real `claude` CLI in bedrock mode (CLAUDE_CODE_USE_BEDROCK=1)
|
||||
with ANTHROPIC_BEDROCK_BASE_URL aimed at the proxy's `/bedrock`
|
||||
passthrough route. The CLI speaks the native InvokeModel wire --
|
||||
`POST /model/{model}/invoke-with-response-stream` -- with the proxy
|
||||
alias in the model segment; the proxy resolves the alias through its
|
||||
router, rewrites the path to the deployment's upstream model id, and
|
||||
SigV4-signs the forwarded request with its own AWS credentials
|
||||
(CLAUDE_CODE_SKIP_BEDROCK_AUTH=1 keeps the CLI from signing).
|
||||
|
||||
The (feature, provider) for this cell is inferred from the file path by
|
||||
`tests/e2e/claude_code/conftest.py`:
|
||||
|
||||
tests/e2e/claude_code/passthrough/test_bedrock_invoke.py
|
||||
^^^^^^^^^^^ ^^^^^^^^^^^^^^
|
||||
feature_id provider
|
||||
|
||||
The CLI also fires a best-effort `GET /bedrock/inference-profiles`
|
||||
listing at startup; its failure is non-fatal and does not gate this
|
||||
cell.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from claude_code._passthrough import bedrock_extra_env, run_passthrough_cell
|
||||
|
||||
BEDROCK_INVOKE_MODELS = [
|
||||
"claude-haiku-4-5-bedrock-invoke",
|
||||
"claude-sonnet-4-6-bedrock-invoke",
|
||||
"claude-opus-4-7-bedrock-invoke",
|
||||
]
|
||||
|
||||
|
||||
def test_passthrough_bedrock_invoke(compat_result):
|
||||
"""Drive the `claude` CLI through `{proxy}/bedrock` and assert a reply."""
|
||||
run_passthrough_cell(
|
||||
compat_result=compat_result,
|
||||
models=BEDROCK_INVOKE_MODELS,
|
||||
prompt="Reply with the single word 'pong' and nothing else.",
|
||||
build_extra_env=bedrock_extra_env,
|
||||
)
|
||||
45
tests/e2e/claude_code/passthrough/test_vertex_ai.py
Normal file
45
tests/e2e/claude_code/passthrough/test_vertex_ai.py
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
"""passthrough x Vertex AI.
|
||||
|
||||
Drive the real `claude` CLI in vertex mode (CLAUDE_CODE_USE_VERTEX=1)
|
||||
with ANTHROPIC_VERTEX_BASE_URL aimed at the proxy's `/vertex_ai`
|
||||
passthrough route. The CLI speaks the native rawPredict wire --
|
||||
`POST .../projects/{p}/locations/{l}/publishers/anthropic/models/{model}:streamRawPredict`
|
||||
-- with the proxy alias in the model segment; the proxy resolves the
|
||||
alias through its router, replaces the placeholder project/location
|
||||
path segments with the deployment's `vertex_project` /
|
||||
`vertex_location`, and attaches its own Google credentials
|
||||
(CLAUDE_CODE_SKIP_VERTEX_AUTH=1 keeps the CLI from minting a token).
|
||||
|
||||
The (feature, provider) for this cell is inferred from the file path by
|
||||
`tests/e2e/claude_code/conftest.py`:
|
||||
|
||||
tests/e2e/claude_code/passthrough/test_vertex_ai.py
|
||||
^^^^^^^^^^^ ^^^^^^^^^
|
||||
feature_id provider
|
||||
|
||||
This cell requires the vertex deployments in the proxy config to carry
|
||||
`use_in_pass_through: true` (see test_config.yaml) -- that is what
|
||||
registers their credentials with the passthrough router. Without it
|
||||
the proxy forwards the CLI's own headers (the virtual-key bearer) to
|
||||
Google and every tier fails with a 401.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from claude_code._passthrough import run_passthrough_cell, vertex_extra_env
|
||||
|
||||
VERTEX_MODELS = [
|
||||
"claude-haiku-4-5-vertex",
|
||||
"claude-sonnet-4-6-vertex",
|
||||
"claude-opus-4-7-vertex",
|
||||
]
|
||||
|
||||
|
||||
def test_passthrough_vertex_ai(compat_result):
|
||||
"""Drive the `claude` CLI through `{proxy}/vertex_ai` and assert a reply."""
|
||||
run_passthrough_cell(
|
||||
compat_result=compat_result,
|
||||
models=VERTEX_MODELS,
|
||||
prompt="Reply with the single word 'pong' and nothing else.",
|
||||
build_extra_env=vertex_extra_env,
|
||||
)
|
||||
|
|
@ -59,21 +59,31 @@ model_list:
|
|||
aws_region_name: us-east-1
|
||||
|
||||
# ---- Vertex AI ----
|
||||
# `use_in_pass_through: true` registers each deployment's
|
||||
# project/location/credentials with the /vertex_ai passthrough
|
||||
# router, which the `passthrough` row needs to resolve
|
||||
# .../models/{alias}:streamRawPredict URLs. That registration only
|
||||
# reads the canonical `vertex_project`/`vertex_location` param names
|
||||
# (not the `vertex_ai_*` aliases); the chat translation path accepts
|
||||
# both.
|
||||
- model_name: claude-haiku-4-5-vertex
|
||||
litellm_params:
|
||||
model: vertex_ai/claude-haiku-4-5
|
||||
vertex_ai_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_ai_location: os.environ/VERTEXAI_LOCATION
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: os.environ/VERTEXAI_LOCATION
|
||||
use_in_pass_through: true
|
||||
- model_name: claude-sonnet-4-6-vertex
|
||||
litellm_params:
|
||||
model: vertex_ai/claude-sonnet-4-6
|
||||
vertex_ai_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_ai_location: os.environ/VERTEXAI_LOCATION
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: os.environ/VERTEXAI_LOCATION
|
||||
use_in_pass_through: true
|
||||
- model_name: claude-opus-4-7-vertex
|
||||
litellm_params:
|
||||
model: vertex_ai/claude-opus-4-7
|
||||
vertex_ai_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_ai_location: os.environ/VERTEXAI_LOCATION
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: os.environ/VERTEXAI_LOCATION
|
||||
use_in_pass_through: true
|
||||
|
||||
# ---- Microsoft Foundry (Anthropic deployments on Azure) ----
|
||||
- model_name: claude-haiku-4-5-azure
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@
|
|||
- {id: logging.datadog.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions], source: "integrations/datadog/datadog.py", rationale: "Failure metrics for alerting/SLO"}
|
||||
- {id: logging.prometheus.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, embeddings], source: "integrations/prometheus.py", rationale: "Standard OSS metrics; per-key cardinality (existing e2e)"}
|
||||
- {id: logging.otel.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/otel/logger.py", rationale: "OTEL spans on every call path"}
|
||||
- {id: logging.otel.stream.exports_metric, module: logging, tier: P0, event: stream, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses], source: "integrations/otel/logger.py", rationale: "Streaming closes the LLM span from the stream path; historically prone to duplicate/orphaned spans"}
|
||||
- {id: logging.otel.failure.exports_metric, module: logging, tier: P0, event: failure, assertions: [exports_metric], exercised_on: [chat_completions, messages], source: "integrations/otel/logger.py", rationale: "Error spans for observability continuity"}
|
||||
- {id: logging.braintrust.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/braintrust_logging.py", rationale: "Evals platform spend"}
|
||||
- {id: logging.langsmith.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions, messages], source: "integrations/langsmith.py", rationale: "LangChain ecosystem"}
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@
|
|||
- {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"}
|
||||
- {id: reliability.routing.usage_based.picks_under_tpm, module: reliability, tier: P0, behavior: routing, variant: usage_based, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "router_strategy/lowest_tpm_rpm_v2.py", rationale: "Routes to lowest-TPM deployment; prevents over-allocation"}
|
||||
- {id: reliability.routing.least_busy.picks_lowest_traffic, module: reliability, tier: P1, behavior: routing, variant: least_busy, assertions: [picks_lowest_traffic], exercised_on: [chat_completions, messages], source: "router_strategy/least_busy.py", rationale: "Fewest in-flight requests"}
|
||||
- {id: reliability.routing.complexity_llm_classifier.routes_by_llm_tier, module: reliability, tier: P1, behavior: routing, variant: complexity_llm_classifier, assertions: [routes_by_llm_tier], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py", fail_before_fix: proven, rationale: "v2 auto-router LLM complexity classifier runs over the proxy and routes by semantic tier instead of silently crashing on absent litellm_metadata and falling back to heuristic scoring"}
|
||||
- {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"}
|
||||
- {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"}
|
||||
- {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"}
|
||||
|
|
|
|||
|
|
@ -66,6 +66,23 @@ configs:
|
|||
model: openai/text-embedding-3-small
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
# v2 auto-router with the LLM complexity classifier. SIMPLE stays on the
|
||||
# openai backend; every higher tier routes to the anthropic backend, so the
|
||||
# served deployment (read back from the spend log's model) reveals whether
|
||||
# the LLM classifier actually ran or silently fell back to heuristic scoring.
|
||||
- model_name: complexity-smart-router
|
||||
litellm_params:
|
||||
model: auto_router/complexity_router
|
||||
complexity_router_config:
|
||||
classifier_type: llm
|
||||
classifier_llm_config:
|
||||
model: gpt-5.5
|
||||
tiers:
|
||||
SIMPLE: gpt-5.5
|
||||
MEDIUM: claude-haiku-4-5
|
||||
COMPLEX: claude-haiku-4-5
|
||||
REASONING: claude-haiku-4-5
|
||||
|
||||
services:
|
||||
litellm:
|
||||
image: ghcr.io/berriai/litellm:main-latest
|
||||
|
|
@ -79,6 +96,7 @@ services:
|
|||
env_file: .env
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: sk-1234
|
||||
STORE_MODEL_IN_DB: "True"
|
||||
LITELLM_OTEL_V2: "true"
|
||||
PHOENIX_COLLECTOR_HTTP_ENDPOINT: http://jaeger:4318/v1/traces
|
||||
PHOENIX_API_KEY: local-jaeger-noauth
|
||||
|
|
|
|||
|
|
@ -121,6 +121,11 @@ class StreamingResponse(BaseModel):
|
|||
headers: dict[str, str] = {}
|
||||
body: str
|
||||
chunks: int = 0 # streamed events (0 for non-streaming)
|
||||
# First in-stream error event, if any. A streamed call commits its HTTP 200
|
||||
# before the upstream completes, so upstream failures (e.g. insufficient
|
||||
# quota) arrive as SSE error events inside an otherwise-successful response;
|
||||
# the consumed body is elided, so this is the only place they surface.
|
||||
stream_error: str | None = None
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
|
|
@ -289,7 +294,19 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon
|
|||
body=resp.text,
|
||||
)
|
||||
lines = cast("Iterator[bytes]", resp.iter_lines())
|
||||
chunks = sum(1 for line in lines if line)
|
||||
chunks = 0
|
||||
stream_error: str | None = None
|
||||
for line in lines:
|
||||
if not line:
|
||||
continue
|
||||
chunks += 1
|
||||
if stream_error is None and (
|
||||
line.startswith(b"event: error")
|
||||
or b'"type":"error"' in line
|
||||
or b'"type": "error"' in line
|
||||
or line.startswith(b'data: {"error"')
|
||||
):
|
||||
stream_error = line.decode(errors="replace")[:300]
|
||||
return StreamingResponse(
|
||||
status_code=resp.status_code,
|
||||
call_id=call_id,
|
||||
|
|
@ -298,6 +315,7 @@ def _streaming_outcome(resp: requests.Response, stream: bool) -> StreamingRespon
|
|||
headers=headers,
|
||||
body="<streamed>",
|
||||
chunks=chunks,
|
||||
stream_error=stream_error,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -77,11 +77,12 @@ WEATHER_TOOL = ChatTool(
|
|||
|
||||
|
||||
class ResponsesRequestBody(BaseModel):
|
||||
"""OpenAI Responses API /v1/responses request (non-streaming)."""
|
||||
"""OpenAI Responses API /v1/responses request."""
|
||||
|
||||
model: str
|
||||
input: str
|
||||
max_output_tokens: int
|
||||
stream: bool | None = None
|
||||
|
||||
|
||||
class TeamCallbackBody(BaseModel):
|
||||
|
|
@ -464,30 +465,43 @@ class LoggingClient:
|
|||
json=body,
|
||||
)
|
||||
|
||||
def messages_raw(self, key: str, model: str, text: str, *, max_tokens: int = 16) -> StreamingResponse:
|
||||
"""Non-streaming POST /v1/messages (Anthropic-native body): raw outcome
|
||||
judged by status/body/headers, for tests that need x-litellm-call-id."""
|
||||
def messages_raw(
|
||||
self, key: str, model: str, text: str, *, max_tokens: int = 16, stream: bool = False
|
||||
) -> StreamingResponse:
|
||||
"""POST /v1/messages (Anthropic-native body): raw outcome judged by
|
||||
status/body/headers, for tests that need x-litellm-call-id. With
|
||||
``stream=True`` the SSE body is consumed and its events counted."""
|
||||
body = AnthropicMessagesBody(
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
stream=True if stream else None,
|
||||
)
|
||||
if stream:
|
||||
return self.gateway.transport.stream(
|
||||
"/v1/messages", headers=self.gateway.transport.bearer(key), json=body
|
||||
)
|
||||
return self.gateway.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
),
|
||||
"/v1/messages", headers=self.gateway.transport.bearer(key), json=body
|
||||
)
|
||||
|
||||
def responses_raw(
|
||||
self, key: str, model: str, text: str, *, max_output_tokens: int = 64
|
||||
self, key: str, model: str, text: str, *, max_output_tokens: int = 64, stream: bool = False
|
||||
) -> StreamingResponse:
|
||||
"""Non-streaming POST /v1/responses (OpenAI Responses API): raw outcome
|
||||
judged by status/body/headers, for tests that need x-litellm-call-id.
|
||||
"""POST /v1/responses (OpenAI Responses API): raw outcome judged by
|
||||
status/body/headers, for tests that need x-litellm-call-id.
|
||||
max_output_tokens caps reasoning-model output cost; a capped response is
|
||||
still a 200 and still exports the trace."""
|
||||
still a 200 and still exports the trace. With ``stream=True`` the SSE
|
||||
body is consumed and its events counted."""
|
||||
body = ResponsesRequestBody(
|
||||
model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None
|
||||
)
|
||||
if stream:
|
||||
return self.gateway.transport.stream(
|
||||
"/v1/responses", headers=self.gateway.transport.bearer(key), json=body
|
||||
)
|
||||
return self.gateway.transport.send(
|
||||
"/v1/responses",
|
||||
headers=self.gateway.transport.bearer(key),
|
||||
json=ResponsesRequestBody(model=model, input=text, max_output_tokens=max_output_tokens),
|
||||
"/v1/responses", headers=self.gateway.transport.bearer(key), json=body
|
||||
)
|
||||
|
||||
def scrape_metrics(self) -> str:
|
||||
|
|
|
|||
|
|
@ -21,13 +21,14 @@ import time
|
|||
from collections.abc import Callable
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker
|
||||
from e2e_http import NoBody, StreamingResponse, require_successful_call
|
||||
from lifecycle import ResourceManager
|
||||
from logging_client import LoggingClient
|
||||
from otel_client import JaegerTrace, OtelReader
|
||||
from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient
|
||||
from models import LiteLLMParamsBody
|
||||
from otel_client import JaegerSpan, JaegerTrace, OtelReader
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -96,7 +97,9 @@ def _chain_reaches(span_id: str, root_id: str, trace: JaegerTrace) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: str) -> None:
|
||||
def _assert_complete_trace(
|
||||
hits: list[JaegerTrace], *, route: str, genai_span: str, require_cost_span: bool = True
|
||||
) -> None:
|
||||
"""The enforced behavior: the destination holds exactly one trace for the
|
||||
call, rooted at the SERVER span, with auth/db/cost children and the gen-AI
|
||||
span all connected into that one tree - no dangling parent references."""
|
||||
|
|
@ -135,7 +138,8 @@ def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: s
|
|||
assert any(name.startswith(DB_SPAN_PREFIX) for name in names), (
|
||||
f"no db ('{DB_SPAN_PREFIX}*') span in the trace; spans: {names}"
|
||||
)
|
||||
assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}"
|
||||
if require_cost_span:
|
||||
assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}"
|
||||
|
||||
genai = next((span for span in trace.spans if span.operation_name == genai_span), None)
|
||||
assert genai is not None, f"gen-AI span {genai_span!r} missing; spans: {names}"
|
||||
|
|
@ -146,8 +150,79 @@ def _assert_complete_trace(hits: list[JaegerTrace], *, route: str, genai_span: s
|
|||
)
|
||||
|
||||
|
||||
def _settled_names(*, route: str, genai_span: str) -> set[str]:
|
||||
return {f"POST {route}", f"auth {route}", COST_SPAN, genai_span}
|
||||
def _settled_names(*, route: str, genai_span: str, require_cost_span: bool = True) -> set[str]:
|
||||
names = {f"POST {route}", f"auth {route}", genai_span}
|
||||
return (names | {COST_SPAN}) if require_cost_span else names
|
||||
|
||||
|
||||
def _tag(span: JaegerSpan, key: str) -> str | int | float | bool | None:
|
||||
for tag in span.tags:
|
||||
if tag.key == key:
|
||||
return tag.value
|
||||
return None
|
||||
|
||||
|
||||
#: The attribute contract a failed call's gen-AI span must carry (LIT-4179), as
|
||||
#: one reviewable payload. Exact-match values; error.message is additionally
|
||||
#: proven untruncated by _assert_error_span_contract, which parses the provider
|
||||
#: error JSON embedded in it - a truncated message stops parsing.
|
||||
EXPECTED_ERROR_SPAN_ATTRIBUTES: dict[str, str] = {
|
||||
"error": "True",
|
||||
"error.type": "AuthenticationError",
|
||||
"otel.status_code": "ERROR",
|
||||
"litellm.provider.error.code": "401",
|
||||
"litellm.provider.error.llm_provider": "anthropic",
|
||||
}
|
||||
|
||||
|
||||
class _ProviderErrorDetail(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
type: str
|
||||
message: str
|
||||
|
||||
|
||||
class _ProviderError(BaseModel):
|
||||
"""The provider error JSON embedded in error.message; validating it proves
|
||||
the attribute survived untruncated (a cut-off message stops parsing)."""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
error: _ProviderErrorDetail
|
||||
|
||||
|
||||
def _assert_error_span_contract(span: JaegerSpan) -> None:
|
||||
"""The failed call's gen-AI span carries the LIT-4179 error contract: the
|
||||
exact attributes in EXPECTED_ERROR_SPAN_ATTRIBUTES, plus an untruncated
|
||||
error.message whose embedded provider error JSON still parses and whose
|
||||
text also rides the span status description."""
|
||||
for key, expected in EXPECTED_ERROR_SPAN_ATTRIBUTES.items():
|
||||
actual = _tag(span, key)
|
||||
assert str(actual) == expected, f"error span attribute {key!r} must be {expected!r}, got {actual!r}"
|
||||
|
||||
message = _tag(span, "error.message")
|
||||
assert isinstance(message, str) and message, "error span must carry a non-empty error.message"
|
||||
assert "AnthropicException" in message, (
|
||||
f"error.message must carry the upstream provider exception, got: {message[:200]}"
|
||||
)
|
||||
start, end = message.find("{"), message.rfind("}")
|
||||
assert start != -1 and end > start, (
|
||||
f"error.message carries no parseable provider error JSON (truncated?): {message[:200]}"
|
||||
)
|
||||
try:
|
||||
provider_error = _ProviderError.model_validate_json(message[start : end + 1])
|
||||
except ValidationError:
|
||||
pytest.fail(f"the embedded provider error JSON does not parse (truncated?): {message[:300]}")
|
||||
assert provider_error.error.message == "invalid x-api-key", (
|
||||
f"the embedded provider error must survive untruncated; parsed: {provider_error}"
|
||||
)
|
||||
assert _tag(span, "otel.status_description") == message, (
|
||||
"the span status description must carry the same untruncated message as error.message"
|
||||
)
|
||||
stack = _tag(span, "litellm.provider.error.stack_trace")
|
||||
assert isinstance(stack, str) and stack, (
|
||||
"the error span must carry a non-empty litellm.provider.error.stack_trace"
|
||||
)
|
||||
|
||||
|
||||
class TestOtelTraceCompleteness:
|
||||
|
|
@ -257,3 +332,238 @@ class TestOtelTraceCompleteness:
|
|||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span)
|
||||
|
||||
@pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["chat_completions"])
|
||||
def test_chat_completions_stream_exports_complete_trace(
|
||||
self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager
|
||||
) -> None:
|
||||
"""A successful streamed `/chat/completions` request should export one
|
||||
complete OTEL trace. The trace must contain a single root `SERVER`
|
||||
span, with the auth, database, cost, and gen-AI `CLIENT` spans all
|
||||
connected back to that root.
|
||||
|
||||
Streaming has an additional lifecycle risk because the gen-AI span is
|
||||
closed by the stream-consumption path after the final chunk has
|
||||
arrived and usage has been aggregated. Historically, this has caused
|
||||
duplicate or orphaned spans.
|
||||
|
||||
The test therefore confirms that:
|
||||
|
||||
* The response actually streams.
|
||||
* Exactly one gen-AI span is created for the request.
|
||||
* The gen-AI span contains `litellm.request.streaming=true`.
|
||||
"""
|
||||
route = "/chat/completions"
|
||||
_assert_otel_destination_configured(client)
|
||||
|
||||
key = client.key_with_alias(f"otel-stream-chat-{unique_marker()}", models=[MODEL])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
marker = unique_marker()
|
||||
outcome = _first_ok(
|
||||
client,
|
||||
lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", stream=True, max_tokens=16),
|
||||
)
|
||||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}"
|
||||
assert outcome.chunks > 0, "the stream must deliver at least one event"
|
||||
assert outcome.stream_error is None, (
|
||||
f"the stream carried an upstream error event despite the 200: {outcome.stream_error}"
|
||||
)
|
||||
|
||||
genai_span = f"chat {MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span)
|
||||
|
||||
genai_spans = [span for span in hits[0].spans if span.operation_name == genai_span]
|
||||
assert len(genai_spans) == 1, (
|
||||
f"a streamed call must produce exactly ONE gen-AI span, got {len(genai_spans)}; "
|
||||
f"spans: {hits[0].span_names()}"
|
||||
)
|
||||
assert _tag(genai_spans[0], "litellm.request.streaming") is True, (
|
||||
"the gen-AI span must record litellm.request.streaming=true; its absence means "
|
||||
"the stream flag was dropped before the model call"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["messages"])
|
||||
def test_messages_stream_exports_complete_trace(
|
||||
self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager
|
||||
) -> None:
|
||||
"""A successful streamed `/v1/messages` request should export one
|
||||
complete OTEL trace. The trace must contain a single root `SERVER`
|
||||
span, with the auth, database, cost, and gen-AI `CLIENT` spans all
|
||||
connected back to that root.
|
||||
|
||||
This endpoint has the same streaming lifecycle risk as
|
||||
`/chat/completions`: the gen-AI span is closed by the
|
||||
stream-consumption path after the final chunk has arrived and usage
|
||||
has been aggregated.
|
||||
|
||||
The test therefore confirms that:
|
||||
|
||||
* The response actually streams.
|
||||
* Exactly one gen-AI span is created for the request.
|
||||
* The gen-AI span contains `litellm.request.streaming=true`.
|
||||
"""
|
||||
route = "/v1/messages"
|
||||
_assert_otel_destination_configured(client)
|
||||
|
||||
key = client.key_with_alias(f"otel-stream-messages-{unique_marker()}", models=[MODEL])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
marker = unique_marker()
|
||||
outcome = _first_ok(
|
||||
client,
|
||||
lambda: client.messages_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16, stream=True),
|
||||
)
|
||||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}"
|
||||
assert outcome.chunks > 0, "the stream must deliver at least one event"
|
||||
assert outcome.stream_error is None, (
|
||||
f"the stream carried an upstream error event despite the 200: {outcome.stream_error}"
|
||||
)
|
||||
|
||||
genai_span = f"chat {MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span)
|
||||
|
||||
genai_spans = [span for span in hits[0].spans if span.operation_name == genai_span]
|
||||
assert len(genai_spans) == 1, (
|
||||
f"a streamed call must produce exactly ONE gen-AI span, got {len(genai_spans)}; "
|
||||
f"spans: {hits[0].span_names()}"
|
||||
)
|
||||
assert _tag(genai_spans[0], "litellm.request.streaming") is True, (
|
||||
"the gen-AI span must record litellm.request.streaming=true; its absence means "
|
||||
"the stream flag was dropped before the model call"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["responses"])
|
||||
def test_responses_stream_exports_complete_trace(
|
||||
self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager
|
||||
) -> None:
|
||||
"""A successful streamed /v1/responses request should export one complete
|
||||
OTEL trace. The trace must contain a single root SERVER span, with the
|
||||
auth, database, and gen-AI CLIENT spans all connected back to that
|
||||
root.
|
||||
|
||||
This endpoint has the same streaming lifecycle risk as the other
|
||||
streaming surfaces: the gen-AI span is closed by the
|
||||
stream-consumption path after the final event has arrived and usage
|
||||
has been aggregated.
|
||||
|
||||
The test therefore confirms that:
|
||||
|
||||
* The response actually streams.
|
||||
* Exactly one gen-AI span is created for the request.
|
||||
* Spend is recorded correctly.
|
||||
"""
|
||||
route = "/v1/responses"
|
||||
_assert_otel_destination_configured(client)
|
||||
|
||||
key = client.key_with_alias(
|
||||
f"otel-stream-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
marker = unique_marker()
|
||||
outcome = _first_ok(
|
||||
client,
|
||||
lambda: client.responses_raw(key, CHEAP_OPENAI_MODEL, f"reply with one word {marker}", stream=True),
|
||||
)
|
||||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}"
|
||||
assert outcome.chunks > 0, "the stream must deliver at least one event"
|
||||
assert outcome.stream_error is None, (
|
||||
f"the stream carried an upstream error event despite the 200: {outcome.stream_error}"
|
||||
)
|
||||
|
||||
genai_span = f"chat {CHEAP_OPENAI_MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
|
||||
genai_spans = [span for span in hits[0].spans if span.operation_name == genai_span]
|
||||
assert len(genai_spans) == 1, (
|
||||
f"a streamed call must produce exactly ONE gen-AI span, got {len(genai_spans)}; "
|
||||
f"spans: {hits[0].span_names()}"
|
||||
)
|
||||
|
||||
spend_row = client.poll_proxy_spend_for_key(key)
|
||||
assert spend_row is not None and spend_row.spend is not None and spend_row.spend > 0, (
|
||||
"a successful streamed responses call must record a positive-spend row in /spend/logs "
|
||||
"(the cost-write SPAN is knowingly absent on this surface, LIT-4428, but the spend "
|
||||
f"itself must land); got {spend_row!r}"
|
||||
)
|
||||
assert spend_row.call_type == "aresponses", (
|
||||
f"the spend row must be attributed to the responses call type, got {spend_row.call_type!r}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("logging.otel.failure.exports_metric", exercised_on=["chat_completions"])
|
||||
def test_failed_chat_completions_error_span_attributes(
|
||||
self, client: LoggingClient, otel_reader: OtelReader, resources: ResourceManager
|
||||
) -> None:
|
||||
"""This test checks that a failed `/chat/completions` request produces a
|
||||
single, complete OTEL trace. The model-call span should include all
|
||||
expected error attributes, a non-empty stack trace, and the full,
|
||||
untruncated `error.message`. The provider error embedded in that
|
||||
message must also remain valid JSON. The root server span should
|
||||
record the same 401 response returned to the client.
|
||||
|
||||
The test uses a deployment with an invalid upstream API key. This
|
||||
allows the request to pass LiteLLM’s proxy authentication and fail at
|
||||
the provider, which is necessary to generate a model-call error span.
|
||||
There should be no cost-write span because failed requests are not
|
||||
billed."""
|
||||
route = "/chat/completions"
|
||||
_assert_otel_destination_configured(client)
|
||||
|
||||
model_name = f"otel-err-{unique_marker()}"
|
||||
model_id = client.create_model(
|
||||
model_name,
|
||||
LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY),
|
||||
)
|
||||
resources.defer(lambda: client.delete_model(model_id))
|
||||
key = client.key_with_alias(f"otel-err-{unique_marker()}", models=[model_name])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
deadline = time.monotonic() + client.gateway.poll_timeout
|
||||
while True:
|
||||
outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16)
|
||||
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
|
||||
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
|
||||
break
|
||||
time.sleep(client.gateway.poll_interval)
|
||||
assert "AnthropicException" in outcome.body, (
|
||||
"never saw the upstream provider failure before the deadline; the key may still be "
|
||||
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
|
||||
)
|
||||
assert outcome.status_code == 401, (
|
||||
f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}"
|
||||
)
|
||||
assert outcome.call_id is not None, "failed responses must still carry x-litellm-call-id"
|
||||
|
||||
genai_span = f"chat {model_name}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
|
||||
root = next(span for span in hits[0].spans if not span.references)
|
||||
assert str(_tag(root, "http.status_code")) == "401", (
|
||||
f"the SERVER span must record the 401 the client received, got {_tag(root, 'http.status_code')!r}"
|
||||
)
|
||||
genai = next(span for span in hits[0].spans if span.operation_name == genai_span)
|
||||
_assert_error_span_contract(genai)
|
||||
|
|
|
|||
|
|
@ -154,6 +154,7 @@ class AnthropicMessagesBody(BaseModel):
|
|||
model: str
|
||||
messages: list[ChatMessage]
|
||||
max_tokens: int
|
||||
stream: bool | None = None
|
||||
|
||||
|
||||
class AnthropicMessagesResponse(BaseModel):
|
||||
|
|
|
|||
20
tests/e2e/router/complexity_router_client.py
Normal file
20
tests/e2e/router/complexity_router_client.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
"""Client for the complexity auto-router e2e tests.
|
||||
|
||||
The suite drives the shared /chat/completions and spend-log reads on the Gateway,
|
||||
so this client only carries the Gateway the shared lifecycle needs for cleanup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ComplexityRouterClient:
|
||||
gateway: Gateway
|
||||
|
||||
|
||||
def build_client() -> ComplexityRouterClient:
|
||||
return ComplexityRouterClient(gateway=build_gateway())
|
||||
15
tests/e2e/router/conftest.py
Normal file
15
tests/e2e/router/conftest.py
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
"""Router suite's `client` fixture.
|
||||
|
||||
The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker
|
||||
live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared
|
||||
Gateway, so the `resources` fixture cleans up keys this suite creates.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from complexity_router_client import ComplexityRouterClient, build_client
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def client() -> ComplexityRouterClient:
|
||||
return build_client()
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue