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

# Conflicts:
#	tests/test_litellm/router_strategy/test_complexity_router.py
This commit is contained in:
moe-berri 2026-09-08 15:49:53 -07:00
commit c7b80f1966
145 changed files with 8234 additions and 475 deletions

View file

@ -1,5 +1,11 @@
#!/bin/sh
# stale samples from a previous container incarnation would be summed into the aggregate
if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then
mkdir -p "$PROMETHEUS_MULTIPROC_DIR"
rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"

View file

@ -142,6 +142,40 @@ class CheckBatchCost:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None:
org_id = getattr(job, "org_id", None)
if org_id:
return org_id
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
)
key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
if key_org_id:
return key_org_id
except Exception as e:
verbose_proxy_logger.error(
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
f"still trying the team's: {e}"
)
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
)
return getattr(team_row, "organization_id", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
return None
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
) -> dict[str, object]:
@ -153,6 +187,10 @@ class CheckBatchCost:
user_api_key_alias; when it has no alias, or the key has since been rotated or
deleted, the field keeps the creating user's alias that _get_user_info filled in,
because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id must be resolved here too: the spend update writer reads it
off this metadata to increment organization spend, so leaving it out silently
drops batch cost from org accounting for keys and teams that belong to one.
"""
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
@ -172,6 +210,9 @@ class CheckBatchCost:
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
org_id: Final = await self._get_org_id(job, batch_id)
if org_id is not None:
metadata["user_api_key_org_id"] = org_id
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
@ -641,7 +682,7 @@ class CheckBatchCost:
from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -805,6 +846,7 @@ class CheckBatchCost:
function_id=str(uuid.uuid4()),
)
deployment_api_base: Final = deployment_info.litellm_params.api_base
logging_obj.update_environment_variables(
litellm_params={
# set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@ -813,9 +855,17 @@ class CheckBatchCost:
"user-agent": CHECK_BATCH_COST_USER_AGENT,
}
},
"metadata": await self._build_creator_attribution_metadata(job, batch_id),
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
"metadata": {
**(await self._build_creator_attribution_metadata(job, batch_id)),
# spend logs read the deployment identity off these metadata keys, so
# without them the batch cost row carries no model_id or model_group
"model_info": {"id": model_id},
"model_group": deployment_info.model_name,
},
},
optional_params={},
custom_llm_provider=str(llm_provider) if llm_provider else None,
)
if not await self._claim_job_for_costing(job):
@ -833,6 +883,8 @@ class CheckBatchCost:
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)

View file

@ -280,6 +280,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
if user_api_key_dict.org_id:
return user_api_key_dict.org_id
if not user_api_key_dict.team_id:
return None
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
try:
team: Final = await get_team_object(
team_id=user_api_key_dict.team_id,
prisma_client=self.prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
return team.organization_id
except Exception as e:
verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
return None
async def store_unified_object_id(
self,
unified_object_id: str,
@ -352,6 +373,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_purpose": file_purpose,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"org_id": await self._resolve_creator_org_id(user_api_key_dict),
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
**attribution_columns,

View file

@ -152,6 +152,13 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
@ -189,6 +196,11 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.metricsServer.enabled }}
- name: metrics
containerPort: {{ .Values.metricsServer.port }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}

View file

@ -0,0 +1,17 @@
{{- if .Values.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.fullname" . }}-metrics
labels:
{{- include "litellm.labels" . | nindent 4 }}
spec:
type: ClusterIP
ports:
- port: {{ .Values.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.selectorLabels" . | nindent 4 }}
{{- end }}

View file

@ -26,7 +26,7 @@ spec:
{{- toYaml .namespaceSelector.matchNames | nindent 4 }}
{{- end }}
endpoints:
- port: http
- port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
path: /metrics/
interval: {{ .interval }}
scrapeTimeout: {{ .scrapeTimeout }}

View file

@ -0,0 +1,106 @@
suite: separate metrics server
templates:
- configmap-litellm.yaml
- deployment.yaml
- service.yaml
- service-metrics.yaml
- servicemonitor.yaml
tests:
- it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default
asserts:
- notContains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
any: true
template: deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- hasDocuments:
count: 0
template: service-metrics.yaml
- it: should scrape the proxy port when the metrics server is disabled
template: servicemonitor.yaml
set:
serviceMonitor.enabled: true
asserts:
- equal:
path: spec.endpoints[0].port
value: http
- it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor
set:
metricsServer.enabled: true
metricsServer.port: 4101
serviceMonitor.enabled: true
service.type: LoadBalancer
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "4101"
template: deployment.yaml
- contains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
containerPort: 4101
protocol: TCP
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-metrics
template: service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
template: service-metrics.yaml
- equal:
path: spec.endpoints[0].port
value: metrics
template: servicemonitor.yaml
- equal:
path: spec.endpoints[0].path
value: /metrics/
template: servicemonitor.yaml
- it: should reject a metrics port equal to the proxy port
template: deployment.yaml
set:
metricsServer.enabled: true
metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: metricsServer.port must differ from service.port

View file

@ -180,6 +180,16 @@ proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT)
# so a scrape never runs on an inference worker. Adds a `metrics` port to the
# container and a dedicated ClusterIP `<release>-metrics` Service, and the
# ServiceMonitor scrapes it instead of the proxy port. The separate port has
# no virtual-key auth: keep it off public ingress. Needs the proxy image
# v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an

View file

@ -441,3 +441,5 @@ ImplementationSpecific
{{- .pathType -}}
{{- end -}}
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}

View file

@ -64,14 +64,25 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
{{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
{{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
{{- end }}
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -97,16 +108,54 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: metrics
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- {{ .Values.gateway.metricsServer.port | quote }}
env:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
ports:
- name: metrics
containerPort: {{ .Values.gateway.metricsServer.port }}
protocol: TCP
volumeMounts:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
readinessProbe:
tcpSocket: { port: metrics }
periodSeconds: 10
livenessProbe:
tcpSocket: { port: metrics }
periodSeconds: 15
failureThreshold: 6
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

View file

@ -0,0 +1,18 @@
{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.gateway.fullname" . }}-metrics
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
spec:
type: ClusterIP
ports:
- port: {{ .Values.gateway.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
{{- end }}

View file

@ -0,0 +1,148 @@
suite: test gateway metrics sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/service.yaml
- gateway/service-metrics.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, volume, env or service port when the metrics server is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
any: true
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- hasDocuments:
count: 0
template: gateway/service-metrics.yaml
- it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4101
gateway.service.type: LoadBalancer
gateway.image.tag: v1.101.0
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.101.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- "4101"
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].env
value:
- name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].ports
value:
- name: metrics
containerPort: 4101
protocol: TCP
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].volumeMounts
value:
- name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 50m
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
emptyDir: {}
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: gateway/service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway-metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: gateway/service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
template: gateway/service-metrics.yaml
- it: rejects a metrics port equal to the gateway port
template: gateway/deployment.yaml
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: gateway.metricsServer.port must differ from the gateway port 4000

View file

@ -268,6 +268,22 @@ gateway:
config:
create: true
proxy_config: {}
# Serve Prometheus /metrics from a `metrics` sidecar container (same image,
# `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the
# workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a
# scrape never runs on an inference worker. Adds a `metrics` port to the pod
# and a dedicated ClusterIP `<gateway>-metrics` Service; point your scrape
# config at it. The port has no virtual-key auth: keep it off public ingress.
# Needs the gateway image v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
requests:
cpu: 50m
memory: 128Mi
limits:
memory: 512Mi
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion

View file

@ -0,0 +1,4 @@
-- Add org_id column to LiteLLM_ManagedObjectTable
-- Snapshots the creating key's organization at submission time, like team_id,
-- so CheckBatchCost can bill organization spend hours later without re-resolving
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT;

View file

@ -1036,6 +1036,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
created_at DateTime @default(now())
created_by String?
team_id String?
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
updated_at DateTime @updatedAt

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.94"
version = "0.4.95"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.94"
version = "0.4.95"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -47,6 +47,7 @@ from typing import (
)
from litellm.types.integrations.datadog import DatadogInitParams
from litellm.types.integrations.newrelic import NewRelicInitParams
from litellm.litellm_core_utils.core_helpers import drop_params_env_flag
from litellm._logging import (
set_verbose,
_turn_on_debug,
@ -238,7 +239,7 @@ token: Optional[str] = (
)
telemetry = True
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False))
drop_params = drop_params_env_flag(os.environ, verbose_logger)
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
use_chat_completions_url_for_anthropic_messages: bool = bool(
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)

View file

@ -10,7 +10,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.types.llms.openai import Batch
from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import token_counter
@ -23,6 +23,8 @@ class BatchCostUsageResult:
models: list[str]
successful_requests: int
failed_requests: int
prompt_cost: float = 0.0
completion_cost: float = 0.0
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
@ -151,7 +153,8 @@ class _LineOutcome(Enum):
@dataclass(frozen=True, slots=True)
class _BatchOutputLineStats:
cost: float
prompt_cost: float
completion_cost: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
@ -214,15 +217,16 @@ def _compute_output_line_stats(
raw_model: Final = response_body.get("model")
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
completion_details: Final = usage.completion_tokens_details
line_prompt_cost, line_completion_cost = _output_line_cost(
usage=usage,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
response_model=response_model,
model_info=model_info,
)
return _BatchOutputLineStats(
cost=_output_line_cost(
response_body=response_body,
usage=usage,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
response_model=response_model,
model_info=model_info,
),
prompt_cost=line_prompt_cost,
completion_cost=line_completion_cost,
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
@ -234,31 +238,24 @@ def _compute_output_line_stats(
def _output_line_cost(
response_body: Mapping[str, object],
usage: Usage,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None,
response_model: str | None,
model_info: ModelInfo | None,
) -> float:
) -> tuple[float, float]:
"""(prompt_cost, completion_cost) for one output line, priced at batch rates."""
from litellm.cost_calculator import batch_cost_calculator
if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
return litellm.completion_cost(
completion_response=response_body,
custom_llm_provider=custom_llm_provider,
call_type=CallTypes.aretrieve_batch.value,
)
cost_model: Final = (
model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
)
prompt_cost, completion_cost = batch_cost_calculator(
return batch_cost_calculator(
usage=usage,
model=cost_model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
)
return prompt_cost + completion_cost
def _aggregate_batch_cost_usage_models(
@ -291,7 +288,9 @@ def _aggregate_batch_cost_usage_models(
**cache_token_params,
)
batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
total_cost: Final = sum((stats.cost for stats in line_stats), 0.0)
total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0)
total_completion_cost: Final = sum((stats.completion_cost for stats in line_stats), 0.0)
total_cost: Final = total_prompt_cost + total_completion_cost
verbose_logger.debug(
"batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
total_cost,
@ -306,6 +305,8 @@ def _aggregate_batch_cost_usage_models(
models=batch_models,
successful_requests=successful_requests,
failed_requests=failed_requests,
prompt_cost=total_prompt_cost,
completion_cost=total_completion_cost,
)
@ -330,7 +331,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
"""
from litellm.cost_calculator import batch_cost_calculator
total_cost = 0.0
total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below
total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below
total_tokens = 0
prompt_tokens = 0
completion_tokens = 0
@ -362,7 +364,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
model=actual_model_name,
custom_llm_provider="vertex_ai",
)
total_cost += p_cost + c_cost
total_prompt_cost += p_cost
total_completion_cost += c_cost
except Exception as e:
verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e))
@ -370,6 +373,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
completion_tokens += _completion
total_tokens += _total
total_cost: Final = total_prompt_cost + total_completion_cost
verbose_logger.info(
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
total_cost,
@ -390,6 +394,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
models=[actual_model_name],
successful_requests=successful_requests,
failed_requests=failed_requests,
prompt_cost=total_prompt_cost,
completion_cost=total_completion_cost,
)

View file

@ -370,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
and isinstance(tool_call.get("custom"), dict)
)
for msg in messages:
leading_system_count: Final = next(
(index for index, msg in enumerate(messages) if msg.get("role") != "system"),
len(messages),
)
for index, msg in enumerate(messages):
role = msg.get("role")
content = msg.get("content", "")
tool_calls = msg.get("tool_calls")
@ -378,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if role == "system":
# Extract system message as instructions
if isinstance(content, str):
if isinstance(content, str) and index < leading_system_count:
if instructions:
# Concatenate multiple system prompts with a space
instructions = f"{instructions} {content}"

View file

@ -1461,7 +1461,10 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"})
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"

View file

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

View file

@ -356,6 +356,12 @@
"description": "OpenTelemetry collector endpoint URL",
"required": true
},
"otel_traces_endpoint": {
"type": "text",
"ui_name": "Traces Endpoint URL",
"description": "Complete trace export URL used verbatim when the collector does not serve /v1/traces (OTel v2 only)",
"required": false
},
"otel_headers": {
"type": "text",
"ui_name": "Headers",

View file

@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger):
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None,
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
allowed_hooks: Final = frozenset(supported_event_hooks) | (
frozenset((GuardrailEventHooks.logging_only,))
if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
else frozenset()
)
def _validate_event_hook_list_is_in_supported_event_hooks(
event_hook: list[GuardrailEventHooks] | list[str],
supported_event_hooks: list[GuardrailEventHooks],
@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger):
for hook in event_hook:
if isinstance(hook, str):
hook = GuardrailEventHooks(hook)
if hook not in supported_event_hooks:
if hook not in allowed_hooks:
raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}")
if event_hook is None:
@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger):
default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default]
_validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks)
elif isinstance(event_hook, GuardrailEventHooks):
if event_hook not in supported_event_hooks:
if event_hook not in allowed_hooks:
raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}")
@staticmethod
@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger):
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
def _deployment_pre_call_target(self) -> "CustomLogger":
def _deployment_hook_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
return self
try:
@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger):
# CHECK IF GUARDRAIL REJECTS THE REQUEST
if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
target: Final = self._deployment_pre_call_target()
target: Final = self._deployment_hook_target()
if target is not self:
kwargs["guardrail_to_apply"] = self
result: Final = await target.async_pre_call_hook(
@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger):
return None
# CHECK IF GUARDRAIL REJECTS THE REQUEST
result: Final = await self.async_post_call_success_hook(
target: Final = self._deployment_hook_target()
hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger):
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=request_data,
data=hook_request_data,
response=response,
)

View file

@ -72,6 +72,14 @@ class ExporterSpec(BaseModel):
description="console | in_memory | otlp_http | otlp_grpc | <factory kind>",
)
endpoint: str | None = None
traces_endpoint: str | None = Field(
default=None,
description=(
"Complete OTLP/HTTP trace URL, used verbatim. Set this when the "
"collector serves traces on a path other than ``/v1/traces``; "
"``endpoint`` is a base URL the signal path is appended to."
),
)
headers: str | None = None
owner: ExporterOwner | None = Field(
default=None,
@ -127,6 +135,14 @@ class OpenTelemetryV2Config(BaseSettings):
default=None,
validation_alias=AliasChoices("OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"),
)
traces_endpoint: str | None = Field(
default=None,
validation_alias=AliasChoices("OTEL_TRACES_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"),
description=(
"Complete OTLP/HTTP trace URL for the single-destination shorthand, "
"used verbatim instead of ``endpoint`` + ``/v1/traces``."
),
)
headers: str | None = Field(
default=None,
validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"),
@ -250,7 +266,7 @@ class OpenTelemetryV2Config(BaseSettings):
@model_validator(mode="after")
def _normalize(self) -> "OpenTelemetryV2Config":
# An endpoint with the default exporter kind implies OTLP/HTTP.
if self.endpoint and self.exporter == "console":
if (self.endpoint or self.traces_endpoint) and self.exporter == "console":
self.exporter = "otlp_http"
# When no explicit destinations are given, fold the single-destination
# shorthand into one spec so the provider always has a destination.
@ -259,6 +275,7 @@ class OpenTelemetryV2Config(BaseSettings):
ExporterSpec(
kind=self.exporter,
endpoint=self.endpoint,
traces_endpoint=self.traces_endpoint,
headers=self.headers,
)
]

View file

@ -170,7 +170,7 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
)
return HTTPExporter(
endpoint=_otlp_traces_endpoint(spec.endpoint),
endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint),
headers=parse_headers(spec.headers),
)
if kind in _OTLP_GRPC_KINDS:
@ -201,7 +201,14 @@ def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple
exporters, populate ``config.exporters`` directly.
"""
return _exporter_from_spec(ExporterSpec(kind=config.exporter, endpoint=config.endpoint, headers=config.headers))
return _exporter_from_spec(
ExporterSpec(
kind=config.exporter,
endpoint=config.endpoint,
traces_endpoint=config.traces_endpoint,
headers=config.headers,
)
)
def _otlp_metrics_endpoint(endpoint: str | None) -> str | None:

View file

@ -1,10 +1,12 @@
# What is this?
## Helper utilities
import copy
import logging
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason
@ -37,6 +39,41 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non
return float(seconds / denominator)
_DROP_PARAMS_BOOL: Final = TypeAdapter(bool)
def normalize_drop_params(value: object) -> bool | None:
if value is None or isinstance(value, bool):
return value
try:
return _DROP_PARAMS_BOOL.validate_python(value.strip() if isinstance(value, str) else value)
except ValidationError:
return None
def drop_params_flag(value: object, source: str, logger: logging.Logger) -> bool:
normalized: Final = normalize_drop_params(value)
if normalized is None and value is not None:
logger.warning("%s=%r is not a flag value, treating it as off", source, value)
return bool(normalized)
DROP_PARAMS_ENV_VAR: Final = "LITELLM_DROP_PARAMS"
def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool:
configured: Final = environ.get(DROP_PARAMS_ENV_VAR, "").strip()
if configured == "":
return False
normalized: Final = normalize_drop_params(configured)
if normalized is None:
logger.warning(
"%s=%r is not a flag value, treating it as on. Set it to true or false", DROP_PARAMS_ENV_VAR, configured
)
return True
return normalized
def safe_divide(
numerator: float,
denominator: float,

View file

@ -2,6 +2,7 @@ from collections.abc import Mapping, MutableMapping
from types import MappingProxyType
from typing import Final
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.llms.openai.data_residency import infer_openai_data_residency
AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
@ -113,7 +114,7 @@ def get_litellm_params(
custom_prompt_dict: dict | None = None,
litellm_metadata: dict | None = None,
disable_add_transform_inline_image_block: bool | None = None,
drop_params: bool | None = None,
drop_params: bool | str | None = None,
prompt_id: str | None = None,
prompt_variables: dict | None = None,
async_call: bool | None = None,
@ -175,7 +176,7 @@ def get_litellm_params(
"custom_prompt_dict": custom_prompt_dict,
"litellm_metadata": litellm_metadata,
"disable_add_transform_inline_image_block": disable_add_transform_inline_image_block,
"drop_params": drop_params,
"drop_params": normalize_drop_params(drop_params),
"prompt_id": prompt_id,
"prompt_variables": prompt_variables,
"async_call": async_call,

View file

@ -48,6 +48,7 @@ from litellm.constants import (
)
from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
_select_model_name_for_cost_calc,
)
from litellm.exceptions import (
@ -447,6 +448,13 @@ def _provider_response_id(source: object) -> str | None:
return candidate if isinstance(candidate, str) and candidate else None
def mask_api_base_credentials(api_base: str) -> str:
if "key=" not in api_base:
return api_base
key_end: Final = api_base.find("key=") + 4
return api_base[:key_end] + "*" * 5 + api_base[-4:]
class Logging(LiteLLMLoggingBaseClass):
global \
supabaseClient, \
@ -1189,14 +1197,7 @@ class Logging(LiteLLMLoggingBaseClass):
return data
def _get_masked_api_base(self, api_base: str) -> str:
if "key=" in api_base:
# Find the position of "key=" in the string
key_index: Final = api_base.find("key=") + 4
# Mask the last 5 characters after "key="
masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:]
else:
masked_api_base = api_base
return str(masked_api_base)
return str(mask_api_base_credentials(api_base))
def _pre_call(self, input, api_key, model=None, additional_args={}):
"""
@ -2028,6 +2029,17 @@ class Logging(LiteLLMLoggingBaseClass):
results=result,
)
elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
combined_ws_usage: Final = (
ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
)
logging_result = LiteLLMRealtimeStreamLoggingObject(
usage=combined_ws_usage,
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
elif (
self.call_type == CallTypes.llm_passthrough_route.value
or self.call_type == CallTypes.allm_passthrough_route.value
@ -2938,6 +2950,19 @@ class Logging(LiteLLMLoggingBaseClass):
result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above
result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result.usage = batch_usage
batch_prompt_cost: Final = kwargs.get("batch_prompt_cost", None)
batch_completion_cost: Final = kwargs.get("batch_completion_cost", None)
if (
isinstance(batch_prompt_cost, float)
and isinstance(batch_completion_cost, float)
and isinstance(batch_cost, float)
):
self.set_cost_breakdown(
input_cost=batch_prompt_cost,
output_cost=batch_completion_cost,
total_cost=batch_cost,
cost_for_built_in_tools_cost_usd_dollar=0.0,
)
elif should_compute_batch_data:
batch_result: Final = await _handle_completed_batch(
@ -2953,6 +2978,12 @@ class Logging(LiteLLMLoggingBaseClass):
result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result.usage = batch_result.usage
self.set_cost_breakdown(
input_cost=batch_result.prompt_cost,
output_cost=batch_result.completion_cost,
total_cost=batch_result.cost,
cost_for_built_in_tools_cost_usd_dollar=0.0,
)
self.truncated_messages_for_logging = await truncate_base64_in_messages_async(
StandardLoggingPayloadSetup.append_system_prompt_messages(

View file

@ -13,9 +13,10 @@ Pattern Overview:
"""
import json
from collections.abc import Mapping, Sequence
from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from itertools import chain, repeat
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import ReadOnly, TypedDict, assert_never
@ -29,6 +30,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
anthropic_tool_name,
@ -168,6 +170,8 @@ class AnthropicMessagesHandler(BaseTranslation):
them through guardrail rewrites; downstream provider handling is out of scope.
"""
delivers_ended_stream_text_rewrites = True
def __init__(self):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
@ -1014,11 +1018,17 @@ class AnthropicMessagesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Sequence[object]:
"""
Process output streaming response by applying guardrails to text content.
Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far.
With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite
written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked);
a rewrite on a stream that never reported a ``stop_reason`` has no write-back and is reported as
undeliverable, so the pipeline executor discards it and releases the original chunks.
"""
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -1065,6 +1075,15 @@ class AnthropicMessagesHandler(BaseTranslation):
responses_so_far, request_data
)
raise
guardrailed_texts: Final = _guardrailed_inputs.get("texts")
if (
deliver_ended_stream_rewrites
and isinstance(string_so_far, str)
and string_so_far
and guardrailed_texts
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
return responses_so_far
@ -1087,6 +1106,11 @@ class AnthropicMessagesHandler(BaseTranslation):
if e.original_response is None:
e.original_response = self._build_streaming_usage_response(responses_so_far, request_data)
raise
unended_texts: Final = _guardrailed_inputs.get("texts")
if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
return responses_so_far
def _prepare_request_data(
@ -1180,6 +1204,63 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
@staticmethod
def _write_ended_stream_text_rewrite(
responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: str,
) -> None:
"""Deliver an ended-stream guardrail text rewrite by rewriting the
buffered chunks in place: the first ``text_delta`` carries the full
rewritten text and every later one is blanked, leaving the surrounding
message and content-block framing untouched. Handles both chunk formats
this stream carries (parsed event dicts and raw SSE bytes)."""
replacements: Final = chain((rewritten_text,), repeat(""))
for idx, item in enumerate(responses_so_far):
if isinstance(item, dict):
delta = item.get("delta")
if item.get("type") == "content_block_delta" and isinstance(delta, dict):
if delta.get("type") == "text_delta":
delta["text"] = next(replacements)
elif isinstance(item, (bytes, bytearray)):
responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer
AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
)
@staticmethod
def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
"""Rewrite every ``text_delta`` data line in one SSE chunk with the next
replacement text, leaving all other events and framing byte-identical."""
try:
decoded: Final = sse_bytes.decode("utf-8")
except UnicodeDecodeError:
return sse_bytes
return "\n\n".join(
AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
).encode("utf-8")
@staticmethod
def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
@staticmethod
def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
if not line.startswith("data:"):
return line
try:
data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads(
line[len("data:") :].strip()
)
except json.JSONDecodeError:
return line
if not isinstance(data, dict) or data.get("type") != "content_block_delta":
return line
delta: Final = data.get("delta")
if not isinstance(delta, dict) or delta.get("type") != "text_delta":
return line
return "data: " + json.dumps(
{**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts
)
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
return StreamingScanKey(

View file

@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, Optional
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
if TYPE_CHECKING:
from fastapi import HTTPException
@ -52,6 +52,14 @@ class StreamingScanKey:
class BaseTranslation(ABC):
delivers_ended_stream_text_rewrites: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` accepts
``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
stream, writes guardrail text rewrites back across ``responses_so_far`` so
a buffered pipeline can release rewritten chunks. Tool-call rewrites, and
text rewrites on every other translation, are undeliverable: the pipeline
executor discards them and releases the original chunks."""
@staticmethod
def transform_user_api_key_dict_to_metadata(
user_api_key_dict: Any | None,
@ -157,6 +165,7 @@ class BaseTranslation(ABC):
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Any:
"""
Process output streaming response with guardrails.
@ -164,6 +173,11 @@ class BaseTranslation(ABC):
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.
``deliver_ended_stream_rewrites`` is passed True only when the caller
holds the whole buffered stream and the subclass declares
``delivers_ended_stream_text_rewrites``: the handler then writes
guardrail text rewrites back across ``responses_so_far`` instead of
discarding them.
"""
return responses_so_far

View file

@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
return self._transform_ocr_response(
provider_config=provider_config,
model=model,
@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
# Use async response transform for async operations
return await provider_config.async_transform_ocr_response(
model=model,

View file

@ -78,6 +78,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
Convert chat completions request data to OpenAI-spec structured messages.
@ -453,6 +455,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list["ModelResponseStream"]:
"""
Process output streaming responses by applying guardrails to text content.
@ -467,6 +470,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
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.
deliver_ended_stream_rewrites: When True and the buffered stream has
ended, guardrail text rewrites are written back across
``responses_so_far`` (full rewritten text in each choice's first
content-carrying chunk, the rest blanked) instead of discarded.
Returns:
The (unmodified) list of responses.
@ -492,6 +499,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
)
async def _process_streaming_block_only(
@ -502,27 +510,23 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict | None,
deliver_ended_stream_rewrites: bool = False,
) -> 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)."""
(see ``_process_streaming_transform`` for the incremental_diff path) unless
``deliver_ended_stream_rewrites`` opts the ended-stream branch in."""
has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
if has_stream_ended:
# convert to model response
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
# run process_output_response
await self.process_output_response(
response=model_response,
await self._process_ended_stream(
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,
deliver_ended_stream_rewrites=deliver_ended_stream_rewrites,
)
return responses_so_far
# Step 0: Check if any response has text content to process
@ -595,6 +599,39 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return responses_so_far
async def _process_ended_stream(
self,
*,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None",
user_api_key_dict: "UserAPIKeyAuth | None",
request_data: dict[str, object] | None, # mutable-ok: same request-payload shape the hooks take
deliver_ended_stream_rewrites: bool,
) -> None:
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text rewrite
back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
await self.process_output_response(
response=model_response,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if deliver_ended_stream_rewrites:
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
def build_stream_error_items(
self,
exc: "HTTPException",
@ -745,8 +782,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""
combined_texts: Final[dict[tuple[int, int | None], str]] = {}
for response_idx, response in enumerate(responses_so_far):
for choice_idx, choice in enumerate(response.choices):
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):
@ -759,7 +796,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(content, str):
# String content - accumulate for this choice
str_key: tuple[int, int | None] = (choice_idx, None)
str_key: tuple[int, int | None] = (choice.index, None)
if str_key not in combined_texts:
combined_texts[str_key] = ""
combined_texts[str_key] += content
@ -770,7 +807,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
text_str = content_item.get("text")
if text_str:
list_key: tuple[int, int | None] = (
choice_idx,
choice.index,
content_idx,
)
if list_key not in combined_texts:
@ -960,6 +997,52 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "name" in func_dict:
existing_tool_call.function.name = func_dict["name"]
@staticmethod
def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]:
return tuple(
choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices
)
async def _write_ended_stream_text_rewrites(
self,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrailed_response: "ModelResponse",
pre_guardrail_texts: tuple[str | None, ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail text rewrites back across the buffered
chunks: the full rewritten text lands in the choice's first
content-carrying chunk and the rest are blanked, the same shape the
in-flight write-back uses. Chunks carrying only finish_reason or usage
stay untouched. A rewrite on a stream carrying more than one distinct
choice index is reported as undeliverable, so the pipeline executor
discards it and releases the original chunks."""
post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response)
changed: Final = tuple(
after
for before, after in zip(pre_guardrail_texts, post_guardrail_texts)
if before is not None and after is not None and after != before
)
if not changed:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
if len(stream_choice_indices) != 1:
# stream_chunk_builder collapses every choice into one index-0
# choice, so a rewrite of the rebuilt response cannot be attributed
# back to a single choice on an n>1 stream: report it undeliverable
# rather than deliver the rewrite on the wrong choice
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
target_choice_index: Final = next(iter(stream_choice_indices))
await self._apply_guardrail_responses_to_output_streaming(
responses=responses_so_far,
guardrailed_texts=list(changed), # mutable-ok: callee takes lists
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
)
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],
@ -975,7 +1058,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Args:
responses: List of ModelResponseStream objects to modify
guardrailed_texts: List of guardrailed text responses (combined from all chunks)
task_mappings: List of tuples (choice_idx, content_idx)
task_mappings: List of tuples (choice_idx, content_idx), where choice_idx
is the choice's ``index`` field, not its position in a chunk's list
Override this method to customize how responses are applied to streaming responses.
"""
@ -991,9 +1075,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
# Key: (choice_idx, content_idx), Value: boolean (True if already set)
already_set: Final[dict[tuple[int, int | None], bool]] = {}
# Iterate through all responses and update content
for response_idx, response in enumerate(responses):
for choice_idx_in_response, choice in enumerate(response.choices):
# Iterate through all responses and update content, matching each chunk's
# choice by its index field: on n>1 streams a chunk usually carries one
# choice at list position 0 whose index names the logical choice.
for response in responses:
for choice in response.choices:
if isinstance(choice, litellm.StreamingChoices):
content = choice.delta.content
elif isinstance(choice, litellm.Choices):
@ -1006,7 +1092,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if isinstance(content, str):
# String content
str_key: tuple[int, int | None] = (choice_idx_in_response, None)
str_key: tuple[int, int | None] = (choice.index, None)
if str_key in guardrail_map:
if str_key not in already_set:
# First chunk - set the complete guardrailed text
@ -1027,7 +1113,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
for content_idx, content_item in enumerate(content):
if "text" in content_item:
list_key: tuple[int, int | None] = (
choice_idx_in_response,
choice.index,
content_idx,
)
if list_key in guardrail_map:

View file

@ -33,7 +33,7 @@ import time
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from itertools import accumulate, chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast
@ -49,6 +49,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_stream_usage,
@ -118,6 +119,15 @@ class ResponsesStreamChunk(TypedDict, total=False):
content_index: ReadOnly[int]
_TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
{
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
}
)
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
)
@ -330,6 +340,8 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
"""
Convert Responses API request data to OpenAI-spec structured messages.
@ -667,6 +679,8 @@ class OpenAIResponsesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> list[Any]:
"""
Process output streaming response by applying guardrails to text content.
@ -675,10 +689,18 @@ class OpenAIResponsesHandler(BaseTranslation):
chunk, apply the guardrail, then write the result back in-place so the
caller sees the modified content (e.g. PII tokens replaced).
For ``response.completed`` events (the normal end-of-stream signal) we
use the same per-item extraction + task-mapping approach as
``process_output_response`` so that unmasking / blocking works correctly
for every output item.
For terminal envelope events (``response.completed``, and equally
``response.incomplete`` / ``response.failed``, whose envelopes carry the
partial output) we use the same per-item extraction + task-mapping
approach as ``process_output_response`` so that unmasking / blocking
works correctly for every output item. With
``deliver_ended_stream_rewrites`` the earlier text-carrying events
(``response.output_text.delta`` / ``.done``,
``response.content_part.done``, ``response.output_item.done``) are synced
to the rewritten envelope too, so a client reading deltas sees the
rewrite instead of the raw model output; a rewrite observed where no
write-back is possible is reported as undeliverable, so the pipeline
executor discards it and releases the original events.
"""
if not responses_so_far:
return responses_so_far
@ -690,14 +712,16 @@ class OpenAIResponsesHandler(BaseTranslation):
return responses_so_far
# ------------------------------------------------------------------ #
# Case 1: response.completed — full response is available in the #
# final chunk; iterate output items, apply guardrail, write back. #
# Case 1: terminal envelope events (completed/incomplete/failed). #
# the accumulated response is available in the final chunk; iterate #
# output items, apply guardrail, write back. Falls through to the #
# string fallback when the envelope yields nothing to check. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.completed":
if final_chunk.get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES:
response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {}
if not hasattr(response_obj, "get"):
return responses_so_far
outputs: Final[Sequence[object]] = response_obj.get("output") or []
outputs: Final[Sequence[object]] = (
(response_obj.get("output") or []) if hasattr(response_obj, "get") else []
)
texts_to_check: Final[list[str]] = []
tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = []
@ -747,11 +771,25 @@ class OpenAIResponsesHandler(BaseTranslation):
responses=guardrailed_texts,
task_mappings=task_mappings,
)
return responses_so_far
if deliver_ended_stream_rewrites:
rewrites_by_position: Final = MappingProxyType(
{
task_mappings[task_idx]: rewritten
for task_idx, rewritten in enumerate(guardrailed_texts)
if task_idx < len(texts_to_check) and rewritten != texts_to_check[task_idx]
}
)
if rewrites_by_position:
self._sync_stream_events_with_rewrites(
stream_events=responses_so_far[:-1],
rewrites_by_position=rewrites_by_position,
)
return responses_so_far
# ------------------------------------------------------------------ #
# Case 2: response.output_item.done — extract tool calls only. #
# Case 2: response.output_item.done — extract tool calls only, then #
# fall through to the text fallback when a caller expects rewrites #
# delivered, so a truncated buffer still reports text undeliverable. #
# ------------------------------------------------------------------ #
if final_chunk.get("type") == "response.output_item.done":
model_response_stream: Final = (
@ -769,12 +807,14 @@ class OpenAIResponsesHandler(BaseTranslation):
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
if not deliver_ended_stream_rewrites:
return responses_so_far
# ------------------------------------------------------------------ #
# Fallback: apply guardrail to the accumulated text string. #
# No structured write-back is possible here; guardrails that only #
# need to block/flag (not rewrite) still work correctly. #
# need to block/flag (not rewrite) still work correctly, and a #
# rewrite a caller expects delivered is reported undeliverable. #
# ------------------------------------------------------------------ #
string_so_far: Final = self.get_streaming_string_so_far(responses_so_far)
if string_so_far:
@ -784,28 +824,83 @@ class OpenAIResponsesHandler(BaseTranslation):
)
if response_model:
fallback_inputs["model"] = response_model
await guardrail_to_apply.apply_guardrail(
fallback_outputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=fallback_inputs,
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,
)
fallback_texts: Final = fallback_outputs.get("texts")
if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown")
return responses_so_far
@staticmethod
def _write_event_field(event: object, field: str, value: str) -> None:
if isinstance(event, dict):
event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place
else:
setattr(event, field, value)
def _sync_stream_events_with_rewrites(
self,
stream_events: Sequence[Any],
rewrites_by_position: Mapping[tuple[int, int], str],
) -> None:
"""Sync pre-completion stream events with the rewritten completed
response, keyed by ``(output_index, content_index)``: the first
``output_text.delta`` for a rewritten item carries the full rewritten
text and the rest are blanked, while ``output_text.done``,
``content_part.done``, and ``output_item.done`` events carry the full
rewritten text, so every event a client may read agrees with the
rewritten ``response.completed`` payload."""
delta_replacements: Final = MappingProxyType(
{position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()}
)
for event in stream_events:
if not (isinstance(event, dict) or hasattr(event, "get")):
continue
event_type = event.get("type")
output_index = event.get("output_index")
content_index = event.get("content_index")
if event_type == "response.output_item.done" and isinstance(output_index, int):
self._sync_output_item_done_event(event.get("item"), output_index, rewrites_by_position)
continue
if not isinstance(output_index, int) or not isinstance(content_index, int):
continue
position = (output_index, content_index)
if event_type == "response.output_text.delta" and position in delta_replacements:
self._write_event_field(event, "delta", next(delta_replacements[position]))
elif event_type == "response.output_text.done" and position in rewrites_by_position:
self._write_event_field(event, "text", rewrites_by_position[position])
elif event_type == "response.content_part.done" and position in rewrites_by_position:
part = event.get("part")
if isinstance(part, dict) or hasattr(part, "text"):
self._write_event_field(part, "text", rewrites_by_position[position])
@staticmethod
def _sync_output_item_done_event(
item: object,
output_index: int,
rewrites_by_position: Mapping[tuple[int, int], str],
) -> None:
content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None)
if not isinstance(content, list):
return
for (item_idx, content_idx), rewritten in rewrites_by_position.items():
if item_idx != output_index or content_idx >= len(content):
continue
OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if the streaming has ended.
"""
if not responses_so_far:
return False
terminal_types: Final = frozenset(
(
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
)
)
return stream_item_field(responses_so_far[-1], "type") in terminal_types
return stream_item_field(responses_so_far[-1], "type") in _TERMINAL_ENVELOPE_EVENT_TYPES
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
if not responses_so_far or not hasattr(responses_so_far[-1], "get"):

View file

@ -32,6 +32,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse
created_by: str | None = None
team_id: str | None = None
org_id: str | None = None
class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase):

View file

@ -3585,6 +3585,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
litellm_callback_params=[
"OTEL_EXPORTER",
"OTEL_ENDPOINT",
"OTEL_TRACES_ENDPOINT",
"OTEL_HEADERS",
],
)

View file

@ -2038,7 +2038,7 @@ async def _user_api_key_auth_builder(
fallback_spend=team_member_spend,
max_budget=team_member_budget,
)
if team_member_spend > team_member_budget:
if team_member_spend >= team_member_budget:
_entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
raise litellm.BudgetExceededError(
current_cost=team_member_spend,

View file

@ -10,8 +10,10 @@ import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
@ -507,6 +509,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
CLIENT_OUTPUT_CEILING_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

View file

@ -19,7 +19,7 @@ from litellm.cost_calculator import _infer_call_type
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.llms import load_guardrail_translation_mappings
from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import (
@ -69,6 +69,36 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran
return translation
def resolve_endpoint_translation(
user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None
) -> "tuple[str, BaseTranslation] | None":
"""
Resolve the endpoint guardrail translation for a streamed response: the
request route wins, falling back to inferring the call type from the first
response chunk (the same resolution order the streaming iterator hook uses).
Returns None when the call type is unresolvable or has no translation.
"""
route_call_types: Final = (
get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None
)
call_type: Final = (
route_call_types[0].value
if route_call_types
else (
_infer_call_type(call_type=None, completion_response=first_response_item)
if first_response_item is not None
else None
)
)
if call_type is None:
return None
try:
handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type))
except ValueError:
return None
return call_type, handler_cls()
def _chunk_choices(item: object) -> Sequence[object]:
choices: Final[Sequence[object]] = getattr(item, "choices", None) or []
return choices
@ -343,7 +373,7 @@ class UnifiedLLMGuardrails(CustomLogger):
return response
async def _handle_streaming_block(
async def handle_streaming_block(
self,
exc: "ModifyResponseException",
endpoint_translation: _EndpointTranslation,
@ -399,7 +429,7 @@ class UnifiedLLMGuardrails(CustomLogger):
return None
return call_type
async def _emit_streaming_http_error(
async def emit_streaming_http_error(
self,
exc: HTTPException,
call_type: str | None,
@ -592,7 +622,7 @@ class UnifiedLLMGuardrails(CustomLogger):
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(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=bool(responses_yielded),
@ -601,7 +631,7 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
raise _StreamTerminated()
except HTTPException as e:
async for error_item in self._emit_streaming_http_error(
async for error_item in self.emit_streaming_http_error(
e,
call_type,
responses_so_far,
@ -781,7 +811,7 @@ class UnifiedLLMGuardrails(CustomLogger):
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(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=bool(responses_yielded),
@ -869,6 +899,14 @@ class UnifiedLLMGuardrails(CustomLogger):
choices: Final = _chunk_choices(item)
return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object:
"""Streaming flag resolution order (later wins): default < guardrail
attribute < guardrail_config dict < this callback's optional_params."""
attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default)
config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None)
config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value
return self.optional_params.get(name, config_value)
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -897,17 +935,8 @@ class UnifiedLLMGuardrails(CustomLogger):
if guardrail_to_apply is None:
guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
# Get streaming configuration. Resolution order (later wins): default
# < guardrail attribute < guardrail_config dict < this callback's
# optional_params.
def _streaming_flag(name: str, default: object) -> Any:
value = default
if guardrail_to_apply is not None:
value = getattr(guardrail_to_apply, name, value)
config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {})
if isinstance(config, dict):
value = config.get(name, value)
return self.optional_params.get(name, value)
return self.resolve_streaming_flag(guardrail_to_apply, name, default)
sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
# Only apply the guardrail at end of stream (not per chunk).
@ -1091,7 +1120,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# The current chunk was appended to responses_so_far but not
# yet yielded, so exclude it: the continuation must reflect
# only what the client has actually received.
async for block_chunk in self._handle_streaming_block(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=chunks_yielded,
@ -1101,7 +1130,7 @@ class UnifiedLLMGuardrails(CustomLogger):
return
except HTTPException as e:
# Response already started (we already yielded chunks); cannot send 400.
async for error_item in self._emit_streaming_http_error(
async for error_item in self.emit_streaming_http_error(
e,
call_type,
responses_so_far,
@ -1175,7 +1204,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# terminating SSE sequence with the block message rather than
# propagating into a bare error blob that truncates the stream.
# The withheld original chunks are never released.
async for block_chunk in self._handle_streaming_block(
async for block_chunk in self.handle_streaming_block(
e,
endpoint_translation,
stream_started=bool(responses_yielded),
@ -1184,7 +1213,7 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
return
except HTTPException as e:
async for error_item in self._emit_streaming_http_error(
async for error_item in self.emit_streaming_http_error(
e,
call_type,
responses_so_far,

View file

@ -18,11 +18,13 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
OTEL_SERVICE_NAME_METADATA_KEYS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
SESSION_ID_OMITTED_METADATA_KEY,
@ -289,6 +291,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
"standard_logging_object",
"proxy_server_request",
@ -325,7 +328,9 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
# ``attempted_fallbacks`` and ``original_model_group`` are written by the router
# and read by spend logs as fact; a client value has no legitimate meaning and no
# key or team setting keeps it, so the strip is never gated.
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset({"attempted_fallbacks", "original_model_group"})
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
)
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
# Request fields whose value, when URL-valued, becomes the outbound destination
@ -3075,10 +3080,9 @@ def _apply_resolved_guardrails_to_metadata(
if metadata_variable_name not in data:
data[metadata_variable_name] = {}
# Track pipeline-managed guardrails to exclude from independent execution
pipeline_managed_guardrails: set = set()
# Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode
if pipelines:
pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines)
data[metadata_variable_name]["_guardrail_pipelines"] = pipelines
data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails
verbose_proxy_logger.debug(
@ -3095,10 +3099,8 @@ def _apply_resolved_guardrails_to_metadata(
existing_guardrails = []
# Combine existing guardrails with policy-resolved guardrails (no duplicates)
# Exclude pipeline-managed guardrails from the flat list
combined = set(existing_guardrails)
combined.update(resolved_guardrails)
combined -= pipeline_managed_guardrails
data[metadata_variable_name]["guardrails"] = list(combined)
verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined))

View file

@ -5,12 +5,16 @@ Runs guardrails sequentially per pipeline step definitions, handling
pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import copy
import time
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar
from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import LOGS_GUARDRAIL_INFORMATION_MARKER
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -21,6 +25,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
independent_snapshot,
)
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -29,7 +34,14 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineStep,
PipelineStepResult,
)
from litellm.types.utils import StandardLoggingGuardrailInformation
from litellm.types.utils import GenericGuardrailAPIInputs, StandardLoggingGuardrailInformation
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
)
from litellm.proxy._types import UserAPIKeyAuth
try:
from fastapi.exceptions import HTTPException
@ -37,6 +49,121 @@ except ImportError:
HTTPException = None
class UndeliverableStreamRewrite(Exception):
def __init__(self, guardrail_name: str) -> None:
super().__init__(
f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's "
"streaming pipeline cannot deliver"
)
self.guardrail_name: Final = guardrail_name
def _tool_call_shape(tool_call: object) -> tuple[object, object]:
plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
function: Final = plain.get("function") if isinstance(plain, Mapping) else None
if not isinstance(function, Mapping):
return (None, None)
return (function.get("name"), function.get("arguments"))
def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None:
return None if texts is None else tuple(texts)
def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None:
return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls)
def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool:
return sent is not None and returned is not None and returned != sent
_GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object])
def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT:
vars(method)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the method the class body just defined
return method
class _StreamRewriteObserver(CustomGuardrail):
"""Stand-in handed to the endpoint translation in place of a streaming pipeline step's
guardrail. It records whether the guardrail returned different output than it was given,
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text
rewrites are deliverable on translations that write them back across the buffered chunks
(``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any
other translation are discarded by the executor, which releases the original chunks.
The inner guardrail's ``apply_guardrail`` already records the guardrail information
and span, so the observer's stays out of ``log_guardrail_information``."""
def __init__(self, inner: CustomGuardrail) -> None:
super().__init__(guardrail_name=inner.guardrail_name)
self.inner: Final = inner
self.rewrote_texts = False
self.rewrote_tool_calls = False
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@_logged_by_inner_guardrail
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
sent_texts: Final = _text_snapshot(inputs.get("texts"))
sent_tool_shapes: Final = _tool_call_shapes(inputs.get("tool_calls"))
outputs: Final = await self.inner.apply_guardrail(
inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
)
self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(
sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))
)
return outputs
def _prepare_hook_input(
step: PipelineStep,
callback: CustomGuardrail,
data: dict, # mutable-ok: same request-payload shape the hooks mutate
raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data
) -> tuple[dict, bool]: # mutable-ok: returns that same request-payload dict
"""Inject the step's guardrail name into metadata so should_run_guardrail() allows it,
and pick the payload the step scans: a scan_raw_request step evaluates the pristine
pre-pipeline snapshot instead of `data` (which earlier pass_data steps in this same
pipeline may have already rewritten), same reason the normal sequential/parallel
guardrail loops do this."""
if "metadata" not in data:
data["metadata"] = {} # mutable-ok: request metadata bucket, hooks mutate it
data["metadata"]["guardrails"] = [
step.guardrail
] # mutable-ok: guardrails list is part of the request-payload shape
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] # mutable-ok: request metadata shape
return hook_input, scans_raw_request
def _release_original_chunks(
guardrail_name: str,
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place
originals: Sequence[object],
) -> None:
streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives
verbose_proxy_logger.warning(
"Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming "
"pipeline cannot deliver yet; the rewrite was discarded and the original stream released",
guardrail_name,
)
class PipelineExecutor:
"""Executes guardrail pipelines with ordered, conditional step logic."""
@ -49,6 +176,8 @@ class PipelineExecutor:
call_type: str,
policy_name: str,
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step
endpoint_translation: "BaseTranslation | None" = None,
) -> PipelineExecutionResult:
"""
Execute pipeline steps sequentially with conditional actions.
@ -65,6 +194,12 @@ class PipelineExecutor:
step whose guardrail opted into ``scan_raw_request`` evaluates
the original request instead of whatever an earlier
``pass_data`` step in this same pipeline already rewrote.
streaming_chunks: buffered chunks of a completed stream. When set
(with ``endpoint_translation``), post_call steps scan the
assembled streamed output through the endpoint translation
instead of calling ``async_post_call_success_hook``.
endpoint_translation: the guardrail translation for the streamed
endpoint, resolved by the caller.
Returns:
PipelineExecutionResult with terminal action and step results
@ -89,6 +224,8 @@ class PipelineExecutor:
user_api_key_dict=user_api_key_dict,
call_type=call_type,
raw_request_snapshot=raw_request_snapshot,
streaming_chunks=streaming_chunks,
endpoint_translation=endpoint_translation,
)
duration = time.perf_counter() - start_time
@ -114,8 +251,10 @@ class PipelineExecutor:
action,
)
# Forward modified data to next step if pass_data is True
if step.pass_data and modified_data is not None:
# Forward modified data to the next step if pass_data is True;
# post_call response replacements always chain, matching the flat
# callback loop where each hook sees the previous hook's response
if modified_data is not None and (step.pass_data or mode == "post_call"):
working_data = {**working_data, **modified_data}
# Handle terminal actions
@ -129,6 +268,7 @@ class PipelineExecutor:
step_results=step_results,
error_message=error_detail,
original_exception=original_exception,
modified_data=working_data if working_data != data else None,
)
if action == "modify_response":
@ -137,6 +277,7 @@ class PipelineExecutor:
terminal_action="modify_response",
step_results=step_results,
modify_response_message=step.modify_response_message or error_detail,
modified_data=working_data if working_data != data else None,
)
# action == "next" → continue to next step
@ -144,6 +285,51 @@ class PipelineExecutor:
# Ran out of steps without a terminal action → default allow
return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
@staticmethod
async def _run_streaming_step(
step: PipelineStep,
callback: CustomGuardrail,
endpoint_translation: "BaseTranslation",
streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place
hook_input: dict[str, object], # mutable-ok: same request-payload shape as data
user_api_key_dict: "UserAPIKeyAuth | None",
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Run one streaming post_call step through the endpoint translation, delivering
text rewrites on translations that support ended-stream write-back. A rewrite that
cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation
without write-back, or one the translation refused with
``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the
originals and the step passes, so the client gets the stream the merge base sent."""
observer: Final = _StreamRewriteObserver(callback)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
originals: Final = copy.deepcopy(streaming_chunks)
try:
if deliver_rewrites:
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
deliver_ended_stream_rewrites=True,
)
else:
await endpoint_translation.process_output_streaming_response(
responses_so_far=streaming_chunks,
guardrail_to_apply=observer,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=hook_input,
)
except UndeliverableStreamRewrite:
_release_original_chunks(step.guardrail, streaming_chunks, originals)
else:
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
if not callback.records_own_guardrail_information:
add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail)
@staticmethod
async def _run_step(
step: PipelineStep,
@ -152,6 +338,8 @@ class PipelineExecutor:
user_api_key_dict: Any,
call_type: str,
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step
endpoint_translation: "BaseTranslation | None" = None,
) -> tuple[
Literal["pass", "fail", "error"],
dict | None,
@ -175,29 +363,13 @@ class PipelineExecutor:
verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail)
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot)
snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input))
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
if use_unified:
use_unified: Final = PipelineExecutor.supports_unified_execution(callback)
if use_unified and streaming_chunks is None:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
@ -213,6 +385,24 @@ class PipelineExecutor:
callback.mark_pre_call_hook_ran(data)
if isinstance(response, dict):
callback.mark_pre_call_hook_ran(response)
elif mode == "post_call" and streaming_chunks is not None:
if not use_unified or endpoint_translation is None:
return (
"error",
None,
f"Guardrail '{step.guardrail}' does not support streaming pipeline execution",
None,
)
await PipelineExecutor._run_streaming_step(
step=step,
callback=callback,
endpoint_translation=endpoint_translation,
streaming_chunks=streaming_chunks,
hook_input=hook_input,
user_api_key_dict=user_api_key_dict,
litellm_logging_obj=data.get("litellm_logging_obj"),
)
response = None
elif mode == "post_call":
response = await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
@ -226,11 +416,19 @@ class PipelineExecutor:
# same contract as run_in_parallel/scan_raw_request elsewhere: any
# data it returned is discarded, since applying it on top of the
# raw snapshot would silently undo whatever an earlier step in
# this pipeline already did.
modified_data = None
if response is not None and isinstance(response, dict) and not scans_raw_request:
modified_data = response
return ("pass", modified_data, None, None)
# this pipeline already did. A post_call hook's non-None return is
# a replacement response (the flat callback-loop contract), carried
# under the same "response" key the step input uses.
if response is None or scans_raw_request:
return ("pass", None, None, None)
if mode == "post_call":
return (
"pass",
{"response": response},
None,
None,
) # mutable-ok: modified-data contract is a plain dict
return ("pass", response if isinstance(response, dict) else None, None, None)
except Exception as e:
if CustomGuardrail._is_guardrail_intervention(e):
@ -246,6 +444,12 @@ class PipelineExecutor:
entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:],
)
@staticmethod
def supports_unified_execution(callback: CustomGuardrail) -> bool:
"""Whether this guardrail runs through the unified apply_guardrail path,
the interface streaming pipeline execution requires."""
return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
"""Look up an initialized guardrail callback by name from litellm.callbacks."""

View file

@ -8,10 +8,13 @@ from __future__ import annotations
import glob
import os
import re
from typing import Final
from litellm._logging import verbose_proxy_logger
_LIVE_GAUGE_PID: Final = re.compile(r"gauge_live[a-z]*_(\d+)\.db$")
def wipe_directory(directory: str) -> None:
"""Delete all .db files in the directory. Called once before workers fork."""
@ -38,3 +41,35 @@ def mark_worker_exit(worker_pid: int) -> None:
verbose_proxy_logger.info("Prometheus cleanup: marked worker %s as dead", worker_pid)
except Exception as e:
verbose_proxy_logger.warning("Failed to mark prometheus worker %s as dead: %s", worker_pid, e)
def _is_running(pid: int) -> bool:
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
except PermissionError:
return True
return True
def mark_dead_workers(directory: str) -> tuple[int, ...]:
"""Drop the live-gauge files of workers that no longer exist and return their pids.
Uvicorn's multi-worker supervisor has no exit hook, so a replacement worker calls this at startup; without it
a crashed worker's in-flight gauges stay in the aggregate forever.
"""
owners: Final = frozenset(
int(match.group(1))
for match in map(_LIVE_GAUGE_PID.search, glob.glob(os.path.join(directory, "gauge_live*_*.db")))
if match is not None
)
dead: Final = tuple(sorted(pid for pid in owners if pid != os.getpid() and not _is_running(pid)))
if not dead:
return dead
from prometheus_client import multiprocess
for pid in dead:
multiprocess.mark_process_dead(pid, path=directory)
verbose_proxy_logger.info("Prometheus cleanup: marked dead workers %s in %s", dead, directory)
return dead

View file

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

View file

@ -278,6 +278,7 @@ from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
drop_params_flag,
get_litellm_metadata_from_kwargs,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
@ -631,6 +632,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
router as pass_through_router,
)
from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit
from litellm.proxy.public_endpoints import router as public_endpoints_router
from litellm.proxy.public_endpoints.public_v1 import router as public_v1_router
from litellm.proxy.rag_endpoints.endpoints import router as rag_router
@ -1059,6 +1061,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
init_verbose_loggers()
prometheus_multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR")
if prometheus_multiproc_dir:
mark_dead_workers(prometheus_multiproc_dir)
## RUN WORKER STARTUP HOOKS (e.g., gflags initialization) ##
_startup_hooks_env: Final = os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "")
if _startup_hooks_env:
@ -1365,6 +1371,9 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
await proxy_shutdown_event(worker_heartbeat=worker_heartbeat)
if prometheus_multiproc_dir:
mark_worker_exit(os.getpid())
def _generate_stable_operation_id(route: "APIRoute") -> str:
operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}")
@ -2574,10 +2583,11 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
)
async def reseed_spend_counter_from_db(counter_key: str) -> None:
async def reseed_spend_counter_from_db(counter_key: str) -> bool:
"""Recover a counter that the reservation reconcile found in an inconsistent
state (missing, or where applying the reconcile delta would drive it
negative) by reseeding it from the DB instead of deleting it.
negative) by reseeding it from the DB instead of deleting it. Returns
whether a DB row was found and the counter was reseeded.
The DB row is a LAGGING authoritative floor, not post-request truth: the
entity .spend column is flushed in batches (every PROXY_BATCH_WRITE_AT), so
@ -2592,8 +2602,9 @@ async def reseed_spend_counter_from_db(counter_key: str) -> None:
"""
db_spend: Final = await SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key)
if db_spend is None:
return
return False
await _repair_stale_spend_counter(counter_key=counter_key, db_spend=db_spend)
return True
async def _floor_spend_from_db(
@ -2655,21 +2666,14 @@ async def _authoritative_floor_spend(
return db_spend
async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]:
"""Return (spend, authoritative). ``authoritative`` is True when the value
came from Redis or a fresh DB read (cross-pod truth), False when it came
from the per-pod in-memory copy or the caller's fallback. Only the
fail-closed path reads the flag; normal callers ignore it."""
# 1. Redis first (cross-pod authoritative). On clean miss, skip
# in-memory: per-pod in-memory only has this pod's writes, so it
# would mask cross-pod increments.
redis_clean_miss = False
async def read_spend_counter_cache_value(counter_key: str) -> tuple[float | None, bool]:
"""Return (value, authoritative) for the live counter, None when absent. A clean
Redis miss is final: the per-pod in-memory copy outlives the Redis TTL and only
holds this pod's writes, so it is consulted only when Redis is unreachable."""
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val), True
redis_clean_miss = True
redis_val: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
return (float(redis_val) if redis_val is not None else None), True
except Exception as e:
verbose_proxy_logger.debug(
"get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
@ -2677,13 +2681,20 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float)
e,
)
# 2. In-memory only when Redis is unreachable.
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val), False
in_memory_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
return (float(in_memory_val) if in_memory_val is not None else None), False
# 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]:
"""Return (spend, authoritative). ``authoritative`` is True when the value
came from Redis or a fresh DB read (cross-pod truth), False when it came
from the per-pod in-memory copy or the caller's fallback. Only the
fail-closed path reads the flag; normal callers ignore it."""
cached_val, cached_authoritative = await read_spend_counter_cache_value(counter_key=counter_key)
if cached_val is not None:
return cached_val, cached_authoritative
# Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
db_spend: Final = await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
@ -5509,6 +5520,8 @@ class ProxyConfig:
parse_budget_reset_time(value)
setattr(litellm, key, value)
elif key == "drop_params":
litellm.drop_params = drop_params_flag(value, "litellm_settings.drop_params", verbose_proxy_logger)
else:
verbose_proxy_logger.debug(
"%s setting litellm.%s=%s%s",

View file

@ -1036,6 +1036,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
created_at DateTime @default(now())
created_by String?
team_id String?
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
updated_at DateTime @updatedAt

View file

@ -905,13 +905,13 @@ async def _set_reserved_entry_actual_cost(
increment=adjustment,
)
elif reseed_on_inconsistent:
# Post-call reconcile / release: the counter was flushed or reseeded
# between reservation and reconcile (Redis restart / cross-pod reset),
# so the optimistic delta no longer applies. Recover by reseeding from
# the DB's lagging authoritative floor rather than deleting the counter
# and failing open — deleting it is what left budgets unenforced after a
# Redis reload.
await reseed_spend_counter_from_db(counter_key=counter_key)
# Post-call reconcile / release: the counter was flushed, expired or reseeded
# between reservation and reconcile, so the optimistic delta no longer applies.
# Reseed from the DB floor (which cannot include this request's cost yet) and
# add the settled cost, since increment_spend_counters skips reserved keys.
reseeded: Final = await reseed_spend_counter_from_db(counter_key=counter_key)
if reseeded and actual_cost > 0:
await _increment_spend_counter_cache(counter_key=counter_key, increment=actual_cost)
else:
# Pre-call admission resize: the in-flight reservation cost is not yet
# persisted, so the DB floor would discard it. Keep the original
@ -925,18 +925,16 @@ async def _counter_can_apply_adjustment(
counter_key: str,
adjustment: float,
) -> bool:
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import read_spend_counter_cache_value
current_value: Final = await spend_counter_cache.async_get_cache(key=counter_key)
try:
current_value, _ = await read_spend_counter_cache_value(counter_key=counter_key)
except (TypeError, ValueError):
return False
if current_value is None:
return False
try:
current_float: Final = float(current_value)
except (TypeError, ValueError):
return False
return not (adjustment < 0 and current_float + adjustment < -1e-12)
return not (adjustment < 0 and current_value + adjustment < -1e-12)
async def _release_applied_entries_best_effort(

View file

@ -590,6 +590,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
metadata=metadata,
standard_logging_payload=standard_logging_payload,
omit_when_missing=_omits_session_id_when_missing(metadata),
batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id),
),
request_duration_ms=_get_request_duration_ms(start_time, end_time),
status=_get_status_for_spend_log(
@ -628,20 +629,44 @@ def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> boo
return general_settings.get("missing_session_id") == "omit"
_BATCH_TRACE_CALL_TYPES: Final = frozenset(
{
CallTypes.create_batch.value,
CallTypes.acreate_batch.value,
CallTypes.retrieve_batch.value,
CallTypes.aretrieve_batch.value,
}
)
def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None:
"""A batch's create row and its poller-written cost row both derive their request id
from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using
that id as the session groups the batch lifecycle into one trace on the logs UI. The
poller builds its own logging context, so per-request trace ids can never link them."""
if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id:
return None
return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX)
def _get_session_id_for_spend_log(
kwargs: Mapping[str, object],
metadata: Mapping[str, object] | None,
standard_logging_payload: StandardLoggingPayload | None,
omit_when_missing: bool,
batch_trace_session_id: str | None = None,
) -> str | None:
"""Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may
be a copied trace id."""
be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks
the per-request trace ids because those differ between the create call and the cost poller's row."""
if omit_when_missing:
session_id: Final = metadata.get("session_id") if metadata else None
return str(session_id) if session_id else None
from litellm._uuid import uuid
if batch_trace_session_id is not None:
return batch_trace_session_id
if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
return str(standard_logging_payload.get("trace_id"))
if kwargs.get("litellm_trace_id") is not None:

View file

@ -11,7 +11,7 @@ import sys
import threading
import time
import traceback
from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import date, datetime, timedelta, timezone
from email.mime.multipart import MIMEMultipart
@ -139,6 +139,7 @@ from litellm.proxy.db.token_auth import (
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
resolve_endpoint_translation,
)
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
@ -177,6 +178,9 @@ from litellm.types.mcp import (
)
from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
from litellm.utils import (
_add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper
)
if TYPE_CHECKING:
from mcp.types import CallToolResult
@ -446,12 +450,161 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail
)
def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]:
managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails")
return (
frozenset(cast("Collection[str]", managed)) # cast-ok: the policy engine wrote these guardrail names
if managed
else frozenset()
def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> frozenset[str]:
return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
def _pipeline_managed_guardrail_names(
data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
) -> frozenset[str]:
return _pipeline_step_guardrail_names(
tuple((policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == mode)
)
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
resolved: Final = tuple(
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
cast( # cast-ok: the resolver returns None for unknown names, filtered below
_custom_logger_compatible_callbacks_literal, callback
)
)
if isinstance(callback, str)
else callback
for callback in litellm.callbacks
)
present: Final = tuple(callback for callback in resolved if callback is not None)
guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail))
others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger
"tuple[CustomLogger, ...]",
tuple(callback for callback in present if not isinstance(callback, CustomGuardrail)),
)
return (guardrails, others)
def _merge_pipeline_metadata_bucket(
data: dict, bucket_key: str, modified_bucket_value: object
) -> None: # mutable-ok: request payload dict, written in place
if not isinstance(modified_bucket_value, dict):
return
modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed
surviving_writes: Final = {
key: value for key, value in modified_bucket.items() if key != "guardrails"
} # mutable-ok: merged into the live request metadata bucket in place
existing_bucket: Final = data.get(bucket_key)
if isinstance(existing_bucket, dict):
cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed
else:
data[bucket_key] = surviving_writes
def _merge_pipeline_metadata_writes(
data: dict, modified_data: Mapping[str, object]
) -> None: # mutable-ok: request payload dict, written in place
"""
Copy metadata-bucket writes from a pipeline's working copy back onto the request.
Post_call pipelines run step hooks against a copied request dict so the payload
already sent upstream stays untouched, but hooks record proxy-internal logging
state in the metadata buckets (``applied_guardrails`` for response headers,
``standard_logging_guardrail_information`` for spend logs), and those writes
must reach the request dict the proxy keeps reading after the pipeline returns.
The ``guardrails`` key is the executor's per-step activation flag for
``should_run_guardrail``, not a hook write, so it stays in the working copy.
"""
for bucket_key in ("metadata", "litellm_metadata"):
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
return tuple(
(policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call"
)
def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None:
if data.get("background") is not True:
return
policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data))
if not policy_names:
return
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines do not run on background responses yet; "
"the response is released ungoverned by them: %s",
", ".join(policy_names),
)
def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
unsupported: Final = tuple(
dict.fromkeys(
step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail)
)
)
if not unsupported:
return True
verbose_proxy_logger.warning(
"Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, "
"which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s",
policy_name,
", ".join(unsupported),
)
return False
def _route_supports_streaming_pipelines(user_api_key_dict: UserAPIKeyAuth) -> bool:
return not user_api_key_dict.request_route or resolve_endpoint_translation(user_api_key_dict, None) is not None
def _stream_gated_guardrail_names(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> frozenset[str]:
if not _route_supports_streaming_pipelines(user_api_key_dict):
return frozenset()
return _pipeline_step_guardrail_names(
tuple(
(policy_name, pipeline)
for policy_name, pipeline in _post_call_pipelines(request_data)
if all(_pipeline_step_supports_unified_streaming(step.guardrail) for step in pipeline.steps)
)
)
def _streamable_post_call_pipelines(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
"""
The post_call pipelines a streaming response can be gated through.
Streaming pipelines scan the buffered stream through the endpoint guardrail
translation of the request route, so every step's guardrail needs the
unified apply_guardrail interface and the route needs a translation. A
pipeline that cannot be run that way yet is left out and its guardrails
run on the stream on their own, the way they did before pipelines ran on
streams at all, with a warning naming the pipeline.
"""
post_call_pipelines: Final = _post_call_pipelines(request_data)
if not post_call_pipelines:
return ()
if not _route_supports_streaming_pipelines(user_api_key_dict):
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
"on their own: %s",
user_api_key_dict.request_route,
", ".join(policy_name for policy_name, _pipeline in post_call_pipelines),
)
return ()
return tuple(
(policy_name, pipeline)
for policy_name, pipeline in post_call_pipelines
if _pipeline_is_streamable(policy_name, pipeline)
)
@ -857,6 +1010,14 @@ class ProxyLogging:
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
# Runs after load_config applied every litellm_settings key: logger __init__s read e.g. s3_callback_params
success_callbacks: Final = tuple(cb for cb in litellm.success_callback if isinstance(cb, str))
failure_callbacks: Final = tuple(cb for cb in litellm.failure_callback if isinstance(cb, str))
for callback in success_callbacks:
_add_custom_logger_callback_to_specific_event(callback, "success")
for callback in failure_callbacks:
_add_custom_logger_callback_to_specific_event(callback, "failure")
async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]):
# only use this if slack alerting is being used
if self.alerting is None:
@ -1585,7 +1746,8 @@ class ProxyLogging:
call_type: str,
event_hook: str,
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
) -> dict:
response: LLMResponseTypes | None = None,
) -> tuple[dict, LLMResponseTypes | None]: # mutable-ok: returns the request-payload dict onward
"""
Execute guardrail pipelines if any are configured for this request.
@ -1597,20 +1759,27 @@ class ProxyLogging:
``scan_raw_request`` evaluates the pristine request, not whatever an
earlier ``pass_data`` step in the same pipeline already rewrote.
Returns the (possibly modified) data dict.
Returns the (possibly modified) data dict, plus the replacement
response when a post_call pipeline step returned one (None when the
response is unchanged), matching the flat callback-loop contract.
"""
pipelines: Final = _policy_pipelines(data)
if not pipelines:
return data
return data, None
current_response = response # rebind-ok: chains each pipeline's replacement response into the next
for policy_name, pipeline in pipelines:
if pipeline.mode != event_hook:
continue
step_input: dict = (
{**data, "response": current_response} if current_response is not None else data
) # mutable-ok: same request-payload shape as data
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode=pipeline.mode,
data=data,
data=step_input,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
policy_name=policy_name,
@ -1621,26 +1790,46 @@ class ProxyLogging:
result=result,
data=data,
policy_name=policy_name,
original_response=current_response,
)
return data
if current_response is not None and result.modified_data is not None:
current_response = result.modified_data.get("response", current_response)
return data, current_response if current_response is not response else None
@staticmethod
def _handle_pipeline_result(
result: PipelineExecutionResult,
data: dict,
policy_name: str,
original_response: "LLMResponseTypes | Sequence[object] | None" = None,
) -> dict:
"""
Handle a PipelineExecutionResult allow, block, or modify_response.
Returns data dict if allowed, raises on block/modify_response.
``original_response`` is set on the post_call path, where the request
payload (already sent upstream) must stay untouched; a replacement
response carried in ``modified_data`` is adopted by the caller, and
metadata-bucket writes (applied guardrails, guardrail logging info)
are merged back so headers and spend logs still see them, on block
and modify_response too, so failure spend records keep guardrail
cost and status. On the
streaming path it is the buffered chunk list, carried into
``ModifyResponseException.original_response`` for usage reporting.
"""
if result.terminal_action == "allow":
if result.modified_data is not None:
data.update(result.modified_data)
if original_response is None:
data.update(result.modified_data)
else:
_merge_pipeline_metadata_writes(data, result.modified_data)
return data
if result.modified_data is not None:
_merge_pipeline_metadata_writes(data, result.modified_data)
if result.terminal_action == "block":
original_exception: Final = result.original_exception
if original_exception is not None and not _exception_changes_request_flow(original_exception):
@ -1678,6 +1867,7 @@ class ProxyLogging:
request_data=data,
guardrail_name=f"pipeline:{policy_name}",
detection_info=None,
original_response=original_response,
)
return data
@ -1794,8 +1984,10 @@ class ProxyLogging:
)
try:
_warn_background_skips_post_call_pipelines(data)
# Execute guardrail pipelines before the normal callback loop
data = await self._maybe_execute_pipelines(
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
@ -1804,7 +1996,7 @@ class ProxyLogging:
)
# Get pipeline-managed guardrails to skip in normal loop
pipeline_managed: Final = _pipeline_managed_guardrail_names(data)
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call")
caps: Final = ProxyLogging._callback_capabilities()
# Skip the per-request callback walk entirely when nothing in
@ -2782,36 +2974,35 @@ class ProxyLogging:
from litellm.proxy.proxy_server import llm_router
from litellm.types.guardrails import GuardrailEventHooks
guardrail_callbacks: Final[list[CustomGuardrail]] = []
other_callbacks: Final[list[CustomLogger]] = []
_, pipeline_response = await self._maybe_execute_pipelines(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
event_hook="post_call",
response=response,
)
if pipeline_response is not None:
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call")
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
try:
for callback in litellm.callbacks:
_callback: CustomLogger | None = None
if isinstance(callback, str):
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
cast(_custom_logger_compatible_callbacks_literal, callback)
)
else:
_callback = callback
if _callback is not None:
if isinstance(_callback, CustomGuardrail):
guardrail_callbacks.append(_callback)
else:
other_callbacks.append(_callback)
############## Handle Guardrails ########################################
#############################################################################
# Merge model-level guardrails before checking which guardrails to run
guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False)
callback
for callback in guardrail_callbacks
if getattr(callback, "run_in_parallel", False)
and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed)
)
for callback in guardrail_callbacks:
# Main - V2 Guardrails implementation
if callback.guardrail_name and callback.guardrail_name in pipeline_managed:
continue
if getattr(callback, "run_in_parallel", False):
continue
@ -3108,11 +3299,16 @@ class ProxyLogging:
# dict lookups + llm_router.get_deployment() per callback per chunk.
_cached_guardrail_data: dict | None = None
_guardrail_data_computed = False
pipeline_gated: Final = (
_stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
)
for callback in litellm.callbacks:
try:
_callback: CustomLogger | None = None
if isinstance(callback, CustomGuardrail):
if callback.guardrail_name in pipeline_gated:
continue
# Main - V2 Guardrails implementation
from litellm.types.guardrails import GuardrailEventHooks
@ -3169,12 +3365,13 @@ class ProxyLogging:
1. /chat/completions
"""
caps: Final = ProxyLogging._callback_capabilities()
post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
# Fast path: no real overrides. Internal proxy CustomLogger callbacks
# (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default
# ``async for chunk: yield chunk`` body, so wrapping the iterator
# through each of them adds N pass-through trampolines per chunk for
# zero behavior change. Skip the chain entirely and stream through.
if not caps.iterator_overrides:
if not caps.iterator_overrides and not post_call_pipelines:
try:
async for chunk in response:
yield chunk
@ -3194,8 +3391,11 @@ class ProxyLogging:
current_response = response
stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines)
for resolved_callback, kind in caps.iterator_overrides:
if isinstance(resolved_callback, CustomGuardrail):
if resolved_callback.guardrail_name in pipeline_gated_names:
continue
if (
resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call)
is not True
@ -3235,6 +3435,14 @@ class ProxyLogging:
),
)
if post_call_pipelines:
current_response = self._pipeline_gated_stream(
response=current_response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
pipelines=post_call_pipelines,
)
try:
async for chunk in current_response:
yield chunk
@ -3250,6 +3458,81 @@ class ProxyLogging:
# we reach this point the metadata is fully populated.
ProxyLogging._fire_deferred_stream_logging(request_data)
async def _pipeline_gated_stream(
self,
response: "AsyncGenerator[object, None]",
user_api_key_dict: UserAPIKeyAuth,
request_data: dict, # mutable-ok: same request-payload shape the hooks mutate
pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
) -> "AsyncGenerator[Any, None]":
"""
Execute post_call policy pipelines against a streamed response.
Buffers the whole stream (nothing reaches the client until every
pipeline allows it), then runs each pipeline's steps against the
assembled output through the endpoint guardrail translation, the same
machinery flat post_call guardrails use at end of stream. An allow
releases the buffered chunks: verbatim when no guardrail rewrote the
output, rewritten in place when one rewrote text and the translation
delivers ended-stream rewrites (later steps then re-scan the rewritten
chunks, so rewrites chain). A rewrite the translation cannot deliver
yet (a tool-call rewrite, or a text rewrite on a route without
write-back) is discarded by the executor and the original chunks are
released, as is a buffered shape no translation resolves; a block or
modify_response terminates with the translation's block chunks or the
raised error.
"""
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
async for item in response:
buffered.append(item)
if not buffered:
return
resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
if resolved is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; "
"the stream is released ungoverned by them: %s",
", ".join(policy_name for policy_name, _pipeline in pipelines),
)
for buffered_item in buffered:
yield buffered_item
return
call_type, endpoint_translation = resolved
for policy_name, pipeline in pipelines:
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
steps=pipeline.steps,
mode="post_call",
data=request_data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
policy_name=policy_name,
streaming_chunks=buffered,
endpoint_translation=endpoint_translation,
)
try:
ProxyLogging._handle_pipeline_result(
result, data=request_data, policy_name=policy_name, original_response=buffered
)
except ModifyResponseException as e:
if e.original_response is None:
e.original_response = buffered
async for block_chunk in unified_guardrail.handle_streaming_block(
e, endpoint_translation, stream_started=False, responses_so_far=()
):
yield block_chunk
return
except HTTPException as e:
async for error_chunk in unified_guardrail.emit_streaming_http_error(
e, call_type, buffered, request_data
):
yield error_chunk
return
for buffered_item in buffered:
yield buffered_item
@staticmethod
def _fire_deferred_stream_logging(request_data: dict) -> None:
"""

View file

@ -17,6 +17,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
from litellm.constants import request_timeout
from litellm.integrations.anthropic_cache_control_hook import CARRY_UNMATCHED_MESSAGE_POINTS
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.prompt_templates.common_utils import (
update_responses_input_with_model_file_ids,
@ -1257,7 +1258,7 @@ def responses(
responses_api_provider_config=responses_api_provider_config,
response_api_optional_params=response_api_optional_params,
allowed_openai_params=allowed_openai_params,
drop_params=request_drop_params if isinstance(request_drop_params, bool) else None,
drop_params=normalize_drop_params(request_drop_params),
)
litellm_logging_obj.update_from_kwargs(
@ -2085,7 +2086,7 @@ def compact_responses(
responses_api_provider_config=responses_api_provider_config,
response_api_optional_params=response_api_optional_params,
allowed_openai_params=None,
drop_params=request_drop_params if isinstance(request_drop_params, bool) else None,
drop_params=normalize_drop_params(request_drop_params),
)
# Pre Call logging

View file

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

View file

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

View file

@ -20,7 +20,7 @@ import random
import re
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import accumulate, islice, takewhile
from itertools import accumulate, chain, islice, takewhile
from threading import Lock
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
@ -31,6 +31,7 @@ from litellm._logging import verbose_router_logger
from litellm.constants import (
EMPTY_MAPPING,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
OUTPUT_TOKEN_CEILING_PARAMS,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
@ -82,6 +83,7 @@ from .config import (
ClassificationRubric,
ComplexityRouterConfig,
ComplexityTier,
CustomDimension,
TierDefinition,
)
from .stall_detector import detect_stalled_task
@ -891,6 +893,15 @@ class DimensionScore:
self.signal = signal
class _CustomDimensionMatchers(NamedTuple):
"""One custom dimension's distinct matchers and the number of hits that saturates its score."""
dimension: CustomDimension
keywords: tuple[str, ...]
patterns: tuple[re.Pattern[str], ...]
saturation: int
class KeywordOverride(NamedTuple):
"""A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired."""
@ -1133,7 +1144,12 @@ class ComplexityRouter(CustomLogger):
)
self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
self._custom_dimensions = tuple(
(dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns))
_CustomDimensionMatchers(
dimension,
tuple(dict.fromkeys(keyword.lower() for keyword in dimension.keywords)),
tuple(re.compile(pattern, re.IGNORECASE) for pattern in dict.fromkeys(dimension.patterns)),
2 if dimension.scoring_mode == "match_count" else 1,
)
for dimension in self.config.custom_dimensions
)
if self.config.has_custom_tiers:
@ -1337,15 +1353,26 @@ class ComplexityRouter(CustomLogger):
score: Final = score_high if match_count >= high_threshold else score_low
return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
def _count_custom_hits(self, matchers: _CustomDimensionMatchers, user_text: str, scanned: str) -> int:
hits: Final = chain(
(self._keyword_matches(user_text, keyword) for keyword in matchers.keywords),
(pattern.search(scanned) is not None for pattern in matchers.patterns),
)
return sum(islice((1 for hit in hits if hit), matchers.saturation))
def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]:
if not self._custom_dimensions:
return ()
scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS]
return tuple(
(DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight)
for dimension, patterns in self._custom_dimensions
if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords)
or any(pattern.search(scanned) is not None for pattern in patterns)
(
DimensionScore(
matchers.dimension.name, hits / matchers.saturation, f"custom ({matchers.dimension.name})"
),
matchers.dimension.weight,
)
for matchers in self._custom_dimensions
if (hits := self._count_custom_hits(matchers, user_text, scanned))
)
def _score_multi_step(self, text: str) -> DimensionScore:
@ -2100,11 +2127,15 @@ class ComplexityRouter(CustomLogger):
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]:
if tier is None:
return MappingProxyType({})
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ())
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else ()
entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None)
return entry.litellm_params if entry is not None else MappingProxyType({})
explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({})
if not self.config.max_tokens_from_tier_model or not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(explicit):
return explicit
ceiling: Final = self._group_output_ceiling(model)
if ceiling is None:
return explicit
return MappingProxyType({**explicit, "max_tokens": ceiling})
@staticmethod
def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str:
@ -2438,12 +2469,15 @@ class ComplexityRouter(CustomLogger):
return name if self.config.has_custom_tiers else ComplexityTier(name)
def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None:
return self._deployment_limit(group, deployment, "max_input_tokens")
def _deployment_limit(
self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"]
) -> int | None:
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
deployment_model_info: Final = deployment.get("model_info")
declared: Final = (
deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None
)
declared: Final = deployment_model_info.get(key) if isinstance(deployment_model_info, Mapping) else None
if isinstance(declared, int):
return declared
litellm_params: Final = deployment.get("litellm_params")
@ -2460,18 +2494,34 @@ class ComplexityRouter(CustomLogger):
deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts
received_model_name=group,
)
window: Final = model_info.get("max_input_tokens")
limit: Final = model_info.get(key)
except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others
return None
return window if isinstance(window, int) else None
return limit if isinstance(limit, int) else None
def _group_deployments(self, group: str) -> Sequence[Mapping[str, object]]:
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
deployments: Final = list_models(model_name=group) if callable(list_models) else None
return tuple(deployments) if isinstance(deployments, list) else ()
def _group_output_ceiling(self, group: str) -> int | None:
"""Smallest max_output_tokens across the group's deployments, or None when any deployment
declares none: the core router picks within the group without a fit check, and a ceiling
above an unmapped member's real limit is a provider 400 on that member."""
deployments: Final = self._group_deployments(group)
ceilings: Final = tuple(
ceiling
for deployment in deployments
if (ceiling := self._deployment_limit(group, deployment, "max_output_tokens")) is not None
)
return min(ceilings) if ceilings and len(ceilings) == len(deployments) else None
def _group_window_facts(self, group: str) -> tuple[int | None, bool]:
"""(smallest declared context window across the group's deployments, whether any deployment
declares none). The core router picks a deployment within the group without a fit check, so
the group is only as safe as its smallest member."""
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
deployments: Final = list_models(model_name=group) if callable(list_models) else None
if not isinstance(deployments, list) or not deployments:
deployments: Final = self._group_deployments(group)
if not deployments:
return (None, True)
windows: Final = tuple(
window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None
@ -3526,6 +3576,7 @@ class ComplexityRouter(CustomLogger):
ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
)
fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM
default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model)
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
@ -3534,7 +3585,9 @@ class ComplexityRouter(CustomLogger):
cause="default_fallback",
tier=fallback_tier,
conversation_continuing=conversation_continuing,
tier_litellm_params=default_tier_params,
),
litellm_params=default_tier_params,
)
ask: Final = user_message or ""
@ -3561,6 +3614,7 @@ class ComplexityRouter(CustomLogger):
_tier_name(plan_floor),
routed_model,
)
plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model)
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
@ -3572,7 +3626,9 @@ class ComplexityRouter(CustomLogger):
matched_keyword=plan_mode_sentinel,
escalation_keyword=escalation_keyword,
escalated=False,
tier_litellm_params=plan_tier_params,
),
litellm_params=plan_tier_params,
)
override: Final = await self._resolve_keyword_tier_override(ask, request_kwargs)
@ -3665,6 +3721,7 @@ class ComplexityRouter(CustomLogger):
outcome.signals,
fallback_model,
)
fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model)
return PreRoutingHookResponse(
model=fallback_model,
messages=messages if has_original_messages else None,
@ -3675,7 +3732,9 @@ class ComplexityRouter(CustomLogger):
signals=outcome.signals,
escalation_keyword=escalation_keyword,
escalated=False,
tier_litellm_params=fallback_tier_params,
),
litellm_params=fallback_tier_params,
)
if self.config.adaptive:
# hard_floor rather than a hard pick, and passed whenever the sentinel is present

View file

@ -681,6 +681,14 @@ class CustomDimension(BaseModel):
weight: float = Field(gt=0, le=1, allow_inf_nan=False)
keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
scoring_mode: Literal["binary", "match_count"] = Field(
default="binary",
description=(
"'binary' scores 1 when any matcher hits. 'match_count' scores 0.5 when one distinct matcher hits and 1 "
"when two or more do; repeated occurrences of one matcher never raise it. Keywords are distinct "
"case-insensitively, patterns by source, and a keyword and a pattern are always distinct from each other."
),
)
@model_validator(mode="after")
def _validate_matchers(self) -> "CustomDimension":
@ -822,8 +830,9 @@ class ComplexityRouterConfig(BaseModel):
default=(),
max_length=16,
description=(
"Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
"when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. "
"Named dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
"when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters; "
"scoring_mode 'match_count' instead grades half weight for one distinct matcher and full for two or more. "
"Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, "
"backreferences and lookarounds are rejected. Conservative work limits include alternation paths, "
"repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. "
@ -1108,6 +1117,20 @@ class ComplexityRouterConfig(BaseModel):
"wording the built-ins don't cover, or after a client release changes its strings."
),
)
max_tokens_from_tier_model: bool = Field(
default=True,
description=(
"Set max_tokens on every routed request to the output ceiling of the tier model it "
"lands on, replacing whatever the caller sent. A caller behind an auto-router cannot "
"pick one value that fits every tier: the smallest tier's ceiling starves a bigger "
"tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The "
"ceiling is the smallest max_output_tokens across the tier model's deployments, read "
"from each deployment's model_info and then the model cost map; a tier model with a "
"deployment whose ceiling is unknown keeps the caller's value. A max_tokens, "
"max_completion_tokens or max_output_tokens in the tier's own litellm_params still "
"wins. Set false to forward the caller's value unchanged."
),
)
route_housekeeping_to_cheapest_tier: bool = Field(
default=True,
description=(

View file

@ -39,7 +39,7 @@ class LowestCostLoggingHandler(CustomLogger):
# ------------
"""
{
{model_group}_map: {
cost_map:{model_group}: {
id: {
f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
}
@ -50,7 +50,7 @@ class LowestCostLoggingHandler(CustomLogger):
current_hour: Final = datetime.now().strftime("%H")
current_minute: Final = datetime.now().strftime("%M")
precise_minute: Final = f"{current_date}-{current_hour}-{current_minute}"
cost_key: Final = f"{model_group}_map"
cost_key: Final = f"cost_map:{model_group}"
total_tokens = 0
@ -112,15 +112,14 @@ class LowestCostLoggingHandler(CustomLogger):
# ------------
"""
{
{model_group}_map: {
cost_map:{model_group}: {
id: {
"cost": [..]
f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
}
}
}
"""
cost_key: Final = f"{model_group}_map"
cost_key: Final = f"cost_map:{model_group}"
current_date: Final = datetime.now().strftime("%Y-%m-%d")
current_hour: Final = datetime.now().strftime("%H")
@ -176,7 +175,7 @@ class LowestCostLoggingHandler(CustomLogger):
"""
Returns a deployment with the lowest cost
"""
cost_key: Final = f"{model_group}_map"
cost_key: Final = f"cost_map:{model_group}"
request_count_dict: Final = await self.router_cache.async_get_cache(key=cost_key) or {}

View file

@ -41,8 +41,7 @@ def simple_shuffle(
############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick #################
for weight_by in ["weight", "rpm", "tpm"]:
weight = healthy_deployments[0].get("litellm_params").get(weight_by, None)
if weight is not None:
if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments):
weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments]
verbose_router_logger.debug("\nweight %s", weights)
total_weight = sum(weights)

View file

@ -13,7 +13,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload
from litellm._logging import verbose_logger
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.types.router import ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors
@ -461,7 +461,10 @@ def _request_tags_after_router_consumption(metadata: object, model: str) -> Sequ
if not isinstance(metadata, Mapping):
return None
typed_metadata: Final[Mapping[str, object]] = metadata
request_tags: Final = _tags_in_metadata(typed_metadata)
request_tags: Final = _tags_in_metadata(
typed_metadata,
key=ROUTING_REQUEST_TAGS_METADATA_KEY if ROUTING_REQUEST_TAGS_METADATA_KEY in typed_metadata else "tags",
)
stamp: Final = typed_metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY)
if not isinstance(stamp, ConsumedRequestTagsStamp) or stamp.model_group != model:
return request_tags
@ -646,7 +649,7 @@ async def get_deployments_for_tag(
return healthy_deployments
def _tags_in_metadata(metadata: object) -> list[str]:
def _tags_in_metadata(metadata: object, key: str = "tags") -> list[str]:
"""
Tags out of a metadata bucket the caller controls the shape of.
@ -657,7 +660,7 @@ def _tags_in_metadata(metadata: object) -> list[str]:
if not isinstance(metadata, Mapping):
return []
typed_metadata: Final[Mapping[str, object]] = metadata
tags: Final = typed_metadata.get("tags")
tags: Final = typed_metadata.get(key)
if isinstance(tags, str) or not isinstance(tags, Sequence):
return []
typed_tags: Final[Sequence[object]] = tags

View file

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

View file

@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import (
DEFAULT_COOLDOWN_TIME_SECONDS,
DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS,
@ -558,9 +559,12 @@ def should_cooldown_based_on_allowed_fails_policy(
When *allowed_fails_override* / *cooldown_time_override* are supplied they
take precedence over the router-level values (used by deployment-level overrides).
The counter lives in the router's shared ``DualCache`` (Redis when configured), so
every worker process increments the same key and the threshold applies fleet-wide.
When *cache_key_suffix* is supplied the fail counter is keyed as
``{deployment}:{cache_key_suffix}`` so that different exception types are
tracked independently per deployment.
``deployment:{deployment}:allowed_fails:{cache_key_suffix}`` so that different
exception types are tracked independently per deployment.
Returns:
- True if fails exceed the allowed limit (should cooldown)
@ -584,16 +588,25 @@ def should_cooldown_based_on_allowed_fails_policy(
else (litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS)
)
cache_key: Final = f"{deployment}:{cache_key_suffix}" if cache_key_suffix else deployment
current_fails: Final = litellm_router_instance.failed_calls.get_cache(key=cache_key) or 0
updated_fails: Final = current_fails + 1
base_key: Final = f"deployment:{deployment}:allowed_fails"
cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key
updated_fails: Final = _increment_allowed_fails(
cache=litellm_router_instance.cache, cache_key=cache_key, ttl=cooldown_time
)
return updated_fails > allowed_fails
if updated_fails > allowed_fails:
return True
else:
litellm_router_instance.failed_calls.set_cache(key=cache_key, value=updated_fails, ttl=cooldown_time)
return False
def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> int:
"""
Return the fleet-wide fail count. ``DualCache.increment_cache`` bumps the in-memory tier
before Redis and re-raises a Redis error, so a Redis outage degrades to this worker's own count.
"""
try:
return cache.increment_cache(key=cache_key, value=1, ttl=ttl)
except Exception as e: # noqa: BLE001 # a Redis outage must not stop failing deployments from cooling down
verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e)
local_fails: Final = cache.get_cache(key=cache_key, local_only=True)
return local_fails if isinstance(local_fails, int) else 0
def _is_allowed_fails_set_on_router(

View file

@ -12,7 +12,9 @@ import httpx
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
if TYPE_CHECKING:
from litellm.router import Router
@ -314,6 +316,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/
stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/
max_retries: int | None = None
drop_params: bool | str | None = None
organization: str | None = None # for openai orgs
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None
litellm_credential_name: str | None = None
@ -404,6 +407,18 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
return filtered
return data
@field_validator("drop_params", mode="before")
@classmethod
def coerce_drop_params(cls, value: object) -> bool | str | None:
normalized: Final = normalize_drop_params(value)
if normalized is not None:
return normalized
if isinstance(value, str):
return value
if value is not None:
verbose_logger.warning("drop_params=%r is not a flag value, treating it as unset", value)
return None
def __contains__(self, key) -> bool:
# Define custom behavior for the 'in' operator
return hasattr(self, key)

View file

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

View file

@ -80,6 +80,7 @@ from litellm.constants import (
PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO,
TOOL_CHOICE_OBJECT_TOKEN_COUNT,
)
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.fallback_generalizations import (
match_capability_generalizations,
)
@ -3239,7 +3240,7 @@ def get_optional_params_transcription(
passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS")
custom_llm_provider = passed_params.pop("custom_llm_provider")
drop_params = passed_params.pop("drop_params")
drop_params = normalize_drop_params(passed_params.pop("drop_params"))
special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs")
for k, v in special_params.items():
passed_params[k] = v
@ -3347,7 +3348,7 @@ def get_optional_params_image_gen(
model = passed_params.pop("model", None)
custom_llm_provider = passed_params.pop("custom_llm_provider")
provider_config = passed_params.pop("provider_config", None)
drop_params = passed_params.pop("drop_params", None)
drop_params = normalize_drop_params(passed_params.pop("drop_params", None))
additional_drop_params = passed_params.pop("additional_drop_params", None)
special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs")
for k, v in special_params.items():
@ -3475,7 +3476,7 @@ def get_optional_params_embeddings(
custom_llm_provider = passed_params.pop("custom_llm_provider", None)
special_params: Final = passed_params.pop("kwargs")
drop_params = passed_params.pop("drop_params", None)
drop_params = normalize_drop_params(passed_params.pop("drop_params", None))
additional_drop_params = passed_params.pop("additional_drop_params", None)
allowed_openai_params = passed_params.pop("allowed_openai_params", None) or []
# Remove function objects from passed_params to avoid JSON serialization errors
@ -4202,6 +4203,7 @@ def get_optional_params(
base_model: str | None = None,
**kwargs,
):
drop_params = normalize_drop_params(drop_params) # rebind-ok: config and DB deployments pass "true" as a string
passed_params: Final = locals().copy()
special_params: Final = passed_params.pop("kwargs")
# Remove base_model from passed_params so it doesn't interfere with
@ -4279,20 +4281,20 @@ def get_optional_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "anthropic_text":
optional_params = litellm.AnthropicTextConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
optional_params = litellm.AnthropicTextConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
@ -4301,14 +4303,14 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "triton":
optional_params = litellm.TritonConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params if drop_params is not None else False,
drop_params=bool(drop_params),
)
elif custom_llm_provider == "maritalk":
@ -4316,35 +4318,35 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "replicate":
optional_params = litellm.ReplicateConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "predibase":
optional_params = litellm.PredibaseConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "huggingface":
optional_params = litellm.HuggingFaceChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "together_ai":
optional_params = litellm.TogetherAIChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "vertex_ai" and (
model in litellm.vertex_chat_models
@ -4358,7 +4360,7 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "gemini":
@ -4366,21 +4368,21 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "vertex_ai_beta" or (custom_llm_provider == "vertex_ai" and "gemini" in model):
optional_params = litellm.VertexGeminiConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif litellm.VertexAIAnthropicConfig.is_supported_model(model=model, custom_llm_provider=custom_llm_provider):
optional_params = litellm.VertexAIAnthropicConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "vertex_ai":
if model in litellm.vertex_mistral_models:
@ -4389,35 +4391,35 @@ def get_optional_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
else:
optional_params = litellm.MistralConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif model in litellm.vertex_ai_ai21_models:
optional_params = litellm.VertexAIAi21Config().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif provider_config is not None:
optional_params = provider_config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
else: # use generic openai-like param mapping
optional_params = litellm.VertexAILlama3Config().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "sagemaker":
@ -4426,7 +4428,7 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "bedrock":
BedrockModelInfo: Final = getattr(sys.modules[__name__], "BedrockModelInfo")
@ -4437,14 +4439,14 @@ def get_optional_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif bedrock_route == "openai":
optional_params = litellm.AmazonBedrockOpenAIConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif "anthropic" in bedrock_base_model and bedrock_route == "invoke":
if bedrock_base_model in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names():
@ -4452,21 +4454,21 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
else:
optional_params = litellm.AmazonAnthropicClaudeConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif provider_config is not None:
optional_params = provider_config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
if bedrock_route == "claude_platform":
optional_params = BedrockModelInfo.map_claude_platform_auth_params(
@ -4477,28 +4479,28 @@ def get_optional_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "ollama":
optional_params = litellm.OllamaConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "ollama_chat":
optional_params = litellm.OllamaChatConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "nlp_cloud":
optional_params = litellm.NLPCloudConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "petals":
@ -4506,35 +4508,35 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "deepinfra":
optional_params = litellm.DeepInfraConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "perplexity" and provider_config is not None:
optional_params = provider_config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "mistral" or custom_llm_provider == "codestral":
optional_params = litellm.MistralConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "text-completion-codestral":
optional_params = litellm.CodestralTextCompletionConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "text-completion-inception":
@ -4542,7 +4544,7 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "databricks":
@ -4550,21 +4552,21 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "nvidia_nim":
optional_params = litellm.NvidiaNimConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "cerebras":
optional_params = litellm.CerebrasConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "xai":
optional_params = litellm.XAIChatConfig().map_openai_params(
@ -4577,77 +4579,77 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "fireworks_ai":
optional_params = litellm.FireworksAIConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "volcengine":
optional_params = litellm.VolcEngineConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "hosted_vllm":
optional_params = litellm.HostedVLLMChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "vllm":
optional_params = litellm.VLLMConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "groq":
optional_params = litellm.GroqChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "bedrock_mantle":
optional_params = litellm.BedrockMantleChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "deepseek":
optional_params = litellm.DeepSeekChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "tencent":
optional_params = litellm.TencentChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "openrouter":
optional_params = litellm.OpenrouterConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "watsonx":
optional_params = litellm.IBMWatsonXChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
# WatsonX-text param check
for param in passed_params:
@ -4660,21 +4662,21 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "openai":
optional_params = litellm.OpenAIConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "nebius":
optional_params = litellm.NebiusConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif custom_llm_provider == "azure":
_azure_detection_model: Final = base_model or model
@ -4683,14 +4685,14 @@ def get_optional_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=_azure_detection_model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=_azure_detection_model):
optional_params = litellm.AzureOpenAIGPT5Config().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=_azure_detection_model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
else:
verbose_logger.debug(
@ -4709,21 +4711,21 @@ def get_optional_params(
optional_params=optional_params,
model=_azure_detection_model,
api_version=api_version,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
elif provider_config is not None:
optional_params = provider_config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
else: # assume passing in params for openai-like api
optional_params = litellm.OpenAILikeChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
drop_params=bool(drop_params),
)
# if user passed in non-default kwargs for specific providers/models, pass them along
optional_params = add_provider_specific_params_to_optional_params(

View file

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

View file

@ -1036,6 +1036,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
created_at DateTime @default(now())
created_by String?
team_id String?
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
updated_at DateTime @updatedAt

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -144,18 +144,20 @@ def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
@pytest.mark.asyncio
async def test_batch_cost_calculator(sample_file_content_dict):
"""
mock litellm.completion_cost to return 0.5
mock batch_cost_calculator to return (0.3, 0.2) per line
we know sample_file_content_dict has 2 successful responses
so we expect the cost to be 0.5 * 2 = 1.0
so we expect the cost to be (0.3 + 0.2) * 2 = 1.0, split 0.6 / 0.4
"""
with patch("litellm.completion_cost", return_value=0.5):
with patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.3, 0.2)):
result = _aggregate_batch_cost_usage_models(
entries=sample_file_content_dict,
custom_llm_provider="openai",
)
assert result.cost == 1.0 # 0.5 * 2 successful responses
assert result.cost == pytest.approx(1.0) # (0.3 + 0.2) * 2 successful responses
assert result.prompt_cost == pytest.approx(0.6)
assert result.completion_cost == pytest.approx(0.4)
def test_get_response_from_batch_job_output_file(sample_file_content_dict):
@ -402,6 +404,56 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data():
assert mock_batch.usage == explicit_usage
@pytest.mark.asyncio
async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown():
"""The poller passes the batch's prompt/completion cost split so the spend row's
cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown
card renders blank for every batch. Regression for the split being dropped."""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.utils import CallTypes, LiteLLMBatch
mock_batch = LiteLLMBatch(
id="batch-breakdown-1",
object="batch",
endpoint="/v1/chat/completions",
errors=None,
input_file_id="file-input-1",
completion_window="24h",
status="completed",
output_file_id="file-output-1",
created_at=1234567890,
)
mock_batch._hidden_params = {}
logging_obj = Logging(
model="gpt-5-mini",
messages=[{"role": "user", "content": "test"}],
stream=False,
call_type=CallTypes.aretrieve_batch.value,
litellm_call_id="test-call-breakdown",
function_id="test-function",
start_time=time.time(),
dynamic_success_callbacks=[],
)
logging_obj.custom_llm_provider = "openai"
await logging_obj.async_success_handler(
result=mock_batch,
start_time=time.time(),
end_time=time.time() + 1,
batch_cost=0.10,
batch_usage=litellm.Usage(prompt_tokens=200, completion_tokens=100, total_tokens=300),
batch_models=["gpt-5-mini"],
batch_prompt_cost=0.06,
batch_completion_cost=0.04,
)
assert logging_obj.cost_breakdown is not None
assert logging_obj.cost_breakdown["input_cost"] == 0.06
assert logging_obj.cost_breakdown["output_cost"] == 0.04
assert logging_obj.cost_breakdown["total_cost"] == 0.10
@pytest.mark.asyncio
async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch():
"""

View file

@ -583,6 +583,90 @@ class TestCheckBatchCost:
assert passed_model_info["input_cost_per_token_batches"] == 2e-06
assert passed_model_info["output_cost_per_token_batches"] == 4e-06
@pytest.mark.asyncio
async def test_poller_masks_api_base_credentials_before_logging(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""Request rows mask `key=` query credentials out of api_base before it is
logged, but the poller skips that pre-call step, so an unmasked deployment
api_base would land verbatim on the batch cost row: regression test for the
poller masking the same way.
"""
import base64
from unittest.mock import patch
import httpx
import respx
from litellm.litellm_core_utils.litellm_logging import Logging
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_job = MagicMock()
mock_job.id = "job-masked-api-base-1"
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
mock_job.created_by = "user-1"
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_response = MagicMock()
mock_response.status = "completed"
mock_response.output_file_id = "file-output-123"
mock_response.error_file_id = None
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
mock_deployment = MagicMock()
mock_deployment.litellm_params.custom_llm_provider = "openai"
mock_deployment.litellm_params.model = "gpt-5.4-mini"
mock_deployment.litellm_params.api_base = "https://gateway.example.com/v1?key=AIzaSyVERYSECRET7890"
mock_deployment.model_info.model_dump.return_value = {}
mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment)
output_line = json.dumps(
{
"custom_id": "req-1",
"response": {
"status_code": 200,
"body": {
"id": "chatcmpl-1",
"object": "chat.completion",
"model": "gpt-5.4-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
},
"error": None,
}
)
with (
respx.mock(assert_all_called=True) as provider,
patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to the row it logs
Logging, "async_success_handler", autospec=True
) as success_handler,
):
provider.get("https://api.openai.com/v1/files/file-output-123/content").mock(
return_value=httpx.Response(200, content=f"{output_line}\n".encode())
)
await check_batch_cost_instance.check_batch_cost()
cost_row_calls = [call for call in success_handler.await_args_list if "batch_cost" in call.kwargs]
assert len(cost_row_calls) == 1
logged_api_base = cost_row_calls[0].args[0].litellm_params["api_base"]
assert logged_api_base == "https://gateway.example.com/v1?key=*****7890"
assert "VERYSECRET" not in logged_api_base
@pytest.mark.asyncio
async def test_primary_path_completion_update_includes_batch_processed(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
@ -2583,6 +2667,85 @@ class TestBatchCostAttribution:
assert metadata["user_api_key_alias"] == "prod-key"
@pytest.mark.asyncio
async def test_org_id_snapshotted_on_the_row_wins(self):
"""The org_id column captures the creating key's organization at submission time,
like team_id, so a key later moved to another org still bills the original one."""
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-moved-to"),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
metadata = await instance._build_creator_attribution_metadata(
self._job(org_id="org-at-creation"), "batch-1"
)
assert metadata["user_api_key_org_id"] == "org-at-creation"
@pytest.mark.asyncio
async def test_org_id_comes_from_the_creating_key(self):
"""The spend update writer increments organization spend from user_api_key_org_id.
A legacy row without the org_id column falls back to the creating key's org."""
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-42"),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-42"
@pytest.mark.asyncio
async def test_org_id_falls_back_to_the_team_organization(self):
"""A key with no org of its own still books batch spend against its team's
organization, matching how the request path resolves org attribution."""
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", organization_id=None),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-team"
@pytest.mark.asyncio
async def test_key_lookup_failure_still_bills_the_team_org(self):
"""A key-table error while resolving a legacy row's org must not drop the team's
organization: the two lookups fail independently, so org spend still lands."""
from types import SimpleNamespace
instance = self._instance(
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
side_effect=Exception("db down")
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-team"
@pytest.mark.asyncio
async def test_no_org_leaves_the_key_unset(self):
"""Without any org the key is absent entirely, so the spend writer's org update
stays skipped instead of matching an empty-string organization."""
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", organization_id=None),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None),
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert "user_api_key_org_id" not in metadata
@pytest.mark.asyncio
async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self):
"""

View file

@ -204,8 +204,8 @@ class TestExceptionTypeCountersTrackedIndependently:
cache_key_suffix="RateLimitError",
)
rl_counter = router.failed_calls.get_cache(key="primary:RateLimitError") or 0
generic_counter = router.failed_calls.get_cache(key="primary:generic") or 0
rl_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
generic_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
assert rl_counter == 3, "RateLimitError counter should be 3"
assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments"
@ -218,8 +218,8 @@ class TestExceptionTypeCountersTrackedIndependently:
cache_key_suffix="generic",
)
generic_counter_after = router.failed_calls.get_cache(key="primary:generic") or 0
rl_counter_after = router.failed_calls.get_cache(key="primary:RateLimitError") or 0
generic_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
rl_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
assert generic_counter_after == 1, "generic counter should now be 1"
assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError"

View file

@ -490,7 +490,9 @@ def test_aggregate_counts_successful_and_failed_requests(monkeypatch):
def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0)
import litellm.cost_calculator as cc
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6))
result = bu._aggregate_batch_cost_usage_models(
entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai"
)
@ -501,6 +503,7 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
1,
0,
)
assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6)
# =========================================================================== #
@ -508,15 +511,17 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
# =========================================================================== #
def test_cost_from_content_completion_cost_path(monkeypatch):
# model_info is None -> litellm.completion_cost per successful row.
def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch):
# model_info is None -> batch_cost_calculator per successful row, model from the response body.
import litellm.cost_calculator as cc
calls = []
def _completion_cost(**kw):
def _batch_cost(**kw):
calls.append(kw)
return 0.5
return (0.3, 0.2)
monkeypatch.setattr(litellm, "completion_cost", _completion_cost)
monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost)
rows = [
_success_row(usage=_usage(10, 5)),
_failed_row(), # excluded -> not costed
@ -525,8 +530,10 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert result.cost == 1.0 # 2 successful * 0.5
assert result.cost == pytest.approx(1.0) # 2 successful * (0.3 + 0.2)
assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4))
assert len(calls) == 2 # failed row not costed
assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls)
assert result.successful_requests == 2
assert result.failed_requests == 1
@ -579,7 +586,9 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
"""A one-shot generator: any implementation that iterates the entries twice
(e.g. separate cost and usage passes) sees nothing on the second pass and
returns wrong totals for at least one of cost/usage/models."""
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5)
import litellm.cost_calculator as cc
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25))
one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))])
result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
@ -754,12 +763,15 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
@pytest.mark.asyncio
async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
import litellm.cost_calculator as cc
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5)
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0))
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")
assert result.cost == 2.5
assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0)
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
assert result.models == ["gpt-4o"]
@ -1108,8 +1120,10 @@ async def test_handle_completed_batch_orchestration(monkeypatch):
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return _vertex_jsonl(rows)
import litellm.cost_calculator as cc
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3)
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3))
result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")

View file

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

View file

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

View file

@ -375,6 +375,7 @@ def _in_memory_managed_files():
table.upsert = AsyncMock(side_effect=_upsert)
prisma = MagicMock()
prisma.db.litellm_managedobjecttable = table
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
cache = MagicMock()
cache.async_set_cache = AsyncMock()
@ -390,7 +391,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create():
"""Regression (spend loss): the batch create persists the creating key hash and tags so
CheckBatchCost can write an attributed spend row instead of a blank one the DB drops."""
instance, store = _in_memory_managed_files()
creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice")
creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice", org_id="org-acme")
await instance.store_unified_object_id(
unified_object_id="unified-b",
@ -407,9 +408,70 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create():
assert row["api_key"] == "hash-alice"
assert row["created_by"] == "alice"
assert row["team_id"] == "team-alpha"
assert row["org_id"] == "org-acme"
assert row["request_tags"].data == ["env:prod"]
@pytest.mark.asyncio
async def test_store_unified_object_id_resolves_org_through_the_cached_team():
"""Most keys belong to an org only through their team, so the auth object carries no
org_id. The create reads the team that auth already cached, so org spend is snapshotted
at submission time without a database query in the request path."""
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy.proxy_server import user_api_key_cache
instance, store = _in_memory_managed_files()
creator = UserAPIKeyAuth(user_id="alice", team_id="team-cached", api_key="hash-alice")
await user_api_key_cache.async_set_cache(
key="team_id:team-cached",
value=LiteLLM_TeamTableCachedObj(team_id="team-cached", organization_id="org-via-team"),
model_type=LiteLLM_TeamTableCachedObj,
)
try:
await instance.store_unified_object_id(
unified_object_id="unified-b",
file_object=_build_batch_response(batch_id="b", status="validating"),
litellm_parent_otel_span=None,
model_object_id="b",
file_purpose="batch",
user_api_key_dict=creator,
persist_attribution=True,
)
finally:
user_api_key_cache.delete_cache(key="team_id:team-cached")
assert store["unified-b"]["org_id"] == "org-via-team"
instance.prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_store_unified_object_id_resolves_org_from_the_db_when_the_team_is_not_cached():
"""A team no request has run under yet is absent from the auth cache; its organization
still comes back from the table so the org is billed rather than dropped."""
from litellm.models.team import LiteLLM_TeamTable
from litellm.proxy.proxy_server import user_api_key_cache
instance, store = _in_memory_managed_files()
instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=LiteLLM_TeamTable(team_id="team-uncached", organization_id="org-via-db")
)
creator = UserAPIKeyAuth(user_id="alice", team_id="team-uncached", api_key="hash-alice")
try:
await instance.store_unified_object_id(
unified_object_id="unified-b",
file_object=_build_batch_response(batch_id="b", status="validating"),
litellm_parent_otel_span=None,
model_object_id="b",
file_purpose="batch",
user_api_key_dict=creator,
persist_attribution=True,
)
finally:
user_api_key_cache.delete_cache(key="team_id:team-uncached")
assert store["unified-b"]["org_id"] == "org-via-db"
@pytest.mark.asyncio
async def test_store_unified_object_id_omits_key_and_tags_without_persist_attribution():
"""Regression (spend redirect): a caller that is not the batch create (a poll, or the
@ -471,6 +533,7 @@ async def test_store_unified_object_id_attribution_columns_are_write_once():
upsert_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"]
assert "api_key" not in upsert_data["update"]
assert "request_tags" not in upsert_data["update"]
assert "org_id" not in upsert_data["update"]
@pytest.mark.asyncio

View file

@ -3,7 +3,10 @@ baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name
builders, and the registry validator's failure paths. Needs the OTel SDK."""
import json
import threading
from collections.abc import Iterator
from dataclasses import replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import pytest
@ -538,6 +541,77 @@ def test_build_span_exporter_variants():
assert "OTLPSpanExporter" in type(http_exporter).__name__
@pytest.fixture
def otlp_collector() -> Iterator[tuple[str, list[str]]]:
received_paths: list[str] = []
class RecordingHandler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
self.rfile.read(int(self.headers.get("Content-Length", "0")))
received_paths.append(self.path)
self.send_response(200)
self.end_headers()
def log_message(self, format: str, *args: object) -> None:
pass
server = ThreadingHTTPServer(("127.0.0.1", 0), RecordingHandler)
threading.Thread(target=server.serve_forever, daemon=True).start()
try:
yield f"http://127.0.0.1:{server.server_port}", received_paths
finally:
server.shutdown()
server.server_close()
def _export_one_span(cfg: OpenTelemetryV2Config) -> None:
provider = providers.build_tracer_provider(cfg)
provider.get_tracer("probe").start_span("probe").end()
assert provider.force_flush()
provider.shutdown()
def test_traces_endpoint_env_posts_to_the_configured_url_verbatim(monkeypatch, otlp_collector):
base_url, received_paths = otlp_collector
for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_EXPORTER_OTLP_ENDPOINT"):
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("OTEL_ENDPOINT", f"{base_url}/services/collector")
monkeypatch.setenv("OTEL_TRACES_ENDPOINT", f"{base_url}/services/collector/traces")
cfg = OpenTelemetryV2Config.from_env()
assert cfg.exporter == "otlp_http"
_export_one_span(cfg)
assert received_paths == ["/services/collector/traces"]
def test_traces_endpoint_alias_alone_implies_otlp_http(monkeypatch, otlp_collector):
base_url, received_paths = otlp_collector
for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"):
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", f"{base_url}/custom/traces")
cfg = OpenTelemetryV2Config.from_env()
assert cfg.exporter == "otlp_http"
_export_one_span(cfg)
assert received_paths == ["/custom/traces"]
def test_traces_endpoint_per_exporter_coexists_with_default_normalization(otlp_collector):
base_url, received_paths = otlp_collector
cfg = OpenTelemetryV2Config(
exporters=[
{"kind": "otlp_http", "endpoint": base_url},
{
"kind": "otlp_http",
"endpoint": f"{base_url}/services/collector",
"traces_endpoint": f"{base_url}/services/collector/traces",
},
]
)
_export_one_span(cfg)
assert sorted(received_paths) == ["/services/collector/traces", "/v1/traces"]
def test_otlp_metric_exporter_uses_cumulative_histogram_temporality():
"""Histograms must export as cumulative, not delta.

View file

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

View file

@ -1,11 +1,16 @@
"""Tests for litellm_core_utils.core_helpers module."""
import logging
import pytest
from litellm.litellm_core_utils.core_helpers import (
_FINISH_REASON_MAP,
drop_params_env_flag,
drop_params_flag,
get_or_create_metadata_bucket,
map_finish_reason,
normalize_drop_params,
reconstruct_model_name,
redact_nested_match_and_regex_keys,
)
@ -257,6 +262,73 @@ class TestRedactNestedMatchAndRegexKeys:
assert redact_nested_match_and_regex_keys("plain") == "plain"
@pytest.mark.parametrize(
"value, expected",
[
(True, True),
(False, False),
("true", True),
("True", True),
(" TRUE ", True),
("false", False),
("False", False),
("yes", True),
("off", False),
("1", True),
(1, True),
(0, False),
(None, None),
("", None),
("os.environ/DROP_PARAMS", None),
("v2:gcm:not-a-flag", None),
(2, None),
],
)
def test_normalize_drop_params(value, expected):
assert normalize_drop_params(value) is expected
@pytest.mark.parametrize("value, expected", [("true", True), ("off", False), (None, False)])
def test_drop_params_flag_returns_a_bool_without_a_warning(value, expected, caplog):
with caplog.at_level(logging.WARNING, logger="drop-params-test"):
assert drop_params_flag(value, "LITELLM_DROP_PARAMS", logging.getLogger("drop-params-test")) is expected
assert caplog.text == ""
@pytest.mark.parametrize("value", ["temperature", "ture", 2])
def test_drop_params_flag_treats_non_flag_values_as_off_with_a_warning(value, caplog):
with caplog.at_level(logging.WARNING, logger="drop-params-test"):
assert drop_params_flag(value, "LITELLM_DROP_PARAMS", logging.getLogger("drop-params-test")) is False
assert f"LITELLM_DROP_PARAMS={value!r} is not a flag value, treating it as off" in caplog.text
@pytest.mark.parametrize(
"environ, expected",
[
({}, False),
({"LITELLM_DROP_PARAMS": ""}, False),
({"LITELLM_DROP_PARAMS": " "}, False),
({"LITELLM_DROP_PARAMS": "true"}, True),
({"LITELLM_DROP_PARAMS": " False "}, False),
({"LITELLM_DROP_PARAMS": "0"}, False),
],
)
def test_drop_params_env_flag_reads_a_flag_without_a_warning(environ, expected, caplog):
with caplog.at_level(logging.WARNING, logger="drop-params-test"):
assert drop_params_env_flag(environ, logging.getLogger("drop-params-test")) is expected
assert caplog.text == ""
@pytest.mark.parametrize("configured", ["temperature", "temperature,top_p", "enabled"])
def test_drop_params_env_flag_keeps_a_non_flag_value_on_with_a_warning(configured, caplog):
with caplog.at_level(logging.WARNING, logger="drop-params-test"):
assert drop_params_env_flag({"LITELLM_DROP_PARAMS": configured}, logging.getLogger("drop-params-test")) is True
assert (
f"LITELLM_DROP_PARAMS={configured!r} is not a flag value, treating it as on. Set it to true or false"
in caplog.text
)
class TestIsExpectedClientError:
def test_status_ranges(self):
from litellm.litellm_core_utils.core_helpers import is_expected_client_error

View file

@ -215,3 +215,11 @@ class TestMetadataFallsBackToLitellmMetadata:
assert result["metadata"] is not litellm_metadata
result["metadata"].pop("trace_id")
assert litellm_metadata == {"trace_id": "trace-1"}
@pytest.mark.parametrize(
"value, expected",
[("true", True), ("false", False), (" TRUE ", True), (True, True), (None, None), ("os.environ/DROP_PARAMS", None)],
)
def test_drop_params_strings_reach_litellm_params_as_flags(value, expected):
assert get_litellm_params(drop_params=value)["drop_params"] is expected

View file

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

View file

@ -264,6 +264,117 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
# Should return the responses unchanged
assert result == responses_so_far
@staticmethod
def _ended_sse_chunks() -> list:
events = [
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@staticmethod
def _masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs.get("texts", [])]}
return MaskWorld(guardrail_name="test")
@staticmethod
def _delta_texts(chunks: list) -> list:
texts = []
for chunk in chunks:
for line in chunk.decode().split("\n"):
if not line.startswith("data:"):
continue
data = json.loads(line[len("data:") :].strip())
if data.get("type") == "content_block_delta":
texts.append(data["delta"]["text"])
return texts
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert self._delta_texts(chunks) == ["hello [MASKED]", ""]
raw = b"".join(chunks).decode()
assert "event: message_start" in raw and "event: message_stop" in raw
assert '"stop_reason": "end_turn"' in raw
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()
original = [bytes(chunk) for chunk in chunks]
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
original = [bytes(chunk) for chunk in chunks]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks == original
@pytest.mark.asyncio
async def test_unended_stream_rewrite_without_delivery_expected_does_not_raise(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_sse_chunks()[:-2]
original = [bytes(chunk) for chunk in chunks]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert result is chunks
assert chunks == original
class TestAnthropicMessagesHandlerInputProcessing:
"""Test input processing preserves litellm_metadata for dynamic guardrails."""

View file

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

View file

@ -1074,6 +1074,154 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
# Should return the responses
assert result == responses_so_far
@staticmethod
def _ended_stream_chunks() -> list:
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
return [
ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=0, delta=Delta(content="Hello"), finish_reason=None)],
),
ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")],
),
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_text_back_into_chunks(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_stream_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks[0].choices[0].delta.content == "HELLO WORLD"
assert chunks[1].choices[0].delta.content in (None, "")
assert chunks[1].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert chunks[0].choices[0].delta.content == "Hello"
assert chunks[1].choices[0].delta.content == " world"
assert chunks[1].choices[0].finish_reason == "stop"
@staticmethod
def _two_choice_stream_chunks() -> list:
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)],
)
return [
chunk(0, "safe "),
chunk(1, "hello "),
chunk(0, "text", "stop"),
chunk(1, "world", "stop"),
]
@staticmethod
def _world_masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
texts = inputs.get("texts", [])
return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
return MaskWorld(guardrail_name="test-mask")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_stream_chunks()
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._world_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_stream_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
handler = OpenAIChatCompletionsHandler()
def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)],
)
chunks = [chunk("hello ", None), chunk("world", "stop")]
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._world_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert chunks[0].choices[0].delta.content == "hello [MASKED]"
assert chunks[1].choices[0].delta.content in (None, "")
class TestUndecoratedGuardrailIsRecorded:
"""LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail

View file

@ -1128,6 +1128,209 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
output_text = result[-1]["response"]["output"][0]["content"][0]["text"]
assert output_text == original_text
@staticmethod
def _ended_stream_events() -> List[dict]:
content = [{"type": "output_text", "text": "hello world"}]
item = {
"type": "message",
"id": "msg_123",
"status": "completed",
"role": "assistant",
"content": content,
}
return [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
{
"type": "response.content_part.done",
"output_index": 0,
"content_index": 0,
"part": {"type": "output_text", "text": "hello world"},
},
{"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"model": "gpt-4o",
"output": [{**item, "content": [dict(c) for c in content]}],
"status": "completed",
},
},
]
@staticmethod
def _masking_guardrail() -> CustomGuardrail:
class MaskWorld(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
texts = inputs.get("texts", [])
return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]}
return MaskWorld(guardrail_name="test-mask")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_all_stream_events(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["delta"] == ""
assert events[2]["text"] == "hello [MASKED]"
assert events[3]["part"]["text"] == "hello [MASKED]"
assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
@pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
events[-1]["type"] = terminal_type
events[-1]["response"]["status"] = terminal_type.split(".")[-1]
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["delta"] == "hello [MASKED]"
assert events[1]["delta"] == ""
assert events[2]["text"] == "hello [MASKED]"
assert events[3]["part"]["text"] == "hello [MASKED]"
assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_fallback_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "},
{"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"},
]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_output_item_done_last_rewrite_with_delivery_expected_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_output_item_done_last_scans_text_with_delivery_expected(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
guardrail = MockRecordingGuardrail(guardrail_name="test")
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]]
@pytest.mark.asyncio
async def test_output_item_done_last_without_delivery_expected_skips_text(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()[:-1]
guardrail = MockRecordingGuardrail(guardrail_name="test")
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
assert result is events
assert guardrail.seen_inputs == []
@pytest.mark.asyncio
async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self):
handler = OpenAIResponsesHandler()
events = [
{"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"},
]
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
)
assert result is events
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self):
handler = OpenAIResponsesHandler()
events = self._ended_stream_events()
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._masking_guardrail(),
litellm_logging_obj=None,
)
assert events[0]["delta"] == "hello "
assert events[1]["delta"] == "world"
assert events[2]["text"] == "hello world"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@pytest.mark.asyncio
async def test_failed_stream_scans_delta_text(self):
"""A stream ending in response.failed has text only in delta events; the

View file

@ -6767,6 +6767,109 @@ async def test_temp_budget_increase_applied_for_cached_key():
assert cached_after.max_budget == 2.0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"team_member_spend, expect_blocked",
[
(2.4, True),
(2.4000000000000004, True),
(2.39, False),
],
)
async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spend, expect_blocked):
"""A team member counter sitting exactly at the cap (where a resized reservation
lands it) must be rejected by the cached-key auth path like every other budget check."""
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
from litellm.proxy.utils import hash_token
api_key = "sk-team-member-exact-cap"
hashed_token = hash_token(api_key)
team_id = "team-exact-cap"
user_id = "user-exact-cap"
max_budget = 2.4
user_api_key_cache = DualCache()
await _cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=UserAPIKeyAuth(
token=hashed_token,
team_id=team_id,
user_id=user_id,
team_member_spend=team_member_spend,
),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=None,
)
await user_api_key_cache.async_set_cache(
key=f"team_id:{team_id}",
value=LiteLLM_TeamTableCachedObj(team_id=team_id),
)
await user_api_key_cache.async_set_cache(
key=user_id,
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
)
await user_api_key_cache.async_set_cache(
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
value=LiteLLM_TeamMembership(
user_id=user_id,
team_id=team_id,
spend=team_member_spend,
budget_id="budget-exact-cap",
litellm_budget_table=LiteLLM_BudgetTable(max_budget=max_budget),
),
)
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {api_key}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
async def _auth():
return await _user_api_key_auth_builder(
request=mock_request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]},
)
with (
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
"litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True}
),
patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state
patch( # test-quality-ok: seed the cached key, team and membership without a DB
"litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache
),
patch( # test-quality-ok: module-global proxy state
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
),
patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=team_member_spend),
),
):
if not expect_blocked:
result = await _auth()
assert result.team_member_spend == team_member_spend
return
with pytest.raises(ProxyException) as exc_info:
await _auth()
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
async def _proxy_exception_for_key(
api_key: str,
general_settings: dict[str, bool],

View file

@ -1034,7 +1034,7 @@ class TestStreamingTransform:
)
emitted = []
async for item in handler._emit_streaming_http_error(
async for item in handler.emit_streaming_http_error(
exc,
call_type=CallTypes.asend_message.value,
responses_so_far=[{"id": "req-1"}],

View file

@ -30,7 +30,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
)
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams
async def _passthrough_row(update_data):
@ -3070,6 +3070,31 @@ class TestUpdateDBModelBlocked:
assert "blocked" not in result
class TestUpdateDBModelKeepsLegacyDropParams:
def test_partial_patch_keeps_encrypted_string_drop_params(self, monkeypatch):
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
legacy_row = Deployment(
model_name="gpt-5-nano",
litellm_params=LiteLLM_Params(
model="openai/gpt-5-nano",
api_key=encrypt_value_helper(value="sk-old"),
drop_params=encrypt_value_helper(value="true"),
),
model_info=ModelInfo(id="legacy-row"),
)
result = update_db_model(
db_model=legacy_row,
updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(api_key="sk-new")),
)
stored = json.loads(result["litellm_params"])
assert decrypt_value_helper(value=stored["drop_params"], key="drop_params") == "true"
def _build_db_model_with_pricing():
"""Wildcard deployment with custom pricing in litellm_params; Deployment.__init__
mirrors SPECIAL_MODEL_INFO_PARAMS into model_info, so both blobs hold the rate."""

View file

@ -5,6 +5,7 @@ Uses mock guardrails to validate pipeline execution without external services.
"""
import copy
import logging
from typing import Literal
from unittest.mock import MagicMock
@ -17,7 +18,7 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeGuardrail,
)
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
@ -1052,3 +1053,244 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch):
assert outcome == "pass"
assert guardrail.native_pre_call_ran is True
assert "guardrail_to_apply" not in data
class _TextReturningGuardrail(CustomGuardrail):
def __init__(self, returned_texts):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
self.returned_texts = returned_texts
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": self.returned_texts}
class _TextTranslation:
delivers_ended_stream_text_rewrites = False
def __init__(self):
self.seen_guardrail_names = []
async def process_output_streaming_response(
self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None
):
self.seen_guardrail_names.append(guardrail_to_apply.guardrail_name)
await guardrail_to_apply.apply_guardrail(
inputs={"texts": ["hello world"]},
request_data=request_data or {},
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
class _WritingTranslation:
"""Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the
chat/Responses/Messages handlers do on an ended stream."""
delivers_ended_stream_text_rewrites = True
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj=None,
user_api_key_dict=None,
request_data=None,
deliver_ended_stream_rewrites=False,
):
assert deliver_ended_stream_rewrites is True
outputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]},
request_data=request_data or {},
input_type="response",
logging_obj=litellm_logging_obj,
)
responses_so_far[0]["text"] = outputs["texts"][0]
responses_so_far[0]["tool_call"] = outputs["tool_calls"][0]
return responses_so_far
class _RefusingTranslation:
delivers_ended_stream_text_rewrites = True
async def process_output_streaming_response(
self,
responses_so_far,
guardrail_to_apply,
litellm_logging_obj=None,
user_api_key_dict=None,
request_data=None,
deliver_ended_stream_rewrites=False,
):
responses_so_far[0]["text"] = "half-written"
raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name)
def _chunk():
return {"text": "hello world", "tool_call": {"function": {"name": "lookup", "arguments": '{"ssn": "123"}'}}}
async def _run_streaming_step(translation, streaming_chunks=None):
chunks = [object()] if streaming_chunks is None else streaming_chunks
return await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")],
mode="post_call",
data={"model": "m"},
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="p",
streaming_chunks=chunks,
endpoint_translation=translation,
)
def _assert_passed_with_discard_warning(result, caplog):
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
translation = _TextTranslation()
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(translation, chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
assert translation.seen_guardrail_names == ["masker"]
@pytest.mark.asyncio
async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))])
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_TextTranslation())
assert result.terminal_action == "allow"
assert [step.outcome for step in result.step_results] == ["pass"]
assert not any("discarded" in record.getMessage() for record in caplog.records)
class _InPlaceMutatingGuardrail(CustomGuardrail):
"""Rewrites like bedrock/presidio do: rebinds inputs["texts"] on the dict it was handed
and returns that same dict, so a post-call comparison against inputs sees no change."""
def __init__(self):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
inputs["texts"] = ["hello [MASKED]"]
return inputs
@pytest.mark.asyncio
async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_InPlaceMutatingGuardrail()])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_TextTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
class _TextAndToolCallRewritingGuardrail(CustomGuardrail):
def __init__(self, rewrite_tool_call):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
self.rewrite_tool_call = rewrite_tool_call
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
tool_calls = (
[{"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}]
if self.rewrite_tool_call
else inputs["tool_calls"]
)
return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": tool_calls}
@pytest.mark.asyncio
async def test_streaming_step_delivers_text_rewrite_through_writing_translation(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=False)])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_WritingTranslation(), chunks)
assert result.terminal_action == "allow"
assert chunks[0]["text"] == "hello [MASKED]"
assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "123"}'
assert not any("discarded" in record.getMessage() for record in caplog.records)
@pytest.mark.asyncio
async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_WritingTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
class _BlockingStreamGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
raise HTTPException(status_code=400, detail={"error": "output blocked"})
def _recorded_guardrail_statuses(result):
return [
entry["guardrail_status"]
for entry in result.modified_data["metadata"]["standard_logging_guardrail_information"]
]
@pytest.mark.asyncio
async def test_streaming_step_records_guardrail_information_once_on_mask(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
assert result.terminal_action == "allow"
assert _recorded_guardrail_statuses(result) == ["success"]
@pytest.mark.asyncio
async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_header(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"]
@pytest.mark.asyncio
async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()])
result = await _run_streaming_step(_WritingTranslation(), [_chunk()])
assert [step.outcome for step in result.step_results] == ["fail"]
assert _recorded_guardrail_statuses(result) == ["guardrail_intervened"]
@pytest.mark.asyncio
async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewrite(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_RefusingTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]

View file

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

View file

@ -20,6 +20,7 @@ import pytest
import litellm
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.proxy_server import (
ProxyConfig,
_is_remote_module_url,
@ -2428,6 +2429,111 @@ def test_ProxyConfig__add_deployment_resolves_env_refs_on_arbitrary_field(monkey
assert deployment.litellm_params.some_future_field == "resolved-custom-value"
@pytest.mark.parametrize(
"stored_drop_params",
["true", "os.environ/DROP_PARAMS_FLAG"],
)
def test_ProxyConfig__add_deployment_turns_stored_drop_params_string_into_bool(monkeypatch, stored_drop_params):
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
monkeypatch.setenv("DROP_PARAMS_FLAG", "true")
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
db_model = SimpleNamespace(
model_id="model-1",
model_name="gpt-5-nano",
model_info={"id": "model-1"},
litellm_params={
"model": encrypt_value_helper(value="openai/gpt-5-nano"),
"drop_params": encrypt_value_helper(value=stored_drop_params),
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model])
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.drop_params is True
def test_ProxyConfig__add_deployment_keeps_loading_rows_after_a_non_flag_drop_params(monkeypatch):
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
def db_model(model_id, drop_params):
return SimpleNamespace(
model_id=model_id,
model_name="gpt-5-nano",
model_info={"id": model_id},
litellm_params={
"model": encrypt_value_helper(value="openai/gpt-5-nano"),
"drop_params": encrypt_value_helper(value=drop_params),
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model("bad-row", 2), db_model("good-after", "true")])
deployments = [call.kwargs["deployment"] for call in fake_router.upsert_deployment.call_args_list]
assert added == 2
assert [d.litellm_params.drop_params for d in deployments] == [None, True]
@pytest.mark.asyncio
@pytest.mark.parametrize("configured, expected", [("true", True), ("false", False)])
async def test_ProxyConfig_load_config_turns_litellm_settings_drop_params_string_into_bool(
tmp_path, monkeypatch, configured, expected
):
f = tmp_path / "c.yaml"
f.write_text(f'model_list: []\nlitellm_settings:\n drop_params: "{configured}"\n')
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr(litellm, "drop_params", not expected)
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.drop_params is expected
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_resolves_a_litellm_settings_drop_params_env_ref(tmp_path, monkeypatch):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\nlitellm_settings:\n drop_params: os.environ/DROP_PARAMS_FROM_ENV\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setenv("DROP_PARAMS_FROM_ENV", "true")
monkeypatch.setattr(litellm, "drop_params", False)
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.drop_params is True
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_warns_and_turns_off_a_non_flag_litellm_settings_drop_params(
tmp_path, monkeypatch, caplog
):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\nlitellm_settings:\n drop_params: ture\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr(litellm, "drop_params", True)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.drop_params is False
assert "litellm_settings.drop_params='ture' is not a flag value, treating it as off" in caplog.text
# ---------------------------------------------------------------------------
# ProxyConfig.decrypt_model_list_from_db
# ---------------------------------------------------------------------------

View file

@ -132,6 +132,68 @@ def test_legacy_policy_keeps_trace_id_fallback():
assert len(str(generated)) == 36
def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id():
"""The create call's request id IS the batch id and the poller's cost row appends
_batch_cost to it, so deriving the session from the request id lands both rows in one
trace on the logs UI even though the poller builds a fresh logging context per cycle."""
from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id
create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1")
cost_session: Final = _get_batch_trace_session_id(
call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost"
)
assert create_session == cost_session == "batch-uid-1"
def test_non_batch_call_types_derive_no_batch_session():
from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id
assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None
def test_batch_session_outranks_the_per_request_trace_id():
"""Each batch lifecycle call carries its own auto-generated trace id, so letting the
trace id win would scatter the rows across sessions again."""
session_id: Final = _get_session_id_for_spend_log(
kwargs={"litellm_trace_id": "trace-abc"},
metadata={"trace_id": "trace-abc"},
standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING,
omit_when_missing=False,
batch_trace_session_id="batch-uid-1",
)
assert session_id == "batch-uid-1"
def test_omit_policy_still_suppresses_batch_sessions():
session_id: Final = _get_session_id_for_spend_log(
kwargs={},
metadata=None,
standard_logging_payload=None,
omit_when_missing=True,
batch_trace_session_id="batch-uid-1",
)
assert session_id is None
def test_get_logging_payload_groups_batch_create_and_cost_rows_in_one_session():
def _payload(call_type: str) -> SpendLogsPayload:
return get_logging_payload(
kwargs={
"call_type": call_type,
"model": "gpt-4o-mini",
"litellm_params": {"metadata": {"user_api_key": "test-key"}},
},
response_obj=litellm.ModelResponse(id="batch-uid-1", choices=[], usage=litellm.Usage()),
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
create_payload: Final = _payload("acreate_batch")
cost_payload: Final = _payload("aretrieve_batch")
assert cost_payload["request_id"] == "batch-uid-1_batch_cost"
assert create_payload["session_id"] == cost_payload["session_id"] == "batch-uid-1"
@pytest.mark.parametrize(
("request_metadata", "expected"),
[

View file

@ -2198,6 +2198,74 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state):
assert reservation["finalized"] is True
class _ExpiringRedisCache:
def __init__(self) -> None:
self.store: dict[str, float] = {}
async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None:
return self.store.get(key)
async def async_increment(self, key: str, value: float, **kwargs: object) -> float:
self.store[key] = self.store.get(key, 0.0) + float(value)
return self.store[key]
async def async_set_max(self, key: str, value: float, **kwargs: object) -> float:
self.store[key] = max(self.store.get(key, float("-inf")), float(value))
return self.store[key]
async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool:
self.store[key] = float(value)
return True
async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None:
self.store.pop(key, None)
@pytest.mark.asyncio
async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced(
spend_counter_state,
):
"""Redis key expired mid-stream while the pod's in-memory copy still holds the
reserved value: reconcile must reseed from the DB floor plus the settled cost
instead of applying ``actual - reserved`` to the empty key."""
import litellm.proxy.proxy_server as ps
counter_cache, _ = spend_counter_state
counter_key = "spend:team_member:user-expiry:team-expiry"
redis_cache = _ExpiringRedisCache()
counter_cache.redis_cache = redis_cache
counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6)
reservation = {
"reserved_cost": 0.6,
"entries": [
{
"counter_key": counter_key,
"entity_type": "TeamMember",
"entity_id": "user-expiry:team-expiry",
"reserved_cost": 0.6,
"applied_adjustment": 0.0,
}
],
"finalized": False,
}
with patch.object( # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for
ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.3)
):
await ps.increment_spend_counters(
token="key-expiry",
team_id="team-expiry",
user_id="user-expiry",
response_cost=0.05,
budget_reservation=reservation,
)
assert redis_cache.store[counter_key] == pytest.approx(0.35)
assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.35)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_should_invalidate_reserved_counters_after_persisted_spend_failure(
spend_counter_state,

View file

@ -4148,6 +4148,48 @@ async def test_add_guardrails_from_policy_engine():
attachment_registry._initialized = False
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps():
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import (
GuardrailPipeline,
PipelineStep,
Policy,
PolicyAttachment,
PolicyGuardrails,
)
data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "metadata": {}}
policy_registry = get_policy_registry()
policy_registry._policies = {
"response-governance": Policy(
guardrails=PolicyGuardrails(add=["pii_blocker"]),
pipeline=GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii_blocker")]),
),
}
policy_registry._initialized = True
attachment_registry = get_attachment_registry()
attachment_registry._attachments = [PolicyAttachment(policy="response-governance", scope="*")]
attachment_registry._initialized = True
try:
await add_guardrails_from_policy_engine(
data=data,
metadata_variable_name="metadata",
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
)
finally:
policy_registry._policies = {}
policy_registry._initialized = False
attachment_registry._attachments = []
attachment_registry._initialized = False
assert data["metadata"]["guardrails"] == ["pii_blocker"]
assert data["metadata"]["_pipeline_managed_guardrails"] == {"pii_blocker"}
assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"]
@pytest.mark.asyncio
async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data():
"""
@ -7272,7 +7314,12 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth:
)
_PLANTED_STAMPS = {"attempted_fallbacks": 99, "original_model_group": "spoofed-group", "client_key": "client_value"}
_PLANTED_STAMPS = {
"attempted_fallbacks": 99,
"original_model_group": "spoofed-group",
"_client_output_ceiling": {"api_base": "https://attacker.example"},
"client_key": "client_value",
}
@pytest.mark.asyncio
@ -7301,6 +7348,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
assert "litellm_metadata" not in updated
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
assert "_client_output_ceiling" not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"

View file

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

View file

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

View file

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

View file

@ -3166,6 +3166,47 @@ async def test_custom_ui_sso_sign_in_handler_config_loading():
os.unlink(config_file_path)
@pytest.mark.asyncio
async def test_startup_initializes_string_callbacks_after_all_litellm_settings_load(tmp_path, monkeypatch):
from litellm.integrations.s3_v2 import S3Logger
from litellm.litellm_core_utils import litellm_logging
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import ProxyLogging
config_file = tmp_path / "config.yaml"
config_file.write_text(
"model_list: []\n"
"litellm_settings:\n"
" success_callback:\n"
" - s3_v2\n"
" failure_callback:\n"
" - s3_v2\n"
" s3_callback_params:\n"
" s3_bucket_name: ordering-regression-bucket\n"
" s3_region_name: us-west-2\n"
)
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm, "s3_callback_params", None)
monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [])
await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file))
ProxyLogging(user_api_key_cache=MagicMock())._init_litellm_callbacks(llm_router=None)
success_loggers = [cb for cb in litellm._async_success_callback if isinstance(cb, S3Logger)]
failure_loggers = [cb for cb in litellm._async_failure_callback if isinstance(cb, S3Logger)]
assert len(success_loggers) == 1
assert len(failure_loggers) == 1
assert success_loggers[0].s3_bucket_name == "ordering-regression-bucket"
assert success_loggers[0].s3_region_name == "us-west-2"
assert "s3_v2" not in litellm.success_callback
assert "s3_v2" not in litellm.failure_callback
@pytest.mark.asyncio
async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch):
"""
@ -8315,8 +8356,8 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter(
"""When the reservation reconcile finds the counter in an inconsistent state
(here: missing), it must NOT delete the counter and fail open (the old
behavior, which left the counter unenforced after a Redis reload). It reseeds
from the authoritative DB so the counter reflects the recorded total and
budget gating continues."""
from the authoritative DB and adds this request's settled cost, which the
async spend flush has not written yet, so budget gating continues."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import increment_spend_counters
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
@ -8353,9 +8394,7 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter(
)
assert budget_reservation["finalized"] is True
# counter reseeded to the authoritative DB value, not deleted/left None
# and not double-counted via a direct increment
assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.6)
assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.85)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma

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