mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit_7174_stream_tool_call_rewrites
This commit is contained in:
commit
49ab5fa7fb
127 changed files with 5165 additions and 354 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -152,6 +152,13 @@ spec:
|
|||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- if .Values.metricsServer.enabled }}
|
||||
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
|
||||
{{- fail "metricsServer.port must differ from service.port" }}
|
||||
{{- end }}
|
||||
- name: PROMETHEUS_METRICS_PORT
|
||||
value: {{ .Values.metricsServer.port | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.migrationJob.enabled }}
|
||||
# Schema updates are owned by the dedicated migrations Job; skip
|
||||
# the proxy's startup `prisma db push` so N replicas don't race
|
||||
|
|
@ -189,6 +196,11 @@ spec:
|
|||
- name: http
|
||||
containerPort: {{ .Values.service.port }}
|
||||
protocol: TCP
|
||||
{{- if .Values.metricsServer.enabled }}
|
||||
- name: metrics
|
||||
containerPort: {{ .Values.metricsServer.port }}
|
||||
protocol: TCP
|
||||
{{- end }}
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: {{ .Values.livenessProbe.path | quote }}
|
||||
|
|
|
|||
17
helm/litellm-helm/templates/service-metrics.yaml
Normal file
17
helm/litellm-helm/templates/service-metrics.yaml
Normal 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 }}
|
||||
|
|
@ -26,7 +26,7 @@ spec:
|
|||
{{- toYaml .namespaceSelector.matchNames | nindent 4 }}
|
||||
{{- end }}
|
||||
endpoints:
|
||||
- port: http
|
||||
- port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
|
||||
path: /metrics/
|
||||
interval: {{ .interval }}
|
||||
scrapeTimeout: {{ .scrapeTimeout }}
|
||||
|
|
|
|||
106
helm/litellm-helm/tests/metrics_server_tests.yaml
Normal file
106
helm/litellm-helm/tests/metrics_server_tests.yaml
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -441,3 +441,5 @@ ImplementationSpecific
|
|||
{{- .pathType -}}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}
|
||||
|
|
|
|||
|
|
@ -64,14 +64,25 @@ spec:
|
|||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- if .Values.gateway.metricsServer.enabled }}
|
||||
{{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
|
||||
{{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
|
||||
{{- end }}
|
||||
- name: PROMETHEUS_MULTIPROC_DIR
|
||||
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
|
||||
volumeMounts:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
{{- end }}
|
||||
{{- if .Values.gateway.metricsServer.enabled }}
|
||||
- name: prometheus-multiproc
|
||||
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
|
||||
{{- end }}
|
||||
|
|
@ -97,16 +108,54 @@ spec:
|
|||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.gateway.resources | nindent 12 }}
|
||||
{{- if .Values.gateway.metricsServer.enabled }}
|
||||
- name: metrics
|
||||
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
|
||||
{{- with .Values.gateway.securityContext }}
|
||||
securityContext:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
command:
|
||||
- python
|
||||
- -m
|
||||
- litellm.proxy.prometheus_metrics_server
|
||||
- --port
|
||||
- {{ .Values.gateway.metricsServer.port | quote }}
|
||||
env:
|
||||
- name: PROMETHEUS_MULTIPROC_DIR
|
||||
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
|
||||
ports:
|
||||
- name: metrics
|
||||
containerPort: {{ .Values.gateway.metricsServer.port }}
|
||||
protocol: TCP
|
||||
volumeMounts:
|
||||
- name: prometheus-multiproc
|
||||
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
|
||||
readinessProbe:
|
||||
tcpSocket: { port: metrics }
|
||||
periodSeconds: 10
|
||||
livenessProbe:
|
||||
tcpSocket: { port: metrics }
|
||||
periodSeconds: 15
|
||||
failureThreshold: 6
|
||||
resources:
|
||||
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- with .Values.gateway.extraContainers }}
|
||||
{{- tpl (toYaml .) $ | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
|
||||
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
|
||||
volumes:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: gateway-config
|
||||
configMap:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-config
|
||||
{{- end }}
|
||||
{{- if .Values.gateway.metricsServer.enabled }}
|
||||
- name: prometheus-multiproc
|
||||
emptyDir: {}
|
||||
{{- end }}
|
||||
{{- if .Values.billingMetrics.enabled }}
|
||||
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
|
|
|||
18
helm/litellm/templates/gateway/service-metrics.yaml
Normal file
18
helm/litellm/templates/gateway/service-metrics.yaml
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-metrics
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: gateway
|
||||
spec:
|
||||
type: ClusterIP
|
||||
ports:
|
||||
- port: {{ .Values.gateway.metricsServer.port }}
|
||||
targetPort: metrics
|
||||
protocol: TCP
|
||||
name: metrics
|
||||
selector:
|
||||
{{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
|
||||
{{- end }}
|
||||
148
helm/litellm/tests/metrics_server_tests.yaml
Normal file
148
helm/litellm/tests/metrics_server_tests.yaml
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -3585,6 +3585,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
|
|||
litellm_callback_params=[
|
||||
"OTEL_EXPORTER",
|
||||
"OTEL_ENDPOINT",
|
||||
"OTEL_TRACES_ENDPOINT",
|
||||
"OTEL_HEADERS",
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -178,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
|
||||
|
|
@ -555,6 +558,24 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") ->
|
|||
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"], ...]:
|
||||
|
|
@ -571,13 +592,12 @@ def _streamable_post_call_pipelines(
|
|||
post_call_pipelines: Final = _post_call_pipelines(request_data)
|
||||
if not post_call_pipelines:
|
||||
return ()
|
||||
route: Final = user_api_key_dict.request_route
|
||||
if route and resolve_endpoint_translation(user_api_key_dict, None) is None:
|
||||
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",
|
||||
route,
|
||||
user_api_key_dict.request_route,
|
||||
", ".join(policy_name for policy_name, _pipeline in post_call_pipelines),
|
||||
)
|
||||
return ()
|
||||
|
|
@ -990,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:
|
||||
|
|
@ -3271,15 +3299,15 @@ class ProxyLogging:
|
|||
# dict lookups + llm_router.get_deployment() per callback per chunk.
|
||||
_cached_guardrail_data: dict | None = None
|
||||
_guardrail_data_computed = False
|
||||
pipeline_managed: Final = (
|
||||
_pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset()
|
||||
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_managed:
|
||||
if callback.guardrail_name in pipeline_gated:
|
||||
continue
|
||||
# Main - V2 Guardrails implementation
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import random
|
|||
import re
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from itertools import accumulate, islice, takewhile
|
||||
from itertools import accumulate, chain, islice, takewhile
|
||||
from threading import Lock
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
|
||||
|
|
@ -31,6 +31,7 @@ from litellm._logging import verbose_router_logger
|
|||
from litellm.constants import (
|
||||
EMPTY_MAPPING,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
OUTPUT_TOKEN_CEILING_PARAMS,
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
)
|
||||
|
|
@ -82,6 +83,7 @@ from .config import (
|
|||
ClassificationRubric,
|
||||
ComplexityRouterConfig,
|
||||
ComplexityTier,
|
||||
CustomDimension,
|
||||
TierDefinition,
|
||||
)
|
||||
from .stall_detector import detect_stalled_task
|
||||
|
|
@ -879,6 +881,15 @@ class DimensionScore:
|
|||
self.signal = signal
|
||||
|
||||
|
||||
class _CustomDimensionMatchers(NamedTuple):
|
||||
"""One custom dimension's distinct matchers and the number of hits that saturates its score."""
|
||||
|
||||
dimension: CustomDimension
|
||||
keywords: tuple[str, ...]
|
||||
patterns: tuple[re.Pattern[str], ...]
|
||||
saturation: int
|
||||
|
||||
|
||||
class KeywordOverride(NamedTuple):
|
||||
"""A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired."""
|
||||
|
||||
|
|
@ -1121,7 +1132,12 @@ class ComplexityRouter(CustomLogger):
|
|||
)
|
||||
self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS
|
||||
self._custom_dimensions = tuple(
|
||||
(dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns))
|
||||
_CustomDimensionMatchers(
|
||||
dimension,
|
||||
tuple(dict.fromkeys(keyword.lower() for keyword in dimension.keywords)),
|
||||
tuple(re.compile(pattern, re.IGNORECASE) for pattern in dict.fromkeys(dimension.patterns)),
|
||||
2 if dimension.scoring_mode == "match_count" else 1,
|
||||
)
|
||||
for dimension in self.config.custom_dimensions
|
||||
)
|
||||
if self.config.has_custom_tiers:
|
||||
|
|
@ -1325,15 +1341,26 @@ class ComplexityRouter(CustomLogger):
|
|||
score: Final = score_high if match_count >= high_threshold else score_low
|
||||
return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
|
||||
|
||||
def _count_custom_hits(self, matchers: _CustomDimensionMatchers, user_text: str, scanned: str) -> int:
|
||||
hits: Final = chain(
|
||||
(self._keyword_matches(user_text, keyword) for keyword in matchers.keywords),
|
||||
(pattern.search(scanned) is not None for pattern in matchers.patterns),
|
||||
)
|
||||
return sum(islice((1 for hit in hits if hit), matchers.saturation))
|
||||
|
||||
def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]:
|
||||
if not self._custom_dimensions:
|
||||
return ()
|
||||
scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS]
|
||||
return tuple(
|
||||
(DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight)
|
||||
for dimension, patterns in self._custom_dimensions
|
||||
if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords)
|
||||
or any(pattern.search(scanned) is not None for pattern in patterns)
|
||||
(
|
||||
DimensionScore(
|
||||
matchers.dimension.name, hits / matchers.saturation, f"custom ({matchers.dimension.name})"
|
||||
),
|
||||
matchers.dimension.weight,
|
||||
)
|
||||
for matchers in self._custom_dimensions
|
||||
if (hits := self._count_custom_hits(matchers, user_text, scanned))
|
||||
)
|
||||
|
||||
def _score_multi_step(self, text: str) -> DimensionScore:
|
||||
|
|
@ -2084,11 +2111,15 @@ class ComplexityRouter(CustomLogger):
|
|||
raise ValueError(f"No model configured for tier {tier_key} and no default_model set")
|
||||
|
||||
def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]:
|
||||
if tier is None:
|
||||
return MappingProxyType({})
|
||||
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ())
|
||||
entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else ()
|
||||
entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None)
|
||||
return entry.litellm_params if entry is not None else MappingProxyType({})
|
||||
explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({})
|
||||
if not self.config.max_tokens_from_tier_model or not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(explicit):
|
||||
return explicit
|
||||
ceiling: Final = self._group_output_ceiling(model)
|
||||
if ceiling is None:
|
||||
return explicit
|
||||
return MappingProxyType({**explicit, "max_tokens": ceiling})
|
||||
|
||||
@staticmethod
|
||||
def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str:
|
||||
|
|
@ -2423,12 +2454,15 @@ class ComplexityRouter(CustomLogger):
|
|||
return name if self.config.has_custom_tiers else ComplexityTier(name)
|
||||
|
||||
def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None:
|
||||
return self._deployment_limit(group, deployment, "max_input_tokens")
|
||||
|
||||
def _deployment_limit(
|
||||
self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"]
|
||||
) -> int | None:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
|
||||
deployment_model_info: Final = deployment.get("model_info")
|
||||
declared: Final = (
|
||||
deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None
|
||||
)
|
||||
declared: Final = deployment_model_info.get(key) if isinstance(deployment_model_info, Mapping) else None
|
||||
if isinstance(declared, int):
|
||||
return declared
|
||||
litellm_params: Final = deployment.get("litellm_params")
|
||||
|
|
@ -2445,18 +2479,34 @@ class ComplexityRouter(CustomLogger):
|
|||
deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts
|
||||
received_model_name=group,
|
||||
)
|
||||
window: Final = model_info.get("max_input_tokens")
|
||||
limit: Final = model_info.get(key)
|
||||
except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others
|
||||
return None
|
||||
return window if isinstance(window, int) else None
|
||||
return limit if isinstance(limit, int) else None
|
||||
|
||||
def _group_deployments(self, group: str) -> Sequence[Mapping[str, object]]:
|
||||
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
|
||||
deployments: Final = list_models(model_name=group) if callable(list_models) else None
|
||||
return tuple(deployments) if isinstance(deployments, list) else ()
|
||||
|
||||
def _group_output_ceiling(self, group: str) -> int | None:
|
||||
"""Smallest max_output_tokens across the group's deployments, or None when any deployment
|
||||
declares none: the core router picks within the group without a fit check, and a ceiling
|
||||
above an unmapped member's real limit is a provider 400 on that member."""
|
||||
deployments: Final = self._group_deployments(group)
|
||||
ceilings: Final = tuple(
|
||||
ceiling
|
||||
for deployment in deployments
|
||||
if (ceiling := self._deployment_limit(group, deployment, "max_output_tokens")) is not None
|
||||
)
|
||||
return min(ceilings) if ceilings and len(ceilings) == len(deployments) else None
|
||||
|
||||
def _group_window_facts(self, group: str) -> tuple[int | None, bool]:
|
||||
"""(smallest declared context window across the group's deployments, whether any deployment
|
||||
declares none). The core router picks a deployment within the group without a fit check, so
|
||||
the group is only as safe as its smallest member."""
|
||||
list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None)
|
||||
deployments: Final = list_models(model_name=group) if callable(list_models) else None
|
||||
if not isinstance(deployments, list) or not deployments:
|
||||
deployments: Final = self._group_deployments(group)
|
||||
if not deployments:
|
||||
return (None, True)
|
||||
windows: Final = tuple(
|
||||
window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None
|
||||
|
|
@ -3505,6 +3555,7 @@ class ComplexityRouter(CustomLogger):
|
|||
ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs
|
||||
)
|
||||
fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM
|
||||
default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model)
|
||||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
|
|
@ -3513,7 +3564,9 @@ class ComplexityRouter(CustomLogger):
|
|||
cause="default_fallback",
|
||||
tier=fallback_tier,
|
||||
conversation_continuing=conversation_continuing,
|
||||
tier_litellm_params=default_tier_params,
|
||||
),
|
||||
litellm_params=default_tier_params,
|
||||
)
|
||||
|
||||
ask: Final = user_message or ""
|
||||
|
|
@ -3540,6 +3593,7 @@ class ComplexityRouter(CustomLogger):
|
|||
_tier_name(plan_floor),
|
||||
routed_model,
|
||||
)
|
||||
plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model)
|
||||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
|
|
@ -3551,7 +3605,9 @@ class ComplexityRouter(CustomLogger):
|
|||
matched_keyword=plan_mode_sentinel,
|
||||
escalation_keyword=escalation_keyword,
|
||||
escalated=False,
|
||||
tier_litellm_params=plan_tier_params,
|
||||
),
|
||||
litellm_params=plan_tier_params,
|
||||
)
|
||||
|
||||
override: Final = await self._resolve_keyword_tier_override(ask, request_kwargs)
|
||||
|
|
@ -3644,6 +3700,7 @@ class ComplexityRouter(CustomLogger):
|
|||
outcome.signals,
|
||||
fallback_model,
|
||||
)
|
||||
fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model)
|
||||
return PreRoutingHookResponse(
|
||||
model=fallback_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
|
|
@ -3654,7 +3711,9 @@ class ComplexityRouter(CustomLogger):
|
|||
signals=outcome.signals,
|
||||
escalation_keyword=escalation_keyword,
|
||||
escalated=False,
|
||||
tier_litellm_params=fallback_tier_params,
|
||||
),
|
||||
litellm_params=fallback_tier_params,
|
||||
)
|
||||
if self.config.adaptive:
|
||||
# hard_floor rather than a hard pick, and passed whenever the sentinel is present
|
||||
|
|
|
|||
|
|
@ -667,6 +667,14 @@ class CustomDimension(BaseModel):
|
|||
weight: float = Field(gt=0, le=1, allow_inf_nan=False)
|
||||
keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
|
||||
patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32)
|
||||
scoring_mode: Literal["binary", "match_count"] = Field(
|
||||
default="binary",
|
||||
description=(
|
||||
"'binary' scores 1 when any matcher hits. 'match_count' scores 0.5 when one distinct matcher hits and 1 "
|
||||
"when two or more do; repeated occurrences of one matcher never raise it. Keywords are distinct "
|
||||
"case-insensitively, patterns by source, and a keyword and a pattern are always distinct from each other."
|
||||
),
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_matchers(self) -> "CustomDimension":
|
||||
|
|
@ -794,8 +802,9 @@ class ComplexityRouterConfig(BaseModel):
|
|||
default=(),
|
||||
max_length=16,
|
||||
description=(
|
||||
"Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
|
||||
"when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. "
|
||||
"Named dimensions added to the heuristic-v1 score. Each contributes its inline weight once "
|
||||
"when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters; "
|
||||
"scoring_mode 'match_count' instead grades half weight for one distinct matcher and full for two or more. "
|
||||
"Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, "
|
||||
"backreferences and lookarounds are rejected. Conservative work limits include alternation paths, "
|
||||
"repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. "
|
||||
|
|
@ -1080,6 +1089,20 @@ class ComplexityRouterConfig(BaseModel):
|
|||
"wording the built-ins don't cover, or after a client release changes its strings."
|
||||
),
|
||||
)
|
||||
max_tokens_from_tier_model: bool = Field(
|
||||
default=True,
|
||||
description=(
|
||||
"Set max_tokens on every routed request to the output ceiling of the tier model it "
|
||||
"lands on, replacing whatever the caller sent. A caller behind an auto-router cannot "
|
||||
"pick one value that fits every tier: the smallest tier's ceiling starves a bigger "
|
||||
"tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The "
|
||||
"ceiling is the smallest max_output_tokens across the tier model's deployments, read "
|
||||
"from each deployment's model_info and then the model cost map; a tier model with a "
|
||||
"deployment whose ceiling is unknown keeps the caller's value. A max_tokens, "
|
||||
"max_completion_tokens or max_output_tokens in the tier's own litellm_params still "
|
||||
"wins. Set false to forward the caller's value unchanged."
|
||||
),
|
||||
)
|
||||
route_housekeeping_to_cheapest_tier: bool = Field(
|
||||
default=True,
|
||||
description=(
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -580,6 +580,7 @@ CallTypesLiteral = Literal[
|
|||
"search",
|
||||
"asearch",
|
||||
"_arealtime",
|
||||
"_aresponses_websocket",
|
||||
"create_batch",
|
||||
"acreate_batch",
|
||||
"create_file",
|
||||
|
|
|
|||
122
litellm/utils.py
122
litellm/utils.py
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
}
|
||||
|
|
|
|||
|
|
@ -156,6 +156,17 @@ resource "aws_security_group" "tasks" {
|
|||
security_groups = [aws_security_group.alb.id]
|
||||
}
|
||||
|
||||
dynamic "ingress" {
|
||||
for_each = local.metrics_enabled && length(var.gateway_metrics_scrape_cidrs) > 0 ? [1] : []
|
||||
content {
|
||||
description = "Prometheus scrapers to the gateway metrics sidecar"
|
||||
from_port = var.gateway_metrics_port
|
||||
to_port = var.gateway_metrics_port
|
||||
protocol = "tcp"
|
||||
cidr_blocks = var.gateway_metrics_scrape_cidrs
|
||||
}
|
||||
}
|
||||
|
||||
egress {
|
||||
description = "All egress (LLM providers, RDS, Redis)"
|
||||
from_port = 0
|
||||
|
|
|
|||
108
terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl
Normal file
108
terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl
Normal 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]
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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.")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)."""
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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``.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -7314,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
|
||||
|
|
@ -7343,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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -2131,7 +2131,11 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail(
|
|||
seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1
|
||||
return None
|
||||
|
||||
managed = RecordingGuardrail(
|
||||
class UnifiedRecordingGuardrail(RecordingGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
return inputs
|
||||
|
||||
managed = UnifiedRecordingGuardrail(
|
||||
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True
|
||||
)
|
||||
free = RecordingGuardrail(
|
||||
|
|
@ -2150,3 +2154,34 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail(
|
|||
assert result is not None
|
||||
assert seen.get("gr-post") is None
|
||||
assert seen["gr-free"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_stream(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
|
||||
class ChunkHookGuardrail(CustomGuardrail):
|
||||
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
|
||||
seen["count"] = seen.get("count", 0) + 1
|
||||
seen["response"] = response
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[ChunkHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
|
||||
result = await proxy_logging.async_post_call_streaming_hook(
|
||||
data=data,
|
||||
response=_stream_chunks()[0],
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert seen["count"] == 1
|
||||
assert seen["response"] == "hello "
|
||||
|
|
|
|||
|
|
@ -246,8 +246,10 @@ async def test_aresponses_keeps_include_obfuscation_in_stream_options():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("drop_params", [True, "true"])
|
||||
async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service_tier(
|
||||
monkeypatch,
|
||||
drop_params,
|
||||
):
|
||||
"""
|
||||
Request-level drop_params=True (as the proxy injects for agentic CLIs) must
|
||||
|
|
@ -271,7 +273,7 @@ async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service
|
|||
aws_region_name="us-east-1",
|
||||
input="hi",
|
||||
service_tier="priority",
|
||||
drop_params=True,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from pydantic import ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_utils.auto_router_model_naming import (
|
||||
CUSTOMIZATION_CAPABILITY,
|
||||
GATED_AUTO_ROUTER_CAPABILITIES,
|
||||
|
|
@ -24,7 +25,12 @@ from litellm.router_utils.auto_router_model_naming import (
|
|||
)
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
|
||||
from litellm.constants import (
|
||||
OUTPUT_TOKEN_CEILING_PARAMS,
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
)
|
||||
from litellm.router import as_output_cap
|
||||
from litellm.router_strategy.complexity_router.complexity_router import (
|
||||
_CLASSIFICATION_CURRENT_MESSAGE_ONLY,
|
||||
_CLASSIFICATION_WITH_CONVERSATION,
|
||||
|
|
@ -826,6 +832,8 @@ class TestCustomDimensions:
|
|||
pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"),
|
||||
pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"),
|
||||
pytest.param({"unknown": True}, {}, id="extra-field"),
|
||||
pytest.param({"scoring_mode": "graded"}, {}, id="unknown-scoring-mode"),
|
||||
pytest.param({"scoring_mode": None}, {}, id="null-scoring-mode"),
|
||||
],
|
||||
)
|
||||
def test_custom_dimension_invalid_configuration_rejected(
|
||||
|
|
@ -878,43 +886,125 @@ class TestCustomDimensions:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh"))
|
||||
@pytest.mark.parametrize("scoring_mode", ("binary", "match_count"))
|
||||
@pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh", "orbitmesh fluxgate"))
|
||||
async def test_custom_dimensions_public_hook_scores_only_current_ask(
|
||||
self, mock_router_instance: MagicMock, current_ask: str
|
||||
self, mock_router_instance: MagicMock, current_ask: str, scoring_mode: str
|
||||
) -> None:
|
||||
router: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{
|
||||
"tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"},
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}],
|
||||
"dimension_weights": {},
|
||||
"custom_dimensions": [
|
||||
{
|
||||
"name": "internalFrameworks",
|
||||
"weight": 0.8,
|
||||
"keywords": ["orbitmesh", "fluxgate"],
|
||||
"scoring_mode": scoring_mode,
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
result: Final = await router.async_pre_routing_hook(
|
||||
model="test-router",
|
||||
request_kwargs={},
|
||||
messages=[
|
||||
{"role": "system", "content": "orbitmesh"},
|
||||
{"role": "user", "content": "orbitmesh"},
|
||||
{"role": "assistant", "content": "orbitmesh is ready"},
|
||||
{"role": "system", "content": "orbitmesh fluxgate"},
|
||||
{"role": "user", "content": "orbitmesh fluxgate"},
|
||||
{"role": "assistant", "content": "orbitmesh fluxgate is ready"},
|
||||
{"role": "user", "content": current_ask},
|
||||
{"role": "tool", "tool_call_id": "previous", "content": "orbitmesh"},
|
||||
{"role": "tool", "tool_call_id": "previous", "content": "orbitmesh fluxgate"},
|
||||
],
|
||||
)
|
||||
assert result is not None
|
||||
assert result.routing_decision is not None
|
||||
assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (current_ask == "orbitmesh")
|
||||
assert result.model == ("top" if current_ask == "orbitmesh" else "cheap")
|
||||
expected_score: Final = (
|
||||
0.0
|
||||
if current_ask == "Hello!"
|
||||
else 0.4
|
||||
if scoring_mode == "match_count" and current_ask == "orbitmesh"
|
||||
else 0.8
|
||||
)
|
||||
assert result.routing_decision["score"] == expected_score
|
||||
assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (expected_score > 0)
|
||||
assert result.model == ("cheap" if expected_score == 0 else "strong" if expected_score == 0.4 else "top")
|
||||
assert "orbitmesh" not in " ".join(result.routing_decision["signals"])
|
||||
|
||||
def test_custom_patterns_scan_only_the_first_2048_characters(self, mock_router_instance: MagicMock) -> None:
|
||||
@pytest.mark.parametrize("scoring_mode", ("binary", "match_count"))
|
||||
def test_custom_patterns_scan_only_the_first_2048_characters(
|
||||
self, mock_router_instance: MagicMock, scoring_mode: str
|
||||
) -> None:
|
||||
router: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{"custom_dimensions": [{"name": "late", "weight": 0.7, "patterns": [r"zzz{1,3}"]}]},
|
||||
{
|
||||
"custom_dimensions": [
|
||||
{
|
||||
"name": "late",
|
||||
"weight": 0.7,
|
||||
"patterns": [r"zzz{1,3}", r"yyy{1,3}"],
|
||||
"scoring_mode": scoring_mode,
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
baseline: Final = ComplexityRouter("test-router", mock_router_instance)
|
||||
assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2]
|
||||
assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2]
|
||||
second_hit_past_the_bound: Final = "yyy " + "a" * 2044 + " zzz"
|
||||
contribution: Final = (
|
||||
router.classify(second_hit_past_the_bound)[1] - baseline.classify(second_hit_past_the_bound)[1]
|
||||
)
|
||||
assert contribution == pytest.approx(0.7 if scoring_mode == "binary" else 0.35)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prompt,expected_score",
|
||||
[
|
||||
pytest.param("Hello!", 0.0, id="no-hit"),
|
||||
pytest.param("orbitmesh orbitmesh ORBITMESH again", 0.5, id="one-keyword-repeated"),
|
||||
pytest.param("create table a; CREATE TABLE b; create table c", 0.5, id="one-pattern-repeated"),
|
||||
pytest.param("orbitmesh and fluxgate", 1.0, id="two-keywords"),
|
||||
pytest.param("orbitmesh then create table t", 1.0, id="keyword-plus-pattern"),
|
||||
pytest.param("create table a; alter table b", 1.0, id="two-patterns"),
|
||||
pytest.param("orbitmesh fluxgate create table a alter table b", 1.0, id="all-matchers"),
|
||||
],
|
||||
)
|
||||
def test_match_count_grades_distinct_matchers(
|
||||
self, mock_router_instance: MagicMock, prompt: str, expected_score: float
|
||||
) -> None:
|
||||
dimension: Final = {
|
||||
"name": "graded",
|
||||
"weight": 0.6,
|
||||
"keywords": ["orbitmesh", "ORBITMESH", "fluxgate"],
|
||||
"patterns": [r"\bcreate\s{1,4}table\b", r"\bcreate\s{1,4}table\b", r"\balter\s{1,4}table\b"],
|
||||
}
|
||||
baseline: Final = ComplexityRouter("test-router", mock_router_instance)
|
||||
binary: Final = ComplexityRouter("test-router", mock_router_instance, {"custom_dimensions": [dimension]})
|
||||
graded: Final = ComplexityRouter(
|
||||
"test-router",
|
||||
mock_router_instance,
|
||||
{"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]},
|
||||
)
|
||||
_, baseline_score, baseline_signals = baseline.classify(prompt)
|
||||
_, binary_score, binary_signals = binary.classify(prompt)
|
||||
_, graded_score, graded_signals = graded.classify(prompt)
|
||||
assert graded_score == pytest.approx(baseline_score + 0.6 * expected_score)
|
||||
assert binary_score == pytest.approx(baseline_score + (0.6 if expected_score else 0.0))
|
||||
expected_signals: Final = [*baseline_signals, *(["custom (graded)"] if expected_score else [])]
|
||||
assert graded_signals == expected_signals
|
||||
assert binary_signals == expected_signals
|
||||
|
||||
def test_scoring_mode_round_trips_and_defaults_to_binary(self) -> None:
|
||||
dimension: Final = {"name": "graded", "weight": 0.6, "keywords": ["orbitmesh"]}
|
||||
legacy: Final = ComplexityRouterConfig.model_validate({"custom_dimensions": [dimension]})
|
||||
graded: Final = ComplexityRouterConfig.model_validate(
|
||||
{"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]}
|
||||
)
|
||||
assert legacy.custom_dimensions[0].scoring_mode == "binary"
|
||||
assert graded.model_dump(mode="json")["custom_dimensions"][0]["scoring_mode"] == "match_count"
|
||||
assert ComplexityRouterConfig.model_validate(graded.model_dump(mode="json")) == graded
|
||||
|
||||
def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None:
|
||||
heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]}
|
||||
|
|
@ -1511,6 +1601,7 @@ class TestRouterComplexityDeploymentMethods:
|
|||
def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None:
|
||||
"""Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no
|
||||
prompt at all, leaves a router unmetered, so several of them register under a ceiling of one."""
|
||||
|
||||
def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]:
|
||||
llm_config: dict[str, object] = {"model": "gpt-4o-mini"}
|
||||
if preset is not None:
|
||||
|
|
@ -1647,6 +1738,7 @@ class TestRouterComplexityDeploymentMethods:
|
|||
def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None:
|
||||
"""tier_labels renames the built-in ladder without defining one, so it stays ungated: two such
|
||||
routers register under a ceiling of one."""
|
||||
|
||||
def labeled(model_name: str, model_id: str) -> dict[str, object]:
|
||||
row = self._router_row(model_name, model_id, "heuristic")
|
||||
row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"}
|
||||
|
|
@ -2520,9 +2612,7 @@ class TestLLMClassifier:
|
|||
assert outcome.classifier_cost == pytest.approx(1.35e-05)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(
|
||||
self, llm_classifier_config
|
||||
):
|
||||
async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(self, llm_classifier_config):
|
||||
real_router = Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -2565,9 +2655,7 @@ class TestLLMClassifier:
|
|||
assert real_router.total_calls["openai/mock-backup-classifier"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclassify_enforces_total_classifier_deadline(
|
||||
self, mock_router_instance, llm_classifier_config
|
||||
):
|
||||
async def test_aclassify_enforces_total_classifier_deadline(self, mock_router_instance, llm_classifier_config):
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def slow_classifier(**_kwargs: object) -> None:
|
||||
|
|
@ -3333,11 +3421,11 @@ class TestRouterPreRoutingAliasOverrides:
|
|||
|
||||
def test_drop_client_effort_carriers_helper_edge_shapes(self):
|
||||
no_pin: Dict = {"thinking": {"type": "adaptive"}}
|
||||
Router._drop_client_effort_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1})
|
||||
Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1})
|
||||
assert no_pin == {"thinking": {"type": "adaptive"}}
|
||||
|
||||
non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3}
|
||||
Router._drop_client_effort_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"})
|
||||
Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"})
|
||||
assert non_dict_carriers == {"output_config": "max", "reasoning": 3}
|
||||
|
||||
effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}}
|
||||
|
|
@ -3791,11 +3879,11 @@ class TestRouterPreRoutingSharedAliasName:
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
async def _routed_call_kwargs(router: Router, **request_params) -> dict:
|
||||
async def _routed_call_kwargs(router: Router, prompt: str = "hi", **request_params) -> dict:
|
||||
mock_acompletion = AsyncMock(return_value=litellm.ModelResponse(choices=[{"message": {"content": "hi"}}]))
|
||||
with patch.object(litellm, "acompletion", mock_acompletion):
|
||||
await router.acompletion(
|
||||
model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_params
|
||||
model="smart-router", messages=[{"role": "user", "content": prompt}], **request_params
|
||||
)
|
||||
return mock_acompletion.call_args.kwargs
|
||||
|
||||
|
|
@ -12414,9 +12502,7 @@ class TestTierHealthFailover:
|
|||
llm_provider="",
|
||||
)
|
||||
filtered = (*cooling, *blocked, *excluded)
|
||||
healthy = [
|
||||
{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered
|
||||
]
|
||||
healthy = [{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered]
|
||||
if not healthy:
|
||||
raise RouterRateLimitError(
|
||||
model=model, cooldown_time=60.0, enable_pre_call_checks=False, cooldown_list=[]
|
||||
|
|
@ -12845,9 +12931,7 @@ class TestTierHealthFailover:
|
|||
assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(
|
||||
self, mock_router_instance
|
||||
):
|
||||
async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(self, mock_router_instance):
|
||||
"""RPM exhaustion is its own verdict from the owner (RouterRateLimitErrorBasic). A peer
|
||||
in that state would be rejected downstream, so it cannot be the substitute."""
|
||||
from litellm.types.router import RouterRateLimitErrorBasic
|
||||
|
|
@ -12880,9 +12964,7 @@ class TestTierHealthFailover:
|
|||
assert {r.model for r in results} == {"live-c"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(
|
||||
self, mock_router_instance
|
||||
):
|
||||
async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(self, mock_router_instance):
|
||||
"""The Responses API carries its prompt as `input`, never as messages. The owner only
|
||||
runs its context-window pre-call check when one of them is present, so dropping `input`
|
||||
would silently skip window filtering on that whole surface."""
|
||||
|
|
@ -12908,9 +12990,7 @@ class TestTierHealthFailover:
|
|||
), "the eligibility probe must forward `input` to the owner"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(
|
||||
self, mock_router_instance
|
||||
):
|
||||
async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance):
|
||||
"""The owner answers an unconfigured group with BadRequestError. Reading that as live
|
||||
would both skip failover off it and let it be chosen as a substitute."""
|
||||
router = self._router(
|
||||
|
|
@ -13099,9 +13179,7 @@ class TestClassifierVision:
|
|||
routed as default_fallback on text the request never contained.
|
||||
"""
|
||||
router = self._router(mock_router_instance, vision={"enabled": True})
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="m", request_kwargs={}, messages=self._turn(IMG_PART)
|
||||
)
|
||||
response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART))
|
||||
assert response.routing_decision["cause"] == "llm_classifier"
|
||||
assert response.model == "t-complex"
|
||||
assert [block["type"] for block in self._classifier_user_content(mock_router_instance)] == [
|
||||
|
|
@ -13112,9 +13190,7 @@ class TestClassifierVision:
|
|||
@pytest.mark.asyncio
|
||||
async def test_image_only_turn_still_falls_back_when_vision_is_off(self, mock_router_instance):
|
||||
router = self._router(mock_router_instance, vision={"enabled": False})
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="m", request_kwargs={}, messages=self._turn(IMG_PART)
|
||||
)
|
||||
response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART))
|
||||
assert response.routing_decision["cause"] == "default_fallback"
|
||||
mock_router_instance.acompletion.assert_not_awaited()
|
||||
|
||||
|
|
@ -13184,9 +13260,7 @@ class TestClassifierVision:
|
|||
makes the image the only variable; a margin loose enough to leave the score undecided
|
||||
would pass whether or not the guard exists.
|
||||
"""
|
||||
router = self._router(
|
||||
mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra
|
||||
)
|
||||
router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra)
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="m", request_kwargs={}, messages=self._turn({"type": "text", "text": "what is this"}, IMG_PART)
|
||||
)
|
||||
|
|
@ -13200,9 +13274,7 @@ class TestClassifierVision:
|
|||
self, mock_router_instance, classifier_type, extra, short_circuit_cause
|
||||
):
|
||||
"""The negative class: same router, same text, no image, and the scorer still decides."""
|
||||
router = self._router(
|
||||
mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra
|
||||
)
|
||||
router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra)
|
||||
response = await router.async_pre_routing_hook(
|
||||
model="m", request_kwargs={}, messages=[{"role": "user", "content": "what is this"}]
|
||||
)
|
||||
|
|
@ -13212,3 +13284,379 @@ class TestClassifierVision:
|
|||
def test_max_images_must_be_positive(self):
|
||||
with pytest.raises(ValidationError):
|
||||
ClassifierLLMConfig(model="clf", vision={"enabled": True, "max_images": 0})
|
||||
|
||||
|
||||
class TestMaxTokensFromTierModel:
|
||||
"""The auto-router replaces the caller's output ceiling with the tier model's own, so one
|
||||
client-side value no longer starves a bigger tier or gets rejected by a smaller one."""
|
||||
|
||||
COMPLEX_PROMPT: Final = (
|
||||
"Design a distributed rate limiter with Redis, sharding and failover. Analyze the consistency "
|
||||
"tradeoffs and implement the algorithm step by step with tests."
|
||||
)
|
||||
SMALL: Final = {
|
||||
"model_name": "small",
|
||||
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"},
|
||||
"model_info": {"max_output_tokens": 8192},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _router(
|
||||
tier_litellm_params: dict | None = None,
|
||||
max_tokens_from_tier_model: bool | None = None,
|
||||
simple_deployments: list[dict] | None = None,
|
||||
extra_config: dict | None = None,
|
||||
) -> Router:
|
||||
simple_tier: dict = {"model_name": "small"}
|
||||
if tier_litellm_params:
|
||||
simple_tier["litellm_params"] = tier_litellm_params
|
||||
config: dict = {
|
||||
"tiers": {"SIMPLE": simple_tier, "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"},
|
||||
**(extra_config or {}),
|
||||
}
|
||||
if max_tokens_from_tier_model is not None:
|
||||
config["max_tokens_from_tier_model"] = max_tokens_from_tier_model
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config},
|
||||
},
|
||||
*(simple_deployments or [TestMaxTokensFromTierModel.SMALL]),
|
||||
{
|
||||
"model_name": "big",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "k"},
|
||||
"model_info": {"max_output_tokens": 64000},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _routed(router: Router, prompt: str = "hi", **request_kwargs) -> dict:
|
||||
"""Drive the real routing entry point and return the request kwargs it leaves behind."""
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="smart-router", request_kwargs=request_kwargs, messages=[{"role": "user", "content": prompt}]
|
||||
)
|
||||
return {"model": deployment["litellm_params"]["model"], **request_kwargs}
|
||||
|
||||
@staticmethod
|
||||
async def _routed_responses(router: Router, prompt: str = "hi", **request_kwargs) -> dict:
|
||||
"""The Responses surface hands the router `input` both as the prompt argument and inside the
|
||||
request kwargs, so the hook sees the same shape the real call carries."""
|
||||
routed: dict = {"input": prompt, **request_kwargs}
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="smart-router", request_kwargs=routed, input=prompt
|
||||
)
|
||||
return {"model": deployment["litellm_params"]["model"], **routed}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_ceiling_is_replaced_by_the_routed_tier_models_ceiling(self):
|
||||
router = self._router()
|
||||
|
||||
simple = await self._routed(router, max_tokens=8192)
|
||||
complex_ = await self._routed(router, self.COMPLEX_PROMPT, max_tokens=8192)
|
||||
|
||||
assert (simple["model"], simple["max_tokens"]) == ("anthropic/claude-haiku-4-5", 8192)
|
||||
assert (complex_["model"], complex_["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
|
||||
assert "max_output_tokens" not in complex_
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_every_client_carrier_of_the_ceiling_is_replaced(self):
|
||||
sent = await self._routed(self._router(), self.COMPLEX_PROMPT, max_completion_tokens=8192)
|
||||
|
||||
assert sent["max_tokens"] == 64000
|
||||
assert "max_completion_tokens" not in sent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_surface_gets_the_ceiling_under_its_own_name(self):
|
||||
sent = await self._routed_responses(self._router(), self.COMPLEX_PROMPT, max_output_tokens=8192)
|
||||
|
||||
assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
|
||||
assert "max_tokens" not in sent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"tier_params, responses_call",
|
||||
[
|
||||
({"max_tokens": 4321}, False),
|
||||
({"max_tokens": 4321}, True),
|
||||
({"max_completion_tokens": 4321}, False),
|
||||
({"max_completion_tokens": 4321}, True),
|
||||
({"max_output_tokens": 4321}, False),
|
||||
],
|
||||
)
|
||||
async def test_operators_own_tier_ceiling_wins_under_the_surface_name(self, tier_params, responses_call):
|
||||
router = self._router(tier_litellm_params=tier_params)
|
||||
if responses_call:
|
||||
sent = await self._routed_responses(router, max_output_tokens=8192)
|
||||
else:
|
||||
sent = await self._routed(router, max_tokens=8192)
|
||||
|
||||
surface_key = "max_output_tokens" if responses_call else "max_tokens"
|
||||
assert sent[surface_key] == 4321
|
||||
assert not (OUTPUT_TOKEN_CEILING_PARAMS - {surface_key}) & sent.keys()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_opting_out_forwards_the_client_value_unchanged(self):
|
||||
sent = await self._routed(self._router(max_tokens_from_tier_model=False), self.COMPLEX_PROMPT, max_tokens=8192)
|
||||
|
||||
assert sent["max_tokens"] == 8192
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self):
|
||||
unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}}
|
||||
|
||||
sent = await self._routed(self._router(simple_deployments=[self.SMALL, unmapped]), max_tokens=4000)
|
||||
|
||||
assert sent["max_tokens"] == 4000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_multi_deployment_tier_model_uses_its_smallest_ceiling(self):
|
||||
smaller: dict = {
|
||||
**self.SMALL,
|
||||
"litellm_params": {**self.SMALL["litellm_params"], "api_key": "k2"},
|
||||
"model_info": {"max_output_tokens": 4096},
|
||||
}
|
||||
|
||||
sent = await self._routed(self._router(simple_deployments=[self.SMALL, smaller]), max_tokens=100000)
|
||||
|
||||
assert sent["max_tokens"] == 4096
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ceiling_falls_back_to_the_cost_map(self, monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"auto-cap-probe-model",
|
||||
{"litellm_provider": "openai", "mode": "chat", "max_output_tokens": 4242, "max_input_tokens": 100000},
|
||||
)
|
||||
mapped_only: dict = {
|
||||
"model_name": "small",
|
||||
"litellm_params": {"model": "openai/auto-cap-probe-model", "api_key": "k"},
|
||||
}
|
||||
|
||||
sent = await self._routed(self._router(simple_deployments=[mapped_only]), max_tokens=8192)
|
||||
|
||||
assert sent["max_tokens"] == 4242
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("client_kwargs", [{}, {"max_tokens": 0}], ids=["omitted", "zero"])
|
||||
async def test_omitted_and_zero_are_replaced_like_any_other_value(self, client_kwargs):
|
||||
sent = await self._routed(self._router(), self.COMPLEX_PROMPT, **client_kwargs)
|
||||
|
||||
assert sent["max_tokens"] == 64000
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tier_params, responses_call, expected",
|
||||
[
|
||||
({"max_tokens": 1, "temperature": 0.2}, False, {"max_tokens": 1, "temperature": 0.2}),
|
||||
({"max_tokens": 1}, True, {"max_output_tokens": 1}),
|
||||
({"max_completion_tokens": 2}, False, {"max_tokens": 2}),
|
||||
({"max_completion_tokens": 2}, True, {"max_output_tokens": 2}),
|
||||
({"max_output_tokens": 3}, False, {"max_tokens": 3}),
|
||||
({"max_output_tokens": 3}, True, {"max_output_tokens": 3}),
|
||||
({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 1}),
|
||||
({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, True, {"max_output_tokens": 3}),
|
||||
({"max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 2}),
|
||||
({"reasoning_effort": "low"}, True, {"reasoning_effort": "low"}),
|
||||
],
|
||||
)
|
||||
def test_every_tier_alias_collapses_onto_the_surface_key(self, tier_params, responses_call, expected):
|
||||
assert dict(Router._tier_ceiling_under_the_surface_name(tier_params, responses_call=responses_call)) == expected
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_default_fallback_exit_carries_the_ceiling(self):
|
||||
routed: dict = {"max_tokens": 8192}
|
||||
deployment = await self._router().async_get_available_deployment(
|
||||
model="smart-router", request_kwargs=routed, messages=[{"role": "system", "content": "be nice"}]
|
||||
)
|
||||
|
||||
assert routed["metadata"]["routing_decision"]["cause"] == "default_fallback"
|
||||
assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_plan_mode_exit_carries_the_ceiling(self):
|
||||
routed: dict = {"max_tokens": 8192}
|
||||
deployment = await self._router(
|
||||
extra_config={"plan_mode_min_tier": "REASONING"}
|
||||
).async_get_available_deployment(
|
||||
model="smart-router",
|
||||
request_kwargs=routed,
|
||||
messages=[
|
||||
{"role": "user", "content": "plan the refactor"},
|
||||
{"role": "system", "content": "Plan mode is active"},
|
||||
],
|
||||
)
|
||||
|
||||
assert routed["metadata"]["routing_decision"]["cause"] == "plan_mode"
|
||||
assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_default_model_landing_with_no_tier_still_gets_its_ceiling(self):
|
||||
strategy = ComplexityRouter(
|
||||
model_name="smart-router",
|
||||
litellm_router_instance=self._router(),
|
||||
complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"},
|
||||
)
|
||||
|
||||
assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self):
|
||||
"""A model-group fallback re-enters routing with the same kwargs; a Sonnet-sized ceiling
|
||||
must not ride onto the plain group the caller configured as the fallback."""
|
||||
big: dict = {
|
||||
"model_name": "big",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-5",
|
||||
"api_key": "k",
|
||||
"mock_response": "litellm.InternalServerError",
|
||||
},
|
||||
"model_info": {"max_output_tokens": 64000},
|
||||
}
|
||||
plain: dict = {
|
||||
"model_name": "plain",
|
||||
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k", "mock_response": "ok"},
|
||||
}
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {"SIMPLE": "big", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"}
|
||||
},
|
||||
},
|
||||
},
|
||||
big,
|
||||
plain,
|
||||
],
|
||||
fallbacks=[{"smart-router": ["plain"]}],
|
||||
num_retries=0,
|
||||
)
|
||||
recorder = _OutputCeilingRecorder()
|
||||
litellm.callbacks.append(recorder)
|
||||
try:
|
||||
await router.acompletion(
|
||||
model="smart-router", messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], max_tokens=8192
|
||||
)
|
||||
finally:
|
||||
litellm.callbacks.remove(recorder)
|
||||
|
||||
assert recorder.seen == [("claude-sonnet-5", 64000), ("claude-haiku-4-5", 8192)]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_caller_seeded_stamp_cannot_inject_kwargs_on_a_plain_group(self):
|
||||
"""The stamp sits in a metadata bucket a caller can write; a planted one must yield
|
||||
nothing but integer ceiling carriers, never a redirected api_base or credential."""
|
||||
planted: dict = {
|
||||
"api_base": "https://attacker.example",
|
||||
"api_key": "stolen",
|
||||
"max_tokens": "not-an-int",
|
||||
"max_completion_tokens": True,
|
||||
"max_output_tokens": 321,
|
||||
}
|
||||
routed: dict = {"max_tokens": 8192, "metadata": {"_client_output_ceiling": planted}}
|
||||
|
||||
await self._router().async_get_available_deployment(
|
||||
model="big", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert {k: v for k, v in routed.items() if k not in ("metadata", "model_info")} == {"max_output_tokens": 321}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_pass_through_routing_entry_point_pins_and_restores_the_same_way(self):
|
||||
pass_through: dict = {**self.SMALL["litellm_params"], "use_in_pass_through": True}
|
||||
small: dict = {**self.SMALL, "litellm_params": pass_through}
|
||||
plain: dict = {**small, "model_name": "plain"}
|
||||
router = self._router(simple_deployments=[small, plain])
|
||||
for deployment in router.model_list:
|
||||
deployment["litellm_params"]["use_in_pass_through"] = True
|
||||
routed: dict = {"max_tokens": 8192}
|
||||
|
||||
deployment = await router.async_get_available_deployment_for_pass_through(
|
||||
model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": self.COMPLEX_PROMPT}]
|
||||
)
|
||||
pinned = routed["max_tokens"]
|
||||
await router.async_get_available_deployment_for_pass_through(
|
||||
model="plain", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert (deployment["litellm_params"]["model"], pinned, routed["max_tokens"]) == (
|
||||
"anthropic/claude-sonnet-5",
|
||||
64000,
|
||||
8192,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_classifier_fallback_exit_carries_the_ceiling(self):
|
||||
router = self._router(
|
||||
extra_config={
|
||||
"classifier_type": "llm",
|
||||
"classifier_llm_config": {"model": "no-such-classifier", "timeout_ms": 400},
|
||||
"classifier_fallback": "default_model",
|
||||
"default_model": "big",
|
||||
}
|
||||
)
|
||||
routed: dict = {"max_tokens": 8192}
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
|
||||
assert routed["metadata"]["routing_decision"]["cause"] == "default_model_fallback"
|
||||
assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, expected",
|
||||
[(8192, 8192), ("8192", 8192), (100.9, 100), (0, 0), (-1, None), (True, None), ("x", None), (None, None)],
|
||||
)
|
||||
def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected):
|
||||
assert as_output_cap(value) == expected
|
||||
|
||||
def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self):
|
||||
stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}}
|
||||
Router._restore_client_ceiling_no_tier_pins(stamped)
|
||||
assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192}
|
||||
|
||||
coerced: dict = {
|
||||
"max_tokens": 64000,
|
||||
"metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}},
|
||||
}
|
||||
Router._restore_client_ceiling_no_tier_pins(coerced)
|
||||
assert {k: v for k, v in coerced.items() if k != "metadata"} == {
|
||||
"max_tokens": 8192,
|
||||
"max_completion_tokens": 100,
|
||||
}
|
||||
|
||||
unstamped: dict = {"max_tokens": 64000, "metadata": {}}
|
||||
Router._restore_client_ceiling_no_tier_pins(unstamped)
|
||||
assert unstamped["max_tokens"] == 64000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pinning_stamps_the_callers_carriers_once(self):
|
||||
router = self._router()
|
||||
request_kwargs: dict = {"max_completion_tokens": 8192}
|
||||
|
||||
first = router._pin_tier_params_onto_request(
|
||||
model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False
|
||||
)
|
||||
second = router._pin_tier_params_onto_request(
|
||||
model="big", tier_litellm_params={"max_tokens": 32000}, request_kwargs=request_kwargs, responses_call=False
|
||||
)
|
||||
none = router._pin_tier_params_onto_request(
|
||||
model="big", tier_litellm_params=None, request_kwargs=request_kwargs, responses_call=False
|
||||
)
|
||||
|
||||
assert (first, second, none) == (True, True, False)
|
||||
assert request_kwargs["max_tokens"] == 32000
|
||||
assert request_kwargs["metadata"]["_client_output_ceiling"] == {"max_completion_tokens": 8192}
|
||||
|
||||
|
||||
class _OutputCeilingRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.seen: list[tuple[str, int | None]] = []
|
||||
|
||||
def log_pre_api_call(self, model, messages, kwargs):
|
||||
self.seen.append((model, kwargs.get("optional_params", {}).get("max_tokens")))
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import copy
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
|
@ -7,6 +8,8 @@ from litellm.caching.caching import DualCache
|
|||
from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
|
||||
|
||||
DEPLOYMENT_ID = "9876"
|
||||
COST_KEY = "cost_map:gpt-5.5-pool"
|
||||
LATENCY_KEYS = ("gpt-5.5-pool_map", "gpt-5.5-pool_cost_map")
|
||||
KWARGS = {
|
||||
"litellm_params": {
|
||||
"metadata": {"model_group": "gpt-5.5-pool"},
|
||||
|
|
@ -24,7 +27,7 @@ def _chat_response_with_no_completion_tokens() -> litellm.ModelResponse:
|
|||
|
||||
|
||||
def _recorded_minute_counters(cache: DualCache) -> dict[str, int]:
|
||||
cached = cache.get_cache(key="gpt-5.5-pool_map") or {}
|
||||
cached = cache.get_cache(key=COST_KEY) or {}
|
||||
minute_buckets = cached.get(DEPLOYMENT_ID, {})
|
||||
assert len(minute_buckets) == 1, f"expected one minute bucket, got {minute_buckets}"
|
||||
return next(iter(minute_buckets.values()))
|
||||
|
|
@ -44,6 +47,47 @@ def test_log_success_event_counts_a_response_with_no_completion_tokens():
|
|||
assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"])
|
||||
async def test_log_success_event_keeps_cost_bookkeeping_out_of_the_latency_routing_entry(use_async: bool):
|
||||
cache = DualCache()
|
||||
latency_entry = {DEPLOYMENT_ID: {"latency": [0.5], "time_to_first_token": [0.1]}}
|
||||
for latency_key in LATENCY_KEYS:
|
||||
cache.set_cache(key=latency_key, value=copy.deepcopy(latency_entry))
|
||||
handler = LowestCostLoggingHandler(router_cache=cache)
|
||||
call_args = {
|
||||
"kwargs": KWARGS,
|
||||
"response_obj": _chat_response_with_no_completion_tokens(),
|
||||
"start_time": datetime(2026, 1, 1, 12, 0, 0),
|
||||
"end_time": datetime(2026, 1, 1, 12, 0, 2),
|
||||
}
|
||||
|
||||
if use_async:
|
||||
await handler.async_log_success_event(**call_args)
|
||||
else:
|
||||
handler.log_success_event(**call_args)
|
||||
|
||||
assert [cache.get_cache(key=latency_key) for latency_key in LATENCY_KEYS] == [latency_entry, latency_entry]
|
||||
assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_available_deployments_applies_rpm_limit_from_the_cost_entry():
|
||||
cache = DualCache()
|
||||
handler = LowestCostLoggingHandler(router_cache=cache)
|
||||
precise_minute = datetime.now().strftime("%Y-%m-%d-%H-%M")
|
||||
cache.set_cache(key=COST_KEY, value={DEPLOYMENT_ID: {precise_minute: {"tpm": 12, "rpm": 1}}})
|
||||
healthy_deployments = [{"model_info": {"id": DEPLOYMENT_ID}, "litellm_params": {"model": "gpt-5.5", "rpm": 1}}]
|
||||
|
||||
picked = await handler.async_get_available_deployments(
|
||||
model_group="gpt-5.5-pool",
|
||||
healthy_deployments=healthy_deployments,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert picked is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_log_success_event_counts_a_response_with_no_completion_tokens():
|
||||
cache = DualCache()
|
||||
|
|
|
|||
|
|
@ -792,15 +792,168 @@ def test_strategy_reinit_unregisters_override_selectors():
|
|||
router = _build_router(routing_strategy="least-busy")
|
||||
override_selector = router._get_override_strategy_selector("latency-based-routing")
|
||||
assert override_selector is not None
|
||||
assert any(id(cb) == id(override_selector) for cb in litellm.callbacks)
|
||||
assert not any(cb is override_selector for cb in litellm.callbacks)
|
||||
|
||||
router.update_settings(routing_strategy="latency-based-routing")
|
||||
|
||||
assert router._override_selectors == {}
|
||||
assert not any(id(cb) == id(override_selector) for cb in litellm.callbacks)
|
||||
assert not any(cb is override_selector for cb in litellm.callbacks)
|
||||
assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger
|
||||
|
||||
|
||||
def test_override_selectors_are_not_registered_process_wide(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
router = _build_router(routing_strategy="simple-shuffle")
|
||||
v2_selector = router._get_override_strategy_selector("usage-based-routing-v2")
|
||||
least_busy_selector = router._get_override_strategy_selector("least-busy")
|
||||
assert v2_selector is not None and least_busy_selector is not None
|
||||
|
||||
assert litellm.callbacks == []
|
||||
assert litellm.input_callback == []
|
||||
|
||||
|
||||
def _rpm_limited_model_list():
|
||||
return [
|
||||
{
|
||||
"model_name": "other-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_key": "sk-test-3",
|
||||
"api_base": "https://example.invalid",
|
||||
"rpm": 1,
|
||||
},
|
||||
"model_info": {"id": "deploy-3"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def _mock_completion(router, **override):
|
||||
return await router.acompletion(
|
||||
model="other-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **override
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it():
|
||||
router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0)
|
||||
|
||||
await _mock_completion(router, routing_strategy="usage-based-routing-v2")
|
||||
override_selector = router._override_selectors["usage-based-routing-v2"]
|
||||
assert not any(cb is override_selector for cb in litellm.callbacks)
|
||||
|
||||
with patch.object(
|
||||
override_selector, "async_pre_call_check", wraps=override_selector.async_pre_call_check
|
||||
) as pre_call_spy:
|
||||
for _ in range(2):
|
||||
plain = await _mock_completion(router)
|
||||
assert plain.choices[0].message.content == "ok"
|
||||
assert not pre_call_spy.called
|
||||
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await _mock_completion(router, routing_strategy="usage-based-routing-v2")
|
||||
|
||||
|
||||
def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it():
|
||||
router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0)
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
|
||||
router.completion(
|
||||
model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2"
|
||||
)
|
||||
override_selector = router._override_selectors["usage-based-routing-v2"]
|
||||
assert not any(cb is override_selector for cb in litellm.callbacks)
|
||||
|
||||
for _ in range(2):
|
||||
plain = router.completion(model="other-model", messages=messages, mock_response="ok")
|
||||
assert plain.choices[0].message.content == "ok"
|
||||
|
||||
with pytest.raises(ValueError, match="No deployments available"):
|
||||
router.completion(
|
||||
model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_selector_pre_call_check_only_runs_for_override_selectors(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
deployment = _rpm_limited_model_list()[0]
|
||||
|
||||
override_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle")
|
||||
override_selector = override_router._get_override_strategy_selector("usage-based-routing-v2")
|
||||
await override_router._async_override_selector_pre_call_check(
|
||||
"usage-based-routing-v2", override_selector, deployment, None
|
||||
)
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
override_router._override_selector_pre_call_check("usage-based-routing-v2", override_selector, deployment)
|
||||
|
||||
default_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="usage-based-routing-v2")
|
||||
for _ in range(2):
|
||||
await default_router._async_override_selector_pre_call_check(
|
||||
"usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment, None
|
||||
)
|
||||
default_router._override_selector_pre_call_check(
|
||||
"usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment
|
||||
)
|
||||
await default_router._async_override_selector_pre_call_check(None, None, deployment, None)
|
||||
default_router._override_selector_pre_call_check(None, None, deployment)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested():
|
||||
router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0)
|
||||
kwargs = {"model": "deploy-3", "messages": [{"role": "user", "content": "hi"}], "mock_response": "ok"}
|
||||
|
||||
first = await router.acompletion(**kwargs, routing_strategy="usage-based-routing-v2")
|
||||
assert first.choices[0].message.content == "ok"
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(**kwargs, routing_strategy="usage-based-routing-v2")
|
||||
assert (await router.acompletion(**kwargs)).choices[0].message.content == "ok"
|
||||
|
||||
|
||||
def test_sync_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested():
|
||||
router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0)
|
||||
kwargs = {"model": "deploy-3", "messages": [{"role": "user", "content": "hi"}], "mock_response": "ok"}
|
||||
|
||||
first = router.completion(**kwargs, routing_strategy="usage-based-routing-v2")
|
||||
assert first.choices[0].message.content == "ok"
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
router.completion(**kwargs, routing_strategy="usage-based-routing-v2")
|
||||
assert router.completion(**kwargs).choices[0].message.content == "ok"
|
||||
|
||||
|
||||
def _pass_through_rpm_limited_model_list():
|
||||
deployment = _rpm_limited_model_list()[0]
|
||||
return [{**deployment, "litellm_params": {**deployment["litellm_params"], "use_in_pass_through": True}}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_early_return_paths_run_the_override_pre_call_check():
|
||||
router = Router(model_list=_pass_through_rpm_limited_model_list(), routing_strategy="simple-shuffle")
|
||||
override = {"routing_strategy": "usage-based-routing-v2"}
|
||||
|
||||
pinned = await router.async_get_available_deployment(
|
||||
model="other-model", request_kwargs={**override, "_encrypted_content_affinity_pinned": True}
|
||||
)
|
||||
assert pinned["model_info"]["id"] == "deploy-3"
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.async_get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override)
|
||||
plain = await router.async_get_available_deployment_for_pass_through(model="deploy-3", request_kwargs={})
|
||||
assert plain["model_info"]["id"] == "deploy-3"
|
||||
|
||||
|
||||
def test_sync_pass_through_specific_deployment_runs_the_override_pre_call_check():
|
||||
router = Router(model_list=_pass_through_rpm_limited_model_list(), routing_strategy="simple-shuffle")
|
||||
override = {"routing_strategy": "usage-based-routing-v2"}
|
||||
|
||||
first = router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override)
|
||||
assert first["model_info"]["id"] == "deploy-3"
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override)
|
||||
plain = router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs={})
|
||||
assert plain["model_info"]["id"] == "deploy-3"
|
||||
|
||||
|
||||
def _quality_group(strategy="latency-based-routing"):
|
||||
return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}]
|
||||
|
||||
|
|
|
|||
|
|
@ -3031,6 +3031,21 @@ def test_request_tags_after_router_consumption_drops_only_the_consumed_tags():
|
|||
assert _request_tags_after_router_consumption(partially_consumed, "gemini-flash") == ("deploy:us",)
|
||||
|
||||
|
||||
def test_request_tags_after_router_consumption_ignores_tags_merged_from_prior_deployments():
|
||||
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY
|
||||
from litellm.router_strategy.tag_based_routing import _request_tags_after_router_consumption
|
||||
from litellm.types.router import ConsumedRequestTagsStamp
|
||||
|
||||
metadata = {
|
||||
"tags": ["route", "®ion:eu", "free"],
|
||||
ROUTING_REQUEST_TAGS_METADATA_KEY: ("route", "®ion:eu"),
|
||||
"inherited_tags": ["®ion:eu"],
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY: ConsumedRequestTagsStamp(model_group="gemini-flash", tags=("route",)),
|
||||
}
|
||||
assert _request_tags_after_router_consumption(metadata, "gemini-flash") == ("®ion:eu",)
|
||||
assert _request_tags_after_router_consumption(metadata, "other-group") == ["route", "®ion:eu"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_non_router_tags_still_pick_the_matching_tier_deployment():
|
||||
# tags=["route", "deploy:us"]: "route" picks the router and is spent there,
|
||||
|
|
|
|||
54
tests/test_litellm/router_strategy/test_simple_shuffle.py
Normal file
54
tests/test_litellm/router_strategy/test_simple_shuffle.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
from collections import Counter
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict
|
||||
|
||||
DRAWS = 200
|
||||
|
||||
|
||||
def _deployment(dep_id: str, metric: LiteLLMParamsTypedDict | None = None) -> DeploymentTypedDict:
|
||||
params: LiteLLMParamsTypedDict = {"model": "gpt-4o", "api_key": "key", "mock_response": f"from {dep_id}"}
|
||||
return {
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {**params, **(metric or {})},
|
||||
"model_info": {"id": dep_id},
|
||||
}
|
||||
|
||||
|
||||
async def _draw_model_ids(router: Router) -> Counter[str]:
|
||||
counts: Counter[str] = Counter()
|
||||
for _ in range(DRAWS):
|
||||
response = await router.acompletion(model="test-model", messages=[{"role": "user", "content": "hi"}])
|
||||
counts[response._hidden_params["model_id"]] += 1
|
||||
return counts
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("metric", [{"weight": 5}, {"rpm": 5}, {"tpm": 5}], ids=["weight", "rpm", "tpm"])
|
||||
async def test_weighted_pick_when_only_a_later_deployment_carries_the_metric(metric: LiteLLMParamsTypedDict):
|
||||
router = Router(
|
||||
model_list=[_deployment("unweighted"), _deployment("weighted", metric)],
|
||||
routing_strategy="simple-shuffle",
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
counts = await _draw_model_ids(router)
|
||||
|
||||
assert counts["weighted"] == DRAWS
|
||||
assert counts["unweighted"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uniform_pick_when_every_configured_weight_is_zero():
|
||||
router = Router(
|
||||
model_list=[_deployment("unweighted"), _deployment("standby", {"weight": 0})],
|
||||
routing_strategy="simple-shuffle",
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
counts = await _draw_model_ids(router)
|
||||
|
||||
assert counts["unweighted"] > 0
|
||||
assert counts["standby"] > 0
|
||||
|
|
@ -7,6 +7,7 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
|
||||
from litellm.router_utils.auto_router_tuning_baseline import (
|
||||
DEFAULT_TUNING_FINGERPRINT,
|
||||
HEURISTIC_V1_TUNING_FIELDS,
|
||||
|
|
@ -21,6 +22,33 @@ from litellm.router_utils.auto_router_tuning_baseline import (
|
|||
|
||||
_TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}
|
||||
_ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"}
|
||||
_KEYWORD_DIMENSION: Final = {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}
|
||||
_HISTORICAL_FINGERPRINTS: Final = (
|
||||
({}, "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"),
|
||||
(
|
||||
{"custom_dimensions": [_KEYWORD_DIMENSION]},
|
||||
"b5c3c3f3be6341a8a16148d68d9067e03f94ed01a0bbfcde7955763042744372",
|
||||
),
|
||||
(
|
||||
{"custom_dimensions": [{"name": "sqlDdl", "weight": 0.4, "patterns": [r"\bCREATE\s{1,4}TABLE\b"]}]},
|
||||
"814ce0017fc7f60a160b262f658d910e9bdf784e6139a4ba4f1e2657aa203950",
|
||||
),
|
||||
(
|
||||
{
|
||||
"tiers": _TIERS,
|
||||
"dimension_weights": {"codePresence": 0.3},
|
||||
"custom_dimensions": [
|
||||
{
|
||||
"name": "internalFrameworks",
|
||||
"weight": 0.2,
|
||||
"keywords": ["orbitmesh", "fluxgate"],
|
||||
"patterns": [r"\bALTER\s{1,4}TABLE\b"],
|
||||
}
|
||||
],
|
||||
},
|
||||
"38970dc9224e265ab38c89674563d8d0537822591f9239b45251db6f5ca6cc39",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _router(
|
||||
|
|
@ -77,6 +105,27 @@ class TestTuningFingerprint:
|
|||
def test_explicit_empty_tier_model_configs_follow_omission(self) -> None:
|
||||
assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT
|
||||
|
||||
@pytest.mark.parametrize(("config", "fingerprint"), _HISTORICAL_FINGERPRINTS)
|
||||
def test_fingerprints_recorded_before_scoring_mode_existed_are_preserved(
|
||||
self, config: Mapping[str, object], fingerprint: str
|
||||
) -> None:
|
||||
"""Literal hashes captured from the merged implementation at 9bc9104102, before CustomDimension.scoring_mode."""
|
||||
assert tuning_fingerprint(config) == fingerprint
|
||||
|
||||
def test_binary_scoring_mode_hashes_like_its_absence(self) -> None:
|
||||
historical: Final = tuning_fingerprint({"custom_dimensions": [_KEYWORD_DIMENSION]})
|
||||
explicit: Final = tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "binary"}]})
|
||||
reserialized: Final = ComplexityRouterConfig.model_validate(
|
||||
{"custom_dimensions": [_KEYWORD_DIMENSION]}
|
||||
).model_dump(mode="json", include={"custom_dimensions"})
|
||||
assert reserialized["custom_dimensions"][0]["scoring_mode"] == "binary"
|
||||
assert reserialized["custom_dimensions"][0]["patterns"] == []
|
||||
assert historical == explicit == tuning_fingerprint(reserialized)
|
||||
assert (
|
||||
tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]})
|
||||
!= historical
|
||||
)
|
||||
|
||||
def test_tier_model_overrides_change_the_fingerprint(self) -> None:
|
||||
plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}})
|
||||
with_override = tuning_fingerprint(
|
||||
|
|
@ -218,15 +267,18 @@ class TestQuota:
|
|||
is None
|
||||
)
|
||||
|
||||
def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self) -> None:
|
||||
@pytest.mark.parametrize(
|
||||
"edit",
|
||||
[
|
||||
pytest.param({"weight": 0.9}, id="weight"),
|
||||
pytest.param({"scoring_mode": "match_count"}, id="scoring-mode"),
|
||||
],
|
||||
)
|
||||
def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self, edit: Mapping[str, object]) -> None:
|
||||
baselines: Final = snapshot_tuning_baselines(())
|
||||
original: Final = _router("a", {})
|
||||
config: Final = {
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}]
|
||||
}
|
||||
edited_config: Final = {
|
||||
"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.9, "keywords": ["orbitmesh"]}]
|
||||
}
|
||||
config: Final = {"custom_dimensions": [_KEYWORD_DIMENSION]}
|
||||
edited_config: Final = {"custom_dimensions": [{**_KEYWORD_DIMENSION, **edit}]}
|
||||
added: Final = _router("a", config)
|
||||
edited: Final = _router("a", edited_config)
|
||||
second: Final = _router("b", config)
|
||||
|
|
@ -240,6 +292,13 @@ class TestQuota:
|
|||
assert mutable_tuned_identities((original,), baselines) == frozenset()
|
||||
assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None
|
||||
|
||||
def test_graded_dimension_recorded_at_snapshot_is_its_own_baseline(self) -> None:
|
||||
graded: Final = _router("a", {"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]})
|
||||
baselines: Final = snapshot_tuning_baselines((graded,))
|
||||
assert mutable_tuned_identities((graded,), baselines) == frozenset()
|
||||
reverted_to_binary: Final = _router("a", {"custom_dimensions": [_KEYWORD_DIMENSION]})
|
||||
assert mutable_tuned_identities((reverted_to_binary,), baselines) == {router_identity(graded)}
|
||||
|
||||
def test_violation_message_names_the_limit_and_remedy(self) -> None:
|
||||
message = tuning_limit_violation(held=2, limit=1)
|
||||
assert message is not None
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
_get_deployment_cooldown_policy,
|
||||
_resolve_allowed_fails_from_policy,
|
||||
|
|
@ -269,18 +271,20 @@ class TestShouldCooldownBasedOnDeploymentPolicy:
|
|||
|
||||
|
||||
class TestShouldCooldownBasedOnAllowedFailsPolicy:
|
||||
def _make_router(self, cooldown_time: float = 60.0) -> MagicMock:
|
||||
def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | None = None) -> MagicMock:
|
||||
router = MagicMock()
|
||||
router.cooldown_time = cooldown_time
|
||||
router.allowed_fails = 0
|
||||
router.allowed_fails_policy = None
|
||||
router.get_allowed_fails_from_policy.return_value = None
|
||||
router.failed_calls.get_cache.return_value = None
|
||||
router.cache = cache if cache is not None else DualCache(in_memory_cache=InMemoryCache())
|
||||
return router
|
||||
|
||||
def test_cooldown_time_override_zero_is_not_falsy(self):
|
||||
"""cooldown_time_override=0 must be honored; it must not fall through to the router-level value."""
|
||||
router = self._make_router(cooldown_time=60.0)
|
||||
router.cache = MagicMock()
|
||||
router.cache.increment_cache.return_value = 1
|
||||
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
||||
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
|
|
@ -291,12 +295,68 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy:
|
|||
cooldown_time_override=0.0,
|
||||
)
|
||||
|
||||
set_cache_call = router.failed_calls.set_cache.call_args
|
||||
assert set_cache_call is not None
|
||||
assert set_cache_call[1]["ttl"] == 0.0, (
|
||||
increment_call = router.cache.increment_cache.call_args
|
||||
assert increment_call is not None
|
||||
assert increment_call[1]["ttl"] == 0.0, (
|
||||
"cooldown_time_override=0 should be used as TTL, not the router-level 60.0"
|
||||
)
|
||||
|
||||
def test_fail_counter_is_shared_across_router_instances(self):
|
||||
"""Two workers (two Router objects over one shared cache) must pool their failures toward allowed_fails."""
|
||||
shared_cache = DualCache(in_memory_cache=InMemoryCache())
|
||||
workers = (self._make_router(cache=shared_cache), self._make_router(cache=shared_cache))
|
||||
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
|
||||
results = [
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=workers[i % 2],
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
)
|
||||
for i in range(6)
|
||||
]
|
||||
|
||||
assert results == [False, False, False, False, False, True]
|
||||
assert shared_cache.get_cache(key="deployment:dep-1:allowed_fails") == 6
|
||||
|
||||
def test_fleet_wide_count_from_redis_decides_cooldown(self):
|
||||
"""The Redis (fleet-wide) count decides, even when this process has only seen one failure."""
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.increment_cache.return_value = 6
|
||||
router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
|
||||
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
)
|
||||
|
||||
assert result is True
|
||||
redis_cache.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0)
|
||||
|
||||
def test_redis_outage_falls_back_to_this_workers_count(self):
|
||||
"""When every Redis increment fails, the worker's own in-memory count must still cool the deployment down."""
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.increment_cache.side_effect = ConnectionError("redis down")
|
||||
router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
|
||||
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
|
||||
results = [
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
)
|
||||
for _ in range(6)
|
||||
]
|
||||
|
||||
assert results == [False, False, False, False, False, True]
|
||||
assert redis_cache.increment_cache.call_count == 6
|
||||
|
||||
|
||||
class TestRoutingGroupCooldownAlternatives:
|
||||
def _router(self, routing_groups=None):
|
||||
|
|
|
|||
|
|
@ -179,7 +179,7 @@ class TestHealthCheckCooldownIntegration:
|
|||
assert result is False
|
||||
|
||||
# Check counter was incremented
|
||||
current_fails = router.failed_calls.get_cache(key="deploy-1")
|
||||
current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails")
|
||||
assert current_fails == 1
|
||||
|
||||
def test_health_check_failure_triggers_cooldown_at_threshold(self):
|
||||
|
|
@ -263,7 +263,7 @@ class TestHealthCheckCooldownIntegration:
|
|||
assert "exception" not in healthy_endpoint
|
||||
|
||||
# Verify failed_calls counter is untouched
|
||||
current_fails = router.failed_calls.get_cache(key="deploy-1")
|
||||
current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails")
|
||||
assert current_fails is None
|
||||
|
||||
def test_disable_cooldowns_prevents_health_check_cooldown(self):
|
||||
|
|
|
|||
|
|
@ -223,6 +223,59 @@ def test_gating_matches_the_monolithic_entrypoint_and_get_secret_bool(
|
|||
assert monolith[1] == ("args=litellm --port 4000" if traced else "args=--port 4000")
|
||||
|
||||
|
||||
def test_wipes_the_prometheus_multiproc_dir_before_uvicorn_forks(tmp_path: Path) -> None:
|
||||
"""A restarted container inherits the emptyDir of its predecessor, whose worker pids it may reuse, so the
|
||||
stale .db files must be gone before any worker opens the one carrying its own pid."""
|
||||
multiproc_dir = tmp_path / "multiproc"
|
||||
multiproc_dir.mkdir()
|
||||
(multiproc_dir / "gauge_livesum_7.db").write_bytes(b"stale")
|
||||
(multiproc_dir / "counter_7.db").write_bytes(b"stale")
|
||||
(multiproc_dir / "keep.txt").write_text("not a sample")
|
||||
|
||||
bin_dir = tmp_path / "bin"
|
||||
bin_dir.mkdir()
|
||||
_write_stubs(bin_dir, ("uvicorn",))
|
||||
record = tmp_path / "record.txt"
|
||||
env = {
|
||||
**os.environ,
|
||||
"PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
|
||||
"RECORD": str(record),
|
||||
"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir),
|
||||
}
|
||||
env.pop("USE_DDTRACE", None)
|
||||
result = subprocess.run(
|
||||
["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
|
||||
assert sorted(p.name for p in multiproc_dir.iterdir()) == ["keep.txt"]
|
||||
assert record.read_text().splitlines()[0] == "exec=uvicorn"
|
||||
|
||||
|
||||
def test_creates_a_missing_prometheus_multiproc_dir(tmp_path: Path) -> None:
|
||||
bin_dir = tmp_path / "bin"
|
||||
bin_dir.mkdir()
|
||||
_write_stubs(bin_dir, ("uvicorn",))
|
||||
missing = tmp_path / "multiproc"
|
||||
env = {
|
||||
**os.environ,
|
||||
"PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
|
||||
"RECORD": str(tmp_path / "record.txt"),
|
||||
"PROMETHEUS_MULTIPROC_DIR": str(missing),
|
||||
}
|
||||
env.pop("USE_DDTRACE", None)
|
||||
result = subprocess.run(
|
||||
["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"], env=env, capture_output=True, text=True
|
||||
)
|
||||
|
||||
assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
|
||||
assert missing.is_dir()
|
||||
|
||||
|
||||
def _copied_script(dockerfile: Path, image_path: str) -> Path:
|
||||
"""Resolve the repo file a Dockerfile `COPY`s to `image_path`, so tests run what the image ships."""
|
||||
matches = _COPY_RE.findall(dockerfile.read_text())
|
||||
|
|
|
|||
33
tests/test_litellm/test_drop_params_env_var.py
Normal file
33
tests/test_litellm/test_drop_params_env_var.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _import_litellm_with(configured: str) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
[sys.executable, "-c", "import litellm; print(litellm.drop_params)"],
|
||||
env={**os.environ, "LITELLM_DROP_PARAMS": configured},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True"), ("", "False")])
|
||||
def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected):
|
||||
result = _import_litellm_with(configured)
|
||||
|
||||
assert result.stdout.strip() == expected
|
||||
assert "is not a flag value" not in result.stderr
|
||||
|
||||
|
||||
def test_litellm_drop_params_env_var_non_flag_value_stays_on_with_a_warning():
|
||||
result = _import_litellm_with("temperature")
|
||||
|
||||
assert result.stdout.strip() == "True"
|
||||
assert (
|
||||
"LITELLM_DROP_PARAMS='temperature' is not a flag value, treating it as on. Set it to true or false"
|
||||
in result.stderr
|
||||
)
|
||||
|
|
@ -40,7 +40,7 @@ from litellm.router import (
|
|||
_is_retriable_anthropic_status,
|
||||
)
|
||||
from litellm.router_strategy import simple_shuffle
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
from litellm.types.router import DeploymentTypedDict, RetryPolicy
|
||||
|
||||
|
||||
def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata():
|
||||
|
|
@ -6132,6 +6132,50 @@ def test_update_kwargs_with_deployment_no_tags():
|
|||
assert "tags" not in kwargs["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_does_not_narrow_tag_filtered_group_to_failed_deployments_tags():
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "tagged-group",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.5",
|
||||
"api_key": "fake-key",
|
||||
"tags": ["free"],
|
||||
"input_cost_per_token": 0.000001,
|
||||
"output_cost_per_token": 0.000001,
|
||||
"mock_response": "litellm.ContextWindowExceededError",
|
||||
},
|
||||
"model_info": {"id": "tagged-failing"},
|
||||
},
|
||||
{
|
||||
"model_name": "tagged-group",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.5",
|
||||
"api_key": "fake-key",
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.001,
|
||||
"mock_response": "ok",
|
||||
},
|
||||
"model_info": {"id": "untagged-healthy"},
|
||||
},
|
||||
],
|
||||
routing_strategy="cost-based-routing",
|
||||
enable_tag_filtering=True,
|
||||
num_retries=2,
|
||||
retry_after=0,
|
||||
retry_policy=RetryPolicy(BadRequestErrorRetries=2),
|
||||
)
|
||||
metadata: Final[dict[str, object]] = {}
|
||||
|
||||
response = await router.acompletion(
|
||||
model="tagged-group", messages=[{"role": "user", "content": "hi"}], metadata=metadata
|
||||
)
|
||||
|
||||
assert response._hidden_params["model_id"] == "untagged-healthy"
|
||||
assert metadata["tags"] == ["free"]
|
||||
|
||||
|
||||
def test_update_kwargs_with_deployment_merges_tools():
|
||||
"""
|
||||
Test that when both deployment litellm_params and request have tools,
|
||||
|
|
@ -14390,3 +14434,70 @@ async def test_router_max_parallel_requests_slot_released_when_stream_closed_ear
|
|||
|
||||
assert tracker.peak == 1
|
||||
assert tracker.current == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch):
|
||||
from litellm import Router
|
||||
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5-nano",
|
||||
"api_key": "sk-fake",
|
||||
"temperature": 1,
|
||||
"reasoning_effort": "minimal",
|
||||
"drop_params": "true",
|
||||
"mock_response": "Hello, world!",
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-nano")
|
||||
assert deployment is not None
|
||||
assert deployment.litellm_params.drop_params is True
|
||||
|
||||
response = await router.acompletion(
|
||||
model="gpt-5-nano",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
temperature=0.1,
|
||||
)
|
||||
assert response.choices[0].message.content == "Hello, world!"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["ture", "enabled"])
|
||||
def test_router_warns_when_a_deployment_drop_params_string_is_not_a_flag(value, caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {"model": "openai/gpt-5-nano", "api_key": "sk-fake", "drop_params": value},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-nano")
|
||||
assert deployment is not None
|
||||
assert deployment.litellm_params.drop_params == value
|
||||
assert f"model=gpt-5-nano drop_params={value!r} is not a flag value, treating it as unset" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [True, "true", "off", None])
|
||||
def test_router_stays_quiet_when_a_deployment_drop_params_is_a_flag(value, caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {"model": "openai/gpt-5-nano", "api_key": "sk-fake", "drop_params": value},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert "is not a flag value" not in caplog.text
|
||||
|
|
|
|||
|
|
@ -6094,6 +6094,34 @@ class TestFinalOptionalParamsLineRedaction:
|
|||
assert "'temperature': 0.25" in printed
|
||||
|
||||
|
||||
class TestDropParamsStringCoercion:
|
||||
@pytest.mark.parametrize("drop_params", ["true", "True", True])
|
||||
def test_truthy_drop_params_drops_unsupported_temperature(self, drop_params, monkeypatch):
|
||||
from litellm.utils import get_optional_params
|
||||
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
result = get_optional_params(
|
||||
model="gpt-5-nano",
|
||||
custom_llm_provider="openai",
|
||||
temperature=0.1,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
assert "temperature" not in result
|
||||
|
||||
@pytest.mark.parametrize("drop_params", ["false", False, None])
|
||||
def test_falsy_drop_params_still_raises(self, drop_params, monkeypatch):
|
||||
from litellm.utils import get_optional_params
|
||||
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
with pytest.raises(litellm.UnsupportedParamsError):
|
||||
get_optional_params(
|
||||
model="gpt-5-nano",
|
||||
custom_llm_provider="openai",
|
||||
temperature=0.1,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
|
||||
def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
|
||||
return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()]
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.types.router import (
|
||||
SPECIAL_MODEL_INFO_PARAMS,
|
||||
Deployment,
|
||||
GenericLiteLLMParams,
|
||||
LiteLLM_Params,
|
||||
ModelInfo,
|
||||
)
|
||||
|
|
@ -89,3 +92,33 @@ def test_pricing_strings_are_coerced_to_float():
|
|||
def test_invalid_pricing_is_rejected():
|
||||
with pytest.raises(ValueError, match='validation error for ModelInfo'):
|
||||
ModelInfo(id="x", input_cost_per_token="free")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, expected",
|
||||
[
|
||||
(True, True),
|
||||
("true", True),
|
||||
(" False ", False),
|
||||
("yes", True),
|
||||
(None, None),
|
||||
("os.environ/DROP_PARAMS", "os.environ/DROP_PARAMS"),
|
||||
("v2:gcm:ciphertext-from-a-pre-fix-row", "v2:gcm:ciphertext-from-a-pre-fix-row"),
|
||||
],
|
||||
)
|
||||
def test_drop_params_coerces_flags_and_keeps_unresolved_strings(value, expected):
|
||||
assert GenericLiteLLMParams(drop_params=value).drop_params == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [2, 2.5, [], {}])
|
||||
def test_drop_params_ignores_non_flag_non_string_values_with_a_warning(value, caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
assert GenericLiteLLMParams(drop_params=value).drop_params is None
|
||||
assert f"drop_params={value!r} is not a flag value" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [True, "true", None, "os.environ/DROP_PARAMS", "v2:gcm:ciphertext-from-a-pre-fix-row"])
|
||||
def test_drop_params_flags_and_strings_log_nothing(value, caplog):
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
GenericLiteLLMParams(drop_params=value)
|
||||
assert caplog.text == ""
|
||||
|
|
|
|||
|
|
@ -44,8 +44,9 @@ import {
|
|||
} from "./ComplexityRouterConfig";
|
||||
|
||||
const DEFAULT_SCORING_EXPLANATION =
|
||||
"The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " +
|
||||
"terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:";
|
||||
"The router scores each request across 7 built-in dimensions: token count, code presence, reasoning markers, technical " +
|
||||
"terms, simple indicators, multi-step patterns, and question complexity, plus any custom dimensions you add. " +
|
||||
"The weighted score determines the tier:";
|
||||
|
||||
const HEURISTIC_V2_EXPLANATION =
|
||||
"The router estimates success probability for all four tiers with the bundled calibrated model, then selects " +
|
||||
|
|
|
|||
|
|
@ -1521,14 +1521,16 @@ describe("ComplexityRouterConfig tier editing", () => {
|
|||
renderWithProviders(<ComplexityRouterConfig {...baseProps} value={customValue} onEditingTiersChange={vi.fn()} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
expect(screen.queryByText("How Classification Works")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("scores each request across 7 dimensions", { exact: false })).not.toBeInTheDocument();
|
||||
expect(
|
||||
screen.queryByText("scores each request across 7 built-in dimensions", { exact: false }),
|
||||
).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the scorer card on a built-in router, whose tiers the score still decides", () => {
|
||||
renderWithProviders(<ComplexityRouterConfig {...baseProps} onEditingTiersChange={vi.fn()} />);
|
||||
fireEvent.click(screen.getByText("Advanced: Classification Method"));
|
||||
expect(screen.getByText("How Classification Works")).toBeInTheDocument();
|
||||
expect(screen.getByText("scores each request across 7 dimensions", { exact: false })).toBeInTheDocument();
|
||||
expect(screen.getByText("scores each request across 7 built-in dimensions", { exact: false })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("says why a custom row is blocked instead of only reddening its border", () => {
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ import EscalationKeywords from "./EscalationKeywords";
|
|||
import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules";
|
||||
import SemanticKeywordMatching from "./SemanticKeywordMatching";
|
||||
import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs";
|
||||
import { type CustomDimensionRow } from "./custom_dimensions";
|
||||
import CompressionControls from "./CompressionControls";
|
||||
import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression";
|
||||
|
||||
|
|
@ -432,6 +433,11 @@ export interface ComplexityRouterConfigValue {
|
|||
tier_boundaries?: TierBoundaries;
|
||||
token_thresholds?: TokenThresholds;
|
||||
dimension_weights?: DimensionWeights;
|
||||
/**
|
||||
* Operator-added scoring dimensions, each carrying its own inline weight. Undefined means the router has
|
||||
* none and keeps the key out of the payload; an empty array is a real "the last row was removed" state.
|
||||
*/
|
||||
custom_dimensions?: CustomDimensionRow[];
|
||||
/**
|
||||
* Score floor the reasoning-marker override must clear. Undefined keeps the key out of the payload, so the
|
||||
* floor tracks tier_boundaries.simple_medium; an explicit 0 is a real floor that promotes on the markers alone.
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue