mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-20 05:04:24 +00:00
merge(litellm_internal_staging): resolve _delta_has_content rename conflict
Keep staging's _delta_has_content helper/name while preserving peek-path tool_name_mapping restore and stop-only usage-hold fixes from this branch. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
commit
8db49269d3
479 changed files with 42161 additions and 6936 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
|
||||
|
|
|
|||
2
.github/workflows/test-unit-proxy-db.yml
vendored
2
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -5,6 +5,8 @@ on:
|
|||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
|
|||
6
.github/workflows/zizmor.yml
vendored
6
.github/workflows/zizmor.yml
vendored
|
|
@ -4,7 +4,11 @@ on:
|
|||
push:
|
||||
branches: [main, litellm_internal_staging]
|
||||
pull_request:
|
||||
branches: [main, litellm_internal_staging]
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
|
|
|
|||
7
.gitignore
vendored
7
.gitignore
vendored
|
|
@ -106,6 +106,13 @@ STABILIZATION_TODO.md
|
|||
**/coverage
|
||||
test-config
|
||||
|
||||
# Claude Code compatibility-matrix pytest artifact (CI-only output).
|
||||
compat-results.json
|
||||
compat-results.json.shards/
|
||||
compat-rate-limit-summary.json
|
||||
# Matrix JSON produced by the daily-cron publisher (pushed to litellm-docs).
|
||||
compatibility-matrix.json
|
||||
|
||||
# ---------- Terraform ----------
|
||||
# Provider binaries + module cache — regenerated by `terraform init`.
|
||||
**/.terraform/
|
||||
|
|
|
|||
|
|
@ -114,6 +114,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
|
|||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
|
||||
# whole /root/.cache drags in the uv build cache (~660 MB, includes a
|
||||
|
|
|
|||
|
|
@ -111,6 +111,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
|
|||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
|
|
|
|||
|
|
@ -137,6 +137,7 @@ COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/pr
|
|||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
COPY --from=builder /app/litellm-proxy-extras /app/litellm-proxy-extras
|
||||
COPY --from=builder /app/.cache /app/.cache
|
||||
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
|
||||
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -113,6 +113,10 @@ class PagerDutyAlerting(SlackAlerting):
|
|||
user_api_key_spend=_meta.get("user_api_key_spend"),
|
||||
user_api_key_max_budget=_meta.get("user_api_key_max_budget"),
|
||||
user_api_key_budget_reset_at=_meta.get("user_api_key_budget_reset_at"),
|
||||
user_api_key_user_spend=_meta.get("user_api_key_user_spend"),
|
||||
user_api_key_user_max_budget=_meta.get("user_api_key_user_max_budget"),
|
||||
user_api_key_team_spend=_meta.get("user_api_key_team_spend"),
|
||||
user_api_key_team_max_budget=_meta.get("user_api_key_team_max_budget"),
|
||||
user_api_key_org_id=_meta.get("user_api_key_org_id"),
|
||||
user_api_key_org_alias=_meta.get("user_api_key_org_alias"),
|
||||
user_api_key_team_id=_meta.get("user_api_key_team_id"),
|
||||
|
|
@ -196,6 +200,10 @@ class PagerDutyAlerting(SlackAlerting):
|
|||
if user_api_key_dict.budget_reset_at
|
||||
else None
|
||||
),
|
||||
user_api_key_user_spend=user_api_key_dict.user_spend,
|
||||
user_api_key_user_max_budget=user_api_key_dict.user_max_budget,
|
||||
user_api_key_team_spend=user_api_key_dict.team_spend,
|
||||
user_api_key_team_max_budget=user_api_key_dict.team_max_budget,
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
user_api_key_org_alias=user_api_key_dict.organization_alias,
|
||||
user_api_key_team_id=user_api_key_dict.team_id,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.50"
|
||||
version = "0.1.51"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.50"
|
||||
version = "0.1.51"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -46,4 +46,9 @@ Reminders:
|
|||
- gateway.config.proxy_config (rendered into a ConfigMap and mounted at
|
||||
/app/config/config.yaml; gateway reads it via
|
||||
CONFIG_FILE_PATH)
|
||||
- {component}.pdb.{enabled,minAvailable,maxUnavailable} (per-component PodDisruptionBudget; disabled by
|
||||
default — with hpa.minReplicas of 1, minAvailable: 1
|
||||
would block node drains)
|
||||
- {component}.topologySpreadConstraints (standard k8s list, e.g. spread replicas across
|
||||
topology.kubernetes.io/zone)
|
||||
- Enable ingress.enabled=true to dispatch / → ui, gateway data-plane prefixes → gateway, and the catch-all → backend.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
*/}}
|
||||
|
|
@ -244,6 +295,52 @@ harmless no-op for the Job and authoritative for the app pods.
|
|||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
PodDisruptionBudget shared by gateway, backend, and ui.
|
||||
|
||||
Invoke with a dict:
|
||||
(dict "root" $ "component" .Values.gateway "componentName" "gateway"
|
||||
"fullname" (include "litellm.gateway.fullname" .)
|
||||
"selectorLabels" (include "litellm.gateway.selectorLabels" .))
|
||||
|
||||
Renders nothing unless both the component and its `pdb.enabled` are on.
|
||||
Only one of minAvailable / maxUnavailable should be set; if both are,
|
||||
minAvailable wins. If neither is set, falls back to `maxUnavailable: 1` so
|
||||
an enabled-but-unconfigured PDB still permits node drains.
|
||||
|
||||
"Set" means non-nil and non-empty-string, so an explicit 0 (e.g.
|
||||
`maxUnavailable: 0` to forbid all voluntary disruptions) is honored rather
|
||||
than silently replaced by the fallback.
|
||||
*/}}
|
||||
{{- define "litellm.pdb" -}}
|
||||
{{- $root := .root -}}
|
||||
{{- $component := .component -}}
|
||||
{{- $min := $component.pdb.minAvailable -}}
|
||||
{{- $max := $component.pdb.maxUnavailable -}}
|
||||
{{- $minSet := not (or (kindIs "invalid" $min) (eq (printf "%v" $min) "")) -}}
|
||||
{{- $maxSet := not (or (kindIs "invalid" $max) (eq (printf "%v" $max) "")) -}}
|
||||
{{- if and $component.enabled $component.pdb $component.pdb.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ .fullname }}
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" $root | nindent 4 }}
|
||||
app.kubernetes.io/component: {{ .componentName }}
|
||||
spec:
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- .selectorLabels | nindent 6 }}
|
||||
{{- if $minSet }}
|
||||
minAvailable: {{ $min }}
|
||||
{{- else if $maxSet }}
|
||||
maxUnavailable: {{ $max }}
|
||||
{{- else }}
|
||||
maxUnavailable: 1
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
|
||||
{{/*
|
||||
Renders `envFrom:` block for a component's `envConfigMaps` / `envSecrets`
|
||||
lists. Each entry is a resource name; the chart wires the whole ConfigMap /
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
@ -89,4 +98,8 @@ spec:
|
|||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.topologySpreadConstraints }}
|
||||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
6
helm/litellm/templates/backend/poddisruptionbudget.yaml
Normal file
6
helm/litellm/templates/backend/poddisruptionbudget.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{{- include "litellm.pdb" (dict
|
||||
"root" $
|
||||
"component" .Values.backend
|
||||
"componentName" "backend"
|
||||
"fullname" (include "litellm.backend.fullname" .)
|
||||
"selectorLabels" (include "litellm.backend.selectorLabels" .)) }}
|
||||
|
|
@ -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 }}
|
||||
|
|
@ -91,4 +100,8 @@ spec:
|
|||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.topologySpreadConstraints }}
|
||||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
6
helm/litellm/templates/gateway/poddisruptionbudget.yaml
Normal file
6
helm/litellm/templates/gateway/poddisruptionbudget.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{{- include "litellm.pdb" (dict
|
||||
"root" $
|
||||
"component" .Values.gateway
|
||||
"componentName" "gateway"
|
||||
"fullname" (include "litellm.gateway.fullname" .)
|
||||
"selectorLabels" (include "litellm.gateway.selectorLabels" .)) }}
|
||||
|
|
@ -76,4 +76,8 @@ spec:
|
|||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.ui.topologySpreadConstraints }}
|
||||
topologySpreadConstraints:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
|
|
|||
6
helm/litellm/templates/ui/poddisruptionbudget.yaml
Normal file
6
helm/litellm/templates/ui/poddisruptionbudget.yaml
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
{{- include "litellm.pdb" (dict
|
||||
"root" $
|
||||
"component" .Values.ui
|
||||
"componentName" "ui"
|
||||
"fullname" (include "litellm.ui.fullname" .)
|
||||
"selectorLabels" (include "litellm.ui.selectorLabels" .)) }}
|
||||
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
|
||||
188
helm/litellm/tests/pdb_topology_spread_tests.yaml
Normal file
188
helm/litellm/tests/pdb_topology_spread_tests.yaml
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
suite: test pod disruption budgets and topology spread constraints
|
||||
templates:
|
||||
- gateway/poddisruptionbudget.yaml
|
||||
- backend/poddisruptionbudget.yaml
|
||||
- ui/poddisruptionbudget.yaml
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: renders no PDB by default
|
||||
templates:
|
||||
- gateway/poddisruptionbudget.yaml
|
||||
- backend/poddisruptionbudget.yaml
|
||||
- ui/poddisruptionbudget.yaml
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: gateway PDB uses minAvailable and matches the gateway selector labels
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.pdb.enabled: true
|
||||
gateway.pdb.minAvailable: 1
|
||||
asserts:
|
||||
- isKind:
|
||||
of: PodDisruptionBudget
|
||||
- equal:
|
||||
path: apiVersion
|
||||
value: policy/v1
|
||||
- equal:
|
||||
path: metadata.name
|
||||
value: RELEASE-NAME-litellm-gateway
|
||||
- equal:
|
||||
path: spec.minAvailable
|
||||
value: 1
|
||||
- notExists:
|
||||
path: spec.maxUnavailable
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: gateway
|
||||
|
||||
- it: backend PDB uses maxUnavailable when minAvailable is unset
|
||||
template: backend/poddisruptionbudget.yaml
|
||||
set:
|
||||
backend.pdb.enabled: true
|
||||
backend.pdb.maxUnavailable: 25%
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.maxUnavailable
|
||||
value: 25%
|
||||
- notExists:
|
||||
path: spec.minAvailable
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: backend
|
||||
|
||||
- it: minAvailable wins when both minAvailable and maxUnavailable are set
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.pdb.enabled: true
|
||||
gateway.pdb.minAvailable: 2
|
||||
gateway.pdb.maxUnavailable: 1
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.minAvailable
|
||||
value: 2
|
||||
- notExists:
|
||||
path: spec.maxUnavailable
|
||||
|
||||
- it: an explicit maxUnavailable 0 is honored instead of the fallback
|
||||
template: backend/poddisruptionbudget.yaml
|
||||
set:
|
||||
backend.pdb.enabled: true
|
||||
backend.pdb.maxUnavailable: 0
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.maxUnavailable
|
||||
value: 0
|
||||
- notExists:
|
||||
path: spec.minAvailable
|
||||
|
||||
- it: an explicit minAvailable 0 is honored and beats a set maxUnavailable
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.pdb.enabled: true
|
||||
gateway.pdb.minAvailable: 0
|
||||
gateway.pdb.maxUnavailable: 1
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.minAvailable
|
||||
value: 0
|
||||
- notExists:
|
||||
path: spec.maxUnavailable
|
||||
|
||||
- it: enabled PDB with neither knob set falls back to maxUnavailable 1
|
||||
template: ui/poddisruptionbudget.yaml
|
||||
set:
|
||||
ui.pdb.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.maxUnavailable
|
||||
value: 1
|
||||
- notExists:
|
||||
path: spec.minAvailable
|
||||
- equal:
|
||||
path: spec.selector.matchLabels
|
||||
value:
|
||||
app.kubernetes.io/name: litellm
|
||||
app.kubernetes.io/instance: RELEASE-NAME
|
||||
app.kubernetes.io/component: ui
|
||||
|
||||
- it: renders no PDB for a disabled component even when its pdb is enabled
|
||||
template: gateway/poddisruptionbudget.yaml
|
||||
set:
|
||||
gateway.enabled: false
|
||||
gateway.pdb.enabled: true
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: deployments omit topologySpreadConstraints by default
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
asserts:
|
||||
- notExists:
|
||||
path: spec.template.spec.topologySpreadConstraints
|
||||
|
||||
- it: gateway deployment renders configured topologySpreadConstraints
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
gateway.topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: topology.kubernetes.io/zone
|
||||
whenUnsatisfiable: ScheduleAnyway
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/component: gateway
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints
|
||||
value:
|
||||
- maxSkew: 1
|
||||
topologyKey: topology.kubernetes.io/zone
|
||||
whenUnsatisfiable: ScheduleAnyway
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/component: gateway
|
||||
|
||||
- it: backend deployment renders configured topologySpreadConstraints
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: kubernetes.io/hostname
|
||||
whenUnsatisfiable: DoNotSchedule
|
||||
labelSelector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/component: backend
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
|
||||
value: kubernetes.io/hostname
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints[0].whenUnsatisfiable
|
||||
value: DoNotSchedule
|
||||
|
||||
- it: ui deployment renders configured topologySpreadConstraints
|
||||
template: ui/deployment.yaml
|
||||
set:
|
||||
ui.topologySpreadConstraints:
|
||||
- maxSkew: 1
|
||||
topologyKey: topology.kubernetes.io/zone
|
||||
whenUnsatisfiable: ScheduleAnyway
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.topologySpreadConstraints[0].topologyKey
|
||||
value: topology.kubernetes.io/zone
|
||||
|
|
@ -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:
|
||||
|
|
@ -171,10 +190,28 @@ gateway:
|
|||
maxReplicas: 10
|
||||
targetCPUUtilizationPercentage: 70
|
||||
targetMemoryUtilizationPercentage: 80
|
||||
# PodDisruptionBudget for the gateway pods. Set exactly one of
|
||||
# `minAvailable` / `maxUnavailable` (minAvailable wins if both are set;
|
||||
# enabling without either falls back to `maxUnavailable: 1`). Disabled by
|
||||
# default: with the default hpa.minReplicas of 1, a `minAvailable: 1` PDB
|
||||
# would block node drains entirely.
|
||||
pdb:
|
||||
enabled: false
|
||||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
# Standard k8s topologySpreadConstraints for the gateway pods, e.g. to
|
||||
# spread replicas across zones:
|
||||
# - maxSkew: 1
|
||||
# topologyKey: topology.kubernetes.io/zone
|
||||
# whenUnsatisfiable: ScheduleAnyway
|
||||
# labelSelector:
|
||||
# matchLabels:
|
||||
# app.kubernetes.io/component: gateway
|
||||
topologySpreadConstraints: []
|
||||
|
||||
# ---------- backend (UI / management API) ----------
|
||||
backend:
|
||||
|
|
@ -214,10 +251,17 @@ backend:
|
|||
minReplicas: 1
|
||||
maxReplicas: 4
|
||||
targetCPUUtilizationPercentage: 70
|
||||
# Same shape as gateway.pdb.
|
||||
pdb:
|
||||
enabled: false
|
||||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
# Same shape as gateway.topologySpreadConstraints.
|
||||
topologySpreadConstraints: []
|
||||
|
||||
# ---------- ui (Next.js static dashboard) ----------
|
||||
ui:
|
||||
|
|
@ -260,7 +304,14 @@ ui:
|
|||
minReplicas: 1
|
||||
maxReplicas: 3
|
||||
targetCPUUtilizationPercentage: 80
|
||||
# Same shape as gateway.pdb.
|
||||
pdb:
|
||||
enabled: false
|
||||
minAvailable: ""
|
||||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
# Same shape as gateway.topologySpreadConstraints.
|
||||
topologySpreadConstraints: []
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "issuer" TEXT;
|
||||
|
|
@ -325,6 +325,7 @@ model LiteLLM_MCPServerTable {
|
|||
command String?
|
||||
args String[] @default([])
|
||||
env Json? @default("{}")
|
||||
issuer String?
|
||||
authorization_url String?
|
||||
token_url String?
|
||||
registration_url String?
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.77"
|
||||
version = "0.4.78"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.77"
|
||||
version = "0.4.78"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -688,10 +688,8 @@ def get_redis_connection_pool(
|
|||
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
|
||||
connection_class = async_redis.Connection
|
||||
if redis_kwargs.pop("ssl", False):
|
||||
connection_class = async_redis.SSLConnection
|
||||
redis_kwargs["connection_class"] = connection_class
|
||||
if redis_kwargs.pop("ssl", None):
|
||||
redis_kwargs["connection_class"] = async_redis.SSLConnection
|
||||
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -103,6 +103,18 @@ class DualCache(BaseCache):
|
|||
if default_redis_ttl is not None:
|
||||
self.default_redis_ttl = default_redis_ttl
|
||||
|
||||
def _backfill_kwargs(self, kwargs: "dict[str, object]") -> "dict[str, object]":
|
||||
"""
|
||||
Kwargs for writing a Redis read result into the in-memory tier.
|
||||
|
||||
Applies ``default_in_memory_ttl`` exactly like the write paths do;
|
||||
without it, backfilled entries fall to ``InMemoryCache``'s own default
|
||||
TTL and can outlive the TTL this cache was configured with.
|
||||
"""
|
||||
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
|
||||
return {**kwargs, "ttl": self.default_in_memory_ttl}
|
||||
return kwargs
|
||||
|
||||
def set_cache(self, key, value, local_only: bool = False, **kwargs):
|
||||
# Update both Redis and in-memory cache
|
||||
try:
|
||||
|
|
@ -160,7 +172,7 @@ class DualCache(BaseCache):
|
|||
|
||||
if redis_result is not None:
|
||||
# Update in-memory cache with the value from Redis
|
||||
self.in_memory_cache.set_cache(key, redis_result, **kwargs)
|
||||
self.in_memory_cache.set_cache(key, redis_result, **self._backfill_kwargs(kwargs))
|
||||
|
||||
result = redis_result
|
||||
|
||||
|
|
@ -226,7 +238,7 @@ class DualCache(BaseCache):
|
|||
|
||||
if redis_result is not None:
|
||||
# Update in-memory cache with the value from Redis
|
||||
await self.in_memory_cache.async_set_cache(key, redis_result, **kwargs)
|
||||
await self.in_memory_cache.async_set_cache(key, redis_result, **self._backfill_kwargs(kwargs))
|
||||
|
||||
result = redis_result
|
||||
|
||||
|
|
@ -318,7 +330,7 @@ class DualCache(BaseCache):
|
|||
result[key_to_index[key]] = value
|
||||
|
||||
if value is not None and self.in_memory_cache is not None:
|
||||
await self.in_memory_cache.async_set_cache(key, value, **kwargs)
|
||||
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
|
||||
|
||||
return result
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -966,11 +966,15 @@ class BudgetExceededError(Exception):
|
|||
max_budget: float,
|
||||
message: Optional[str] = None,
|
||||
llm_provider: Optional[str] = None,
|
||||
entity_type: Optional[str] = None,
|
||||
entity_id: Optional[str] = None,
|
||||
):
|
||||
self.current_cost = current_cost
|
||||
self.max_budget = max_budget
|
||||
self.status_code = 429
|
||||
self.llm_provider = llm_provider or ""
|
||||
self.entity_type = entity_type
|
||||
self.entity_id = entity_id
|
||||
# Surface unified rate-limit fields without joining the RateLimitError
|
||||
# hierarchy so existing `except BudgetExceededError:` handlers keep
|
||||
# working; custom callbacks reading StandardLoggingPayload pick these
|
||||
|
|
|
|||
|
|
@ -267,29 +267,10 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
# If no keys, return default from env (likely logging to console or something else)
|
||||
return OpenTelemetryConfig.from_env()
|
||||
|
||||
# Determine endpoint - default to US cloud
|
||||
langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host()
|
||||
|
||||
if langfuse_host:
|
||||
# If LANGFUSE_HOST is provided, construct OTEL endpoint from it
|
||||
if not langfuse_host.startswith("http"):
|
||||
langfuse_host = "https://" + langfuse_host
|
||||
endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel"
|
||||
verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}")
|
||||
else:
|
||||
# Default to US cloud endpoint
|
||||
endpoint = LANGFUSE_CLOUD_US_ENDPOINT
|
||||
verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}")
|
||||
|
||||
auth_header = LangfuseOtelLogger._get_langfuse_authorization_header(
|
||||
public_key=public_key, secret_key=secret_key
|
||||
)
|
||||
otlp_auth_headers = f"Authorization={auth_header}"
|
||||
|
||||
return OpenTelemetryConfig(
|
||||
exporter="otlp_http",
|
||||
endpoint=endpoint,
|
||||
headers=otlp_auth_headers,
|
||||
return LangfuseOtelLogger._build_langfuse_otel_config(
|
||||
public_key=public_key,
|
||||
secret_key=secret_key,
|
||||
langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -316,33 +297,36 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
"LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set for Langfuse OpenTelemetry integration."
|
||||
)
|
||||
|
||||
# Determine endpoint - default to US cloud
|
||||
langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host()
|
||||
return LangfuseOtelLogger._build_langfuse_otel_config(
|
||||
public_key=public_key,
|
||||
secret_key=secret_key,
|
||||
langfuse_host=LangfuseOtelLogger._get_langfuse_otel_host(),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_langfuse_otel_config(
|
||||
public_key: str, secret_key: str, langfuse_host: Optional[str]
|
||||
) -> "OpenTelemetryConfig":
|
||||
"""
|
||||
Builds an OTLP HTTP config pointing at the Langfuse OTEL endpoint for the
|
||||
given host (US cloud when no host is provided), authorized with the given keys.
|
||||
"""
|
||||
if langfuse_host:
|
||||
# If LANGFUSE_HOST is provided, construct OTEL endpoint from it
|
||||
if not langfuse_host.startswith("http"):
|
||||
langfuse_host = "https://" + langfuse_host
|
||||
endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel"
|
||||
normalized_host = langfuse_host if langfuse_host.startswith("http") else f"https://{langfuse_host}"
|
||||
endpoint = f"{normalized_host.rstrip('/')}/api/public/otel"
|
||||
verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}")
|
||||
else:
|
||||
# Default to US cloud endpoint
|
||||
endpoint = LANGFUSE_CLOUD_US_ENDPOINT
|
||||
verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}")
|
||||
|
||||
auth_header = LangfuseOtelLogger._get_langfuse_authorization_header(
|
||||
public_key=public_key, secret_key=secret_key
|
||||
)
|
||||
otlp_auth_headers = f"Authorization={auth_header}"
|
||||
|
||||
# Prevent modification of global env vars which causes leakage
|
||||
# os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint
|
||||
# os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers
|
||||
|
||||
return OpenTelemetryConfig(
|
||||
exporter="otlp_http",
|
||||
endpoint=endpoint,
|
||||
headers=otlp_auth_headers,
|
||||
headers=f"Authorization={auth_header}",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -378,6 +362,29 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
|
||||
return dynamic_headers
|
||||
|
||||
def construct_dynamic_otel_config(
|
||||
self, standard_callback_dynamic_params: StandardCallbackDynamicParams
|
||||
) -> Optional["OpenTelemetryConfig"]:
|
||||
"""
|
||||
Build a full per-request OTLP config from team/key dynamic Langfuse credentials.
|
||||
|
||||
Key-scoped credentials must define the export target, not just the auth
|
||||
headers: without this, a proxy with no global LANGFUSE_* env vars keeps its
|
||||
init-time fallback exporter (console), so key-level langfuse_otel silently
|
||||
never reaches Langfuse.
|
||||
"""
|
||||
public_key = standard_callback_dynamic_params.get("langfuse_public_key")
|
||||
secret_key = standard_callback_dynamic_params.get("langfuse_secret_key")
|
||||
if not public_key or not secret_key:
|
||||
return None
|
||||
|
||||
langfuse_host = standard_callback_dynamic_params.get("langfuse_host") or self._get_langfuse_otel_host()
|
||||
return LangfuseOtelLogger._build_langfuse_otel_config(
|
||||
public_key=public_key,
|
||||
secret_key=secret_key,
|
||||
langfuse_host=langfuse_host,
|
||||
)
|
||||
|
||||
def create_litellm_proxy_request_started_span(
|
||||
self,
|
||||
start_time: datetime,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
|
|||
parse_semconv_opt_in,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.secret_managers.main import get_secret_bool, str_to_bool
|
||||
from litellm.types.services import ServiceLoggerPayload
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -948,12 +949,22 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
Returns:
|
||||
Tracer: The tracer to use for this request
|
||||
"""
|
||||
dynamic_config = self._get_dynamic_otel_config_from_kwargs(kwargs)
|
||||
if dynamic_config is not None:
|
||||
verbose_logger.debug(
|
||||
"[OTEL DEBUG] Using DYNAMIC config tracer with endpoint: %s",
|
||||
dynamic_config.endpoint,
|
||||
)
|
||||
return self._get_tracer_with_dynamic_config(dynamic_config)
|
||||
|
||||
dynamic_headers = self._get_dynamic_otel_headers_from_kwargs(kwargs)
|
||||
|
||||
if dynamic_headers is not None:
|
||||
# Create spans using a temporary tracer with dynamic headers
|
||||
tracer_to_use = self._get_tracer_with_dynamic_headers(dynamic_headers)
|
||||
verbose_logger.debug("[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", dynamic_headers)
|
||||
verbose_logger.debug(
|
||||
"[OTEL DEBUG] Using DYNAMIC tracer with headers: %s", redact_string(str(dynamic_headers))
|
||||
)
|
||||
else:
|
||||
# For langfuse_otel without dynamic headers, create a provider with env var credentials
|
||||
if hasattr(self, "callback_name") and self.callback_name == "langfuse_otel":
|
||||
|
|
@ -989,6 +1000,32 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
return dynamic_headers if dynamic_headers else None
|
||||
|
||||
def _get_dynamic_otel_config_from_kwargs(self, kwargs: dict) -> Optional[OpenTelemetryConfig]:
|
||||
"""Extract a full dynamic exporter config from kwargs if available."""
|
||||
standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = kwargs.get(
|
||||
"standard_callback_dynamic_params"
|
||||
)
|
||||
|
||||
if not standard_callback_dynamic_params:
|
||||
return None
|
||||
|
||||
return self.construct_dynamic_otel_config(standard_callback_dynamic_params=standard_callback_dynamic_params)
|
||||
|
||||
def _get_tracer_with_dynamic_config(self, dynamic_config: OpenTelemetryConfig):
|
||||
"""Create (or reuse) a tracer whose exporter target comes from a per-request config."""
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
||||
cache_key = f"dynamic_config:{dynamic_config.exporter}:{dynamic_config.endpoint}:{dynamic_config.headers}"
|
||||
if cache_key in self._tracer_provider_cache:
|
||||
return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME)
|
||||
|
||||
temp_provider = TracerProvider(resource=self._get_litellm_resource(self.config))
|
||||
temp_provider.add_span_processor(self._get_span_processor(config_override=dynamic_config))
|
||||
|
||||
self._tracer_provider_cache[cache_key] = temp_provider
|
||||
|
||||
return temp_provider.get_tracer(LITELLM_TRACER_NAME)
|
||||
|
||||
def _get_tracer_with_dynamic_headers(self, dynamic_headers: dict):
|
||||
"""Create a temporary tracer with dynamic headers for this request only."""
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
|
@ -1020,6 +1057,19 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
"""
|
||||
return None
|
||||
|
||||
def construct_dynamic_otel_config(
|
||||
self, standard_callback_dynamic_params: StandardCallbackDynamicParams
|
||||
) -> Optional[OpenTelemetryConfig]:
|
||||
"""
|
||||
Construct a full exporter config from standard callback dynamic params.
|
||||
|
||||
Override this when team/key dynamic params must control the export
|
||||
target (exporter kind + endpoint), not just the request headers. When
|
||||
this returns a config, it takes precedence over
|
||||
construct_dynamic_otel_headers for the request.
|
||||
"""
|
||||
return None
|
||||
|
||||
#########################################################
|
||||
# End of Team/Key Based Logging Control Flow
|
||||
#########################################################
|
||||
|
|
@ -2747,7 +2797,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
verbose_logger.debug("OpenTelemetry: No parent context found, creating root span")
|
||||
return None, None
|
||||
|
||||
def _get_span_processor(self, dynamic_headers: Optional[dict] = None):
|
||||
def _get_span_processor(
|
||||
self,
|
||||
dynamic_headers: Optional[dict] = None,
|
||||
config_override: Optional[OpenTelemetryConfig] = None,
|
||||
):
|
||||
from opentelemetry.sdk.trace.export import (
|
||||
BatchSpanProcessor,
|
||||
ConsoleSpanExporter,
|
||||
|
|
@ -2755,40 +2809,45 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
SpanExporter,
|
||||
)
|
||||
|
||||
otel_exporter = config_override.exporter if config_override else self.OTEL_EXPORTER
|
||||
otel_endpoint = config_override.endpoint if config_override else self.OTEL_ENDPOINT
|
||||
otel_headers = config_override.headers if config_override else self.OTEL_HEADERS
|
||||
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry Logger, initializing span processor \nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
self.OTEL_ENDPOINT,
|
||||
self.OTEL_HEADERS,
|
||||
"OpenTelemetry Logger, initializing span processor \nexporter: %s\nendpoint: %s\nheaders: %s",
|
||||
otel_exporter,
|
||||
otel_endpoint,
|
||||
redact_string(str(otel_headers)),
|
||||
)
|
||||
_split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or self.OTEL_HEADERS)
|
||||
_split_otel_headers = OpenTelemetry._get_headers_dictionary(headers=dynamic_headers or otel_headers)
|
||||
|
||||
if dynamic_headers:
|
||||
verbose_logger.debug(
|
||||
"[OTEL DEBUG] Creating span processor with DYNAMIC headers: %s",
|
||||
{k: v[:20] + "..." if len(str(v)) > 20 else v for k, v in _split_otel_headers.items()},
|
||||
redact_string(str(_split_otel_headers)),
|
||||
)
|
||||
elif config_override:
|
||||
verbose_logger.debug(
|
||||
"[OTEL DEBUG] Creating span processor with DYNAMIC config, endpoint: %s",
|
||||
otel_endpoint,
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug("[OTEL DEBUG] Creating span processor with GLOBAL headers")
|
||||
|
||||
if hasattr(self.OTEL_EXPORTER, "export"): # Check if it has the export method that SpanExporter requires
|
||||
if hasattr(otel_exporter, "export"): # Check if it has the export method that SpanExporter requires
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: intiializing SpanExporter. Value of OTEL_EXPORTER: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
otel_exporter,
|
||||
)
|
||||
return SimpleSpanProcessor(cast(SpanExporter, self.OTEL_EXPORTER))
|
||||
return SimpleSpanProcessor(cast(SpanExporter, otel_exporter))
|
||||
|
||||
if self.OTEL_EXPORTER == "console":
|
||||
if otel_exporter == "console":
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
otel_exporter,
|
||||
)
|
||||
return BatchSpanProcessor(ConsoleSpanExporter())
|
||||
elif (
|
||||
self.OTEL_EXPORTER == "otlp_http"
|
||||
or self.OTEL_EXPORTER == "http/protobuf"
|
||||
or self.OTEL_EXPORTER == "http/json"
|
||||
):
|
||||
elif otel_exporter == "otlp_http" or otel_exporter == "http/protobuf" or otel_exporter == "http/json":
|
||||
try:
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
|
||||
OTLPSpanExporter as OTLPSpanExporterHTTP,
|
||||
|
|
@ -2801,13 +2860,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: intiializing http exporter. Value of OTEL_EXPORTER: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
otel_exporter,
|
||||
)
|
||||
normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces")
|
||||
normalized_endpoint = self._normalize_otel_endpoint(otel_endpoint, "traces")
|
||||
return BatchSpanProcessor(
|
||||
OTLPSpanExporterHTTP(endpoint=normalized_endpoint, headers=_split_otel_headers),
|
||||
)
|
||||
elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc":
|
||||
elif otel_exporter == "otlp_grpc" or otel_exporter == "grpc":
|
||||
try:
|
||||
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
|
||||
OTLPSpanExporter as OTLPSpanExporterGRPC,
|
||||
|
|
@ -2820,16 +2879,16 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: intiializing grpc exporter. Value of OTEL_EXPORTER: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
otel_exporter,
|
||||
)
|
||||
normalized_endpoint = self._normalize_otel_endpoint(self.OTEL_ENDPOINT, "traces")
|
||||
normalized_endpoint = self._normalize_otel_endpoint(otel_endpoint, "traces")
|
||||
return BatchSpanProcessor(
|
||||
OTLPSpanExporterGRPC(endpoint=normalized_endpoint, headers=_split_otel_headers),
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
otel_exporter,
|
||||
)
|
||||
return BatchSpanProcessor(ConsoleSpanExporter())
|
||||
|
||||
|
|
@ -2841,7 +2900,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
"OpenTelemetry Logger, initializing log exporter \nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
self.OTEL_ENDPOINT,
|
||||
self.OTEL_HEADERS,
|
||||
redact_string(str(self.OTEL_HEADERS)),
|
||||
)
|
||||
|
||||
_split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS)
|
||||
|
|
@ -2928,7 +2987,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
"OpenTelemetry Logger, initializing metric reader\nself.OTEL_EXPORTER: %s\nself.OTEL_ENDPOINT: %s\nself.OTEL_HEADERS: %s",
|
||||
self.OTEL_EXPORTER,
|
||||
self.OTEL_ENDPOINT,
|
||||
self.OTEL_HEADERS,
|
||||
redact_string(str(self.OTEL_HEADERS)),
|
||||
)
|
||||
|
||||
_split_otel_headers = OpenTelemetry._get_headers_dictionary(self.OTEL_HEADERS)
|
||||
|
|
|
|||
|
|
@ -161,8 +161,13 @@ def get_s3_object_key(
|
|||
start_time: datetime,
|
||||
s3_file_name: str,
|
||||
) -> str:
|
||||
sanitized_s3_file_name = s3_file_name.replace("/", "_")
|
||||
s3_object_key = (
|
||||
(s3_path.rstrip("/") + "/" if s3_path else "") + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name
|
||||
(s3_path.rstrip("/") + "/" if s3_path else "")
|
||||
+ prefix
|
||||
+ start_time.strftime("%Y-%m-%d")
|
||||
+ "/"
|
||||
+ sanitized_s3_file_name
|
||||
) # we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key += ".json"
|
||||
return s3_object_key
|
||||
|
|
|
|||
|
|
@ -19,9 +19,11 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.websearch_interception.tools import (
|
||||
get_litellm_web_search_tool,
|
||||
get_litellm_web_search_tool_openai,
|
||||
get_litellm_web_search_tool_responses,
|
||||
is_anthropic_native_web_search_tool,
|
||||
is_web_search_tool,
|
||||
is_web_search_tool_chat_completion,
|
||||
is_web_search_tool_responses,
|
||||
)
|
||||
from litellm.integrations.websearch_interception.transformation import (
|
||||
WebSearchTransformation,
|
||||
|
|
@ -32,11 +34,12 @@ from litellm.types.integrations.websearch_interception import (
|
|||
)
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
RESPONSES_AGENTIC_SURFACE,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.types.utils import CallTypes, LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
# Key used to flag, on per-request kwargs, that the originating client sent
|
||||
|
|
@ -251,6 +254,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if not tools:
|
||||
return None
|
||||
|
||||
if call_type in (CallTypes.responses, CallTypes.aresponses):
|
||||
return self._convert_responses_tools(kwargs=kwargs, tools=tools)
|
||||
|
||||
# Check if any tool is a web search tool (native or already LiteLLM standard)
|
||||
has_websearch = any(is_web_search_tool(t) for t in tools)
|
||||
|
||||
|
|
@ -291,6 +297,26 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
return kwargs
|
||||
|
||||
def _convert_responses_tools(self, kwargs: dict[str, Any], tools: list[dict[str, Any]]) -> dict | None:
|
||||
"""Convert Responses API web search tools to the LiteLLM standard function tool."""
|
||||
if not any(is_web_search_tool_responses(tool) for tool in tools):
|
||||
return None
|
||||
|
||||
verbose_logger.debug("WebSearchInterception: Converting Responses web_search tools to LiteLLM standard")
|
||||
|
||||
converted_tools = [
|
||||
get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools
|
||||
]
|
||||
|
||||
converted_kwargs = {**kwargs, "tools": converted_tools}
|
||||
|
||||
if kwargs.get("stream"):
|
||||
verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
|
||||
converted_kwargs["stream"] = False
|
||||
converted_kwargs["_websearch_interception_converted_stream"] = True
|
||||
|
||||
return converted_kwargs
|
||||
|
||||
@classmethod
|
||||
def from_config_yaml(cls, config: WebSearchInterceptionConfig) -> "WebSearchInterceptionLogger":
|
||||
"""
|
||||
|
|
@ -461,6 +487,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
if kwargs.get("_agentic_loop_api_surface") == RESPONSES_AGENTIC_SURFACE:
|
||||
return await self.async_should_run_responses_agentic_loop(
|
||||
response=response,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}")
|
||||
verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}")
|
||||
|
||||
|
|
@ -597,6 +634,54 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
}
|
||||
return True, tools_dict
|
||||
|
||||
async def async_should_run_responses_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
) -> tuple[bool, dict]:
|
||||
"""Check if WebSearch interception is needed for the Responses API."""
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Responses hook called! provider={custom_llm_provider}, stream={stream}"
|
||||
)
|
||||
|
||||
if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers:
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Skipping provider {custom_llm_provider} (not in enabled list: {self.enabled_providers})"
|
||||
)
|
||||
return False, {}
|
||||
|
||||
has_websearch_tool = any(is_web_search_tool_responses(t) for t in (tools or []))
|
||||
if not has_websearch_tool:
|
||||
verbose_logger.debug("WebSearchInterception: No litellm_web_search tool in responses request")
|
||||
return False, {}
|
||||
|
||||
should_intercept, tool_calls = WebSearchTransformation.transform_request(
|
||||
response=response,
|
||||
stream=stream,
|
||||
response_format="responses",
|
||||
)
|
||||
|
||||
if not should_intercept:
|
||||
verbose_logger.debug("WebSearchInterception: No WebSearch function_call detected in responses output")
|
||||
return False, {}
|
||||
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Detected {len(tool_calls)} WebSearch function_call(s), executing agentic loop"
|
||||
)
|
||||
|
||||
tools_dict = {
|
||||
"tool_calls": tool_calls,
|
||||
"tool_type": "websearch",
|
||||
"provider": custom_llm_provider,
|
||||
"response_format": "responses",
|
||||
}
|
||||
return True, tools_dict
|
||||
|
||||
async def async_run_agentic_loop(
|
||||
self,
|
||||
tools: Dict,
|
||||
|
|
@ -655,6 +740,18 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
if kwargs.get("_agentic_loop_api_surface") == RESPONSES_AGENTIC_SURFACE:
|
||||
return await self.async_build_responses_agentic_loop_plan(
|
||||
tools=tools,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
tool_calls = tools["tool_calls"]
|
||||
thinking_blocks = tools.get("thinking_blocks", [])
|
||||
request_patch, structured_results = await self._build_anthropic_request_patch(
|
||||
|
|
@ -809,6 +906,133 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
metadata={"tool_type": "websearch", "response_format": response_format},
|
||||
)
|
||||
|
||||
async def async_build_responses_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
optional_params: dict,
|
||||
logging_obj: Any,
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
) -> AgenticLoopPlan:
|
||||
tool_calls = tools["tool_calls"]
|
||||
request_patch = await self._build_responses_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={"tool_type": "websearch", "response_format": "responses"},
|
||||
)
|
||||
|
||||
async def _build_responses_request_patch(
|
||||
self,
|
||||
model: str,
|
||||
messages: Union[str, list[dict]],
|
||||
tool_calls: list[dict],
|
||||
optional_params: dict,
|
||||
kwargs: dict,
|
||||
) -> AgenticLoopRequestPatch:
|
||||
"""Execute litellm.asearch() and build a Responses API rerun patch."""
|
||||
search_tasks = [
|
||||
(
|
||||
self._execute_search(tool_call["input"]["query"], kwargs=kwargs)
|
||||
if isinstance(tool_call.get("input"), dict) and tool_call["input"].get("query")
|
||||
else self._create_empty_search_result()
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
|
||||
verbose_logger.debug(f"WebSearchInterception: Executing {len(search_tasks)} responses search(es) in parallel")
|
||||
search_results = await asyncio.gather(*search_tasks, return_exceptions=True)
|
||||
|
||||
search_texts = [self._extract_search_text(result) for result in search_results]
|
||||
|
||||
followup_items = [
|
||||
item
|
||||
for tool_call, search_text in zip(tool_calls, search_texts)
|
||||
for item in (
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": tool_call.get("call_id"),
|
||||
"name": LITELLM_WEB_SEARCH_TOOL_NAME,
|
||||
"arguments": tool_call.get("arguments", ""),
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_call.get("call_id"),
|
||||
"output": search_text,
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
input_list = self._normalize_responses_input(messages) + followup_items
|
||||
|
||||
tools_param = optional_params.get("tools")
|
||||
optional_params_clean = {
|
||||
k: v
|
||||
for k, v in optional_params.items()
|
||||
if k not in {"tools", "tool_choice", "stream", "model_alias_map", "stream_response", "custom_prompt_dict"}
|
||||
}
|
||||
|
||||
kwargs_for_followup = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_websearch_interception")
|
||||
and k
|
||||
not in {
|
||||
"_agentic_loop_api_surface",
|
||||
"litellm_logging_obj",
|
||||
"acompletion",
|
||||
"custom_llm_provider",
|
||||
"model_alias_map",
|
||||
}
|
||||
}
|
||||
|
||||
full_model_name = model
|
||||
if "/" not in model and isinstance(kwargs.get("custom_llm_provider"), str):
|
||||
full_model_name = f"{kwargs['custom_llm_provider']}/{model}"
|
||||
|
||||
verbose_logger.debug(
|
||||
"WebSearchInterception: Built responses request patch model=%s input_items=%d searches=%d",
|
||||
full_model_name,
|
||||
len(input_list),
|
||||
len(search_texts),
|
||||
)
|
||||
|
||||
return AgenticLoopRequestPatch(
|
||||
model=full_model_name,
|
||||
messages=input_list,
|
||||
tools=tools_param if isinstance(tools_param, list) else None,
|
||||
optional_params=optional_params_clean,
|
||||
kwargs=kwargs_for_followup,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_responses_input(messages: Union[str, list[dict]]) -> list[dict]:
|
||||
if isinstance(messages, str):
|
||||
return [{"role": "user", "content": messages}]
|
||||
if isinstance(messages, list):
|
||||
return list(messages)
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _extract_search_text(result: Any) -> str:
|
||||
if isinstance(result, Exception):
|
||||
verbose_logger.error(f"WebSearchInterception: Responses search failed with error: {str(result)}")
|
||||
return f"Search failed: {str(result)}"
|
||||
if isinstance(result, tuple) and len(result) == 2:
|
||||
text_value, _ = result
|
||||
return text_value if isinstance(text_value, str) else str(text_value)
|
||||
verbose_logger.debug(f"WebSearchInterception: Unexpected search result type {type(result)}")
|
||||
return str(result)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_max_tokens(
|
||||
optional_params: Dict,
|
||||
|
|
|
|||
|
|
@ -82,6 +82,75 @@ def get_litellm_web_search_tool_openai() -> Dict[str, Any]:
|
|||
}
|
||||
|
||||
|
||||
def get_litellm_web_search_tool_responses() -> dict[str, Any]:
|
||||
"""
|
||||
Get the standard LiteLLM web search tool definition in Responses API format.
|
||||
|
||||
Used by async_pre_call_deployment_hook on the Responses API path, where a
|
||||
function tool is a flat object (``type: "function"`` with a top-level
|
||||
``name`` and ``parameters``) rather than the nested ``function`` wrapper
|
||||
used by Chat Completions.
|
||||
|
||||
Returns:
|
||||
Dict containing the Responses-style function tool definition.
|
||||
"""
|
||||
return {
|
||||
"type": "function",
|
||||
"name": LITELLM_WEB_SEARCH_TOOL_NAME,
|
||||
"description": (
|
||||
"Search the web for information. Use this when you need current "
|
||||
"information or answers to questions that require up-to-date data."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "The search query to execute",
|
||||
}
|
||||
},
|
||||
"required": ["query"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def is_web_search_tool_responses(tool: dict[str, Any]) -> bool:
|
||||
"""
|
||||
Check if a tool is a web search tool for the Responses API.
|
||||
|
||||
Detects:
|
||||
- OpenAI native Responses web search tools, whose ``type`` is one of
|
||||
``web_search``, ``web_search_2025_08_26``, ``web_search_preview``,
|
||||
``web_search_preview_2025_03_11`` (matched by the ``web_search`` prefix)
|
||||
- The LiteLLM standard function tool in Responses shape:
|
||||
``{"type": "function", "name": "litellm_web_search"}``
|
||||
|
||||
Args:
|
||||
tool: Tool dictionary to check
|
||||
|
||||
Returns:
|
||||
True if tool is a Responses-API web search tool
|
||||
|
||||
Example:
|
||||
>>> is_web_search_tool_responses({"type": "web_search"})
|
||||
True
|
||||
>>> is_web_search_tool_responses({"type": "web_search_preview"})
|
||||
True
|
||||
>>> is_web_search_tool_responses({"type": "function", "name": "litellm_web_search"})
|
||||
True
|
||||
>>> is_web_search_tool_responses({"type": "function", "name": "get_weather"})
|
||||
False
|
||||
"""
|
||||
tool_type = tool.get("type", "")
|
||||
if not isinstance(tool_type, str):
|
||||
return False
|
||||
|
||||
if tool_type == "function":
|
||||
return tool.get("name") == LITELLM_WEB_SEARCH_TOOL_NAME
|
||||
|
||||
return tool_type == "web_search" or tool_type.startswith("web_search_")
|
||||
|
||||
|
||||
def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Check if a tool is a web search tool for Chat Completions API (strict check).
|
||||
|
|
|
|||
|
|
@ -59,9 +59,73 @@ class WebSearchTransformation:
|
|||
# Parse non-streaming response based on format
|
||||
if response_format == "openai":
|
||||
return WebSearchTransformation._detect_from_openai_response(response)
|
||||
elif response_format == "responses":
|
||||
return WebSearchTransformation._detect_from_responses_response(response)
|
||||
else:
|
||||
return WebSearchTransformation._detect_from_non_streaming_response(response)
|
||||
|
||||
@staticmethod
|
||||
def _detect_from_responses_response(
|
||||
response: Any,
|
||||
) -> tuple[bool, list[dict]]:
|
||||
"""Parse a Responses API response for ``litellm_web_search`` function calls.
|
||||
|
||||
After pre-request conversion the native web search tool is replaced by a
|
||||
``litellm_web_search`` function tool, so the model emits ``function_call``
|
||||
items in ``response.output`` instead of a native ``web_search_call``.
|
||||
"""
|
||||
if isinstance(response, dict):
|
||||
output = response.get("output", [])
|
||||
else:
|
||||
output = getattr(response, "output", None) or []
|
||||
|
||||
if not isinstance(output, list):
|
||||
return False, []
|
||||
|
||||
tool_calls: list[dict] = []
|
||||
for item in output:
|
||||
if isinstance(item, dict):
|
||||
item_type = item.get("type")
|
||||
item_name = item.get("name")
|
||||
call_id = item.get("call_id")
|
||||
arguments = item.get("arguments", "")
|
||||
else:
|
||||
item_type = getattr(item, "type", None)
|
||||
item_name = getattr(item, "name", None)
|
||||
call_id = getattr(item, "call_id", None)
|
||||
arguments = getattr(item, "arguments", "")
|
||||
|
||||
if item_type != "function_call" or item_name != LITELLM_WEB_SEARCH_TOOL_NAME:
|
||||
continue
|
||||
|
||||
if isinstance(arguments, str):
|
||||
try:
|
||||
parsed_input = json.loads(arguments) if arguments else {}
|
||||
except json.JSONDecodeError:
|
||||
verbose_logger.warning(
|
||||
f"WebSearchInterception: Failed to parse function_call arguments: {arguments}"
|
||||
)
|
||||
parsed_input = {}
|
||||
elif isinstance(arguments, dict):
|
||||
parsed_input = arguments
|
||||
else:
|
||||
parsed_input = {}
|
||||
|
||||
arguments_str = arguments if isinstance(arguments, str) else json.dumps(parsed_input)
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": call_id,
|
||||
"call_id": call_id,
|
||||
"type": "function_call",
|
||||
"name": item_name,
|
||||
"arguments": arguments_str,
|
||||
"input": parsed_input,
|
||||
}
|
||||
)
|
||||
verbose_logger.debug(f"WebSearchInterception: Found {item_name} function_call with call_id={call_id}")
|
||||
|
||||
return len(tool_calls) > 0, tool_calls
|
||||
|
||||
@staticmethod
|
||||
def _detect_from_non_streaming_response(
|
||||
response: Any,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm import (
|
|||
)
|
||||
from litellm._logging import _is_debugging_on, _redact_string, verbose_logger
|
||||
from litellm.exceptions import (
|
||||
BudgetExceededError,
|
||||
validate_rate_limit_category,
|
||||
validate_rate_limit_type,
|
||||
)
|
||||
|
|
@ -925,7 +926,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 +1135,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 +3073,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:
|
||||
"""
|
||||
|
|
@ -4599,6 +4598,10 @@ class StandardLoggingPayloadSetup:
|
|||
user_api_key_spend=None,
|
||||
user_api_key_max_budget=None,
|
||||
user_api_key_budget_reset_at=None,
|
||||
user_api_key_user_spend=None,
|
||||
user_api_key_user_max_budget=None,
|
||||
user_api_key_team_spend=None,
|
||||
user_api_key_team_max_budget=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_org_alias=None,
|
||||
|
|
@ -4945,6 +4948,7 @@ class StandardLoggingPayloadSetup:
|
|||
|
||||
rate_limit_category = validate_rate_limit_category(getattr(original_exception, "category", None))
|
||||
rate_limit_type = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None))
|
||||
budget_error = original_exception if isinstance(original_exception, BudgetExceededError) else None
|
||||
|
||||
return StandardLoggingPayloadErrorInformation(
|
||||
error_code=error_status,
|
||||
|
|
@ -4954,6 +4958,10 @@ class StandardLoggingPayloadSetup:
|
|||
error_message=error_message,
|
||||
error_rate_limit_category=rate_limit_category,
|
||||
error_rate_limit_type=rate_limit_type,
|
||||
error_budget_entity_type=budget_error.entity_type if budget_error else None,
|
||||
error_budget_entity_id=budget_error.entity_id if budget_error else None,
|
||||
error_budget_limit=budget_error.max_budget if budget_error else None,
|
||||
error_budget_spend=budget_error.current_cost if budget_error else None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -5430,6 +5438,10 @@ def get_standard_logging_metadata(
|
|||
user_api_key_spend=None,
|
||||
user_api_key_max_budget=None,
|
||||
user_api_key_budget_reset_at=None,
|
||||
user_api_key_user_spend=None,
|
||||
user_api_key_user_max_budget=None,
|
||||
user_api_key_team_spend=None,
|
||||
user_api_key_team_max_budget=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_org_alias=None,
|
||||
|
|
@ -5529,6 +5541,10 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
|
|||
user_api_key_team_id=str("test_team"),
|
||||
user_api_key_user_id=str("test_user"),
|
||||
user_api_key_team_alias=str("test_team_alias"),
|
||||
user_api_key_user_spend=None,
|
||||
user_api_key_user_max_budget=None,
|
||||
user_api_key_team_spend=None,
|
||||
user_api_key_team_max_budget=None,
|
||||
user_api_key_org_id=None,
|
||||
spend_logs_metadata=None,
|
||||
requester_ip_address=str("127.0.0.1"),
|
||||
|
|
|
|||
|
|
@ -445,6 +445,7 @@ class PromptTokensDetailsResult(TypedDict):
|
|||
text_tokens: int
|
||||
audio_tokens: int
|
||||
image_tokens: int
|
||||
video_tokens: int
|
||||
character_count: int
|
||||
image_count: int
|
||||
video_length_seconds: float
|
||||
|
|
@ -473,6 +474,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
)
|
||||
audio_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0
|
||||
image_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0
|
||||
video_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0))
|
||||
character_count = (
|
||||
cast(
|
||||
Optional[int],
|
||||
|
|
@ -503,6 +505,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult:
|
|||
text_tokens=text_tokens,
|
||||
audio_tokens=audio_tokens,
|
||||
image_tokens=image_tokens,
|
||||
video_tokens=video_tokens,
|
||||
character_count=character_count,
|
||||
image_count=image_count,
|
||||
video_length_seconds=float(video_length_seconds),
|
||||
|
|
@ -515,6 +518,7 @@ class CompletionTokensDetailsResult(TypedDict):
|
|||
text_tokens: int
|
||||
reasoning_tokens: int
|
||||
image_tokens: int
|
||||
video_tokens: int
|
||||
|
||||
|
||||
def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
|
||||
|
|
@ -546,12 +550,14 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes
|
|||
)
|
||||
or 0
|
||||
)
|
||||
video_tokens = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
|
||||
|
||||
return CompletionTokensDetailsResult(
|
||||
audio_tokens=audio_tokens,
|
||||
text_tokens=text_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
image_tokens=image_tokens,
|
||||
video_tokens=video_tokens,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -586,6 +592,13 @@ def _calculate_input_cost(
|
|||
image_token_cost_key = "input_cost_per_token"
|
||||
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
|
||||
|
||||
### VIDEO TOKEN COST
|
||||
if prompt_tokens_details["video_tokens"]:
|
||||
video_token_cost_key = "input_cost_per_video_token"
|
||||
if model_info.get(video_token_cost_key) is None:
|
||||
video_token_cost_key = "input_cost_per_token"
|
||||
prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"])
|
||||
|
||||
### CACHE WRITING COST - Now uses tiered pricing
|
||||
if (
|
||||
prompt_tokens_details["cache_creation_tokens"]
|
||||
|
|
@ -698,6 +711,7 @@ def generic_cost_per_token(
|
|||
text_tokens=usage.prompt_tokens,
|
||||
audio_tokens=0,
|
||||
image_tokens=0,
|
||||
video_tokens=0,
|
||||
character_count=0,
|
||||
image_count=0,
|
||||
video_length_seconds=0.0,
|
||||
|
|
@ -716,13 +730,14 @@ def generic_cost_per_token(
|
|||
audio_tokens = prompt_tokens_details["audio_tokens"]
|
||||
cache_creation = prompt_tokens_details["cache_creation_tokens"]
|
||||
image_tokens = prompt_tokens_details["image_tokens"]
|
||||
video_tokens = prompt_tokens_details["video_tokens"]
|
||||
|
||||
# Check for double-counting: sum of details > prompt_tokens means overlap
|
||||
total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens
|
||||
total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens
|
||||
has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens
|
||||
|
||||
if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting:
|
||||
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens
|
||||
text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens
|
||||
# Clamp to zero: inconsistent streaming usage
|
||||
if text_tokens < 0:
|
||||
text_tokens = 0
|
||||
|
|
@ -751,6 +766,7 @@ def generic_cost_per_token(
|
|||
audio_tokens = 0
|
||||
reasoning_tokens = 0
|
||||
image_tokens = 0
|
||||
video_tokens = 0
|
||||
is_text_tokens_total = False
|
||||
if usage.completion_tokens_details is not None:
|
||||
completion_tokens_details = _parse_completion_tokens_details(usage)
|
||||
|
|
@ -758,19 +774,20 @@ def generic_cost_per_token(
|
|||
text_tokens = completion_tokens_details["text_tokens"]
|
||||
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
|
||||
image_tokens = completion_tokens_details["image_tokens"]
|
||||
video_tokens = completion_tokens_details["video_tokens"]
|
||||
|
||||
# Handle text_tokens calculation:
|
||||
# 1. If text_tokens is explicitly provided and > 0, use it
|
||||
# 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder
|
||||
# 2. If there's a breakdown (reasoning/audio/image/video tokens), calculate text_tokens as the remainder
|
||||
# 3. If no breakdown at all, assume all completion_tokens are text_tokens
|
||||
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0
|
||||
has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 or video_tokens > 0
|
||||
if text_tokens == 0:
|
||||
if has_token_breakdown:
|
||||
# Calculate text tokens as remainder when we have a breakdown
|
||||
# This handles cases like OpenAI's reasoning models where text_tokens isn't provided
|
||||
text_tokens = max(
|
||||
0,
|
||||
usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens,
|
||||
usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens - video_tokens,
|
||||
)
|
||||
else:
|
||||
# No breakdown at all, all tokens are text tokens
|
||||
|
|
@ -803,6 +820,14 @@ def generic_cost_per_token(
|
|||
)
|
||||
completion_cost += float(image_tokens) * _output_cost_per_image_token
|
||||
|
||||
## VIDEO COST
|
||||
if not is_text_tokens_total and video_tokens and video_tokens > 0:
|
||||
_output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None)
|
||||
_output_cost_per_video_token = (
|
||||
_output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost
|
||||
)
|
||||
completion_cost += float(video_tokens) * _output_cost_per_video_token
|
||||
|
||||
## REGIONAL DATA-RESIDENCY UPLIFT
|
||||
# Applied as a flat multiplier across all token costs for the request
|
||||
# when the upstream is a regionalized OpenAI host (eu./us.api.openai.com).
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -467,6 +467,7 @@ class ChunkProcessor:
|
|||
cache_read_input_tokens: Optional[int] = None
|
||||
completion_tokens_details: Optional[CompletionTokensDetails] = None
|
||||
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
|
||||
cost: Optional[float] = None
|
||||
|
||||
if "prompt_tokens" in usage_chunk:
|
||||
prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0
|
||||
|
|
@ -476,6 +477,8 @@ class ChunkProcessor:
|
|||
cache_creation_input_tokens = usage_chunk.get("cache_creation_input_tokens")
|
||||
if "cache_read_input_tokens" in usage_chunk:
|
||||
cache_read_input_tokens = usage_chunk.get("cache_read_input_tokens")
|
||||
if "cost" in usage_chunk:
|
||||
cost = usage_chunk.get("cost")
|
||||
if hasattr(usage_chunk, "completion_tokens_details"):
|
||||
if isinstance(usage_chunk.completion_tokens_details, dict):
|
||||
completion_tokens_details = CompletionTokensDetails(**usage_chunk.completion_tokens_details)
|
||||
|
|
@ -494,6 +497,7 @@ class ChunkProcessor:
|
|||
"cache_read_input_tokens": cache_read_input_tokens,
|
||||
"completion_tokens_details": completion_tokens_details,
|
||||
"prompt_tokens_details": prompt_tokens_details,
|
||||
"cost": cost,
|
||||
}
|
||||
|
||||
def count_reasoning_tokens(self, response: ModelResponse) -> Optional[int]:
|
||||
|
|
@ -512,6 +516,22 @@ class ChunkProcessor:
|
|||
|
||||
return reasoning_tokens
|
||||
|
||||
@staticmethod
|
||||
def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
|
||||
usage_chunk: Usage | dict[str, Any] | None = None
|
||||
if hasattr(chunk, "usage") and chunk.usage is not None:
|
||||
usage_chunk = chunk.usage
|
||||
elif "usage" in chunk:
|
||||
usage_chunk = chunk["usage"]
|
||||
elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
|
||||
chunk, "_hidden_params"
|
||||
):
|
||||
usage_chunk = chunk._hidden_params.get("usage", None)
|
||||
|
||||
if isinstance(usage_chunk, dict):
|
||||
return Usage(**usage_chunk)
|
||||
return usage_chunk
|
||||
|
||||
def _calculate_usage_per_chunk(
|
||||
self,
|
||||
chunks: List[Union[Dict[str, Any], ModelResponse]],
|
||||
|
|
@ -548,18 +568,12 @@ class ChunkProcessor:
|
|||
# is last-wins, so without preserving this separately the 1h breakdown is
|
||||
# lost and 1h cache writes get billed at the 5m rate.
|
||||
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
|
||||
cost: Optional[float] = None
|
||||
|
||||
for chunk in chunks:
|
||||
usage_chunk: Optional[Usage] = None
|
||||
if "usage" in chunk:
|
||||
usage_chunk = chunk["usage"]
|
||||
elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
|
||||
chunk, "_hidden_params"
|
||||
):
|
||||
usage_chunk = chunk._hidden_params.get("usage", None)
|
||||
usage_chunk = self._extract_usage_chunk(chunk)
|
||||
|
||||
if usage_chunk is not None:
|
||||
if isinstance(usage_chunk, dict):
|
||||
usage_chunk = Usage(**usage_chunk)
|
||||
usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk)
|
||||
if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0:
|
||||
prompt_tokens = usage_chunk_dict["prompt_tokens"]
|
||||
|
|
@ -610,6 +624,9 @@ class ChunkProcessor:
|
|||
prompt_tokens_details, cache_creation_token_details
|
||||
)
|
||||
|
||||
if usage_chunk_dict["cost"] is not None:
|
||||
cost = usage_chunk_dict["cost"]
|
||||
|
||||
prompt_tokens_details = self._attach_cache_creation_token_details(
|
||||
prompt_tokens_details, cache_creation_token_details
|
||||
)
|
||||
|
|
@ -629,6 +646,7 @@ class ChunkProcessor:
|
|||
web_search_requests=web_search_requests,
|
||||
completion_tokens_details=completion_tokens_details,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
cost=cost,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -727,6 +745,7 @@ class ChunkProcessor:
|
|||
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = calculated_usage_per_chunk[
|
||||
"prompt_tokens_details"
|
||||
]
|
||||
cost: Optional[float] = calculated_usage_per_chunk["cost"]
|
||||
|
||||
try:
|
||||
returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages)
|
||||
|
|
@ -784,6 +803,9 @@ class ChunkProcessor:
|
|||
else:
|
||||
returned_usage.prompt_tokens_details.web_search_requests = web_search_requests
|
||||
|
||||
if cost is not None:
|
||||
setattr(returned_usage, "cost", cost)
|
||||
|
||||
# Return a new usage object with the new values
|
||||
|
||||
returned_usage = Usage(**returned_usage.model_dump())
|
||||
|
|
|
|||
|
|
@ -962,10 +962,11 @@ class CustomStreamWrapper:
|
|||
if self.custom_llm_provider == "bedrock" and "trace" in model_response:
|
||||
return model_response
|
||||
|
||||
# Default - return StopIteration
|
||||
if hasattr(model_response, "usage"):
|
||||
self.chunks.append(model_response)
|
||||
raise StopIteration
|
||||
# Don't raise StopIteration here - some providers (like OpenRouter)
|
||||
# send usage/cost data in chunks after the finish_reason chunk
|
||||
if hasattr(model_response, "usage") and model_response.usage is not None:
|
||||
return model_response
|
||||
return
|
||||
# flush any remaining holding chunk
|
||||
if len(self.holding_chunk) > 0:
|
||||
if model_response.choices[0].delta.content is None:
|
||||
|
|
@ -1474,12 +1475,16 @@ class CustomStreamWrapper:
|
|||
|
||||
self.tool_call = True
|
||||
|
||||
if hasattr(chunk, "usage") and chunk.usage is not None:
|
||||
model_response.usage = chunk.usage
|
||||
|
||||
## RETURN ARG
|
||||
return self.return_processed_chunk_logic(
|
||||
result = self.return_processed_chunk_logic(
|
||||
completion_obj=completion_obj,
|
||||
model_response=model_response, # type: ignore
|
||||
response_obj=response_obj,
|
||||
)
|
||||
return result
|
||||
|
||||
except StopIteration:
|
||||
raise StopIteration
|
||||
|
|
@ -1686,6 +1691,21 @@ class CustomStreamWrapper:
|
|||
model_response.choices[0].finish_reason = "tool_calls"
|
||||
return model_response
|
||||
|
||||
@staticmethod
|
||||
def _propagate_usage_cost_to_hidden_params(
|
||||
response: "ModelResponse",
|
||||
) -> None:
|
||||
"""
|
||||
If the assembled response carries a provider-reported cost on
|
||||
usage.cost, copy it into _hidden_params so litellm's cost
|
||||
calculator uses it instead of a token-based estimate.
|
||||
"""
|
||||
_usage = getattr(response, "usage", None)
|
||||
if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None:
|
||||
if "additional_headers" not in response._hidden_params:
|
||||
response._hidden_params["additional_headers"] = {}
|
||||
response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost)
|
||||
|
||||
def __next__(self) -> "ModelResponseStream":
|
||||
cache_hit = False
|
||||
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
|
||||
|
|
@ -1741,6 +1761,10 @@ class CustomStreamWrapper:
|
|||
# hasattr(response, "usage") is always True — must check
|
||||
# `is not None` to avoid running this path on every chunk.
|
||||
if getattr(response, "usage", None) is not None:
|
||||
usage_to_preserve = response.usage
|
||||
if usage_to_preserve:
|
||||
response._hidden_params["usage"] = usage_to_preserve
|
||||
|
||||
obj_dict = response.model_dump()
|
||||
|
||||
if "usage" in obj_dict:
|
||||
|
|
@ -1789,6 +1813,8 @@ class CustomStreamWrapper:
|
|||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
|
||||
|
||||
setattr(
|
||||
response,
|
||||
"usage",
|
||||
|
|
@ -1974,97 +2000,7 @@ class CustomStreamWrapper:
|
|||
self.chunks.append(processed_chunk)
|
||||
return processed_chunk
|
||||
except (StopAsyncIteration, StopIteration):
|
||||
if self.sent_last_chunk is True:
|
||||
# log the final chunk with accurate streaming values
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
# see sync __next__: a raise from stream_chunk_builder inside this
|
||||
# except handler escapes __anext__ and drops the request from SpendLogs.
|
||||
# Recover best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
setattr(
|
||||
response,
|
||||
"usage",
|
||||
getattr(complete_streaming_response, "usage"),
|
||||
)
|
||||
try:
|
||||
_copy = complete_streaming_response.model_copy(deep=True)
|
||||
except RuntimeError:
|
||||
_copy = complete_streaming_response.model_copy()
|
||||
asyncio.create_task(
|
||||
self.async_cache_streaming_response(
|
||||
processed_chunk=_copy,
|
||||
cache_hit=cache_hit,
|
||||
)
|
||||
)
|
||||
# Update hidden_params with final usage from
|
||||
# stream_chunk_builder (see sync __next__ for full comment).
|
||||
if (
|
||||
self.stream_options is None
|
||||
and complete_streaming_response is not None
|
||||
and self._last_returned_hidden_params is not None
|
||||
):
|
||||
final_usage = getattr(complete_streaming_response, "usage", None)
|
||||
if final_usage is not None:
|
||||
self._last_returned_hidden_params["usage"] = final_usage
|
||||
|
||||
if self.sent_stream_usage is False and self.send_stream_usage is True:
|
||||
self.sent_stream_usage = True
|
||||
return response
|
||||
|
||||
_deferred_cb = getattr(
|
||||
self.logging_obj,
|
||||
"_on_deferred_stream_complete",
|
||||
None,
|
||||
)
|
||||
if _deferred_cb is not None:
|
||||
# Proxy has post-call guardrails. Store the assembled
|
||||
# response so the outer streaming consumer
|
||||
# (ProxyLogging.async_post_call_streaming_iterator_hook)
|
||||
# can fire the deferred callback AFTER all guardrail
|
||||
# end-of-stream blocks complete. Scheduling here via
|
||||
# create_task would race with unified_guardrail's
|
||||
# end-of-stream block for short-stream providers.
|
||||
self.logging_obj._deferred_stream_complete_args = ( # type: ignore[attr-defined]
|
||||
complete_streaming_response,
|
||||
cache_hit,
|
||||
)
|
||||
else:
|
||||
# prefer_async_handlers routes CustomLogger to async_success_handler
|
||||
# when consumers use ``async for`` on sync-SDK streams. Legacy string
|
||||
# callbacks still run via executor.submit inside dispatch_success_handlers.
|
||||
asyncio.create_task(
|
||||
self.logging_obj.dispatch_success_handlers(
|
||||
complete_streaming_response,
|
||||
cache_hit=cache_hit,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
|
||||
raise StopAsyncIteration # Re-raise StopIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
processed_chunk = self.finish_reason_handler()
|
||||
return processed_chunk
|
||||
return await self._finalize_completed_stream(cache_hit=cache_hit)
|
||||
except httpx.TimeoutException as e: # if httpx read timeout error occues
|
||||
traceback_exception = traceback.format_exc()
|
||||
## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT
|
||||
|
|
@ -2079,20 +2015,122 @@ class CustomStreamWrapper:
|
|||
# Handle any exceptions that might occur during streaming
|
||||
asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception))
|
||||
self._handle_stream_fallback_error(e)
|
||||
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
|
||||
if self.received_finish_reason is None:
|
||||
self._log_stream_failure_and_raise(e)
|
||||
return await self._finalize_completed_stream(cache_hit=cache_hit)
|
||||
except Exception as e:
|
||||
traceback_exception = traceback.format_exc()
|
||||
if self.logging_obj is not None:
|
||||
self._record_partial_usage_for_failure()
|
||||
## LOGGING
|
||||
threading.Thread(
|
||||
target=self.logging_obj.failure_handler,
|
||||
args=(e, traceback_exception),
|
||||
).start() # log response
|
||||
# Handle any exceptions that might occur during streaming
|
||||
asyncio.create_task(
|
||||
self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore
|
||||
self._log_stream_failure_and_raise(e)
|
||||
|
||||
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
|
||||
if self.sent_last_chunk is True:
|
||||
# log the final chunk with accurate streaming values
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
self._handle_stream_fallback_error(e)
|
||||
except Exception as e:
|
||||
# see sync __next__: a raise from stream_chunk_builder inside this
|
||||
# except handler escapes __anext__ and drops the request from SpendLogs.
|
||||
# Recover best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
|
||||
|
||||
setattr(
|
||||
response,
|
||||
"usage",
|
||||
getattr(complete_streaming_response, "usage"),
|
||||
)
|
||||
try:
|
||||
_copy = complete_streaming_response.model_copy(deep=True)
|
||||
except RuntimeError:
|
||||
_copy = complete_streaming_response.model_copy()
|
||||
asyncio.create_task(
|
||||
self.async_cache_streaming_response(
|
||||
processed_chunk=_copy,
|
||||
cache_hit=cache_hit,
|
||||
)
|
||||
)
|
||||
# Update hidden_params with final usage from
|
||||
# stream_chunk_builder (see sync __next__ for full comment).
|
||||
if (
|
||||
self.stream_options is None
|
||||
and complete_streaming_response is not None
|
||||
and self._last_returned_hidden_params is not None
|
||||
):
|
||||
final_usage = getattr(complete_streaming_response, "usage", None)
|
||||
if final_usage is not None:
|
||||
self._last_returned_hidden_params["usage"] = final_usage
|
||||
|
||||
if self.sent_stream_usage is False and self.send_stream_usage is True:
|
||||
self.sent_stream_usage = True
|
||||
return response
|
||||
|
||||
_deferred_cb = getattr(
|
||||
self.logging_obj,
|
||||
"_on_deferred_stream_complete",
|
||||
None,
|
||||
)
|
||||
if _deferred_cb is not None:
|
||||
# Proxy has post-call guardrails. Store the assembled
|
||||
# response so the outer streaming consumer
|
||||
# (ProxyLogging.async_post_call_streaming_iterator_hook)
|
||||
# can fire the deferred callback AFTER all guardrail
|
||||
# end-of-stream blocks complete. Scheduling here via
|
||||
# create_task would race with unified_guardrail's
|
||||
# end-of-stream block for short-stream providers.
|
||||
self.logging_obj._deferred_stream_complete_args = ( # type: ignore[attr-defined]
|
||||
complete_streaming_response,
|
||||
cache_hit,
|
||||
)
|
||||
else:
|
||||
# prefer_async_handlers routes CustomLogger to async_success_handler
|
||||
# when consumers use ``async for`` on sync-SDK streams. Legacy string
|
||||
# callbacks still run via executor.submit inside dispatch_success_handlers.
|
||||
asyncio.create_task(
|
||||
self.logging_obj.dispatch_success_handlers(
|
||||
complete_streaming_response,
|
||||
cache_hit=cache_hit,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
|
||||
raise StopAsyncIteration # Re-raise StopIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
processed_chunk = self.finish_reason_handler()
|
||||
return processed_chunk
|
||||
|
||||
def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn:
|
||||
traceback_exception = traceback.format_exc()
|
||||
if self.logging_obj is not None:
|
||||
self._record_partial_usage_for_failure()
|
||||
## LOGGING
|
||||
threading.Thread(
|
||||
target=self.logging_obj.failure_handler,
|
||||
args=(e, traceback_exception),
|
||||
).start() # log response
|
||||
# Handle any exceptions that might occur during streaming
|
||||
asyncio.create_task(
|
||||
self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore
|
||||
)
|
||||
self._handle_stream_fallback_error(e)
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
"""
|
||||
|
|
@ -2228,12 +2266,16 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
|
|||
"""Assume most recent usage chunk has total usage uptil then."""
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
latest_usage_chunk = None
|
||||
|
||||
for chunk in chunks:
|
||||
if "usage" in chunk and chunk["usage"] is not None:
|
||||
if "prompt_tokens" in chunk["usage"]:
|
||||
prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0
|
||||
if "completion_tokens" in chunk["usage"]:
|
||||
completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0
|
||||
usage = chunk["usage"]
|
||||
latest_usage_chunk = usage
|
||||
if "prompt_tokens" in usage:
|
||||
prompt_tokens = usage.get("prompt_tokens", 0) or 0
|
||||
if "completion_tokens" in usage:
|
||||
completion_tokens = usage.get("completion_tokens", 0) or 0
|
||||
|
||||
returned_usage_chunk = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
|
|
@ -2241,6 +2283,15 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
|
|||
total_tokens=prompt_tokens + completion_tokens,
|
||||
)
|
||||
|
||||
if latest_usage_chunk is not None:
|
||||
latest_cost = (
|
||||
latest_usage_chunk.get("cost")
|
||||
if isinstance(latest_usage_chunk, dict)
|
||||
else getattr(latest_usage_chunk, "cost", None)
|
||||
)
|
||||
if latest_cost is not None:
|
||||
returned_usage_chunk.cost = latest_cost
|
||||
|
||||
return returned_usage_chunk
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -13,14 +13,18 @@ from typing import (
|
|||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
get_args,
|
||||
)
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.types.llms.anthropic import (
|
||||
AppliedEdit,
|
||||
CompactionBlock,
|
||||
ContextManagementResponse,
|
||||
StreamingContentBlockDeltaType,
|
||||
UsageDelta,
|
||||
UsageIteration,
|
||||
)
|
||||
|
|
@ -30,6 +34,23 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
|
||||
_STREAMING_DELTA_TYPES = frozenset(get_args(StreamingContentBlockDeltaType))
|
||||
|
||||
|
||||
def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str:
|
||||
match delta_type:
|
||||
case "text_delta":
|
||||
return "text"
|
||||
case "input_json_delta":
|
||||
return "partial_json"
|
||||
case "thinking_delta":
|
||||
return "thinking"
|
||||
case "signature_delta":
|
||||
return "signature"
|
||||
case _:
|
||||
assert_never(delta_type)
|
||||
|
||||
|
||||
class _CombinedChunkSplitter:
|
||||
"""
|
||||
Splits a streaming chunk that carries BOTH response content and a
|
||||
|
|
@ -428,7 +449,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
else:
|
||||
processed_first = self._augment_message_delta_usage(processed_first)
|
||||
self.chunk_queue.append(processed_first)
|
||||
elif self._trigger_delta_has_content(processed_first):
|
||||
elif self._delta_has_content(processed_first):
|
||||
self.chunk_queue.append(processed_first)
|
||||
else:
|
||||
self.chunk_queue.append(
|
||||
|
|
@ -500,12 +521,15 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
|
||||
# 3. If the trigger chunk carries delta content, queue it
|
||||
# so the first delta of the new block is not silently dropped.
|
||||
if self._trigger_delta_has_content(processed_chunk):
|
||||
if self._delta_has_content(processed_chunk):
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
|
||||
self.sent_content_block_finish = False
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
if processed_chunk["type"] == "content_block_delta" and not self._delta_has_content(processed_chunk):
|
||||
continue
|
||||
|
||||
if processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False:
|
||||
# Queue both the content_block_stop and the message_delta
|
||||
self.chunk_queue.append(
|
||||
|
|
@ -686,7 +710,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
else:
|
||||
processed_first = self._augment_message_delta_usage(processed_first)
|
||||
self.chunk_queue.append(processed_first)
|
||||
elif self._trigger_delta_has_content(processed_first):
|
||||
elif self._delta_has_content(processed_first):
|
||||
self.chunk_queue.append(processed_first)
|
||||
else:
|
||||
self.chunk_queue.append(
|
||||
|
|
@ -749,11 +773,19 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
"content_block": self.current_content_block_start,
|
||||
}
|
||||
)
|
||||
if self._trigger_delta_has_content(processed_chunk):
|
||||
|
||||
# 3. If the trigger chunk carries delta content, queue it
|
||||
# so the first delta of the new block is not silently dropped.
|
||||
if self._delta_has_content(processed_chunk):
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
self.sent_content_block_finish = False
|
||||
return self.chunk_queue.popleft()
|
||||
|
||||
if processed_chunk["type"] == "content_block_delta" and not self._delta_has_content(
|
||||
processed_chunk
|
||||
):
|
||||
continue
|
||||
|
||||
if processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False:
|
||||
# Queue both the content_block_stop and the holding chunk
|
||||
self.chunk_queue.append(
|
||||
|
|
@ -885,20 +917,33 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
self.current_content_block_index += 1
|
||||
|
||||
@staticmethod
|
||||
def _trigger_delta_has_content(processed_chunk: Dict[str, Any]) -> bool:
|
||||
"""Return True if a translated trigger chunk carries a non-empty
|
||||
``content_block_delta`` payload that must be re-emitted after a
|
||||
block transition.
|
||||
def _delta_has_content(processed_chunk: Dict[str, Any]) -> bool:
|
||||
"""Return True if a translated chunk carries a non-empty
|
||||
``content_block_delta`` payload.
|
||||
|
||||
When an upstream chunk both *triggers* a new content block (its type
|
||||
differs from the active block) and *carries* delta content, that
|
||||
content belongs to the new block. The synthesized
|
||||
``content_block_start`` only ever carries an empty body — see
|
||||
Gates every ``content_block_delta`` emission. An empty delta carries
|
||||
no information, and the translate fallback types empty deltas as
|
||||
``text_delta`` regardless of the active block's type — emitting one
|
||||
into an open ``thinking`` block (e.g. Bedrock Converse sends an empty
|
||||
reasoning delta mid-block) crashes strict Anthropic SDK clients with
|
||||
"Content block is not a text block".
|
||||
|
||||
Also gates re-emission after a block transition: when an upstream
|
||||
chunk both *triggers* a new content block (its type differs from the
|
||||
active block) and *carries* delta content, that content belongs to
|
||||
the new block. The synthesized ``content_block_start`` only ever
|
||||
carries an empty body — see
|
||||
``_translate_streaming_openai_chunk_to_anthropic_content_block``,
|
||||
which returns an empty ``TextBlock``/``ToolUseBlock``/thinking block —
|
||||
so the trigger chunk's delta must be re-queued or the first token of
|
||||
the new block (the first non-empty text/thinking delta, or bundled
|
||||
tool arguments) is silently dropped.
|
||||
|
||||
Delta types outside ``StreamingContentBlockDeltaType`` — the closed
|
||||
set the translate layer can produce — are treated as empty. The
|
||||
per-type payload lookup is exhaustively matched against that set in
|
||||
``_delta_payload_field``, so extending the translate layer with a new
|
||||
delta type fails type-checking here until it is handled.
|
||||
"""
|
||||
if processed_chunk.get("type") != "content_block_delta":
|
||||
return False
|
||||
|
|
@ -906,15 +951,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
if not isinstance(delta, dict):
|
||||
return False
|
||||
delta_type = delta.get("type")
|
||||
if delta_type == "text_delta":
|
||||
return bool(delta.get("text"))
|
||||
if delta_type == "input_json_delta":
|
||||
return bool(delta.get("partial_json"))
|
||||
if delta_type == "thinking_delta":
|
||||
return bool(delta.get("thinking"))
|
||||
if delta_type == "signature_delta":
|
||||
return bool(delta.get("signature"))
|
||||
return False
|
||||
if delta_type not in _STREAMING_DELTA_TYPES:
|
||||
return False
|
||||
return bool(delta.get(_delta_payload_field(delta_type)))
|
||||
|
||||
def _restore_tool_name_mapping(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -104,6 +104,7 @@ from litellm.types.llms.anthropic import (
|
|||
ContextManagementResponse,
|
||||
MessageBlockDelta,
|
||||
MessageDelta,
|
||||
StreamingContentBlockDeltaType,
|
||||
UsageDelta,
|
||||
UsageIteration,
|
||||
)
|
||||
|
|
@ -1402,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
|
||||
)
|
||||
|
|
@ -1419,7 +1415,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
def _translate_streaming_openai_chunk_to_anthropic(
|
||||
self, choices: List[Union[OpenAIStreamingChoice, StreamingChoices]]
|
||||
) -> Tuple[
|
||||
Literal["text_delta", "input_json_delta", "thinking_delta", "signature_delta"],
|
||||
StreamingContentBlockDeltaType,
|
||||
Union[
|
||||
ContentTextBlockDelta,
|
||||
ContentJsonBlockDelta,
|
||||
|
|
@ -1457,17 +1453,14 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
if choice.delta.reasoning_content:
|
||||
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)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -11,10 +12,30 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class StreamTransformSink:
|
||||
"""Out-parameter used by ``process_output_streaming_response`` to hand the
|
||||
guardrailed streaming state back to the caller.
|
||||
|
||||
The streaming text-transform path must not mutate ``responses_so_far`` (it is
|
||||
the raw accumulator the guardrail re-reads every round), so the guardrailed
|
||||
accumulated text per choice (``mutated_text_per_choice``, keyed by
|
||||
``StreamingChoices.index``) and the per-choice trailing holdback the guardrail
|
||||
requested (``holdback_per_choice``, from ``stream_holdback_chars``) are
|
||||
reported here instead of in place. Only the OpenAI chat handler populates this
|
||||
today; the hook passes a fresh sink per round and reads it afterwards. A
|
||||
mutable dataclass is deliberate: it is a write-once output parameter for a
|
||||
single call, not shared state.
|
||||
"""
|
||||
|
||||
mutated_text_per_choice: dict[int, str] = field(default_factory=dict)
|
||||
holdback_per_choice: dict[int, int] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BaseTranslation(ABC):
|
||||
@staticmethod
|
||||
def transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict: Optional[Any],
|
||||
user_api_key_dict: Any | None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Transform user_api_key_dict to a metadata dict with prefixed keys.
|
||||
|
|
@ -73,7 +94,7 @@ class BaseTranslation(ABC):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
request_data: dict | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Process output response with guardrails.
|
||||
|
|
@ -92,12 +113,15 @@ class BaseTranslation(ABC):
|
|||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
request_data: dict | None = None,
|
||||
stream_transform_sink: StreamTransformSink | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Process output streaming response with guardrails.
|
||||
|
||||
Optional to override in subclasses.
|
||||
Optional to override in subclasses. ``stream_transform_sink`` is the
|
||||
out-parameter used by handlers that support streaming text
|
||||
transformations (see ``StreamTransformSink``); base handlers ignore it.
|
||||
"""
|
||||
return responses_so_far
|
||||
|
||||
|
|
@ -105,8 +129,8 @@ class BaseTranslation(ABC):
|
|||
self,
|
||||
exc: "ModifyResponseException",
|
||||
stream_started: bool = False,
|
||||
responses_so_far: Optional[list[Any]] = None,
|
||||
) -> Optional[list[bytes]]:
|
||||
responses_so_far: list[Any] | None = None,
|
||||
) -> list[bytes] | None:
|
||||
"""
|
||||
Build the streaming chunks that deliver a guardrail block message and
|
||||
cleanly terminate the stream in this provider's wire format.
|
||||
|
|
@ -125,7 +149,7 @@ class BaseTranslation(ABC):
|
|||
"""
|
||||
return None
|
||||
|
||||
def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]:
|
||||
def get_structured_messages(self, data: dict) -> List["AllMessageValues"] | None:
|
||||
"""
|
||||
Convert request data to OpenAI-spec structured messages.
|
||||
|
||||
|
|
|
|||
|
|
@ -85,29 +85,12 @@ class AiohttpResponseStream(httpx.AsyncByteStream):
|
|||
try:
|
||||
async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE):
|
||||
yield chunk
|
||||
except (
|
||||
aiohttp.ClientPayloadError,
|
||||
aiohttp.client_exceptions.ClientPayloadError,
|
||||
) as e:
|
||||
# Handle incomplete transfers more gracefully
|
||||
# Log the error but don't re-raise if we've already yielded some data
|
||||
verbose_logger.debug(f"Transfer incomplete, but continuing: {e}")
|
||||
# If the error is due to incomplete transfer encoding, we can still
|
||||
# return what we've received so far, similar to how httpx handles it
|
||||
return
|
||||
except RuntimeError as e:
|
||||
# Some providers (e.g., SSE streams) may close the connection
|
||||
# causing aiohttp StreamReader to raise a generic RuntimeError
|
||||
# with message "Connection closed.". Treat this as a graceful
|
||||
# end-of-stream so downstream consumers don't error.
|
||||
if "Connection closed" in str(e):
|
||||
verbose_logger.debug("Upstream closed streaming connection; ending iterator gracefully")
|
||||
return
|
||||
raise
|
||||
if "Connection closed" not in str(e):
|
||||
raise
|
||||
raise httpx.ReadError(str(e)) from e
|
||||
except aiohttp.http_exceptions.TransferEncodingError as e:
|
||||
# Handle transfer encoding errors gracefully
|
||||
verbose_logger.debug(f"Transfer encoding error, but continuing: {e}")
|
||||
return
|
||||
raise httpx.ReadError(str(e)) from e
|
||||
except Exception:
|
||||
# For other exceptions, use the normal mapping
|
||||
with map_aiohttp_exceptions():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -2657,9 +2690,10 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
result = final_response if final_response is not None else initial_response
|
||||
if litellm_params.get("_code_interpreter_interception_converted_stream") and not litellm_params.get(
|
||||
"_agentic_loop_depth"
|
||||
):
|
||||
interception_converted_stream = litellm_params.get(
|
||||
"_code_interpreter_interception_converted_stream"
|
||||
) or litellm_params.get("_websearch_interception_converted_stream")
|
||||
if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"):
|
||||
return self._wrap_responses_response_as_fake_stream(
|
||||
result=result,
|
||||
model=model,
|
||||
|
|
@ -5224,6 +5258,8 @@ class BaseLLMHTTPHandler:
|
|||
tools = anthropic_messages_optional_request_params.get("tools", [])
|
||||
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
|
||||
|
||||
hook_kwargs = {**kwargs, "_agentic_loop_api_surface": api_surface}
|
||||
|
||||
for callback in callbacks:
|
||||
if not isinstance(callback, CustomLogger):
|
||||
continue
|
||||
|
|
@ -5244,7 +5280,7 @@ class BaseLLMHTTPHandler:
|
|||
tools=tools,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
kwargs=hook_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
|
|
@ -5270,7 +5306,7 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
|
||||
try:
|
||||
kwargs_with_provider = kwargs.copy() if kwargs else {}
|
||||
kwargs_with_provider = hook_kwargs.copy()
|
||||
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
|
||||
build_plan_overridden = (
|
||||
callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan
|
||||
|
|
|
|||
|
|
@ -14,11 +14,14 @@ Pattern Overview:
|
|||
This pattern can be replicated for other message formats (e.g., Anthropic).
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
StreamTransformSink,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_skip_system_message_for_guardrail,
|
||||
effective_skip_tool_message_for_guardrail,
|
||||
|
|
@ -27,6 +30,9 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
)
|
||||
from litellm.main import stream_chunk_builder
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
coerce_stream_holdback_value,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
Choices,
|
||||
GenericGuardrailAPIInputs,
|
||||
|
|
@ -50,7 +56,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
Methods can be overridden to customize behavior for different message formats.
|
||||
"""
|
||||
|
||||
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
|
||||
def get_structured_messages(self, data: dict) -> List[AllMessageValues] | None:
|
||||
"""
|
||||
Convert chat completions request data to OpenAI-spec structured messages.
|
||||
|
||||
|
|
@ -65,7 +71,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
self,
|
||||
data: dict,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional[Any] = None,
|
||||
litellm_logging_obj: Any | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Process input messages by applying guardrails to text content.
|
||||
|
|
@ -80,7 +86,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
texts_to_check: List[str] = []
|
||||
images_to_check: List[str] = []
|
||||
tool_calls_to_check: List[ChatCompletionToolParam] = []
|
||||
text_task_mappings: List[Tuple[int, Optional[int]]] = []
|
||||
text_task_mappings: List[Tuple[int, int | None]] = []
|
||||
tool_call_task_mappings: List[Tuple[int, int]] = []
|
||||
|
||||
# Step 1: Extract all text content, images, and tool calls
|
||||
|
|
@ -184,7 +190,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
texts_to_check: List[str],
|
||||
images_to_check: List[str],
|
||||
tool_calls_to_check: List[ChatCompletionToolParam],
|
||||
text_task_mappings: List[Tuple[int, Optional[int]]],
|
||||
text_task_mappings: List[Tuple[int, int | None]],
|
||||
tool_call_task_mappings: List[Tuple[int, int]],
|
||||
skip_system_message: bool = False,
|
||||
skip_tool_message: bool = False,
|
||||
|
|
@ -239,7 +245,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
responses: List[str],
|
||||
task_mappings: List[Tuple[int, Optional[int]]],
|
||||
task_mappings: List[Tuple[int, int | None]],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail responses back to input message text content.
|
||||
|
|
@ -249,7 +255,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
for task_idx, guardrail_response in enumerate(responses):
|
||||
mapping = task_mappings[task_idx]
|
||||
msg_idx = cast(int, mapping[0])
|
||||
content_idx_optional = cast(Optional[int], mapping[1])
|
||||
content_idx_optional = cast(int | None, mapping[1])
|
||||
|
||||
# Handle content
|
||||
content = messages[msg_idx].get("content", None)
|
||||
|
|
@ -291,9 +297,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
self,
|
||||
response: "ModelResponse",
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional[Any] = None,
|
||||
user_api_key_dict: Optional[Any] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
litellm_logging_obj: Any | None = None,
|
||||
user_api_key_dict: Any | None = None,
|
||||
request_data: dict | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Process output response by applying guardrails to text content.
|
||||
|
|
@ -320,7 +326,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
texts_to_check: List[str] = []
|
||||
images_to_check: List[str] = []
|
||||
tool_calls_to_check: List[Dict[str, Any]] = []
|
||||
text_task_mappings: List[Tuple[int, Optional[int]]] = []
|
||||
text_task_mappings: List[Tuple[int, int | None]] = []
|
||||
tool_call_task_mappings: List[Tuple[int, int]] = []
|
||||
# text_task_mappings: Track (choice_index, content_index) for each text
|
||||
# content_index is None for string content, int for list content
|
||||
|
|
@ -402,9 +408,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
self,
|
||||
responses_so_far: List["ModelResponseStream"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Optional[Any] = None,
|
||||
user_api_key_dict: Optional[Any] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
litellm_logging_obj: Any | None = None,
|
||||
user_api_key_dict: Any | None = None,
|
||||
request_data: dict | None = None,
|
||||
stream_transform_sink: StreamTransformSink | None = None,
|
||||
) -> List["ModelResponseStream"]:
|
||||
"""
|
||||
Process output streaming responses by applying guardrails to text content.
|
||||
|
|
@ -414,14 +421,50 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
guardrail_to_apply: The guardrail instance to apply
|
||||
litellm_logging_obj: Optional logging object
|
||||
user_api_key_dict: User API key metadata to pass to guardrails
|
||||
stream_transform_sink: Optional out-parameter for the streaming text
|
||||
transformation path. When provided, the guardrail runs over the raw
|
||||
accumulated text (``responses_so_far`` is left untouched so it stays
|
||||
a correct raw accumulator across rounds) and the guardrailed text
|
||||
plus requested holdback are reported per choice on the sink.
|
||||
|
||||
Returns:
|
||||
Modified list of responses with guardrail applied to content
|
||||
The (unmodified) list of responses.
|
||||
|
||||
Response Format Support:
|
||||
- String content: choice.message.content = "text here"
|
||||
- List content: choice.message.content = [{"type": "text", "text": "text here"}, ...]
|
||||
"""
|
||||
if stream_transform_sink is not None:
|
||||
await self._process_streaming_transform(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
sink=stream_transform_sink,
|
||||
)
|
||||
return responses_so_far
|
||||
|
||||
return await self._process_streaming_block_only(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
async def _process_streaming_block_only(
|
||||
self,
|
||||
*,
|
||||
responses_so_far: list["ModelResponseStream"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Any | None,
|
||||
user_api_key_dict: Any | None,
|
||||
request_data: dict | None,
|
||||
) -> list["ModelResponseStream"]:
|
||||
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
|
||||
terminate the stream. Text rewrites are not propagated to the client here
|
||||
(see ``_process_streaming_transform`` for the incremental_diff path)."""
|
||||
# check if the stream has ended
|
||||
has_stream_ended = False
|
||||
for chunk in responses_so_far:
|
||||
|
|
@ -467,7 +510,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
# Step 2: Create lists for guardrail processing
|
||||
texts_to_check: List[str] = []
|
||||
images_to_check: List[str] = []
|
||||
task_mappings: List[Tuple[int, Optional[int]]] = []
|
||||
task_mappings: List[Tuple[int, int | None]] = []
|
||||
# Track (choice_index, content_index) for each combined text
|
||||
|
||||
for (map_choice_idx, map_content_idx), combined_text in combined_texts.items():
|
||||
|
|
@ -520,9 +563,109 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
return responses_so_far
|
||||
|
||||
@staticmethod
|
||||
def _accumulate_string_content_by_choice_index(
|
||||
responses_so_far: list["ModelResponseStream"],
|
||||
) -> dict[int, str]:
|
||||
"""Accumulate raw string ``delta.content`` per choice, keyed by
|
||||
``StreamingChoices.index`` (not enumerate position, which collapses to 0
|
||||
when each chunk carries a single non-zero-indexed choice for ``n > 1``).
|
||||
|
||||
Only string content participates; list-of-blocks content is out of scope
|
||||
for the incremental transform path. Reads ``responses_so_far`` without
|
||||
mutating it so it stays a correct raw accumulator across rounds.
|
||||
"""
|
||||
accumulated: dict[int, str] = {}
|
||||
for response in responses_so_far:
|
||||
for choice in response.choices:
|
||||
if isinstance(choice, litellm.StreamingChoices):
|
||||
content = choice.delta.content
|
||||
elif isinstance(choice, litellm.Choices):
|
||||
content = choice.message.content
|
||||
else:
|
||||
continue
|
||||
if isinstance(content, str) and content:
|
||||
idx = getattr(choice, "index", 0) or 0
|
||||
accumulated[idx] = accumulated.get(idx, "") + content
|
||||
return accumulated
|
||||
|
||||
async def _process_streaming_transform(
|
||||
self,
|
||||
*,
|
||||
responses_so_far: list["ModelResponseStream"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
litellm_logging_obj: Any | None,
|
||||
user_api_key_dict: Any | None,
|
||||
request_data: dict | None,
|
||||
sink: StreamTransformSink,
|
||||
) -> None:
|
||||
"""Run the guardrail over the raw accumulated text and report the
|
||||
guardrailed text plus requested holdback per choice on ``sink``.
|
||||
|
||||
Unlike the block-only path this never mutates ``responses_so_far``: it
|
||||
re-derives the raw accumulated text every round (so a rewrite guardrail
|
||||
always sees consistent input) and hands the result back out of band.
|
||||
"""
|
||||
raw_by_index = self._accumulate_string_content_by_choice_index(responses_so_far)
|
||||
if not raw_by_index:
|
||||
sink.mutated_text_per_choice = {}
|
||||
sink.holdback_per_choice = {}
|
||||
return
|
||||
|
||||
# Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits
|
||||
# choice 1 before choice 0 still hands the guardrail texts in a
|
||||
# deterministic index order. Without this, the guardrail's returned
|
||||
# texts (aligned to the input order it received) would map back to the
|
||||
# wrong choice indices when we rebuild the sink dicts by
|
||||
# ``enumerate(indices)``.
|
||||
indices = sorted(raw_by_index.keys())
|
||||
texts_to_check = [raw_by_index[i] for i in indices]
|
||||
|
||||
if request_data is None:
|
||||
request_data = {"responses": responses_so_far}
|
||||
elif "responses" not in request_data:
|
||||
request_data["responses"] = responses_so_far
|
||||
if "litellm_metadata" not in request_data:
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(texts=texts_to_check)
|
||||
if responses_so_far and getattr(responses_so_far[0], "model", None):
|
||||
inputs["model"] = responses_so_far[0].model
|
||||
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
returned_texts = guardrailed_inputs.get("texts")
|
||||
# No "texts" key means the guardrail made no change (action NONE): the raw
|
||||
# accumulated text is the guardrailed text. A present-but-shorter list is a
|
||||
# guardrail contract violation; those choices are omitted below (withheld,
|
||||
# not emitted raw) so a malformed response fails closed instead of leaking.
|
||||
if returned_texts is None:
|
||||
returned_texts = texts_to_check
|
||||
elif len(returned_texts) < len(texts_to_check):
|
||||
verbose_proxy_logger.warning(
|
||||
"OpenAI Chat Completions: guardrail returned %s transformed texts for %s inputs on the "
|
||||
"streaming transform path; withholding the unmatched choices to fail closed.",
|
||||
len(returned_texts),
|
||||
len(texts_to_check),
|
||||
)
|
||||
|
||||
holdback = guardrailed_inputs.get("stream_holdback_chars") or []
|
||||
sink.mutated_text_per_choice = {
|
||||
idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)
|
||||
}
|
||||
sink.holdback_per_choice = {
|
||||
indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
|
||||
}
|
||||
|
||||
def _combine_streaming_texts(
|
||||
self, responses_so_far: List["ModelResponseStream"]
|
||||
) -> Dict[Tuple[int, Optional[int]], str]:
|
||||
) -> Dict[Tuple[int, int | None], str]:
|
||||
"""
|
||||
Combine all streaming chunks into complete text per choice.
|
||||
|
||||
|
|
@ -534,7 +677,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
Returns:
|
||||
Dict mapping (choice_idx, content_idx) to combined text string
|
||||
"""
|
||||
combined_texts: Dict[Tuple[int, Optional[int]], str] = {}
|
||||
combined_texts: Dict[Tuple[int, int | None], str] = {}
|
||||
|
||||
for response_idx, response in enumerate(responses_so_far):
|
||||
for choice_idx, choice in enumerate(response.choices):
|
||||
|
|
@ -550,7 +693,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
if isinstance(content, str):
|
||||
# String content - accumulate for this choice
|
||||
str_key: Tuple[int, Optional[int]] = (choice_idx, None)
|
||||
str_key: Tuple[int, int | None] = (choice_idx, None)
|
||||
if str_key not in combined_texts:
|
||||
combined_texts[str_key] = ""
|
||||
combined_texts[str_key] += content
|
||||
|
|
@ -560,7 +703,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
for content_idx, content_item in enumerate(content):
|
||||
text_str = content_item.get("text")
|
||||
if text_str:
|
||||
list_key: Tuple[int, Optional[int]] = (
|
||||
list_key: Tuple[int, int | None] = (
|
||||
choice_idx,
|
||||
content_idx,
|
||||
)
|
||||
|
|
@ -607,7 +750,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
texts_to_check: List[str],
|
||||
images_to_check: List[str],
|
||||
tool_calls_to_check: List[Dict[str, Any]],
|
||||
text_task_mappings: List[Tuple[int, Optional[int]]],
|
||||
text_task_mappings: List[Tuple[int, int | None]],
|
||||
tool_call_task_mappings: List[Tuple[int, int]],
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -619,7 +762,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
# Determine content source and tool calls based on choice type
|
||||
content = None
|
||||
tool_calls: Optional[List[Any]] = None
|
||||
tool_calls: List[Any] | None = None
|
||||
if isinstance(choice, litellm.Choices):
|
||||
content = choice.message.content
|
||||
tool_calls = choice.message.tool_calls
|
||||
|
|
@ -662,7 +805,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
tool_calls_to_check.append(tool_call_dict)
|
||||
tool_call_task_mappings.append((choice_idx, int(tool_call_idx)))
|
||||
|
||||
def _convert_tool_call_to_dict(self, tool_call: Union[Dict[str, Any], Any]) -> Optional[Dict[str, Any]]:
|
||||
def _convert_tool_call_to_dict(self, tool_call: Union[Dict[str, Any], Any]) -> Dict[str, Any] | None:
|
||||
"""
|
||||
Convert a tool call object to dictionary format.
|
||||
|
||||
|
|
@ -691,7 +834,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
self,
|
||||
response: "ModelResponse",
|
||||
responses: List[str],
|
||||
task_mappings: List[Tuple[int, Optional[int]]],
|
||||
task_mappings: List[Tuple[int, int | None]],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail text responses back to output response.
|
||||
|
|
@ -701,7 +844,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
for task_idx, guardrail_response in enumerate(responses):
|
||||
mapping = task_mappings[task_idx]
|
||||
choice_idx = cast(int, mapping[0])
|
||||
content_idx_optional = cast(Optional[int], mapping[1])
|
||||
content_idx_optional = cast(int | None, mapping[1])
|
||||
|
||||
choice = cast(Choices, response.choices[choice_idx])
|
||||
|
||||
|
|
@ -755,7 +898,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
self,
|
||||
responses: List["ModelResponseStream"],
|
||||
guardrailed_texts: List[str],
|
||||
task_mappings: List[Tuple[int, Optional[int]]],
|
||||
task_mappings: List[Tuple[int, int | None]],
|
||||
) -> None:
|
||||
"""
|
||||
Apply guardrail responses back to output streaming responses.
|
||||
|
|
@ -771,16 +914,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
Override this method to customize how responses are applied to streaming responses.
|
||||
"""
|
||||
# Build a mapping of what guardrailed text to use for each (choice_idx, content_idx)
|
||||
guardrail_map: Dict[Tuple[int, Optional[int]], str] = {}
|
||||
guardrail_map: Dict[Tuple[int, int | None], str] = {}
|
||||
for task_idx, guardrail_response in enumerate(guardrailed_texts):
|
||||
mapping = task_mappings[task_idx]
|
||||
choice_idx = cast(int, mapping[0])
|
||||
content_idx_optional = cast(Optional[int], mapping[1])
|
||||
content_idx_optional = cast(int | None, mapping[1])
|
||||
guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response
|
||||
|
||||
# Track which choices we've already set the guardrailed text for
|
||||
# Key: (choice_idx, content_idx), Value: boolean (True if already set)
|
||||
already_set: Dict[Tuple[int, Optional[int]], bool] = {}
|
||||
already_set: Dict[Tuple[int, int | None], bool] = {}
|
||||
|
||||
# Iterate through all responses and update content
|
||||
for response_idx, response in enumerate(responses):
|
||||
|
|
@ -797,7 +940,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
|
||||
if isinstance(content, str):
|
||||
# String content
|
||||
str_key: Tuple[int, Optional[int]] = (choice_idx_in_response, None)
|
||||
str_key: Tuple[int, int | None] = (choice_idx_in_response, None)
|
||||
if str_key in guardrail_map:
|
||||
if str_key not in already_set:
|
||||
# First chunk - set the complete guardrailed text
|
||||
|
|
@ -817,7 +960,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
# List content - handle each content item
|
||||
for content_idx, content_item in enumerate(content):
|
||||
if "text" in content_item:
|
||||
list_key: Tuple[int, Optional[int]] = (
|
||||
list_key: Tuple[int, int | None] = (
|
||||
choice_idx_in_response,
|
||||
content_idx,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -998,6 +998,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
response_modalities.append("IMAGE")
|
||||
elif modality == "audio":
|
||||
response_modalities.append("AUDIO")
|
||||
elif modality == "video":
|
||||
response_modalities.append("VIDEO")
|
||||
else:
|
||||
response_modalities.append("MODALITY_UNSPECIFIED")
|
||||
return response_modalities
|
||||
|
|
|
|||
|
|
@ -19612,6 +19612,39 @@
|
|||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini/gemini-omni-flash-preview": {
|
||||
"input_cost_per_audio_token": 1.5e-06,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "gemini",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_tokens": 65535,
|
||||
"mode": "chat",
|
||||
"output_cost_per_reasoning_token": 9e-06,
|
||||
"output_cost_per_token": 9e-06,
|
||||
"output_cost_per_video_token": 1.75e-05,
|
||||
"rpm": 2000,
|
||||
"source": "https://ai.google.dev/gemini-api/docs/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"video"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true,
|
||||
"tpm": 800000
|
||||
},
|
||||
"gemini/gemini-3.1-pro-preview": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
|
||||
|
|
@ -19776,6 +19809,37 @@
|
|||
},
|
||||
"web_search_billing_unit": "per_query"
|
||||
},
|
||||
"gemini-omni-flash-preview": {
|
||||
"input_cost_per_audio_token": 1.5e-06,
|
||||
"input_cost_per_token": 1.5e-06,
|
||||
"litellm_provider": "vertex_ai-language-models",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 65535,
|
||||
"max_tokens": 65535,
|
||||
"mode": "chat",
|
||||
"output_cost_per_reasoning_token": 9e-06,
|
||||
"output_cost_per_token": 9e-06,
|
||||
"output_cost_per_video_token": 1.75e-05,
|
||||
"source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview",
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image",
|
||||
"audio",
|
||||
"video"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text",
|
||||
"video"
|
||||
],
|
||||
"supports_audio_input": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"gemini-3.5-flash": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"input_cost_per_audio_token": 1e-06,
|
||||
|
|
@ -44197,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,
|
||||
|
|
@ -44309,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,
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
|
|||
command: Optional[str] = None
|
||||
args: List[str] = Field(default_factory=list)
|
||||
env: Dict[str, str] = Field(default_factory=dict)
|
||||
issuer: Optional[str] = None
|
||||
authorization_url: Optional[str] = None
|
||||
token_url: Optional[str] = None
|
||||
registration_url: Optional[str] = None
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ if TYPE_CHECKING:
|
|||
|
||||
_AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset(
|
||||
{
|
||||
"issuer",
|
||||
"authorization_url",
|
||||
"token_url",
|
||||
"registration_url",
|
||||
|
|
@ -60,6 +61,13 @@ _AUTH_FLOW_SCOPED_FIELDS: frozenset = frozenset(
|
|||
}
|
||||
)
|
||||
|
||||
|
||||
def _blank_to_none(value: Optional[str]) -> Optional[str]:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
return value.strip() or None
|
||||
|
||||
|
||||
# Token-exchange settings with dedicated columns that also exist on
|
||||
# ``MCPCredentials`` as a legacy shape (rows and REST callers that predate the
|
||||
# columns). Every write lifts blob values into the columns and strips them from
|
||||
|
|
@ -697,13 +705,15 @@ async def update_mcp_server(
|
|||
# of being reset to a schema default (transport=sse, allow_all_keys=False...).
|
||||
data_dict = _prepare_mcp_server_data(data, exclude_unset=True, fields_set=fields_set)
|
||||
|
||||
# Pre-fetch existing record once if we need it for auth_type or credential logic
|
||||
# Pre-fetch existing record once if we need it for auth_type, url, or credential logic
|
||||
existing = None
|
||||
has_credentials = "credentials" in data_dict and data_dict["credentials"] is not None
|
||||
# An explicit token-exchange column write (set or clear) also migrates the
|
||||
# legacy blob copies below, so the existing row is needed for those updates.
|
||||
explicit_te_write = bool(_TOKEN_EXCHANGE_COLUMN_FIELDS & data_dict.keys())
|
||||
if data.auth_type or has_credentials or explicit_te_write:
|
||||
url_provided = "url" in data_dict and data_dict["url"] is not None
|
||||
issuer_provided = "issuer" in data_dict
|
||||
if data.auth_type or has_credentials or explicit_te_write or url_provided or issuer_provided:
|
||||
existing = await MCPServerRepository(prisma_client).table.find_unique(where={"server_id": data.server_id})
|
||||
|
||||
auth_type_changed = bool(
|
||||
|
|
@ -711,13 +721,30 @@ async def update_mcp_server(
|
|||
and existing
|
||||
and _credential_auth_class(existing.auth_type) != _credential_auth_class(data.auth_type)
|
||||
)
|
||||
# A url change re-points the server at a potentially different upstream, so any discovered or
|
||||
# trust-on-first-use OAuth endpoints/issuer belong to the old upstream and must re-discover.
|
||||
url_changed = bool(url_provided and existing and existing.url != data_dict["url"])
|
||||
old_issuer = _blank_to_none(getattr(existing, "issuer", None)) if existing else None
|
||||
issuer_changed = bool(
|
||||
issuer_provided and old_issuer is not None and _blank_to_none(data_dict.get("issuer")) != old_issuer
|
||||
)
|
||||
|
||||
# Clear stale credentials when auth_type changes but no new credentials provided
|
||||
if auth_type_changed and "credentials" not in data_dict:
|
||||
data_dict["credentials"] = None
|
||||
|
||||
if auth_type_changed:
|
||||
data_dict.update({field: None for field in _AUTH_FLOW_SCOPED_FIELDS if field not in data_dict})
|
||||
if auth_type_changed or url_changed or issuer_changed:
|
||||
# Clear each auth-flow-scoped field that the caller either omitted (partial update) or
|
||||
# resubmitted unchanged. The edit form re-sends every field, so a stale issuer/endpoint
|
||||
# belonging to the old upstream would otherwise survive a url/auth_type change and win in the
|
||||
# resolution merge; only a genuinely new submitted value is kept.
|
||||
data_dict.update(
|
||||
{
|
||||
field: None
|
||||
for field in _AUTH_FLOW_SCOPED_FIELDS
|
||||
if field not in data_dict or data_dict[field] == getattr(existing, field, None)
|
||||
}
|
||||
)
|
||||
|
||||
# An explicit column write that does not touch credentials must still migrate
|
||||
# the row's legacy blob copies: lift values for columns the caller left
|
||||
|
|
@ -1181,6 +1208,7 @@ def mcp_oauth_token_identity(server: object) -> tuple[object, ...]:
|
|||
getattr(server, "spec_path", None),
|
||||
getattr(server, "auth_type", None),
|
||||
getattr(server, "oauth2_flow", None),
|
||||
getattr(server, "issuer", None),
|
||||
getattr(server, "authorization_url", None),
|
||||
getattr(server, "token_url", None),
|
||||
getattr(server, "registration_url", None),
|
||||
|
|
|
|||
|
|
@ -186,6 +186,203 @@ _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 _uses_issuer_anchor(manual_issuer: str | None, is_discovery_auth_type: bool) -> bool:
|
||||
"""Whether the endpoints are authoritatively anchored to an admin-pinned issuer (RFC 8414 §3.3).
|
||||
|
||||
This is the trust/provenance property, distinct from whether the ``issuer`` field is merely
|
||||
populated: a trust-on-first-use discovered issuer sets ``issuer`` for token identity but is NOT
|
||||
anchored, so its endpoints stay resource-rooted. Anchoring holds only when the issuer was pinned
|
||||
(present on the row/config) on a discovery auth type. Every consumer of "is this anchored" reads
|
||||
this one definition, so the answer cannot diverge across build paths.
|
||||
"""
|
||||
return _blank_to_none(manual_issuer) is not None and is_discovery_auth_type
|
||||
|
||||
|
||||
def _endpoints_yield_to_issuer(
|
||||
issuer: str | None,
|
||||
is_discovery_auth_type: bool,
|
||||
authorization_url: str | None,
|
||||
token_url: str | None,
|
||||
registration_url: str | None,
|
||||
) -> tuple[str | None, str | None, str | None]:
|
||||
"""The single rule that makes an admin-configured ``issuer`` the sole authoritative endpoint
|
||||
source (RFC 8414 §3.3): when it is set for a discovery auth type, the stored/manual
|
||||
``authorization_url``/``token_url``/``registration_url`` do not apply. They neither anchor nor
|
||||
short-circuit discovery, never override the issuer document in the merge, and never substitute for
|
||||
it when the issuer fetch fails (fail-closed). Returns the endpoint values that remain in force,
|
||||
i.e. all ``None`` when issuer-anchored, else the inputs unchanged. Called at every resolution site
|
||||
so the invariant holds in one place instead of being re-derived per merge.
|
||||
"""
|
||||
if issuer is not None and is_discovery_auth_type:
|
||||
return None, None, None
|
||||
return authorization_url, token_url, registration_url
|
||||
|
||||
|
||||
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 _issuer_matches(claimed_issuer: object, configured_issuer: str) -> bool:
|
||||
"""RFC 8414 §3.3 issuer equality between the metadata document's self-attested ``issuer`` and the
|
||||
admin-configured issuer, tolerant only of URL-insignificant differences (scheme/host case, the
|
||||
default port, a trailing slash). A non-string or empty claimed issuer never matches, so a
|
||||
document that omits ``issuer`` fails closed under issuer-anchored discovery.
|
||||
"""
|
||||
if not isinstance(claimed_issuer, str) or not claimed_issuer:
|
||||
return False
|
||||
return _normalized_authorize_endpoint(claimed_issuer) == _normalized_authorize_endpoint(configured_issuer)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
A rebuild wholesale-replaces the registry entry, so without this a transient upstream outage
|
||||
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 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.
|
||||
|
||||
When the server is issuer-anchored (``issuer_is_anchored`` -- a pinned issuer on a discovery auth
|
||||
type), the endpoints come solely from the §3.3-validated issuer document, so carry-forward is
|
||||
skipped entirely for its endpoints: a failed issuer fetch leaves them ``None`` and must stay
|
||||
``None`` (fail-closed), never resurrected from the previous registry entry. A merely discovered
|
||||
(trust-on-first-use) issuer is NOT anchored -- ``issuer`` is set for token identity but the
|
||||
endpoints are resource-rooted, so they still carry forward as last-known-good, gated by the
|
||||
corroboration check below like any other resource-rooted server. Scopes stay resource-driven and
|
||||
can carry either way.
|
||||
"""
|
||||
if previous_server is None:
|
||||
return
|
||||
if previous_server.url != new_server.url or previous_server.auth_type != new_server.auth_type:
|
||||
return
|
||||
if new_server.issuer_is_anchored:
|
||||
# Endpoints come solely from the §3.3-validated issuer document; a failed fetch stays
|
||||
# fail-closed and must not be resurrected from the previous entry. Only the resource-driven
|
||||
# scopes carry as last-known-good.
|
||||
if not new_server.scopes and previous_server.scopes:
|
||||
new_server.scopes = previous_server.scopes
|
||||
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 may_carry and new_server.token_url is None and previous_server.token_url:
|
||||
new_server.token_url = previous_server.token_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 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."""
|
||||
|
|
@ -999,36 +1196,68 @@ class MCPServerManager:
|
|||
)
|
||||
|
||||
auth_type = server_config.get("auth_type", None)
|
||||
if server_url and (
|
||||
auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
|
||||
manual_issuer = _blank_to_none(server_config.get("issuer"))
|
||||
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"))
|
||||
is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
|
||||
use_issuer_anchor = _uses_issuer_anchor(manual_issuer, is_discovery_auth_type)
|
||||
manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer(
|
||||
manual_issuer,
|
||||
is_discovery_auth_type,
|
||||
manual_authorization_url,
|
||||
manual_token_url,
|
||||
manual_registration_url,
|
||||
)
|
||||
should_discover = bool(server_url) and (
|
||||
is_discovery_auth_type
|
||||
or self._obo_needs_endpoint_discovery(
|
||||
auth_type,
|
||||
server_config.get("token_exchange_endpoint"),
|
||||
server_config.get("token_url"),
|
||||
manual_token_url,
|
||||
)
|
||||
):
|
||||
)
|
||||
if not should_discover:
|
||||
mcp_oauth_metadata = None
|
||||
elif manual_issuer is not None and is_discovery_auth_type:
|
||||
mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url)
|
||||
else:
|
||||
mcp_oauth_metadata = await self._descovery_metadata(
|
||||
server_url=server_url,
|
||||
allow_origin_fallback=auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
|
||||
allow_origin_fallback=is_discovery_auth_type,
|
||||
)
|
||||
|
||||
if use_issuer_anchor:
|
||||
gated_oauth_metadata = mcp_oauth_metadata
|
||||
elif is_discovery_auth_type:
|
||||
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")),
|
||||
)
|
||||
else:
|
||||
mcp_oauth_metadata = None
|
||||
gated_oauth_metadata = 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_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
|
||||
)
|
||||
resolved_registration_url = server_config.get("registration_url") or (
|
||||
mcp_oauth_metadata.registration_url if mcp_oauth_metadata else None
|
||||
discovered_issuer = (
|
||||
gated_oauth_metadata.discovered_issuer
|
||||
if gated_oauth_metadata and not gated_oauth_metadata.from_origin_fallback
|
||||
else None
|
||||
)
|
||||
effective_issuer = manual_issuer or discovered_issuer
|
||||
|
||||
config_oauth2_flow = server_config.get("oauth2_flow", None)
|
||||
if auth_type == MCPAuth.oauth2 and config_oauth2_flow not in (
|
||||
|
|
@ -1077,6 +1306,8 @@ class MCPServerManager:
|
|||
client_secret=server_config.get("client_secret", None),
|
||||
oauth2_flow=self._explicit_oauth2_flow(config_oauth2_flow),
|
||||
scopes=resolved_scopes,
|
||||
issuer=effective_issuer,
|
||||
issuer_is_anchored=use_issuer_anchor,
|
||||
authorization_url=resolved_authorization_url,
|
||||
token_url=resolved_token_url,
|
||||
registration_url=resolved_registration_url,
|
||||
|
|
@ -1337,12 +1568,59 @@ class MCPServerManager:
|
|||
decrypt_global_env_var_values(env_vars_list)
|
||||
return env_vars_list
|
||||
|
||||
async def _resolve_table_oauth_metadata(
|
||||
self,
|
||||
*,
|
||||
mcp_server: LiteLLM_MCPServerTable,
|
||||
auth_type: MCPAuthType,
|
||||
server_url: Optional[str],
|
||||
manual_issuer: Optional[str],
|
||||
manual_authorization_url: Optional[str],
|
||||
manual_token_url: Optional[str],
|
||||
is_discovery_auth_type: bool,
|
||||
use_issuer_anchor: bool,
|
||||
scopes: Optional[list[str]],
|
||||
token_exchange_endpoint: Optional[str],
|
||||
) -> Optional[MCPOAuthMetadata]:
|
||||
has_all_upstream_oauth_fields = bool(manual_authorization_url and manual_token_url and scopes)
|
||||
needs_discovery = bool(server_url) and (
|
||||
(is_discovery_auth_type and not has_all_upstream_oauth_fields)
|
||||
or self._obo_needs_endpoint_discovery(auth_type, token_exchange_endpoint, manual_token_url)
|
||||
)
|
||||
if not needs_discovery:
|
||||
mcp_oauth_metadata: Optional[MCPOAuthMetadata] = None
|
||||
elif use_issuer_anchor and manual_issuer is not None:
|
||||
mcp_oauth_metadata = await self._fetch_issuer_anchored_oauth_metadata(manual_issuer, server_url)
|
||||
else:
|
||||
mcp_oauth_metadata = await self._descovery_metadata(
|
||||
server_url=server_url, # type: ignore[arg-type]
|
||||
allow_origin_fallback=is_discovery_auth_type,
|
||||
)
|
||||
if needs_discovery and not use_issuer_anchor and mcp_oauth_metadata is None:
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery yielded no metadata for server %s (%s); "
|
||||
"OAuth endpoints/scopes stay unresolved until a rebuild succeeds",
|
||||
mcp_server.server_id,
|
||||
server_url,
|
||||
)
|
||||
if use_issuer_anchor:
|
||||
return mcp_oauth_metadata
|
||||
if is_discovery_auth_type:
|
||||
return _restrict_discovery_to_corroborated_authorization_server(
|
||||
mcp_oauth_metadata,
|
||||
manual_authorization_url,
|
||||
mcp_server.server_id,
|
||||
bool(getattr(mcp_server, "dcr_bridge", None)),
|
||||
)
|
||||
return mcp_oauth_metadata
|
||||
|
||||
async def build_mcp_server_from_table(
|
||||
self,
|
||||
mcp_server: LiteLLM_MCPServerTable,
|
||||
*,
|
||||
credentials_are_encrypted: bool = True,
|
||||
env_vars_are_encrypted: Optional[bool] = None,
|
||||
persist_discovered_endpoints: bool = True,
|
||||
) -> MCPServer:
|
||||
_mcp_info: MCPInfo = mcp_server.mcp_info or {}
|
||||
env_dict = _deserialize_json_dict(getattr(mcp_server, "env", None))
|
||||
|
|
@ -1419,25 +1697,38 @@ class MCPServerManager:
|
|||
|
||||
auth_type = cast(MCPAuthType, mcp_server.auth_type)
|
||||
server_url = mcp_server.url
|
||||
needs_discovery = bool(server_url) and (
|
||||
(auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and not mcp_server.authorization_url)
|
||||
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_issuer = _blank_to_none(mcp_server.issuer)
|
||||
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)
|
||||
is_discovery_auth_type = auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
|
||||
use_issuer_anchor = _uses_issuer_anchor(manual_issuer, is_discovery_auth_type)
|
||||
manual_authorization_url, manual_token_url, manual_registration_url = _endpoints_yield_to_issuer(
|
||||
manual_issuer, is_discovery_auth_type, manual_authorization_url, manual_token_url, manual_registration_url
|
||||
)
|
||||
mcp_oauth_metadata = (
|
||||
await self._descovery_metadata(
|
||||
server_url=server_url, # type: ignore[arg-type]
|
||||
allow_origin_fallback=auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
|
||||
)
|
||||
if needs_discovery
|
||||
else None
|
||||
token_exchange_endpoint = mcp_server.token_exchange_endpoint or (
|
||||
credentials_dict.get("token_exchange_endpoint") if credentials_dict else None
|
||||
)
|
||||
gated_oauth_metadata = await self._resolve_table_oauth_metadata(
|
||||
mcp_server=mcp_server,
|
||||
auth_type=auth_type,
|
||||
server_url=server_url,
|
||||
manual_issuer=manual_issuer,
|
||||
manual_authorization_url=manual_authorization_url,
|
||||
manual_token_url=manual_token_url,
|
||||
is_discovery_auth_type=is_discovery_auth_type,
|
||||
use_issuer_anchor=use_issuer_anchor,
|
||||
scopes=scopes,
|
||||
token_exchange_endpoint=token_exchange_endpoint,
|
||||
)
|
||||
|
||||
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)
|
||||
discovered_issuer = (
|
||||
gated_oauth_metadata.discovered_issuer
|
||||
if gated_oauth_metadata and not gated_oauth_metadata.from_origin_fallback
|
||||
else None
|
||||
)
|
||||
effective_issuer = manual_issuer or discovered_issuer
|
||||
|
||||
new_server = MCPServer(
|
||||
server_id=mcp_server.server_id,
|
||||
|
|
@ -1457,9 +1748,11 @@ 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),
|
||||
issuer=effective_issuer,
|
||||
issuer_is_anchored=use_issuer_anchor,
|
||||
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
|
||||
),
|
||||
|
|
@ -1506,12 +1799,23 @@ class MCPServerManager:
|
|||
max_concurrent_requests=getattr(mcp_server, "max_concurrent_requests", None),
|
||||
)
|
||||
_warn_internal_delegate_pkce_if_applicable(new_server, source="database")
|
||||
await self._persist_discovered_obo_token_url(
|
||||
server_id=mcp_server.server_id,
|
||||
auth_type=auth_type,
|
||||
existing_token_url=mcp_server.token_url,
|
||||
discovered_token_url=new_server.token_url,
|
||||
)
|
||||
if persist_discovered_endpoints:
|
||||
await self._persist_discovered_obo_token_url(
|
||||
server_id=mcp_server.server_id,
|
||||
auth_type=auth_type,
|
||||
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_issuer=manual_issuer,
|
||||
existing_authorization_url=manual_authorization_url,
|
||||
existing_token_url=manual_token_url,
|
||||
existing_scopes=scopes,
|
||||
metadata=gated_oauth_metadata,
|
||||
is_issuer_anchored=use_issuer_anchor,
|
||||
)
|
||||
return new_server
|
||||
|
||||
async def _persist_discovered_obo_token_url(
|
||||
|
|
@ -1549,6 +1853,89 @@ class MCPServerManager:
|
|||
except Exception as exc: # noqa: BLE001 - best-effort; a failed write re-discovers next build
|
||||
verbose_logger.warning("Failed to persist discovered OBO token_url for MCP server %s: %s", server_id, exc)
|
||||
|
||||
async def _persist_discovered_oauth_endpoints(
|
||||
self,
|
||||
*,
|
||||
server_id: str,
|
||||
auth_type: MCPAuthType | None,
|
||||
existing_issuer: str | None,
|
||||
existing_authorization_url: str | None,
|
||||
existing_token_url: str | None,
|
||||
existing_scopes: list[str] | None,
|
||||
metadata: MCPOAuthMetadata | None,
|
||||
is_issuer_anchored: bool = False,
|
||||
) -> None:
|
||||
"""Write freshly discovered OAuth endpoints back onto the DB row.
|
||||
|
||||
Same rationale as ``_persist_discovered_obo_token_url`` but for the interactive oauth2
|
||||
family: discovered ``authorization_url``/``token_url``/``scopes`` otherwise live only on
|
||||
the in-memory registry entry, which is rebuilt on every client connect (the DCR reuse path
|
||||
calls ``update_server``) and on every post-write DB reload, so one failed re-discovery
|
||||
serves 400 "authorization url is not set" from /authorize until a later rebuild succeeds.
|
||||
Only fills row fields that are currently empty, never persists origin-fallback guesses
|
||||
(RFC 9728/8414-advertised metadata only), and deliberately skips ``registration_url``
|
||||
because ``_dcr_bridge_relays_client_registration`` keys off that column. Best-effort: a
|
||||
failed write re-discovers on the next build. Scopes go through ``update_mcp_server`` so
|
||||
they merge into the credentials blob without touching the stored client credentials.
|
||||
|
||||
For an issuer-anchored server (``is_issuer_anchored``) the endpoints are re-derived from the
|
||||
§3.3-validated issuer document on every build, so they are NOT persisted into the endpoint
|
||||
columns: persisting them would make the next build see populated endpoints and treat them as
|
||||
authoritative stored values, defeating the "endpoints come solely from the issuer" invariant.
|
||||
Only the resource-driven scopes are persisted for such servers.
|
||||
"""
|
||||
if auth_type not in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES:
|
||||
return
|
||||
if metadata is None or metadata.from_origin_fallback:
|
||||
return
|
||||
issuer_update = (
|
||||
{"issuer": metadata.discovered_issuer} if metadata.discovered_issuer and not existing_issuer else {}
|
||||
)
|
||||
authorization_url_update = (
|
||||
{"authorization_url": metadata.authorization_url}
|
||||
if metadata.authorization_url and not existing_authorization_url and not is_issuer_anchored
|
||||
else {}
|
||||
)
|
||||
token_url_update = (
|
||||
{"token_url": metadata.token_url}
|
||||
if metadata.token_url and not existing_token_url and not is_issuer_anchored
|
||||
else {}
|
||||
)
|
||||
scopes_update = {"credentials": {"scopes": metadata.scopes}} if metadata.scopes and not existing_scopes else {}
|
||||
updates: dict[str, object] = {
|
||||
**issuer_update,
|
||||
**authorization_url_update,
|
||||
**token_url_update,
|
||||
**scopes_update,
|
||||
}
|
||||
if not updates:
|
||||
return
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # db.py imports this module at load
|
||||
update_mcp_server,
|
||||
)
|
||||
from litellm.proxy._types import UpdateMCPServerRequest # noqa: PLC0415 # heavy module; import at call time
|
||||
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime value, set after startup
|
||||
|
||||
if prisma_client is None:
|
||||
return
|
||||
try:
|
||||
await update_mcp_server(
|
||||
prisma_client=prisma_client,
|
||||
data=UpdateMCPServerRequest.model_validate({"server_id": server_id, **updates}),
|
||||
touched_by="mcp_oauth_discovery",
|
||||
)
|
||||
verbose_logger.info(
|
||||
"Persisted discovered OAuth endpoints for MCP server %s: %s",
|
||||
server_id,
|
||||
sorted(updates),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - best-effort; a failed write re-discovers next build
|
||||
verbose_logger.warning(
|
||||
"Failed to persist discovered OAuth endpoints for MCP server %s: %s",
|
||||
server_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
async def _maybe_register_openapi_tools(self, server: MCPServer, *, initialize_mapping: bool = True):
|
||||
"""Register OpenAPI tools if the server has a spec_path configured."""
|
||||
if server.spec_path:
|
||||
|
|
@ -1607,6 +1994,10 @@ class MCPServerManager:
|
|||
existing_prefix = self.registry[mcp_server.server_id].short_prefix
|
||||
if existing_prefix and not new_server.short_prefix:
|
||||
new_server.short_prefix = existing_prefix
|
||||
_carry_forward_resolved_oauth_endpoints(
|
||||
new_server=new_server,
|
||||
previous_server=self.registry[mcp_server.server_id],
|
||||
)
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.registry[mcp_server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
|
|
@ -2969,16 +3360,20 @@ class MCPServerManager:
|
|||
) = await self._attempt_well_known_discovery(server_url)
|
||||
|
||||
metadata = None
|
||||
used_origin_fallback = False
|
||||
if allow_origin_fallback and not authorization_servers:
|
||||
try:
|
||||
parsed_url = urlparse(server_url)
|
||||
if parsed_url.scheme and parsed_url.netloc:
|
||||
authorization_servers = [f"{parsed_url.scheme}://{parsed_url.netloc}"]
|
||||
used_origin_fallback = True
|
||||
except Exception:
|
||||
authorization_servers = []
|
||||
|
||||
if authorization_servers:
|
||||
metadata = await self._fetch_authorization_server_metadata(authorization_servers, server_url)
|
||||
if metadata is not None and used_origin_fallback:
|
||||
metadata.from_origin_fallback = True
|
||||
|
||||
preferred_scopes = scopes or resource_scopes
|
||||
if metadata is None and preferred_scopes:
|
||||
|
|
@ -3085,8 +3480,41 @@ class MCPServerManager:
|
|||
return metadata
|
||||
return None
|
||||
|
||||
async def _fetch_issuer_anchored_oauth_metadata(
|
||||
self, issuer: str, server_url: Optional[str]
|
||||
) -> Optional[MCPOAuthMetadata]:
|
||||
"""RFC 8414 issuer-anchored discovery for the OAuth endpoints, with resource-driven scopes.
|
||||
|
||||
Fetch authorization-server metadata from the admin-configured issuer's own origin and adopt
|
||||
its ``token_endpoint``/``registration_endpoint`` only when the document self-attests that same
|
||||
issuer (RFC 8414 §3.3). Because the trust anchor is the pinned issuer rather than anything the
|
||||
MCP resource advertises, the endpoints are authoritative for that issuer and cannot be
|
||||
substituted by a compromised resource. Fails closed (returns None) on a §3.3 mismatch or a
|
||||
fetch failure. The issuer is passed as its own ``server_url`` so the endpoint fetch is treated
|
||||
as same-authority and is not subject to the resource-scoped SSRF shortcut.
|
||||
|
||||
Scopes are NOT taken from the issuer document. Per the MCP authorization spec Scope Selection
|
||||
Strategy and RFC 9728, the scopes a client requests are resource-driven (the WWW-Authenticate
|
||||
challenge or the protected-resource ``scopes_supported``), so the resource's advertised scopes
|
||||
are fetched separately and used; the resource can influence only the requested scope, which
|
||||
the authorization server and user consent bound (RFC 6749 §3.3), never the token endpoint.
|
||||
"""
|
||||
metadata = await self._fetch_single_authorization_server_metadata(issuer, issuer, require_issuer=issuer)
|
||||
if metadata is None:
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth issuer-anchored discovery for issuer %s yielded no metadata whose issuer "
|
||||
"matched (RFC 8414 §3.3); OAuth endpoints stay unresolved until a rebuild succeeds",
|
||||
issuer,
|
||||
)
|
||||
return None
|
||||
resource_metadata = (
|
||||
await self._descovery_metadata(server_url, allow_origin_fallback=False) if server_url else None
|
||||
)
|
||||
resource_scopes = resource_metadata.scopes if resource_metadata else None
|
||||
return metadata.model_copy(update={"scopes": resource_scopes})
|
||||
|
||||
async def _fetch_single_authorization_server_metadata(
|
||||
self, issuer_url: str, server_url: str
|
||||
self, issuer_url: str, server_url: str, require_issuer: Optional[str] = None
|
||||
) -> Optional[MCPOAuthMetadata]:
|
||||
try:
|
||||
parsed = urlparse(issuer_url)
|
||||
|
|
@ -3130,20 +3558,33 @@ class MCPServerManager:
|
|||
)
|
||||
continue
|
||||
|
||||
scopes = self._extract_scopes(data.get("scopes_supported"))
|
||||
claimed_issuer = data.get("issuer")
|
||||
verbose_logger.debug(
|
||||
"Authorization server metadata from %s: issuer=%s grant_types_supported=%s "
|
||||
"token_endpoint_auth_methods_supported=%s",
|
||||
url,
|
||||
data.get("issuer"),
|
||||
claimed_issuer,
|
||||
data.get("grant_types_supported"),
|
||||
data.get("token_endpoint_auth_methods_supported"),
|
||||
)
|
||||
if require_issuer is not None and not _issuer_matches(claimed_issuer, require_issuer):
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth issuer-anchored discovery: metadata at %s self-attests issuer %r, which "
|
||||
"does not match the configured issuer %r (RFC 8414 §3.3); rejecting so a compromised "
|
||||
"resource cannot substitute an attacker authorization server",
|
||||
url,
|
||||
claimed_issuer,
|
||||
require_issuer,
|
||||
)
|
||||
continue
|
||||
|
||||
scopes = self._extract_scopes(data.get("scopes_supported"))
|
||||
metadata = MCPOAuthMetadata(
|
||||
scopes=scopes,
|
||||
authorization_url=data.get("authorization_endpoint"),
|
||||
token_url=data.get("token_endpoint"),
|
||||
registration_url=data.get("registration_endpoint"),
|
||||
discovered_issuer=claimed_issuer if isinstance(claimed_issuer, str) and claimed_issuer else None,
|
||||
)
|
||||
|
||||
if any(
|
||||
|
|
@ -4489,6 +4930,7 @@ class MCPServerManager:
|
|||
# (if any) so the prefix is stable across reloads.
|
||||
if existing_server is not None and existing_server.short_prefix:
|
||||
new_server.short_prefix = existing_server.short_prefix
|
||||
_carry_forward_resolved_oauth_endpoints(new_server=new_server, previous_server=existing_server)
|
||||
new_registry[server.server_id] = new_server
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
|
|
@ -4863,6 +5305,7 @@ class MCPServerManager:
|
|||
command=getattr(server, "command", None),
|
||||
args=getattr(server, "args", None) or [],
|
||||
env=getattr(server, "env", None) or {},
|
||||
issuer=server.issuer,
|
||||
authorization_url=server.authorization_url,
|
||||
token_url=server.token_url,
|
||||
registration_url=server.registration_url,
|
||||
|
|
@ -4972,6 +5415,7 @@ class MCPServerManager:
|
|||
command=getattr(server, "command", None),
|
||||
args=getattr(server, "args", None) or [],
|
||||
env=getattr(server, "env", None) or {},
|
||||
issuer=server.issuer,
|
||||
authorization_url=server.authorization_url,
|
||||
token_url=server.token_url,
|
||||
registration_url=server.registration_url,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1138,6 +1138,7 @@ if MCP_AVAILABLE:
|
|||
static_headers=request.static_headers,
|
||||
client_id=client_id,
|
||||
client_secret=client_secret,
|
||||
issuer=request.issuer,
|
||||
token_url=request.token_url,
|
||||
scopes=scopes,
|
||||
authorization_url=request.authorization_url,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Semantic MCP Tool Filtering using semantic-router
|
|||
Filters MCP tools semantically for /chat/completions and /responses endpoints.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -76,6 +77,7 @@ class SemanticMCPToolFilter:
|
|||
self.tool_router: Optional["SemanticRouter"] = None
|
||||
self.context_window_error: Optional[str] = None
|
||||
self._tool_map: Dict[str, Any] = {} # MCPTool objects or OpenAI function dicts
|
||||
self._index_sync_lock = asyncio.Lock()
|
||||
|
||||
async def build_router_from_mcp_registry(self) -> None:
|
||||
"""Build semantic router from all MCP tools in the registry (no auth checks)."""
|
||||
|
|
@ -182,6 +184,81 @@ class SemanticMCPToolFilter:
|
|||
return
|
||||
raise
|
||||
|
||||
def _has_tools_missing_from_index(self, tools: list[Any]) -> bool:
|
||||
"""Allocation-free check for any named tool not yet in the semantic index."""
|
||||
return any(name and name not in self._tool_map for name in (self._extract_tool_info(t)[0] for t in tools))
|
||||
|
||||
def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]:
|
||||
"""Map name -> tool for every named tool not yet in the semantic index."""
|
||||
return {
|
||||
name: tool
|
||||
for name, tool in ((self._extract_tool_info(t)[0], t) for t in tools)
|
||||
if name and name not in self._tool_map
|
||||
}
|
||||
|
||||
async def _ensure_tools_indexed(self, available_tools: list[Any]) -> None:
|
||||
"""
|
||||
Index request-time tools the startup build never saw.
|
||||
|
||||
The startup index lists every registered MCP server WITHOUT per-user
|
||||
credentials, so servers requiring per-user auth (interactive OAuth
|
||||
tokens, user-scoped env vars) contribute zero routes. Tools reaching
|
||||
the filter came through an authenticated expansion; without indexing
|
||||
them here they can never be selected, so requests either bypass
|
||||
filtering entirely (N->N) or lose every tool to unrelated matches.
|
||||
|
||||
Runs async-only (no synchronous embedding on the request path) and
|
||||
never writes shared error state: an embedding failure here raises and
|
||||
is scoped to the requesting call, so one request's oversized tool
|
||||
description cannot poison the filter for other users on the worker.
|
||||
"""
|
||||
from semantic_router.routers import SemanticRouter
|
||||
from semantic_router.routers.base import Route
|
||||
|
||||
from litellm.router_strategy.auto_router.litellm_encoder import (
|
||||
LiteLLMRouterEncoder,
|
||||
)
|
||||
|
||||
if not self._has_tools_missing_from_index(available_tools):
|
||||
return
|
||||
|
||||
async with self._index_sync_lock:
|
||||
missing = self._tools_missing_from_index(available_tools)
|
||||
if not missing:
|
||||
return
|
||||
|
||||
descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()}
|
||||
routes = [
|
||||
Route(
|
||||
name=name,
|
||||
description=description,
|
||||
utterances=[description],
|
||||
score_threshold=self.similarity_threshold,
|
||||
)
|
||||
for name, description in descriptions.items()
|
||||
]
|
||||
|
||||
if self.tool_router is None:
|
||||
router = SemanticRouter(
|
||||
routes=[],
|
||||
encoder=LiteLLMRouterEncoder(
|
||||
litellm_router_instance=self.router_instance,
|
||||
model_name=self.embedding_model,
|
||||
score_threshold=self.similarity_threshold,
|
||||
),
|
||||
auto_sync="local",
|
||||
top_k=self.top_k,
|
||||
)
|
||||
await router.aadd(routes)
|
||||
self.tool_router = router
|
||||
else:
|
||||
await self.tool_router.aadd(routes)
|
||||
|
||||
self._tool_map.update(missing)
|
||||
verbose_logger.info(
|
||||
f"Semantic tool filter indexed {len(routes)} request-time tools missing from the startup index"
|
||||
)
|
||||
|
||||
async def filter_tools(
|
||||
self,
|
||||
query: str,
|
||||
|
|
@ -216,22 +293,34 @@ class SemanticMCPToolFilter:
|
|||
if not query or not query.strip():
|
||||
return available_tools
|
||||
|
||||
# Router should be built on startup - if not, something went wrong
|
||||
if self.tool_router is None:
|
||||
verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?")
|
||||
return available_tools
|
||||
|
||||
# Run semantic filtering
|
||||
try:
|
||||
await self._ensure_tools_indexed(available_tools)
|
||||
|
||||
if self.tool_router is None:
|
||||
verbose_logger.warning("Semantic router could not be built from the request's tools")
|
||||
return available_tools
|
||||
|
||||
available_names = [name for name in (self._extract_tool_info(t)[0] for t in available_tools) if name]
|
||||
if not available_names:
|
||||
return available_tools
|
||||
|
||||
limit = top_k or self.top_k
|
||||
matches = self.tool_router(text=query, limit=limit)
|
||||
if self.tool_router.top_k < limit:
|
||||
self.tool_router.top_k = limit
|
||||
matches = self.tool_router(text=query, limit=limit, route_filter=available_names)
|
||||
matched_tool_names = self._extract_tool_names_from_matches(matches)
|
||||
|
||||
if not matched_tool_names:
|
||||
return available_tools
|
||||
|
||||
return self._get_tools_by_names(matched_tool_names, available_tools)
|
||||
filtered_tools = self._get_tools_by_names(matched_tool_names, available_tools)
|
||||
if not filtered_tools:
|
||||
return available_tools
|
||||
return filtered_tools
|
||||
|
||||
except SemanticToolFilterContextWindowError:
|
||||
raise
|
||||
except Exception as e:
|
||||
if _is_context_window_error(e):
|
||||
verbose_logger.error(
|
||||
|
|
@ -240,7 +329,7 @@ class SemanticMCPToolFilter:
|
|||
)
|
||||
raise SemanticToolFilterContextWindowError(
|
||||
embedding_model=self.embedding_model,
|
||||
stage="the user query",
|
||||
stage="the user query or the MCP tool descriptions being indexed",
|
||||
original_error=str(e),
|
||||
) from e
|
||||
verbose_logger.error(f"Semantic tool filter failed: {e}", exc_info=True)
|
||||
|
|
|
|||
|
|
@ -7613,6 +7613,18 @@
|
|||
],
|
||||
"title": "Messages"
|
||||
},
|
||||
"metadata": {
|
||||
"anyOf": [
|
||||
{
|
||||
"additionalProperties": true,
|
||||
"type": "object"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Metadata"
|
||||
},
|
||||
"text": {
|
||||
"title": "Text",
|
||||
"type": "string"
|
||||
|
|
|
|||
|
|
@ -1263,6 +1263,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
command: Optional[str] = None
|
||||
args: List[str] = Field(default_factory=list)
|
||||
env: Dict[str, str] = Field(default_factory=dict)
|
||||
issuer: Optional[str] = None
|
||||
authorization_url: Optional[str] = None
|
||||
token_url: Optional[str] = None
|
||||
registration_url: Optional[str] = None
|
||||
|
|
@ -1368,6 +1369,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
command: Optional[str] = None
|
||||
args: List[str] = Field(default_factory=list)
|
||||
env: Dict[str, str] = Field(default_factory=dict)
|
||||
issuer: Optional[str] = None
|
||||
authorization_url: Optional[str] = None
|
||||
token_url: Optional[str] = None
|
||||
registration_url: Optional[str] = None
|
||||
|
|
@ -2367,6 +2369,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
None,
|
||||
description="If True, stores request messages and responses in spend logs. Default is False.",
|
||||
)
|
||||
disable_auto_add_proxy_admin_to_teams: bool | None = Field(
|
||||
None,
|
||||
description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.",
|
||||
)
|
||||
maximum_spend_logs_retention_period: Optional[str] = Field(
|
||||
None,
|
||||
description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.",
|
||||
|
|
|
|||
|
|
@ -361,7 +361,11 @@ def _global_proxy_budget_check(global_proxy_spend: Optional[float], skip_budget_
|
|||
and route != "/models"
|
||||
):
|
||||
if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget:
|
||||
raise litellm.BudgetExceededError(current_cost=global_proxy_spend, max_budget=litellm.max_budget)
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=global_proxy_spend,
|
||||
max_budget=litellm.max_budget,
|
||||
entity_type=Litellm_EntityType.PROXY.value,
|
||||
)
|
||||
|
||||
|
||||
_GUARDRAIL_MODIFICATION_KEYS: tuple = (
|
||||
|
|
@ -648,6 +652,8 @@ async def common_checks(
|
|||
current_cost=user_spend,
|
||||
max_budget=user_budget,
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
entity_type=Litellm_EntityType.USER.value,
|
||||
entity_id=user_object.user_id,
|
||||
)
|
||||
|
||||
# Each scope reads a distinct counter key with no cross-scope ordering
|
||||
|
|
@ -1093,6 +1099,8 @@ async def _check_end_user_budget(
|
|||
current_cost=end_user_spend,
|
||||
max_budget=end_user_budget,
|
||||
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
|
||||
entity_type=Litellm_EntityType.END_USER.value,
|
||||
entity_id=end_user_obj.user_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3552,6 +3560,8 @@ async def _virtual_key_max_budget_check(
|
|||
current_cost=spend,
|
||||
max_budget=valid_token.max_budget,
|
||||
message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}",
|
||||
entity_type=Litellm_EntityType.KEY.value,
|
||||
entity_id=valid_token.token,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3593,6 +3603,8 @@ async def _virtual_key_multi_budget_check(
|
|||
f"ExceededBudget: Key over {w['budget_duration']} budget. "
|
||||
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
|
||||
),
|
||||
entity_type=Litellm_EntityType.KEY.value,
|
||||
entity_id=valid_token.token,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3824,6 +3836,8 @@ async def _check_team_member_budget(
|
|||
current_cost=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}",
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
|
||||
entity_id=f"{valid_token.user_id}:{team_object.team_id}",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3923,6 +3937,8 @@ async def _team_max_budget_check(
|
|||
current_cost=spend,
|
||||
max_budget=team_object.max_budget,
|
||||
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}",
|
||||
entity_type=Litellm_EntityType.TEAM.value,
|
||||
entity_id=team_object.team_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -3960,6 +3976,8 @@ async def _team_multi_budget_check(
|
|||
f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. "
|
||||
f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
|
||||
),
|
||||
entity_type=Litellm_EntityType.TEAM.value,
|
||||
entity_id=team_object.team_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4081,6 +4099,8 @@ async def _project_max_budget_check(
|
|||
current_cost=project_object.spend,
|
||||
max_budget=max_budget,
|
||||
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}",
|
||||
entity_type=Litellm_EntityType.PROJECT.value,
|
||||
entity_id=project_object.project_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4269,6 +4289,8 @@ async def _organization_max_budget_check(
|
|||
current_cost=org_spend,
|
||||
max_budget=org_max_budget,
|
||||
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
|
||||
entity_type=Litellm_EntityType.ORGANIZATION.value,
|
||||
entity_id=org_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4326,6 +4348,8 @@ async def _tag_max_budget_check(
|
|||
current_cost=tag_spend,
|
||||
max_budget=tag_object.litellm_budget_table.max_budget,
|
||||
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
|
||||
entity_type=Litellm_EntityType.TAG.value,
|
||||
entity_id=tag_name,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:]}"
|
||||
|
|
|
|||
|
|
@ -1191,13 +1191,15 @@ async def _user_api_key_auth_builder(
|
|||
return await handle_oauth2_proxy_request(request=request)
|
||||
|
||||
if general_settings.get("enable_jwt_auth", False) is True:
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
if premium_user is not True:
|
||||
raise ValueError(f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}")
|
||||
is_jwt = jwt_handler.is_jwt(token=api_key)
|
||||
verbose_proxy_logger.debug("is_jwt: %s", is_jwt)
|
||||
if is_jwt:
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
if premium_user is not True:
|
||||
raise ValueError(
|
||||
f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}"
|
||||
)
|
||||
# Try JWT-to-Virtual-Key mapping first to avoid
|
||||
# unnecessary DB queries in auth_builder
|
||||
do_standard_jwt_auth = True
|
||||
|
|
@ -1795,6 +1797,8 @@ async def _user_api_key_auth_builder(
|
|||
raise litellm.BudgetExceededError(
|
||||
current_cost=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
|
||||
entity_id=f"{valid_token.user_id}:{valid_token.team_id}",
|
||||
)
|
||||
|
||||
# Check 3. If token is expired
|
||||
|
|
@ -1992,16 +1996,6 @@ async def _user_api_key_auth_builder(
|
|||
raise HTTPException(401, detail="Invalid API key, no token associated")
|
||||
api_key = valid_token.token
|
||||
|
||||
# Add hashed token to cache
|
||||
asyncio.create_task(
|
||||
_cache_key_object(
|
||||
hashed_token=api_key,
|
||||
user_api_key_obj=valid_token,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
|
||||
valid_token_dict = valid_token.model_dump(exclude_none=True)
|
||||
valid_token_dict.pop("token", None)
|
||||
# budget_throttle_pct is excluded from model_dump (it must not leak
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None:
|
|||
console = Console()
|
||||
|
||||
if not teams:
|
||||
console.print("❌ No teams found for your user.")
|
||||
console.print("No teams found for your user.")
|
||||
return
|
||||
|
||||
table = Table(title="Available Teams")
|
||||
|
|
@ -162,7 +162,7 @@ def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_ind
|
|||
# Clear the screen using Rich's method
|
||||
console.clear()
|
||||
|
||||
console.print("🎯 Select a Team (Use ↑↓ arrows, Enter to select, 'q' to skip):\n")
|
||||
console.print("Select a Team (Use up/down arrows, Enter to select, 'q' to skip):\n")
|
||||
|
||||
for i, team in enumerate(teams):
|
||||
team_alias = team.get("team_alias") or "N/A"
|
||||
|
|
@ -184,7 +184,7 @@ def display_interactive_team_selection(teams: List[Dict[str, Any]], selected_ind
|
|||
|
||||
# Highlight the selected item
|
||||
if i == selected_index:
|
||||
console.print(f"➤ [bold cyan]{team_alias}[/bold cyan] ({team_id})")
|
||||
console.print(f"> [bold cyan]{team_alias}[/bold cyan] ({team_id})")
|
||||
console.print(f" Models: [yellow]{models_str}[/yellow]")
|
||||
console.print(f" Budget: [blue]{budget_str}[/blue]\n")
|
||||
else:
|
||||
|
|
@ -220,15 +220,13 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
|
|||
# Clear screen and show selection
|
||||
console = Console()
|
||||
console.clear()
|
||||
click.echo(
|
||||
f"✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
|
||||
)
|
||||
click.echo(f"Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})")
|
||||
return selected_team
|
||||
elif key == "quit" or key == "escape":
|
||||
# Clear screen
|
||||
console = Console()
|
||||
console.clear()
|
||||
click.echo("ℹ️ Team selection skipped.")
|
||||
click.echo("Team selection skipped.")
|
||||
return None
|
||||
elif key is None:
|
||||
# If we can't get key input, fall back to simple selection
|
||||
|
|
@ -237,7 +235,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
|
|||
except KeyboardInterrupt:
|
||||
console = Console()
|
||||
console.clear()
|
||||
click.echo("\n❌ Team selection cancelled.")
|
||||
click.echo("\nTeam selection cancelled.")
|
||||
return None
|
||||
except Exception:
|
||||
# If interactive mode fails, fall back to simple selection
|
||||
|
|
@ -265,18 +263,40 @@ def prompt_team_selection_fallback(
|
|||
if 0 <= index < len(teams):
|
||||
selected_team = teams[index]
|
||||
click.echo(
|
||||
f"\n✅ Selected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
|
||||
f"\nSelected team: {selected_team.get('team_alias', 'N/A')} ({selected_team.get('team_id')})"
|
||||
)
|
||||
return selected_team
|
||||
else:
|
||||
click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}")
|
||||
click.echo(f"Invalid selection. Please enter a number between 1 and {len(teams)}")
|
||||
except ValueError:
|
||||
click.echo("❌ Invalid input. Please enter a number or 'skip'")
|
||||
click.echo("Invalid input. Please enter a number or 'skip'")
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\n❌ Team selection cancelled.")
|
||||
click.echo("\nTeam selection cancelled.")
|
||||
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 +328,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 +368,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
|
||||
|
||||
|
||||
|
|
@ -376,7 +435,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op
|
|||
user_id = data.get("user_id")
|
||||
normalized_teams: List[Dict[str, Any]] = _normalize_teams(teams, team_details)
|
||||
if not normalized_teams:
|
||||
click.echo("⚠️ No teams available for selection.")
|
||||
click.echo("Warning: No teams available for selection.")
|
||||
return None
|
||||
|
||||
# User has multiple teams - let them select
|
||||
|
|
@ -396,7 +455,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op
|
|||
"team_id": None, # Set by server in JWT
|
||||
}
|
||||
|
||||
click.echo("❌ Team selection cancelled or JWT generation failed.")
|
||||
click.echo("Team selection cancelled or JWT generation failed.")
|
||||
return None
|
||||
|
||||
# JWT is ready (single team or team already selected)
|
||||
|
|
@ -407,7 +466,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> Op
|
|||
|
||||
# Show which team was assigned
|
||||
if team_id and len(teams) == 1:
|
||||
click.echo(f"\n✅ Automatically assigned to team: {team_id}")
|
||||
click.echo(f"\nAutomatically assigned to team: {team_id}")
|
||||
|
||||
if api_key:
|
||||
return {
|
||||
|
|
@ -433,19 +492,19 @@ def _handle_team_selection_during_polling(
|
|||
The JWT token with the selected team, or None if selection was skipped
|
||||
"""
|
||||
if not teams:
|
||||
click.echo("ℹ️ No teams found. You can create or join teams using the web interface.")
|
||||
click.echo("No teams found. You can create or join teams using the web interface.")
|
||||
return None
|
||||
|
||||
click.echo("\n" + "=" * 60)
|
||||
click.echo("📋 Select a team for your CLI session...")
|
||||
click.echo("Select a team for your CLI session...")
|
||||
|
||||
team_id = _render_and_prompt_for_team_selection(teams)
|
||||
|
||||
if not team_id:
|
||||
click.echo("ℹ️ No team selected.")
|
||||
click.echo("No team selected.")
|
||||
return None
|
||||
|
||||
click.echo(f"\n🔄 Generating JWT for team: {team_id}")
|
||||
click.echo(f"\nGenerating JWT for team: {team_id}")
|
||||
|
||||
poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
|
||||
data = _poll_for_ready_data(
|
||||
|
|
@ -459,7 +518,7 @@ def _handle_team_selection_during_polling(
|
|||
return None
|
||||
jwt_token = data.get("key")
|
||||
if jwt_token:
|
||||
click.echo(f"✅ Successfully generated JWT for team: {team_id}")
|
||||
click.echo(f"Successfully generated JWT for team: {team_id}")
|
||||
return jwt_token
|
||||
|
||||
return None
|
||||
|
|
@ -507,14 +566,14 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
|
|||
selected_team = teams[index]
|
||||
team_id = str(selected_team.get("team_id"))
|
||||
team_alias = selected_team.get("team_alias") or team_id
|
||||
click.echo(f"\n✅ Selected team: {team_alias} ({team_id})")
|
||||
click.echo(f"\nSelected team: {team_alias} ({team_id})")
|
||||
return team_id
|
||||
|
||||
click.echo(f"❌ Invalid selection. Please enter a number between 1 and {len(teams)}")
|
||||
click.echo(f"Invalid selection. Please enter a number between 1 and {len(teams)}")
|
||||
except ValueError:
|
||||
click.echo("❌ Invalid input. Please enter a number or 'skip'")
|
||||
click.echo("Invalid input. Please enter a number or 'skip'")
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\n❌ Team selection cancelled.")
|
||||
click.echo("\nTeam selection cancelled.")
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -567,7 +626,7 @@ def login(ctx: click.Context):
|
|||
}
|
||||
)
|
||||
|
||||
click.echo("\n✅ Login successful!")
|
||||
click.echo("\nLogin successful!")
|
||||
click.echo(f"JWT Token: {api_key[:20]}...")
|
||||
click.echo("You can now use the CLI without specifying --api-key")
|
||||
|
||||
|
|
@ -576,14 +635,18 @@ def login(ctx: click.Context):
|
|||
show_commands()
|
||||
return
|
||||
else:
|
||||
click.echo("❌ Authentication timed out. Please try again.")
|
||||
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:
|
||||
click.echo("\n❌ Authentication cancelled by user.")
|
||||
click.echo("\nAuthentication cancelled by user.")
|
||||
return
|
||||
except Exception as e:
|
||||
click.echo(f"❌ Authentication failed: {e}")
|
||||
click.echo(f"Authentication failed: {e}")
|
||||
return
|
||||
|
||||
|
||||
|
|
@ -591,7 +654,7 @@ def login(ctx: click.Context):
|
|||
def logout():
|
||||
"""Logout and clear stored authentication"""
|
||||
clear_token()
|
||||
click.echo("✅ Logged out successfully. Authentication token cleared.")
|
||||
click.echo("Logged out successfully. Authentication token cleared.")
|
||||
|
||||
|
||||
@click.command(name="print-token")
|
||||
|
|
@ -638,10 +701,10 @@ def whoami():
|
|||
token_data = load_token()
|
||||
|
||||
if not token_data:
|
||||
click.echo("❌ Not authenticated. Run 'lite login' to authenticate.")
|
||||
click.echo("Not authenticated. Run 'lite login' to authenticate.")
|
||||
return
|
||||
|
||||
click.echo("✅ Authenticated")
|
||||
click.echo("Authenticated")
|
||||
click.echo(f"User Email: {token_data.get('user_email', 'Unknown')}")
|
||||
click.echo(f"User ID: {token_data.get('user_id', 'Unknown')}")
|
||||
click.echo(f"User Role: {token_data.get('user_role', 'Unknown')}")
|
||||
|
|
@ -652,7 +715,7 @@ def whoami():
|
|||
click.echo(f"Token age: {age_hours:.1f} hours")
|
||||
|
||||
if age_hours > CLI_JWT_EXPIRATION_HOURS:
|
||||
click.echo(f"⚠️ Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.")
|
||||
click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.")
|
||||
|
||||
|
||||
@click.group(name="auth")
|
||||
|
|
|
|||
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"]
|
||||
|
|
@ -150,7 +150,7 @@ def chat(
|
|||
f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n"
|
||||
f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n"
|
||||
f"Type '/help' for more commands.",
|
||||
title="🤖 Chat Session",
|
||||
title="Chat Session",
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ def migrate(ctx: click.Context, check_only: bool, dry_run: bool):
|
|||
|
||||
Requires the proxy to be started with
|
||||
``general_settings.encryption_algorithm: aes-256-gcm``. Idempotent and
|
||||
resumable — safe to re-run after an interruption.
|
||||
resumable; safe to re-run after an interruption.
|
||||
|
||||
Examples:
|
||||
litellm-proxy encryption migrate --check # attestation scan, no writes
|
||||
|
|
|
|||
|
|
@ -309,12 +309,12 @@ def _import_keys_to_destination(
|
|||
imported_count += 1
|
||||
|
||||
key_alias = key.get("key_alias", "N/A")
|
||||
click.echo(f"✓ Imported key: {key_alias}")
|
||||
click.echo(f"Imported key: {key_alias}")
|
||||
|
||||
except Exception as e:
|
||||
failed_count += 1
|
||||
key_alias = key.get("key_alias", "N/A")
|
||||
click.echo(f"✗ Failed to import key {key_alias}: {str(e)}", err=True)
|
||||
click.echo(f"Failed to import key {key_alias}: {str(e)}", err=True)
|
||||
|
||||
return imported_count, failed_count
|
||||
|
||||
|
|
|
|||
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"]
|
||||
|
|
@ -21,7 +21,7 @@ def display_teams_table(teams: List[Dict[str, Any]]) -> None:
|
|||
console = Console()
|
||||
|
||||
if not teams:
|
||||
console.print("❌ No teams found for your user.")
|
||||
console.print("No teams found for your user.")
|
||||
return
|
||||
|
||||
table = Table(title="Available Teams")
|
||||
|
|
@ -91,10 +91,10 @@ def available(ctx: click.Context):
|
|||
teams = client.teams.get_available()
|
||||
if teams:
|
||||
console = Console()
|
||||
console.print("\n🎯 Available Teams to Join:")
|
||||
console.print("\nAvailable Teams to Join:")
|
||||
display_teams_table(teams)
|
||||
else:
|
||||
click.echo("ℹ️ No available teams to join.")
|
||||
click.echo("No available teams to join.")
|
||||
except requests.exceptions.HTTPError as e:
|
||||
click.echo(f"Error: HTTP {e.response.status_code}", err=True)
|
||||
error_body = e.response.json()
|
||||
|
|
@ -113,7 +113,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
|
|||
api_key = ctx.obj["api_key"]
|
||||
|
||||
if not api_key:
|
||||
click.echo("❌ No API key found. Please login first using 'litellm login'")
|
||||
click.echo("No API key found. Please login first using 'litellm login'")
|
||||
raise click.Abort()
|
||||
|
||||
try:
|
||||
|
|
@ -122,7 +122,7 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
|
|||
teams = client.teams.list()
|
||||
|
||||
if not teams:
|
||||
click.echo("❌ No teams found for your user.")
|
||||
click.echo("No teams found for your user.")
|
||||
return
|
||||
|
||||
# Use interactive selection from auth module
|
||||
|
|
@ -133,14 +133,14 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
|
|||
if selected_team:
|
||||
team_id = selected_team.get("team_id")
|
||||
else:
|
||||
click.echo("❌ Operation cancelled.")
|
||||
click.echo("Operation cancelled.")
|
||||
return
|
||||
|
||||
# Update the key with the selected team
|
||||
if team_id:
|
||||
click.echo(f"\n🔄 Assigning your key to team: {team_id}")
|
||||
click.echo(f"\nAssigning your key to team: {team_id}")
|
||||
client.keys.update(key=api_key, team_id=team_id)
|
||||
click.echo(f"✅ Successfully assigned key to team: {team_id}")
|
||||
click.echo(f"Successfully assigned key to team: {team_id}")
|
||||
|
||||
# Show team details if available
|
||||
teams = client.teams.list()
|
||||
|
|
@ -148,9 +148,9 @@ def assign_key(ctx: click.Context, team_id: Optional[str]):
|
|||
if team.get("team_id") == team_id:
|
||||
models = team.get("models", [])
|
||||
if models:
|
||||
click.echo(f"🎯 You can now access models: {', '.join(models)}")
|
||||
click.echo(f"You can now access models: {', '.join(models)}")
|
||||
else:
|
||||
click.echo("🎯 You can now access all available models")
|
||||
click.echo("You can now access all available models")
|
||||
break
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
|
|||
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",
|
||||
]
|
||||
|
|
@ -27,13 +27,13 @@ def styled_prompt():
|
|||
verbose_logger.debug(f"Error getting terminal size: {e}")
|
||||
click.echo("\n" * 3)
|
||||
|
||||
# Unicode box drawing characters
|
||||
top_left = "┌"
|
||||
top_right = "┐"
|
||||
bottom_left = "└"
|
||||
bottom_right = "┘"
|
||||
horizontal = "─"
|
||||
vertical = "│"
|
||||
# ASCII box drawing characters
|
||||
top_left = "+"
|
||||
top_right = "+"
|
||||
bottom_left = "+"
|
||||
bottom_right = "+"
|
||||
horizontal = "-"
|
||||
vertical = "|"
|
||||
|
||||
# Create the box with increased width
|
||||
width = 80
|
||||
|
|
|
|||
|
|
@ -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__":
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ from typing import (
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import DualCache, RedisCache
|
||||
from litellm.caching import RedisCache
|
||||
from litellm.constants import (
|
||||
DB_SPEND_UPDATE_JOB_NAME,
|
||||
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
|
||||
|
|
@ -44,7 +44,6 @@ from litellm.proxy._types import (
|
|||
DailyUserSpendTransaction,
|
||||
DBSpendUpdateTransactions,
|
||||
Litellm_EntityType,
|
||||
LiteLLM_UserTable,
|
||||
SpendLogsMetadata,
|
||||
SpendLogsPayload,
|
||||
SpendUpdateQueueItem,
|
||||
|
|
@ -137,7 +136,6 @@ class DBSpendUpdateWriter:
|
|||
disable_spend_logs,
|
||||
litellm_proxy_budget_name,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyUpdateSpend, hash_token
|
||||
|
||||
|
|
@ -195,7 +193,6 @@ class DBSpendUpdateWriter:
|
|||
org_id=org_id,
|
||||
end_user_id=end_user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
litellm_proxy_budget_name=litellm_proxy_budget_name,
|
||||
payload=payload,
|
||||
)
|
||||
|
|
@ -326,7 +323,6 @@ class DBSpendUpdateWriter:
|
|||
org_id: Optional[str],
|
||||
end_user_id: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
litellm_proxy_budget_name: Optional[str],
|
||||
payload: SpendLogsPayload,
|
||||
):
|
||||
|
|
@ -345,7 +341,6 @@ class DBSpendUpdateWriter:
|
|||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
litellm_proxy_budget_name=litellm_proxy_budget_name,
|
||||
end_user_id=end_user_id,
|
||||
)
|
||||
|
|
@ -510,7 +505,6 @@ class DBSpendUpdateWriter:
|
|||
response_cost: Optional[float],
|
||||
user_id: Optional[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
litellm_proxy_budget_name: Optional[str],
|
||||
end_user_id: Optional[str] = None,
|
||||
):
|
||||
|
|
@ -518,10 +512,6 @@ class DBSpendUpdateWriter:
|
|||
- Update that user's row
|
||||
- Update litellm-proxy-budget row (global proxy spend)
|
||||
"""
|
||||
## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db
|
||||
existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
|
||||
if existing_user_obj is not None and isinstance(existing_user_obj, dict):
|
||||
existing_user_obj = LiteLLM_UserTable(**existing_user_obj)
|
||||
try:
|
||||
if prisma_client is not None: # update
|
||||
user_ids = [user_id]
|
||||
|
|
|
|||
212
litellm/proxy/db/query_engine_reaper.py
Normal file
212
litellm/proxy/db/query_engine_reaper.py
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
"""Supervisor-side reaper for orphaned Prisma query-engine processes.
|
||||
|
||||
Each proxy worker owns a Prisma query-engine subprocess whose only cleanup
|
||||
hook is an in-process ``atexit`` handler. When a multi-worker supervisor
|
||||
(uvicorn's multiprocess manager, the gunicorn arbiter) force-kills a hung or
|
||||
crashed worker, that handler never runs: the engine reparents to the nearest
|
||||
subreaper (PID 1 in a container, which is the supervisor itself under the
|
||||
standard docker entrypoint) and keeps its database connection pool
|
||||
established forever, while the replacement worker opens a fresh pool. Over
|
||||
repeated worker deaths the active database connections grow without bound.
|
||||
|
||||
The reaper runs only in the supervisor process, where a query-engine process
|
||||
can never be a legitimate direct child: workers own their engines, and the
|
||||
supervisor never starts one. Any direct child whose command name begins with
|
||||
``query-engine`` is therefore an adopted orphan and is terminated
|
||||
(SIGTERM, bounded grace, SIGKILL) and reaped. On Linux the supervisor also
|
||||
marks itself a child subreaper so orphans reparent to it even when it is not
|
||||
PID 1.
|
||||
|
||||
Linux-only by construction (``/proc`` scan, ``prctl``); a no-op elsewhere.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
QUERY_ENGINE_COMM_PREFIX = "query-engine"
|
||||
REAPER_SCAN_INTERVAL_SECONDS = 5.0
|
||||
SIGTERM_GRACE_SECONDS = 10.0
|
||||
PR_SET_CHILD_SUBREAPER = 36
|
||||
|
||||
|
||||
def set_child_subreaper() -> bool:
|
||||
"""Mark this process as a child subreaper so orphaned descendants
|
||||
reparent to it instead of PID 1. Best-effort: when it fails (or on
|
||||
non-Linux) the reaper still covers the containerized case where the
|
||||
supervisor already is PID 1."""
|
||||
if not sys.platform.startswith("linux"):
|
||||
return False
|
||||
try:
|
||||
libc = ctypes.CDLL(None, use_errno=True)
|
||||
result: int = libc.prctl( # pyright: ignore[reportAny] # ctypes types foreign calls as Any; default restype is c_int
|
||||
PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0
|
||||
)
|
||||
return result == 0
|
||||
except (OSError, AttributeError):
|
||||
return False
|
||||
|
||||
|
||||
def _read_comm_and_ppid(pid: int, proc_root: str) -> Optional[tuple[str, int]]:
|
||||
try:
|
||||
with open(f"{proc_root}/{pid}/stat", encoding="ascii", errors="replace") as stat_file:
|
||||
data = stat_file.read()
|
||||
except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
|
||||
return None
|
||||
lparen = data.find("(")
|
||||
rparen = data.rfind(")")
|
||||
if lparen == -1 or rparen == -1 or rparen < lparen:
|
||||
return None
|
||||
comm = data[lparen + 1 : rparen]
|
||||
fields = data[rparen + 2 :].split()
|
||||
if len(fields) < 2:
|
||||
return None
|
||||
try:
|
||||
ppid = int(fields[1])
|
||||
except ValueError:
|
||||
return None
|
||||
return comm, ppid
|
||||
|
||||
|
||||
def list_orphaned_engine_pids(parent_pid: int, proc_root: str = "/proc") -> tuple[int, ...]:
|
||||
"""PIDs of direct children of ``parent_pid`` whose command name marks
|
||||
them as Prisma query engines. In the supervisor these are always
|
||||
adopted orphans: live engines are children of workers, not of the
|
||||
supervisor."""
|
||||
try:
|
||||
entries = os.listdir(proc_root)
|
||||
except (FileNotFoundError, OSError):
|
||||
return ()
|
||||
candidate_pids = (int(entry) for entry in entries if entry.isdigit())
|
||||
return tuple(
|
||||
pid
|
||||
for pid in candidate_pids
|
||||
if (info := _read_comm_and_ppid(pid, proc_root)) is not None
|
||||
and info[1] == parent_pid
|
||||
and info[0].startswith(QUERY_ENGINE_COMM_PREFIX)
|
||||
)
|
||||
|
||||
|
||||
def _try_reap(pid: int) -> bool:
|
||||
try:
|
||||
reaped_pid, _ = os.waitpid(pid, os.WNOHANG)
|
||||
except ChildProcessError:
|
||||
return True
|
||||
except OSError:
|
||||
return True
|
||||
return reaped_pid == pid
|
||||
|
||||
|
||||
def _send_signal(pid: int, signum: int) -> None:
|
||||
try:
|
||||
os.kill(pid, signum)
|
||||
except (ProcessLookupError, PermissionError, OSError):
|
||||
pass
|
||||
|
||||
|
||||
def _await_reaped(pids: tuple[int, ...], timeout_seconds: float) -> tuple[int, ...]:
|
||||
"""Poll until every PID is reaped or the shared deadline passes.
|
||||
Returns the PIDs still alive at the deadline."""
|
||||
deadline = time.monotonic() + timeout_seconds
|
||||
remaining = pids
|
||||
while remaining and time.monotonic() < deadline:
|
||||
remaining = tuple(pid for pid in remaining if not _try_reap(pid))
|
||||
if remaining:
|
||||
time.sleep(0.2)
|
||||
return remaining
|
||||
|
||||
|
||||
def terminate_and_reap(pid: int, grace_seconds: float = SIGTERM_GRACE_SECONDS) -> None:
|
||||
"""SIGTERM the orphaned engine, escalate to SIGKILL after the grace
|
||||
period, and reap it so it does not linger as a zombie."""
|
||||
terminate_and_reap_all((pid,), grace_seconds=grace_seconds)
|
||||
|
||||
|
||||
def terminate_and_reap_all(
|
||||
pids: tuple[int, ...],
|
||||
grace_seconds: float = SIGTERM_GRACE_SECONDS,
|
||||
) -> None:
|
||||
"""Terminate a batch of orphaned engines concurrently: SIGTERM all of
|
||||
them, share one grace period, SIGKILL the stragglers, and reap. The
|
||||
shared deadline keeps cleanup time bounded when several workers die
|
||||
at once instead of paying the grace period once per orphan."""
|
||||
for pid in pids:
|
||||
verbose_proxy_logger.warning(
|
||||
"Reaping orphaned prisma query-engine PID %s (its worker process exited without cleanup).",
|
||||
pid,
|
||||
)
|
||||
_send_signal(pid, signal.SIGTERM)
|
||||
survivors = _await_reaped(pids, grace_seconds)
|
||||
if not survivors:
|
||||
return
|
||||
for pid in survivors:
|
||||
verbose_proxy_logger.warning(
|
||||
"Orphaned prisma query-engine PID %s did not exit within %.1fs of SIGTERM; sending SIGKILL.",
|
||||
pid,
|
||||
grace_seconds,
|
||||
)
|
||||
_send_signal(pid, signal.SIGKILL)
|
||||
unkillable = _await_reaped(survivors, 5.0)
|
||||
for pid in unkillable:
|
||||
verbose_proxy_logger.error(
|
||||
"Orphaned prisma query-engine PID %s survived SIGKILL; will retry on the next scan.",
|
||||
pid,
|
||||
)
|
||||
|
||||
|
||||
def reap_orphaned_engines(parent_pid: int, proc_root: str = "/proc") -> tuple[int, ...]:
|
||||
"""One scan-and-reap pass. Returns the PIDs it acted on."""
|
||||
orphaned_pids = list_orphaned_engine_pids(parent_pid, proc_root=proc_root)
|
||||
if orphaned_pids:
|
||||
terminate_and_reap_all(orphaned_pids)
|
||||
return orphaned_pids
|
||||
|
||||
|
||||
def _reaper_loop(parent_pid: int) -> None:
|
||||
while True:
|
||||
try:
|
||||
reap_orphaned_engines(parent_pid)
|
||||
except Exception as scan_error: # noqa: BLE001 # reaper thread must survive any scan failure
|
||||
verbose_proxy_logger.debug("Orphaned query-engine scan failed: %s", scan_error)
|
||||
time.sleep(REAPER_SCAN_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
REAPER_THREAD_NAME = "litellm-orphan-query-engine-reaper"
|
||||
|
||||
|
||||
def start_query_engine_reaper() -> Optional[threading.Thread]:
|
||||
"""Start the reaper daemon thread in the supervisor process.
|
||||
|
||||
Must only be called from a process that never hosts the proxy app
|
||||
itself (uvicorn with ``workers > 1``, the gunicorn arbiter): with a
|
||||
single in-process uvicorn worker the query engine is a legitimate
|
||||
direct child and must not be touched. Idempotent: a reaper already
|
||||
running in this process is returned instead of starting a second one.
|
||||
"""
|
||||
if not sys.platform.startswith("linux"):
|
||||
return None
|
||||
existing = next(
|
||||
(thread for thread in threading.enumerate() if thread.name == REAPER_THREAD_NAME),
|
||||
None,
|
||||
)
|
||||
if existing is not None:
|
||||
return existing
|
||||
set_child_subreaper()
|
||||
reaper_thread = threading.Thread(
|
||||
target=_reaper_loop,
|
||||
args=(os.getpid(),),
|
||||
daemon=True,
|
||||
name=REAPER_THREAD_NAME,
|
||||
)
|
||||
reaper_thread.start()
|
||||
verbose_proxy_logger.info(
|
||||
"Started orphaned prisma query-engine reaper in supervisor process %s.",
|
||||
os.getpid(),
|
||||
)
|
||||
return reaper_thread
|
||||
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]},
|
||||
|
|
|
|||
|
|
@ -26,8 +26,13 @@ from typing import (
|
|||
cast,
|
||||
)
|
||||
|
||||
import copy
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -42,10 +47,13 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
||||
BedrockChecksMessage,
|
||||
BedrockChecksViolation,
|
||||
BedrockContentItem,
|
||||
BedrockGuardrailChecksResponse,
|
||||
BedrockGuardrailOutput,
|
||||
BedrockGuardrailQualifier,
|
||||
BedrockGuardrailResponse,
|
||||
|
|
@ -55,6 +63,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
|||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botocore.awsrequest import AWSPreparedRequest
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -72,6 +82,22 @@ from litellm.types.utils import (
|
|||
|
||||
GUARDRAIL_NAME = "bedrock"
|
||||
_BEDROCK_DYNAMIC_BODY_DENYLIST = frozenset({"content", "source"})
|
||||
# Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required).
|
||||
_BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH = "/guardrail-checks/invoke"
|
||||
# InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with
|
||||
# more text blocks is split across multiple messages so ALL content is scanned --
|
||||
# never truncated (truncation would let a user hide content past the limit).
|
||||
_BEDROCK_CHECKS_MAX_CONTENT_BLOCKS = 10
|
||||
_BEDROCK_CHECKS_KNOWN_KEYS = frozenset({"contentFilter", "promptAttack", "sensitiveInformation"})
|
||||
# Keys in a sensitiveInformation result that pinpoint the PII location. They are
|
||||
# stripped before the response is handed to standard logging / telemetry so the
|
||||
# detected PII span cannot be reconstructed from logs.
|
||||
_BEDROCK_CHECKS_PII_LOCATION_KEYS = (
|
||||
"beginOffset",
|
||||
"endOffset",
|
||||
"messageIndex",
|
||||
"contentIndex",
|
||||
)
|
||||
|
||||
# Maps an OpenAI message content-block ``type`` to the Bedrock guardrail qualifier
|
||||
# it represents, so callers can drive contextual grounding by tagging their content.
|
||||
|
|
@ -149,6 +175,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
guardrailIdentifier: Optional[str] = None,
|
||||
guardrailVersion: Optional[str] = None,
|
||||
disable_exception_on_block: Optional[bool] = False,
|
||||
checks: BedrockChecksConfigModel | Mapping[str, object] | None = None,
|
||||
content_filter_threshold: float | None = 0.5,
|
||||
prompt_attack_threshold: float | None = 0.5,
|
||||
pii_confidence_threshold: float | None = 0.5,
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
|
|
@ -157,6 +187,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self.guardrail_provider = "bedrock"
|
||||
self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
|
||||
|
||||
# Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
|
||||
# routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail.
|
||||
self.checks: dict[str, Any] | None = self._normalize_checks(checks)
|
||||
# Per-check block thresholds; a score >= threshold blocks. None => the
|
||||
# check is detect-only (logged, never blocks).
|
||||
self.content_filter_threshold = content_filter_threshold
|
||||
self.prompt_attack_threshold = prompt_attack_threshold
|
||||
self.pii_confidence_threshold = pii_confidence_threshold
|
||||
|
||||
# store kwargs as optional_params
|
||||
self.optional_params = kwargs
|
||||
|
||||
|
|
@ -165,16 +204,35 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
If True, will not raise an exception when the guardrail is blocked.
|
||||
"""
|
||||
|
||||
# `checks` (InvokeGuardrailChecks) and `guardrailIdentifier`/`guardrailVersion`
|
||||
# (ApplyGuardrail) are two different APIs; configuring both is ambiguous.
|
||||
if self.checks is not None and (self.guardrailIdentifier is not None or self.guardrailVersion is not None):
|
||||
raise ValueError(
|
||||
"Bedrock guardrail accepts either 'guardrailIdentifier'/'guardrailVersion' (ApplyGuardrail) "
|
||||
"or 'checks' (InvokeGuardrailChecks), not both."
|
||||
)
|
||||
|
||||
# Set supported event hooks to include MCP hooks
|
||||
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
|
||||
|
||||
super().__init__(**kwargs)
|
||||
BaseAWSLLM.__init__(self)
|
||||
|
||||
# InvokeGuardrailChecks is detect-only: it never returns rewritten content,
|
||||
# so masking has no effect in checks mode.
|
||||
if self.checks is not None and (
|
||||
getattr(self, "mask_request_content", False) or getattr(self, "mask_response_content", False)
|
||||
):
|
||||
verbose_proxy_logger.warning(
|
||||
"Bedrock Guardrail: mask_request_content/mask_response_content have no "
|
||||
"effect with 'checks' (InvokeGuardrailChecks is detect-only)."
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Guardrail initialized with guardrailIdentifier: %s, guardrailVersion: %s",
|
||||
"Bedrock Guardrail initialized with guardrailIdentifier: %s, guardrailVersion: %s, checks: %s",
|
||||
self.guardrailIdentifier,
|
||||
self.guardrailVersion,
|
||||
list(self.checks.keys()) if self.checks else None,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
|
|
@ -187,6 +245,34 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
GuardrailEventHooks.during_mcp_call,
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _normalize_checks(checks: BedrockChecksConfigModel | Mapping[str, object] | None) -> dict[str, Any] | None:
|
||||
"""Normalize the configured `checks` into a plain dict for the API body.
|
||||
|
||||
Accepts a pydantic ``BedrockChecksConfigModel`` or a raw dict; drops None /
|
||||
unknown keys. Returns None when no usable check is configured (=> ApplyGuardrail).
|
||||
"""
|
||||
if checks is None:
|
||||
return None
|
||||
raw = checks.model_dump(exclude_none=True) if isinstance(checks, BedrockChecksConfigModel) else dict(checks)
|
||||
unknown_keys = set(raw.keys()) - _BEDROCK_CHECKS_KNOWN_KEYS
|
||||
if unknown_keys:
|
||||
verbose_proxy_logger.warning(
|
||||
"BedrockGuardrail: unrecognized check key(s) %s will be ignored; "
|
||||
"recognized keys will still be used for InvokeGuardrailChecks. "
|
||||
"Known keys: %s.",
|
||||
sorted(unknown_keys),
|
||||
sorted(_BEDROCK_CHECKS_KNOWN_KEYS),
|
||||
)
|
||||
cleaned = {key: value for key, value in raw.items() if key in _BEDROCK_CHECKS_KNOWN_KEYS and value is not None}
|
||||
if not cleaned and raw:
|
||||
raise ValueError(
|
||||
f"BedrockGuardrail: 'checks' block contained only unrecognized or empty keys {sorted(raw.keys())}. "
|
||||
f"Known keys: {sorted(_BEDROCK_CHECKS_KNOWN_KEYS)}. "
|
||||
"Fix the guardrail config or remove the 'checks' block to use ApplyGuardrail mode."
|
||||
)
|
||||
return cleaned or None
|
||||
|
||||
def _create_bedrock_input_content_request(self, messages: Optional[List[AllMessageValues]]) -> BedrockRequest:
|
||||
"""
|
||||
Create a bedrock request for the input content - the LLM request.
|
||||
|
|
@ -574,6 +660,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
aws_region_name: str,
|
||||
api_key: Optional[str] = None,
|
||||
extra_headers: Optional[dict] = None,
|
||||
request_path: str | None = None,
|
||||
):
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if extra_headers is not None:
|
||||
|
|
@ -585,10 +672,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
proxy_endpoint_url = (
|
||||
f"{proxy_endpoint_url}/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
|
||||
)
|
||||
# api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
|
||||
# Default to the ApplyGuardrail resource path. Callers pass an explicit
|
||||
# request_path for the resource-less InvokeGuardrailChecks endpoint (where
|
||||
# guardrailIdentifier/guardrailVersion are None and must not be interpolated).
|
||||
if request_path is None:
|
||||
request_path = f"/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
|
||||
proxy_endpoint_url = f"{proxy_endpoint_url}{request_path}"
|
||||
encoded_data = json.dumps(data).encode("utf-8")
|
||||
|
||||
# first check api-key, if none, fall back to sigV4
|
||||
|
|
@ -635,14 +724,43 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
async def make_bedrock_api_request(
|
||||
self,
|
||||
source: Literal["INPUT", "OUTPUT"],
|
||||
messages: Optional[List[AllMessageValues]] = None,
|
||||
response: Optional[Union[Any, litellm.ModelResponse]] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
logging_event_type: Optional[GuardrailEventHooks] = None,
|
||||
messages: list[AllMessageValues] | None = None,
|
||||
response: litellm.ModelResponse | None = None,
|
||||
request_data: dict | None = None,
|
||||
logging_event_type: GuardrailEventHooks | None = None,
|
||||
) -> BedrockGuardrailResponse:
|
||||
from datetime import datetime
|
||||
"""Dispatch to the configured Bedrock guardrail API.
|
||||
|
||||
start_time = datetime.now()
|
||||
``checks`` selects the resource-less, detect-only InvokeGuardrailChecks API;
|
||||
otherwise the ApplyGuardrail API is used. Both return a ``BedrockGuardrailResponse``
|
||||
(the checks path returns an empty one on a pass, which downstream masking treats
|
||||
as a no-op) and raise on a blocked request.
|
||||
"""
|
||||
if self.checks is not None:
|
||||
return await self._make_invoke_guardrail_checks_request(
|
||||
source=source,
|
||||
messages=messages,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
logging_event_type=logging_event_type,
|
||||
)
|
||||
return await self._make_apply_guardrail_request(
|
||||
source=source,
|
||||
messages=messages,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
logging_event_type=logging_event_type,
|
||||
)
|
||||
|
||||
async def _make_apply_guardrail_request(
|
||||
self,
|
||||
source: Literal["INPUT", "OUTPUT"],
|
||||
messages: list[AllMessageValues] | None = None,
|
||||
response: litellm.ModelResponse | None = None,
|
||||
request_data: dict | None = None,
|
||||
logging_event_type: GuardrailEventHooks | None = None,
|
||||
) -> BedrockGuardrailResponse:
|
||||
start_time = datetime.now(timezone.utc)
|
||||
credentials, aws_region_name = self._load_credentials()
|
||||
bedrock_request_data: dict = dict(
|
||||
self.convert_to_bedrock_format(source=source, messages=messages, response=response)
|
||||
|
|
@ -683,51 +801,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
else:
|
||||
event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call
|
||||
|
||||
try:
|
||||
httpx_response = await self.async_handler.post(
|
||||
url=prepared_request.url,
|
||||
data=prepared_request.body, # type: ignore
|
||||
headers=prepared_request.headers, # type: ignore
|
||||
)
|
||||
except HTTPException:
|
||||
# Propagate HTTPException (e.g. from non-200 path) as-is
|
||||
raise
|
||||
except Exception as e:
|
||||
# If this is an HTTP error with a response body (e.g. httpx.HTTPStatusError),
|
||||
# extract the AWS error message and propagate it
|
||||
response = getattr(e, "response", None)
|
||||
if isinstance(response, httpx.Response):
|
||||
try:
|
||||
(
|
||||
status_code,
|
||||
detail_message,
|
||||
) = self._parse_bedrock_guardrail_error_response(response)
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": detail_message},
|
||||
request_data=request_data or {},
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now().timestamp(),
|
||||
duration=(datetime.now() - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
)
|
||||
raise HTTPException(status_code=status_code, detail=detail_message) from e
|
||||
except HTTPException:
|
||||
raise
|
||||
# Endpoint down, timeout, or other HTTP/network errors
|
||||
verbose_proxy_logger.error("Bedrock AI: failed to make guardrail request: %s", str(e))
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": str(e)},
|
||||
request_data=request_data or {},
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now().timestamp(),
|
||||
duration=(datetime.now() - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
)
|
||||
raise
|
||||
httpx_response = await self._sign_and_post(
|
||||
prepared_request=prepared_request,
|
||||
request_data=request_data,
|
||||
event_type=event_type,
|
||||
start_time=start_time,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Add guardrail information to request trace
|
||||
|
|
@ -743,8 +822,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
request_data=request_data or {},
|
||||
guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now().timestamp(),
|
||||
duration=(datetime.now() - start_time).total_seconds(),
|
||||
end_time=datetime.now(timezone.utc).timestamp(),
|
||||
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
tracing_detail=tracing_detail or None,
|
||||
)
|
||||
|
|
@ -771,6 +850,338 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
|
||||
return bedrock_guardrail_response
|
||||
|
||||
async def _sign_and_post(
|
||||
self,
|
||||
prepared_request: "AWSPreparedRequest",
|
||||
request_data: dict | None,
|
||||
event_type: GuardrailEventHooks,
|
||||
start_time: "datetime",
|
||||
) -> httpx.Response:
|
||||
"""POST a signed Bedrock request, logging+raising on network/HTTP errors.
|
||||
|
||||
Shared by both the ApplyGuardrail and InvokeGuardrailChecks paths so their
|
||||
transport-error handling cannot drift. Returns the raw ``httpx.Response`` on
|
||||
success (including non-2xx that httpx did not raise on); the 200-path logging,
|
||||
status and tracing stay with each caller because the two APIs report differently.
|
||||
"""
|
||||
try:
|
||||
return await self.async_handler.post(
|
||||
url=prepared_request.url,
|
||||
data=prepared_request.body,
|
||||
headers=prepared_request.headers,
|
||||
)
|
||||
except HTTPException:
|
||||
# Propagate HTTPException (e.g. from non-200 path) as-is
|
||||
raise
|
||||
except Exception as e:
|
||||
# If this is an HTTP error with a response body (e.g. httpx.HTTPStatusError),
|
||||
# extract the AWS error message and propagate it
|
||||
err_response = getattr(e, "response", None)
|
||||
if isinstance(err_response, httpx.Response):
|
||||
try:
|
||||
(
|
||||
status_code,
|
||||
detail_message,
|
||||
) = self._parse_bedrock_guardrail_error_response(err_response)
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": detail_message},
|
||||
request_data=request_data or {},
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now(timezone.utc).timestamp(),
|
||||
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
)
|
||||
raise HTTPException(status_code=status_code, detail=detail_message) from e
|
||||
except HTTPException:
|
||||
raise
|
||||
# Endpoint down, timeout, or other HTTP/network errors
|
||||
verbose_proxy_logger.error("Bedrock AI: failed to make guardrail request: %s", str(e))
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": str(e)},
|
||||
request_data=request_data or {},
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now(timezone.utc).timestamp(),
|
||||
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
)
|
||||
raise
|
||||
|
||||
########### InvokeGuardrailChecks (resource-less, detect-only) ############
|
||||
|
||||
@staticmethod
|
||||
def _chunk_texts_into_checks_messages(
|
||||
role: Literal["user", "assistant", "system"], texts: list[str]
|
||||
) -> list[BedrockChecksMessage]:
|
||||
"""Group ``texts`` into role-tagged messages of <= the API content-block cap.
|
||||
|
||||
A source message with more text blocks than the per-message limit is split
|
||||
across multiple messages so EVERY block is scanned. Truncating instead would
|
||||
let a user hide prohibited content past the limit (guardrail bypass).
|
||||
"""
|
||||
cap = _BEDROCK_CHECKS_MAX_CONTENT_BLOCKS
|
||||
return [
|
||||
BedrockChecksMessage(
|
||||
role=role,
|
||||
content=[{"text": text} for text in texts[start : start + cap]],
|
||||
)
|
||||
for start in range(0, len(texts), cap)
|
||||
]
|
||||
|
||||
def _build_invoke_guardrail_checks_messages(
|
||||
self,
|
||||
source: Literal["INPUT", "OUTPUT"],
|
||||
messages: list[AllMessageValues] | None = None,
|
||||
response: litellm.ModelResponse | None = None,
|
||||
) -> list[BedrockChecksMessage]:
|
||||
"""Build the role-tagged `messages` array for InvokeGuardrailChecks.
|
||||
|
||||
INPUT scans the request messages, OUTPUT scans the model response as an
|
||||
``assistant`` turn. Every non-empty text block of every message is scanned;
|
||||
messages exceeding the per-message content-block cap are split into multiple
|
||||
messages rather than truncated.
|
||||
|
||||
INPUT content is tagged ``user`` regardless of the caller-supplied role.
|
||||
Bedrock excludes ``system`` content from prompt-attack evaluation, so
|
||||
trusting a caller's ``system``/``developer`` label would let an injection
|
||||
avoid the promptAttack check. At the proxy every INPUT message is
|
||||
caller-controlled, so all of it is treated as untrusted user input, matching
|
||||
AWS guidance to tag untrusted content as user input.
|
||||
"""
|
||||
if source == "OUTPUT":
|
||||
# Reuse the ApplyGuardrail output extractor (single source of truth for
|
||||
# pulling assistant text out of a ModelResponse), then re-tag as an
|
||||
# assistant turn for the role-based InvokeGuardrailChecks payload.
|
||||
output_request = self._create_bedrock_output_content_request(response=response)
|
||||
output_texts = [
|
||||
text for item in output_request.get("content") or [] if (text := (item.get("text") or {}).get("text"))
|
||||
]
|
||||
return self._chunk_texts_into_checks_messages("assistant", output_texts)
|
||||
|
||||
return [
|
||||
checks_message
|
||||
for message in messages or []
|
||||
for checks_message in self._chunk_texts_into_checks_messages(
|
||||
"user",
|
||||
[block.text for block in self.get_content_items_for_message(message) or [] if block.text],
|
||||
)
|
||||
]
|
||||
|
||||
async def _make_invoke_guardrail_checks_request(
|
||||
self,
|
||||
source: Literal["INPUT", "OUTPUT"],
|
||||
messages: list[AllMessageValues] | None = None,
|
||||
response: litellm.ModelResponse | None = None,
|
||||
request_data: dict | None = None,
|
||||
logging_event_type: GuardrailEventHooks | None = None,
|
||||
) -> BedrockGuardrailResponse:
|
||||
"""Run the resource-less InvokeGuardrailChecks API and enforce thresholds.
|
||||
|
||||
Detect-only: the API returns scores, never rewritten content. We map scores
|
||||
to a block decision via the configured thresholds. On a pass we return an
|
||||
empty ``BedrockGuardrailResponse`` (downstream masking treats it as a no-op).
|
||||
"""
|
||||
start_time = datetime.now(timezone.utc)
|
||||
|
||||
checks_messages = self._build_invoke_guardrail_checks_messages(
|
||||
source=source, messages=messages, response=response
|
||||
)
|
||||
if not checks_messages:
|
||||
# Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does.
|
||||
return BedrockGuardrailResponse()
|
||||
|
||||
credentials, aws_region_name = self._load_credentials()
|
||||
body: dict[str, Any] = {"messages": checks_messages, "checks": self.checks}
|
||||
api_key: str | None = request_data.get("api_key") if request_data else None
|
||||
|
||||
prepared_request = self._prepare_request(
|
||||
credentials=credentials,
|
||||
data=body,
|
||||
optional_params=self.optional_params,
|
||||
aws_region_name=aws_region_name,
|
||||
api_key=api_key,
|
||||
request_path=_BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH,
|
||||
)
|
||||
verbose_proxy_logger.debug("Bedrock InvokeGuardrailChecks request url: %s", prepared_request.url)
|
||||
|
||||
event_type = logging_event_type or (
|
||||
GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call
|
||||
)
|
||||
|
||||
httpx_response = await self._sign_and_post(
|
||||
prepared_request=prepared_request,
|
||||
request_data=request_data,
|
||||
event_type=event_type,
|
||||
start_time=start_time,
|
||||
)
|
||||
|
||||
if httpx_response.status_code != 200:
|
||||
status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
|
||||
verbose_proxy_logger.error(
|
||||
"Bedrock InvokeGuardrailChecks: error response. Status %s: %s",
|
||||
httpx_response.status_code,
|
||||
detail_message,
|
||||
)
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": detail_message},
|
||||
request_data=request_data or {},
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now(timezone.utc).timestamp(),
|
||||
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
)
|
||||
raise HTTPException(status_code=status_code, detail=detail_message)
|
||||
|
||||
try:
|
||||
json_response = TypeAdapter(BedrockGuardrailChecksResponse).validate_python(httpx_response.json())
|
||||
except (ValidationError, ValueError) as e:
|
||||
verbose_proxy_logger.error("Bedrock InvokeGuardrailChecks: unparseable 200 response: %s", str(e))
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": str(e)},
|
||||
request_data=request_data or {},
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now(timezone.utc).timestamp(),
|
||||
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": "Bedrock InvokeGuardrailChecks returned an unexpected response shape"},
|
||||
) from e
|
||||
violations = self._collect_invoke_checks_violations(json_response)
|
||||
|
||||
# Log a copy with PII location offsets stripped: offsets + the (separately
|
||||
# logged) request messages would otherwise reconstruct the detected PII span.
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response),
|
||||
request_data=request_data or {},
|
||||
guardrail_status=self._get_invoke_checks_status(bool(violations)),
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=datetime.now(timezone.utc).timestamp(),
|
||||
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
|
||||
event_type=event_type,
|
||||
tracing_detail=self._build_invoke_checks_tracing_detail(violations) if violations else None,
|
||||
)
|
||||
|
||||
if violations:
|
||||
raise self._get_block_exception_for_checks(violations, request_data=request_data)
|
||||
|
||||
return BedrockGuardrailResponse()
|
||||
|
||||
def _collect_invoke_checks_violations(
|
||||
self, response: BedrockGuardrailChecksResponse | None
|
||||
) -> list[BedrockChecksViolation]:
|
||||
"""Return the check results whose score meets/exceeds the configured threshold.
|
||||
|
||||
Only checks present in the configured ``checks`` block are evaluated; a
|
||||
threshold of ``None`` makes that check detect-only (never contributes a
|
||||
violation). A truncated sensitiveInformation result counts as a violation
|
||||
(fail closed: omitted detections were never scored). Only the non-sensitive
|
||||
label (category/type) and the numeric score are kept -- never offsets or
|
||||
matched text.
|
||||
"""
|
||||
results: dict[str, Any] = dict((response or {}).get("results") or {})
|
||||
# (results key, score field, label field, threshold). PII uses
|
||||
# confidenceScore/type; the other two use severityScore/category.
|
||||
check_specs = [
|
||||
(
|
||||
"contentFilter",
|
||||
"severityScore",
|
||||
"category",
|
||||
self.content_filter_threshold,
|
||||
),
|
||||
("promptAttack", "severityScore", "category", self.prompt_attack_threshold),
|
||||
(
|
||||
"sensitiveInformation",
|
||||
"confidenceScore",
|
||||
"type",
|
||||
self.pii_confidence_threshold,
|
||||
),
|
||||
]
|
||||
|
||||
configured_checks = self.checks or {}
|
||||
violations: list[BedrockChecksViolation] = []
|
||||
for check_key, score_field, label_field, threshold in check_specs:
|
||||
if threshold is None or check_key not in configured_checks:
|
||||
continue
|
||||
check_result = results.get(check_key) or {}
|
||||
if check_key == "sensitiveInformation" and check_result.get("truncated"):
|
||||
violations.append({"check": check_key, "truncated": True})
|
||||
for entry in check_result.get("results") or []:
|
||||
score = entry.get(score_field)
|
||||
if isinstance(score, (int, float)) and float(score) >= threshold:
|
||||
violation: BedrockChecksViolation = (
|
||||
{"check": check_key, "category": entry.get("category"), "severityScore": float(score)}
|
||||
if score_field == "severityScore"
|
||||
else {"check": check_key, "type": entry.get("type"), "confidenceScore": float(score)}
|
||||
)
|
||||
violations.append(violation)
|
||||
return violations
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_invoke_checks_response_for_logging(
|
||||
response: BedrockGuardrailChecksResponse,
|
||||
) -> dict[str, Any]:
|
||||
"""Strip PII location offsets from a checks response before it is logged."""
|
||||
sanitized: dict[str, Any] = copy.deepcopy(dict(response))
|
||||
sensitive = (sanitized.get("results") or {}).get("sensitiveInformation") or {}
|
||||
for entry in sensitive.get("results") or []:
|
||||
if isinstance(entry, dict):
|
||||
for key in _BEDROCK_CHECKS_PII_LOCATION_KEYS:
|
||||
entry.pop(key, None)
|
||||
return sanitized
|
||||
|
||||
@staticmethod
|
||||
def _get_invoke_checks_status(over_threshold: bool) -> GuardrailStatus:
|
||||
return "guardrail_intervened" if over_threshold else "success"
|
||||
|
||||
@staticmethod
|
||||
def _build_invoke_checks_tracing_detail(
|
||||
violations: list[BedrockChecksViolation],
|
||||
) -> GuardrailTracingDetail:
|
||||
tracing_detail: GuardrailTracingDetail = {}
|
||||
categories = [
|
||||
label
|
||||
for label in (v.get("category") or v.get("type") for v in violations)
|
||||
if isinstance(label, str) and label
|
||||
]
|
||||
if categories:
|
||||
tracing_detail["violation_categories"] = categories
|
||||
tracing_detail["guardrail_action"] = "GUARDRAIL_INTERVENED" if violations else "NONE"
|
||||
return tracing_detail
|
||||
|
||||
def _get_block_exception_for_checks(
|
||||
self, violations: list[BedrockChecksViolation], request_data: dict | None = None
|
||||
) -> Union[HTTPException, ModifyResponseException]:
|
||||
"""Build the block exception for an over-threshold InvokeGuardrailChecks result.
|
||||
|
||||
Mirrors ``_get_http_exception_for_blocked_guardrail``'s return-type branching.
|
||||
The detail carries only non-sensitive labels + scores (no offsets / raw input).
|
||||
"""
|
||||
if self.disable_exception_on_block is True:
|
||||
_request_data = request_data or {}
|
||||
return ModifyResponseException(
|
||||
message="Violated guardrail policy",
|
||||
model=_request_data.get("model", "bedrock-guardrail"),
|
||||
request_data=_request_data,
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
return HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Violated guardrail policy",
|
||||
"bedrock_guardrail_checks": violations,
|
||||
},
|
||||
)
|
||||
|
||||
def _check_bedrock_response_for_exception(self, response) -> bool:
|
||||
"""
|
||||
Return True if the Bedrock ApplyGuardrail response indicates an exception.
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -38,6 +38,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
default_on=litellm_params.default_on,
|
||||
streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"),
|
||||
streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"),
|
||||
streaming_transform_mode=_get_config_value(litellm_params, optional_params, "streaming_transform_mode"),
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback)
|
||||
|
|
|
|||
|
|
@ -58,7 +58,7 @@ _HEADER_PRESENT_PLACEHOLDER = "[present]"
|
|||
|
||||
def _header_value_allowed(
|
||||
header_name: str,
|
||||
extra_allowlist: Optional[Set[str]] = None,
|
||||
extra_allowlist: Set[str] | None = None,
|
||||
) -> bool:
|
||||
"""Return True if this header's value may be forwarded (allowlist, including globs and extra_headers)."""
|
||||
lower = header_name.lower()
|
||||
|
|
@ -74,8 +74,8 @@ def _header_value_allowed(
|
|||
|
||||
def _sanitize_inbound_headers(
|
||||
headers: Any,
|
||||
extra_allowlist: Optional[Set[str]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
extra_allowlist: Set[str] | None = None,
|
||||
) -> Dict[str, str] | None:
|
||||
"""
|
||||
Sanitize inbound headers before passing them to a 3rd party guardrail service.
|
||||
|
||||
|
|
@ -105,8 +105,8 @@ def _sanitize_inbound_headers(
|
|||
def _extract_inbound_headers(
|
||||
request_data: dict,
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
extra_allowlist: Optional[Set[str]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
extra_allowlist: Set[str] | None = None,
|
||||
) -> Dict[str, str] | None:
|
||||
"""
|
||||
Extract inbound headers from available request context.
|
||||
|
||||
|
|
@ -172,15 +172,16 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
headers: Optional[Dict[str, Any]] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
additional_provider_specific_params: Optional[Dict[str, Any]] = None,
|
||||
headers: Dict[str, Any] | None = None,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
additional_provider_specific_params: Dict[str, Any] | None = None,
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
fail_on_error: Optional[bool] = True,
|
||||
extra_headers: Optional[list] = None,
|
||||
streaming_end_of_stream_only: Optional[bool] = None,
|
||||
streaming_sampling_rate: Optional[int] = None,
|
||||
fail_on_error: bool | None = True,
|
||||
extra_headers: list | None = None,
|
||||
streaming_end_of_stream_only: bool | None = None,
|
||||
streaming_sampling_rate: int | None = None,
|
||||
streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
|
|
@ -221,6 +222,13 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
raise ValueError(f"streaming_sampling_rate must be >= 1 (got {streaming_sampling_rate})")
|
||||
self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate
|
||||
|
||||
# Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook.
|
||||
# "block_only" (default) drops text rewrites on the streaming path;
|
||||
# "incremental_diff" emits them as synthetic deltas.
|
||||
self.streaming_transform_mode: Literal["block_only", "incremental_diff"] = (
|
||||
"block_only" if streaming_transform_mode is None else streaming_transform_mode
|
||||
)
|
||||
|
||||
# Set supported event hooks
|
||||
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
|
||||
|
||||
|
|
@ -280,7 +288,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
error: Exception,
|
||||
http_status_code: Optional[int] = None,
|
||||
http_status_code: int | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
status_suffix = f" http_status_code={http_status_code}" if http_status_code else ""
|
||||
verbose_proxy_logger.critical(
|
||||
|
|
@ -326,6 +334,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
return_inputs["tools"] = guardrail_response.tools
|
||||
elif tools:
|
||||
return_inputs["tools"] = tools
|
||||
if guardrail_response.stream_holdback_chars is not None:
|
||||
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
|
||||
return return_inputs
|
||||
|
||||
def _handle_guardrail_request_error(
|
||||
|
|
@ -479,7 +489,7 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj, is_unreachable=False)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[type["GuardrailConfigModel"]]:
|
||||
def get_config_model() -> type["GuardrailConfigModel"] | None:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIConfigModel,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
|
|||
|
||||
import copy
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Union
|
||||
from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Union
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -21,7 +21,13 @@ from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_fo
|
|||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CallTypes, CallTypesLiteral
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
CallTypesLiteral,
|
||||
Delta,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# Imported lazily at runtime (inside the streaming hook) to avoid a
|
||||
|
|
@ -34,7 +40,12 @@ A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message)
|
|||
GUARDRAIL_NAME = "unified_llm_guardrails"
|
||||
|
||||
|
||||
def _get_a2a_request_id(responses_so_far: List[Any], request_data: dict) -> Optional[str]:
|
||||
class _StreamTerminated(Exception):
|
||||
"""Internal signal that the incremental transform stream has already emitted
|
||||
its terminal chunks (block message or in-stream error) and must stop."""
|
||||
|
||||
|
||||
def _get_a2a_request_id(responses_so_far: List[Any], request_data: dict) -> str | None:
|
||||
"""Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting."""
|
||||
for item in responses_so_far:
|
||||
if isinstance(item, dict) and "id" in item:
|
||||
|
|
@ -216,7 +227,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
|
||||
verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response)
|
||||
|
||||
call_type: Optional[CallTypesLiteral] = None
|
||||
call_type: CallTypesLiteral | None = None
|
||||
if user_api_key_dict.request_route is not None:
|
||||
call_types = get_call_types_for_route(user_api_key_dict.request_route)
|
||||
if call_types is not None and len(call_types) > 0: # type: ignore
|
||||
|
|
@ -292,6 +303,498 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
for chunk in block_chunks:
|
||||
yield chunk
|
||||
|
||||
@staticmethod
|
||||
def _resolve_transform_call_type(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
mappings: dict,
|
||||
) -> str | None:
|
||||
"""Resolve the call type for the incremental_diff path, or None if the
|
||||
route is unresolvable / unsupported.
|
||||
|
||||
Incremental transformation needs a route we can resolve before the first
|
||||
chunk and a handler that supports the streaming text-diff protocol (v1:
|
||||
the OpenAI chat completions handler only). Returning None makes the caller
|
||||
fall back to block_only.
|
||||
"""
|
||||
from litellm.llms.openai.chat.guardrail_translation.handler import (
|
||||
OpenAIChatCompletionsHandler,
|
||||
)
|
||||
|
||||
if user_api_key_dict.request_route is None:
|
||||
return None
|
||||
call_types = get_call_types_for_route(user_api_key_dict.request_route)
|
||||
if not call_types:
|
||||
return None
|
||||
call_type = call_types[0].value
|
||||
try:
|
||||
mapped = CallTypes(call_type)
|
||||
except ValueError:
|
||||
return None
|
||||
handler_cls = mappings.get(mapped)
|
||||
if handler_cls is None or not issubclass(handler_cls, OpenAIChatCompletionsHandler):
|
||||
return None
|
||||
return call_type
|
||||
|
||||
async def _emit_streaming_http_error(
|
||||
self,
|
||||
exc: HTTPException,
|
||||
call_type: str | None,
|
||||
responses_so_far: list[Any],
|
||||
request_data: dict,
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Surface a mid-stream HTTPException. For A2A (NDJSON) call types the
|
||||
response has already started, so emit an in-stream JSON-RPC error chunk;
|
||||
otherwise re-raise so the proxy can report it.
|
||||
"""
|
||||
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
|
||||
request_id = _get_a2a_request_id(responses_so_far, request_data)
|
||||
detail = exc.detail if isinstance(exc.detail, dict) else {"message": str(exc.detail)}
|
||||
error_chunk = (
|
||||
json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"error": {
|
||||
"code": -32603,
|
||||
"message": detail.get("error", detail.get("message", str(exc.detail))),
|
||||
"data": {k: v for k, v in detail.items() if k not in ("error", "message")},
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
yield error_chunk
|
||||
return
|
||||
raise exc
|
||||
|
||||
def _build_transform_chunk(
|
||||
self,
|
||||
*,
|
||||
reference_chunk: Any,
|
||||
mutated_text_per_choice: dict[int, str],
|
||||
emitted_text_per_choice: dict[int, str],
|
||||
holdback_per_choice: dict[int, int],
|
||||
finish_reason_per_choice: dict[int, str | None],
|
||||
is_final: bool,
|
||||
) -> ModelResponseStream | None:
|
||||
"""Build the synthetic chunk carrying the newly-guardrailed deltas.
|
||||
|
||||
For each choice, the new delta is the mutated accumulated text past what
|
||||
has already been emitted, minus a trailing holdback (forced to 0 on the
|
||||
final flush). ``emitted_text_per_choice`` holds the exact bytes already
|
||||
sent per choice and is extended in place. Returns None when there is no
|
||||
text to emit (e.g. a tool-call-only turn) or nothing new and this is not
|
||||
the final chunk.
|
||||
|
||||
Raises HTTPException(400, stream_transform_underflow) when the guardrail's
|
||||
transform is not a forward extension of what has already been streamed
|
||||
(shorter than, or rewrites, the already-sent prefix), since emitted bytes
|
||||
cannot be retracted. This makes the framework fail closed rather than
|
||||
silently leave un-transformed text on the wire; a guardrail that needs to
|
||||
rewrite recent output must withhold it first via ``stream_holdback_chars``.
|
||||
"""
|
||||
if not mutated_text_per_choice:
|
||||
# Fix #4 — on the final flush a deferred finish_reason (from a mixed
|
||||
# content+tool_calls chunk whose passthrough suppressed it) still
|
||||
# needs to reach the client, even if the guardrail returned no text
|
||||
# to emit. Build a terminator chunk carrying finish_reason per choice.
|
||||
if is_final and finish_reason_per_choice:
|
||||
terminator_choices: list[StreamingChoices] = []
|
||||
for choice_idx, finish_reason in finish_reason_per_choice.items():
|
||||
if finish_reason is None:
|
||||
continue
|
||||
terminator_choices.append(
|
||||
StreamingChoices(
|
||||
index=choice_idx,
|
||||
delta=Delta(content="", role=None, tool_calls=None),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
)
|
||||
if terminator_choices:
|
||||
return ModelResponseStream(
|
||||
id=getattr(reference_chunk, "id", None),
|
||||
created=getattr(reference_chunk, "created", None),
|
||||
model=getattr(reference_chunk, "model", None),
|
||||
choices=terminator_choices,
|
||||
)
|
||||
return None
|
||||
|
||||
deltas: dict[int, str] = {}
|
||||
for choice_idx, text in mutated_text_per_choice.items():
|
||||
already = emitted_text_per_choice.get(choice_idx, "")
|
||||
if not text.startswith(already):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "stream_transform_underflow",
|
||||
"message": (
|
||||
f"Guardrail streaming transform for choice {choice_idx} is not a forward "
|
||||
f"extension of the {len(already)} chars already streamed to the client "
|
||||
"(it is shorter than, or rewrites, the emitted prefix); emitted bytes "
|
||||
"cannot be retracted. Withhold recent output via stream_holdback_chars "
|
||||
"before rewriting it."
|
||||
),
|
||||
},
|
||||
)
|
||||
holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0))
|
||||
end = max(len(already), len(text) - holdback)
|
||||
deltas[choice_idx] = text[len(already) : end]
|
||||
|
||||
# Iterate the mutated choices (not just those in reference_chunk) so a
|
||||
# choice with pending text is never dropped for n > 1. finish_reason is
|
||||
# taken per choice from the accumulated map (a choice can finish in an
|
||||
# earlier chunk than the stream's last one); tool_calls are dropped since
|
||||
# v1 does not transform streamed tool calls (they pass through raw).
|
||||
synthetic_choices: list[StreamingChoices] = []
|
||||
for choice_idx in mutated_text_per_choice:
|
||||
delta_text = deltas.get(choice_idx, "")
|
||||
finish_reason = finish_reason_per_choice.get(choice_idx) if is_final else None
|
||||
# Skip a choice with nothing to say: no new content and no
|
||||
# finish_reason to deliver. This avoids emitting an empty delta for an
|
||||
# already-finished choice (e.g. one that terminated via a passed-through
|
||||
# tool-call chunk, which already carried its own finish_reason).
|
||||
if not delta_text and finish_reason is None:
|
||||
continue
|
||||
# role="assistant" on this choice's first emitted delta only.
|
||||
role = "assistant" if not emitted_text_per_choice.get(choice_idx) else None
|
||||
synthetic_choices.append(
|
||||
StreamingChoices(
|
||||
index=choice_idx,
|
||||
delta=Delta(content=delta_text, role=role, tool_calls=None),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
)
|
||||
|
||||
if not synthetic_choices:
|
||||
return None
|
||||
|
||||
for choice_idx in mutated_text_per_choice:
|
||||
emitted_text_per_choice[choice_idx] = emitted_text_per_choice.get(choice_idx, "") + deltas.get(
|
||||
choice_idx, ""
|
||||
)
|
||||
|
||||
return ModelResponseStream(
|
||||
id=getattr(reference_chunk, "id", None),
|
||||
created=getattr(reference_chunk, "created", None),
|
||||
model=getattr(reference_chunk, "model", None),
|
||||
choices=synthetic_choices,
|
||||
)
|
||||
|
||||
async def _emit_transform_round(
|
||||
self,
|
||||
*,
|
||||
endpoint_translation: Any,
|
||||
guardrail_to_apply: CustomGuardrail,
|
||||
request_data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: str,
|
||||
reference_chunk: Any,
|
||||
responses_so_far: list[Any],
|
||||
responses_yielded: list[Any],
|
||||
emitted_text_per_choice: dict[int, str],
|
||||
finish_reason_per_choice: dict[int, str | None],
|
||||
is_final: bool,
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Run one guardrail processing round and emit the resulting diff chunk.
|
||||
|
||||
Raises ``_StreamTerminated`` (after emitting the terminal block message or
|
||||
in-stream error) when the guardrail blocks or an underflow occurs.
|
||||
"""
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
StreamTransformSink,
|
||||
)
|
||||
|
||||
sink = StreamTransformSink()
|
||||
try:
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
stream_transform_sink=sink,
|
||||
)
|
||||
synthetic = self._build_transform_chunk(
|
||||
reference_chunk=reference_chunk,
|
||||
mutated_text_per_choice=sink.mutated_text_per_choice,
|
||||
emitted_text_per_choice=emitted_text_per_choice,
|
||||
holdback_per_choice=sink.holdback_per_choice,
|
||||
finish_reason_per_choice=finish_reason_per_choice,
|
||||
is_final=is_final,
|
||||
)
|
||||
except ModifyResponseException as e:
|
||||
if e.original_response is None:
|
||||
e.original_response = responses_so_far
|
||||
async for block_chunk in self._handle_streaming_block(
|
||||
e,
|
||||
endpoint_translation,
|
||||
stream_started=bool(responses_yielded),
|
||||
responses_so_far=responses_yielded,
|
||||
):
|
||||
yield block_chunk
|
||||
raise _StreamTerminated()
|
||||
except HTTPException as e:
|
||||
async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data):
|
||||
yield error_item
|
||||
raise _StreamTerminated()
|
||||
|
||||
if synthetic is not None:
|
||||
responses_yielded.append(synthetic)
|
||||
yield synthetic
|
||||
|
||||
async def _run_incremental_transform_stream(
|
||||
self,
|
||||
*,
|
||||
guardrail_to_apply: CustomGuardrail,
|
||||
response: Any,
|
||||
request_data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: str,
|
||||
sampling_rate: int,
|
||||
end_of_stream_only: bool,
|
||||
mappings: dict,
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Emit guardrail text transformations as new deltas on the stream.
|
||||
|
||||
Raw chunks are withheld and accumulated; on each sampled processing round
|
||||
(and once at end of stream) the guardrailed accumulated text is diffed
|
||||
against what has already been emitted and the new portion is sent as a
|
||||
synthetic chunk. A BLOCK terminates the stream via the shared block
|
||||
handler; an underflow surfaces as an HTTPException.
|
||||
"""
|
||||
endpoint_translation = mappings[CallTypes(call_type)]()
|
||||
responses_so_far: list[Any] = []
|
||||
responses_yielded: list[Any] = []
|
||||
emitted_text_per_choice: dict[int, str] = {}
|
||||
finish_reason_per_choice: dict[int, str | None] = {}
|
||||
chunk_counter = 0
|
||||
last_chunk: Any | None = None
|
||||
|
||||
def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]:
|
||||
return self._emit_transform_round(
|
||||
endpoint_translation=endpoint_translation,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
request_data=request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
reference_chunk=reference_chunk,
|
||||
responses_so_far=responses_so_far,
|
||||
responses_yielded=responses_yielded,
|
||||
emitted_text_per_choice=emitted_text_per_choice,
|
||||
finish_reason_per_choice=finish_reason_per_choice,
|
||||
is_final=is_final,
|
||||
)
|
||||
|
||||
saw_tool_calls = False
|
||||
saw_text_content = False
|
||||
|
||||
try:
|
||||
async for item in response:
|
||||
# v1 transforms only text. A chunk carrying tool_calls is passed
|
||||
# through raw so function-calling turns are not dropped, but ONLY
|
||||
# its tool-call fields are forwarded: content is stripped so any
|
||||
# response text (in the same delta, or in another choice of an n>1
|
||||
# chunk) can never bypass the transform. The original chunk is kept
|
||||
# in responses_so_far so its text is still accumulated + redacted +
|
||||
# emitted as synthetic deltas, and so the guardrail inspects the
|
||||
# assembled tool calls at end of stream (see the block inspection
|
||||
# below), matching block_only. finish_reason rides on the raw
|
||||
# tool-only chunk, so it is not recorded for the text flush.
|
||||
if self._chunk_has_tool_calls(item):
|
||||
saw_tool_calls = True
|
||||
responses_so_far.append(item)
|
||||
last_chunk = item
|
||||
# Fix #3 — flush accumulated text BEFORE the tool-call
|
||||
# passthrough. Without this, a stream of text chunks that
|
||||
# hasn't yet hit a sampled round can be trailed by a
|
||||
# tool-call chunk carrying finish_reason="tool_calls"; an
|
||||
# SSE-compliant client stops reading at that finish_reason
|
||||
# and drops the end-of-stream text flush that would follow.
|
||||
if saw_text_content:
|
||||
async for out in _round(item, is_final=False):
|
||||
yield out
|
||||
# Fix #1 — pass finish_reason_per_choice into the
|
||||
# passthrough so a mixed content+tool_call chunk defers its
|
||||
# finish_reason to the final text terminator (see the
|
||||
# _tool_call_passthrough_chunk docstring).
|
||||
tool_only = self._tool_call_passthrough_chunk(
|
||||
item, finish_reason_per_choice=finish_reason_per_choice
|
||||
)
|
||||
responses_yielded.append(tool_only)
|
||||
yield tool_only
|
||||
continue
|
||||
|
||||
chunk_counter += 1
|
||||
responses_so_far.append(item)
|
||||
last_chunk = item
|
||||
self._record_finish_reasons(item, finish_reason_per_choice)
|
||||
if self._chunk_carries_text(item):
|
||||
saw_text_content = True
|
||||
# Skip the sampled round for a terminal chunk: the end-of-stream
|
||||
# flush below processes it once with holdback forced to 0, so a
|
||||
# sampled round here would guardrail the same content twice.
|
||||
if (
|
||||
not end_of_stream_only
|
||||
and not self._chunk_has_finish_reason(item)
|
||||
and chunk_counter % sampling_rate == 0
|
||||
):
|
||||
async for out in _round(item, is_final=False):
|
||||
yield out
|
||||
|
||||
# v1 does not transform streamed tool calls, but they must still go
|
||||
# through the guardrail's block decision. Run the block_only inspection
|
||||
# over the full assembled response so tool calls cannot bypass it.
|
||||
#
|
||||
# Pass a deep copy of responses_so_far — the block path routes through
|
||||
# ``_process_streaming_block_only`` which mutates ``delta.content``
|
||||
# in-place on the chunk objects it receives. For an n>1 chunk carrying
|
||||
# text on one choice and tool_calls (with finish_reason) on another,
|
||||
# ``has_stream_ended`` reads ``choices[0]`` alone and can miss the
|
||||
# terminal signal, letting the block path rewrite the raw accumulator.
|
||||
# The subsequent final ``_round`` would then re-read the already-mutated
|
||||
# text, producing double-application for a non-idempotent guardrail or a
|
||||
# ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow
|
||||
# list copy wouldn't help — the mutation is on the chunk objects
|
||||
# themselves — so we deepcopy.
|
||||
if saw_tool_calls:
|
||||
async for out in self._inspect_full_response_for_block(
|
||||
endpoint_translation=endpoint_translation,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
request_data=request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
responses_so_far=copy.deepcopy(responses_so_far),
|
||||
responses_yielded=responses_yielded,
|
||||
):
|
||||
yield out
|
||||
|
||||
if last_chunk is not None:
|
||||
async for out in _round(last_chunk, is_final=True):
|
||||
yield out
|
||||
except _StreamTerminated:
|
||||
return
|
||||
|
||||
async def _inspect_full_response_for_block(
|
||||
self,
|
||||
*,
|
||||
endpoint_translation: Any,
|
||||
guardrail_to_apply: CustomGuardrail,
|
||||
request_data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
responses_so_far: list[Any],
|
||||
responses_yielded: list[Any],
|
||||
) -> AsyncGenerator[Any, None]:
|
||||
"""Run the block-only guardrail inspection over the full assembled
|
||||
response (text + tool calls) so nothing bypasses the block decision.
|
||||
|
||||
The guardrail's returned transforms are discarded here (v1 does not
|
||||
transform tool calls); only its block decision matters. A block is
|
||||
surfaced the same way as elsewhere: ModifyResponseException terminates the
|
||||
stream via the shared block handler; a GenericGuardrailAPI block raises and
|
||||
propagates, matching block_only.
|
||||
"""
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
|
||||
try:
|
||||
await endpoint_translation.process_output_streaming_response(
|
||||
responses_so_far=responses_so_far,
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
stream_transform_sink=None,
|
||||
)
|
||||
except ModifyResponseException as e:
|
||||
if e.original_response is None:
|
||||
e.original_response = responses_so_far
|
||||
async for block_chunk in self._handle_streaming_block(
|
||||
e,
|
||||
endpoint_translation,
|
||||
stream_started=bool(responses_yielded),
|
||||
responses_so_far=responses_yielded,
|
||||
):
|
||||
yield block_chunk
|
||||
raise _StreamTerminated()
|
||||
|
||||
@staticmethod
|
||||
def _chunk_has_tool_calls(item: Any) -> bool:
|
||||
for choice in getattr(item, "choices", None) or []:
|
||||
delta = getattr(choice, "delta", None)
|
||||
if getattr(delta, "tool_calls", None):
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _chunk_carries_text(item: Any) -> bool:
|
||||
"""True if any choice in this chunk has non-empty string ``delta.content``."""
|
||||
for choice in getattr(item, "choices", None) or []:
|
||||
delta = getattr(choice, "delta", None)
|
||||
content = getattr(delta, "content", None)
|
||||
if isinstance(content, str) and content != "":
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _tool_call_passthrough_chunk(
|
||||
item: Any,
|
||||
finish_reason_per_choice: "dict[int, str | None] | None" = None,
|
||||
) -> ModelResponseStream:
|
||||
"""Copy of a chunk carrying tool calls with all text content stripped.
|
||||
|
||||
Only tool_calls, role and finish_reason are forwarded; content is set to
|
||||
None so response text can never be delivered raw (it flows through the
|
||||
transform instead). Applies per choice so an n>1 chunk mixing a text
|
||||
choice and a tool-call choice does not leak the text choice.
|
||||
|
||||
For a choice that carries BOTH text content AND tool_calls, ``finish_reason``
|
||||
is suppressed on the passthrough and recorded on
|
||||
``finish_reason_per_choice`` (when provided) so the final synthetic text
|
||||
chunk delivers it. Emitting the passthrough's ``finish_reason`` before the
|
||||
text flush would let a spec-compliant SSE client stop reading at
|
||||
``finish_reason`` and silently drop the guardrailed text, defeating the
|
||||
redaction purpose.
|
||||
"""
|
||||
synthetic_choices: list[StreamingChoices] = []
|
||||
for choice in getattr(item, "choices", None) or []:
|
||||
delta = getattr(choice, "delta", None)
|
||||
idx = getattr(choice, "index", 0) or 0
|
||||
original_finish = getattr(choice, "finish_reason", None)
|
||||
has_text = isinstance(getattr(delta, "content", None), str) and getattr(delta, "content", "") != ""
|
||||
if has_text and original_finish is not None and finish_reason_per_choice is not None:
|
||||
finish_reason_per_choice[idx] = original_finish
|
||||
passthrough_finish: str | None = None
|
||||
else:
|
||||
passthrough_finish = original_finish
|
||||
synthetic_choices.append(
|
||||
StreamingChoices(
|
||||
index=idx,
|
||||
delta=Delta(
|
||||
content=None,
|
||||
role=getattr(delta, "role", None),
|
||||
tool_calls=getattr(delta, "tool_calls", None),
|
||||
),
|
||||
finish_reason=passthrough_finish,
|
||||
)
|
||||
)
|
||||
return ModelResponseStream(
|
||||
id=getattr(item, "id", None),
|
||||
created=getattr(item, "created", None),
|
||||
model=getattr(item, "model", None),
|
||||
choices=synthetic_choices,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None:
|
||||
for choice in getattr(item, "choices", None) or []:
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
if finish_reason is not None:
|
||||
finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason
|
||||
|
||||
@staticmethod
|
||||
def _chunk_has_finish_reason(item: Any) -> bool:
|
||||
choices = getattr(item, "choices", None) or []
|
||||
return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -334,6 +837,10 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
sampling_rate = _streaming_flag("streaming_sampling_rate", 5)
|
||||
# Only apply the guardrail at end of stream (not per chunk).
|
||||
end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False)
|
||||
# "block_only" (default) drops guardrail text rewrites on the streaming
|
||||
# path; "incremental_diff" emits them as synthetic deltas (see
|
||||
# _run_incremental_transform_stream).
|
||||
streaming_transform_mode = _streaming_flag("streaming_transform_mode", "block_only")
|
||||
# Withhold every chunk until end-of-stream moderation passes, then
|
||||
# release the original chunks (clean) or only the block message
|
||||
# (blocked) -- moderating the whole response *before* any content
|
||||
|
|
@ -380,6 +887,35 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
if endpoint_guardrail_translation_mappings is None:
|
||||
endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
|
||||
|
||||
# Streaming text transformation (incremental_diff) diverges enough from the
|
||||
# block_only path that it runs as its own iterator. It requires a route we
|
||||
# can resolve up front to an OpenAI-chat handler (the only supported v1
|
||||
# surface); anything else falls back to the block_only behavior below.
|
||||
if streaming_transform_mode == "incremental_diff":
|
||||
transform_call_type = self._resolve_transform_call_type(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
mappings=endpoint_guardrail_translation_mappings,
|
||||
)
|
||||
if transform_call_type is not None:
|
||||
async for transformed_item in self._run_incremental_transform_stream(
|
||||
guardrail_to_apply=guardrail_to_apply,
|
||||
response=response,
|
||||
request_data=request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=transform_call_type,
|
||||
sampling_rate=sampling_rate,
|
||||
end_of_stream_only=end_of_stream_only,
|
||||
mappings=endpoint_guardrail_translation_mappings,
|
||||
):
|
||||
yield transformed_item
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
"UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported "
|
||||
"for the OpenAI chat completions streaming path with a resolvable request route; "
|
||||
"falling back to block_only for %s",
|
||||
getattr(guardrail_to_apply, "guardrail_name", None),
|
||||
)
|
||||
|
||||
# Infer call type from first chunk
|
||||
call_type = None
|
||||
chunk_counter = 0
|
||||
|
|
|
|||
|
|
@ -16,6 +16,10 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
|
|||
event_hook=litellm_params.mode,
|
||||
guardrailIdentifier=litellm_params.guardrailIdentifier,
|
||||
guardrailVersion=litellm_params.guardrailVersion,
|
||||
checks=litellm_params.checks,
|
||||
content_filter_threshold=litellm_params.content_filter_threshold,
|
||||
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
|
||||
pii_confidence_threshold=litellm_params.pii_confidence_threshold,
|
||||
default_on=litellm_params.default_on,
|
||||
disable_exception_on_block=litellm_params.disable_exception_on_block,
|
||||
mask_request_content=litellm_params.mask_request_content,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import Span
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -76,6 +76,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
|
||||
current_cost=_current_spend,
|
||||
max_budget=_current_model_budget_info.max_budget,
|
||||
entity_type=Litellm_EntityType.KEY.value,
|
||||
entity_id=user_api_key_dict.token,
|
||||
)
|
||||
|
||||
return True
|
||||
|
|
@ -140,6 +142,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}",
|
||||
current_cost=_current_spend,
|
||||
max_budget=_current_model_budget_info.max_budget,
|
||||
entity_type=Litellm_EntityType.END_USER.value,
|
||||
entity_id=end_user_id,
|
||||
)
|
||||
|
||||
return True
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue