Merge branch 'litellm_internal_staging' into litellm_/determined-faraday-03159e
Some checks failed
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled

This commit is contained in:
yuneng-jiang 2026-08-18 22:12:24 -07:00 committed by GitHub
commit a8d060fa23
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
897 changed files with 67008 additions and 21693 deletions

View file

@ -43,6 +43,7 @@ jobs:
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a

View file

@ -35,6 +35,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
# Models & routing config
"/model/",
"/v1/model/info",
"/v1/model/deprecations",
"/v2/model/",
"/model_group",
"/model_access_group/",

View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 22344
"limit": 22343
},
"reportArgumentType": {
"limit": 2578
@ -57,7 +57,7 @@
"limit": 5681
},
"reportMissingTypeArgument": {
"limit": 15609
"limit": 15605
},
"reportMissingTypeStubs": {
"limit": 40
@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1824
"limit": 1823
},
"reportRedeclaration": {
"limit": 8
@ -108,7 +108,7 @@
"limit": 39154
},
"reportUnknownParameterType": {
"limit": 19947
"limit": 19944
},
"reportUnknownVariableType": {
"limit": 30772
@ -132,7 +132,7 @@
"limit": 27
},
"reportUnusedClass": {
"limit": 23
"limit": 21
},
"reportUnusedFunction": {
"limit": 139

View file

@ -41,6 +41,11 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
},
"additionalProperties": False,
},
"guardrail_cost_per_unit": {
"type": "object",
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
"additionalProperties": NONNEG_NUMBER,
},
"metadata": {
"type": "object",
"description": "Free-form notes about the entry (e.g. pricing derivation).",

View file

@ -24,12 +24,16 @@ if TYPE_CHECKING:
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
"complete",
"failed",
"expired",
"cancelled",
)
TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
*PROVIDER_TERMINAL_BATCH_STATUSES,
"stale_expired",
)
@ -286,6 +290,57 @@ class CheckBatchCost:
404 must not retire the row; the staleness sweep bounds it instead."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
"""A 404 naming the output file means there is nothing to fetch on this or any
later poll: providers like Vertex AI advertise an output path for every batch,
including terminal ones that never wrote it. Any other failure may be
transient, so it keeps retrying until the staleness sweep bounds it."""
import openai
from litellm.exceptions import NotFoundError
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
async def _finalize_unbilled_terminal_job(
self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
) -> None:
"""Persist a terminal batch that has nothing billable, converting any raw
provider file ids to managed ids, and take it out of the poll page."""
try:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
ensure_batch_response_managed_file_ids,
)
response.id = job.unified_object_id
await ensure_batch_response_managed_file_ids(
response=response,
managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"),
prisma_client=self.prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await self.prisma_client.db.litellm_managedobjecttable.update(
where={"id": job.id},
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
)
@staticmethod
def _record_error(
prom_logger: Optional["PrometheusLogger"], error_type: str
@ -528,6 +583,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.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -648,15 +704,20 @@ class CheckBatchCost:
f"{_file_attr}={_raw_file_id!r}: {_e}"
)
# Pass deployment model_info so custom batch pricing
# (input_cost_per_token_batches etc.) is used for cost calc
deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {}
# Pass the deployment's router-registered pricing (litellm_params custom
# rates merged with the model's published rates) so custom batch pricing
# (input_cost_per_token_batches etc.) is used for cost calc, exactly as
# the inline retrieve path does.
deployment_model_info = deployment_pricing_model_info(
model_id=model_id,
deployment_model=litellm_model_name,
)
batch_cost, batch_usage, batch_models = (
await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info, # type: ignore[arg-type]
model_info=deployment_model_info,
)
)
logging_obj = LiteLLMLogging(
@ -796,7 +857,7 @@ class CheckBatchCost:
## RETRIEVE THE BATCH JOB OUTPUT FILE
if (
response.status in ("completed", "complete", "expired")
response.status in PROVIDER_TERMINAL_BATCH_STATUSES
and response.output_file_id is not None
):
try:
@ -808,6 +869,15 @@ class CheckBatchCost:
prom_logger=prom_logger,
)
except Exception as tracking_err:
if self._is_output_file_gone_at_provider(
tracking_err, response.output_file_id
) and self._batch_deployment_exists(model_id):
verbose_proxy_logger.warning(
f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} "
f"does not exist at the provider; retiring job {job.id} unbilled"
)
await self._finalize_unbilled_terminal_job(job, response)
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
@ -837,45 +907,8 @@ class CheckBatchCost:
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
elif response.status in (
"completed",
"complete",
"failed",
"expired",
"cancelled",
):
try:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
ensure_batch_response_managed_file_ids,
)
response.id = job.unified_object_id
await ensure_batch_response_managed_file_ids(
response=response,
managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"),
prisma_client=self.prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data = {
"status": response.status,
"file_object": response.model_dump_json(),
}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
await self.prisma_client.db.litellm_managedobjecttable.update(
where={"id": job.id},
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
)
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
await self._finalize_unbilled_terminal_job(job, response)
# Record polling run metrics (always, even if nothing was processed)
if prom_logger:

View file

@ -41,6 +41,7 @@ from litellm.proxy._types import (
CallTypes,
LiteLLM_ManagedFileTable,
LiteLLM_ManagedObjectTable,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
@ -423,13 +424,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# This is because the encoded object ids stored in the managed objects table do not contain the provider information
# To support provider filtering, we would need to store the provider information in the encoded object ids
if provider:
raise Exception("Filtering by 'provider' is not supported when using managed batches.")
raise ProxyException(
message="Filtering by 'provider' is not supported when using managed batches.",
type="invalid_request_error",
param="provider",
code=400,
)
# Model name filtering is not supported for managed batches
# This is because the encoded object ids stored in the managed objects table do not contain the model name
# A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids.
if target_model_names:
raise Exception("Filtering by 'target_model_names' is not supported when using managed batches.")
raise ProxyException(
message="Filtering by 'target_model_names' is not supported when using managed batches.",
type="invalid_request_error",
param="target_model_names",
code=400,
)
if limit == 0:
return build_list_page([])
owner_filter = build_owner_filter(user_api_key_dict)
if owner_filter is None:

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.56"
version = "0.1.57"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.56"
version = "0.1.57"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -83,6 +83,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/azure_ai/",
"/aws/",
"/bedrock/",
"/comprehendmedical",
"/cohere/",
"/gemini/",
"/google/",

View file

@ -119,4 +119,7 @@ spec:
{{- end }}
ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }}
backoffLimit: {{ .Values.migrationJob.backoffLimit }}
{{- with .Values.migrationJob.activeDeadlineSeconds }}
activeDeadlineSeconds: {{ . }}
{{- end }}
{{- end }}

View file

@ -314,3 +314,31 @@ tests:
operator: Equal
value: litellm-e2e
effect: NoSchedule
- it: bounds the Job with a deadline by default, so a blocked migration cannot stall the release forever
set:
migrationJob:
enabled: true
asserts:
- equal:
path: spec.activeDeadlineSeconds
value: 1800
- it: honours an operator-supplied deadline
set:
migrationJob:
enabled: true
activeDeadlineSeconds: 600
asserts:
- equal:
path: spec.activeDeadlineSeconds
value: 600
- it: omits the deadline entirely when it is nulled out, restoring the unbounded behaviour
set:
migrationJob:
enabled: true
activeDeadlineSeconds: null
asserts:
- notExists:
path: spec.activeDeadlineSeconds

View file

@ -427,6 +427,13 @@ migrationJob:
enabled: true # Enable or disable the schema migration Job
retries: 3 # Number of retries for the Job in case of failure
backoffLimit: 4 # Backoff limit for Job restarts
# Wall-clock budget for the whole Job, shared across every `backoffLimit`
# retry rather than granted per attempt. Without it a migration that blocks
# on the database never fails, and when the Helm hook is enabled the release
# waits on it forever: `helm upgrade` and any GitOps controller driving it
# stop reconciling the whole chart until someone deletes the Job by hand.
# Set to null to opt out and restore the unbounded behaviour.
activeDeadlineSeconds: 1800
disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0.
# Optional service account for the migration job.
# Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true.

View file

@ -24,7 +24,7 @@
"/v1/video" "/v1/videos" "/video" "/videos" "/v1/search" "/search"
"/v1/containers" "/containers" "/v1/evals" "/v1/memory" "/queue/chat"
"/v1beta" "/interactions"
"/anthropic" "/azure" "/azure_ai" "/aws" "/bedrock" "/cohere" "/gemini" "/google"
"/anthropic" "/azure" "/azure_ai" "/aws" "/bedrock" "/comprehendmedical" "/cohere" "/gemini" "/google"
"/vertex_ai" "/vertex-ai" "/assemblyai" "/eu.assemblyai" "/langfuse" "/vllm"
"/mistral" "/groq" "/voyage" "/cursor" "/milvus" "/openai_passthrough"
"/toolset"

View file

@ -21,6 +21,9 @@ metadata:
spec:
backoffLimit: {{ .Values.migrationJob.backoffLimit }}
ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }}
{{- with .Values.migrationJob.activeDeadlineSeconds }}
activeDeadlineSeconds: {{ . }}
{{- end }}
template:
metadata:
{{- /* The Job's selector is generated by the controller rather than

View file

@ -167,3 +167,24 @@ tests:
- equal:
path: spec.template.metadata.labels['app.kubernetes.io/component']
value: batch-migrations
- it: bounds the Job with a deadline by default, so a blocked migration cannot stall the release forever
asserts:
- equal:
path: spec.activeDeadlineSeconds
value: 1800
- it: honours an operator-supplied deadline
set:
migrationJob.activeDeadlineSeconds: 600
asserts:
- equal:
path: spec.activeDeadlineSeconds
value: 600
- it: omits the deadline entirely when it is nulled out, restoring the unbounded behaviour
set:
migrationJob.activeDeadlineSeconds: null
asserts:
- notExists:
path: spec.activeDeadlineSeconds

View file

@ -56,6 +56,15 @@ migrationJob:
enabled: true
backoffLimit: 4
ttlSecondsAfterFinished: 120
# Wall-clock budget for the whole Job, shared across every `backoffLimit`
# retry rather than granted per attempt. Without it a migration that blocks
# on the database never fails, and because this is a pre-upgrade hook the
# release waits on it forever: `helm upgrade` and any GitOps controller
# driving it stop reconciling the whole chart until someone deletes the Job
# by hand. A migration that has exhausted its retries is not going to
# succeed on the next one, so failing is strictly better than hanging.
# Set to null to opt out and restore the unbounded behaviour.
activeDeadlineSeconds: 1800
resources: {}
# ServiceAccount for the Job pod only.
#

View file

@ -0,0 +1,16 @@
-- CreateTable
CREATE TABLE "LiteLLM_DailyGuardrailUsageUnits" (
"guardrail_id" TEXT NOT NULL,
"date" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"api_key" TEXT NOT NULL,
"usage_unit" TEXT NOT NULL,
"units" BIGINT NOT NULL DEFAULT 0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_DailyGuardrailUsageUnits_pkey" PRIMARY KEY ("guardrail_id","date","team_id","api_key","usage_unit")
);
-- CreateIndex
CREATE INDEX "LiteLLM_DailyGuardrailUsageUnits_date_idx" ON "LiteLLM_DailyGuardrailUsageUnits"("date");

View file

@ -1069,6 +1069,21 @@ model LiteLLM_DailyGuardrailMetrics {
@@index([guardrail_id])
}
// Daily guardrail billable usage units (one row per guardrail/day/team/key/unit type)
model LiteLLM_DailyGuardrailUsageUnits {
guardrail_id String
date String // YYYY-MM-DD
team_id String // empty string when the request had no team
api_key String // hashed virtual key; empty string when unknown
usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits
units BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@id([guardrail_id, date, team_id, api_key, usage_unit])
@@index([date])
}
// Daily policy metrics for usage dashboard (one row per policy per day)
model LiteLLM_DailyPolicyMetrics {
policy_id String

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.86"
version = "0.4.87"
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.86"
version = "0.4.87"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -792,6 +792,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
nlp_cloud_models.add(key)
elif value.get("litellm_provider") == "aleph_alpha":
aleph_alpha_models.add(key)
elif value.get("litellm_provider") == "bedrock" and value.get("mode") == "guardrail":
pass
elif value.get("litellm_provider") == "bedrock" and not is_bedrock_pricing_only_model(key):
bedrock_models.add(key)
elif value.get("litellm_provider") == "bedrock_converse":

View file

@ -10,7 +10,7 @@ from typing import Any, Final
import litellm
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.secret_redaction import redact_string, redact_structured_value
set_verbose = False
@ -59,6 +59,12 @@ def _redact_string(value: str) -> str:
return redact_string(value)
def _redact_structured_value(key: str | None, value: str) -> str:
if not _ENABLE_SECRET_REDACTION:
return value
return redact_structured_value(key, value)
def redact_secrets(value: str) -> str:
"""Public API: redact known secret/credential patterns from an arbitrary string.
@ -265,7 +271,7 @@ class JsonFormatter(Formatter):
if record.exc_info:
json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info)
return safe_dumps(json_record)
return safe_dumps(json_record, value_transform=_redact_structured_value)
class CorrelationPlainFormatter(logging.Formatter):
@ -276,7 +282,7 @@ class CorrelationPlainFormatter(logging.Formatter):
"""
def format(self, record: logging.LogRecord) -> str:
formatted: Final = super().format(record)
formatted: Final = _redact_string(super().format(record))
trace_id: Final = getattr(record, "trace_id", None)
session_id: Final = getattr(record, "session_id", None)
if not trace_id and not session_id:

View file

@ -48,6 +48,7 @@ async def _handle_completed_batch(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: str | None = None,
litellm_params: dict | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
@ -58,6 +59,9 @@ async def _handle_completed_batch(
custom_llm_provider: The LLM provider
model_name: Optional model name
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
model_info: Optional deployment-level model info with custom pricing,
threaded through so a deployment's configured rates win over the
global cost map.
"""
# A completed batch whose request lines all failed has no output file - the
# results are written to a separate error_file_id and output_file_id is None.
@ -86,6 +90,7 @@ async def _handle_completed_batch(
entries=_iter_batch_input_entries(file_content),
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
)
@ -441,11 +446,23 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
"""
if custom_llm_provider in ("anthropic", "bedrock"):
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
return AnthropicConfig().calculate_usage(
usage_object=response_body.get("usage", None) or {},
usage_object: Final = response_body.get("usage", None) or {}
if custom_llm_provider == "bedrock" and AmazonConverseConfig.is_converse_usage_shape(usage_object):
return AmazonConverseConfig().usage_from_batch_output(usage_object)
anthropic_usage: Final = AnthropicConfig().calculate_usage(
usage_object=usage_object,
reasoning_content=None,
)
if usage_object and anthropic_usage.total_tokens == 0:
verbose_logger.warning(
"batch output line reported usage this parser does not understand, so it will be billed at $0. "
"provider=%s usage_keys=%s",
custom_llm_provider,
sorted(usage_object.keys()),
)
return anthropic_usage
from litellm.responses.utils import ResponseAPILoggingUtils
_usage_dict: Final = response_body.get("usage", None) or {}

View file

@ -107,7 +107,7 @@ async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -157,7 +157,7 @@ def create_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -339,7 +339,9 @@ def create_batch(
@client
async def aretrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
] = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -385,7 +387,9 @@ def _handle_retrieve_batch_providers_without_provider_config(
litellm_params: dict,
_retrieve_batch_request: RetrieveBatchRequest,
_is_async: bool,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
] = "openai",
logging_obj: Any | None = None,
):
api_base: str | None = None
@ -508,7 +512,9 @@ def _handle_retrieve_batch_providers_without_provider_config(
@client
def retrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
] = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -826,7 +832,7 @@ def list_batches(
async def acancel_batch(
batch_id: str,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy"] = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -872,7 +878,7 @@ async def acancel_batch(
def cancel_batch(
batch_id: str,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] | str = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "litellm_proxy"] | str = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,

View file

@ -12,8 +12,11 @@ This module is dependency-injected: callers pass the proxy ``llm_router`` and
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final
import litellm
if TYPE_CHECKING:
from litellm.router import Router
@ -41,3 +44,28 @@ def build_router_embedding_metadata(
metadata: Final[dict[str, Any]] = dict(request_metadata or {})
metadata["semantic-cache-embedding"] = True
return metadata
def resolve_embedding_max_input_tokens(
configured_max_input_tokens: int | None,
embedding_model: str,
router: Router | None,
) -> int | None:
"""Explicit cache setting first, else the Router deployment's configured ``max_input_tokens``."""
if configured_max_input_tokens is not None:
return configured_max_input_tokens
if router is None:
return None
deployment_max_input_tokens, _ = router.get_configured_token_limits(embedding_model)
return deployment_max_input_tokens
def truncate_embedding_input(prompt: str, embedding_model: str, max_input_tokens: int | None) -> str:
"""Keep only the first ``max_input_tokens`` tokens of ``prompt`` for the embedding call."""
if max_input_tokens is None:
return prompt
tokens: Final[Sequence[int]] = litellm.encode(model=embedding_model, text=prompt)
if len(tokens) <= max_input_tokens:
return prompt
truncated: Final[str] = litellm.decode(model=embedding_model, tokens=tokens[:max_input_tokens])
return truncated

View file

@ -97,6 +97,7 @@ class Cache:
qdrant_quantization_config: str | None = None,
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
qdrant_semantic_cache_vector_size: int | None = None,
semantic_cache_embedding_max_input_tokens: int | None = None,
# GCP IAM authentication parameters
gcp_service_account: str | None = None,
gcp_ssl_ca_certs: str | None = None,
@ -122,6 +123,7 @@ class Cache:
qdrant_api_key (str, optional): The api_key for the local or cloud qdrant cluster.
qdrant_collection_name (str, optional): The name for your qdrant collection. Required if type is "qdrant-semantic".
similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
# Disk Cache Args
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
@ -192,6 +194,7 @@ class Cache:
similarity_threshold=similarity_threshold,
embedding_model=redis_semantic_cache_embedding_model,
index_name=redis_semantic_cache_index_name,
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
**kwargs,
)
elif type == LiteLLMCacheType.VALKEY_SEMANTIC:
@ -207,6 +210,7 @@ class Cache:
embedding_model=valkey_semantic_cache_embedding_model,
index_name=valkey_semantic_cache_index_name,
startup_nodes=redis_startup_nodes,
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
**kwargs,
)
elif type == LiteLLMCacheType.QDRANT_SEMANTIC:
@ -218,6 +222,7 @@ class Cache:
quantization_config=qdrant_quantization_config,
embedding_model=qdrant_semantic_cache_embedding_model,
vector_size=qdrant_semantic_cache_vector_size,
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
)
elif type == LiteLLMCacheType.LOCAL:
self.cache = InMemoryCache()

View file

@ -12,7 +12,7 @@ import ast
import asyncio
import json
import os
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import print_verbose
@ -22,12 +22,21 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.types.utils import EmbeddingResponse
from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router
from ._embedding_router import (
build_router_embedding_metadata,
resolve_embedding_max_input_tokens,
resolve_embedding_router,
truncate_embedding_input,
)
from .base_cache import BaseCache
if TYPE_CHECKING:
from litellm.router import Router
class QdrantSemanticCache(BaseCache):
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
embedding_max_input_tokens: int | None = None
def __init__(
self,
@ -39,6 +48,7 @@ class QdrantSemanticCache(BaseCache):
embedding_model="text-embedding-ada-002",
host_type=None,
vector_size=None,
embedding_max_input_tokens: int | None = None,
):
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
@ -57,6 +67,7 @@ class QdrantSemanticCache(BaseCache):
raise Exception("similarity_threshold must be provided, passed None")
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
self.embedding_max_input_tokens = embedding_max_input_tokens
self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
headers = {}
@ -188,6 +199,13 @@ class QdrantSemanticCache(BaseCache):
cached_key: Final = payload.get(self.CACHE_KEY_FIELD_NAME)
return cached_key is not None and str(cached_key) == str(key)
def _embedding_input(self, prompt: str, router: "Router | None") -> str:
return truncate_embedding_input(
prompt,
self.embedding_model,
resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router),
)
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
"""Embed via the proxy Router when it serves the model, else direct."""
try:
@ -197,16 +215,17 @@ class QdrantSemanticCache(BaseCache):
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
if router is not None:
return router.embedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
)
return litellm.embedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
)
@ -218,17 +237,18 @@ class QdrantSemanticCache(BaseCache):
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
if router is not None:
return await router.aembedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
)
return await litellm.aembedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
)

View file

@ -14,7 +14,7 @@ import asyncio
import json
import os
from collections.abc import Callable, Mapping
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -23,9 +23,17 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.types.utils import EmbeddingResponse
from ._embedding_router import build_router_embedding_metadata, resolve_embedding_router
from ._embedding_router import (
build_router_embedding_metadata,
resolve_embedding_max_input_tokens,
resolve_embedding_router,
truncate_embedding_input,
)
from .base_cache import BaseCache
if TYPE_CHECKING:
from litellm.router import Router
class RedisSemanticCache(BaseCache):
"""
@ -38,6 +46,7 @@ class RedisSemanticCache(BaseCache):
DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index"
CACHE_KEY_FIELD_NAME: str = "litellm_cache_key"
embedding_max_input_tokens: int | None = None
def __init__(
self,
@ -48,6 +57,7 @@ class RedisSemanticCache(BaseCache):
similarity_threshold: float | None = None,
embedding_model: str = "text-embedding-ada-002",
index_name: str | None = None,
embedding_max_input_tokens: int | None = None,
**kwargs: object,
):
"""
@ -62,6 +72,8 @@ class RedisSemanticCache(BaseCache):
where 1.0 requires exact matches and 0.0 accepts any match
embedding_model: Model to use for generating embeddings
index_name: Name for the Redis index
embedding_max_input_tokens: Truncate prompts to this many tokens before
embedding; defaults to the Router deployment's configured max_input_tokens
ttl: Default time-to-live for cache entries in seconds
**kwargs: Additional arguments passed to the Redis client
@ -86,6 +98,7 @@ class RedisSemanticCache(BaseCache):
# While similarity: 1 = most similar, 0 = least similar
self.distance_threshold = 1 - similarity_threshold
self.embedding_model = embedding_model
self.embedding_max_input_tokens = embedding_max_input_tokens
# Set up Redis connection
if redis_url is None:
@ -307,6 +320,13 @@ class RedisSemanticCache(BaseCache):
return dict_method()
return value
def _embedding_input(self, prompt: str, router: "Router | None") -> str:
return truncate_embedding_input(
prompt,
self.embedding_model,
resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router),
)
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]:
"""
Routes through the proxy Router when the embedding model is a Router
@ -320,12 +340,13 @@ class RedisSemanticCache(BaseCache):
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
if router is not None:
embedding_response = cast(
EmbeddingResponse,
router.embedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
),
@ -335,7 +356,7 @@ class RedisSemanticCache(BaseCache):
EmbeddingResponse,
litellm.embedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
),
)
@ -490,18 +511,19 @@ class RedisSemanticCache(BaseCache):
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
try:
if router is not None:
embedding_response = await router.aembedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
)
else:
embedding_response = await litellm.aembedding(
model=self.embedding_model,
input=prompt,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
)
return embedding_response["data"][0]["embedding"]

View file

@ -17,7 +17,6 @@ RedisSemanticCache since those are backend agnostic.
import asyncio
import hashlib
import os
import struct
from dataclasses import dataclass
from typing import Any, Final
@ -29,6 +28,7 @@ from redis.commands.search.query import Query
from litellm._logging import print_verbose
from litellm._uuid import uuid
from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
from .redis_semantic_cache import RedisSemanticCache
@ -61,6 +61,7 @@ class ValkeySemanticCache(RedisSemanticCache):
startup_nodes: list | None = None,
sync_client: Redis | None = None,
async_client: AsyncRedis | None = None,
embedding_max_input_tokens: int | None = None,
**kwargs: Any,
):
if similarity_threshold is None:
@ -78,6 +79,7 @@ class ValkeySemanticCache(RedisSemanticCache):
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
self.embedding_max_input_tokens = embedding_max_input_tokens
self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME
self.key_prefix = f"{self.index_name}:"
self._index_dim: int | None = None
@ -92,19 +94,17 @@ class ValkeySemanticCache(RedisSemanticCache):
@staticmethod
def _build_valkey_url(host: str | None, port: str | None, password: str | None, ssl: bool = False) -> str:
host = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST")
port = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT")
password = password or os.environ.get("VALKEY_PASSWORD") or os.environ.get("REDIS_PASSWORD")
resolved_host: Final = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST")
resolved_port: Final = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT")
resolved_password: Final = password or os.environ.get("VALKEY_PASSWORD") or os.environ.get("REDIS_PASSWORD")
if not host or not port:
if not resolved_host or not resolved_port:
raise ValueError(
"Missing required Valkey configuration. Provide host and port "
"(or VALKEY_HOST/VALKEY_PORT), or pass redis_url."
)
credentials: Final = f":{password}@" if password else ""
scheme: Final = "rediss" if ssl else "redis"
return f"{scheme}://{credentials}{host}:{port}"
return build_valkey_url(host=resolved_host, port=resolved_port, password=resolved_password, ssl=ssl)
@classmethod
def _scope_tag(cls, key: str) -> str:
@ -116,7 +116,7 @@ class ValkeySemanticCache(RedisSemanticCache):
@staticmethod
def _embedding_to_bytes(embedding: list[float]) -> bytes:
return struct.pack(f"<{len(embedding)}f", *embedding)
return pack_vector(embedding)
def _index_schema(self, dim: int) -> tuple[TagField, VectorField]:
return (

View file

@ -185,6 +185,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if not isinstance(tool_choice, dict):
return tool_choice
choice_type: Final = tool_choice.get("type")
if isinstance(choice_type, str) and choice_type in ("auto", "none", "required"):
return choice_type
if choice_type not in ("function", "custom"):
return tool_choice
if isinstance(tool_choice.get("name"), str) and tool_choice.get("name"):

View file

@ -1,5 +1,6 @@
import os
import sys
from types import MappingProxyType
from typing import Final, Literal
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
@ -1487,6 +1488,7 @@ WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
SPEND_LOG_RUN_LOOPS: Final = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
SPEND_LOG_CLEANUP_BATCH_SIZE: Final = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3))
@ -1763,3 +1765,7 @@ PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
# one run delete a charge another just wrote. A stale row is hours old and a concurrent
# one is seconds old, so a few minutes separates them.
PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300
# Shared read-only empty mapping, for defaulting optional Mapping parameters without
# constructing a fresh mutable dict at each call site.
EMPTY_MAPPING: Final = MappingProxyType({})

View file

@ -2160,7 +2160,7 @@ def batch_cost_calculator(
output_cost_per_token: Final = model_info.get("output_cost_per_token")
total_prompt_cost = 0.0
total_completion_cost = 0.0
if input_cost_per_token_batches:
if input_cost_per_token_batches is not None:
total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches
elif input_cost_per_token:
details: Final = parse_prompt_tokens_details(usage)
@ -2180,7 +2180,7 @@ def batch_cost_calculator(
cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token
total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2
if output_cost_per_token_batches:
if output_cost_per_token_batches is not None:
total_completion_cost = usage.completion_tokens * output_cost_per_token_batches
elif output_cost_per_token:
total_completion_cost = (

View file

@ -23,12 +23,15 @@ FileCreateProvider = Literal[
"vertex_ai",
"bedrock",
"hosted_vllm",
"litellm_proxy",
"manus",
"anthropic",
]
FileRetrieveProvider = Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic"]
FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "manus", "anthropic"]
FileRetrieveProvider = Literal[
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic"
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"]
import litellm
from litellm import get_secret_str
from litellm.files.streaming import FileContentStreamingResponse

View file

@ -1,7 +1,9 @@
from collections.abc import AsyncIterator, Iterator
from typing import Literal, NamedTuple
FileContentProvider = Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"]
FileContentProvider = Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus"
]
class FileContentStreamingResult(NamedTuple):

View file

@ -5,6 +5,7 @@ import datetime
import os
import random
import time
from collections.abc import Callable
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, Literal
@ -17,7 +18,11 @@ import litellm.litellm_core_utils.litellm_logging
import litellm.types
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.constants import HOURS_IN_A_DAY, SLACK_DAILY_REPORT_LOCK_ID
from litellm.constants import (
HOURS_IN_A_DAY,
SLACK_DAILY_REPORT_LOCK_ID,
SLACK_MODEL_DEPRECATION_LOCK_ID,
)
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type
from litellm.integrations.SlackAlerting.hanging_request_check import (
@ -45,6 +50,10 @@ from litellm.repositories.table_repositories import InvitationLinkRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.types.integrations.slack_alerting import *
from litellm.types.proxy.model_deprecation import (
DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
DEPRECATION_IDLE_POLL_SECONDS,
)
from ..email_templates.templates import *
from .batching_handler import send_to_webhook, squash_payloads
@ -59,6 +68,12 @@ else:
Router = Any
def _proxy_llm_router() -> Router | None:
from litellm.proxy.proxy_server import llm_router
return llm_router
class SlackAlerting(CustomBatchLogger):
"""
Class for sending Slack Alerts
@ -1044,6 +1059,99 @@ Model Info:
async def model_removed_alert(self, model_name: str):
pass
def _deprecation_alerts_enabled(self) -> bool:
return self.alerting is not None and AlertType.model_deprecation_warnings in self.alert_types
async def send_model_deprecation_alert(
self,
llm_router: Router | None = None,
pod_lock_manager: "PodLockManager | None" = None,
) -> bool:
"""Alert on the router's deprecated and imminent models, True when one was sent
The daily lock is claimed only once there is something to say, so an empty pass never blocks a
later real one, and a sent alert is stamped in the shared cache for a day so sibling pods stop asking
"""
if not self._deprecation_alerts_enabled():
return False
from litellm.proxy.common_utils.model_deprecation import (
collect_model_deprecations,
format_deprecation_alert_message,
)
snapshot: Final = collect_model_deprecations(llm_router=llm_router)
message: Final = format_deprecation_alert_message(snapshot)
if message is None:
return False
if not await self._claimed_deprecation_alert_window(pod_lock_manager):
return False
level: Final[Literal["Low", "Medium", "High"]] = "High" if snapshot.deprecated else "Medium"
await self.send_alert(
message=message,
level=level,
alert_type=AlertType.model_deprecation_warnings,
alerting_metadata={ # mutable-ok: send_alert takes a dict payload
"deprecated_count": len(snapshot.deprecated),
"imminent_count": len(snapshot.imminent),
"upcoming_count": len(snapshot.upcoming),
},
)
await self.internal_usage_cache.async_set_cache(
key=SlackAlertingCacheKeys.deprecation_alert_sent_key.value,
value=time.time(),
ttl=DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
)
return True
async def _claimed_deprecation_alert_window(self, pod_lock_manager: "PodLockManager | None") -> bool:
"""Without a redis backed lock there is no fleet to coordinate, so a lone pod always alerts"""
if pod_lock_manager is None:
return True
return (
await pod_lock_manager.acquire_lock(
cronjob_id=SLACK_MODEL_DEPRECATION_LOCK_ID,
ttl=DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
allow_reentrant=False,
)
) is not False
async def _deprecation_alert_sent_within_a_day(self) -> bool:
return (
await self.internal_usage_cache.async_get_cache(key=SlackAlertingCacheKeys.deprecation_alert_sent_key.value)
) is not None
async def _run_deprecation_alert_pass(
self, llm_router: Router | None, pod_lock_manager: "PodLockManager | None"
) -> bool:
if llm_router is None or not self._deprecation_alerts_enabled():
return False
if await self._deprecation_alert_sent_within_a_day():
return False
return await self.send_model_deprecation_alert(llm_router=llm_router, pod_lock_manager=pod_lock_manager)
async def run_scheduled_deprecation_check(
self,
get_llm_router: Callable[[], Router | None] = _proxy_llm_router,
pod_lock_manager: "PodLockManager | None" = None,
) -> None:
"""Poll every pass for a loaded router, the alert being on, and no alert in the last day, then alert
A pass that could not alert (no router yet, alert type off, a sibling pod holds the daily lock, or a
redis blip at claim time) is retried on the next poll instead of costing a day, while a pass that
raised (a missing webhook, say) backs off a full day so a misconfiguration logs once, not every poll
"""
while True:
try:
await self._run_deprecation_alert_pass(get_llm_router(), pod_lock_manager)
except Exception as e: # noqa: BLE001 # a failed alert must not kill the loop
verbose_proxy_logger.exception("Error in model deprecation alert loop: %s", e)
await asyncio.sleep(DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS)
continue
await asyncio.sleep(DEPRECATION_IDLE_POLL_SECONDS)
async def send_webhook_alert(self, webhook_event: WebhookEvent) -> bool:
"""
Sends structured alert to webhook, if set.

View file

@ -1,5 +1,8 @@
import os
from collections.abc import Mapping
import threading
from collections import OrderedDict
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
@ -17,6 +20,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
OTELSemconvCategory,
parse_semconv_opt_in,
)
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
from litellm.integrations.otel.model.semconv import Metric
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.secret_redaction import redact_string
@ -83,6 +87,12 @@ class _ResponseWithUsageView(TypedDict, total=False):
usage: "_UsageCompletionTokensView | None"
# Cap on credential-scoped providers held at once; each one owns an exporter thread.
_MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256
# Dedicated so a slow exporter shutdown cannot starve the shared logging executor.
_PROVIDER_SHUTDOWN_EXECUTOR: Final = ThreadPoolExecutor(max_workers=4, thread_name_prefix="OtelProviderShutdown")
LITELLM_TRACER_NAME: Final = os.getenv("OTEL_TRACER_NAME", "litellm")
LITELLM_METER_NAME: Final = os.getenv("LITELLM_METER_NAME", "litellm")
LITELLM_LOGGER_NAME: Final = os.getenv("LITELLM_LOGGER_NAME", "litellm")
@ -227,6 +237,34 @@ def _freeze_for_dedupe(value: object, _depth: int = 0) -> HashableScope:
return repr(value)
def _shutdown_tracer_provider(provider: "_SDKTracerProvider") -> None:
"""Flush and stop a dropped provider so its exporter thread is reclaimed."""
try:
provider.shutdown()
except Exception as e: # noqa: BLE001 # exporter shutdown must not fail the request that dropped it
verbose_logger.debug("OpenTelemetry: error shutting down dropped tracer provider: %s", e)
@dataclass(frozen=True, slots=True)
class _CachedTracerProvider:
"""A cached credential-scoped provider plus whether it may be shut down when dropped."""
provider: "_SDKTracerProvider"
owns_exporter: bool
def _provider_owns_exporter(exporter: "str | _SpanExporter") -> bool:
"""Whether a provider built for ``exporter`` may be shut down when it is dropped.
``_get_span_processor`` builds a fresh exporter for a named kind, but wraps a
caller-supplied ``SpanExporter`` instance as-is, and that instance is shared with the
logger's own provider. Shutting a dropped provider down would then stop exporting for
the whole process. The shared case also uses ``SimpleSpanProcessor``, so it owns no
thread and there is nothing to reclaim.
"""
return not hasattr(exporter, "export")
@dataclass
class OpenTelemetryConfig:
exporter: str | SpanExporter = "console"
@ -322,6 +360,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
tracer_provider: object | None = None,
logger_provider: object | None = None,
meter_provider: object | None = None,
max_dynamic_tracer_providers: int = _MAX_DYNAMIC_TRACER_PROVIDERS,
**kwargs,
):
team_metadata_keys_override: Final = kwargs.pop("baggage_team_metadata_keys", None)
@ -347,7 +386,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
self.OTEL_EXPORTER = self.config.exporter
self.OTEL_ENDPOINT = self.config.endpoint
self.OTEL_HEADERS = self.config.headers
self._tracer_provider_cache: dict[str, _SDKTracerProvider] = {}
self._tracer_provider_cache: OrderedDict[str, _CachedTracerProvider] = OrderedDict()
self._tracer_provider_cache_lock: Final = threading.Lock()
self._max_dynamic_tracer_providers: Final = max(1, max_dynamic_tracer_providers)
self._init_tracing(tracer_provider)
_debug_otel: Final = str(os.getenv("DEBUG_OTEL", "False")).lower()
@ -678,6 +719,28 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
self._handle_failure(kwargs, response_obj, start_time, end_time)
def _start_service_span(self, payload: ServiceLoggerPayload, parent_otel_span: Span, start_time_ns: int) -> Span:
"""Open a service span, named and classified by what the service is.
A datastore call is an outbound CLIENT span carrying ``db.*`` semconv.
Without those a Postgres span says only ``service=postgres``, so the
backend falls back to the transport peer, which for Prisma is the local
query engine on loopback. Everything else stays an INTERNAL span.
"""
from opentelemetry import trace
from opentelemetry.trace import SpanKind
attributes: Final = db_span_attributes(payload.service.value, payload.call_type)
span: Final = self.tracer.start_span(
name=payload.service,
context=trace.set_span_in_context(parent_otel_span),
start_time=start_time_ns,
kind=SpanKind.CLIENT if attributes else SpanKind.INTERNAL,
)
for key, value in attributes.items():
self.safe_set_attribute(span=span, key=key, value=value)
return span
async def async_service_success_hook(
self,
payload: ServiceLoggerPayload,
@ -686,7 +749,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
end_time: datetime | float | None = None,
event_metadata: dict | None = None,
):
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
_start_time_ns = 0
@ -703,12 +765,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
_end_time_ns = self._to_ns(end_time)
if parent_otel_span is not None:
_span_name: Final = payload.service
service_logging_span: Final = self.tracer.start_span(
name=_span_name,
context=trace.set_span_in_context(parent_otel_span),
start_time=_start_time_ns,
)
service_logging_span: Final = self._start_service_span(payload, parent_otel_span, _start_time_ns)
self.safe_set_attribute(
span=service_logging_span,
key="call_type",
@ -746,7 +803,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
end_time: float | datetime | None = None,
event_metadata: dict | None = None,
):
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
_start_time_ns = 0
@ -763,12 +819,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
_end_time_ns = self._to_ns(end_time)
if parent_otel_span is not None:
_span_name: Final = payload.service
service_logging_span: Final = self.tracer.start_span(
name=_span_name,
context=trace.set_span_in_context(parent_otel_span),
start_time=_start_time_ns,
)
service_logging_span: Final = self._start_service_span(payload, parent_otel_span, _start_time_ns)
self.safe_set_attribute(
span=service_logging_span,
key="call_type",
@ -1027,38 +1078,98 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
return self.construct_dynamic_otel_config(standard_callback_dynamic_params=standard_callback_dynamic_params)
def _get_tracer_with_dynamic_config(self, dynamic_config: OpenTelemetryConfig):
def _insert_or_drop(
self, cache_key: str, built: "_CachedTracerProvider"
) -> "tuple[_CachedTracerProvider, _CachedTracerProvider | None]":
"""Cache ``built`` under ``cache_key``, returning the entry to use and what to drop.
Caller holds ``_tracer_provider_cache_lock``. The drop is either the loser of a
concurrent build for this key or the LRU victim its insertion pushed out.
"""
raced: Final = self._tracer_provider_cache.get(cache_key)
if raced is not None:
self._tracer_provider_cache.move_to_end(cache_key)
return raced, built
self._tracer_provider_cache[cache_key] = built
if len(self._tracer_provider_cache) > self._max_dynamic_tracer_providers:
return built, self._tracer_provider_cache.popitem(last=False)[1]
return built, None
def _cached_dynamic_tracer(
self,
cache_key: str,
build: Callable[[], "_SDKTracerProvider"],
owns_exporter: bool,
) -> "_Tracer":
"""Return the tracer for ``cache_key``, building and caching a provider on miss.
A provider that owns its exporter also owns a ``BatchSpanProcessor`` worker thread
that only stops on ``shutdown()``, so the cache is a bounded LRU and whatever it
drops is shut down. Without both, a proxy serving key-scoped credentials accumulates
one live thread per credential set for the life of the process.
``owns_exporter`` also decides ``shutdown_on_exit`` at build time: a provider we may
never shut down must not hold an interpreter-exit hook, which would both pin it in
memory for the life of the process and stop the shared exporter at exit. Those
providers use ``SimpleSpanProcessor``, which buffers nothing, so the hook costs them
no flush.
``owns_exporter`` describes the provider being built, and is cached with it, because
the two dynamic entry points share this cache and can disagree: whether the LRU
victim may be shut down is a property of the victim, never of the request that
happened to evict it.
"""
with self._tracer_provider_cache_lock:
cached: Final = self._tracer_provider_cache.get(cache_key)
if cached is not None:
self._tracer_provider_cache.move_to_end(cache_key)
return cached.provider.get_tracer(LITELLM_TRACER_NAME)
# Built outside the lock: exporter construction can block on DNS/TLS.
built: Final = _CachedTracerProvider(provider=build(), owns_exporter=owns_exporter)
with self._tracer_provider_cache_lock:
winner, dropped = self._insert_or_drop(cache_key, built)
if dropped is not None and dropped.owns_exporter:
# Off the caller's thread: shutdown joins the exporter worker.
_PROVIDER_SHUTDOWN_EXECUTOR.submit(_shutdown_tracer_provider, dropped.provider)
return winner.provider.get_tracer(LITELLM_TRACER_NAME)
def _get_tracer_with_dynamic_config(self, dynamic_config: OpenTelemetryConfig) -> "_Tracer":
"""Create (or reuse) a tracer whose exporter target comes from a per-request config."""
from opentelemetry.sdk.trace import TracerProvider
cache_key = f"dynamic_config:{dynamic_config.exporter}:{dynamic_config.endpoint}:{dynamic_config.headers}"
if cache_key in self._tracer_provider_cache:
return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME)
owns_exporter: Final = _provider_owns_exporter(dynamic_config.exporter)
temp_provider: Final = TracerProvider(resource=self._get_litellm_resource(self.config))
temp_provider.add_span_processor(self._get_span_processor(config_override=dynamic_config))
def _build() -> "_SDKTracerProvider":
provider: Final = TracerProvider(
resource=self._get_litellm_resource(self.config), shutdown_on_exit=owns_exporter
)
provider.add_span_processor(self._get_span_processor(config_override=dynamic_config))
return provider
self._tracer_provider_cache[cache_key] = temp_provider
cache_key: Final = (
f"dynamic_config:{dynamic_config.exporter}:{dynamic_config.endpoint}:{dynamic_config.headers}"
)
return self._cached_dynamic_tracer(cache_key, _build, owns_exporter)
return temp_provider.get_tracer(LITELLM_TRACER_NAME)
def _get_tracer_with_dynamic_headers(self, dynamic_headers: dict):
"""Create a temporary tracer with dynamic headers for this request only."""
def _get_tracer_with_dynamic_headers(self, dynamic_headers: Mapping[str, str]) -> "_Tracer":
"""Create (or reuse) a tracer whose OTLP headers come from a per-request credential set."""
from opentelemetry.sdk.trace import TracerProvider
# Prevents thread exhaustion by reusing providers for the same credential sets (e.g. per-team keys)
owns_exporter: Final = _provider_owns_exporter(self.OTEL_EXPORTER)
def _build() -> "_SDKTracerProvider":
provider: Final = TracerProvider(
resource=self._get_litellm_resource(self.config), shutdown_on_exit=owns_exporter
)
provider.add_span_processor(self._get_span_processor(dynamic_headers=dynamic_headers))
return provider
cache_key: Final = str(sorted(dynamic_headers.items()))
if cache_key in self._tracer_provider_cache:
return self._tracer_provider_cache[cache_key].get_tracer(LITELLM_TRACER_NAME)
# Create a temporary tracer provider with dynamic headers
temp_provider: Final = TracerProvider(resource=self._get_litellm_resource(self.config))
temp_provider.add_span_processor(self._get_span_processor(dynamic_headers=dynamic_headers))
# Store in cache for reuse
self._tracer_provider_cache[cache_key] = temp_provider
return temp_provider.get_tracer(LITELLM_TRACER_NAME)
return self._cached_dynamic_tracer(cache_key, _build, owns_exporter)
def construct_dynamic_otel_headers(
self, standard_callback_dynamic_params: StandardCallbackDynamicParams
@ -2832,7 +2943,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
def _get_span_processor(
self,
dynamic_headers: dict | None = None,
dynamic_headers: Mapping[str, str] | None = None,
config_override: OpenTelemetryConfig | None = None,
):
from opentelemetry.sdk.trace.export import (
@ -3144,7 +3255,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
@staticmethod
def _get_headers_dictionary(
headers: str | dict | None,
headers: "str | Mapping[str, str] | None",
) -> dict[str, str]:
"""
Convert a string or dictionary of headers into a dictionary of headers.
@ -3158,8 +3269,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
for part in parts:
key, value = part.split("=", 1)
_split_otel_headers[key] = value
elif isinstance(headers, dict):
_split_otel_headers = headers
elif isinstance(headers, Mapping):
_split_otel_headers.update(headers)
return _split_otel_headers
async def async_management_endpoint_success_hook(

View file

@ -18,6 +18,7 @@ from litellm.integrations.otel.mappers.utils import (
serialize_messages,
tool_definition_attrs,
)
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
from litellm.integrations.otel.model.payloads import (
GuardrailSpanData,
LLMCallSpanData,
@ -27,7 +28,6 @@ from litellm.integrations.otel.model.payloads import (
ToolDefinition,
)
from litellm.integrations.otel.model.semconv import (
DB,
MCP,
Error,
GenAI,
@ -36,7 +36,6 @@ from litellm.integrations.otel.model.semconv import (
RpcSystem,
Server,
)
from litellm.integrations.otel.model.spans import db_system
class GenAIMapper:
@ -182,12 +181,8 @@ class GenAIMapper:
def _service(cls, data: ServiceSpanData) -> AttributeMap:
attrs: Final = collect(cls._SERVICE_ATTRS, data)
# An outbound datastore call (DB_CALL / CLIENT span) also carries db.*
# semconv. Internal services (router, budget jobs, …) have no db.system,
# so they get only the litellm.service.* keys above.
system: Final = db_system(data.service_name)
if system is not None:
attrs[DB.SYSTEM_NAME] = system
if data.call_type:
attrs[DB.OPERATION_NAME] = data.call_type
# semconv naming the server it reached. Internal services (router, budget
# jobs, …) have no db.system, so they get only the litellm.service.* keys.
attrs.update(db_span_attributes(data.service_name, data.call_type))
attrs.update({f"{LiteLLM.METADATA_PREFIX}{key}": value for key, value in data.event_metadata.items()})
return attrs

View file

@ -0,0 +1,164 @@
"""OTel ``db.*`` / ``server.*`` attributes naming the database litellm talks to.
Prisma reaches PostgreSQL through a query engine listening on loopback, so
transport-level instrumentation attributes the work to ``localhost`` and an
operator cannot tell it is a PostgreSQL call or correlate it with the database's
own metrics. These attributes name the real server on litellm's DB spans.
Only the host, port, database and schema of the DSN are read, so no credential
can reach an exporter.
"""
from __future__ import annotations
import os
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from urllib.parse import ParseResult, parse_qs, unquote, urlparse
from litellm.integrations.otel.model.semconv import DB, Server
from litellm.integrations.otel.model.spans import POSTGRESQL, db_system
_DATABASE_URL_ENV: Final = "DATABASE_URL"
_READ_REPLICA_ENV: Final = "DATABASE_URL_READ_REPLICA"
_DEFAULT_POSTGRES_PORT: Final = 5432
_DEFAULT_POSTGRES_SCHEMA: Final = "public"
_POSTGRES_SCHEMES: Final = frozenset({"postgres", "postgresql"})
_EMPTY_ATTRIBUTES: Final[Mapping[str, str | int]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class DatabaseEndpoint:
"""The non-sensitive identity of a PostgreSQL server, parsed from a DSN."""
address: str | None
port: int | None
namespace: str | None
def parse_database_endpoint(url: str | None) -> DatabaseEndpoint | None:
"""Parse a PostgreSQL DSN into its exportable endpoint identity.
Returns ``None`` for an absent, malformed or non-PostgreSQL URL rather than
raising: an unparseable DSN must degrade to a span without endpoint
attributes, never break the request that emitted it.
"""
if not url:
return None
try:
parsed: Final = urlparse(url)
if parsed.scheme not in _POSTGRES_SCHEMES:
return None
query: Final = parse_qs(parsed.query)
raw_database: Final = (parsed.path or "").lstrip("/")
if _is_misparsed_authority(parsed, url, raw_database):
return None
# ``host=`` beats the netloc: it is how libpq names a Unix socket
# directory and how the Cloud SQL connector sits behind a localhost
# netloc, where the netloc is the very answer this module replaces.
address: Final = _first(query.get("host")) or parsed.hostname
# ``port=`` accompanies ``host=`` in a libpq URI, so honour it the same way.
port: Final = _port(_first(query.get("port")), parsed.port) if address else None
namespace: Final = _namespace(unquote(raw_database), _first(query.get("schema")))
except ValueError:
return None
if address is None and namespace is None:
return None
return DatabaseEndpoint(address=address, port=port, namespace=namespace)
def _is_misparsed_authority(parsed: ParseResult, url: str, raw_database: str) -> bool:
"""Whether the URL authority may have been truncated by an unencoded character.
``/``, ``#`` or ``?`` in a password ends the netloc early, so urlparse hands
back the username as the host, the leading digits of the password as the
port, and the rest of the credential as the path, query or fragment. The
stranded userinfo ``@`` is the only surviving evidence.
A database name cannot hold an unencoded slash either, so a second path
segment is the same evidence.
A DSN that carries the at-sign in a query parameter instead, such as
``?application_name=svc@prod``, is indistinguishable from a mis-split by any
property of the parse: both leave no userinfo, a host, a port and a path.
Since guessing wrong publishes a credential fragment to a tracing backend,
that ambiguity resolves to refusing the endpoint. Such a DSN loses
``server.address`` and ``db.namespace`` and keeps the rest of the span,
which is the cheaper error of the two. Percent-encode the at-sign to keep
them.
"""
if "/" in raw_database:
return True
return "@" in url and "@" not in parsed.netloc
def _first(values: Sequence[str] | None) -> str:
return values[0] if values else ""
def _port(from_query: str, from_netloc: int | None) -> int:
return int(from_query) if from_query.isdigit() else (from_netloc or _DEFAULT_POSTGRES_PORT)
def _namespace(database: str, schema: str) -> str | None:
"""``{database}|{schema}`` per the PostgreSQL semconv, dropping absent halves.
Only Prisma's literal default schema stays implicit. The match is
case-sensitive because Prisma quotes the name, so ``?schema=PUBLIC`` builds
a second schema alongside ``public`` and the two must not collapse to one
namespace.
"""
qualifier: Final = "" if schema == _DEFAULT_POSTGRES_SCHEMA else schema
return "|".join(part for part in (database, qualifier) if part) or None
def postgres_endpoint() -> DatabaseEndpoint | None:
"""The PostgreSQL endpoint the process is currently connected to.
Read from ``os.environ`` on every span, deliberately, on both counts.
The environment is what Prisma itself connects with, so the span cannot
disagree with the connection; ``get_secret_str`` would consult a configured
secret manager first and could name a different server than the one serving
the query. And the value is not static: the RDS IAM refresh rebuilds the URL
from ``DATABASE_HOST``/``PORT``/``NAME``/``SCHEMA`` every rotation, the
reconnect path re-reads ``DATABASE_URL``, and the DB-backed
``environment_variables`` config overlay can rewrite any of them after
startup, so a value cached for the process lifetime goes stale against a
connection that has genuinely moved. Nothing is memoized either: a cache
keyed on the URL would hold a rotated credential past its rotation, and the
parse is a single ``urlparse`` on a short string.
A configured read replica yields ``None``: ``RoutingPrismaWrapper`` picks
reader or writer per Prisma call, underneath the span, so naming the writer
would attribute replica reads to the primary.
"""
if os.environ.get(_READ_REPLICA_ENV):
return None
return parse_database_endpoint(os.environ.get(_DATABASE_URL_ENV, ""))
def db_span_attributes(service_name: str, call_type: str | None = None) -> Mapping[str, str | int]:
"""The ``db.*``/``server.*`` attributes for a datastore service call.
Empty for services that are not outbound datastore calls. Endpoint
attributes are PostgreSQL-only: ``DATABASE_URL`` says nothing about where
the redis-backed services point. ``db.system`` rides alongside the current
``db.system.name`` because Datadog's OTLP intake still types a database span
from the older key.
"""
system: Final = db_system(service_name)
if system is None:
return _EMPTY_ATTRIBUTES
endpoint: Final = postgres_endpoint() if system == POSTGRESQL else None
pairs: Final[tuple[tuple[str, str | int | None], ...]] = (
(DB.SYSTEM_NAME, system),
(DB.SYSTEM_LEGACY, system),
(DB.OPERATION_NAME, call_type),
(Server.ADDRESS, endpoint.address if endpoint is not None else None),
(Server.PORT, endpoint.port if endpoint is not None else None),
(DB.NAMESPACE, endpoint.namespace if endpoint is not None else None),
)
return MappingProxyType({key: value for key, value in pairs if value})

View file

@ -238,7 +238,11 @@ class DB:
"""
SYSTEM_NAME: Final = "db.system.name"
# Superseded by SYSTEM_NAME, dual-emitted because Datadog's OTLP intake
# still infers a span's database type from this key.
SYSTEM_LEGACY: Final = "db.system"
OPERATION_NAME: Final = "db.operation.name"
NAMESPACE: Final = "db.namespace"
class HTTP:

View file

@ -115,10 +115,12 @@ SPAN_REGISTRY: Final[dict[SpanRole, SpanSpec]] = {
# redis-backed spend queues. Any service not mapped here is litellm-internal work
# and stays an INTERNAL ``SERVICE`` span. This table is the single source of
# datastore knowledge — both the role classifier and the mapper read it.
POSTGRESQL: Final = "postgresql"
_DB_SYSTEM_BY_SERVICE: Final[dict[str, str]] = {
"redis": "redis",
"postgres": "postgresql",
"batch_write_to_db": "postgresql",
"postgres": POSTGRESQL,
"batch_write_to_db": POSTGRESQL,
}

View file

@ -9,13 +9,14 @@ across pods or stop races; the hook reads active jobs through a short-TTL cache.
import asyncio
import hashlib
import random
import traceback
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import groupby
from operator import itemgetter
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Literal
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator, model_validator
@ -32,6 +33,7 @@ from litellm.litellm_core_utils.llm_judge import (
parse_json_verdict,
)
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalDirection
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
@ -55,7 +57,7 @@ _MAX_JUDGE_PROMPT_CHARS: Final = 24_000
# The judge answers with a small JSON object; a tighter budget truncates the JSON
# mid-object and the attempt is lost to an error row.
JUDGE_MAX_OUTPUT_TOKENS: Final = 500
JUDGE_MAX_OUTPUT_TOKENS: Final = 1500
_MAX_ERROR_CHARS: Final = 500
@ -305,16 +307,21 @@ Criteria: correctness, completeness, clarity, conciseness.
Return ONLY valid JSON in this exact format, no other text:
{
"preference": "A" | "B" | "tie",
"confidence": <0.0 to 1.0>,
"reasoning": "<one sentence>"
"confidence": <0.0 to 1.0>
}"""
class PairwiseVerdict(BaseModel):
"""The judge's blind A/B verdict, validated at the parse boundary."""
"""The judge's blind A/B verdict: the response_format schema sent with the judge call
and the validation contract on its reply. Both fields are required and preference is
closed over the prompt's labels, so a malformed or truncated reply is an
unparseable-verdict error row, never a defaulted or fabricated verdict."""
preference: str = "tie"
confidence: float = 0.0
preference: Literal["A", "B", "tie"]
confidence: float
PAIRWISE_JUDGE_RESPONSE_FORMAT: Final = type_to_response_format_param(PairwiseVerdict)
def _sample_hits(request_id: str, job_id: str, percentage: float) -> bool:
@ -325,6 +332,14 @@ def _sample_hits(request_id: str, job_id: str, percentage: float) -> bool:
return bucket * 100.0 < percentage
def _failure_detail(e: BaseException) -> str:
"""Exception class, message, and the raising frame, so an attempt's error row names
the faulty code path without needing debug logs on the pod."""
frames: Final = traceback.extract_tb(e.__traceback__)
location: Final = f" at {frames[-1].filename.rsplit('/', 1)[-1]}:{frames[-1].lineno}" if frames else ""
return f"{type(e).__name__}{location}: {e}"
def _judge_call_cost(response: object) -> float:
"""Price a judge call, treating an unmapped judge model as free rather than fatal."""
import litellm
@ -764,7 +779,9 @@ class ShadowEvalLogger(CustomLogger):
try:
response: Final = await router.acompletion(
model=target_model,
messages=messages, # pyright: ignore[reportArgumentType] # snapshot of the SDK's own message dicts
messages=[ # mutable-ok: provider transforms rewrite messages in place, so the router gets its own copy
dict(m) for m in messages
], # pyright: ignore[reportArgumentType] # snapshot of the SDK's own message dicts
metadata=shadow_metadata,
num_retries=0,
fallbacks=[], # mutable-ok: SDK kwarg; a failed shadow is a recorded error, never a spend multiplier
@ -772,7 +789,7 @@ class ShadowEvalLogger(CustomLogger):
)
except Exception as e: # noqa: BLE001 # provider errors become error rows, not crashes
verbose_logger.debug("shadow_eval: router call failed: %s", e)
return _CallFailure(f"shadow router call failed: {e}")
return _CallFailure(f"shadow router call failed: {_failure_detail(e)}")
text: Final = _chat_final_text(response)
if not text:
return _CallFailure("shadow router returned an empty response")
@ -815,6 +832,7 @@ class ShadowEvalLogger(CustomLogger):
judge_messages, # pyright: ignore[reportArgumentType] # plain SDK message dicts
temperature=0,
max_tokens=JUDGE_MAX_OUTPUT_TOKENS,
response_format=PAIRWISE_JUDGE_RESPONSE_FORMAT,
metadata=judge_metadata,
)
except Exception as e: # noqa: BLE001 # judge outages become error rows, not crashes

View file

@ -64,6 +64,10 @@ from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.core_helpers import reconstruct_model_name
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
cost_breakdown_with_guardrail,
guardrail_information_cost,
)
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
)
@ -108,6 +112,7 @@ from litellm.types.utils import (
LiteLLMBatch,
LiteLLMLoggingBaseClass,
LiteLLMRealtimeStreamLoggingObject,
ModelInfo,
ModelResponse,
ModelResponseStream,
RawRequestTypedDict,
@ -307,6 +312,66 @@ def _get_cached_prometheus_logger():
return _PrometheusLogger
_DEPLOYMENT_PRICING_KEYS: Final = (
"input_cost_per_token",
"output_cost_per_token",
"input_cost_per_token_batches",
"output_cost_per_token_batches",
)
def deployment_pricing_model_info(model_id: str | None, deployment_model: str | None) -> ModelInfo | None:
"""Pricing the router registered under this deployment's model_info.id.
Returns None when the deployment declares no pricing of its own, so the
caller falls back to the global cost map. The raw registration is what
decides that: the router registers an entry for every deployment, and
get_model_info fills absent costs with 0, so asking it directly cannot
tell "configured as free" apart from "no pricing configured". A deployment
may declare only one side of its pricing, so the side it leaves out keeps
the model's published rates instead of billing as zero. Ownership is per
token direction: declaring either rate for a direction takes that whole
direction, so a published batch rate can never displace a standard rate
the deployment configured itself.
"""
if model_id is None:
return None
registered: Final = litellm.model_cost.get(model_id)
if not isinstance(registered, dict) or not any(registered.get(key) is not None for key in _DEPLOYMENT_PRICING_KEYS):
return None
try:
merged: Final = litellm.get_model_info(model=model_id).copy()
except Exception: # noqa: BLE001 # get_model_info raises for ids it cannot resolve a provider for
return None
published: Final = _published_pricing(deployment_model)
if published is None:
return merged
declares_input: Final = (
registered.get("input_cost_per_token") is not None or registered.get("input_cost_per_token_batches") is not None
)
declares_output: Final = (
registered.get("output_cost_per_token") is not None
or registered.get("output_cost_per_token_batches") is not None
)
if not declares_input:
merged["input_cost_per_token"] = published.get("input_cost_per_token")
merged["input_cost_per_token_batches"] = published.get("input_cost_per_token_batches")
if not declares_output:
merged["output_cost_per_token"] = published.get("output_cost_per_token")
merged["output_cost_per_token_batches"] = published.get("output_cost_per_token_batches")
return merged
def _published_pricing(deployment_model: str | None) -> ModelInfo | None:
"""The cost map's own entry for the deployment's model, when it resolves."""
if deployment_model is None:
return None
try:
return litellm.get_model_info(model=deployment_model)
except Exception: # noqa: BLE001 # no published entry to layer the declared rates over
return None
class Logging(LiteLLMLoggingBaseClass):
global \
supabaseClient, \
@ -579,6 +644,28 @@ class Logging(LiteLLMLoggingBaseClass):
return model_id
return None
def get_deployment_model_for_cost(self) -> str | None:
"""The provider-qualified model to price against.
On a batch retrieve both self.model and litellm_params["model"] can be
unset, and self.model can otherwise carry the router's model_group alias,
which no cost map resolves. model_call_details holds the deployment's own
provider-qualified model, so it is preferred.
"""
candidates: Final = (
(self.model_call_details or {}).get("model") if hasattr(self, "model_call_details") else None,
self.litellm_params.get("model") if hasattr(self, "litellm_params") else None,
self.model,
)
return next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None)
def get_router_deployment_model_info(self) -> ModelInfo | None:
"""See deployment_pricing_model_info; None means fall back to the global cost map."""
return deployment_pricing_model_info(
model_id=self.get_router_model_id(),
deployment_model=self.get_deployment_model_for_cost(),
)
def update_environment_variables(
self,
litellm_params: dict,
@ -1007,10 +1094,10 @@ class Logging(LiteLLMLoggingBaseClass):
data=additional_args.get("complete_input_dict", {}),
)
_metadata["raw_request"] = str(curl_command)
_metadata["raw_request"] = _redact_string(str(curl_command))
# split up, so it's easier to parse in the UI
self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict(
raw_request_api_base=str(additional_args.get("api_base") or ""),
raw_request_api_base=self._get_masked_api_base(str(additional_args.get("api_base") or "")),
raw_request_body=self._get_raw_request_body(additional_args.get("complete_input_dict", {})),
# NOTE: setting ignore_sensitive_headers to True will cause
# the Authorization header to be leaked when calls to the health
@ -1024,8 +1111,10 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["raw_request_typed_dict"] = RawRequestTypedDict(
error=str(e),
)
_metadata["raw_request"] = f"Unable to Log \
_metadata["raw_request"] = _redact_string(
f"Unable to Log \
raw request: {e}"
)
if getattr(self, "logger_fn", None) and callable(self.logger_fn):
try:
self.logger_fn(
@ -1119,15 +1208,16 @@ class Logging(LiteLLMLoggingBaseClass):
if _is_debugging_on() or self.litellm_request_debug:
if json_logs:
masked_headers: Final = self._get_masked_headers(headers)
masked_api_base: Final = self._get_masked_api_base(str(api_base or ""))
if self.litellm_request_debug:
verbose_logger.warning( # .warning ensures this shows up in all environments
"POST Request Sent from LiteLLM",
extra={"api_base": {api_base}, **masked_headers},
extra={"api_base": {masked_api_base}, **masked_headers},
)
else:
verbose_logger.debug(
"POST Request Sent from LiteLLM",
extra={"api_base": {api_base}, **masked_headers},
extra={"api_base": {masked_api_base}, **masked_headers},
)
else:
headers = additional_args.get("headers", {})
@ -1167,8 +1257,6 @@ class Logging(LiteLLMLoggingBaseClass):
curl_command = "\nRequest Sent from LiteLLM:\n"
request_str: Final = additional_args.get("request_str", "")
curl_command += request_str
elif api_base == "":
curl_command = str(self.model_call_details)
return curl_command
def _get_masked_headers(self, headers: dict, ignore_sensitive_headers: bool = False) -> dict:
@ -2600,7 +2688,9 @@ class Logging(LiteLLMLoggingBaseClass):
) = await _handle_completed_batch(
batch=result,
custom_llm_provider=self.custom_llm_provider,
model_name=self.get_deployment_model_for_cost(),
litellm_params=self.litellm_params,
model_info=self.get_router_deployment_model_info(),
)
result._hidden_params["response_cost"] = response_cost
@ -5565,12 +5655,14 @@ def get_standard_logging_object_payload(
base_model = metadata.get("deployment")
custom_pricing: Final = use_custom_pricing_for_model(litellm_params=litellm_params)
raw_response_cost: Final = kwargs.get("response_cost")
response_cost: Final[float] = raw_response_cost or 0.0
llm_response_cost: Final[float] = raw_response_cost or 0.0
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))
response_cost: Final[float] = llm_response_cost + guardrail_cost
# clean up litellm hidden params
clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params)
if clean_hidden_params["response_cost"] is None and raw_response_cost is not None:
clean_hidden_params["response_cost"] = response_cost
clean_hidden_params["response_cost"] = llm_response_cost
model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information(
base_model=base_model,
@ -5650,7 +5742,7 @@ def get_standard_logging_object_payload(
metadata=clean_metadata,
cache_key=clean_hidden_params["cache_key"],
response_cost=response_cost,
cost_breakdown=logging_obj.cost_breakdown,
cost_breakdown=cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost),
total_tokens=usage_dict.get("total_tokens", 0),
prompt_tokens=usage_dict.get("prompt_tokens", 0),
completion_tokens=usage_dict.get("completion_tokens", 0),

View file

@ -0,0 +1,78 @@
import math
from collections.abc import Mapping
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_logger
from litellm.types.utils import CostBreakdown
BEDROCK_GUARDRAIL_PRICING_KEY: Final = "bedrock/guardrails"
class GuardrailPricing(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
guardrail_cost_per_unit: Mapping[str, float]
class GuardrailCostEntry(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
guardrail_cost: float | None = None
GuardrailInformationShape = tuple[GuardrailCostEntry, ...] | GuardrailCostEntry | None
_GUARDRAIL_INFORMATION_ADAPTER: Final[TypeAdapter[GuardrailInformationShape]] = TypeAdapter(GuardrailInformationShape)
def _bedrock_guardrail_pricing(aws_region_name: str | None) -> GuardrailPricing | None:
regional_key: Final = f"bedrock/{aws_region_name}/guardrails" if aws_region_name else None
for key in (regional_key, BEDROCK_GUARDRAIL_PRICING_KEY):
if key is None or key not in litellm.model_cost:
continue
try:
return GuardrailPricing.model_validate(litellm.model_cost[key])
except ValidationError as e:
verbose_logger.warning("Ignoring malformed guardrail pricing entry %s: %s", key, e)
return None
def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str | None) -> float:
pricing: Final = _bedrock_guardrail_pricing(aws_region_name)
if pricing is None:
return 0.0
return sum(units * pricing.guardrail_cost_per_unit.get(counter, 0.0) for counter, units in usage_units.items())
def _billable_entry_cost(entry: GuardrailCostEntry) -> float:
cost: Final = entry.guardrail_cost
if cost is None or not math.isfinite(cost) or cost <= 0.0:
return 0.0
return cost
def guardrail_information_cost(guardrail_information: object) -> float:
try:
parsed: Final = _GUARDRAIL_INFORMATION_ADAPTER.validate_python(guardrail_information)
except ValidationError:
return 0.0
if parsed is None:
return 0.0
if isinstance(parsed, GuardrailCostEntry):
return _billable_entry_cost(parsed)
return sum(_billable_entry_cost(entry) for entry in parsed)
def cost_breakdown_with_guardrail(cost_breakdown: CostBreakdown | None, guardrail_cost: float) -> CostBreakdown | None:
if guardrail_cost <= 0.0:
return cost_breakdown
existing: Final[CostBreakdown] = cost_breakdown if cost_breakdown is not None else CostBreakdown()
merged: Final[CostBreakdown] = {
**existing,
"guardrail_cost": guardrail_cost,
"total_cost": existing.get("total_cost", 0.0) + guardrail_cost,
}
return merged

View file

@ -42,14 +42,18 @@ _VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency)
# Pre-resolved service-tier cost-key suffixes (e.g. "_priority"). Used per
# request in the cost-calc path, so the f-strings are built once here instead
# of being rebuilt for every model_info key on every call.
_SERVICE_TIER_SUFFIXES: Final[tuple[str, ...]] = tuple(f"_{st.value}" for st in ServiceTier)
# of being rebuilt for every model_info key on every call. Longest-first so a
# substring match resolves "_ultrafast" before "_fast".
_SERVICE_TIER_SUFFIXES: Final[tuple[str, ...]] = tuple(
sorted((f"_{st.value}" for st in ServiceTier), key=len, reverse=True)
)
_SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType(
{
ServiceTier.FLEX.value: ServiceTier.FLEX.value,
ServiceTier.PRIORITY.value: ServiceTier.PRIORITY.value,
ServiceTier.FAST.value: ServiceTier.PRIORITY.value,
ServiceTier.ULTRAFAST.value: ServiceTier.ULTRAFAST.value,
}
)
@ -191,7 +195,7 @@ def _get_service_tier_cost_key(base_key: str, service_tier: str | None) -> str:
Args:
base_key: The base cost key (e.g., "input_cost_per_token")
service_tier: The service tier ("flex", "priority", "fast", or None for standard)
service_tier: The service tier ("flex", "priority", "fast", "ultrafast", or None for standard)
Returns:
str: The cost key to use (e.g., "input_cost_per_token_flex" or "input_cost_per_token")

View file

@ -1816,16 +1816,19 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]:
This helper uses ``json.JSONDecoder.raw_decode()`` to walk the string
and extract each JSON object individually.
The walk degrades gracefully: if the string is malformed or truncated
(e.g. a stream that ended mid-tool-call), whatever complete objects were
parsed before the bad tail are returned and the remainder is discarded
with a warning, rather than raising. The sole caller
(``_convert_to_bedrock_tool_call_invoke``) treats an empty result as
``input={}`` so the conversation can continue instead of hard-failing.
Returns
-------
list[dict]
A list of parsed dicts one per JSON object found. If *raw* is
empty or whitespace-only, an empty list is returned.
Raises
------
json.JSONDecodeError
If the string contains text that cannot be parsed as JSON at all.
empty, whitespace-only, or wholly unparseable, an empty list is
returned.
"""
import json
@ -1845,7 +1848,17 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]:
if idx >= length:
break
obj, end_idx = decoder.raw_decode(raw, idx)
try:
obj, end_idx = decoder.raw_decode(raw, idx)
except json.JSONDecodeError as e:
verbose_logger.warning(
"split_concatenated_json_objects: discarding unparseable tool-call "
"arguments tail after %d complete object(s); decode_start=%d error=%s",
len(results),
idx,
e,
)
break
if isinstance(obj, dict):
results.append(obj)
else:

View file

@ -3712,7 +3712,13 @@ def _convert_to_bedrock_tool_call_invoke(
_parts_list.append(cache_point_block)
return _parts_list
except Exception as e:
raise Exception(f"Unable to convert openai tool calls={tool_calls} to bedrock tool calls. Received error={e}")
tool_call_ids: Final = tuple(tool.get("id") for tool in tool_calls if isinstance(tool, dict))
raise litellm.BadRequestError(
message=f"Unable to convert openai tool calls with ids={tool_call_ids} to bedrock tool calls. "
f"Received error={e}",
model=model or "",
llm_provider="bedrock",
) from e
def _append_bedrock_tool_result_media_block(

View file

@ -258,6 +258,12 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
# For async objects, return a simple redacted response without deepcopy
return {"text": "redacted-by-litellm"}
if not (
isinstance(result, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse))
or (isinstance(result, dict) and ("choices" in result or "output" in result))
):
return {"text": "redacted-by-litellm"}
_result: Final = copy.deepcopy(result)
if isinstance(_result, litellm.ModelResponse):
if hasattr(_result, "choices") and _result.choices is not None:

View file

@ -1,4 +1,5 @@
import json
from collections.abc import Callable
from typing import Any, Final
from pydantic import BaseModel
@ -11,20 +12,32 @@ def strip_null_bytes(value: str) -> str:
return value.replace("\x00", "")
def safe_dumps(data: Any, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> str:
def safe_dumps(
data: Any,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
value_transform: Callable[[str | None, str], str] | None = None,
) -> str:
"""
Recursively serialize data while detecting circular references.
If a circular reference is detected then a marker string is returned.
NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors.
value_transform, when given, is applied to every string leaf (and to the
str() fallback for non-serializable objects) with the mapping key the leaf
was reached under, so callers can rewrite values without touching structure.
"""
def _serialize(obj: Any, seen: set, depth: int) -> Any:
def _transform(key: str | None, value: str) -> str:
return value if value_transform is None else value_transform(key, value)
def _serialize(obj: Any, seen: set, depth: int, key: str | None = None) -> Any:
# Check for maximum depth.
if depth > max_depth:
return "MaxDepthExceeded"
# Base-case: if it is a primitive, simply return it.
if isinstance(obj, str):
return obj.replace("\x00", "") if "\x00" in obj else obj
cleaned = obj.replace("\x00", "") if "\x00" in obj else obj
return _transform(key, cleaned)
if isinstance(obj, (int, float, bool, type(None))):
return obj
# Check for circular reference.
@ -37,30 +50,30 @@ def safe_dumps(data: Any, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> str:
for k, v in obj.items():
if isinstance(k, (str)):
clean_k = k.replace("\x00", "") if "\x00" in k else k
result[clean_k] = _serialize(v, seen, depth + 1)
result[clean_k] = _serialize(v, seen, depth + 1, clean_k)
seen.remove(id(obj))
return result
elif isinstance(obj, list):
result = [_serialize(item, seen, depth + 1) for item in obj]
result = [_serialize(item, seen, depth + 1, key) for item in obj]
seen.remove(id(obj))
return result
elif isinstance(obj, tuple):
result = tuple(_serialize(item, seen, depth + 1) for item in obj)
result = tuple(_serialize(item, seen, depth + 1, key) for item in obj)
seen.remove(id(obj))
return result
elif isinstance(obj, set):
result = sorted([_serialize(item, seen, depth + 1) for item in obj])
result = sorted([_serialize(item, seen, depth + 1, key) for item in obj])
seen.remove(id(obj))
return result
elif isinstance(obj, BaseModel):
dumped: Final = obj.model_dump()
result = _serialize(dumped, seen, depth + 1)
result = _serialize(dumped, seen, depth + 1, key)
seen.remove(id(obj))
return result
else:
# Fall back to string conversion for non-serializable objects.
try:
return strip_null_bytes(str(obj))
return _transform(key, strip_null_bytes(str(obj)))
except Exception:
return "Unserializable Object"

View file

@ -24,9 +24,6 @@ def _build_secret_patterns() -> "re.Pattern[str]":
r"(?:client_secret|azure_password|azure_username)\s+[^\s,'\"})\]{}>]+",
# AWS access key IDs
r"(?:AKIA|ASIA)[0-9A-Z]{16}",
# AWS secrets / session tokens / access key IDs (key=value)
r"(?:aws_secret_access_key|aws_session_token|aws_access_key_id)"
r"\s*[:=]\s*[A-Za-z0-9/+=]{20,}",
# Bearer tokens (OAuth, JWT, etc.)
r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*",
# Basic auth headers
@ -61,6 +58,7 @@ def _build_secret_patterns() -> "re.Pattern[str]":
# private_key with PEM-aware value capture
r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""",
r"(?:master_key|xai_key|database_url|db_url|connection_string|"
r"aws_secret_access_key|aws_session_token|aws_access_key_id|"
r"signing_key|encryption_key|"
r"auth_token|access_token|refresh_token|"
r"slack_webhook_url|webhook_url|"
@ -83,3 +81,19 @@ _SECRET_RE: Final = _build_secret_patterns()
def redact_string(value: str) -> str:
"""Scrub known secret/credential patterns from *value* and return the result."""
return _SECRET_RE.sub(_REDACTED, value)
def redact_structured_value(key: str | None, value: str) -> str:
"""Scrub *value* as it appeared under *key* inside a structured record.
redact_string() replaces a whole ``key: value`` span with REDACTED, which is
fine inside free text but destroys the surrounding syntax when the span is a
JSON member rather than message content. This renders the pair the way a dict
repr would, so the key-name patterns still fire, but collapses only the value
so the caller's structure survives.
"""
scrubbed: Final = redact_string(value)
if scrubbed != value or key is None:
return scrubbed
rendered: Final = f"'{key}': '{value}'"
return _REDACTED if redact_string(rendered) != rendered else value

View file

@ -1593,6 +1593,7 @@ class CustomStreamWrapper:
if self.stream_options is not None and self.stream_options["include_usage"] is True:
model_response.choices = []
return model_response
self._record_usage_only_chunk(model_response=model_response)
return
## CHECK FOR TOOL USE
@ -1819,6 +1820,16 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason = "tool_calls"
return model_response
def _record_usage_only_chunk(self, model_response: "ModelResponseStream") -> None:
"""
Keep provider usage-only chunks (e.g. OpenRouter's post-finish chunk, which carries a
provider-reported cost) available to cost tracking. They are never returned to the
caller; ``stream_options.include_usage`` only controls what the caller sees.
"""
if getattr(model_response, "usage", None) is None:
return
self.chunks.append(model_response.model_copy(update={"choices": []}))
@staticmethod
def _propagate_usage_cost_to_hidden_params(
response: "ModelResponse",

View file

@ -499,6 +499,39 @@ class AnthropicMessagesHandler(BaseTranslation):
{"role": "system", "content": blocks} if blocks else None # mutable-ok: API message payload
) # mutable-ok: API message payload
@staticmethod
def _fold_leading_systems_into_top_level(
data: dict[str, object], # mutable-ok: API message payload
leading_systems: Sequence[object],
include_existing_system: bool,
) -> None:
"""Deliver leading system rows through Anthropic's top-level system param, which rejects them in messages."""
existing: Final = data.get("system") if include_existing_system else None
existing_blocks: Final[list[object]] = ( # mutable-ok: API message payload
[{"type": "text", "text": existing}]
if isinstance(existing, str) and existing
else list(existing)
if isinstance(existing, list)
else []
)
converted_rows: Final = tuple(
AnthropicMessagesHandler._openai_system_message_to_anthropic(message)
for message in leading_systems
if isinstance(message, dict)
)
folded: Final[list[object]] = existing_blocks + [ # mutable-ok: API message payload
block
for row in converted_rows
if row is not None
for block in (
[{"type": "text", "text": row["content"]}] if isinstance(row["content"], str) else row["content"]
)
]
if folded:
data["system"] = folded # rebind-ok: write-back mutates the request payload in place
else:
data.pop("system", None)
@staticmethod
def _is_hoisted_top_level_system(message: object, hoisted_system_message: object) -> bool:
"""Match the hoisted prompt by identity, or by value after serialization."""
@ -575,9 +608,24 @@ class AnthropicMessagesHandler(BaseTranslation):
)
ordered: Final = AnthropicMessagesHandler._defer_systems_inside_tool_exchanges(structured_messages)
leading_count: Final = next(
(index for index, message in enumerate(ordered) if not _is_system(message)),
len(ordered),
)
leading_systems: Final = ordered[:leading_count]
hoisted_in_leading: Final = any(
AnthropicMessagesHandler._is_hoisted_top_level_system(message, hoisted_system_message)
for message in leading_systems
)
if leading_systems and not (leading_count == 1 and hoisted_in_leading):
AnthropicMessagesHandler._fold_leading_systems_into_top_level(
data,
leading_systems,
include_existing_system=hoisted_system_message is None,
)
run: Final[list] = [] # mutable-ok: API message payload
hoisted_dropped = False # rebind-ok: flips once the hoisted prompt is dropped
for message in ordered:
hoisted_dropped = hoisted_in_leading # rebind-ok: flips once the hoisted prompt is dropped
for message in ordered[leading_count:]:
if not _is_system(message):
run.append(message)
continue

View file

@ -624,6 +624,12 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
return self.chunk_queue.popleft()
if processed_chunk["type"] == "content_block_delta" and not self._delta_has_content(processed_chunk):
# A tool_use block opens with empty arguments (Bedrock Converse's
# ``contentBlockStart``, OpenAI's ``arguments: ""``), so flush the
# block start queued above instead of waiting for the next upstream
# chunk, which on a trailing-burst provider is the whole generation.
if self.chunk_queue:
return self.chunk_queue.popleft()
continue
if processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False:
@ -847,6 +853,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
if processed_chunk["type"] == "content_block_delta" and not self._delta_has_content(
processed_chunk
):
# See ``__next__``: flush the queued block start (issue #32004).
if self.chunk_queue:
return self.chunk_queue.popleft()
continue
if processed_chunk["type"] == "message_delta" and self.sent_content_block_finish is False:

View file

@ -734,6 +734,7 @@ class LiteLLMAnthropicMessagesAdapter:
"input_schema",
"description",
"cache_control",
"strict",
"type",
]
@ -763,6 +764,8 @@ class LiteLLMAnthropicMessagesAdapter:
function_chunk["parameters"] = tool["input_schema"]
if "description" in tool:
function_chunk["description"] = tool["description"]
if "strict" in tool:
function_chunk["strict"] = bool(tool["strict"])
for k, v in tool.items():
if k not in mapped_tool_params: # pass additional computer kwargs

View file

@ -16,7 +16,7 @@ How it works:
import uuid
from collections.abc import AsyncIterator
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
import litellm
import litellm.constants as _c
@ -28,6 +28,9 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
if TYPE_CHECKING:
from litellm.router import Router
ADVISOR_MAX_USES: Final[int] = _c.ADVISOR_MAX_USES
ADVISOR_NATIVE_PROVIDERS: Final[frozenset] = _c.ADVISOR_NATIVE_PROVIDERS
ADVISOR_TOOL_DESCRIPTION: Final[str] = _c.ADVISOR_TOOL_DESCRIPTION
@ -97,6 +100,14 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
parent_request_id: Final[str] = str(kwargs.pop("litellm_call_id", None) or uuid.uuid4())
metadata_base: Final[dict] = dict(kwargs.pop("metadata", None) or {})
advisor_metadata: Final = {
**metadata_base,
"advisor_sub_call": True,
"parent_request_id": parent_request_id,
}
advisor_router: Final = (
None if (advisor_api_key or advisor_api_base) else _resolve_advisor_router(advisor_model)
)
iteration = 0
while True:
@ -138,20 +149,27 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
# --- Advisor sub-call (always non-streaming, no tools) ---
try:
advisor_response: AnthropicMessagesResponse = await _call_messages_handler(
model=advisor_model,
messages=advisor_messages,
tools=None,
stream=False,
max_tokens=max_tokens,
custom_llm_provider=None, # let litellm resolve from model name
metadata={
**metadata_base,
"advisor_sub_call": True,
"parent_request_id": parent_request_id,
},
api_key=advisor_api_key,
api_base=advisor_api_base,
advisor_response: AnthropicMessagesResponse = (
await advisor_router.aanthropic_messages(
model=advisor_model,
messages=advisor_messages,
tools=None,
stream=False,
max_tokens=max_tokens,
metadata=advisor_metadata,
)
if advisor_router is not None
else await _call_messages_handler(
model=advisor_model,
messages=advisor_messages,
tools=None,
stream=False,
max_tokens=max_tokens,
custom_llm_provider=None,
metadata=advisor_metadata,
api_key=advisor_api_key,
api_base=advisor_api_base,
)
)
except Exception as advisor_sub_call_exception:
mark_advisor_orchestration_failure(advisor_sub_call_exception)
@ -284,6 +302,11 @@ def _build_advisor_context(
tool_use blocks are excluded because Anthropic requires tool_use to be
immediately followed by tool_result not the advisor question.
In-sequence system rows (e.g. Claude Code SessionStart hook output) are
excluded: they are executor-directed, and a trailing one becomes invalid
once the question turn is appended after it (a system row must precede an
assistant message or end the array).
"""
question: Final = (advisor_use_block.get("input") or {}).get("question") or (
"Please provide guidance on the current task."
@ -295,7 +318,7 @@ def _build_advisor_context(
for block in raw_content
if isinstance(block, dict) and block.get("type") == "text"
]
result: Final = list(messages)
result: Final = [m for m in messages if m.get("role") != "system"]
if executor_text_blocks:
result.append({"role": "assistant", "content": executor_text_blocks})
result.append({"role": "user", "content": question})
@ -357,6 +380,24 @@ def _inject_max_uses_error(
]
def _resolve_advisor_router(advisor_model: str) -> "Router | None":
"""Return the proxy router when it serves ``advisor_model`` directly or via a wildcard.
Returns ``None`` for SDK callers (no proxy router) and for advisor models the router
doesn't know about, so those keep resolving through ``litellm.anthropic_messages()``
provider inference.
"""
try:
from litellm.proxy.proxy_server import llm_router
except (ImportError, ModuleNotFoundError):
return None
if llm_router is None:
return None
if llm_router.is_recognized_model(advisor_model) or llm_router.pattern_router.route(advisor_model):
return llm_router
return None
async def _call_messages_handler(
model: str,
messages: list[dict],

View file

@ -3,7 +3,7 @@
import json
import traceback
from collections import deque
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final
from litellm import verbose_logger
@ -68,6 +68,19 @@ class AnthropicResponsesStreamWrapper:
self._current_block_index += 1
return self._current_block_index
def _open_block(self, item_id: str | None, content_block: Mapping[str, Any]) -> int:
block_idx = self._next_block_index()
if item_id:
self._item_id_to_block_index[item_id] = block_idx
self._chunk_queue.append(
{
"type": "content_block_start",
"index": block_idx,
"content_block": content_block,
}
)
return block_idx
def _process_event(self, event: Any) -> None:
"""Convert one Responses API event into zero or more Anthropic chunks queued for emission."""
event_type = getattr(event, "type", None)
@ -93,47 +106,22 @@ class AnthropicResponsesStreamWrapper:
item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None)
if item_type == "message":
block_idx = self._next_block_index()
if item_id:
self._item_id_to_block_index[item_id] = block_idx
self._chunk_queue.append(
{
"type": "content_block_start",
"index": block_idx,
"content_block": {"type": "text", "text": ""},
}
)
self._open_block(item_id, {"type": "text", "text": ""})
elif item_type == "function_call":
call_id: Final = (
getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or ""
)
name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or ""
block_idx = self._next_block_index()
if item_id:
self._item_id_to_block_index[item_id] = block_idx
self._pending_tool_ids[item_id] = call_id
self._chunk_queue.append(
self._open_block(
item_id,
{
"type": "content_block_start",
"index": block_idx,
"content_block": {
"type": "tool_use",
"id": call_id,
"name": name,
"input": {},
},
}
)
elif item_type == "reasoning":
block_idx = self._next_block_index()
if item_id:
self._item_id_to_block_index[item_id] = block_idx
self._chunk_queue.append(
{
"type": "content_block_start",
"index": block_idx,
"content_block": {"type": "thinking", "thinking": ""},
}
"type": "tool_use",
"id": call_id,
"name": name,
"input": {},
},
)
return
@ -146,16 +134,7 @@ class AnthropicResponsesStreamWrapper:
# Some providers (e.g. LMStudio) skip response.output_item.added,
# so no text block is open yet; synthesize content_block_start
# instead of emitting a delta with index -1
block_idx = self._next_block_index()
if item_id:
self._item_id_to_block_index[item_id] = block_idx
self._chunk_queue.append(
{
"type": "content_block_start",
"index": block_idx,
"content_block": {"type": "text", "text": ""},
}
)
block_idx = self._open_block(item_id, {"type": "text", "text": ""})
self._chunk_queue.append(
{
"type": "content_block_delta",
@ -169,11 +148,11 @@ class AnthropicResponsesStreamWrapper:
if event_type == "response.reasoning_summary_text.delta":
item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None)
delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "")
block_idx = (
self._item_id_to_block_index.get(item_id, self._current_block_index)
if item_id
else self._current_block_index
)
block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index
if block_idx < 0:
if not delta:
return
block_idx = self._open_block(item_id, {"type": "thinking", "thinking": ""})
self._chunk_queue.append(
{
"type": "content_block_delta",
@ -207,11 +186,9 @@ class AnthropicResponsesStreamWrapper:
item_id = (
getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None
)
block_idx = (
self._item_id_to_block_index.get(item_id, self._current_block_index)
if item_id
else self._current_block_index
)
block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index
if block_idx < 0:
return
self._chunk_queue.append(
{
"type": "content_block_stop",

View file

@ -266,7 +266,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search":
result.append({"type": "web_search_preview"})
continue
func_tool: dict[str, Any] = {"type": "function", "name": tool_name}
# Responses turns strict mode on when `strict` is omitted, silently rewriting
# `required` to every property. Anthropic tools are non-strict unless asked.
func_tool: dict[str, Any] = {
"type": "function",
"name": tool_name,
"strict": bool(tool_dict.get("strict")),
}
if "description" in tool_dict:
func_tool["description"] = tool_dict["description"]
if "input_schema" in tool_dict:

View file

@ -1,3 +1,4 @@
import copy
import enum
import re
from typing import Any, Final, cast
@ -11,6 +12,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
_audio_or_image_in_message_content,
convert_content_list_to_str,
filter_value_from_dict,
)
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
@ -28,6 +30,9 @@ class AzureFoundryErrorStrings(str, enum.Enum):
SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'"
NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ("thinking_blocks", "provider_specific_fields", "cache_control")
class AzureAIStudioConfig(OpenAIConfig):
def get_supported_openai_params(self, model: str) -> list:
model_supports_tool_choice = True # azure ai supports this by default
@ -167,10 +172,23 @@ class AzureAIStudioConfig(OpenAIConfig):
) -> list:
"""
- Azure AI Studio doesn't support content as a list. This handles:
1. Transforms list content to a string.
2. If message contains an image or audio, send as is (user-intended)
1. Strips message fields that are not part of the OpenAI chat-completions
schema (thinking_blocks, provider_specific_fields, cache_control).
Azure AI Foundry backends set additionalProperties=false and reject
these with "Extra inputs are not permitted", which breaks multi-turn
Anthropic-format clients that echo thinking blocks back as history.
2. Transforms list content to a string.
3. If message contains an image or audio, send as is (user-intended)
Operates on a deep copy so the caller's messages keep their thinking blocks
and provider metadata, which a fallback to another provider still needs.
"""
for message in messages:
stripped_messages: Final = copy.deepcopy(messages)
for message in stripped_messages:
message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy
for field in NON_OPENAI_SPEC_MESSAGE_FIELDS:
filter_value_from_dict(message_dict, field)
# Do nothing if the message contains an image or audio
if _audio_or_image_in_message_content(message):
continue
@ -178,7 +196,7 @@ class AzureAIStudioConfig(OpenAIConfig):
texts = convert_content_list_to_str(message=message)
if texts:
message["content"] = texts
return messages
return stripped_messages
def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool:
try:

View file

@ -11,6 +11,7 @@ The operation location must be polled until the analysis completes.
import asyncio
import re
import time
from collections.abc import Mapping
from typing import Any, Final
from urllib.parse import quote
@ -23,15 +24,19 @@ from litellm.constants import (
AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI,
AZURE_OPERATION_POLLING_TIMEOUT,
)
from litellm.exceptions import UnsupportedParamsError
from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin, encode_url_path_segment
from litellm.llms.base_llm.ocr.transformation import (
OCR_REQUEST_FORMAT_PARAM,
BaseOCRConfig,
DocumentType,
OCRPage,
OCRPageDimensions,
OCRRequestData,
OCRRequestFormat,
OCRResponse,
OCRUsageInfo,
parse_ocr_request_format,
)
from litellm.secret_managers.main import get_secret_str
@ -97,8 +102,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
comma-separated string. Other Mistral-specific params (e.g.
`include_image_base64`) are not supported by Azure DI and are
ignored during transformation.
`req_format` selects the response shape: "litellm" (default) returns
the normalized OCR schema, "native" returns Azure DI's own analyze
operation payload as-is.
"""
return ["pages", "features"]
return ["pages", "features", OCR_REQUEST_FORMAT_PARAM]
def map_ocr_params(
self,
@ -117,14 +126,27 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
"""
pages: Final = non_default_params.get("pages")
features: Final = non_default_params.get("features")
request_format: Final = non_default_params.get(OCR_REQUEST_FORMAT_PARAM)
normalized_pages: Final = self._normalize_pages_param(pages) if pages is not None else ""
normalized_features: Final = self._normalize_features_param(features) if features is not None else ""
return {
**optional_params,
**({"pages": normalized_pages} if normalized_pages else {}),
**({"features": normalized_features} if normalized_features else {}),
**(
{OCR_REQUEST_FORMAT_PARAM: self._parse_request_format(request_format, model)}
if request_format is not None
else {}
),
}
@staticmethod
def _parse_request_format(request_format: object, model: str) -> OCRRequestFormat:
try:
return parse_ocr_request_format(request_format)
except ValueError as e:
raise UnsupportedParamsError(message=f"{e}", model=model, llm_provider="azure_ai") from e
@staticmethod
def _normalize_pages_param(pages: Any) -> str:
"""
@ -594,14 +616,33 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
poll_headers = {"Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "")}
return operation_url, poll_headers
def _transform_completed_response(self, model: str, raw_response: httpx.Response) -> OCRResponse:
@staticmethod
def _get_request_format(optional_params: object) -> OCRRequestFormat:
if not isinstance(optional_params, dict):
return "litellm"
request_format: Final = optional_params.get(OCR_REQUEST_FORMAT_PARAM)
if request_format is None:
return "litellm"
return parse_ocr_request_format(request_format)
def _transform_completed_response(
self,
model: str,
raw_response: httpx.Response,
request_format: OCRRequestFormat,
) -> OCRResponse:
"""
Transform a completed Azure Document Intelligence analyze operation
into the Mistral OCR response shape, preserving Azure-native
`analyzeResult` fields (`content`, `tables`, `keyValuePairs`) as
top-level response fields.
When `request_format` is "native", the untouched Azure operation
payload is attached to the response's hidden params so the proxy can
return it verbatim while cost tracking still reads `usage_info`.
"""
operation: Final = AzureDocumentIntelligenceOperation.model_validate(raw_response.json())
raw_operation: Final[Mapping[str, object]] = raw_response.json()
operation: Final = AzureDocumentIntelligenceOperation.model_validate(raw_operation)
verbose_logger.debug("Azure Document Intelligence response status: %s", operation.status)
@ -614,7 +655,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
mistral_pages: Final = [self._transform_azure_page(azure_page) for azure_page in analyze_result.pages]
usage_info: Final = OCRUsageInfo(pages_processed=len(mistral_pages), doc_size_bytes=None)
return OCRResponse(
response: Final = OCRResponse(
pages=mistral_pages,
model=model,
usage_info=usage_info,
@ -624,6 +665,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
keyValuePairs=analyze_result.keyValuePairs,
)
if request_format == "native":
response.set_provider_native_response(raw_operation)
return response
def transform_ocr_response(
self,
model: str,
@ -681,8 +727,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
Returns:
OCRResponse in Mistral format
"""
request_format: Final = self._get_request_format(kwargs.get("optional_params"))
if raw_response.status_code != 202:
return self._transform_completed_response(model=model, raw_response=raw_response)
return self._transform_completed_response(
model=model, raw_response=raw_response, request_format=request_format
)
verbose_logger.debug("Azure DI returned 202 Accepted, polling operation...")
operation_url, poll_headers = self._get_polling_target(raw_response)
@ -691,7 +741,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
headers=poll_headers,
timeout_secs=AZURE_OPERATION_POLLING_TIMEOUT,
)
return self._transform_completed_response(model=model, raw_response=completed_response)
return self._transform_completed_response(
model=model, raw_response=completed_response, request_format=request_format
)
async def async_transform_ocr_response(
self,
@ -714,8 +766,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
Returns:
OCRResponse in Mistral format
"""
request_format: Final = self._get_request_format(kwargs.get("optional_params"))
if raw_response.status_code != 202:
return self._transform_completed_response(model=model, raw_response=raw_response)
return self._transform_completed_response(
model=model, raw_response=raw_response, request_format=request_format
)
verbose_logger.debug("Azure DI returned 202 Accepted, polling operation (async)...")
operation_url, poll_headers = self._get_polling_target(raw_response)
@ -724,4 +780,6 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
headers=poll_headers,
timeout_secs=AZURE_OPERATION_POLLING_TIMEOUT,
)
return self._transform_completed_response(model=model, raw_response=completed_response)
return self._transform_completed_response(
model=model, raw_response=completed_response, request_format=request_format
)

View file

@ -2,7 +2,8 @@
Base OCR transformation configuration.
"""
from typing import TYPE_CHECKING, Any
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from pydantic import PrivateAttr
@ -21,6 +22,26 @@ else:
# File-type inputs are preprocessed to this format in litellm/ocr/main.py.
DocumentType = dict[str, str]
OCRRequestFormat = Literal["litellm", "native"]
OCR_REQUEST_FORMATS: Final[tuple[OCRRequestFormat, ...]] = ("litellm", "native")
OCR_REQUEST_FORMAT_PARAM: Final = "req_format"
OCR_REQUEST_FORMAT_HEADER: Final = "x-req-format"
PROVIDER_NATIVE_RESPONSE_KEY: Final = "provider_native_response"
def parse_ocr_request_format(value: object) -> OCRRequestFormat:
if value == "litellm":
return "litellm"
if value == "native":
return "native"
raise ValueError(
f"Invalid `{OCR_REQUEST_FORMAT_PARAM}`: {value!r}. Expected one of {', '.join(OCR_REQUEST_FORMATS)}."
)
class OCRPageDimensions(LiteLLMPydanticObjectBase):
"""Page dimensions from OCR response."""
@ -80,6 +101,15 @@ class OCRResponse(LiteLLMPydanticObjectBase):
# Define private attributes using PrivateAttr
_hidden_params: dict = PrivateAttr(default_factory=dict)
def set_provider_native_response(self, native_response: Mapping[str, object]) -> None:
"""Keep the provider's own response payload alongside the normalized one."""
self._hidden_params[PROVIDER_NATIVE_RESPONSE_KEY] = native_response
def get_provider_native_response(self) -> Mapping[str, object] | None:
"""The provider's own response payload, when `req_format=native` was requested."""
native_response: Final = self._hidden_params.get(PROVIDER_NATIVE_RESPONSE_KEY)
return native_response if isinstance(native_response, dict) else None
class OCRRequestData(LiteLLMPydanticObjectBase):
"""OCR request data structure."""

View file

@ -1,5 +1,6 @@
from abc import abstractmethod
from typing import TYPE_CHECKING, Any
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, NoReturn
import httpx
@ -154,3 +155,75 @@ class BaseVectorStoreConfig:
response: VectorStoreSearchResponse,
) -> tuple[float, float]:
return 0.0, 0.0
class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
"""
Base config for vector store providers whose datastore has no HTTP API
(e.g. Valkey over RESP). Instead of transforming to an httpx request, the
config executes the search itself via (a)execute_search_vector_store_request.
"""
@abstractmethod
def execute_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
pass
@abstractmethod
async def aexecute_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
pass
def transform_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: Mapping[str, object],
extra_body: Mapping[str, object] | None = None,
) -> NoReturn:
raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP request shape")
def transform_search_vector_store_response(
self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj
) -> NoReturn:
raise NotImplementedError("Direct vector store providers execute the search themselves; no HTTP response shape")
def transform_create_vector_store_request(
self,
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
api_base: str,
) -> NoReturn:
raise NotImplementedError
def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn:
raise NotImplementedError
def get_complete_url(
self,
api_base: str | None,
litellm_params: Mapping[str, object],
) -> str:
return api_base or ""
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
return BaseVectorStoreAuthCredentials()
def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
return VectorStoreIndexEndpoints(read=[], write=[]) # mutable-ok: the TypedDict declares list fields

View file

@ -6,6 +6,7 @@ import copy
import json
import time
import types
from collections.abc import Mapping
from typing import Final, Literal, cast, overload
import httpx
@ -1790,7 +1791,43 @@ class AmazonConverseConfig(BaseConfig):
thinking_blocks_list.append(_redacted_block)
return thinking_blocks_list
def _transform_usage(
@staticmethod
def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool:
"""Converse-family models report camelCase token counts, not Anthropic's snake_case."""
return "inputTokens" in usage_object or "outputTokens" in usage_object
@staticmethod
def _usage_count(usage_object: Mapping[str, object], *keys: str) -> int:
for key in keys:
value = usage_object.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool):
return int(value)
return 0
def usage_from_batch_output(self, usage_object: Mapping[str, object]) -> Usage:
"""Read a Converse-shaped usage block out of a batch output line.
Batch output omits fields the live API always sends, so the block is
completed before going through the same transform, keeping a batch and an
equivalent non-batch call in agreement on tokens.
"""
input_tokens: Final = self._usage_count(usage_object, "inputTokens")
output_tokens: Final = self._usage_count(usage_object, "outputTokens")
cache_read: Final = self._usage_count(usage_object, "cacheReadInputTokens", "cacheReadInputTokenCount")
cache_write: Final = self._usage_count(usage_object, "cacheWriteInputTokens", "cacheWriteInputTokenCount")
return self.transform_usage(
ConverseTokenUsageBlock(
inputTokens=input_tokens,
outputTokens=output_tokens,
totalTokens=self._usage_count(usage_object, "totalTokens") or input_tokens + output_tokens,
cacheReadInputTokenCount=cache_read,
cacheReadInputTokens=cache_read,
cacheWriteInputTokenCount=cache_write,
cacheWriteInputTokens=cache_write,
)
)
def transform_usage(
self,
usage: ConverseTokenUsageBlock,
reasoning_content: str | None = None,
@ -2211,7 +2248,7 @@ class AmazonConverseConfig(BaseConfig):
chat_completion_message["tool_calls"] = filtered_tools
## CALCULATING USAGE - bedrock returns usage in the headers
usage: Final = self._transform_usage(
usage: Final = self.transform_usage(
completion_response["usage"],
reasoning_content=chat_completion_message.get("reasoning_content"),
)

View file

@ -559,7 +559,7 @@ class AWSEventStreamDecoder:
elif "stopReason" in chunk_data:
finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
elif "usage" in chunk_data:
usage = converse_config._transform_usage(chunk_data.get("usage", {}))
usage = converse_config.transform_usage(chunk_data.get("usage", {}))
model_response_provider_specific_fields: Final = {}
if "trace" in chunk_data:

View file

@ -3,6 +3,7 @@ import json
import os
import time
from collections.abc import Iterable, Mapping, MutableMapping, Sequence
from contextlib import suppress
from functools import cache
from itertools import chain
from types import MappingProxyType
@ -149,11 +150,12 @@ class _BedrockS3RequestParams(BaseModel):
class _TrustedS3ModelCredentials(BaseModel):
"""The S3 bucket the server trusts file ids against, from the deployment snapshot."""
"""The S3 buckets the server trusts file ids against, from the deployment snapshot."""
model_config = ConfigDict(extra="ignore")
s3_bucket_name: str | None = None
s3_output_bucket_name: str | None = None
def extract_s3_uri_from_file_id(file_id: str) -> str:
@ -179,6 +181,18 @@ def extract_s3_uri_from_file_id(file_id: str) -> str:
raise ValueError("file_id must be a managed LiteLLM S3 file id")
_S3_BUCKET_REQUIRED_ERROR: Final = "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval."
def _trusted_s3_model_credentials(litellm_params: Mapping[str, object]) -> _TrustedS3ModelCredentials:
trusted_model_credentials: Final = litellm_params.get("_litellm_internal_model_credentials")
if not isinstance(trusted_model_credentials, MappingProxyType):
return _TrustedS3ModelCredentials()
snapshot: Final[dict[str, object]] = {}
snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot
return _TrustedS3ModelCredentials.model_validate(snapshot)
def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str:
"""
Resolve the server-configured S3 bucket for Bedrock file operations.
@ -187,20 +201,50 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str:
environment; never a request-supplied param, since the bucket is what
`validate_managed_cloud_file_id` checks file ids against.
"""
trusted_model_credentials: Final = litellm_params.get("_litellm_internal_model_credentials")
bucket_name: str | None = None
if isinstance(trusted_model_credentials, MappingProxyType):
snapshot: Final[dict[str, object]] = {}
snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot
bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_bucket_name
bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
bucket_name: Final = _trusted_s3_model_credentials(litellm_params).s3_bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
if not bucket_name:
raise ValueError(
"S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval."
)
raise ValueError(_S3_BUCKET_REQUIRED_ERROR)
return bucket_name
def get_configured_s3_bucket_names(litellm_params: Mapping[str, object]) -> tuple[str, ...]:
"""
Resolve the server-configured S3 buckets a Bedrock file id may live in.
Bedrock batch outputs land in ``s3_output_bucket_name`` when it differs from
the input bucket, so retrieval validates against both. Same trust rules as
``get_configured_s3_bucket_name``: only the immutable credential snapshot or
the environment, never a request param.
"""
trusted: Final = _trusted_s3_model_credentials(litellm_params)
input_bucket: Final = trusted.s3_bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
output_bucket: Final = trusted.s3_output_bucket_name or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME")
buckets: Final = tuple(dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket))
if not buckets:
raise ValueError(_S3_BUCKET_REQUIRED_ERROR)
return buckets
def _validate_file_id_against_configured_buckets(
s3_uri: str,
configured_bucket_names: tuple[str, ...],
allow_legacy_cloud_file_ids: bool,
) -> tuple[str, str]:
def validate_against(configured_bucket_name: str) -> tuple[str, str]:
return validate_managed_cloud_file_id(
file_id=s3_uri,
scheme="s3://",
configured_bucket_name=configured_bucket_name,
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids,
)
for candidate_bucket_name in configured_bucket_names[:-1]:
with suppress(ValueError):
return validate_against(candidate_bucket_name)
return validate_against(configured_bucket_names[-1])
def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int:
"""
S3 answers PutObject with an empty body, so the stored object size comes from the
@ -1186,11 +1230,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
raise ValueError("file_id is required for Bedrock file content retrieval")
s3_uri: Final = extract_s3_uri_from_file_id(file_id)
bucket_name, object_key = validate_managed_cloud_file_id(
file_id=s3_uri,
scheme="s3://",
configured_bucket_name=get_configured_s3_bucket_name(litellm_params),
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
bucket_name, object_key = _validate_file_id_against_configured_buckets(
s3_uri=s3_uri,
configured_bucket_names=get_configured_s3_bucket_names(litellm_params),
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
)

View file

@ -2,7 +2,7 @@ import asyncio
import json
import os
import ssl
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from functools import lru_cache
from types import ModuleType
@ -56,7 +56,10 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi
from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse
from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig
from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.base_llm.vector_store.transformation import (
BaseDirectVectorStoreConfig,
BaseVectorStoreConfig,
)
from litellm.llms.base_llm.vector_store_files.transformation import (
BaseVectorStoreFilesConfig,
)
@ -70,6 +73,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
MockResponsesAPIStreamingIterator,
ProjectQuotaCallback,
ResponsesAPIStreamingIterator,
ResponsesWebSocketStreaming,
SyncResponsesAPIStreamingIterator,
@ -253,6 +257,27 @@ def _has_pre_call_deployment_hook(logging_obj: LiteLLMLoggingObj) -> bool:
return False
def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
"""Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM
enforcement, so the Responses WebSocket loop can charge every
``response.create`` frame, not just the connection's first one.
Uses duck-typing on ``litellm.callbacks`` (rather than importing the
proxy hook directly) to avoid a layering violation (SDK importing from
the proxy layer).
"""
import litellm as _litellm
callbacks: Final = cast( # cast-ok: callback registry is inspected before protocol use
Sequence[object], _litellm.callbacks
)
return tuple(
cast(ProjectQuotaCallback, callback) # cast-ok: required callback method is callable
for callback in callbacks
if callable(getattr(callback, "enforce_project_io_token_quota_for_frame", None))
)
class BaseLLMHTTPHandler:
async def _make_common_async_call(
self,
@ -893,6 +918,7 @@ class BaseLLMHTTPHandler:
)
if provider_config is None:
raise ValueError(f"Provider {custom_llm_provider} does not support embedding")
embedding_extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None)
# get config from model, custom llm provider
headers = provider_config.validate_environment(
api_key=api_key,
@ -917,6 +943,8 @@ class BaseLLMHTTPHandler:
optional_params=optional_params,
headers=headers,
)
if embedding_extra_body:
data.update(embedding_extra_body)
# Some providers (e.g. OCI) require request signing after the body is built.
# The default BaseConfig.sign_request returns (headers, None) — a no-op for
@ -1556,12 +1584,14 @@ class BaseLLMHTTPHandler:
model: str,
response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
optional_params: Mapping[str, object],
) -> OCRResponse:
"""Shared logic for transforming OCR responses."""
return provider_config.transform_ocr_response(
model=model,
raw_response=response,
logging_obj=logging_obj,
optional_params=optional_params,
)
def ocr(
@ -1637,6 +1667,7 @@ class BaseLLMHTTPHandler:
model=model,
response=response,
logging_obj=logging_obj,
optional_params=optional_params,
)
async def async_ocr(
@ -1699,6 +1730,7 @@ class BaseLLMHTTPHandler:
model=model,
raw_response=response,
logging_obj=logging_obj,
optional_params=optional_params,
)
def search(
@ -6169,6 +6201,8 @@ class BaseLLMHTTPHandler:
- Uses ManagedResponsesWebSocketHandler which makes HTTP streaming calls
- Forwards events over the websocket connection
"""
_ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks()
if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket():
from litellm.responses.streaming_iterator import (
ManagedResponsesWebSocketHandler,
@ -6185,6 +6219,7 @@ class BaseLLMHTTPHandler:
timeout=timeout,
custom_llm_provider=custom_llm_provider,
first_message=first_message,
quota_callbacks=_ws_quota_callbacks,
**kwargs,
)
await handler.run()
@ -6305,6 +6340,7 @@ class BaseLLMHTTPHandler:
first_message=first_message,
guardrail_callbacks=_ws_guardrail_callbacks,
output_guardrail_callbacks=_ws_output_guardrail_callbacks,
quota_callbacks=_ws_quota_callbacks,
authorized_model=model,
)
await streaming.bidirectional_forward()
@ -9397,6 +9433,27 @@ class BaseLLMHTTPHandler:
)
###### VECTOR STORE HANDLER ######
@staticmethod
def _pre_call_direct_vector_store_search(
logging_obj: LiteLLMLoggingObj,
custom_llm_provider: str,
vector_store_id: str,
query: str | Sequence[str],
) -> None:
"""Direct providers have no HTTP request to echo, and an empty api_base makes the debug
logger fall back to dumping model_call_details, which holds stored provider credentials."""
endpoint: Final = f"{custom_llm_provider}://{vector_store_id}"
logging_obj.pre_call(
input="",
api_key="",
additional_args={ # mutable-ok: pre_call's additional_args contract is a dict
"query": query,
"vector_store_id": vector_store_id,
"api_base": endpoint,
"request_str": f"direct vector store search: {endpoint}",
},
)
async def async_vector_store_search_handler(
self,
vector_store_id: str,
@ -9412,6 +9469,22 @@ class BaseLLMHTTPHandler:
client: HTTPHandler | AsyncHTTPHandler | None = None,
_is_async: bool = False,
) -> VectorStoreSearchResponse:
if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig):
self._pre_call_direct_vector_store_search(
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vector_store_id=vector_store_id,
query=query,
)
return await vector_store_provider_config.aexecute_search_vector_store_request(
vector_store_id=vector_store_id,
query=query,
vector_store_search_optional_params=vector_store_search_optional_params,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape
timeout=timeout,
)
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider),
@ -9525,6 +9598,22 @@ class BaseLLMHTTPHandler:
client=client,
)
if isinstance(vector_store_provider_config, BaseDirectVectorStoreConfig):
self._pre_call_direct_vector_store_search(
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vector_store_id=vector_store_id,
query=query,
)
return vector_store_provider_config.execute_search_vector_store_request(
vector_store_id=vector_store_id,
query=query,
vector_store_search_optional_params=vector_store_search_optional_params,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape
timeout=timeout,
)
if client is None or not isinstance(client, HTTPHandler):
sync_httpx_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)})
else:

View file

@ -29,9 +29,12 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
return None
AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-"
def resolve_fireworks_resource_name(model: str) -> str:
stripped: Final = model.removeprefix("fireworks_ai/")
if stripped.startswith("accounts/") or "#" in stripped:
if stripped.startswith(("accounts/", AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX)) or "#" in stripped:
return stripped
if stripped.startswith(("routers/", "models/")):
return f"accounts/fireworks/{stripped}"

View file

@ -14,6 +14,7 @@ import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.search.transformation import (
@ -22,7 +23,7 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.secret_managers.main import get_secret_str
_UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | bool])
_UrlEncodableParams: Final = TypeAdapter(dict[str, str | int | float | bool])
_StrList: Final = TypeAdapter(list[str])
_StrFrozenSet: Final = TypeAdapter(frozenset[str])
@ -94,16 +95,16 @@ class TinyfishSearchConfig(BaseSearchConfig):
TinyFish equivalents:
- ``query`` (str or list[str]) ``query`` (list joined by spaces)
- ``country`` ``location``
- ``search_domain_filter`` (list[str]) folded into the query as
``(<query>) (site:a OR site:b ...)`` (TinyFish has no first-class
field today; see ML-2084 for the planned ``include_domains``)
- ``search_domain_filter`` (list[str]) folded into the query using
search operators
- ``max_results`` not sent on the wire; stashed on
``self._caller_max_results`` for client-side response truncation
(TinyFish doesn't honor it server-side)
- ``max_tokens_per_page`` silently dropped (no TinyFish equivalent)
Any other ``optional_params`` keys are forwarded to TinyFish as-is.
dict/list values are JSON-encoded so they survive ``urlencode``.
dict and list values are JSON-encoded so structured payloads survive
``urlencode``.
Returns:
``{_TINYFISH_PARAMS_KEY: <dict of querystring entries>}``.
@ -144,14 +145,12 @@ class TinyfishSearchConfig(BaseSearchConfig):
supported_perplexity: Final = _StrFrozenSet.validate_python(raw_supported)
for param, value in optional_params.items():
if param not in supported_perplexity and param not in request_data:
# `fetch` expects a JSON-encoded object on the wire; accept the
# natural Python dict form and serialize here so callers don't
# have to pre-stringify.
if isinstance(value, dict):
# Serialize dicts/lists as JSON so structured params survive urlencode.
if isinstance(value, (dict, list)):
value = json.dumps(value, separators=(",", ":"))
# `urlencode` would render Python bool as "True"/"False"
# (capitalized). ux-labs validators require lowercase
# "true"/"false" (e.g. `include_thumbnail`); normalize here.
# (capitalized). TinyFish Search's bool params require lowercase
# "true"/"false" strings on the wire; normalize here.
elif isinstance(value, bool):
value = "true" if value else "false"
request_data[param] = value
@ -167,17 +166,35 @@ class TinyfishSearchConfig(BaseSearchConfig):
"""
Transform a TinyFish response to LiteLLM's unified ``SearchResponse``.
Mappings (per-result):
- ``title`` ``SearchResult.title`` (defaults to ``""`` if missing/null)
- ``url`` ``SearchResult.url`` (defaults to ``""``)
- ``snippet`` ``SearchResult.snippet`` (defaults to ``""``)
- all other per-result fields (``position``, ``site_name``,
``thumbnail_url``, ``fetch``, ``fetch_error``, ...) ride through as
extras on ``SearchResult`` via its ``extra="allow"`` config.
Per-result field handling:
- ``title``, ``url``, ``snippet`` are declared on ``SearchResult`` and
populated by ``SearchResponse.model_validate`` when present. Missing
or ``None`` values are defaulted to ``""`` beforehand by
``_default_missing_result_fields`` so a degraded result flows through
instead of failing the whole call.
- All undeclared per-result fields (``position``, ``site_name``, and
any others TinyFish returns) ride through as extras via
``SearchResult``'s ``extra="allow"`` config — accessible as
attributes on the result object or enumerable via
``result.model_extra``.
Top-level ``parameter_warnings`` (see ML-2085) is read when present and
each entry is re-fired via ``verbose_logger.warning``. Absent or
malformed entries are silently skipped never throws.
Top-level ``parameter_warnings`` is read when present and each entry
is re-fired via ``verbose_logger.warning``. Absent or malformed
entries are silently skipped never throws.
Top-level extras (``query``, ``total_results``, ``page``, and any
future TinyFish additions) ride through via
``SearchResponse.extra="allow"``. The validated response is returned
in place after truncating ``results`` to the caller's ``max_results``,
so every field pydantic populated survives regardless of which
storage bucket (declared attribute or ``__pydantic_extra__``) holds it.
TinyFish response headers (e.g. ``x-request-id``, ``retry-after``,
``x-ratelimit-limit`` httpx normalizes header names to lowercase)
are stashed on ``response._hidden_params["headers"]`` (raw) and
``response._hidden_params["additional_headers"]`` (sanitized via
``process_response_headers``) so callers can correlate a search with
server-side logs.
Error paths routed through ``self._wrap_error`` for uniform
``"TinyFish Search: <msg>. See <docs> for details."`` wrapping:
@ -223,7 +240,12 @@ class TinyfishSearchConfig(BaseSearchConfig):
_emit_parameter_warnings(parsed)
max_results: Final = self._caller_max_results or _TINYFISH_RESULT_CAP
return SearchResponse(results=list(parsed.results[:max_results]))
parsed.results = parsed.results[:max_results]
raw_headers: Final = dict(raw_response.headers)
hidden: Final = parsed._hidden_params # pyright: ignore[reportPrivateUsage] # sole hidden-params channel
hidden["headers"] = raw_headers
hidden["additional_headers"] = process_response_headers(raw_headers)
return parsed
def _wrap_error(
self,
@ -243,9 +265,9 @@ class TinyfishSearchConfig(BaseSearchConfig):
carry the ``TinyFish Search:`` prefix the bare error already names
the host in the URL, so attribution is implicit there.
"""
# ux-labs frontend wraps every error body as {"error": {"code", "message", "details"?}}.
# TinyFish Search wraps every error body as {"error": {"code", "message", "details"?}}.
# Best-effort unwrap to surface the inner message; fall back to the raw body
# for non-ux-labs responses (CDN HTML pages, other JSON envelopes, plain text).
# for other envelope shapes (CDN HTML pages, other JSON envelopes, plain text).
inner_message = error_message
try:
body: Final[object] = json.loads(error_message) # any-ok: json.loads -> Any
@ -290,7 +312,7 @@ def _default_missing_result_fields(raw_json: object) -> None:
def _emit_parameter_warnings(parsed: SearchResponse) -> None:
"""Re-fire TinyFish-side ``parameter_warnings`` (see ML-2085) as warnings.
"""Re-fire TinyFish-side ``parameter_warnings`` as warnings.
Defensive: skip silently on any shape we don't recognize so a malformed
entry (or an early/partial rollout of the field) never throws.

View file

View file

@ -0,0 +1,18 @@
"""Shared helpers for Valkey integrations (semantic cache, vector stores)."""
import struct
from collections.abc import Sequence
from typing import Final
from urllib.parse import quote
def build_valkey_url(host: str, port: str, password: str | None = None, ssl: bool = False) -> str:
"""Deliberately reads no environment: callers of the vector store control the
host, so an env-sourced password would be sent to a caller-chosen server."""
credentials: Final = f":{quote(password, safe='')}@" if password else ""
scheme: Final = "rediss" if ssl else "redis"
return f"{scheme}://{credentials}{host}:{port}"
def pack_vector(embedding: Sequence[float]) -> bytes:
return struct.pack(f"<{len(embedding)}f", *embedding)

View file

@ -0,0 +1,3 @@
from litellm.llms.valkey.vector_stores.transformation import ValkeyVectorStoreConfig
__all__ = ("ValkeyVectorStoreConfig",)

View file

@ -0,0 +1,299 @@
"""
Valkey vector store provider.
Valkey's vector search (the valkey-search module) speaks RESP only, no HTTP
API, so this config extends BaseDirectVectorStoreConfig and executes the
FT.SEARCH KNN query itself via redis-py instead of shaping an httpx request.
Documents are HASHes indexed by an FT index named after the vector_store_id.
"""
from collections.abc import Awaitable, Callable, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NoReturn
import httpx
from pydantic import BaseModel, ConfigDict
import litellm
from litellm.llms.base_llm.vector_store.transformation import BaseDirectVectorStoreConfig
from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
from litellm.types.utils import EmbeddingResponse
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from redis import Redis
from redis.asyncio import Redis as AsyncRedis
from redis.commands.search.document import Document
from redis.commands.search.query import Query
from redis.commands.search.result import Result
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
DEFAULT_VALKEY_PORT: Final = 6379
DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS: Final = 5.0
DEFAULT_SOCKET_TIMEOUT_SECONDS: Final = 30.0
DEFAULT_MAX_NUM_RESULTS: Final = 10
MIN_MAX_NUM_RESULTS: Final = 1
MAX_MAX_NUM_RESULTS: Final = 50
DEFAULT_EMBEDDING_FIELD_NAME: Final = "embedding"
DEFAULT_TEXT_FIELD_NAME: Final = "text"
DISTANCE_FIELD_NAME: Final = "vector_distance"
_EMPTY_EMBEDDING_CONFIG: Final = MappingProxyType({})
_REDIS_INSTALL_HINT: Final = (
"The Valkey vector store requires the 'redis' package. Run 'pip install redis' to install it."
)
_SEARCH_ONLY_MESSAGE: Final = "Valkey vector store is search-only; create indexes with FT.CREATE directly"
def _import_sync_redis() -> "type[Redis]":
try:
from redis import Redis as SyncRedisClient
except ImportError as e:
raise ValueError(_REDIS_INSTALL_HINT) from e
return SyncRedisClient
def _import_async_redis() -> "type[AsyncRedis]":
try:
from redis.asyncio import Redis as AsyncRedisClient
except ImportError as e:
raise ValueError(_REDIS_INSTALL_HINT) from e
return AsyncRedisClient
def _import_query() -> "type[Query]":
try:
from redis.commands.search.query import Query as RedisQuery
except ImportError as e:
raise ValueError(_REDIS_INSTALL_HINT) from e
return RedisQuery
class _ValkeySearchParams(BaseModel):
"""Typed view over the vector store's litellm_params; unrelated keys are ignored."""
model_config = ConfigDict(frozen=True, extra="ignore")
litellm_embedding_model: str | None = None
litellm_embedding_config: Mapping[str, object] | None = None
valkey_host: str | None = None
valkey_port: int | None = None
valkey_password: str | None = None
valkey_ssl: bool | None = None
valkey_text_field: str | None = None
valkey_embedding_field: str | None = None
@property
def text_field(self) -> str:
return self.valkey_text_field or DEFAULT_TEXT_FIELD_NAME
@property
def embedding_field(self) -> str:
return self.valkey_embedding_field or DEFAULT_EMBEDDING_FIELD_NAME
def require_embedding_model(self) -> str:
if not self.litellm_embedding_model:
raise ValueError(
"litellm_embedding_model is required in litellm_params for the Valkey vector store. "
"Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'"
)
return self.litellm_embedding_model
def connection_url(self) -> str:
if not self.valkey_host:
raise ValueError(
"valkey_host is required in litellm_params for the Valkey vector store. "
"Set it on the vector store's litellm_params, e.g. valkey_host: my-valkey.example.com"
)
return build_valkey_url(
host=self.valkey_host,
port=str(self.valkey_port or DEFAULT_VALKEY_PORT),
password=self.valkey_password,
ssl=bool(self.valkey_ssl),
)
class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
def __init__(
self,
sync_client: "Redis | None" = None,
async_client: "AsyncRedis | None" = None,
embedding_fn: Callable[..., EmbeddingResponse] | None = None,
aembedding_fn: Callable[..., Awaitable[EmbeddingResponse]] | None = None,
) -> None:
super().__init__()
self.sync_client = sync_client
self.async_client = async_client
self.embedding_fn = embedding_fn if embedding_fn is not None else litellm.embedding
self.aembedding_fn = aembedding_fn if aembedding_fn is not None else litellm.aembedding
@staticmethod
def _query_text(query: str | Sequence[str]) -> str:
if isinstance(query, str):
return query
if not query:
raise ValueError("query must not be empty")
return " ".join(query)
@staticmethod
def _socket_timeouts(timeout: float | httpx.Timeout | None) -> tuple[float, float]:
if isinstance(timeout, httpx.Timeout):
return (
timeout.connect or DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS,
timeout.read or DEFAULT_SOCKET_TIMEOUT_SECONDS,
)
if timeout is not None:
return (min(float(timeout), DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS), float(timeout))
return (DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS, DEFAULT_SOCKET_TIMEOUT_SECONDS)
@staticmethod
def _knn_limit(vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams) -> int:
requested: Final = vector_store_search_optional_params.get("max_num_results")
if requested is None:
return DEFAULT_MAX_NUM_RESULTS
if not MIN_MAX_NUM_RESULTS <= requested <= MAX_MAX_NUM_RESULTS:
raise ValueError(
f"max_num_results must be between {MIN_MAX_NUM_RESULTS} and {MAX_MAX_NUM_RESULTS}, got {requested}"
)
return requested
@classmethod
def _knn_query(
cls,
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
embedding_field: str,
text_field: str,
) -> "Query":
if vector_store_search_optional_params.get("filters") is not None:
raise ValueError("Valkey vector store does not support the filters parameter yet")
k: Final = cls._knn_limit(vector_store_search_optional_params)
query_cls: Final = _import_query()
knn_expr: Final = f"*=>[KNN {k} @{embedding_field} $vec AS {DISTANCE_FIELD_NAME}]"
# valkey-search rejects SORTBY on the KNN distance alias, so results are
# re-ordered client-side in _to_response instead.
return query_cls(knn_expr).return_fields(text_field, DISTANCE_FIELD_NAME).paging(0, k).dialect(2)
@staticmethod
def _to_result(doc: "Document", text_field: str) -> VectorStoreSearchResult:
content: Final = [ # mutable-ok: VectorStoreSearchResult declares a list of content parts
VectorStoreResultContent(text=str(getattr(doc, text_field, "")), type="text")
]
return VectorStoreSearchResult(
score=1.0 - float(getattr(doc, DISTANCE_FIELD_NAME)),
content=content,
file_id=getattr(doc, "id", None),
filename=getattr(doc, "id", None),
)
@classmethod
def _to_response(cls, search_result: "Result", query_text: str, text_field: str) -> VectorStoreSearchResponse:
docs: Final = getattr(search_result, "docs", None) or ()
data: Final = sorted(
(cls._to_result(doc, text_field) for doc in docs),
key=lambda result: result.get("score") or 0.0,
reverse=True,
)
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query=query_text,
data=data,
)
def execute_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
params: Final = _ValkeySearchParams.model_validate(litellm_params)
query_text: Final = self._query_text(query)
knn: Final = self._knn_query(
vector_store_search_optional_params,
embedding_field=params.embedding_field,
text_field=params.text_field,
)
embedding_response: Final = self.embedding_fn(
model=params.require_embedding_model(),
input=[query_text], # mutable-ok: litellm.embedding's input contract is a list
**(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
)
vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API
if self.sync_client is not None:
raw: Final = self.sync_client.ft(vector_store_id).search(knn, query_params=vec_params)
return self._to_response(raw, query_text, params.text_field)
connect_timeout, op_timeout = self._socket_timeouts(timeout)
client: Final = _import_sync_redis().from_url(
params.connection_url(),
socket_connect_timeout=connect_timeout,
socket_timeout=op_timeout,
)
try:
raw_result: Final = client.ft(vector_store_id).search(knn, query_params=vec_params)
return self._to_response(raw_result, query_text, params.text_field)
finally:
client.close()
async def aexecute_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
params: Final = _ValkeySearchParams.model_validate(litellm_params)
query_text: Final = self._query_text(query)
knn: Final = self._knn_query(
vector_store_search_optional_params,
embedding_field=params.embedding_field,
text_field=params.text_field,
)
embedding_response: Final = await self.aembedding_fn(
model=params.require_embedding_model(),
input=[query_text], # mutable-ok: litellm.embedding's input contract is a list
**(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
)
vec_params: Final = {"vec": pack_vector(embedding_response.data[0]["embedding"])} # mutable-ok: redis-py API
if self.async_client is not None:
raw: Final = await self.async_client.ft(vector_store_id).search( # pyright: ignore[reportGeneralTypeIssues] # types-redis 4.6 stubs shadow redis 5.3.1 and type the async client's ft() as the sync Search, so search() returns a non-awaitable Result; it is a coroutine at runtime
knn, query_params=vec_params
)
return self._to_response(raw, query_text, params.text_field)
connect_timeout, op_timeout = self._socket_timeouts(timeout)
client: Final = _import_async_redis().from_url(
params.connection_url(),
socket_connect_timeout=connect_timeout,
socket_timeout=op_timeout,
)
try:
raw_result: Final = await client.ft(vector_store_id).search( # pyright: ignore[reportGeneralTypeIssues] # types-redis 4.6 stubs shadow redis 5.3.1 and type the async client's ft() as the sync Search, so search() returns a non-awaitable Result; it is a coroutine at runtime
knn, query_params=vec_params
)
return self._to_response(raw_result, query_text, params.text_field)
finally:
await client.aclose()
def transform_create_vector_store_request(
self,
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
api_base: str,
) -> NoReturn:
raise NotImplementedError(_SEARCH_ONLY_MESSAGE)
def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn:
raise NotImplementedError(_SEARCH_ONLY_MESSAGE)

View file

@ -420,6 +420,8 @@ async def acompletion(
verbosity: Literal["low", "medium", "high"] | None = None,
safety_identifier: str | None = None,
service_tier: str | None = None,
store: bool | None = None,
prompt_cache_key: str | None = None,
# set api_base, api_version, api_key
base_url: str | None = None,
api_version: str | None = None,
@ -585,6 +587,8 @@ async def acompletion(
"verbosity": verbosity,
"safety_identifier": safety_identifier,
"service_tier": service_tier,
"store": store,
"prompt_cache_key": prompt_cache_key,
"extra_headers": extra_headers,
"acompletion": True, # assuming this is a required parameter
"thinking": thinking,
@ -4930,6 +4934,8 @@ def completion(
extra_headers: dict | None = None,
safety_identifier: str | None = None,
service_tier: str | None = None,
store: bool | None = None,
prompt_cache_key: str | None = None,
# soon to be deprecated params by OpenAI
functions: list | None = None,
function_call: str | None = None,
@ -5058,6 +5064,8 @@ def completion(
verbosity=verbosity,
safety_identifier=safety_identifier,
service_tier=service_tier,
store=store,
prompt_cache_key=prompt_cache_key,
base_url=base_url,
api_version=api_version,
api_key=api_key,
@ -5367,6 +5375,8 @@ def completion(
),
"safety_identifier": safety_identifier,
"service_tier": service_tier,
"store": store,
"prompt_cache_key": prompt_cache_key,
"allowed_openai_params": kwargs.get("allowed_openai_params"),
"base_model": base_model,
}

View file

@ -10052,6 +10052,21 @@
"output_cost_per_second": 0.0066027,
"supports_tool_choice": true
},
"bedrock/guardrails": {
"guardrail_cost_per_unit": {
"automatedReasoningPolicyUnits": 0.00017,
"contentPolicyImageUnits": 0.00075,
"contentPolicyUnits": 0.00015,
"contextualGroundingPolicyUnits": 0.0001,
"sensitiveInformationPolicyFreeUnits": 0.0,
"sensitiveInformationPolicyUnits": 0.0001,
"topicPolicyUnits": 0.00015,
"wordPolicyUnits": 0.0
},
"litellm_provider": "bedrock",
"mode": "guardrail",
"source": "https://aws.amazon.com/bedrock/pricing/"
},
"bedrock/ap-northeast-1/1-month-commitment/anthropic.claude-instant-v1": {
"input_cost_per_second": 0.01475,
"litellm_provider": "bedrock",
@ -14441,6 +14456,25 @@
"supports_tool_choice": true,
"supports_output_config": true
},
"databricks/databricks-claude-opus-4-6": {
"input_cost_per_token": 5.00003e-06,
"input_dbu_cost_per_token": 7.1429e-05,
"litellm_provider": "databricks",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 2.5000010000000002e-05,
"output_dbu_cost_per_token": 0.000357143,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4": {
"input_cost_per_token": 2.9999900000000002e-06,
"input_dbu_cost_per_token": 4.2857e-05,
@ -14498,6 +14532,25 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-claude-sonnet-4-6": {
"input_cost_per_token": 2.9999900000000002e-06,
"input_dbu_cost_per_token": 4.2857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 1000000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-2-5-flash": {
"input_cost_per_token": 3.0001999999999996e-07,
"input_dbu_cost_per_token": 4.285999999999999e-06,
@ -14532,6 +14585,74 @@
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-3-1-flash-lite": {
"input_cost_per_token": 3.1248e-07,
"input_dbu_cost_per_token": 4.464e-06,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.87502e-06,
"output_dbu_cost_per_token": 2.6786e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-3-1-pro": {
"input_cost_per_token": 2.49998e-06,
"input_dbu_cost_per_token": 3.5714e-05,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-3-flash": {
"input_cost_per_token": 6.2503e-07,
"input_dbu_cost_per_token": 8.929e-06,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 3.74997e-06,
"output_dbu_cost_per_token": 5.3571e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemini-3-pro": {
"input_cost_per_token": 2.49998e-06,
"input_dbu_cost_per_token": 3.5714e-05,
"litellm_provider": "databricks",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
"supports_function_calling": true,
"supports_tool_choice": true
},
"databricks/databricks-gemma-3-12b": {
"input_cost_per_token": 1.5000999999999998e-07,
"input_dbu_cost_per_token": 2.1429999999999996e-06,
@ -14577,6 +14698,126 @@
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-1-codex-max": {
"input_cost_per_token": 1.24999e-06,
"input_dbu_cost_per_token": 1.7857e-05,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 9.999990000000002e-06,
"output_dbu_cost_per_token": 0.000142857,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-1-codex-mini": {
"input_cost_per_token": 2.4997e-07,
"input_dbu_cost_per_token": 3.571e-06,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.99997e-06,
"output_dbu_cost_per_token": 2.8571e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-2": {
"input_cost_per_token": 1.75e-06,
"input_dbu_cost_per_token": 2.5e-05,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.4e-05,
"output_dbu_cost_per_token": 0.0002,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-2-codex": {
"input_cost_per_token": 1.75e-06,
"input_dbu_cost_per_token": 2.5e-05,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.4e-05,
"output_dbu_cost_per_token": 0.0002,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-3-codex": {
"input_cost_per_token": 1.75e-06,
"input_dbu_cost_per_token": 2.5e-05,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.4e-05,
"output_dbu_cost_per_token": 0.0002,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-4": {
"input_cost_per_token": 2.49998e-06,
"input_dbu_cost_per_token": 3.5714e-05,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-4-mini": {
"input_cost_per_token": 7.4998e-07,
"input_dbu_cost_per_token": 1.0714e-05,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 4.50002e-06,
"output_dbu_cost_per_token": 6.4286e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-4-nano": {
"input_cost_per_token": 1.9999e-07,
"input_dbu_cost_per_token": 2.857e-06,
"litellm_provider": "databricks",
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"metadata": {
"notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation."
},
"mode": "chat",
"output_cost_per_token": 1.24999e-06,
"output_dbu_cost_per_token": 1.7857e-05,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving"
},
"databricks/databricks-gpt-5-mini": {
"input_cost_per_token": 2.4997000000000006e-07,
"input_dbu_cost_per_token": 3.571e-06,
@ -19424,20 +19665,20 @@
"web_search_billing_unit": "per_query"
},
"vertex_ai/gemini-3.6-flash": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_flex": 7.5e-08,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_flex": 3.75e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_reasoning_token": 7.5e-06,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_reasoning_token": 3.75e-06,
"output_cost_per_token": 3.75e-06,
"output_cost_per_token_batches": 1.875e-06,
"output_cost_per_token_flex": 1.875e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
@ -19467,9 +19708,9 @@
"supports_vision": true,
"supports_web_search": true,
"supports_native_streaming": true,
"input_cost_per_token_priority": 2.7e-06,
"output_cost_per_token_priority": 1.35e-05,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_token_priority": 1.35e-06,
"output_cost_per_token_priority": 6.75e-06,
"cache_read_input_token_cost_priority": 1.35e-07,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
@ -21150,20 +21391,20 @@
"web_search_billing_unit": "per_query"
},
"gemini/gemini-3.6-flash": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_flex": 7.5e-08,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_flex": 3.75e-07,
"litellm_provider": "gemini",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_reasoning_token": 7.5e-06,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_reasoning_token": 3.75e-06,
"output_cost_per_token": 3.75e-06,
"output_cost_per_token_batches": 1.875e-06,
"output_cost_per_token_flex": 1.875e-06,
"rpm": 2000,
"source": "https://ai.google.dev/pricing/gemini-3",
"supported_endpoints": [
@ -21196,9 +21437,9 @@
"supports_web_search": true,
"supports_native_streaming": true,
"tpm": 800000,
"input_cost_per_token_priority": 2.7e-06,
"output_cost_per_token_priority": 1.35e-05,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_token_priority": 1.35e-06,
"output_cost_per_token_priority": 6.75e-06,
"cache_read_input_token_cost_priority": 1.35e-07,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,
@ -21544,20 +21785,20 @@
"web_search_billing_unit": "per_query"
},
"gemini-3.6-flash": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost_flex": 7.5e-08,
"input_cost_per_token": 1.5e-06,
"input_cost_per_token_batches": 7.5e-07,
"input_cost_per_token_flex": 7.5e-07,
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_flex": 3.75e-08,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_flex": 3.75e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_reasoning_token": 7.5e-06,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_flex": 3.75e-06,
"output_cost_per_reasoning_token": 3.75e-06,
"output_cost_per_token": 3.75e-06,
"output_cost_per_token_batches": 1.875e-06,
"output_cost_per_token_flex": 1.875e-06,
"source": "https://ai.google.dev/pricing/gemini-3",
"supported_endpoints": [
"/v1/chat/completions",
@ -21588,9 +21829,9 @@
"supports_vision": true,
"supports_web_search": true,
"supports_native_streaming": true,
"input_cost_per_token_priority": 2.7e-06,
"output_cost_per_token_priority": 1.35e-05,
"cache_read_input_token_cost_priority": 2.7e-07,
"input_cost_per_token_priority": 1.35e-06,
"output_cost_per_token_priority": 6.75e-06,
"cache_read_input_token_cost_priority": 1.35e-07,
"search_context_cost_per_query": {
"search_context_size_low": 0.014,
"search_context_size_medium": 0.014,

View file

@ -21,7 +21,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms.azure_ai.ocr.common_utils import (
is_azure_document_intelligence_model,
)
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse
from litellm.llms.base_llm.ocr.transformation import (
OCR_REQUEST_FORMAT_PARAM,
BaseOCRConfig,
OCRResponse,
parse_ocr_request_format,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.types.router import GenericLiteLLMParams
@ -124,6 +129,24 @@ def _prepare_ocr_request(
litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs)
supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model)
requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM)
if requested_format is not None:
try:
parsed_format: Final = parse_ocr_request_format(requested_format)
except ValueError as e:
raise litellm.exceptions.UnsupportedParamsError(
message=f"{e}", model=model, llm_provider=custom_llm_provider
) from e
if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native":
raise litellm.exceptions.UnsupportedParamsError(
message=(
f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, "
f"model: {model}"
),
model=model,
llm_provider=custom_llm_provider,
)
non_default_params: Final = {}
for param in supported_params:
if param in kwargs:
@ -166,6 +189,8 @@ def _prepare_ocr_request(
def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native":
return False
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS

View file

@ -1341,7 +1341,7 @@ async def _persist_dcr_client_registration(
``update_mcp_server`` merges credential blobs: a re-registered public client must not
inherit the previous client's secret or auth method.
"""
if mcp_server.is_true_passthrough or mcp_server.is_oauth_delegate:
if mcp_server.is_client_forwarded_token:
return "skipped"
try:
@ -1678,10 +1678,17 @@ async def authorize(
lookup_name: Final[str | None] = mcp_server_name or client_id
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = (
global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None
await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip)
if lookup_name
else None
)
if mcp_server is None and mcp_server_name is None:
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
mcp_server = (
await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server)
if unresolved_server is not None
else None
)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
_raise_if_not_oauth2(mcp_server)
@ -1761,9 +1768,14 @@ async def token_endpoint(
lookup_name: Final = mcp_server_name or client_id
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip)
mcp_server = await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip)
if mcp_server is None and mcp_server_name is None:
mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
mcp_server = (
await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server)
if unresolved_server is not None
else None
)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
return await exchange_token_with_server(
@ -2175,7 +2187,7 @@ async def _build_oauth_protected_resource_response(
)
if upstream_metadata is not None:
if mcp_server.is_true_passthrough or mcp_server.is_oauth_delegate:
if mcp_server.is_client_forwarded_token:
return upstream_metadata
return {**upstream_metadata, "resource": resource_url}
@ -2397,6 +2409,7 @@ def _build_oauth_authorization_server_response(
request_base_url: Final = get_request_base_url(request)
client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
explicitly_named: Final = mcp_server_name is not None
# When no server name provided, try to resolve the single OAuth2 server
if mcp_server_name is None:
@ -2415,8 +2428,10 @@ def _build_oauth_authorization_server_response(
_raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth authorization server")
issuer: Final = f"{request_base_url}/{mcp_server_name}" if explicitly_named else request_base_url
return {
"issuer": request_base_url, # point to your proxy
"issuer": issuer,
"authorization_endpoint": authorization_endpoint,
"token_endpoint": token_endpoint,
"response_types_supported": ["code"],
@ -2562,9 +2577,10 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
return await register_aggregate_client(request=request, request_body=data)
resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
if resolved:
resolved_server: Final = await global_mcp_server_manager.ensure_oauth_metadata_discovered(resolved)
return await register_client_with_server(
request=request,
mcp_server=resolved,
mcp_server=resolved_server,
client_name=data.get("client_name", ""),
grant_types=data.get("grant_types", []),
response_types=data.get("response_types", []),
@ -2574,7 +2590,10 @@ async def register_client(request: Request, mcp_server_name: str | None = None):
)
return dummy_return
mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip)
mcp_server: Final = await global_mcp_server_manager.get_resolved_mcp_server_by_name(
mcp_server_name,
client_ip=client_ip,
)
if mcp_server is None:
return dummy_return
return await register_client_with_server(

File diff suppressed because it is too large Load diff

View file

@ -1704,7 +1704,7 @@ if MCP_AVAILABLE:
)
extra_headers: dict[str, str] | None = None
is_client_forwarded_mode: Final = server.is_true_passthrough or server.is_oauth_delegate
is_client_forwarded_mode: Final = server.is_client_forwarded_token
# In a multi-server listing scope the request-wide Authorization can only carry one token,
# so it is withheld from a client-forwarded server when another server in scope also consumes
# it (RFC 9700 cross-resource replay); such scopes must bind per-server via
@ -2013,6 +2013,9 @@ if MCP_AVAILABLE:
prefetched_creds=_prefetched_oauth_creds,
)
if server.is_byok and server.auth_type != MCPAuth.oauth2 and server_auth_header is None:
server_auth_header = await _get_byok_credential(server, user_api_key_auth)
try:
tools: Final = await global_mcp_server_manager._get_tools_from_server(
server=server,
@ -2824,6 +2827,7 @@ if MCP_AVAILABLE:
proxy_logging_obj=proxy_logging_obj,
server=mcp_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
# `pre_call_tool_check` may return guardrail-modified
# arguments; honor them on the local path too.
@ -2962,6 +2966,7 @@ if MCP_AVAILABLE:
proxy_logging_obj=proxy_logging_obj,
server=prefix_server,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
if "arguments" in hook_result:
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args
@ -3149,6 +3154,20 @@ if MCP_AVAILABLE:
traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)
from litellm.proxy.proxy_server import proxy_logging_obj
# Ordering is load-bearing. ``_ProxyDBLogger.async_post_call_failure_hook``,
# reached below, writes the failure spend-log row from this logger's
# ``standard_logging_object``, which only exists once the failure handlers
# have run. Flush them first or the row lands with
# ``guardrail_information=None`` and a guardrail block is never counted.
#
# Not double-logged: both handlers gate on ``should_run_logging`` and then
# mark it, so the ``@client`` wrapper's own post-raise logging no-ops on this
# logger, same as ``_fire_mcp_tool_call_logging`` does for ``isError=True``.
if litellm_logging_obj is not None:
end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from
litellm_logging_obj.failure_handler(e, traceback_str, start_time, end_time)
await litellm_logging_obj.async_failure_handler(e, traceback_str, start_time, end_time)
if proxy_logging_obj and user_api_key_auth:
await proxy_logging_obj.post_call_failure_hook(
request_data=kwargs,
@ -3326,6 +3345,7 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
host_progress_callback=host_progress_callback,
litellm_logging_obj=litellm_logging_obj,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
return call_tool_result
@ -3737,6 +3757,14 @@ if MCP_AVAILABLE:
# preemptive challenge and let downstream authorization
# return 403.
continue
if server is not None and server.auth_type == MCPAuth.oauth2 and server.oauth2_flow == "client_credentials":
# Stamped M2M: the challenge decision below never reads discovered
# metadata, so deferred-discovery failures must not 503 this loop.
# Unstamped rows stay on the discover-first path because filling
# authorization_url/token_url can change their inferred flow.
continue
if server is not None:
server = await global_mcp_server_manager.ensure_oauth_metadata_discovered(server)
if server and server.auth_type == MCPAuth.oauth2:
# The challenge decision is per oauth2 sub-mode, not per header:
# gateway-managed modes (M2M and interactive authorization_code)

View file

@ -0,0 +1,6 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg width="64" height="73" viewBox="0 0 64 73" xmlns="http://www.w3.org/2000/svg">
<g id="Group-copy">
<path id="Path" fill="#123678" fill-rule="evenodd" stroke="none" d="M 13.482285 60.694962 L 0.998384 52.884399 L 0.998384 19.502914 L 31.527868 2.001205 L 61.317604 19.532024 L 61.317604 54.64489 L 31.054855 71.68927 L 20.548372 65.115807 L 20.548372 51.041328 L 20.548372 49.119896 L 14.851504 45.555508 L 14.851504 27.453159 L 31.346497 17.99712 L 47.464485 27.482262 L 47.464485 46.451157 L 34.703495 53.638138 L 34.703495 45.998573 C 38.52874 44.52552 41.274452 40.739189 41.274452 36.270489 C 41.274452 30.510658 36.712814 25.88438 31.158138 25.88438 C 25.603172 25.88438 21.041817 30.510658 21.041817 36.270489 C 21.041817 40.739189 23.787249 44.52552 27.612494 45.998573 L 27.612494 60.473576 L 31.261133 62.756348 L 53.635483 50.15464 L 53.635483 23.924595 L 31.477489 10.884869 L 8.680504 23.953705 L 8.680504 48.628967 L 13.482285 51.633297 L 13.482285 60.694962 Z M 31.158138 31.498383 C 33.671822 31.498383 35.660439 33.664162 35.660439 36.270489 C 35.660439 38.876804 33.671822 41.042587 31.158138 41.042587 C 28.644447 41.042587 26.655558 38.876804 26.655558 36.270489 C 26.655558 33.664162 28.644447 31.498383 31.158138 31.498383 Z" />
</g>
</svg>

After

Width:  |  Height:  |  Size: 1.3 KiB

View file

@ -11349,6 +11349,40 @@
"title": "UpdateGuardrailRequest",
"type": "object"
},
"UsageChartPoint": {
"properties": {
"blocked": {
"title": "Blocked",
"type": "integer"
},
"date": {
"title": "Date",
"type": "string"
},
"passed": {
"title": "Passed",
"type": "integer"
},
"score": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Score"
}
},
"required": [
"date",
"passed",
"blocked"
],
"title": "UsageChartPoint",
"type": "object"
},
"UsageDetailResponse": {
"properties": {
"avgLatency": {
@ -11410,8 +11444,7 @@
},
"time_series": {
"items": {
"additionalProperties": true,
"type": "object"
"$ref": "#/components/schemas/UsageChartPoint"
},
"title": "Time Series",
"type": "array"
@ -11423,6 +11456,40 @@
"type": {
"title": "Type",
"type": "string"
},
"usage_units": {
"additionalProperties": {
"type": "integer"
},
"title": "Usage Units",
"type": "object"
},
"usage_units_by_key": {
"additionalProperties": {
"additionalProperties": {
"type": "integer"
},
"type": "object"
},
"title": "Usage Units By Key",
"type": "object"
},
"usage_units_by_team": {
"additionalProperties": {
"additionalProperties": {
"type": "integer"
},
"type": "object"
},
"title": "Usage Units By Team",
"type": "object"
},
"usage_units_daily": {
"items": {
"$ref": "#/components/schemas/UsageUnitsDailyPoint"
},
"title": "Usage Units Daily",
"type": "array"
}
},
"required": [
@ -11437,7 +11504,11 @@
"status",
"trend",
"description",
"time_series"
"time_series",
"usage_units",
"usage_units_daily",
"usage_units_by_team",
"usage_units_by_key"
],
"title": "UsageDetailResponse",
"type": "object"
@ -11572,8 +11643,7 @@
"properties": {
"chart": {
"items": {
"additionalProperties": true,
"type": "object"
"$ref": "#/components/schemas/UsageChartPoint"
},
"title": "Chart",
"type": "array"
@ -11596,6 +11666,13 @@
"totalRequests": {
"title": "Totalrequests",
"type": "integer"
},
"totalUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Totalusageunits",
"type": "object"
}
},
"required": [
@ -11603,7 +11680,8 @@
"chart",
"totalRequests",
"totalBlocked",
"passRate"
"passRate",
"totalUsageUnits"
],
"title": "UsageOverviewResponse",
"type": "object"
@ -11663,6 +11741,13 @@
"type": {
"title": "Type",
"type": "string"
},
"usageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Usageunits",
"type": "object"
}
},
"required": [
@ -11675,11 +11760,33 @@
"avgScore",
"avgLatency",
"status",
"trend"
"trend",
"usageUnits"
],
"title": "UsageOverviewRow",
"type": "object"
},
"UsageUnitsDailyPoint": {
"properties": {
"date": {
"title": "Date",
"type": "string"
},
"units": {
"additionalProperties": {
"type": "integer"
},
"title": "Units",
"type": "object"
}
},
"required": [
"date",
"units"
],
"title": "UsageUnitsDailyPoint",
"type": "object"
},
"ValidationError": {
"properties": {
"loc": {
@ -21477,6 +21584,13 @@
"totalRequests": {
"title": "Totalrequests",
"type": "integer"
},
"totalUsageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Totalusageunits",
"type": "object"
}
},
"required": [
@ -21484,7 +21598,8 @@
"chart",
"totalRequests",
"totalBlocked",
"passRate"
"passRate",
"totalUsageUnits"
],
"title": "UsageOverviewResponse",
"type": "object"
@ -21544,6 +21659,13 @@
"type": {
"title": "Type",
"type": "string"
},
"usageUnits": {
"additionalProperties": {
"type": "integer"
},
"title": "Usageunits",
"type": "object"
}
},
"required": [
@ -21556,7 +21678,8 @@
"avgScore",
"avgLatency",
"status",
"trend"
"trend",
"usageUnits"
],
"title": "UsageOverviewRow",
"type": "object"

View file

@ -287,6 +287,7 @@ class KeyManagementRoutes(str, enum.Enum):
# team usage routes
TEAM_DAILY_ACTIVITY = "/team/daily/activity"
TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated"
# team spend-log viewing
SPEND_LOGS = "/spend/logs"
@ -451,6 +452,7 @@ class LiteLLMRoutes(enum.Enum):
mapped_pass_through_routes = [
"/bedrock",
"/comprehendmedical",
"/vertex-ai",
"/vertex_ai",
"/cohere",
@ -611,6 +613,7 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED.value,
KeyManagementRoutes.SPEND_LOGS.value,
KeyManagementRoutes.SPEND_LOGS_V2.value,
KeyManagementRoutes.KEY_RESET_SPEND.value,
@ -645,6 +648,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/permissions_update",
"/team/permissions_bulk_update",
"/team/daily/activity",
"/team/daily/activity/aggregated",
# gateway request counts (SGR); deployment-wide, admin-only
"/gateway/daily/activity",
# model
@ -680,6 +684,7 @@ class LiteLLMRoutes(enum.Enum):
# permitted teams exactly like /spend/logs/ui — it belongs to the same
# access tier, not to customer management.
"/management/v1/spend_logs/end_users",
"/management/v1/spend_logs/users",
"/cost/estimate",
]
@ -799,12 +804,16 @@ class LiteLLMRoutes(enum.Enum):
"/team/permissions_list",
"/team/permissions_update",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/team/{team_id}/members/me",
"/model/new",
"/model/update",
"/model/delete",
"/user/daily/activity",
"/user/daily/activity/aggregated",
# Endpoint restricts results to organizations the caller is ORG_ADMIN
# of; a caller who administers none gets an empty result set.
"/organization/daily/activity",
"/user/available_roles", # read-only role metadata; any authenticated user may read
"/user/list", # org admins checked in endpoint; non-admins get 403
"/model/{model_id}/update",
@ -861,6 +870,7 @@ class LiteLLMRoutes(enum.Enum):
"/user/available_roles",
"/user/daily/activity",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/tag/daily/activity",
"/tag/list",
"/audit",
@ -872,12 +882,13 @@ class LiteLLMRoutes(enum.Enum):
# PROXY_ADMIN_VIEW_ONLY — the route gate must match).
"/customer/list",
"/customer/info",
# UI Logs page detail drawer (single + session) and the end-user filter
# facet. The list endpoint `/spend/logs/ui` is covered via
# UI Logs page detail drawer (single + session) and the filter facets.
# The list endpoint `/spend/logs/ui` is covered via
# spend_tracking_routes below.
"/spend/logs/ui/{logId}",
"/spend/logs/session/ui",
"/management/v1/spend_logs/end_users",
"/management/v1/spend_logs/users",
# Settings / observability read endpoints exposed in admin-only
# sidebar groups (Logging & Alerts, Admin Settings, Budgets,
# Invitations).
@ -1987,6 +1998,18 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
return values
class TeamCallbackDeleteResponseData(LiteLLMPydanticObjectBase):
team_id: str
success_callbacks: tuple[str, ...]
failure_callbacks: tuple[str, ...]
class TeamCallbackDeleteResponse(LiteLLMPydanticObjectBase):
status: Literal["success"]
message: str
data: TeamCallbackDeleteResponseData
class TeamCallbackMetadata(LiteLLMPydanticObjectBase):
success_callback: list[str] | None = []
failure_callback: list[str] | None = []
@ -3036,6 +3059,8 @@ class NewProjectRequest(LiteLLM_BudgetTable):
models: list[str] = []
model_rpm_limit: dict | None = None
model_tpm_limit: dict | None = None
model_itpm_limit: Mapping[str, int] | None = None
model_otpm_limit: Mapping[str, int] | None = None
blocked: bool = False
object_permission: LiteLLM_ObjectPermissionBase | None = None
@ -3068,6 +3093,8 @@ class UpdateProjectRequest(LiteLLM_BudgetTable):
models: list[str] | None = None
model_rpm_limit: dict | None = None
model_tpm_limit: dict | None = None
model_itpm_limit: Mapping[str, int] | None = None
model_otpm_limit: Mapping[str, int] | None = None
blocked: bool | None = None
budget_id: str | None = None
object_permission: LiteLLM_ObjectPermissionBase | None = None
@ -4219,6 +4246,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict):
LiteLLM_ManagementEndpoint_MetadataFields: Final = [
"model_rpm_limit",
"model_tpm_limit",
"model_itpm_limit",
"model_otpm_limit",
"default_estimated_output_tokens",
"default_estimated_output_tokens_per_model",
"mcp_rpm_limit",

View file

@ -13,6 +13,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
import json
from collections.abc import AsyncGenerator, Mapping
from copy import deepcopy
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from urllib.parse import urlparse
@ -20,6 +21,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.proxy._types import UserAPIKeyAuth
@ -36,6 +38,11 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
)
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING,
coerce_keepalive_interval,
wrap_sse_stream_with_keepalive_pings,
)
from litellm.proxy.utils import ProxyLogging, get_custom_url
from litellm.types.utils import all_litellm_params
@ -46,6 +53,15 @@ if TYPE_CHECKING:
router: Final = APIRouter()
# Mirrors the native seam's own headers: a reverse proxy that batches the whole
# stream would swallow the keepalives this route sends to defeat idle timeouts.
_SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
{
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
)
_PASCAL_TO_WIRE: Final[Mapping[str, str]] = {
"SendMessage": "message/send",
"SendStreamingMessage": "message/stream",
@ -326,7 +342,19 @@ async def _forward_jsonrpc_sse(
generator = _passthrough()
return StreamingResponse(generator, media_type="text/event-stream")
# The upstream agent is only contacted once this generator is first pulled, so
# a slow first event leaves the response body idle for its whole
# time-to-first-token and an intermediary with an idle read timeout drops a
# healthy connection. Off until an operator sets an interval, and the
# buffering hint only goes out when there are keepalives to protect.
keepalive_interval: Final = coerce_keepalive_interval(litellm.sse_keepalive_ping_interval_seconds)
if keepalive_interval is None:
return StreamingResponse(generator, media_type="text/event-stream")
return StreamingResponse(
wrap_sse_stream_with_keepalive_pings(generator, keepalive_interval, ping_chunk=SSE_COMMENT_PING),
media_type="text/event-stream",
headers=_SSE_KEEPALIVE_HEADERS,
)
async def _handle_stream_message(

View file

@ -193,6 +193,9 @@ async def anthropic_response(
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
if isinstance(e, ProxyException):
raise
# Extract model_id from request metadata (same as success path)
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
model_info: Final = litellm_metadata.get("model_info", {}) or {}

View file

@ -1172,7 +1172,7 @@ def enforce_output_token_estimates_are_admin_only(
def get_model_rate_limit_from_metadata(
user_api_key_dict: UserAPIKeyAuth,
metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"],
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit", "model_itpm_limit", "model_otpm_limit"],
) -> dict[str, int] | None:
if getattr(user_api_key_dict, metadata_accessor_key):
return getattr(user_api_key_dict, metadata_accessor_key).get(rate_limit_key)

View file

@ -0,0 +1,18 @@
from litellm.proxy._types import ProxyException
def validate_batch_list_limit(limit: int | None) -> None:
if limit is None or 0 <= limit <= 100:
return
bound, expected, openai_code = (
("below minimum", ">= 0", "integer_below_min_value")
if limit < 0
else ("above maximum", "<= 100", "integer_above_max_value")
)
raise ProxyException(
message=f"Invalid 'limit': integer {bound} value. Expected a value {expected}, but got {limit} instead.",
type="invalid_request_error",
param="limit",
code=400,
openai_code=openai_code,
)

View file

@ -5,6 +5,8 @@
######################################################################
import asyncio
import os
from collections.abc import Mapping
from typing import Any, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
@ -14,6 +16,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
@ -40,6 +43,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
update_batch_in_database,
validate_managed_id_requirement,
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.types.llms.openai import LiteLLMBatchCreateRequest
@ -47,6 +51,23 @@ from litellm.types.llms.openai import LiteLLMBatchCreateRequest
router: Final = APIRouter()
def _raise_not_found_when_openai_fallback_unservable(
requested_provider: "str | None",
data: Mapping[str, object],
not_found_message: str,
) -> None:
if requested_provider is not None:
return
if data.get("api_key") or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY"):
return
raise ProxyException(
message=not_found_message,
type="invalid_request_error",
param=None,
code=404,
)
async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str | None":
"""Resolve a managed (unified) input_file_id to its backend storage_url.
@ -140,6 +161,8 @@ async def create_batch(
)
data["metadata"] = sanitize_openai_provider_metadata(data.get("metadata"))
raise_if_required_body_param_missing(route_type="acreate_batch", data=data)
## check if model is a loadbalanced model
router_model: str | None = None
is_router_model = False
@ -147,12 +170,12 @@ async def create_batch(
router_model = data.get("model", None)
is_router_model = is_known_model(model=router_model, llm_router=llm_router)
custom_llm_provider: Final = (
requested_provider: Final = (
provider
or data.pop("custom_llm_provider", None)
or get_custom_llm_provider_from_request_headers(request=request)
or "openai"
)
custom_llm_provider: Final = requested_provider or "openai"
_create_batch_data: Final = LiteLLMBatchCreateRequest(**data)
# Apply team-level batch output expiry enforcement
@ -314,6 +337,11 @@ async def create_batch(
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
_raise_not_found_when_openai_fallback_unservable(
requested_provider=requested_provider,
data=cast(dict, _create_batch_data), # cast-ok: TypedDict is a dict at runtime
not_found_message=f"No such File object: {input_file_id}",
)
response = await litellm.acreate_batch(
custom_llm_provider=custom_llm_provider,
**_create_batch_data,
@ -563,18 +591,23 @@ async def retrieve_batch(
# SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
else:
custom_llm_provider: Final = (
requested_provider: Final = (
provider
or get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or "openai"
)
custom_llm_provider: Final = requested_provider or "openai"
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
_raise_not_found_when_openai_fallback_unservable(
requested_provider=requested_provider,
data=data,
not_found_message=f"No batch found with id '{batch_id}'.",
)
response = await litellm.aretrieve_batch(
custom_llm_provider=custom_llm_provider,
**data,
@ -679,6 +712,7 @@ async def list_batches(
```
"""
validate_batch_list_limit(limit)
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
@ -967,13 +1001,13 @@ async def cancel_batch(
# SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
else:
body_custom_llm_provider = data.pop("custom_llm_provider", None)
custom_llm_provider: Final = (
requested_provider: Final = (
provider
or body_custom_llm_provider
or get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or "openai"
)
custom_llm_provider: Final = requested_provider or "openai"
# Extract batch_id from data to avoid "multiple values for keyword argument" error
# data was cast from CancelBatchRequest which already contains batch_id
data.pop("batch_id", None)
@ -983,6 +1017,11 @@ async def cancel_batch(
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
_raise_not_found_when_openai_fallback_unservable(
requested_provider=requested_provider,
data=data,
not_found_message=f"No batch found with id '{batch_id}'.",
)
_cancel_batch_data: Final = CancelBatchRequest(batch_id=batch_id, **data)
response = await litellm.acancel_batch(
custom_llm_provider=custom_llm_provider,

View file

@ -1,14 +1,15 @@
import asyncio
import contextlib
import json
import logging
import math
import time
import traceback
from collections.abc import AsyncGenerator, Callable, Mapping
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, overload
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload
import anyio
import httpx
@ -31,11 +32,13 @@ from litellm.constants import (
UNSAFE_PROXY_RESPONSE_HEADERS,
)
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
get_response_headers,
)
@ -47,7 +50,12 @@ from litellm.proxy.common_utils.callback_utils import (
get_logging_caching_headers,
get_remaining_tokens_and_requests_from_request_data,
)
from litellm.proxy.common_utils.sse_keepalive import wrap_sse_stream_with_keepalive_pings
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING_BYTES,
coerce_keepalive_interval,
resolve_ttft_keepalive_interval,
wrap_sse_stream_with_keepalive_pings,
)
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
@ -56,6 +64,100 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
_LlmCallT = TypeVar("_LlmCallT")
ProxyRouteType: TypeAlias = Literal[
"acompletion",
"aembedding",
"aresponses",
"_arealtime",
"_aresponses_websocket",
"acreate_realtime_client_secret",
"arealtime_calls",
"aget_responses",
"adelete_responses",
"acancel_responses",
"acompact_responses",
"acreate_batch",
"aretrieve_batch",
"alist_batches",
"acancel_batch",
"afile_content",
"afile_retrieve",
"afile_delete",
"atext_completion",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
"aretrieve_fine_tuning_job",
"alist_input_items",
"aimage_edit",
"agenerate_content",
"agenerate_content_stream",
"allm_passthrough_route",
"avector_store_search",
"avector_store_create",
"avector_store_retrieve",
"avector_store_list",
"avector_store_update",
"avector_store_delete",
"avector_store_file_create",
"avector_store_file_list",
"avector_store_file_retrieve",
"avector_store_file_content",
"avector_store_file_update",
"avector_store_file_delete",
"aocr",
"asearch",
"avideo_generation",
"avideo_list",
"avideo_status",
"avideo_content",
"avideo_remix",
"avideo_create_character",
"avideo_get_character",
"avideo_edit",
"avideo_extension",
"acreate_container",
"alist_containers",
"aingest",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
"aretrieve_container_file_content",
"acreate_skill",
"alist_skills",
"aget_skill",
"adelete_skill",
"anthropic_messages",
"acreate_interaction",
"aget_interaction",
"adelete_interaction",
"acancel_interaction",
"acreate_agent",
"alist_agents",
"aget_agent",
"adelete_agent",
"alist_agent_versions",
"asend_message",
"call_mcp_tool",
"acreate_eval",
"alist_evals",
"aget_eval",
"aupdate_eval",
"adelete_eval",
"acancel_eval",
"acreate_run",
"alist_runs",
"aget_run",
"acancel_run",
"adelete_run",
]
from litellm.types.utils import ServerToolUse
# Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format)
@ -559,6 +661,11 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
super().__init__(content, status_code=status_code, headers=headers, media_type=media_type)
self._upstream_generator = upstream_generator
@property
def upstream_generator(self) -> AsyncGenerator[str, None] | None:
"""The upstream LLM stream, for a caller that has to run this response's cleanup itself."""
return self._upstream_generator
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
try:
await super().__call__(scope, receive, send)
@ -649,6 +756,39 @@ async def _buffer_first_chunk_honoring_disconnect(
raise _ClientDisconnectedBeforeFirstChunk()
def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
"""Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames.
Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames
are byte-identical.
"""
# Preserve status code from HTTPException (e.g. guardrail blocks)
error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start")
message, structured_fields = _serialize_http_exception_detail(raw_detail)
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None)
# Built in one statement then given its one optional key, rather than spread
# conditionally: the spread form costs two extra dict constructions, which
# type-discipline-budget.json's LIT002 ceiling has no room for.
error_obj: Final = {
"message": message,
"type": getattr(exc, "type", "None"),
"param": getattr(exc, "param", "None"),
"code": str(error_status),
}
if merged_fields:
error_obj["provider_specific_fields"] = merged_fields
return error_status, error_obj
def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
"""The two frames an SSE stream ends with once it can no longer raise."""
return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n"
async def create_response(
generator: AsyncGenerator[str, None],
media_type: str,
@ -740,31 +880,11 @@ async def create_response(
# Unexpected error consuming first chunk.
verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e)
# Preserve status code from HTTPException (e.g., guardrail blocks)
error_status: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
raw_detail: Final = _getattr_object(e, "detail", "Error processing stream start")
message, structured_fields = _serialize_http_exception_detail(raw_detail)
existing_fields: Final = getattr(e, "provider_specific_fields", None) or {}
if structured_fields:
merged_fields: dict | None = {**existing_fields, **structured_fields}
else:
merged_fields = existing_fields or None
# Match ProxyException.to_dict() shape so streaming and non-streaming
# error frames are byte-identical.
error_obj: Final[dict[str, object]] = {
"message": message,
"type": getattr(e, "type", "None"),
"param": getattr(e, "param", "None"),
"code": str(error_status),
}
if merged_fields:
error_obj["provider_specific_fields"] = merged_fields
error_status, error_obj = _sse_error_payload(e)
async def error_gen_message() -> AsyncGenerator[str, None]:
yield f"data: {json.dumps({'error': error_obj})}\n\n"
yield "data: [DONE]\n\n"
for frame in _sse_error_frames(error_obj):
yield frame
return StreamingResponse(
error_gen_message(),
@ -797,6 +917,176 @@ async def create_response(
)
_TTFT_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType(
{
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
)
def ttft_keepalive_interval(request_data: Mapping[str, object], llm_router: Router | None = None) -> float | None:
"""The operator's keepalive interval, but only for a request that asked to stream.
Resolved through the deployments the request could land on, so a deployment's
`keepalive_seconds: 0` stays the hard disable it is documented to be rather
than being switched back on by the global default.
"""
if request_data.get("stream") is not True:
return None
requested_model: Final = request_data.get("model")
deployments: Final = (
llm_router.get_model_list(model_name=requested_model) or ()
if llm_router is not None and isinstance(requested_model, str)
else ()
)
return resolve_ttft_keepalive_interval(deployments, litellm.sse_keepalive_ping_interval_seconds)
async def _aclose_late_response(produced: Response) -> None:
"""Run the cleanup Starlette would have run, for a response it never called.
Closing an already-closed async generator is a no-op, so this is safe to call
from both the relay's own teardown and the outer one.
"""
if not isinstance(produced, StreamingResponse):
return
targets: Final = (
(produced.body_iterator, produced.upstream_generator)
if isinstance(produced, _UpstreamClosingStreamingResponse)
else (produced.body_iterator,)
)
for target in targets:
aclose = getattr(target, "aclose", None)
if aclose is None:
continue
try:
await aclose()
except BaseException as exc: # noqa: BLE001 # teardown must not mask why the stream ended
verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc)
async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None]:
"""Replay a Response that was built after a keepalive had already opened the wire."""
if not isinstance(produced, StreamingResponse):
# The status line is already on the wire, so a non-streaming body, an error
# body included, can only reach the client as an SSE frame.
yield b"data: " + (bytes(produced.body) or b"{}") + b"\n\n"
yield b"data: [DONE]\n\n"
return
try:
async for chunk in produced.body_iterator:
yield chunk.encode("utf-8") if isinstance(chunk, str) else bytes(chunk)
finally:
# Starlette never called this response, so the cleanup its __call__ would
# have run has to happen here or the upstream LLM connection leaks.
with anyio.CancelScope(shield=True):
await _aclose_late_response(produced)
async def _sanitized_late_failure(
exc: Exception,
on_late_failure: "Callable[[Exception], Awaitable[HTTPException | None]] | None",
) -> Exception:
"""Report a late failure and return whatever should reach the client.
``post_call_failure_hook`` lets a callback replace the client-facing error, by
returning a replacement or by raising one, and both are used elsewhere in this
module. Serializing the original would leak provider detail a deployment had
configured away, so the hook's answer wins. A callback that fails some other
way is a bug in the callback, not a reason to lose the real error.
"""
if on_late_failure is None:
return exc
try:
replacement: Final = await on_late_failure(exc)
except HTTPException as raised_replacement:
return raised_replacement
except Exception as hook_failure: # noqa: BLE001 # a broken callback must not replace the real error
verbose_proxy_logger.exception("post_call_failure_hook raised while reporting a late failure: %s", hook_failure)
return exc
return replacement if replacement is not None else exc
async def open_sse_before_first_byte(
produce_response: Awaitable[_LateResponseT],
ping_interval_seconds: float | str | None,
media_type: str = "text/event-stream",
on_late_failure: Callable[[Exception], Awaitable[HTTPException | None]] | None = None,
) -> _LateResponseT | StreamingResponse:
"""Write SSE keepalive comments while the upstream LLM call is still in flight.
The whole time-to-first-token is spent inside `produce_response`: the upstream
withholds its response headers until it emits its first token, so nothing has
entered the ASGI response phase yet and the proxy writes zero bytes. An
intermediary with an idle read timeout (AWS ALB and nginx both default to 60s)
then drops a connection that is perfectly healthy.
When `produce_response` does not finish within one interval, the response is
opened immediately and `: ping` comments, which every conformant SSE client
ignores, fill the wire until the real response is ready to be replayed onto it.
Committing the status line that early is the cost: a failure discovered after
the first ping reaches the client as an SSE error frame under a 200 rather than
as an HTTP error status, and LiteLLM's own `x-litellm-*` response headers are
not yet known. Both are why this stays off until an operator sets an interval.
"""
interval: Final = coerce_keepalive_interval(ping_interval_seconds)
if interval is None:
return await produce_response
produce_task: Final = asyncio.ensure_future(produce_response)
await asyncio.wait((produce_task,), timeout=interval)
if produce_task.done():
# Fast path: the upstream answered inside one interval, so nothing was
# written early and this is byte-identical to not being wrapped at all.
return produce_task.result()
async def keepalive_then_relay() -> AsyncGenerator[bytes, None]:
try:
while not produce_task.done():
yield SSE_COMMENT_PING_BYTES
await asyncio.wait((produce_task,), timeout=interval)
try:
produced: Final = produce_task.result()
except Exception as exc: # noqa: BLE001 # the status line is already sent; surface it as a frame
verbose_proxy_logger.exception(
"request failed after its SSE keepalive had opened the response: %s", exc
)
# The caller's own `except` never sees this, so its failure hook
# would never fire and the failure would go unaudited. The hook
# also gets to sanitize what reaches the client, by returning or
# raising a replacement, so its answer decides the frame.
_, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure))
for frame in _sse_error_frames(error_obj):
yield frame.encode()
return
async for chunk in _relay_late_response(produced):
yield chunk
finally:
if not produce_task.done():
produce_task.cancel()
with anyio.CancelScope(shield=True):
with contextlib.suppress(BaseException):
await produce_task
elif not produce_task.cancelled():
# The upstream may have answered while nobody was draining this
# relay, e.g. the client vanished first. Nothing else holds that
# response, so its stream only gets closed here.
with anyio.CancelScope(shield=True):
with contextlib.suppress(BaseException):
await _aclose_late_response(produce_task.result())
verbose_proxy_logger.info(
"no upstream response after %ss, opening the SSE response early and sending keepalives", interval
)
return StreamingResponse(
keepalive_then_relay(),
media_type=media_type,
headers=_TTFT_KEEPALIVE_HEADERS,
)
def _is_azure_model_router_request(model: str) -> bool:
"""
Check if the requested model is an Azure Model Router.
@ -1043,7 +1333,7 @@ def _log_llm_api_exception(e: Exception) -> None:
async def _cancel_llm_call_on_client_disconnect(
request: Request,
llm_api_call: "asyncio.Future[object]",
llm_api_call: "asyncio.Future[_LlmCallT]",
disconnect_event: asyncio.Event,
) -> None:
try:
@ -1062,8 +1352,8 @@ async def _cancel_llm_call_on_client_disconnect(
async def _await_llm_call_cancelling_on_disconnect(
request: Request,
llm_api_call: "asyncio.Future[Any]",
) -> Any:
llm_api_call: "asyncio.Future[_LlmCallT]",
) -> _LlmCallT:
disconnect_event: Final = asyncio.Event()
monitor: Final = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event))
try:
@ -1714,100 +2004,11 @@ class ProxyBaseLLMRequestProcessing:
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
route_type: Literal[
"acompletion",
"aembedding",
"aresponses",
"_arealtime",
"_aresponses_websocket",
"acreate_realtime_client_secret",
"arealtime_calls",
"aget_responses",
"adelete_responses",
"acancel_responses",
"acompact_responses",
"acreate_batch",
"aretrieve_batch",
"alist_batches",
"acancel_batch",
"afile_content",
"afile_retrieve",
"afile_delete",
"atext_completion",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
"aretrieve_fine_tuning_job",
"alist_input_items",
"aimage_edit",
"agenerate_content",
"agenerate_content_stream",
"allm_passthrough_route",
"avector_store_search",
"avector_store_create",
"avector_store_retrieve",
"avector_store_list",
"avector_store_update",
"avector_store_delete",
"avector_store_file_create",
"avector_store_file_list",
"avector_store_file_retrieve",
"avector_store_file_content",
"avector_store_file_update",
"avector_store_file_delete",
"aocr",
"asearch",
"avideo_generation",
"avideo_list",
"avideo_status",
"avideo_content",
"avideo_remix",
"avideo_create_character",
"avideo_get_character",
"avideo_edit",
"avideo_extension",
"acreate_container",
"alist_containers",
"aingest",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
"aretrieve_container_file_content",
"acreate_skill",
"alist_skills",
"aget_skill",
"adelete_skill",
"anthropic_messages",
"acreate_interaction",
"aget_interaction",
"adelete_interaction",
"acancel_interaction",
"acreate_agent",
"alist_agents",
"aget_agent",
"adelete_agent",
"alist_agent_versions",
"asend_message",
"call_mcp_tool",
"acreate_eval",
"alist_evals",
"aget_eval",
"aupdate_eval",
"adelete_eval",
"acancel_eval",
"acreate_run",
"alist_runs",
"aget_run",
"acancel_run",
"adelete_run",
],
route_type: ProxyRouteType,
proxy_logging_obj: ProxyLogging,
general_settings: dict,
general_settings: dict[str, object],
proxy_config: ProxyConfig,
select_data_generator: Callable | None = None,
select_data_generator: Callable[..., object] | None = None,
llm_router: Router | None = None,
model: str | None = None,
user_model: str | None = None,
@ -1817,7 +2018,72 @@ class ProxyBaseLLMRequestProcessing:
user_api_base: str | None = None,
version: str | None = None,
is_streaming_request: bool | None = False,
contents: list | None = None, # Add contents parameter
contents: list[object] | None = None,
skip_pre_call_logic: bool = False,
) -> Any:
"""Run the request, sending SSE keepalives while the upstream is still silent.
Everything below this point, the upstream call included, happens before the
proxy can write a byte, so a slow time-to-first-token leaves the response
idle. See ``open_sse_before_first_byte``; unwrapped unless an operator sets
``litellm_settings.sse_keepalive_ping_interval_seconds``.
"""
async def _audit_late_failure(exc: Exception) -> HTTPException | None:
# Once a keepalive is on the wire this can no longer raise, so the
# caller's `except` never runs its own post_call_failure_hook.
return await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=exc,
request_data=self.data,
)
return await open_sse_before_first_byte(
self._process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type=route_type,
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
llm_router=llm_router,
model=model,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
is_streaming_request=is_streaming_request,
contents=contents,
skip_pre_call_logic=skip_pre_call_logic,
),
ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router),
on_late_failure=_audit_late_failure,
)
async def _process_llm_request(
self,
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
route_type: ProxyRouteType,
proxy_logging_obj: ProxyLogging,
general_settings: dict[str, object],
proxy_config: ProxyConfig,
select_data_generator: Callable[..., object] | None = None,
llm_router: Router | None = None,
model: str | None = None,
user_model: str | None = None,
user_temperature: float | None = None,
user_request_timeout: float | None = None,
user_max_tokens: int | None = None,
user_api_base: str | None = None,
version: str | None = None,
is_streaming_request: bool | None = False,
contents: list[object] | None = None, # Add contents parameter
skip_pre_call_logic: bool = False,
) -> Any:
"""
@ -2039,6 +2305,7 @@ class ProxyBaseLLMRequestProcessing:
return StreamingResponse(
content=generator,
status_code=status.HTTP_200_OK,
media_type=self._passthrough_event_stream_media_type(),
headers=custom_headers,
)
else:
@ -2197,11 +2464,21 @@ class ProxyBaseLLMRequestProcessing:
additional_headers = hidden_params.get("additional_headers", {}) or {}
recover_response_cost: Final = not response_cost and hidden_params.get("response_cost") is None
response_cost_for_headers: Final = (
llm_cost_for_headers: Final = (
self._response_cost_from_logging_obj(response=response, logging_obj=logging_obj) or ""
if recover_response_cost
else response_cost
)
_, request_metadata_bucket = get_or_create_metadata_bucket(self.data)
guardrail_cost_for_headers: Final = guardrail_information_cost(
request_metadata_bucket.get("standard_logging_guardrail_information")
)
response_cost_for_headers: Final = (
(llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0)
+ guardrail_cost_for_headers
if guardrail_cost_for_headers > 0
else llm_cost_for_headers
)
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
@ -2494,10 +2771,16 @@ class ProxyBaseLLMRequestProcessing:
def _passthrough_event_stream_media_type(self) -> str | None:
"""
Content-type for a buffered passthrough event-stream response, resolved
from the provider handler so the proxy stays provider-agnostic. Mirrors
the upstream content-type the non-streaming path forwards, since the
buffered streaming generator carries no headers of its own.
Content-type for a passthrough event-stream response, resolved from the
provider handler so the proxy stays provider-agnostic. Mirrors the
upstream content-type the non-streaming path forwards, since the
streaming generator carries no headers of its own. Used for both the
buffered (guardrail-rewritten) and the unbuffered relay paths so
clients that enforce the event-stream content-type (e.g. Claude Code on
Bedrock invoke-with-response-stream) see the correct header instead of
no content-type at all, which they fall back to reading as
application/octet-stream. Returns None for providers with no
event-stream media type, leaving the response headers unchanged.
"""
from litellm.llms.pass_through.guardrail_translation.handler import (
LlmPassthroughRouteHandler,

View file

@ -39,10 +39,10 @@ _EXTRA_SENSITIVE_CALLBACK_KEYS: Final = {"gcs_path_service_account"}
# already-encrypted input cheaply (no decrypt-attempt round trip) and
# avoid double-encrypting if `LITELLM_SALT_KEY` is rotated between writes.
_CALLBACK_VAR_ENCRYPTED_PREFIX: Final = "litellm_enc::"
# Metadata slots that hold operator-configured callback setup (and therefore
# integration credentials). Resolved from UserAPIKeyAuth during pre-call setup,
# never read back off the copies stamped into request metadata.
_CALLBACK_CONFIG_SLOTS: Final = frozenset({"logging", "callback_settings"})
# Metadata slots that hold operator-configured callback and secret-manager setup
# (and therefore integration credentials). Resolved from UserAPIKeyAuth during
# pre-call setup, never read back off the copies stamped into request metadata.
_CALLBACK_CONFIG_SLOTS: Final = frozenset({"logging", "callback_settings", "secret_manager_settings"})
blue_color_code: Final = "\033[94m"
reset_color_code: Final = "\033[0m"

View file

@ -0,0 +1,226 @@
from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import date, datetime, timezone
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import litellm
from litellm._logging import verbose_logger
from litellm.types.proxy.model_deprecation import (
DEFAULT_DEPRECATION_WARN_DAYS,
DeprecationStatus,
ModelDeprecationInfo,
ModelDeprecationResponse,
)
if TYPE_CHECKING:
from litellm.router import Router
_NO_MODEL_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class _ResolvedDeprecation:
deprecation_date: date
litellm_model: str | None
litellm_provider: str | None
def _parse_deprecation_date(raw_value: object) -> date | None:
if isinstance(raw_value, datetime):
return raw_value.date()
if isinstance(raw_value, date):
return raw_value
if not isinstance(raw_value, str):
return None
try:
return date.fromisoformat(raw_value.strip())
except ValueError:
return None
def _cost_map_lookup(model_key: object) -> _ResolvedDeprecation | None:
if not isinstance(model_key, str) or not model_key:
return None
entry: Final = litellm.model_cost.get(model_key)
if not isinstance(entry, Mapping):
return None
parsed: Final = _parse_deprecation_date(entry.get("deprecation_date"))
if parsed is None:
return None
provider: Final = entry.get("litellm_provider")
return _ResolvedDeprecation(
deprecation_date=parsed,
litellm_model=model_key,
litellm_provider=provider if isinstance(provider, str) else None,
)
def _mapping_field(deployment: Mapping[str, object], key: str) -> Mapping[str, object]:
value: Final = deployment.get(key)
return value if isinstance(value, Mapping) else _NO_MODEL_METADATA
def _resolve_deployment_deprecation(
deployment: Mapping[str, object],
) -> _ResolvedDeprecation | None:
"""Resolve a deployment's deprecation date, preferring its explicit override"""
model_info: Final = _mapping_field(deployment, "model_info")
raw_model: Final = _mapping_field(deployment, "litellm_params").get("model")
override: Final = _parse_deprecation_date(model_info.get("deprecation_date"))
if override is not None:
provider: Final = model_info.get("litellm_provider")
return _ResolvedDeprecation(
deprecation_date=override,
litellm_model=raw_model if isinstance(raw_model, str) else None,
litellm_provider=provider if isinstance(provider, str) else None,
)
unprefixed: Final = raw_model.split("/", 1)[1] if isinstance(raw_model, str) and "/" in raw_model else None
return next(
(
resolved
for resolved in (
_cost_map_lookup(model_info.get("base_model")),
_cost_map_lookup(raw_model),
_cost_map_lookup(unprefixed),
)
if resolved is not None
),
None,
)
def _classify(days_until: int, warn_within_days: int) -> DeprecationStatus:
if days_until < 0:
return "deprecated"
if days_until <= warn_within_days:
return "imminent"
return "upcoming"
def _build_info(deployment: Mapping[str, object], today: date, warn_within_days: int) -> ModelDeprecationInfo | None:
model_name: Final = deployment.get("model_name")
if not isinstance(model_name, str) or not model_name:
return None
resolved: Final = _resolve_deployment_deprecation(deployment)
if resolved is None:
return None
days_until: Final = (resolved.deprecation_date - today).days
return ModelDeprecationInfo(
model_name=model_name,
litellm_model=resolved.litellm_model,
deprecation_date=resolved.deprecation_date,
days_until_deprecation=days_until,
status=_classify(days_until, warn_within_days),
litellm_provider=resolved.litellm_provider,
)
def _dedupe(
models: Sequence[ModelDeprecationInfo],
) -> tuple[ModelDeprecationInfo, ...]:
"""Report a model group carrying the same date on several deployments once"""
ordered: Final = sorted(models, key=lambda model: (model.model_name, model.deprecation_date))
return tuple(
next(group) for _, group in groupby(ordered, key=lambda model: (model.model_name, model.deprecation_date))
)
def _bucket(models: Sequence[ModelDeprecationInfo], status: DeprecationStatus) -> tuple[ModelDeprecationInfo, ...]:
return tuple(
sorted(
(model for model in models if model.status == status),
key=lambda model: model.deprecation_date,
)
)
def collect_model_deprecations(
llm_router: Router | None,
warn_within_days: int = DEFAULT_DEPRECATION_WARN_DAYS,
today: date | None = None,
) -> ModelDeprecationResponse:
"""Bucket every deployment carrying a deprecation date by how urgent it is"""
snapshot_time: Final = datetime.now(timezone.utc)
effective_today: Final = today or snapshot_time.date()
deployments: Final = (llm_router.get_model_list() or ()) if llm_router is not None else ()
deduped: Final = _dedupe(
tuple(
info
for info in (_build_info(deployment, effective_today, warn_within_days) for deployment in deployments)
if info is not None
)
)
verbose_logger.debug(
"model_deprecation: %d/%d deployments carry a deprecation date",
len(deduped),
len(deployments),
)
return ModelDeprecationResponse(
deprecated=_bucket(deduped, "deprecated"),
imminent=_bucket(deduped, "imminent"),
upcoming=_bucket(deduped, "upcoming"),
warn_within_days=warn_within_days,
checked_at=snapshot_time,
)
def _escape_slack_mrkdwn(value: str) -> str:
"""Neutralize Slack control characters so a model name cannot forge a mention or link"""
return value.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
def _format_entry(info: ModelDeprecationInfo) -> str:
suffix: Final = (
f"already deprecated {abs(info.days_until_deprecation)}d ago"
if info.days_until_deprecation < 0
else f"in {info.days_until_deprecation}d"
)
return (
f"• `{_escape_slack_mrkdwn(info.model_name)}` "
f"(provider: {_escape_slack_mrkdwn(info.litellm_provider) if info.litellm_provider else 'unknown'}, "
f"deprecates {info.deprecation_date.isoformat()}, {suffix})"
)
def format_deprecation_alert_message(
snapshot: ModelDeprecationResponse,
) -> str | None:
"""Render the alert for the deprecated and imminent buckets, None when both are empty
Upcoming models are left out of the alert to keep it actionable.
"""
if not snapshot.deprecated and not snapshot.imminent:
return None
deprecated_section: Final = (
("\n*Already deprecated:*", *(_format_entry(i) for i in snapshot.deprecated)) if snapshot.deprecated else ()
)
imminent_section: Final = (
(
f"\n*Deprecating within {snapshot.warn_within_days} days:*",
*(_format_entry(i) for i in snapshot.imminent),
)
if snapshot.imminent
else ()
)
return "\n".join(
(
"*⚠️ Model Deprecation Warning*",
*deprecated_section,
*imminent_section,
"\nPlan migrations to a supported model. See "
"https://docs.litellm.ai/docs/proxy/model_management for guidance.",
)
)

View file

@ -1,15 +1,23 @@
import asyncio
import contextlib
import math
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Iterable, Mapping
from typing import Final
import anyio
ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n'
SSE_COMMENT_PING: Final = ": ping\n\n"
SSE_COMMENT_PING_BYTES: Final = SSE_COMMENT_PING.encode()
# The byte form of proxy_server._SSE_FRAME_DELIMITERS, CR-only included: SSE
# terminates a line with CRLF, LF or CR, so a blank line is any of these three.
_SSE_FRAME_DELIMITERS: Final = (b"\r\n\r\n", b"\n\n", b"\r\r")
_SSE_DELIMITER_LOOKBACK: Final = max(len(delimiter) for delimiter in _SSE_FRAME_DELIMITERS)
_STREAM_START_TAIL: Final = b"\n\n"
_SSE_MEDIA_TYPE: Final = "text/event-stream"
def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None:
def coerce_keepalive_interval(ping_interval_seconds: float | str | None) -> float | None:
if ping_interval_seconds is None:
return None
try:
@ -28,23 +36,32 @@ def keepalive_ping_has_fired(elapsed_seconds: float, ping_interval_seconds: floa
the status line is already on the wire. With pings disabled nothing flushes early, so a raise
still carries its real status.
"""
interval: Final = _coerce_interval(ping_interval_seconds)
interval: Final = coerce_keepalive_interval(ping_interval_seconds)
return interval is not None and elapsed_seconds >= interval
def wrap_sse_stream_with_keepalive_pings(
stream: AsyncGenerator[str, None],
ping_interval_seconds: float | str | None,
ping_chunk: str = ANTHROPIC_PING_SSE_CHUNK,
) -> AsyncGenerator[str, None]:
interval: Final = _coerce_interval(ping_interval_seconds)
"""Fill idle gaps in an SSE stream, including the one before its first chunk.
``ping_chunk`` is what gets written into those gaps. It defaults to Anthropic's
own ``ping`` event because that is the protocol the first caller speaks; a
stream carrying anything else wants ``SSE_COMMENT_PING``, which is a comment
every conformant SSE client discards rather than a frame it has to understand.
"""
interval: Final = coerce_keepalive_interval(ping_interval_seconds)
if interval is None:
return stream
return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval)
return _keepalive_ping_stream(stream=stream, ping_interval_seconds=interval, ping_chunk=ping_chunk)
async def _keepalive_ping_stream(
stream: AsyncGenerator[str, None],
ping_interval_seconds: float,
ping_chunk: str,
) -> AsyncGenerator[str, None]:
pending = asyncio.ensure_future(
stream.__anext__()
@ -53,7 +70,7 @@ async def _keepalive_ping_stream(
while True:
await asyncio.wait({pending}, timeout=ping_interval_seconds)
if not pending.done():
yield ANTHROPIC_PING_SSE_CHUNK
yield ping_chunk
continue
try:
yield pending.result()
@ -66,3 +83,96 @@ async def _keepalive_ping_stream(
with contextlib.suppress(BaseException):
await pending
await stream.aclose()
def is_sse_content_type(content_type: str | None) -> bool:
return content_type is not None and content_type.split(";", 1)[0].strip().lower() == _SSE_MEDIA_TYPE
def wrap_passthrough_sse_bytes_with_keepalive_pings(
stream: AsyncGenerator[bytes, None],
ping_interval_seconds: float | str | None,
upstream_headers: Mapping[str, str],
) -> AsyncGenerator[bytes, None]:
"""Fill upstream silence on a byte-relaying passthrough stream with SSE comments.
Passthrough routes relay upstream bytes verbatim, so a model that thinks for
longer than an intermediary's idle read timeout has its connection dropped
before the first token. Only streams the upstream itself declares as
``text/event-stream`` are wrapped: a comment spliced into a binary transport
(AWS event streams on ``/bedrock``, protobuf, NDJSON) would corrupt it.
"""
interval: Final = coerce_keepalive_interval(ping_interval_seconds)
if interval is None or not is_sse_content_type(upstream_headers.get("content-type")):
return stream
return _keepalive_ping_byte_stream(stream=stream, ping_interval_seconds=interval)
async def _keepalive_ping_byte_stream(
stream: AsyncGenerator[bytes, None],
ping_interval_seconds: float,
) -> AsyncGenerator[bytes, None]:
pending = asyncio.ensure_future(
stream.__anext__()
) # rebind-ok: re-armed with the next __anext__ after each delivered chunk
# The tail of the bytes relayed so far, long enough to hold any delimiter.
# Seeded as a delimiter because a stream starts at a frame boundary, and kept
# across chunks because a delimiter can be split between two transport reads,
# which testing only the latest chunk would miss for the rest of the stream.
recent_tail = _STREAM_START_TAIL # rebind-ok: rolling window over the relayed bytes
try:
while True:
await asyncio.wait((pending,), timeout=ping_interval_seconds)
if not pending.done():
# The relayed chunks are raw transport reads, not whole SSE
# frames, so an upstream that stalls halfway through a frame
# must not have a comment spliced into it.
if recent_tail.endswith(_SSE_FRAME_DELIMITERS):
yield SSE_COMMENT_PING_BYTES
continue
try:
chunk: bytes = pending.result()
except StopAsyncIteration:
return
if chunk:
recent_tail = (recent_tail + chunk)[-_SSE_DELIMITER_LOOKBACK:]
yield chunk
pending = asyncio.ensure_future(stream.__anext__())
finally:
pending.cancel()
with anyio.CancelScope(shield=True):
with contextlib.suppress(BaseException):
await pending
await stream.aclose()
def resolve_ttft_keepalive_interval(
deployments: Iterable[Mapping[str, object]],
global_interval: float | str | None,
) -> float | None:
"""The keepalive interval to use before the upstream has answered at all.
No deployment has served the request yet, so a per-deployment
``keepalive_seconds`` is only trusted when every candidate under the requested
model carries the same one, which is how the mid-stream engine treats its own
model_name fallback. Otherwise the operator's global default applies.
An explicit ``0`` survives as a disable, since coercion rejects it: that keeps
an operator's documented hard disable working on this path too, rather than
letting the global switch a deployment back on behind their back.
A client-supplied value is deliberately not consulted. Opening the response
early is an operator decision, and a request must not be able to enable it for
a deployment that never did.
"""
configured: Final = frozenset(_keepalive_param(deployment) for deployment in deployments)
agreed: Final = next(iter(configured)) if len(configured) == 1 else None
return coerce_keepalive_interval(global_interval if agreed is None else agreed)
def _keepalive_param(deployment: Mapping[str, object]) -> float | str | None:
params: Final = deployment.get("litellm_params")
if not isinstance(params, Mapping):
return None
value: Final = params.get("keepalive_seconds")
return value if isinstance(value, (int, float, str)) else None

View file

@ -31,6 +31,7 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
@ -872,6 +873,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
credentials, aws_region_name = self._load_credentials()
allow_chunking: Final = not self._content_uses_contextual_grounding(content)
completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
try:
responses: Final = await self._apply_guardrail_content_with_chunking(
content=content,
@ -883,6 +885,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
completed_chunk_usages=completed_chunk_usages,
)
except HTTPException as exc:
if not isinstance(exc.detail, dict):
@ -891,6 +894,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
event_type=event_type,
start_time=start_time,
aws_region_name=aws_region_name,
completed_chunk_usages=completed_chunk_usages,
)
raise
merged_response: Final = self._merge_bedrock_guardrail_responses(responses)
@ -899,6 +904,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
event_type=event_type,
start_time=start_time,
aws_region_name=aws_region_name,
)
return merged_response
@ -913,6 +919,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
event_type: GuardrailEventHooks,
start_time: "datetime",
allow_chunking: bool,
completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: billed-chunk usage accumulator
) -> tuple[BedrockContentChunkResult, ...]:
"""Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large.
@ -959,6 +966,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
event_type=event_type,
start_time=start_time,
completed_chunk_usages=completed_chunk_usages,
)
return (
BedrockContentChunkResult(
@ -989,6 +997,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
completed_chunk_usages=completed_chunk_usages,
)
for batch in batches
]
@ -1015,6 +1024,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
completed_chunk_usages=completed_chunk_usages,
)
second_results: Final = await self._apply_guardrail_content_with_chunking(
content=second_half,
@ -1026,6 +1036,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
event_type=event_type,
start_time=start_time,
allow_chunking=allow_chunking,
completed_chunk_usages=completed_chunk_usages,
)
combined_results: Final = tuple(first_results) + tuple(second_results)
if is_single_item_text_split:
@ -1045,6 +1056,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: passed through to the single-call layer
) -> BedrockGuardrailResponse:
"""Post one ApplyGuardrail call for `content`, retrying with exponential
backoff on AWS ThrottlingException (HTTP 429).
@ -1072,6 +1084,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
event_type=event_type,
start_time=start_time,
completed_chunk_usages=completed_chunk_usages,
)
except HTTPException as exc:
if (
@ -1093,6 +1106,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: billed-chunk usage accumulator
) -> BedrockGuardrailResponse:
"""Make exactly one signed ApplyGuardrail HTTP call for `content` and
parse the result. Raises HTTPException on a guardrail block or any
@ -1108,7 +1122,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
A block is logged here rather than by the caller: it ends the whole chunking
flow immediately, with no further chunks attempted, so there is no later
merged response for the caller to log instead.
merged response for the caller to log instead. The logged usage still spans
the whole logical request: chunks that passed before the block appended what
AWS billed them to ``completed_chunk_usages``, and the attempt log sums those
with the blocking call's own usage.
"""
bedrock_request_data: Final = { # mutable-ok: outbound JSON request body
**base_request_data,
@ -1151,10 +1168,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data=request_data,
event_type=event_type,
start_time=start_time,
aws_region_name=aws_region_name,
completed_chunk_usages=completed_chunk_usages,
)
raise self._get_http_exception_for_blocked_guardrail(
bedrock_guardrail_response, request_data=request_data
)
response_usage: Final = bedrock_guardrail_response.get("usage")
if isinstance(response_usage, dict):
completed_chunk_usages.append(
response_usage
) # rebind-ok: accumulator threaded from make_bedrock_api_request, recording this billed call
return bedrock_guardrail_response
status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
@ -1172,14 +1196,31 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
aws_region_name: str | None,
completed_chunk_usages: Sequence[BedrockGuardrailUsage],
) -> None:
"""Log a single ApplyGuardrail HTTP attempt as-is (its own status,
derived from its own response). Used only for the blocked-content
case, which ends the whole chunking flow immediately."""
tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response))
"""Log the blocking ApplyGuardrail attempt, which ends the whole chunking
flow immediately. Its status derives from its own response, but its usage
(and so its cost) spans every billed call of the logical request: the
chunks that passed before the block plus the blocking call itself."""
blocking_usage: Final = json_response.get("usage")
billed_usages: Final[tuple[BedrockGuardrailUsage, ...]] = tuple(completed_chunk_usages) + (
(blocking_usage,) if isinstance(blocking_usage, dict) else ()
)
logged_json_response: Final = (
{ # mutable-ok: raw AWS JSON payload carrying the total billed usage
**json_response,
"usage": self._sum_usage_counters(billed_usages),
}
if completed_chunk_usages
else json_response
)
tracing_detail: Final = self._build_tracing_detail(
BedrockGuardrailResponse(**logged_json_response), aws_region_name=aws_region_name
)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=json_response,
guardrail_json_response=logged_json_response,
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
start_time=start_time.timestamp(),
@ -1195,6 +1236,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
aws_region_name: str | None,
) -> None:
"""Log one logical ApplyGuardrail call -- possibly several chunk calls
under the hood -- using its final merged response, so a chunked
@ -1205,7 +1247,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
``Output.__type`` with an exception marker. That marker survives the merge,
so the status is derived from the merged response rather than assumed to be
a success, which is what the pre-chunking code reported for that shape."""
tracing_detail: Final = self._build_tracing_detail(merged_response)
tracing_detail: Final = self._build_tracing_detail(merged_response, aws_region_name=aws_region_name)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict
@ -1228,20 +1270,36 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
event_type: GuardrailEventHooks,
start_time: "datetime",
aws_region_name: str | None,
completed_chunk_usages: Sequence[BedrockGuardrailUsage],
) -> None:
"""Log one logical ApplyGuardrail call that failed end-to-end (an
unrecoverable too-large error, a non-size validation error, or
exhausted throttle retries) as a single failure, rather than logging
every failed attempt chunking made along the way."""
every failed attempt chunking made along the way. Chunk calls AWS
billed before the failure still carry their usage and cost."""
billed_usage: Final = self._sum_usage_counters(completed_chunk_usages) if completed_chunk_usages else None
error_payload: Final = {"error": str(detail)} # mutable-ok: logging helper requires a dict
json_response: Final = (
{**error_payload, "usage": billed_usage} # mutable-ok: logging helper requires a dict
if billed_usage is not None
else error_payload
)
tracing_detail: Final = (
self._build_tracing_detail(BedrockGuardrailResponse(usage=billed_usage), aws_region_name=aws_region_name)
if billed_usage is not None
else None
)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict
guardrail_json_response=json_response,
request_data=request_data or {}, # mutable-ok: logging helper requires a dict
guardrail_status="guardrail_failed_to_respond",
start_time=start_time.timestamp(),
end_time=datetime.now(timezone.utc).timestamp(),
duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
event_type=event_type,
tracing_detail=tracing_detail or None,
)
@staticmethod
@ -1504,15 +1562,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
Keys are taken from the responses rather than from a fixed list, so a counter
this code does not know about (AWS has added several) is still summed and
reported instead of being silently dropped to zero."""
chunk_usages: Final = tuple(
chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback
for chunk_result in chunk_results
return BedrockGuardrail._sum_usage_counters(
tuple(
chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback
for chunk_result in chunk_results
)
)
@staticmethod
def _sum_usage_counters(usages: Sequence[BedrockGuardrailUsage]) -> BedrockGuardrailUsage:
return cast( # cast-ok: TypedDict assembled from a comprehension
BedrockGuardrailUsage,
{ # mutable-ok: builds the TypedDict payload
key: sum(usage.get(key) or 0 for usage in chunk_usages)
for key in dict.fromkeys(key for usage in chunk_usages for key in usage)
key: sum(usage.get(key) or 0 for usage in usages)
for key in dict.fromkeys(key for usage in usages for key in usage)
},
)
@ -2036,7 +2099,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return (status_code, err)
return (status_code, message)
def _build_tracing_detail(self, response: BedrockGuardrailResponse) -> GuardrailTracingDetail:
def _build_tracing_detail(
self, response: BedrockGuardrailResponse, aws_region_name: str | None
) -> GuardrailTracingDetail:
"""
Build the tracing detail from the raw Bedrock response, before
redaction, so downstream loggers (OTEL, Langfuse, ...) get the
@ -2053,6 +2118,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
bedrock_action: Final = response.get("action")
if isinstance(bedrock_action, str):
tracing_detail["guardrail_action"] = bedrock_action
usage: Final = response.get("usage")
if isinstance(usage, dict):
usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream
key: value for key, value in usage.items() if isinstance(value, int)
}
if usage_units:
tracing_detail["guardrail_usage"] = usage_units
tracing_detail["guardrail_cost"] = bedrock_guardrail_cost(
usage_units=usage_units, aws_region_name=aws_region_name
)
return tracing_detail
def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> list[str]:

View file

@ -36,6 +36,13 @@ _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail"
_INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls")
_DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname
_KEYS_DUPLICATING_SCAN_INPUTS: Final = ("messages", "input")
_LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + (
"additional_args",
"standard_logging_object",
"original_response",
)
class _Action(str, enum.Enum):
BLOCKED = "BLOCKED"
@ -131,9 +138,20 @@ class NomaV2Guardrail(CustomGuardrail):
logging_obj: Optional["LiteLLMLoggingObj"],
application_id: str | None,
) -> dict:
payload_request_data: Final = self._sanitize_payload_for_transport(request_data)
payload_request_data: Final = self._sanitize_payload_for_transport(
{key: value for key, value in request_data.items() if key not in _KEYS_DUPLICATING_SCAN_INPUTS}
)
if logging_obj is not None:
payload_request_data["litellm_logging_obj"] = getattr(logging_obj, "model_call_details", None)
model_call_details: Final = getattr(logging_obj, "model_call_details", None)
payload_request_data["litellm_logging_obj"] = (
{
key: value
for key, value in model_call_details.items()
if key not in _LOGGING_KEYS_DUPLICATING_SCAN_INPUTS
}
if isinstance(model_call_details, dict)
else model_call_details
)
payload: Final[dict[str, Any]] = {
"inputs": inputs,

View file

@ -4,18 +4,22 @@ GET /guardrails/usage/overview, /guardrails/usage/detail/:id, /guardrails/usage/
"""
import json
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import date, datetime, timedelta, timezone
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, overload
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
DailyGuardrailUsageUnitsRepository,
DailyPolicyMetricsRepository,
GuardrailsRepository,
PolicyRepository,
@ -28,6 +32,7 @@ if TYPE_CHECKING:
from prisma import types as prisma_types
from prisma.actions import (
LiteLLM_DailyGuardrailMetricsActions,
LiteLLM_DailyGuardrailUsageUnitsActions,
LiteLLM_DailyPolicyMetricsActions,
LiteLLM_GuardrailsTableActions,
LiteLLM_PolicyTableActions,
@ -41,6 +46,42 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
_USAGE_MAX_RANGE_DAYS: Final = 366
def _resolve_usage_window(start_date: str | None, end_date: str | None) -> tuple[str, str]:
from fastapi import HTTPException, status
now: Final = datetime.now(timezone.utc)
end: Final = end_date or now.strftime("%Y-%m-%d")
start: Final = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
try:
parsed: Final = (date.fromisoformat(start), date.fromisoformat(end))
except ValueError:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="start_date and end_date must be in YYYY-MM-DD format",
)
start_obj, end_obj = parsed
if (start_obj.isoformat(), end_obj.isoformat()) != (start, end):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="start_date and end_date must be in YYYY-MM-DD format",
)
if end_obj < start_obj:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="start_date must be on or before end_date",
)
if end_obj - start_obj > timedelta(days=_USAGE_MAX_RANGE_DAYS):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Date range too large; maximum is {_USAGE_MAX_RANGE_DAYS} days",
)
return start, end
def _guardrails_table(
prisma_client: "PrismaClient",
@ -92,6 +133,50 @@ async def _find_daily_policy_metrics(
return await _daily_policy_metrics_table(prisma_client).find_many(where=where)
def _daily_guardrail_usage_units_table(
prisma_client: "PrismaClient",
) -> "LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
units_table: Final[LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]] = (
DailyGuardrailUsageUnitsRepository(prisma_client).table
)
return units_table
async def _find_daily_guardrail_usage_units(
prisma_client: "PrismaClient",
where: "prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput",
) -> "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
from prisma.errors import TableNotFoundError
try:
return await _daily_guardrail_usage_units_table(prisma_client).find_many(where=where)
except TableNotFoundError as e:
verbose_proxy_logger.warning(
"Guardrail usage units are unavailable until the LiteLLM_DailyGuardrailUsageUnits migration is applied: %s",
e,
)
return ()
def _counter_name(row: "prisma_models.LiteLLM_DailyGuardrailUsageUnits") -> str:
return row.usage_unit
def _sum_counter_units(rows: "Iterable[prisma_models.LiteLLM_DailyGuardrailUsageUnits]") -> Mapping[str, int]:
ordered: Final = sorted(rows, key=_counter_name)
return MappingProxyType(
{name: sum(int(r.units) for r in group) for name, group in groupby(ordered, key=_counter_name)}
)
def _units_by(
rows: "Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]",
key_of: "Callable[[prisma_models.LiteLLM_DailyGuardrailUsageUnits], str]",
) -> Mapping[str, Mapping[str, int]]:
ordered: Final = sorted(rows, key=key_of)
return MappingProxyType({key: _sum_counter_units(group) for key, group in groupby(ordered, key=key_of)})
# --- Response models ---
@ -140,6 +225,7 @@ class UsageOverviewRow(BaseModel):
avgLatency: float | None
status: str # healthy | warning | critical
trend: str # up | down | stable
usageUnits: Mapping[str, int]
class UsageOverviewResponse(BaseModel):
@ -148,6 +234,12 @@ class UsageOverviewResponse(BaseModel):
totalRequests: int
totalBlocked: int
passRate: float
totalUsageUnits: Mapping[str, int]
class UsageUnitsDailyPoint(BaseModel):
date: str
units: Mapping[str, int]
class UsageDetailResponse(BaseModel):
@ -163,6 +255,10 @@ class UsageDetailResponse(BaseModel):
trend: str
description: str | None
time_series: list[UsageChartPoint]
usage_units: Mapping[str, int]
usage_units_daily: Sequence[UsageUnitsDailyPoint]
usage_units_by_team: Mapping[str, Mapping[str, int]]
usage_units_by_key: Mapping[str, Mapping[str, int]]
class UsageLogEntry(BaseModel):
@ -278,6 +374,7 @@ def _guardrail_overview_rows(
guardrails: "Sequence[_DbOrConfigGuardrail]",
agg: Mapping[str, _MetricTotals],
prev_agg: Mapping[str, float],
units_agg: Mapping[str, Mapping[str, int]],
) -> list[UsageOverviewRow]:
rows: Final[list[UsageOverviewRow]] = []
covered_keys: Final[set[str]] = set()
@ -303,6 +400,7 @@ def _guardrail_overview_rows(
prev_fail = float(prev_agg.get(k, 0.0) or 0.0)
break
trend = _trend_from_comparison(fail_rate, prev_fail)
row_units: Mapping[str, int] = next((units_agg[k] for k in lookup_keys if k in units_agg), _EMPTY_UNITS)
rows.append(
UsageOverviewRow(
id=gid,
@ -315,6 +413,7 @@ def _guardrail_overview_rows(
avgLatency=None,
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=row_units,
)
)
# Add rows for guardrails with metrics but not in guardrails table (e.g. MCP, config)
@ -337,6 +436,7 @@ def _guardrail_overview_rows(
avgLatency=None,
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=units_agg.get(agg_key, _EMPTY_UNITS),
)
)
return rows
@ -366,6 +466,7 @@ def _policy_overview_rows(
avgLatency=None,
status=_status_from_fail_rate(fail_rate),
trend=trend,
usageUnits=_EMPTY_UNITS,
)
)
return rows
@ -386,11 +487,11 @@ async def guardrails_usage_overview(
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return UsageOverviewResponse(rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0)
return UsageOverviewResponse(
rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS
)
now: Final = datetime.now(timezone.utc)
end: Final = end_date or now.strftime("%Y-%m-%d")
start: Final = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
start, end = _resolve_usage_window(start_date, end_date)
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
@ -408,24 +509,33 @@ async def guardrails_usage_overview(
)
# Previous period for trend
start_prev: Final = (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d")
start_prev: Final = (date.fromisoformat(start) - timedelta(days=7)).isoformat()
metrics_prev: Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics] = await _find_daily_guardrail_metrics(
prisma_client, where={"date": {"gte": start_prev, "lt": start}}
)
units_where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput] = {
"date": {"gte": start, "lte": end}
}
units_rows: Final[
Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]
] = await _find_daily_guardrail_usage_units(prisma_client, where=units_where)
agg: Final = _aggregate_daily_metrics(metrics, "guardrail_id")
prev_agg: Final = _prev_fail_rates(metrics_prev, "guardrail_id")
units_agg: Final = _units_by(units_rows, lambda r: r.guardrail_id)
chart: Final = _chart_from_metrics(metrics)
total_requests: Final = sum(a["requests"] for a in agg.values())
total_blocked: Final = sum(a["blocked"] for a in agg.values())
pass_rate: Final = (100.0 * (total_requests - total_blocked) / total_requests) if total_requests else 100.0
rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg)
rows: Final = _guardrail_overview_rows(guardrails, agg, prev_agg, units_agg)
return UsageOverviewResponse(
rows=rows,
chart=chart,
totalRequests=total_requests,
totalBlocked=total_blocked,
passRate=round(pass_rate, 1),
totalUsageUnits=_sum_counter_units(units_rows),
)
except Exception as e:
from litellm.proxy.utils import handle_exception_on_proxy
@ -453,9 +563,7 @@ async def guardrails_usage_detail(
raise HTTPException(status_code=500, detail="Prisma client not initialized")
now: Final = datetime.now(timezone.utc)
end: Final = end_date or now.strftime("%Y-%m-%d")
start: Final = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
start, end = _resolve_usage_window(start_date, end_date)
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
@ -478,13 +586,21 @@ async def guardrails_usage_detail(
"date": {"gte": start, "lte": end},
},
)
start_prev: Final = (date.fromisoformat(start) - timedelta(days=7)).isoformat()
metrics_prev: Final[Sequence[prisma_models.LiteLLM_DailyGuardrailMetrics]] = await _find_daily_guardrail_metrics(
prisma_client,
where={
"guardrail_id": {"in": metric_ids},
"date": {"lt": start},
"date": {"gte": start_prev, "lt": start},
},
)
units_where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereInput] = {
"guardrail_id": {"in": metric_ids},
"date": {"gte": start, "lte": end},
}
units_rows: Final[
Sequence[prisma_models.LiteLLM_DailyGuardrailUsageUnits]
] = await _find_daily_guardrail_usage_units(prisma_client, where=units_where)
requests: Final = sum(int(m.requests_evaluated or 0) for m in metrics)
blocked: Final = sum(int(m.blocked_count or 0) for m in metrics)
@ -510,6 +626,8 @@ async def guardrails_usage_detail(
litellm_params: Final = _to_dict(_get_guardrail_field(guardrail, "litellm_params"))
guardrail_info: Final = _to_dict(_get_guardrail_field(guardrail, "guardrail_info"))
_guardrail_name: Final = _get_guardrail_field(guardrail, "guardrail_name")
daily_unit_sums: Final = sorted(_units_by(units_rows, lambda r: r.date).items())
units_daily: Final = tuple(UsageUnitsDailyPoint(date=d, units=units) for d, units in daily_unit_sums)
return UsageDetailResponse(
guardrail_id=guardrail_id,
@ -524,6 +642,10 @@ async def guardrails_usage_detail(
trend=trend,
description=guardrail_info.get("description"),
time_series=time_series,
usage_units=_sum_counter_units(units_rows),
usage_units_daily=units_daily,
usage_units_by_team=_units_by(units_rows, lambda r: r.team_id),
usage_units_by_key=_units_by(units_rows, lambda r: r.api_key),
)
@ -743,11 +865,11 @@ async def policies_usage_overview(
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return UsageOverviewResponse(rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0)
return UsageOverviewResponse(
rows=[], chart=[], totalRequests=0, totalBlocked=0, passRate=100.0, totalUsageUnits=_EMPTY_UNITS
)
now: Final = datetime.now(timezone.utc)
end: Final = end_date or now.strftime("%Y-%m-%d")
start: Final = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d")
start, end = _resolve_usage_window(start_date, end_date)
try:
policies: Final = await _policies_table(prisma_client).find_many()
@ -758,7 +880,7 @@ async def policies_usage_overview(
prisma_client,
where={
"date": {
"gte": (datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7)).strftime("%Y-%m-%d"),
"gte": (date.fromisoformat(start) - timedelta(days=7)).isoformat(),
"lt": start,
}
},
@ -776,6 +898,7 @@ async def policies_usage_overview(
totalRequests=total_requests,
totalBlocked=total_blocked,
passRate=round(pass_rate, 1),
totalUsageUnits=_EMPTY_UNITS,
)
except Exception as e:
from litellm.proxy.utils import handle_exception_on_proxy

View file

@ -3,18 +3,148 @@ Track guardrail and policy usage for the dashboard: upsert daily metrics and
insert into SpendLogGuardrailIndex when spend logs are written.
"""
import asyncio
import json
from collections import defaultdict
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final
from functools import partial
from itertools import groupby
from operator import itemgetter
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
DailyGuardrailUsageUnitsRepository,
SpendLogGuardrailIndexRepository,
)
if TYPE_CHECKING:
from prisma import types as prisma_types
_UPSERT_RETRY_TIMES: Final = 3
_MAX_PENDING_ROWS: Final = 10_000
_RowKey = TypeVar("_RowKey")
_RowValue = TypeVar("_RowValue")
class _UsageUnitKey(NamedTuple):
guardrail_id: str
date: str
team_id: str
api_key: str
usage_unit: str
class _MetricsKey(NamedTuple):
guardrail_id: str
date: str
class PendingRollups:
"""Rollup rows whose connection-error retries exhausted, held for the next flush."""
def __init__(self) -> None:
self.lock: Final = asyncio.Lock()
self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({})
self.units: Mapping[_UsageUnitKey, int] = MappingProxyType({})
_PENDING_ROLLUPS: Final = PendingRollups()
_NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({})
def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]:
return (*base, *(key for key in extra if key not in base))
def _merged_unit_rows(
base: Mapping[_UsageUnitKey, int], extra: Mapping[_UsageUnitKey, int]
) -> Mapping[_UsageUnitKey, int]:
return MappingProxyType({key: base.get(key, 0) + extra.get(key, 0) for key in _merged_keys(base, extra)})
def _merged_metric_rows(
base: Mapping[_MetricsKey, Mapping[str, int]], extra: Mapping[_MetricsKey, Mapping[str, int]]
) -> Mapping[_MetricsKey, Mapping[str, int]]:
def merged_counters(key: _MetricsKey) -> Mapping[str, int]:
base_counters: Final = base.get(key, _NO_COUNTERS)
extra_counters: Final = extra.get(key, _NO_COUNTERS)
return MappingProxyType(
{
counter: int(base_counters.get(counter, 0)) + int(extra_counters.get(counter, 0))
for counter in _merged_keys(base_counters, extra_counters)
}
)
return MappingProxyType({key: merged_counters(key) for key in _merged_keys(base, extra)})
def _capped(rows: Mapping[_RowKey, _RowValue], label: str) -> Mapping[_RowKey, _RowValue]:
if len(rows) <= _MAX_PENDING_ROWS:
return rows
verbose_proxy_logger.warning(
"Guardrail usage tracking: pending %s requeue exceeds %d rows; dropping the %d oldest (non-fatal)",
label,
_MAX_PENDING_ROWS,
len(rows) - _MAX_PENDING_ROWS,
)
return MappingProxyType(dict(tuple(rows.items())[len(rows) - _MAX_PENDING_ROWS :]))
async def _attempt_upsert(
upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], key: _RowKey, value: _RowValue
) -> Exception | None:
try:
await upsert_row(key, value)
except Exception as error:
return error
return None
async def _upsert_rows_with_retry(
rows: Mapping[_RowKey, _RowValue],
upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]],
label: str,
sleep: Callable[[float], Awaitable[None]],
retries_left: int = _UPSERT_RETRY_TIMES,
) -> Mapping[_RowKey, _RowValue]:
"""Returns the rows still failing with connection errors once retries exhaust, for requeueing."""
outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()}
for key, error in outcomes.items():
if error is not None and not isinstance(error, DB_RETRY_SAFE_ERROR_TYPES):
verbose_proxy_logger.warning(
"Guardrail usage tracking: %s upsert failed for %s and is not safe to retry (non-fatal): %s",
label,
key,
error,
)
retryable: Final = MappingProxyType(
{key: rows[key] for key, error in outcomes.items() if isinstance(error, DB_RETRY_SAFE_ERROR_TYPES)}
)
if not retryable:
return MappingProxyType({})
if retries_left == 0:
for key in retryable:
verbose_proxy_logger.warning(
"Guardrail usage tracking: %s upsert failed for %s after %d retries; requeued for the next flush "
"(non-fatal): %s",
label,
key,
_UPSERT_RETRY_TIMES,
outcomes[key],
)
return retryable
await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left))
return await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1)
def _guardrail_status_to_action(status: str | None) -> str:
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
@ -28,7 +158,7 @@ def _guardrail_status_to_action(status: str | None) -> str:
return "passed"
def _parse_guardrail_info_from_payload(payload: dict[str, Any]) -> list[dict[str, Any]]:
def _parse_guardrail_info_from_payload(payload: Mapping[str, Any]) -> Sequence[Mapping[str, Any]]:
"""Extract guardrail_information from spend log payload metadata."""
meta = payload.get("metadata")
if not meta:
@ -53,9 +183,96 @@ def _date_str(dt: datetime) -> str:
return dt.astimezone(timezone.utc).strftime("%Y-%m-%d")
def _parse_payload_start_time(payload: Mapping[str, Any]) -> datetime | None:
start_time: Final = payload.get("startTime")
if isinstance(start_time, datetime):
return start_time
if not isinstance(start_time, str):
return None
try:
return datetime.fromisoformat(start_time.replace("Z", "+00:00"))
except (ValueError, TypeError):
return None
def _iter_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Iterator[tuple[_UsageUnitKey, int]]:
for payload in logs_to_process:
start_time = _parse_payload_start_time(payload)
if not payload.get("request_id") or start_time is None:
continue
date_key = _date_str(start_time)
team_id = str(payload.get("team_id") or "")
api_key = str(payload.get("api_key") or "")
for entry in _parse_guardrail_info_from_payload(payload):
guardrail_id = str(entry.get("guardrail_id") or entry.get("guardrail_name") or "")
usage = entry.get("guardrail_usage")
if not guardrail_id or not isinstance(usage, dict):
continue
for unit_name, units in usage.items():
if isinstance(units, int) and not isinstance(units, bool) and units > 0:
yield _UsageUnitKey(guardrail_id, date_key, team_id, api_key, str(unit_name)), units
def _sum_usage_unit_increments(logs_to_process: Sequence[Mapping[str, Any]]) -> Mapping[_UsageUnitKey, int]:
ordered: Final = sorted(_iter_usage_unit_increments(logs_to_process), key=itemgetter(0))
return MappingProxyType(
{key: sum(units for _, units in group) for key, group in groupby(ordered, key=itemgetter(0))}
)
async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey, units: int) -> None:
row: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsCreateInput] = {
"guardrail_id": key.guardrail_id,
"date": key.date,
"team_id": key.team_id,
"api_key": key.api_key,
"usage_unit": key.usage_unit,
"units": units,
}
where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereUniqueInput] = {
"guardrail_id_date_team_id_api_key_usage_unit": {
"guardrail_id": key.guardrail_id,
"date": key.date,
"team_id": key.team_id,
"api_key": key.api_key,
"usage_unit": key.usage_unit,
}
}
data: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsUpsertInput] = {
"create": row,
"update": {"units": {"increment": units}},
}
await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data)
async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None:
n: Final = int(agg["requests_evaluated"])
await DailyGuardrailMetricsRepository(prisma_client).table.upsert(
where={"guardrail_id_date": {"guardrail_id": key.guardrail_id, "date": key.date}},
data={
"create": {
"guardrail_id": key.guardrail_id,
"date": key.date,
"requests_evaluated": n,
"passed_count": int(agg["passed_count"]),
"blocked_count": int(agg["blocked_count"]),
"flagged_count": int(agg["flagged_count"]),
},
"update": {
"requests_evaluated": {"increment": n},
"passed_count": {"increment": int(agg["passed_count"])},
"blocked_count": {"increment": int(agg["blocked_count"])},
"flagged_count": {"increment": int(agg["flagged_count"])},
},
},
)
async def process_spend_logs_guardrail_usage(
prisma_client: PrismaClient,
logs_to_process: list[dict[str, Any]],
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
pending: PendingRollups = _PENDING_ROLLUPS,
) -> None:
"""
After spend logs are written: update DailyGuardrailMetrics and insert
@ -64,7 +281,7 @@ async def process_spend_logs_guardrail_usage(
if not logs_to_process:
return
# Aggregate daily metrics by (guardrail_id, date). Latency/score metrics dropped.
daily_guardrail: Final[dict[tuple, dict[str, Any]]] = defaultdict(
daily_guardrail: Final[dict[_MetricsKey, dict[str, Any]]] = defaultdict(
lambda: {
"requests_evaluated": 0,
"passed_count": 0,
@ -76,21 +293,16 @@ async def process_spend_logs_guardrail_usage(
for payload in logs_to_process:
request_id = payload.get("request_id")
start_time = payload.get("startTime")
if not request_id or not start_time:
start_time = _parse_payload_start_time(payload)
if not request_id or start_time is None:
continue
if isinstance(start_time, str):
try:
start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00"))
except (ValueError, TypeError):
continue
date_key = _date_str(start_time)
for entry in _parse_guardrail_info_from_payload(payload):
guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or ""
if not guardrail_id:
continue
key = (guardrail_id, date_key)
key = _MetricsKey(guardrail_id, date_key)
daily_guardrail[key]["requests_evaluated"] += 1
action = _guardrail_status_to_action(entry.get("guardrail_status"))
if action == "passed":
@ -109,64 +321,42 @@ async def process_spend_logs_guardrail_usage(
}
)
if not daily_guardrail and not index_rows:
async with pending.lock:
pending_metrics: Final = pending.metrics
pending_units: Final = pending.units
pending.metrics = MappingProxyType({})
pending.units = MappingProxyType({})
# Upsert daily guardrail metrics (counts only; latency/score dropped)
evaluated_metrics: Final = MappingProxyType(
{key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0}
)
metrics_rows: Final = _merged_metric_rows(pending_metrics, evaluated_metrics)
unit_rows: Final = _merged_unit_rows(pending_units, _sum_usage_unit_increments(logs_to_process))
if not metrics_rows and not index_rows and not unit_rows:
return
try:
# Insert index rows (skip duplicates by request_id + guardrail_id)
if index_rows:
index_data: Final = []
for r in index_rows:
st = r["start_time"]
if isinstance(st, str):
try:
st = datetime.fromisoformat(st.replace("Z", "+00:00"))
except (ValueError, TypeError):
continue
index_data.append(
{
"request_id": r["request_id"],
"guardrail_id": r["guardrail_id"],
"policy_id": r.get("policy_id"),
"start_time": st,
}
)
try:
await SpendLogGuardrailIndexRepository(prisma_client).table.create_many(
data=index_data,
data=index_rows,
skip_duplicates=True,
)
except Exception as e:
verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e)
# Upsert daily guardrail metrics (counts only; latency/score dropped)
for (guardrail_id, date_key), agg in daily_guardrail.items():
n = int(agg["requests_evaluated"])
if n == 0:
continue
await DailyGuardrailMetricsRepository(prisma_client).table.upsert(
where={
"guardrail_id_date": {
"guardrail_id": guardrail_id,
"date": date_key,
}
},
data={
"create": {
"guardrail_id": guardrail_id,
"date": date_key,
"requests_evaluated": n,
"passed_count": int(agg["passed_count"]),
"blocked_count": int(agg["blocked_count"]),
"flagged_count": int(agg["flagged_count"]),
},
"update": {
"requests_evaluated": {"increment": n},
"passed_count": {"increment": int(agg["passed_count"])},
"blocked_count": {"increment": int(agg["blocked_count"])},
"flagged_count": {"increment": int(agg["flagged_count"])},
},
},
)
failed_metrics: Final = await _upsert_rows_with_retry(
metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep
)
failed_units: Final = await _upsert_rows_with_retry(
unit_rows, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep
)
if failed_metrics or failed_units:
async with pending.lock:
pending.metrics = _capped(_merged_metric_rows(pending.metrics, failed_metrics), "daily metrics")
pending.units = _capped(_merged_unit_rows(pending.units, failed_units), "usage unit")
except Exception as e:
verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e)

View file

@ -18,11 +18,12 @@ Quick summary:
"""
import json
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn
from collections.abc import Iterable, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel
from pydantic import BaseModel, Field, TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -40,10 +41,15 @@ from litellm.proxy._types import (
SpecialModelNames,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata
from litellm.proxy.common_utils.proxy_rate_limit_error import (
ProxyRateLimitError,
map_v3_rate_limit_type,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROJECT_ITPM_DESCRIPTOR_KEY,
PROJECT_OTPM_DESCRIPTOR_KEY,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
if TYPE_CHECKING:
@ -76,6 +82,11 @@ else:
RateLimitDescriptor = dict[str, Any]
_BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object])
IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int]
class BatchFileUsage(BaseModel):
"""
Internal model for batch file usage tracking, used for batch rate limiting
@ -83,6 +94,16 @@ class BatchFileUsage(BaseModel):
total_tokens: int
request_count: int
output_tokens: int = 0
# Keyed by each row's own `body.model`, distinct from `total_tokens`/
# `output_tokens` (the whole-file totals charged to the file-bound/
# top-level routing model's key/team/model limits). A batch's rows can
# each target a different model, so the project's per-model ITPM/OTPM
# quota for a row's actual model must be charged with that row's own
# tokens -- see `_create_project_io_descriptors_for_models`.
per_model_usage: dict[str, dict[str, int]] = Field(
default_factory=dict
) # mutable-ok: accumulated incrementally per row while parsing the batch file
class _PROXY_BatchRateLimiter(CustomLogger):
@ -198,6 +219,15 @@ class _PROXY_BatchRateLimiter(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
data: dict,
) -> list["RateLimitDescriptor"]:
"""Build the standard key/user/team/model descriptor list a batch is charged against.
Deliberately excludes the project-scoped ITPM/OTPM descriptors: those
are charged per the JSONL row's own `body.model` once the file is
parsed (`_create_project_io_descriptors_for_models`), not the
file-bound/top-level routing model this function resolves. Charging
project quotas here would let a caller bind the file to a model
without a quota while rows execute against a quota-limited model.
"""
return self.parallel_request_limiter._create_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
@ -206,6 +236,57 @@ class _PROXY_BatchRateLimiter(CustomLogger):
model_has_failures=False,
)
@staticmethod
def _project_has_any_io_token_limits(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""True when the project has any per-model ITPM/OTPM quota configured.
Used to stop the "skip batch input file processing" fast path from
bypassing a project quota configured for a model other than the
batch's file-bound/top-level routing model: the row models that
actually drive execution and billing aren't known until the JSONL
is parsed, so the file must be read whenever *any* model could be
quota-limited, not only when the routing model itself is.
"""
if user_api_key_dict.project_id is None:
return False
return bool(
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit")
) or bool(get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit"))
def _create_project_io_descriptors_for_models(
self,
user_api_key_dict: UserAPIKeyAuth,
per_model_usage: Mapping[str, Mapping[str, int]],
) -> tuple[list["RateLimitDescriptor"], list[IncrementAmounts]]: # mutable-ok: see below
"""Build project ITPM/OTPM descriptors charged against each row's own model.
One descriptor pair per distinct `body.model` found in the JSONL,
each incremented only by that model's own counted usage -- never the
whole-batch total -- so a quota-limited model can't hide behind an
unlimited routing model, and an unrelated model's rows can't inflate
a different model's counter.
"""
extra_descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: see above
extra_increments: Final[list[IncrementAmounts]] = [] # mutable-ok: see above
for model, usage in per_model_usage.items():
model_descriptors: list[RateLimitDescriptor] = [] # mutable-ok: reset per loop iteration, not module state
self.parallel_request_limiter.add_project_io_token_rate_limit_descriptors_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=model,
descriptors=model_descriptors,
)
for descriptor in model_descriptors:
extra_descriptors.append(descriptor)
extra_increments.append(
{ # mutable-ok: atomic limiter API requires mutable increment records
"requests": 0,
"tokens": usage.get("output_tokens", 0)
if descriptor["key"] == PROJECT_OTPM_DESCRIPTOR_KEY
else usage.get("total_tokens", 0),
}
)
return extra_descriptors, extra_increments
def _should_skip_batch_input_file_processing(
self,
data: dict,
@ -232,6 +313,11 @@ class _PROXY_BatchRateLimiter(CustomLogger):
routing deployment's trusted credentials and the batch is constrained
to run on that provider.
The no-limits check also treats any project-configured ITPM/OTPM
quota as an applicable limit, even when it isn't scoped to the
routing model: a row can target a different, quota-limited model,
and that isn't knowable without parsing the JSONL.
Returns ``(should_skip, descriptors)`` where ``descriptors`` is the
rate-limit descriptor list computed for the no-limits check, so the
caller can reuse it for counter enforcement without recomputing.
@ -257,7 +343,9 @@ class _PROXY_BatchRateLimiter(CustomLogger):
user_api_key_dict=user_api_key_dict,
data=data,
)
if not self._has_applicable_batch_rate_limits(descriptors):
if not self._has_applicable_batch_rate_limits(descriptors) and not self._project_has_any_io_token_limits(
user_api_key_dict
):
verbose_proxy_logger.debug("Skipping batch input file processing: no rate limits configured")
return True, None
@ -297,6 +385,58 @@ class _PROXY_BatchRateLimiter(CustomLogger):
return False
return True
def _estimate_entry_output_tokens(
self,
entry: Mapping[str, object],
min_configured_otpm_limit: int | None,
) -> int:
"""Conservative per-row output-token estimate for the project OTPM reservation.
Batch completion never reconciles actual usage back into the rate
limiter, so this pre-call estimate is the only OTPM enforcement a
batch gets. Mirrors the real-time no-``max_tokens`` floor so a row
that omits an output cap can't be used to bypass OTPM the way an
unbounded streaming request could.
Embeddings rows are identified by the row's own ``url`` (the OpenAI
batch schema puts the target route there, e.g. ``/v1/embeddings``),
never by body shape: a `/v1/responses` row also carries `body.input`
with no `messages`/`prompt`, so guessing from body shape alone would
misclassify a token-generating Responses row as a zero-output
embeddings row and let it skip the OTPM reservation entirely.
"""
url: Final = entry.get("url")
if isinstance(url, str) and "embeddings" in url:
return 0 # embeddings: no output tokens
raw_body: Final = entry.get("body")
body: Final[Mapping[str, object]] = (
MappingProxyType(_BATCH_BODY_ADAPTER.validate_python(raw_body))
if isinstance(raw_body, Mapping)
else MappingProxyType({}) # mutable-ok: immediately frozen empty fallback
)
# `max_tokens`/`max_completion_tokens` cap chat completions; `/v1/responses`
# rows cap output with `max_output_tokens` instead -- omitting it here
# would fall through to the floor estimate for every capped Responses row.
explicit_cap: Final = next(
(
v
for v in (
body.get("max_tokens"),
body.get("max_completion_tokens"),
body.get("max_output_tokens"),
)
if v is not None
),
None,
)
candidate_count: Final = self.parallel_request_limiter.get_output_candidate_count(body)
if explicit_cap is not None:
try:
return max(0, int(explicit_cap)) * candidate_count
except (TypeError, ValueError, OverflowError):
pass
return self.parallel_request_limiter.no_max_tokens_output_floor(min_configured_otpm_limit) * candidate_count
@staticmethod
def _has_applicable_batch_rate_limits(
descriptors: list["RateLimitDescriptor"],
@ -382,9 +522,22 @@ class _PROXY_BatchRateLimiter(CustomLogger):
"""Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded."""
from datetime import datetime
# Find the descriptor for this status
# Find the descriptor for this status. Matching on (key, value) is
# required, not key alone: a batch can carry several project ITPM/OTPM
# descriptors sharing one key (e.g. `model_per_project_otpm`) but
# scoped to different models via `value`
# ("{project_id}:{model}") -- key-only matching would always resolve
# to the first same-keyed descriptor regardless of which one was
# actually over its limit. Falls back to key-only matching for
# statuses that predate `descriptor_value` (e.g. from should_rate_limit).
status_descriptor_value: Final = status.get("descriptor_value")
descriptor_index: Final = next(
(i for i, d in enumerate(descriptors) if d.get("key") == status.get("descriptor_key")),
(
i
for i, d in enumerate(descriptors)
if d.get("key") == status.get("descriptor_key")
and (status_descriptor_value is None or d.get("value") == status_descriptor_value)
),
0,
)
descriptor: Final[RateLimitDescriptor] = (
@ -407,9 +560,27 @@ class _PROXY_BatchRateLimiter(CustomLogger):
f"Limit resets at: {reset_time_formatted}"
)
else: # tokens
# Project ITPM/OTPM descriptors are keyed "{project_id}:{model}" and
# charged with that model's own rows (see
# `_create_project_io_descriptors_for_models`), not the whole
# batch's totals -- report the matching per-model figure when one
# is available so the error reflects what was actually charged.
descriptor_model: Final = (
descriptor.get("value", "").split(":", 1)[-1]
if descriptor.get("key") in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY)
else None
)
model_usage: Final = batch_usage.per_model_usage.get(descriptor_model) if descriptor_model else None
batch_token_count: Final = (
(model_usage or {}).get("output_tokens", batch_usage.output_tokens)
if descriptor.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY
else (model_usage or {}).get("total_tokens", batch_usage.total_tokens)
if descriptor.get("key") == PROJECT_ITPM_DESCRIPTOR_KEY
else batch_usage.total_tokens
)
detail = (
f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. "
f"Batch contains {batch_usage.total_tokens} tokens but only {remaining_display} tokens remaining "
f"Batch contains {batch_token_count} tokens but only {remaining_display} tokens remaining "
f"out of {current_limit} TPM limit. "
f"Limit resets at: {reset_time_formatted}"
)
@ -444,7 +615,10 @@ class _PROXY_BatchRateLimiter(CustomLogger):
falls back to a per-process asyncio.Lock + in-memory operation.
``descriptors`` may be passed in by the pre-call hook to reuse the list
already computed when deciding whether to skip file processing.
already computed when deciding whether to skip file processing. It
never contains project ITPM/OTPM descriptors (those are model-specific
and only knowable once ``batch_usage.per_model_usage`` is populated by
parsing the JSONL), so this always builds and appends them here.
"""
if descriptors is None:
descriptors = self._create_batch_rate_limit_descriptors(
@ -452,11 +626,20 @@ class _PROXY_BatchRateLimiter(CustomLogger):
data=data,
)
increment: Final[dict[Literal["requests", "tokens"], int]] = {
"requests": batch_usage.request_count,
"tokens": batch_usage.total_tokens,
}
increments: Final[list[dict[Literal["requests", "tokens"], int]]] = [increment for _ in descriptors]
increments: list[IncrementAmounts] = [ # mutable-ok: reassigned below to append project IO increments
{ # mutable-ok: atomic limiter API requires mutable increment records
"requests": batch_usage.request_count,
"tokens": batch_usage.total_tokens,
}
for _d in descriptors
]
project_io_descriptors, project_io_increments = self._create_project_io_descriptors_for_models(
user_api_key_dict=user_api_key_dict,
per_model_usage=batch_usage.per_model_usage,
)
descriptors = [*descriptors, *project_io_descriptors]
increments = [*increments, *project_io_increments]
rate_limit_response: Final = await self.parallel_request_limiter.atomic_check_and_increment_by_n(
descriptors=descriptors,
@ -482,6 +665,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
user_api_key_dict: UserAPIKeyAuth | None = None,
data: dict | None = None,
descriptors: Sequence["RateLimitDescriptor"] | None = None,
) -> BatchFileUsage:
"""
Count number of requests and tokens in a batch input file.
@ -490,10 +674,37 @@ class _PROXY_BatchRateLimiter(CustomLogger):
file_id: The file ID to read
custom_llm_provider: The custom LLM provider to use for token encoding
user_api_key_dict: User authentication information for file access (required for managed files)
descriptors: Rate limit descriptors already computed for this batch, so the
configured project OTPM limit can scale the no-``max_tokens`` output floor
Returns:
BatchFileUsage with total_tokens and request_count
BatchFileUsage with total_tokens, output_tokens, request_count, and
per_model_usage (each row's own totals, keyed by its `body.model`)
"""
descriptor_otpm_limits: Final = tuple(
int(v)
for d in (descriptors or ())
if d.get("key") == PROJECT_OTPM_DESCRIPTOR_KEY
for rate_limit in (d.get("rate_limit"),)
for v in (rate_limit.get("tokens_per_unit") if rate_limit is not None else None,)
if v is not None
)
# `descriptors` only ever carries the routing model's own OTPM limit
# (see `_create_batch_rate_limit_descriptors`), but a row can target
# any project-configured model. Folding in every configured model's
# OTPM limit keeps the no-`max_tokens` floor from drifting wide just
# because a row's specific model isn't known until parsed below.
project_otpm_limits: Final = (
tuple(int(v) for v in project_otpm_limit_map.values())
if user_api_key_dict is not None
and (
project_otpm_limit_map := get_model_rate_limit_from_metadata(
user_api_key_dict, "project_metadata", "model_otpm_limit"
)
)
else ()
)
min_configured_otpm_limit: Final = min((*descriptor_otpm_limits, *project_otpm_limits), default=None)
try:
# Check if this is a managed file (base64 encoded unified file ID)
from litellm.proxy.openai_files_endpoints.common_utils import (
@ -545,23 +756,51 @@ class _PROXY_BatchRateLimiter(CustomLogger):
# Counting stays best-effort, so a legitimate (e.g. multimodal) row
# the counter can't measure is estimated, not hard-rejected.
models: Final[set] = set()
# Keyed by each row's own `body.model`, so the project ITPM/OTPM
# quota for that model is charged with only its own rows' tokens,
# never the whole batch's -- see `_create_project_io_descriptors_for_models`.
per_model_usage: Final[dict[str, dict[str, int]]] = {}
total_tokens = 0
output_tokens = 0 # rebind-ok: accumulated per JSONL row in the loop below
request_count = 0
for raw_line in _iter_batch_input_lines(file_content_bytes):
request_count += 1
try:
entry = json.loads(raw_line)
except Exception:
total_tokens += _estimate_batch_entry_tokens(raw_line)
entry_total_tokens = _estimate_batch_entry_tokens(raw_line)
entry_output_tokens = self.parallel_request_limiter.no_max_tokens_output_floor(
min_configured_otpm_limit
)
total_tokens += entry_total_tokens
output_tokens += entry_output_tokens
continue
model: str | None = (entry.get("body") or {}).get("model") if isinstance(entry, dict) else None
if model:
models.add(model)
if isinstance(entry, dict):
model = (entry.get("body") or {}).get("model")
if model:
models.add(model)
entry_output_tokens = self._estimate_entry_output_tokens(entry, min_configured_otpm_limit)
else:
entry_output_tokens = self.parallel_request_limiter.no_max_tokens_output_floor(
min_configured_otpm_limit
)
output_tokens += entry_output_tokens
try:
total_tokens += _count_entry_tokens(entry)
entry_total_tokens = _count_entry_tokens(entry)
except Exception:
total_tokens += _estimate_batch_entry_tokens(raw_line)
entry_total_tokens = _estimate_batch_entry_tokens(raw_line)
total_tokens += entry_total_tokens
if model:
model_usage = per_model_usage.setdefault(
model, {"total_tokens": 0, "output_tokens": 0, "request_count": 0}
)
model_usage["total_tokens"] += entry_total_tokens
model_usage["output_tokens"] += entry_output_tokens
model_usage["request_count"] += 1
# Validate every model named in the batch JSONL against the
# caller's per-key model allowlist. Without this, a caller
@ -578,6 +817,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
return BatchFileUsage(
total_tokens=total_tokens,
request_count=request_count,
output_tokens=output_tokens,
per_model_usage=per_model_usage,
)
except HTTPException as e:
@ -814,6 +1055,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
custom_llm_provider=custom_llm_provider,
user_api_key_dict=user_api_key_dict,
data=data,
descriptors=batch_rate_limit_descriptors,
)
verbose_proxy_logger.debug(

View file

@ -454,7 +454,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
)
verbose_proxy_logger.debug("Atomic check+increment response: %s", json.dumps(atomic_response, indent=2))
verbose_proxy_logger.debug(
"Atomic check+increment response: %s", json.dumps(atomic_response, indent=2, default=list)
)
if atomic_response["overall_code"] == "OVER_LIMIT":
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model)

File diff suppressed because it is too large Load diff

View file

@ -11,6 +11,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_litellm_metadata_from_kwargs,
)
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
get_key_object,
@ -125,6 +126,15 @@ class _ProxyDBLogger(CustomLogger):
existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
existing_metadata.update(_metadata)
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
if (
isinstance(litellm_metadata_bucket, dict)
and "standard_logging_guardrail_information" not in existing_metadata
):
guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information")
if guardrail_info is not None:
existing_metadata["standard_logging_guardrail_information"] = guardrail_info
if "litellm_params" not in request_data:
request_data["litellm_params"] = {}
@ -175,9 +185,14 @@ class _ProxyDBLogger(CustomLogger):
# recovered cost onto request_data (the usage rides along in
# ``combined_usage_object`` for the token columns), so attribute the
# real partial spend to this failure row instead of zero.
recovered_response_cost = 0.0
if isinstance(request_data.get("combined_usage_object"), litellm.Usage):
recovered_response_cost = max(float(request_data.get("response_cost") or 0.0), 0.0)
recovered_stream_cost: Final = (
max(float(request_data.get("response_cost") or 0.0), 0.0)
if isinstance(request_data.get("combined_usage_object"), litellm.Usage)
else 0.0
)
recovered_response_cost: Final = recovered_stream_cost + guardrail_information_cost(
existing_metadata.get("standard_logging_guardrail_information")
)
await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key_dict.api_key,

View file

@ -38,6 +38,8 @@ from litellm.proxy._types import (
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
SpecialHeaders,
TeamCallbackMetadata,
UserAPIKeyAuth,
@ -294,7 +296,10 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY: Final = "allow_client_mess
_CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
# ``model_info`` carries the same pricing fields when read by
# ``use_custom_pricing_for_model``; strip from metadata for the same reason.
_CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info"})
# ``standard_logging_guardrail_information`` is proxy-written telemetry summed
# into response_cost and spend; a client seeding it forges (even negative)
# guardrail cost.
_CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logging_guardrail_information"})
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"
# Request fields whose value, when URL-valued, becomes the outbound destination
@ -348,6 +353,36 @@ def reject_url_valued_destination(field: str, value: str) -> None:
)
_METADATA_JSON_TYPE_NAMES: Final[Mapping[type, str]] = MappingProxyType(
{bool: "a boolean", int: "an integer", float: "a number", str: "a string", list: "an array"}
)
def _invalid_metadata_type_error(field: str, value: object) -> ProxyException:
received_type: Final = _METADATA_JSON_TYPE_NAMES.get(type(value), f"a {type(value).__name__}")
return ProxyException(
message=f"Invalid type for '{field}': expected an object, but got {received_type} instead.",
type=ProxyErrorTypes.bad_request_error,
param=field,
code=400,
)
def _normalized_metadata_object(field: str, value: object) -> Mapping[str, Any]:
"""Return ``value`` as a metadata object or raise a 400 like OpenAI does.
A JSON string that parses to an object is accepted because multipart/form-data
and ``extra_body`` callers can only send metadata as a string. The caller pops
the raw value from the request body before validating so the failure-logging
hooks that inspect the body afterwards don't crash on it and mask the 400 as a 500.
"""
if isinstance(value, dict):
return value
if isinstance(value, str) and isinstance((parsed := safe_json_loads(value)), dict):
return parsed
raise _invalid_metadata_type_error(field=field, value=value)
def _strip_untrusted_request_header_controls(
headers: Any,
*,
@ -1271,8 +1306,15 @@ class LiteLLMProxyRequestSetup:
)
if user_api_key_dict.budget_reservation is not None:
data[_metadata_variable_name]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation
# Add the full UserAPIKeyAuth object for MCP server access control
data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict
# UserAPIKeyAuth object for MCP server access control
data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy(
update={
"metadata": strip_callback_config(user_api_key_dict.metadata),
"team_metadata": strip_callback_config(user_api_key_dict.team_metadata),
"project_metadata": strip_callback_config(user_api_key_dict.project_metadata),
"organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata),
}
)
return data
@staticmethod
@ -1294,10 +1336,11 @@ class LiteLLMProxyRequestSetup:
)
# ignore any special fields
added_metadata: Final = {}
for k, v in management_endpoint_metadata.items():
if k not in (LiteLLM_ManagementEndpoint_MetadataFields_Premium + LiteLLM_ManagementEndpoint_MetadataFields):
added_metadata[k] = v
added_metadata: Final = {
k: v
for k, v in (strip_callback_config(management_endpoint_metadata) or {}).items()
if k not in (LiteLLM_ManagementEndpoint_MetadataFields_Premium + LiteLLM_ManagementEndpoint_MetadataFields)
}
if data[_metadata_variable_name].get("user_api_key_auth_metadata") is None:
data[_metadata_variable_name]["user_api_key_auth_metadata"] = {}
data[_metadata_variable_name]["user_api_key_auth_metadata"].update(added_metadata)
@ -1572,6 +1615,13 @@ async def add_litellm_data_to_request(
continue
data.pop(_internal_key, None)
_reject_url_valued_destinations(data)
_raw_metadata_by_field: Final = {
_metadata_field: data.pop(_metadata_field)
for _metadata_field in ("metadata", "litellm_metadata")
if data.get(_metadata_field) is not None
}
for _metadata_field, _raw_metadata in _raw_metadata_by_field.items():
data[_metadata_field] = _normalized_metadata_object(_metadata_field, _raw_metadata)
# Strip spoofable auth metadata from user-supplied metadata dict
_user_metadata = data.get("metadata")
if isinstance(_user_metadata, dict):
@ -1711,29 +1761,10 @@ async def add_litellm_data_to_request(
verbose_proxy_logger.debug("receiving data: %s", data)
# Parse metadata if it's a string (e.g., from multipart/form-data)
if "metadata" in data and data["metadata"] is not None:
if isinstance(data["metadata"], str):
data["metadata"] = safe_json_loads(data["metadata"])
if not isinstance(data["metadata"], dict):
verbose_proxy_logger.warning(
"Failed to parse 'metadata' as JSON dict. Received value: %s", data["metadata"]
)
# requester_metadata is snapshotted AFTER the strip below so
# downstream consumers (e.g. PANW guardrail reading user_ip /
# profile_id) don't see attacker-injected admin slots preserved in
# the deepcopy.
# Parse litellm_metadata if it's a string (e.g., from multipart/form-data or extra_body)
if "litellm_metadata" in data and data["litellm_metadata"] is not None:
if isinstance(data["litellm_metadata"], str):
parsed_litellm_metadata: Final = safe_json_loads(data["litellm_metadata"])
if not isinstance(parsed_litellm_metadata, dict):
verbose_proxy_logger.warning(
"Failed to parse 'litellm_metadata' as JSON dict. Received value: %s", data["litellm_metadata"]
)
else:
data["litellm_metadata"] = parsed_litellm_metadata
# requester_metadata is snapshotted AFTER the strip below so
# downstream consumers (e.g. PANW guardrail reading user_ip /
# profile_id) don't see attacker-injected admin slots preserved in
# the deepcopy.
# Strip internal pipeline state and admin-injection slots from user input.
# Runs AFTER the string-to-dict parse above so JSON-string metadata (sent

View file

@ -574,6 +574,33 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]:
)
_NO_KEY_LABELS: Final[tuple[str | None, str | None]] = (None, None)
async def _with_key_labels(
prisma_client: "PrismaClient", responses: Sequence[ShadowEvalJobResponse]
) -> tuple[ShadowEvalJobResponse, ...]:
"""Resolve each job's key hash to the key's alias and masked name in one batched read,
so the UI can say whose traffic a job shadows. Deleted keys resolve to None."""
if not responses:
return ()
key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": sorted({response.api_key_id for response in responses})}} # mutable-ok: Prisma filter
)
labels: Final[Mapping[str, tuple[str | None, str | None]]] = {
row.token: (row.key_alias, row.key_name) for row in key_rows or ()
}
return tuple(
response.model_copy(
update={ # mutable-ok: pydantic update payload
"key_alias": labels.get(response.api_key_id, _NO_KEY_LABELS)[0],
"key_name": labels.get(response.api_key_id, _NO_KEY_LABELS)[1],
}
)
for response in responses
)
async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> ShadowEvalResult | None:
"""Both stratifications of one job's verdicts. Tier answers "where does the router do
well"; the model stratification groups by whichever model served the real arm, so it
@ -686,7 +713,9 @@ async def start_shadow_eval(
f"Key already has an active {data.direction} shadow eval job (started concurrently). Stop it first."
),
) from e
return ShadowEvalJobResponse.model_validate(job, from_attributes=True)
return ShadowEvalJobResponse.model_validate(job, from_attributes=True).model_copy(
update={"key_alias": key_row.key_alias, "key_name": key_row.key_name} # mutable-ok: pydantic update payload
)
@router.get(
@ -711,7 +740,10 @@ async def list_shadow_eval_jobs(
order={"created_at": "desc"}, # mutable-ok: Prisma order
take=limit,
)
return tuple(ShadowEvalJobResponse.model_validate(record, from_attributes=True) for record in records or ())
return await _with_key_labels(
prisma_client,
tuple(ShadowEvalJobResponse.model_validate(record, from_attributes=True) for record in records or ()),
)
@router.get(
@ -742,7 +774,10 @@ async def get_shadow_eval_job(
where={"job_id": job_id, "outcome": "error"}, # mutable-ok: Prisma filter
order={"created_at": "desc"}, # mutable-ok: Prisma order
)
return ShadowEvalJobResponse.model_validate(record, from_attributes=True).model_copy(
labeled: Final = await _with_key_labels(
prisma_client, (ShadowEvalJobResponse.model_validate(record, from_attributes=True),)
)
return labeled[0].model_copy(
update={ # mutable-ok: pydantic update payload
"judged_count": totals[0].judged_count if totals else 0,
"error_count": totals[0].error_count if totals else 0,
@ -781,4 +816,7 @@ async def stop_shadow_eval_job(
where={"id": job_id}, # mutable-ok: Prisma filter
data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload
)
return ShadowEvalJobResponse.model_validate(updated, from_attributes=True)
labeled: Final = await _with_key_labels(
prisma_client, (ShadowEvalJobResponse.model_validate(updated, from_attributes=True),)
)
return labeled[0]

View file

@ -1,5 +1,6 @@
import asyncio
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, Protocol
@ -142,6 +143,11 @@ class _GroupingSetsRow(SimpleNamespace):
failed_requests: int | None
class _EntityRollupRow(_GroupingSetsRow):
entity_id: str | None
api_key_rolled: int
def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float:
"""Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled.
@ -224,6 +230,15 @@ def compute_tag_metadata_totals(records: Sequence[DailySpendRecord]) -> SpendMet
return metadata_metrics
def _entity_metadata(
entity_metadata_field: Mapping[str, dict[str, object]] | None,
entity_id: str,
) -> dict[str, object]:
"""The metadata payload for one entity breakdown bucket, empty when the caller passed none."""
stored: Final = entity_metadata_field.get(entity_id) if entity_metadata_field else None
return stored if stored is not None else {} # mutable-ok: payload pydantic validates into its own dict
def update_breakdown_metrics(
breakdown: BreakdownMetrics,
record: DailySpendRecord,
@ -395,7 +410,7 @@ def update_breakdown_metrics(
if entity_value not in breakdown.entities:
breakdown.entities[entity_value] = MetricWithMetadata(
metrics=SpendMetrics(),
metadata=(entity_metadata_field.get(entity_value, {}) if entity_metadata_field else {}),
metadata=_entity_metadata(entity_metadata_field, entity_value),
)
breakdown.entities[entity_value].metrics = update_metrics(breakdown.entities[entity_value].metrics, record)
@ -419,7 +434,7 @@ def update_breakdown_metrics(
async def get_api_key_metadata(
prisma_client: PrismaClient,
api_keys: set[str],
api_keys: AbstractSet[str],
) -> dict[str, _KeyMetadataDict]:
"""Get api key metadata, falling back to deleted keys table for keys not found in active table.
@ -555,34 +570,17 @@ def _build_where_conditions(
return where_conditions
def _build_aggregated_sql_query(
def _build_aggregated_where_clause(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None,
start_date: str,
end_date: str,
adjusted_start: str,
adjusted_end: str,
model: str | None,
api_key: str | None,
exclude_entity_ids: list[str] | None = None,
timezone_offset_minutes: int | None = None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None, # mutable-ok: filter union shared with the paginated path
) -> tuple[str, list[str]]:
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
Groups by (date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns.
The entity_id column is intentionally omitted from GROUP BY to collapse
rows across entities this is where the biggest row reduction comes from.
Returns:
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
"""Build the WHERE clause and $N params shared by the aggregated queries."""
sql_conditions: Final[list[str]] = []
sql_params: Final[list[str]] = []
p = 1 # parameter index (1-based for PostgreSQL $N placeholders)
@ -596,13 +594,16 @@ def _build_aggregated_sql_query(
sql_params.append(adjusted_end)
p += 1
# Optional entity filter
# Optional entity filter; an empty list must match nothing, not everything
if entity_id is not None:
if isinstance(entity_id, list):
placeholders = ", ".join(f"${p + i}" for i in range(len(entity_id)))
sql_conditions.append(f'"{entity_id_field}" IN ({placeholders})')
sql_params.extend(entity_id)
p += len(entity_id)
if entity_id:
placeholders = ", ".join(f"${p + i}" for i in range(len(entity_id)))
sql_conditions.append(f'"{entity_id_field}" IN ({placeholders})')
sql_params.extend(entity_id)
p += len(entity_id)
else:
sql_conditions.append("FALSE")
else:
sql_conditions.append(f'"{entity_id_field}" = ${p}')
sql_params.append(entity_id)
@ -621,13 +622,68 @@ def _build_aggregated_sql_query(
sql_params.append(model)
p += 1
# Optional api_key filter
if api_key:
# Optional api_key filter; an empty list must match nothing, not everything
if isinstance(api_key, list):
if api_key:
placeholders = ", ".join(f"${p + i}" for i in range(len(api_key)))
sql_conditions.append(f"api_key IN ({placeholders})")
sql_params.extend(api_key)
p += len(api_key)
else:
sql_conditions.append("FALSE")
elif api_key:
sql_conditions.append(f"api_key = ${p}")
sql_params.append(api_key)
p += 1
where_clause: Final = " AND ".join(sql_conditions)
return " AND ".join(sql_conditions), sql_params
def _ptu_flat_cost_select(table_name: str) -> str:
"""Only LiteLLM_DailyTeamSpend carries ptu_flat_cost; other daily tables emit a
constant zero so the SpendMetrics.flat_cost response shape stays uniform."""
if table_name == "litellm_dailyteamspend":
return "SUM(ptu_flat_cost)::float AS ptu_flat_cost"
return "0::float AS ptu_flat_cost"
def _build_aggregated_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
timezone_offset_minutes: int | None = None,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
Groups by (date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns.
The entity_id column is intentionally omitted from GROUP BY to collapse
rows across entities this is where the biggest row reduction comes from.
Returns:
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
where_clause, sql_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
# Postgres computes every rollup level the response needs — per-date
# totals, per-(date, model), per-(date, model, api_key), per-provider,
@ -641,14 +697,6 @@ def _build_aggregated_sql_query(
# total_successful_requests metadata they feed) once the admin UI reads SGR
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
# api_requests rollups are still served from here.
#
# Only LiteLLM_DailyTeamSpend carries ptu_flat_cost; other daily tables emit a
# constant zero so the SpendMetrics.flat_cost response shape stays uniform.
ptu_flat_cost_select: Final = (
"SUM(ptu_flat_cost)::float AS ptu_flat_cost"
if table_name == "litellm_dailyteamspend"
else "0::float AS ptu_flat_cost"
)
sql_query: Final = f"""
SELECT
date,
@ -662,7 +710,7 @@ def _build_aggregated_sql_query(
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
{ptu_flat_cost_select},
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
@ -696,6 +744,70 @@ def _build_aggregated_sql_query(
return sql_query, sql_params
def _build_entity_rollup_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
timezone_offset_minutes: int | None = None,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""Per-entity companion to _build_aggregated_sql_query.
Two rollup levels over the same WHERE clause (date, entity) and
(date, entity, api_key) told apart by GROUPING(api_key): 1 when the
api_key column is rolled up, 0 when it is part of the key.
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
where_clause, sql_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
sql_query: Final = f"""
SELECT
"{entity_id_field}" AS entity_id,
date,
api_key,
GROUPING(api_key) AS api_key_rolled,
SUM(spend)::float AS spend,
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens,
SUM(compression_saved_tokens)::bigint AS compression_saved_tokens,
SUM(compression_savings_spend)::float AS compression_savings_spend,
SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend,
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date, "{entity_id_field}"),
(date, "{entity_id_field}", api_key)
)
"""
return sql_query, sql_params
def _aggregate_spend_records_sync(
*,
records: Sequence[DailySpendRecord],
@ -1097,6 +1209,40 @@ async def get_daily_activity(
)
def _fold_entity_rollups_sync(
*,
results: Sequence[DailySpendData],
entity_rows: Sequence[_EntityRollupRow],
api_key_metadata: Mapping[str, _KeyMetadataDict],
entity_metadata_field: Mapping[str, dict[str, object]] | None, # mutable-ok: shared field shape
) -> None:
"""Write breakdown.entities onto the already-built per-day results."""
by_date: Final = {day.date.strftime("%Y-%m-%d"): day for day in results} # mutable-ok: local fold index
for row in entity_rows:
day = by_date.get(row.date)
if day is None:
continue
entities = day.breakdown.entities
entity_id = row.entity_id or "Unassigned"
bucket = entities.get(entity_id)
if bucket is None:
bucket = MetricWithMetadata(
metrics=SpendMetrics(),
metadata=_entity_metadata(entity_metadata_field, entity_id),
)
entities[entity_id] = bucket
metrics = _record_to_spend_metrics(row)
if row.api_key_rolled:
bucket.metrics = metrics
elif row.api_key and row.api_key != PTU_SENTINEL_API_KEY:
bucket.api_key_breakdown[row.api_key] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, row.api_key)
)
async def get_daily_activity_aggregated(
prisma_client: PrismaClient | None,
table_name: str,
@ -1106,9 +1252,10 @@ async def get_daily_activity_aggregated(
start_date: str | None,
end_date: str | None,
model: str | None,
api_key: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None,
timezone_offset_minutes: int | None = None,
include_entity_breakdown: bool = False,
) -> SpendAnalyticsPaginatedResponse:
"""Aggregated variant that returns the full result set (no pagination).
@ -1116,6 +1263,9 @@ async def get_daily_activity_aggregated(
all individual rows into Python. This collapses rows across entities
(users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows.
include_entity_breakdown runs a small companion rollup query and folds
`breakdown.entities` onto the response, as entity-scoped views like Team Usage need.
Matches the response model of the paginated endpoint so the UI does not need to transform.
"""
if prisma_client is None:
@ -1143,12 +1293,34 @@ async def get_daily_activity_aggregated(
timezone_offset_minutes=timezone_offset_minutes,
)
# Execute GROUPING SETS query — returns one row per rollup level.
rows = await prisma_client.db.query_raw(sql_query, *sql_params)
if rows is None:
rows = []
entity_query: Final = (
_build_entity_rollup_sql_query(
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
timezone_offset_minutes=timezone_offset_minutes,
)
if include_entity_breakdown
else None
)
records: Final = [_GroupingSetsRow(**row) for row in rows]
# Execute the GROUPING SETS query (one row per rollup level), alongside
# the per-entity companion rollup when the caller wants entities.
raw_rows, raw_entity_rows = (
await asyncio.gather(
prisma_client.db.query_raw(sql_query, *sql_params),
prisma_client.db.query_raw(entity_query[0], *entity_query[1]),
)
if entity_query is not None
else (await prisma_client.db.query_raw(sql_query, *sql_params), None)
)
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])]
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
@ -1157,6 +1329,24 @@ async def get_daily_activity_aggregated(
records=records,
)
if raw_entity_rows:
entity_records: Final = tuple(_EntityRollupRow(**row) for row in raw_entity_rows)
entity_api_keys: Final = frozenset(
r.api_key for r in entity_records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY
)
entity_key_metadata: Final = (
await get_api_key_metadata(prisma_client, entity_api_keys)
if entity_api_keys
else {} # mutable-ok: matches the helper's dict return
)
await asyncio.to_thread(
_fold_entity_rollups_sync,
results=aggregated["results"],
entity_rows=entity_records,
api_key_metadata=entity_key_metadata,
entity_metadata_field=entity_metadata_field,
)
return SpendAnalyticsPaginatedResponse(
results=aggregated["results"],
metadata=DailySpendMetadata(

View file

@ -1,7 +1,7 @@
"""`/management/v1/spend_logs` facets."""
from datetime import datetime, timezone
from typing import Annotated, Any, Final
from typing import Annotated, Any, Final, Literal
from fastapi import APIRouter, Depends, Query, Request
@ -35,7 +35,7 @@ def _as_utc(value: datetime) -> datetime:
return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc)
async def _end_user_scope_clause(
async def _spend_log_scope_clause(
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
next_param_index: int,
@ -43,8 +43,8 @@ async def _end_user_scope_clause(
"""SQL predicate restricting the facet to spend logs this caller may read.
Returns ``(None, ())`` for a proxy admin. Mirrors the scoping ``/spend/logs/ui``
applies, so the dropdown can never offer an end user whose rows the caller
could not open.
applies, so a dropdown can never offer a value from a row the caller could
not open.
"""
from litellm.proxy.spend_tracking.spend_management_endpoints import (
_get_permitted_team_ids_for_spend_logs,
@ -77,6 +77,98 @@ async def _end_user_scope_clause(
return f"({' OR '.join(clauses)})", params
async def _list_spend_log_facet(
request: Request,
user_api_key_dict: UserAPIKeyAuth,
start_time: datetime,
end_time: datetime,
q: str | None,
page: int,
page_size: int,
column: Literal["end_user", "user"],
) -> FacetListResponse:
try:
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}database-not-connected",
title="Database not connected",
status=503,
detail=CommonProxyErrors.db_not_connected_error.value,
)
)
column_sql: Final = "end_user" if column == "end_user" else '"user"'
window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time))
search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else ()
search_clause: Final = (f"{column_sql} ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else ()
scope_clause, scope_params = await _spend_log_scope_clause(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
next_param_index=len(window_params) + len(search_params) + 1,
)
where_parts: Final = (
(
"\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')",
"\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')",
f"{column_sql} IS NOT NULL",
f"{column_sql} != ''",
)
+ search_clause
+ ((scope_clause,) if scope_clause is not None else ())
)
# The inner LIMIT walks the startTime index newest first and bounds the
# rows DISTINCT can inspect. request_id makes the cut-off deterministic,
# and page_size + 1 reveals has_more without a COUNT(*).
params: Final = (
window_params
+ search_params
+ scope_params
+ (SPEND_LOGS_FACET_SCAN_CAP, page_size + 1, (page - 1) * page_size)
)
scan_idx: Final = len(params) - 2
facet_sql: Final = (
f"SELECT DISTINCT {column_sql} FROM ("
f" SELECT {column_sql}"
f' FROM "LiteLLM_SpendLogs"'
f" WHERE {' AND '.join(where_parts)}"
f' ORDER BY "startTime" DESC, request_id DESC'
f" LIMIT ${scan_idx}"
f") recent"
f" ORDER BY {column_sql} ASC"
f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}"
)
rows: Final = await prisma_client.db.query_raw(facet_sql, *params)
values: Final[list[str]] = [row[column] for row in rows if row.get(column)]
has_more: Final = len(values) > page_size
return FacetListResponse(
data=values[:page_size],
meta=PageMeta(page=page, page_size=page_size, has_more=has_more),
links=build_page_links(request=request, page=page, has_more=has_more),
)
except ManagementProblem:
raise
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.management_v1.spend_logs._list_spend_log_facet(): Exception occured - %s",
e,
)
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}internal-server-error",
title="Internal server error",
status=500,
detail=f"Failed to list spend log {column.replace('_', ' ')}s.",
)
)
@router.get(
"/spend_logs/end_users",
tags=["Budget & Spend Tracking"],
@ -116,85 +208,47 @@ async def list_spend_log_end_users(
--header 'Authorization: Bearer sk-1234'
```
"""
try:
from litellm.proxy.proxy_server import prisma_client
return await _list_spend_log_facet(
request=request,
user_api_key_dict=user_api_key_dict,
start_time=start_time,
end_time=end_time,
q=q,
page=page,
page_size=page_size,
column="end_user",
)
if prisma_client is None:
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}database-not-connected",
title="Database not connected",
status=503,
detail=CommonProxyErrors.db_not_connected_error.value,
)
)
window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time))
search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else ()
search_clause: Final = (f"end_user ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else ()
scope_clause, scope_params = await _end_user_scope_clause(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
next_param_index=len(window_params) + len(search_params) + 1,
)
where_parts: Final = (
(
"\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')",
"\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')",
"end_user IS NOT NULL",
"end_user != ''",
)
+ search_clause
+ ((scope_clause,) if scope_clause is not None else ())
)
# The inner LIMIT is the safety bound: it walks the startTime index newest
# first and stops, so DISTINCT never runs over an unbounded row set.
# request_id breaks startTime ties so the cut-off row is deterministic and
# successive OFFSET pages agree on the set they are paging through.
# page_size + 1: one row beyond the page reveals has_more without a COUNT(*).
params: Final = (
window_params
+ search_params
+ scope_params
+ (SPEND_LOGS_FACET_SCAN_CAP, page_size + 1, (page - 1) * page_size)
)
scan_idx: Final = len(params) - 2
facet_sql: Final = (
f"SELECT DISTINCT end_user FROM ("
f" SELECT end_user"
f' FROM "LiteLLM_SpendLogs"'
f" WHERE {' AND '.join(where_parts)}"
f' ORDER BY "startTime" DESC, request_id DESC'
f" LIMIT ${scan_idx}"
f") recent"
f" ORDER BY end_user ASC"
f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}"
)
rows: Final = await prisma_client.db.query_raw(facet_sql, *params)
end_users: Final[list[str]] = [row["end_user"] for row in rows if row.get("end_user")]
has_more: Final = len(end_users) > page_size
return FacetListResponse(
data=end_users[:page_size],
meta=PageMeta(page=page, page_size=page_size, has_more=has_more),
links=build_page_links(request=request, page=page, has_more=has_more),
)
except ManagementProblem:
raise
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.management_endpoints.management_v1.spend_logs.list_spend_log_end_users(): Exception occured - %s",
e,
)
raise ManagementProblem(
ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}internal-server-error",
title="Internal server error",
status=500,
detail="Failed to list spend log end users.",
)
)
@router.get(
"/spend_logs/users",
tags=["Budget & Spend Tracking"],
dependencies=[Depends(user_api_key_auth), Depends(reject_unknown_query_params)],
response_model=FacetListResponse,
)
async def list_spend_log_users(
request: Request,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
start_time: Annotated[
datetime,
Query(alias="filter[startTime][gte]", description="Window start (UTC when no offset is given)"),
],
end_time: Annotated[
datetime,
Query(alias="filter[startTime][lte]", description="Window end (UTC when no offset is given)"),
],
q: Annotated[str | None, Query(description="Case-insensitive partial match on the internal user id")] = None,
page: Annotated[int, Query(ge=1, description="Page number")] = 1,
page_size: Annotated[int, Query(ge=1, le=100, description="Page size")] = 50,
) -> FacetListResponse:
"""The distinct internal users appearing in spend logs the caller can read."""
return await _list_spend_log_facet(
request=request,
user_api_key_dict=user_api_key_dict,
start_time=start_time,
end_time=end_time,
q=q,
page=page,
page_size=page_size,
column="user",
)

View file

@ -559,7 +559,7 @@ async def get_organization_daily_activity(
# Fetch organization aliases for metadata
where_condition: Final = _STR_OBJECT_DICT_ADAPTER.validate_python({})
if org_ids_list:
if org_ids_list is not None:
where_condition["organization_id"] = {"in": list(org_ids_list)}
org_aliases: Final = await _table(OrganizationRepository(prisma_client)).find_many(where=where_condition)

View file

@ -9,7 +9,7 @@ import copy
import json
import traceback
from datetime import datetime, timezone
from typing import Any, Final
from typing import Annotated, Any, Final
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -23,6 +23,8 @@ from litellm.proxy._types import (
LitellmTableNames,
ProxyErrorTypes,
ProxyException,
TeamCallbackDeleteResponse,
TeamCallbackDeleteResponseData,
TeamCallbackMetadata,
UserAPIKeyAuth,
)
@ -209,6 +211,14 @@ async def _emit_team_callback_audit_log(
task.add_done_callback(_log_audit_task_exception)
def _callback_error(status_code: int, message: str) -> HTTPException:
"""Build the ``{"error": ...}`` failure body the team callback endpoints return."""
return HTTPException(
status_code=status_code,
detail={"error": message}, # mutable-ok: the error response body is a JSON object
)
@router.post(
"/team/{team_id:path}/callback",
tags=["team management"],
@ -363,6 +373,151 @@ async def add_team_callbacks(
)
@router.delete(
"/team/{team_id:path}/callback/{callback_name}",
tags=["team management"], # mutable-ok: FastAPI's route decorator takes a list of tags
dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator takes a list of dependencies
response_model=TeamCallbackDeleteResponse,
)
@management_endpoint_wrapper
async def delete_team_callback(
http_request: Request,
team_id: str,
callback_name: str,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
litellm_changed_by: Annotated[
str | None,
Header(
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability"
),
] = None,
):
"""
Remove a single callback from a team
The team's other callbacks stay registered and keep firing. Use this instead of
POST /team/{team_id}/disable_logging, which clears every callback on the team at once.
Every entry registered under this callback_name is removed, across callback types, so a
callback registered for both "success" and "failure" is deregistered by one call.
Parameters:
- team_id (str, required): The unique identifier for the team
- callback_name (str, required): The name of the callback to remove, matched exactly as it was
registered with POST /team/{team_id}/callback (e.g. "langfuse", "langsmith", "gcs")
Example curl:
```
curl -X DELETE 'http://localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback/langsmith' \
-H 'Authorization: Bearer sk-1234'
```
Covers callbacks registered through POST /team/{team_id}/callback and the Admin UI. Teams still
on the deprecated callback_settings metadata shape hold no such entries, so this returns 404 for
them; POST /team/{team_id}/disable_logging remains the way to clear those.
Returns 404 if the team does not exist, or if callback_name is not registered for the team.
"""
try:
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
raise _callback_error(500, CommonProxyErrors.db_not_connected_error.value)
_existing_team: Final = await prisma_client.get_data(
team_id=team_id, table_name="team", query_type="find_unique"
)
if _existing_team is None:
raise _callback_error(404, f"Team id = {team_id} does not exist.")
# IDOR guard: only proxy admins / org admins / team admins of THIS team may
# deregister its callbacks, otherwise any authenticated key holder could
# silence another team's observability integration.
await _verify_team_access(
team_obj=LiteLLM_TeamTable(**_existing_team.model_dump()),
user_api_key_dict=user_api_key_dict,
)
team_metadata: Final = _existing_team.metadata
registered_callbacks: Final = team_metadata.get("logging")
entries: Final = registered_callbacks if isinstance(registered_callbacks, list) else ()
remaining_callbacks: Final = [ # mutable-ok: metadata["logging"] is isinstance-checked for list downstream
entry for entry in entries if not (isinstance(entry, dict) and entry.get("callback_name") == callback_name)
]
if len(remaining_callbacks) == len(entries):
raise _callback_error(404, f"callback_name = {callback_name} is not registered for team_id = {team_id}.")
updated_metadata: Final = {**team_metadata, "logging": remaining_callbacks} # mutable-ok: persisted as JSON
encrypted_metadata: Final = encrypt_callback_vars(updated_metadata)
team_metadata_json: Final = json.dumps(encrypted_metadata)
updated_team: Final = await TeamRepository(prisma_client).table.update(
where={"team_id": team_id}, # mutable-ok: prisma where takes a dict literal
data={"metadata": team_metadata_json}, # mutable-ok: prisma data takes a dict literal
# `object_permission` is included so `_refresh_cached_team` doesn't write a
# cached team with the relation nulled out, see team_model_add for the rationale.
include={"object_permission": True}, # mutable-ok: prisma include takes a dict literal
)
if updated_team is None:
raise _callback_error(404, f"Team id = {team_id} does not exist. Error removing team callback")
# Request-time callback resolution reads the cached team, so without this
# the removed callback keeps firing for live keys until the cache expires.
await _refresh_cached_team(
team_row=updated_team,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _emit_team_callback_audit_log(
team_id=team_id,
before_metadata=team_metadata,
after_metadata=encrypted_metadata,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
# Report what survives with the same resolution the GET endpoint uses, so a
# caller can confirm in one round trip that its other callbacks are intact.
surviving: Final = _resolve_team_callbacks(encrypted_metadata)
response: Final = TeamCallbackDeleteResponse(
status="success",
message=f"Callback {callback_name} removed for team {team_id}",
data=TeamCallbackDeleteResponseData(
team_id=team_id,
success_callbacks=tuple(surviving.success_callback or ()),
failure_callbacks=tuple(surviving.failure_callback or ()),
),
)
except HTTPException:
# Legitimate 4xx (403 from the access guard, 404 for an unknown team or
# an unregistered callback). Re-raise without the error-level log noise
# the catch-all below would produce.
raise
except ProxyException:
raise
except Exception as e:
verbose_proxy_logger.error("litellm.proxy.proxy_server.delete_team_callback(): Exception occurred - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
raise ProxyException(
message="Internal Server Error, " + str(e),
type=ProxyErrorTypes.internal_server_error.value,
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
else:
return response
@router.post(
"/team/{team_id}/disable_logging",
tags=["team management"],

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