mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
* fix(proxy): name the blocking guardrail in x-litellm-applied-guardrails When a guardrail hook raises, the common ProxyLogging dispatch (sequential and parallel pre_call, pipeline block, during_call and post_call metrics wrapper, streaming iterator wrapper) now records that guardrail in applied_guardrails before re-raising, and pre_call_hook folds request-declared guardrails in on its raising path. Buffered streams rebuild their response headers after the first chunk so a post_call block reached while buffering carries the blocker too Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): attribute only the raising layer in stream and pipeline blocks The streaming wrapper caught every exception crossing its boundary and named its own callback, so a block by an inner guardrail or a provider stream failure also named every outer guardrail. The wrapper now runs the hook over an upstream boundary that remembers the exception it raised, and skips attribution when the same exception passes through Pipeline blocks converted from SensitiveDataRouteException or ModifyResponseException into a generic guardrail_pipeline_error now still record the blocking step's guardrail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop explanatory docstrings from the stream attribution helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
8445 lines
366 KiB
Python
8445 lines
366 KiB
Python
import asyncio
|
|
import contextlib
|
|
import copy
|
|
import hashlib
|
|
import inspect
|
|
import json
|
|
import math
|
|
import os
|
|
import smtplib
|
|
import ssl
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from collections.abc import (
|
|
AsyncGenerator,
|
|
AsyncIterable,
|
|
AsyncIterator,
|
|
Awaitable,
|
|
Callable,
|
|
Coroutine,
|
|
Mapping,
|
|
Sequence,
|
|
)
|
|
from dataclasses import dataclass, field
|
|
from datetime import date, datetime, timedelta, timezone
|
|
from email.mime.multipart import MIMEMultipart
|
|
from email.mime.text import MIMEText
|
|
from functools import partial
|
|
from types import MappingProxyType
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
ClassVar,
|
|
Final,
|
|
Generic,
|
|
Literal,
|
|
Optional,
|
|
Protocol,
|
|
TypeVar,
|
|
Union,
|
|
cast,
|
|
overload,
|
|
)
|
|
|
|
from typing_extensions import ReadOnly, TypedDict
|
|
|
|
from litellm import _custom_logger_compatible_callbacks_literal
|
|
from litellm.constants import (
|
|
DEFAULT_MODEL_CREATED_AT_TIME,
|
|
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
|
MAX_TEAM_LIST_LIMIT,
|
|
SPEND_LOG_QUEUE_MAX_BYTES,
|
|
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
|
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
|
|
)
|
|
from litellm.proxy._types import (
|
|
CommonProxyErrors,
|
|
ProxyErrorTypes,
|
|
ProxyException,
|
|
SpendLogsMetadata,
|
|
SpendLogsPayload,
|
|
)
|
|
from litellm.proxy.common_utils.openai_error_payload import (
|
|
litellm_call_id_headers,
|
|
openai_error_param,
|
|
with_litellm_call_id,
|
|
)
|
|
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.types.proxy.model_listing import ModelInfoResponse
|
|
from litellm.types.utils import CallTypes, CallTypesLiteral, ModelInfo, Usage
|
|
|
|
try:
|
|
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
|
|
BaseEmailLogger,
|
|
)
|
|
from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import (
|
|
ResendEmailLogger,
|
|
)
|
|
from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import (
|
|
SendGridEmailLogger,
|
|
)
|
|
from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import (
|
|
SMTPEmailLogger,
|
|
)
|
|
except ImportError:
|
|
BaseEmailLogger = None
|
|
SendGridEmailLogger = None
|
|
SMTPEmailLogger = None
|
|
ResendEmailLogger = None
|
|
|
|
try:
|
|
import backoff
|
|
except ImportError:
|
|
raise ImportError("backoff is not installed. Please install it via 'pip install backoff'")
|
|
|
|
from fastapi import HTTPException, status
|
|
from pydantic import TypeAdapter
|
|
|
|
import litellm
|
|
import litellm.litellm_core_utils
|
|
import litellm.litellm_core_utils.litellm_logging
|
|
from litellm import (
|
|
EmbeddingResponse,
|
|
ImageResponse,
|
|
ModelResponse,
|
|
ModelResponseStream,
|
|
Router,
|
|
)
|
|
from litellm._logging import _redact_string, verbose_proxy_logger
|
|
from litellm._service_logger import ServiceLogging, ServiceTypes
|
|
from litellm.caching.caching import DualCache, RedisCache
|
|
from litellm.caching.dual_cache import LimitedSizeOrderedDict
|
|
from litellm.exceptions import (
|
|
GuardrailRaisedException,
|
|
RejectedRequestError,
|
|
SensitiveDataRouteException,
|
|
)
|
|
from litellm.integrations.custom_guardrail import (
|
|
CustomGuardrail,
|
|
ModifyResponseException,
|
|
)
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.integrations.prometheus import PrometheusLogger
|
|
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
|
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
|
|
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
|
from litellm.litellm_core_utils.core_helpers import (
|
|
coerce_token_limit,
|
|
get_or_create_metadata_bucket,
|
|
independent_snapshot,
|
|
is_expected_client_error,
|
|
)
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
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.token_counter import offload_token_count
|
|
from litellm.llms import load_guardrail_translation_mappings
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
from litellm.proxy._types import (
|
|
AlertType,
|
|
CallInfo,
|
|
LiteLLM_VerificationTokenView,
|
|
Member,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.proxy.auth.route_checks import RouteChecks
|
|
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
|
|
from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.db.create_views import (
|
|
create_missing_views,
|
|
create_view_tolerating_race,
|
|
should_create_missing_views,
|
|
)
|
|
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
|
from litellm.proxy.db.db_url_settings import (
|
|
DatabaseURLSettings,
|
|
add_missing_query_params,
|
|
token_refresh_params_from_url,
|
|
)
|
|
from litellm.proxy.db.exception_handler import (
|
|
PrismaDBExceptionHandler,
|
|
call_with_db_reconnect_retry,
|
|
)
|
|
from litellm.proxy.db.health_check_latest import (
|
|
LatestHealthCheckRow,
|
|
fetch_latest_health_checks,
|
|
fetch_latest_health_checks_for_models,
|
|
)
|
|
from litellm.proxy.db.log_db_metrics import log_db_metrics
|
|
from litellm.proxy.db.pgbouncer import database_url_is_pooled
|
|
from litellm.proxy.db.prisma_client import (
|
|
PrismaWrapper,
|
|
parse_iam_endpoint_from_url,
|
|
)
|
|
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
|
from litellm.proxy.db.spend_log_batching import (
|
|
spend_log_queue_within_budget,
|
|
spend_log_row_bytes,
|
|
spend_log_write_batches,
|
|
)
|
|
from litellm.proxy.db.token_auth import (
|
|
DatabaseTokenAuth,
|
|
mint_database_token,
|
|
resolve_database_token_auth,
|
|
)
|
|
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
|
UnifiedLLMGuardrails,
|
|
resolve_endpoint_translation,
|
|
)
|
|
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
|
|
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
|
|
from litellm.proxy.hooks.parallel_request_limiter import (
|
|
_PROXY_MaxParallelRequestsHandler,
|
|
)
|
|
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
|
_PROXY_MaxParallelRequestsHandler_v3,
|
|
)
|
|
from litellm.proxy.hooks.sensitive_data_routing import (
|
|
_PROXY_SensitiveDataRoutingHandler,
|
|
)
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
|
|
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
|
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
|
|
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
|
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
|
from litellm.repositories.budget_repository import BudgetRepository
|
|
from litellm.repositories.config_repository import ConfigRepository
|
|
from litellm.repositories.table_repositories import (
|
|
EndUserRepository,
|
|
HealthCheckRepository,
|
|
SpendLogsRepository,
|
|
UserNotificationsRepository,
|
|
)
|
|
from litellm.repositories.team_repository import TeamRepository
|
|
from litellm.repositories.user_repository import UserRepository
|
|
from litellm.repositories.verification_token_repository import (
|
|
VerificationTokenRepository,
|
|
)
|
|
from litellm.router_utils.common_utils import resolve_model_group_alias
|
|
from litellm.secret_managers.main import str_to_bool
|
|
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
|
from litellm.types.llms.openai import ResponsesAPIResponse
|
|
from litellm.types.mcp import (
|
|
MCPDuringCallResponseObject,
|
|
MCPPreCallRequestObject,
|
|
MCPPreCallResponseObject,
|
|
)
|
|
from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult
|
|
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
|
|
from litellm.utils import (
|
|
_add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp.types import CallToolResult
|
|
from opentelemetry.trace import Span as _Span
|
|
from prisma import models as prisma_models
|
|
from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions
|
|
from prisma.client import TransactionManager
|
|
from prisma.models import LiteLLM_DeprecatedVerificationToken
|
|
from prisma.types import HttpConfig
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
|
from litellm.models.team import LiteLLM_TeamTableCachedObj
|
|
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
|
|
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
|
|
from litellm.repositories.prisma_protocols import TableActions
|
|
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
|
|
|
|
Span = _Span | object
|
|
else:
|
|
Span = Any
|
|
|
|
_T: Final = TypeVar("_T")
|
|
|
|
|
|
class _ViewCountRow(TypedDict):
|
|
view_count: ReadOnly[int]
|
|
view_names: ReadOnly[Sequence[str] | None]
|
|
|
|
|
|
class _RelTuplesRow(TypedDict):
|
|
reltuples: ReadOnly[int]
|
|
|
|
|
|
class _EndUserBatchTable(Protocol):
|
|
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
|
|
|
|
|
|
class _EndUserSpendBatch(Protocol):
|
|
@property
|
|
def litellm_endusertable(self) -> _EndUserBatchTable: ...
|
|
|
|
|
|
unified_guardrail: Final = UnifiedLLMGuardrails()
|
|
|
|
NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages})
|
|
|
|
|
|
def print_verbose(print_statement: object):
|
|
"""
|
|
Prints the given `print_statement` to the console if `litellm.set_verbose` is True.
|
|
Also logs the `print_statement` at the debug level using `verbose_proxy_logger`.
|
|
|
|
:param print_statement: The statement to be printed and logged.
|
|
:type print_statement: Any
|
|
"""
|
|
import traceback
|
|
|
|
verbose_proxy_logger.debug("%s\n%s", print_statement, traceback.format_exc())
|
|
if litellm.set_verbose:
|
|
print(f"LiteLLM Proxy: {_redact_string(str(print_statement))}") # noqa: T201
|
|
|
|
|
|
def _get_email_logger_class():
|
|
"""
|
|
Determine which email logger class to use based on environment variables.
|
|
Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback)
|
|
|
|
Returns:
|
|
The email logger class to use, or None if BaseEmailLogger is not available
|
|
"""
|
|
if BaseEmailLogger is None:
|
|
return None
|
|
|
|
# Check for SendGrid API key
|
|
if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"):
|
|
return SendGridEmailLogger
|
|
|
|
# Check for Resend API key
|
|
if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"):
|
|
return ResendEmailLogger
|
|
|
|
# Check for SMTP configuration
|
|
if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"):
|
|
return SMTPEmailLogger
|
|
|
|
# Fallback to BaseEmailLogger (though it won't actually send emails)
|
|
return BaseEmailLogger
|
|
|
|
|
|
class InternalUsageCache:
|
|
def __init__(self, dual_cache: DualCache):
|
|
self.dual_cache: DualCache = dual_cache
|
|
|
|
async def async_get_cache(
|
|
self,
|
|
key: str,
|
|
litellm_parent_otel_span: Span | None,
|
|
local_only: bool = False,
|
|
**kwargs: object,
|
|
) -> Any:
|
|
return await self.dual_cache.async_get_cache(
|
|
key=key,
|
|
local_only=local_only,
|
|
parent_otel_span=litellm_parent_otel_span,
|
|
**kwargs,
|
|
)
|
|
|
|
async def async_set_cache(
|
|
self,
|
|
key: str,
|
|
value: object,
|
|
litellm_parent_otel_span: Span | None,
|
|
local_only: bool = False,
|
|
**kwargs: object,
|
|
) -> None:
|
|
return await self.dual_cache.async_set_cache(
|
|
key=key,
|
|
value=value,
|
|
local_only=local_only,
|
|
litellm_parent_otel_span=litellm_parent_otel_span,
|
|
**kwargs,
|
|
)
|
|
|
|
async def async_batch_set_cache(
|
|
self,
|
|
cache_list: list[tuple[str, object]],
|
|
litellm_parent_otel_span: Span | None,
|
|
local_only: bool = False,
|
|
**kwargs: object,
|
|
) -> None:
|
|
return await self.dual_cache.async_set_cache_pipeline(
|
|
cache_list=cache_list,
|
|
local_only=local_only,
|
|
litellm_parent_otel_span=litellm_parent_otel_span,
|
|
**kwargs,
|
|
)
|
|
|
|
async def async_batch_get_cache(
|
|
self,
|
|
keys: Sequence[str | None],
|
|
parent_otel_span: Span | None = None,
|
|
local_only: bool = False,
|
|
):
|
|
return await self.dual_cache.async_batch_get_cache(
|
|
keys=list(keys),
|
|
parent_otel_span=parent_otel_span,
|
|
local_only=local_only,
|
|
)
|
|
|
|
async def async_increment_cache(
|
|
self,
|
|
key: str,
|
|
value: float,
|
|
litellm_parent_otel_span: Span | None,
|
|
local_only: bool = False,
|
|
**kwargs,
|
|
):
|
|
return await self.dual_cache.async_increment_cache(
|
|
key=key,
|
|
value=value,
|
|
local_only=local_only,
|
|
parent_otel_span=litellm_parent_otel_span,
|
|
**kwargs,
|
|
)
|
|
|
|
def set_cache(
|
|
self,
|
|
key: str,
|
|
value: object,
|
|
local_only: bool = False,
|
|
**kwargs: object,
|
|
) -> None:
|
|
return self.dual_cache.set_cache(
|
|
key=key,
|
|
value=value,
|
|
local_only=local_only,
|
|
**kwargs,
|
|
)
|
|
|
|
def get_cache(
|
|
self,
|
|
key: str,
|
|
local_only: bool = False,
|
|
**kwargs: object,
|
|
) -> Any:
|
|
return self.dual_cache.get_cache(
|
|
key=key,
|
|
local_only=local_only,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
### LOGGING ###
|
|
|
|
# Cache for inspect.signature checks — avoids repeated introspection per request
|
|
_CALLBACK_ACCEPTS_CALL_INFO: Final[dict[int, bool]] = {}
|
|
|
|
|
|
def _accepts_litellm_call_info(cb: CustomLogger) -> bool:
|
|
key: Final = id(type(cb))
|
|
if key not in _CALLBACK_ACCEPTS_CALL_INFO:
|
|
sig: Final = inspect.signature(cb.async_post_call_response_headers_hook)
|
|
_CALLBACK_ACCEPTS_CALL_INFO[key] = "litellm_call_info" in sig.parameters
|
|
return _CALLBACK_ACCEPTS_CALL_INFO[key]
|
|
|
|
|
|
def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: object) -> None:
|
|
"""
|
|
If `exc` is an HTTPException with a dict `detail`, mutate it in place to
|
|
add `guardrail_name` and `guardrail_mode` taken from the callback instance.
|
|
|
|
Uses setdefault so guardrails that already populate these fields explicitly
|
|
win over the inferred defaults. No-op for non-HTTPException, non-dict-detail,
|
|
or callbacks without `guardrail_name`. Never raises.
|
|
"""
|
|
if not isinstance(exc, HTTPException):
|
|
return
|
|
detail: Final = getattr(exc, "detail", None)
|
|
if not isinstance(detail, dict):
|
|
return
|
|
guardrail_name: Final[object] = getattr(callback, "guardrail_name", None)
|
|
if guardrail_name:
|
|
detail.setdefault("guardrail_name", guardrail_name)
|
|
event_hook: Final[object] = getattr(callback, "event_hook", None)
|
|
if event_hook:
|
|
detail.setdefault("guardrail_mode", event_hook)
|
|
|
|
|
|
def _record_raising_guardrail(request_data: Mapping[str, object], callback: object) -> None:
|
|
guardrail_name: Final[object] = getattr(callback, "guardrail_name", None)
|
|
if isinstance(request_data, dict) and isinstance(guardrail_name, str):
|
|
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=guardrail_name)
|
|
|
|
|
|
class _UpstreamStreamBoundary(Generic[_T]):
|
|
__slots__ = ("_upstream", "failure")
|
|
|
|
def __init__(self, upstream: AsyncIterable[_T]) -> None:
|
|
self._upstream: Final = upstream.__aiter__()
|
|
self.failure: BaseException | None = None
|
|
|
|
def __aiter__(self) -> "_UpstreamStreamBoundary[_T]":
|
|
return self
|
|
|
|
async def __anext__(self) -> _T:
|
|
try:
|
|
return await self._upstream.__anext__()
|
|
except StopAsyncIteration:
|
|
raise
|
|
except Exception as e:
|
|
self.failure = e
|
|
raise
|
|
|
|
|
|
class _StreamIteratorHook(Protocol[_T]):
|
|
def __call__(self, *, response: AsyncIterator[_T]) -> AsyncGenerator[_T, None]: ...
|
|
|
|
|
|
def _is_client_error_exception(exc: Exception) -> bool:
|
|
if isinstance(exc, HTTPException):
|
|
return exc.status_code < 500
|
|
if isinstance(exc, ProxyException):
|
|
return not (exc.code.isdigit() and int(exc.code) >= 500)
|
|
return False
|
|
|
|
|
|
def _exception_changes_request_flow(exc: BaseException) -> bool:
|
|
"""
|
|
True for guardrail exceptions the proxy turns into an alternate request flow
|
|
(a reroute or a passthrough response) rather than a block. A pipeline step
|
|
configured to block must honor that block, so these are surfaced as the
|
|
generic pipeline block instead of being re-raised verbatim.
|
|
"""
|
|
return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException))
|
|
|
|
|
|
def _policy_state_metadata(data: Mapping[str, object]) -> Mapping[str, object]:
|
|
"""
|
|
Return the metadata bucket the policy engine wrote its pipeline state into.
|
|
|
|
The route decides the bucket (``litellm_metadata`` for ``/v1/messages``,
|
|
responses, batches, files and bedrock, ``metadata`` everywhere else), and both
|
|
buckets can be present at once because callers send their own provider-facing
|
|
``metadata`` (Claude Code sends ``metadata.user_id``) or their own
|
|
``litellm_metadata``. Pipeline slots are stripped from caller input before the
|
|
policy engine runs, so whichever bucket carries them is the proxy's own write.
|
|
"""
|
|
return next(
|
|
(
|
|
bucket
|
|
for bucket in (data.get("metadata"), data.get("litellm_metadata"))
|
|
if isinstance(bucket, dict)
|
|
and ("_guardrail_pipelines" in bucket or "_pipeline_managed_guardrails" in bucket)
|
|
),
|
|
{},
|
|
)
|
|
|
|
|
|
def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
|
|
pipelines: Final = _policy_state_metadata(data).get("_guardrail_pipelines")
|
|
return (
|
|
tuple(cast("Sequence[tuple[str, GuardrailPipeline]]", pipelines)) # cast-ok: the policy engine wrote the slot
|
|
if pipelines
|
|
else ()
|
|
)
|
|
|
|
|
|
def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> frozenset[str]:
|
|
return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps)
|
|
|
|
|
|
def pipeline_managed_guardrail_names(
|
|
data: Mapping[str, object], mode: Literal["pre_call", "post_call"]
|
|
) -> frozenset[str]:
|
|
return _pipeline_step_guardrail_names(
|
|
tuple((policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == mode)
|
|
)
|
|
|
|
|
|
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
|
|
resolved: Final = tuple(
|
|
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast( # cast-ok: the resolver returns None for unknown names, filtered below
|
|
_custom_logger_compatible_callbacks_literal, callback
|
|
)
|
|
)
|
|
if isinstance(callback, str)
|
|
else callback
|
|
for callback in litellm.callbacks
|
|
)
|
|
present: Final = tuple(callback for callback in resolved if callback is not None)
|
|
guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail))
|
|
others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger
|
|
"tuple[CustomLogger, ...]",
|
|
tuple(callback for callback in present if not isinstance(callback, CustomGuardrail)),
|
|
)
|
|
return (guardrails, others)
|
|
|
|
|
|
def _merge_pipeline_metadata_bucket(
|
|
data: dict, bucket_key: str, modified_bucket_value: object
|
|
) -> None: # mutable-ok: request payload dict, written in place
|
|
if not isinstance(modified_bucket_value, dict):
|
|
return
|
|
modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed
|
|
surviving_writes: Final = {
|
|
key: value for key, value in modified_bucket.items() if key != "guardrails"
|
|
} # mutable-ok: merged into the live request metadata bucket in place
|
|
existing_bucket: Final = data.get(bucket_key)
|
|
if isinstance(existing_bucket, dict):
|
|
cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed
|
|
else:
|
|
data[bucket_key] = surviving_writes
|
|
|
|
|
|
def _merge_pipeline_metadata_writes(
|
|
data: dict, modified_data: Mapping[str, object]
|
|
) -> None: # mutable-ok: request payload dict, written in place
|
|
"""
|
|
Copy metadata-bucket writes from a pipeline's working copy back onto the request.
|
|
|
|
Post_call pipelines run step hooks against a copied request dict so the payload
|
|
already sent upstream stays untouched, but hooks record proxy-internal logging
|
|
state in the metadata buckets (``applied_guardrails`` for response headers,
|
|
``standard_logging_guardrail_information`` for spend logs), and those writes
|
|
must reach the request dict the proxy keeps reading after the pipeline returns.
|
|
|
|
The ``guardrails`` key is the executor's per-step activation flag for
|
|
``should_run_guardrail``, not a hook write, so it stays in the working copy.
|
|
"""
|
|
for bucket_key in ("metadata", "litellm_metadata"):
|
|
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
|
|
|
|
|
|
def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool:
|
|
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
|
|
if callback is None:
|
|
return False
|
|
if PipelineExecutor.supports_unified_execution(callback):
|
|
return True
|
|
return (
|
|
translation is not None
|
|
and type(translation).assembles_streamed_response
|
|
and PipelineExecutor.supports_streaming_execution(callback)
|
|
)
|
|
|
|
|
|
def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
|
|
return tuple(
|
|
(policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call"
|
|
)
|
|
|
|
|
|
_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress"))
|
|
|
|
|
|
def _is_pending_background_response(response: LLMResponseTypes) -> bool:
|
|
return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES
|
|
|
|
|
|
def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]:
|
|
resolved: Final = PolicyResolver.resolve_policy_guardrails(
|
|
policy_name=policy_name, policies=get_policy_registry().get_all_policies()
|
|
)
|
|
return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps)
|
|
|
|
|
|
def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]:
|
|
return frozenset(
|
|
callback.guardrail_name
|
|
for callback in litellm.callbacks
|
|
if isinstance(callback, CustomGuardrail)
|
|
and callback.guardrail_name is not None
|
|
and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
|
)
|
|
|
|
|
|
def _without_names(
|
|
bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write
|
|
slot: str,
|
|
names: frozenset[str],
|
|
) -> None:
|
|
claimed: Final = bucket.get(slot)
|
|
if not isinstance(claimed, list):
|
|
return
|
|
remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to
|
|
name for name in claimed if name not in names
|
|
]
|
|
if remaining:
|
|
bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place
|
|
else:
|
|
bucket.pop(slot)
|
|
|
|
|
|
def _withdraw_deferred_claims(
|
|
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
|
|
deferred: Sequence[tuple[str, "GuardrailPipeline"]],
|
|
) -> None:
|
|
outside_by_policy: Final = MappingProxyType(
|
|
{policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
|
|
)
|
|
running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union(
|
|
_guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
|
|
)
|
|
withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
|
|
withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere
|
|
_, bucket = get_or_create_metadata_bucket(data)
|
|
_without_names(bucket, "applied_policies", withdrawn_policies)
|
|
_without_names(bucket, "applied_guardrails", withdrawn_guardrails)
|
|
sources: Final = bucket.get("policy_sources")
|
|
if not isinstance(sources, dict):
|
|
return
|
|
remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place
|
|
name: reason for name, reason in sources.items() if name not in withdrawn_policies
|
|
}
|
|
if remaining_sources:
|
|
bucket["policy_sources"] = remaining_sources
|
|
else:
|
|
bucket.pop("policy_sources")
|
|
|
|
|
|
def _defer_post_call_pipelines(
|
|
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
|
|
response: ResponsesAPIResponse,
|
|
) -> None:
|
|
deferred: Final = _post_call_pipelines(data)
|
|
if not deferred:
|
|
return
|
|
verbose_proxy_logger.debug(
|
|
"Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s",
|
|
response.id,
|
|
response.status,
|
|
", ".join(policy_name for policy_name, _pipeline in deferred),
|
|
)
|
|
tag_matched: Final = _tag_matched_deferrals(data, deferred)
|
|
if tag_matched:
|
|
verbose_proxy_logger.warning(
|
|
"Policy engine: background response %s matched post_call policies through a request tag at submit; "
|
|
"retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body "
|
|
"does not govern the completed response: %s",
|
|
response.id,
|
|
", ".join(tag_matched),
|
|
)
|
|
body_selected: Final = _body_selected_deferrals(data, deferred)
|
|
if body_selected:
|
|
verbose_proxy_logger.warning(
|
|
"Policy engine: background response %s matched post_call policies through the request body's policies "
|
|
"list at submit; retrieval carries no request body, so those policies do not govern the completed "
|
|
"response: %s",
|
|
response.id,
|
|
", ".join(body_selected),
|
|
)
|
|
_withdraw_deferred_claims(data, deferred)
|
|
|
|
|
|
def _tag_matched_deferrals(
|
|
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
|
|
) -> tuple[str, ...]:
|
|
sources: Final = _policy_state_metadata(data).get("policy_sources")
|
|
if not isinstance(sources, dict):
|
|
return ()
|
|
return tuple(
|
|
policy_name
|
|
for policy_name, _pipeline in deferred
|
|
if policy_name in sources and "tag:" in str(sources[policy_name])
|
|
)
|
|
|
|
|
|
def _body_selected_deferrals(
|
|
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
|
|
) -> tuple[str, ...]:
|
|
sources: Final = _policy_state_metadata(data).get("policy_sources")
|
|
attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset()
|
|
return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
|
|
|
|
|
|
def _pipeline_unsupported_streaming_guardrails(
|
|
pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
|
|
) -> tuple[str, ...]:
|
|
return tuple(
|
|
dict.fromkeys(
|
|
step.guardrail
|
|
for step in pipeline.steps
|
|
if not _pipeline_step_supports_streaming(step.guardrail, translation)
|
|
)
|
|
)
|
|
|
|
|
|
def _pipeline_is_streamable(
|
|
policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None"
|
|
) -> bool:
|
|
unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation)
|
|
if not unsupported:
|
|
return True
|
|
verbose_proxy_logger.warning(
|
|
"Policy '%s' has post_call pipeline guardrails a streaming pipeline cannot run on this route yet; they "
|
|
"need the unified apply_guardrail interface, or a post-call hook without a streaming iterator hook on a "
|
|
"route whose translation assembles the streamed response. The stream skips the pipeline and its "
|
|
"guardrails run on their own: %s",
|
|
policy_name,
|
|
", ".join(unsupported),
|
|
)
|
|
return False
|
|
|
|
|
|
def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseTranslation | None":
|
|
resolved: Final = resolve_endpoint_translation(user_api_key_dict, None)
|
|
return None if resolved is None else resolved[1]
|
|
|
|
|
|
def stream_gated_guardrail_names(
|
|
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
|
) -> frozenset[str]:
|
|
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
|
|
if translation is None:
|
|
return frozenset()
|
|
return _pipeline_step_guardrail_names(
|
|
tuple(
|
|
(policy_name, pipeline)
|
|
for policy_name, pipeline in _post_call_pipelines(request_data)
|
|
if not _pipeline_unsupported_streaming_guardrails(pipeline, translation)
|
|
)
|
|
)
|
|
|
|
|
|
def _streamable_post_call_pipelines(
|
|
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
|
|
) -> tuple[tuple[str, "GuardrailPipeline"], ...]:
|
|
"""
|
|
The post_call pipelines a streaming response can be gated through.
|
|
|
|
Streaming pipelines scan the buffered stream through the endpoint guardrail
|
|
translation of the request route, so every step's guardrail needs either the
|
|
unified apply_guardrail interface or, on a route whose translation assembles
|
|
the streamed response, a post-call hook that is its only streaming path, and
|
|
the route needs a translation. A pipeline that
|
|
cannot be run that way yet is left out and its guardrails run on the stream
|
|
on their own, the way they did before pipelines ran on streams at all, with
|
|
a warning naming the pipeline.
|
|
"""
|
|
post_call_pipelines: Final = _post_call_pipelines(request_data)
|
|
if not post_call_pipelines:
|
|
return ()
|
|
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
|
|
if translation is None:
|
|
verbose_proxy_logger.warning(
|
|
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
|
|
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
|
|
"on their own: %s",
|
|
user_api_key_dict.request_route,
|
|
", ".join(policy_name for policy_name, _pipeline in post_call_pipelines),
|
|
)
|
|
return ()
|
|
return tuple(
|
|
(policy_name, pipeline)
|
|
for policy_name, pipeline in post_call_pipelines
|
|
if _pipeline_is_streamable(policy_name, pipeline, translation)
|
|
)
|
|
|
|
|
|
def _prompt_block_text(block: object) -> str:
|
|
if isinstance(block, str):
|
|
return block
|
|
if not isinstance(block, dict):
|
|
return ""
|
|
block_text: Final = block.get("text")
|
|
return block_text if isinstance(block_text, str) else ""
|
|
|
|
|
|
def _system_prompt_text(system_input: object) -> str:
|
|
if isinstance(system_input, str):
|
|
return system_input
|
|
if not isinstance(system_input, list):
|
|
return ""
|
|
return "".join(_prompt_block_text(block) for block in system_input)
|
|
|
|
|
|
def _count_request_input_tokens(model: str, request_input: object, system_input: object) -> int:
|
|
system_text: Final = _system_prompt_text(system_input)
|
|
system_tokens: Final = litellm.token_counter(model=model, text=system_text) if system_text else 0
|
|
if isinstance(request_input, str):
|
|
return system_tokens + litellm.token_counter(model=model, text=request_input)
|
|
if not isinstance(request_input, list) or not request_input:
|
|
return system_tokens
|
|
text_entries: Final = tuple(entry for entry in request_input if isinstance(entry, str))
|
|
if len(text_entries) == len(request_input):
|
|
return system_tokens + litellm.token_counter(model=model, text="".join(text_entries))
|
|
return system_tokens + litellm.token_counter(
|
|
model=model, messages=request_input, use_default_image_token_count=True
|
|
)
|
|
|
|
|
|
def _estimate_dispatched_failure_usage(model: str, request_input: object, system_input: object) -> Usage | None:
|
|
"""A request that failed after dispatch consumed provider-billed input
|
|
tokens, but no provider usage ever came back. Estimate the input side with
|
|
the same tokenizer fallback interrupted streams use, so the spend log's
|
|
failure row records what was sent instead of zero."""
|
|
try:
|
|
input_tokens: Final = _count_request_input_tokens(
|
|
model=model, request_input=request_input, system_input=system_input
|
|
)
|
|
except Exception:
|
|
return None
|
|
if input_tokens <= 0:
|
|
return None
|
|
return Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens)
|
|
|
|
|
|
_INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset(
|
|
call_type.value
|
|
for call_type in (
|
|
CallTypes.completion,
|
|
CallTypes.acompletion,
|
|
CallTypes.text_completion,
|
|
CallTypes.atext_completion,
|
|
CallTypes.anthropic_messages,
|
|
CallTypes.aanthropic_messages,
|
|
CallTypes.responses,
|
|
CallTypes.aresponses,
|
|
CallTypes.embedding,
|
|
CallTypes.aembedding,
|
|
CallTypes.moderation,
|
|
CallTypes.amoderation,
|
|
CallTypes.image_generation,
|
|
CallTypes.aimage_generation,
|
|
CallTypes.speech,
|
|
CallTypes.aspeech,
|
|
CallTypes.rerank,
|
|
CallTypes.arerank,
|
|
CallTypes.generate_content,
|
|
CallTypes.agenerate_content,
|
|
CallTypes.generate_content_stream,
|
|
CallTypes.agenerate_content_stream,
|
|
)
|
|
)
|
|
|
|
|
|
def _failure_usage_to_lift(
|
|
model_call_details: Mapping[str, object],
|
|
request_body: Mapping[str, object],
|
|
dispatched: bool,
|
|
) -> tuple[object, object] | None:
|
|
"""A stream that broke mid-flight still billed the provider for the chunks
|
|
already delivered; the streaming handler stashes that recovered usage and
|
|
cost in model_call_details, so prefer it. Otherwise a request that was
|
|
dispatched to a provider and failed without upstream usage gets an
|
|
estimated input-side Usage with zero cost. The raw request body backfills
|
|
the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a
|
|
chat-completions provider) without filling optional_params. Returns the
|
|
(combined_usage_object, response_cost) pair to lift, or None."""
|
|
recovered_usage: Final = model_call_details.get("combined_usage_object")
|
|
if recovered_usage is not None:
|
|
return recovered_usage, model_call_details.get("response_cost")
|
|
if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL):
|
|
return None
|
|
if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES:
|
|
return None
|
|
optional_params: Final = model_call_details.get("optional_params")
|
|
dispatched_system: Final = (
|
|
(optional_params.get("system") or optional_params.get("instructions"))
|
|
if isinstance(optional_params, dict)
|
|
else None
|
|
)
|
|
system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions")
|
|
estimated_usage: Final = _estimate_dispatched_failure_usage(
|
|
model=str(model_call_details.get("model") or ""),
|
|
request_input=model_call_details.get("messages"),
|
|
system_input=system_input,
|
|
)
|
|
if estimated_usage is None:
|
|
return None
|
|
return estimated_usage, 0.0
|
|
|
|
|
|
_EMPTY_LIFT: Final = MappingProxyType({})
|
|
|
|
|
|
def _call_type_for_route(route: str | None) -> str | None:
|
|
"""The route's call type when it maps to a single operation (its async and sync variants);
|
|
None for routes shared by several operations, since the method is not known here."""
|
|
if route is None:
|
|
return None
|
|
call_types: Final = get_call_types_for_route(route)
|
|
if not call_types:
|
|
return None
|
|
operations: Final = frozenset(call_type.value.removeprefix("a") for call_type in call_types)
|
|
return call_types[0].value if len(operations) == 1 else None
|
|
|
|
|
|
_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
|
|
|
|
|
|
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
|
|
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
|
|
request_data (it is not serialisable), so the caller merges these fields
|
|
onto request_data first: the request start and first-handoff instants for
|
|
duration and preprocessing latency, the call type, recovered or estimated
|
|
usage for token counts, and the standard logging object for deployment
|
|
attribution on failed-request spend logs."""
|
|
_logging_obj: Final = request_data.get("litellm_logging_obj")
|
|
if _logging_obj is None:
|
|
return _EMPTY_LIFT
|
|
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
|
|
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
|
|
_usage_to_lift: Final = _failure_usage_to_lift(
|
|
model_call_details=_model_call_details,
|
|
request_body=request_data,
|
|
dispatched=_first_handoff is not None,
|
|
)
|
|
_entries: Final = (
|
|
("start_time", _model_call_details.get("start_time")),
|
|
("first_api_call_start_time", _first_handoff),
|
|
("call_type", _model_call_details.get("call_type")),
|
|
("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]),
|
|
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
|
|
("standard_logging_object", _model_call_details.get("standard_logging_object")),
|
|
)
|
|
return MappingProxyType({key: value for key, value in _entries if value is not None})
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _CallbackCapabilities:
|
|
"""Cached per-hook capability flags derived from ``litellm.callbacks``.
|
|
|
|
Recomputing this per request walked the callback list and resolved every
|
|
string entry via ``get_custom_logger_compatible_class`` — a measurable
|
|
chunk of overhead on streaming and non-streaming chat completions.
|
|
"""
|
|
|
|
has_post_call_response_headers: bool = False
|
|
has_iterator_override: bool = False
|
|
has_streaming_chunk_override: bool = False
|
|
has_guardrail: bool = False
|
|
has_pre_call_override: bool = False
|
|
has_content_enforcer: bool = False
|
|
# Tuple[(resolved_callback, "override" | "apply_guardrail"), ...]
|
|
# Ordered the same as ``litellm.callbacks``; used to build the streaming
|
|
# iterator chain without re-scanning per request.
|
|
iterator_overrides: tuple[tuple[Any, str], ...] = field(default_factory=tuple)
|
|
# Resolved CustomLogger callbacks in original order. Pre-resolving once
|
|
# avoids the per-request ``get_custom_logger_compatible_class`` walk for
|
|
# every string entry in ``litellm.callbacks``.
|
|
resolved_callbacks: tuple[object, ...] = field(default_factory=tuple)
|
|
|
|
|
|
class ProxyLogging:
|
|
"""
|
|
Logging/Custom Handlers for proxy.
|
|
|
|
Implemented mainly to:
|
|
- log successful/failed db read/writes
|
|
- support the max parallel request integration
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
user_api_key_cache: UserApiKeyCache,
|
|
premium_user: bool = False,
|
|
):
|
|
## INITIALIZE LITELLM CALLBACKS ##
|
|
self.call_details: dict = {}
|
|
self.call_details["user_api_key_cache"] = user_api_key_cache
|
|
self.internal_usage_cache: InternalUsageCache = InternalUsageCache(
|
|
dual_cache=DualCache(default_in_memory_ttl=1) # ping redis cache every 1s
|
|
)
|
|
self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache)
|
|
self.cache_control_check = _PROXY_CacheControlCheck()
|
|
self.alerting: list[str] | None = None
|
|
self.alerting_threshold: float = 300 # default to 5 min. threshold
|
|
self.alert_types: list[AlertType] = DEFAULT_ALERT_TYPES
|
|
self.alert_to_webhook_url: dict | None = None
|
|
self.slack_alerting_instance: SlackAlerting = SlackAlerting(
|
|
alerting_threshold=self.alerting_threshold,
|
|
alerting=self.alerting,
|
|
internal_usage_cache=self.internal_usage_cache.dual_cache,
|
|
)
|
|
self.email_logging_instance: Any | None = None
|
|
if BaseEmailLogger is not None:
|
|
email_logger_class: Final = _get_email_logger_class()
|
|
if email_logger_class is not None:
|
|
# All email logger classes now accept internal_usage_cache
|
|
self.email_logging_instance = email_logger_class(
|
|
internal_usage_cache=self.internal_usage_cache.dual_cache,
|
|
)
|
|
self.premium_user = premium_user
|
|
self.service_logging_obj = ServiceLogging()
|
|
self.db_spend_update_writer = DBSpendUpdateWriter()
|
|
self.proxy_hook_mapping: dict[str, CustomLogger] = {}
|
|
|
|
# Guard flags to prevent duplicate background tasks
|
|
self.daily_report_started: bool = False
|
|
self.hanging_requests_check_started: bool = False
|
|
self.deprecation_check_started: bool = False
|
|
|
|
def startup_event(
|
|
self,
|
|
llm_router: Router | None,
|
|
redis_usage_cache: RedisCache | None,
|
|
):
|
|
"""Initialize logging and alerting on proxy startup"""
|
|
## UPDATE SLACK ALERTING ##
|
|
self.slack_alerting_instance.update_values(llm_router=llm_router)
|
|
|
|
## UPDATE INTERNAL USAGE CACHE ##
|
|
self.update_values(
|
|
redis_cache=redis_usage_cache
|
|
) # used by parallel request limiter for rate limiting keys across instances
|
|
|
|
self._init_litellm_callbacks(
|
|
llm_router=llm_router
|
|
) # INITIALIZE LITELLM CALLBACKS ON SERVER STARTUP <- do this to catch any logging errors on startup, not when calls are being made
|
|
|
|
if (
|
|
self.slack_alerting_instance is not None
|
|
and "daily_reports" in self.slack_alerting_instance.alert_types
|
|
and not self.daily_report_started
|
|
):
|
|
asyncio.create_task(
|
|
self.slack_alerting_instance._run_scheduled_daily_report(
|
|
llm_router=llm_router,
|
|
pod_lock_manager=self.db_spend_update_writer.pod_lock_manager,
|
|
)
|
|
) # RUN DAILY REPORT (if scheduled)
|
|
self.daily_report_started = True
|
|
|
|
if (
|
|
self.slack_alerting_instance is not None
|
|
and AlertType.llm_requests_hanging in self.slack_alerting_instance.alert_types
|
|
and not self.hanging_requests_check_started
|
|
):
|
|
asyncio.create_task(
|
|
self.slack_alerting_instance.hanging_request_check.check_for_hanging_requests()
|
|
) # RUN HANGING REQUEST CHECK (if user wants to alert on hanging requests)
|
|
self.hanging_requests_check_started = True
|
|
|
|
self._ensure_deprecation_check_scheduled()
|
|
|
|
def _ensure_deprecation_check_scheduled(self) -> None:
|
|
"""Alerting can be configured at startup or by a later config reload, so schedule from either path"""
|
|
if self.alerting is None or self.deprecation_check_started:
|
|
return
|
|
|
|
try:
|
|
asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return
|
|
|
|
asyncio.create_task(
|
|
self.slack_alerting_instance.run_scheduled_deprecation_check(
|
|
pod_lock_manager=self.db_spend_update_writer.pod_lock_manager
|
|
)
|
|
)
|
|
self.deprecation_check_started = True
|
|
|
|
def update_values(
|
|
self,
|
|
alerting: list | None = None,
|
|
alerting_threshold: float | None = None,
|
|
redis_cache: RedisCache | None = None,
|
|
alert_types: list[AlertType] | None = None,
|
|
alerting_args: dict | None = None,
|
|
alert_to_webhook_url: dict | None = None,
|
|
alert_type_config: dict | None = None,
|
|
):
|
|
updated_slack_alerting: bool = False
|
|
if alerting is not None:
|
|
self.alerting = alerting
|
|
updated_slack_alerting = True
|
|
if alerting_threshold is not None:
|
|
self.alerting_threshold = alerting_threshold
|
|
updated_slack_alerting = True
|
|
if alert_types is not None:
|
|
self.alert_types = alert_types
|
|
updated_slack_alerting = True
|
|
if alert_to_webhook_url is not None:
|
|
self.alert_to_webhook_url = alert_to_webhook_url
|
|
updated_slack_alerting = True
|
|
if alert_type_config is not None:
|
|
updated_slack_alerting = True
|
|
|
|
if updated_slack_alerting is True:
|
|
self._ensure_deprecation_check_scheduled()
|
|
self.slack_alerting_instance.update_values(
|
|
alerting=self.alerting,
|
|
alerting_threshold=self.alerting_threshold,
|
|
alert_types=self.alert_types,
|
|
alerting_args=alerting_args,
|
|
alert_to_webhook_url=self.alert_to_webhook_url,
|
|
alert_type_config=alert_type_config,
|
|
)
|
|
|
|
if self.alerting is not None and ("slack" in self.alerting or "ms_teams" in self.alerting):
|
|
# NOTE: ENSURE we only add callbacks when alerting is on
|
|
# We should NOT add callbacks when alerting is off
|
|
if (
|
|
"daily_reports" in self.alert_types
|
|
or "outage_alerts" in self.alert_types
|
|
or "region_outage_alerts" in self.alert_types
|
|
):
|
|
litellm.logging_callback_manager.add_litellm_callback(self.slack_alerting_instance)
|
|
litellm.logging_callback_manager.add_litellm_success_callback(
|
|
self.slack_alerting_instance.response_taking_too_long_callback
|
|
)
|
|
|
|
if redis_cache is not None:
|
|
self.internal_usage_cache.dual_cache.redis_cache = redis_cache
|
|
self.db_spend_update_writer.redis_update_buffer.redis_cache = redis_cache
|
|
self.db_spend_update_writer.pod_lock_manager.redis_cache = redis_cache
|
|
|
|
def _add_proxy_hooks(self, llm_router: Router | None = None):
|
|
"""
|
|
Add proxy hooks to litellm.callbacks
|
|
"""
|
|
from litellm.proxy.proxy_server import prisma_client
|
|
|
|
for hook in PROXY_HOOKS:
|
|
proxy_hook = get_proxy_hook(hook)
|
|
expected_args = inspect.getfullargspec(proxy_hook).args
|
|
if "prisma_client" in expected_args and prisma_client is None:
|
|
verbose_proxy_logger.debug(
|
|
"Skipping proxy hook %s: it requires a database and no prisma client is configured", hook
|
|
)
|
|
continue
|
|
passed_in_args: dict[str, Any] = {}
|
|
if "internal_usage_cache" in expected_args:
|
|
passed_in_args["internal_usage_cache"] = self.internal_usage_cache
|
|
if "prisma_client" in expected_args:
|
|
passed_in_args["prisma_client"] = prisma_client
|
|
proxy_hook_obj = cast(CustomLogger, proxy_hook(**passed_in_args))
|
|
litellm.logging_callback_manager.add_litellm_callback(proxy_hook_obj)
|
|
|
|
self.proxy_hook_mapping[hook] = proxy_hook_obj
|
|
|
|
def get_proxy_hook(self, hook: str) -> CustomLogger | None:
|
|
"""
|
|
Get a proxy hook from the proxy_hook_mapping
|
|
"""
|
|
return self.proxy_hook_mapping.get(hook)
|
|
|
|
def _init_litellm_callbacks(self, llm_router: Router | None = None):
|
|
self._add_proxy_hooks(llm_router)
|
|
litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj)
|
|
|
|
# Track string callbacks and their initialized instances so we can
|
|
# replace them in-place, preventing duplicates (string + instance) in
|
|
# litellm.callbacks which caused double-counting of metrics.
|
|
string_callbacks_to_replace: Final[dict[int, CustomLogger]] = {}
|
|
|
|
for idx, callback in enumerate(litellm.callbacks):
|
|
if isinstance(callback, str):
|
|
initialized_callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback),
|
|
internal_usage_cache=self.internal_usage_cache.dual_cache,
|
|
llm_router=llm_router,
|
|
)
|
|
|
|
if initialized_callback is not None:
|
|
string_callbacks_to_replace[idx] = initialized_callback
|
|
|
|
# Replace string entries in litellm.callbacks with initialized instances
|
|
for idx, initialized_callback in string_callbacks_to_replace.items():
|
|
litellm.callbacks[idx] = initialized_callback
|
|
|
|
# Fan ``litellm.callbacks`` (the "all events" registry) out into the
|
|
# success/failure event lists eagerly, at startup. ``completion()`` does
|
|
# this lazily in ``function_setup`` on the first call, but request paths
|
|
# that build their own logging object and never run ``function_setup`` —
|
|
# notably pass-through endpoints — read ``litellm._async_success_callback``
|
|
# directly. Without this, a config-registered logger (e.g. ``otel``) is
|
|
# invisible to pass-through traffic until some other request warms the
|
|
# global lists. The manager dedupes, so this is idempotent with
|
|
# ``function_setup``.
|
|
for callback in litellm.callbacks:
|
|
if isinstance(callback, CustomLogger):
|
|
litellm.logging_callback_manager.add_litellm_success_callback(callback)
|
|
litellm.logging_callback_manager.add_litellm_failure_callback(callback)
|
|
litellm.logging_callback_manager.add_litellm_async_success_callback(callback)
|
|
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback)
|
|
|
|
# Runs after load_config applied every litellm_settings key: logger __init__s read e.g. s3_callback_params
|
|
success_callbacks: Final = tuple(cb for cb in litellm.success_callback if isinstance(cb, str))
|
|
failure_callbacks: Final = tuple(cb for cb in litellm.failure_callback if isinstance(cb, str))
|
|
for callback in success_callbacks:
|
|
_add_custom_logger_callback_to_specific_event(callback, "success")
|
|
for callback in failure_callbacks:
|
|
_add_custom_logger_callback_to_specific_event(callback, "failure")
|
|
|
|
async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]):
|
|
# only use this if slack alerting is being used
|
|
if self.alerting is None:
|
|
return
|
|
|
|
# current alerting threshold
|
|
alerting_threshold: float = self.alerting_threshold
|
|
|
|
# add a 100 second buffer to the alerting threshold
|
|
# ensures we don't send errant hanging request slack alerts
|
|
alerting_threshold += 100
|
|
|
|
await self.internal_usage_cache.async_set_cache(
|
|
key=f"request_status:{litellm_call_id}",
|
|
value=status,
|
|
local_only=True,
|
|
ttl=alerting_threshold,
|
|
litellm_parent_otel_span=None,
|
|
)
|
|
|
|
def _convert_user_api_key_auth_to_dict(self, user_api_key_auth_obj):
|
|
"""
|
|
Helper function to convert UserAPIKeyAuth object to dictionary.
|
|
Handles both Pydantic models and regular objects.
|
|
"""
|
|
if user_api_key_auth_obj is not None:
|
|
if hasattr(user_api_key_auth_obj, "model_dump"):
|
|
# If it's a Pydantic model, convert to dict
|
|
return user_api_key_auth_obj.model_dump()
|
|
elif hasattr(user_api_key_auth_obj, "__dict__"):
|
|
# If it's a regular object, convert to dict
|
|
return user_api_key_auth_obj.__dict__
|
|
return {}
|
|
|
|
def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict:
|
|
"""
|
|
Convert MCP tool call to LLM message format for existing guardrail validation.
|
|
"""
|
|
from litellm.types.llms.openai import ChatCompletionUserMessage
|
|
|
|
guardrail_context: Final = TypeAdapter(Mapping[str, object]).validate_python(
|
|
kwargs.get("guardrail_context") or MappingProxyType({})
|
|
)
|
|
|
|
parent_metadata: Final = copy.deepcopy(
|
|
TypeAdapter(dict[str, object]).validate_python(guardrail_context.get("metadata") or MappingProxyType({}))
|
|
)
|
|
|
|
# Create a synthetic message that represents the tool call
|
|
tool_call_content: Final = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
|
|
|
synthetic_message: Final = ChatCompletionUserMessage(role="user", content=tool_call_content)
|
|
|
|
synthetic_metadata: Final[dict[str, object]] = { # mutable-ok: existing guardrail hooks mutate request metadata
|
|
**MappingProxyType({key: value for key, value in parent_metadata.items() if key != "guardrails"}),
|
|
"headers": kwargs.get("headers") or {},
|
|
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
|
|
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
|
|
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
|
|
}
|
|
|
|
# Create synthetic LLM data that guardrails can process
|
|
synthetic_data: Final = {
|
|
"messages": [synthetic_message],
|
|
"model": guardrail_context.get("model", kwargs.get("model", "mcp-tool-call")),
|
|
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
|
|
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
|
|
"user_api_key_end_user_id": kwargs.get("user_api_key_end_user_id"),
|
|
"user_api_key_hash": kwargs.get("user_api_key_hash"),
|
|
"user_api_key_request_route": kwargs.get("user_api_key_request_route"),
|
|
"mcp_tool_name": request_obj.tool_name, # Keep original for reference
|
|
"mcp_arguments": request_obj.arguments, # Keep original for reference
|
|
# Surface the per-MCP-server rate-limit identity so the
|
|
# ParallelRequestLimiterV3 hook can apply mcp_rpm_limit on the
|
|
# synthetic call_mcp_tool payload (otherwise a key with
|
|
# mcp_rpm_limit could exceed it via the MCP path).
|
|
"mcp_server_name": kwargs.get("mcp_rate_limit_server_name"),
|
|
# Raw Bearer token from the original HTTP request — allows guardrails
|
|
# (e.g. MCPJWTSigner) to independently verify the caller's identity
|
|
# before re-signing an outbound token (FR-5 verify+re-sign).
|
|
"incoming_bearer_token": kwargs.get("incoming_bearer_token"),
|
|
"metadata": synthetic_metadata,
|
|
}
|
|
user_api_key_auth: Final = kwargs.get("user_api_key_auth")
|
|
if isinstance(user_api_key_auth, UserAPIKeyAuth):
|
|
add_guardrails_from_auth_metadata(
|
|
user_api_key_dict=user_api_key_auth,
|
|
data=synthetic_data,
|
|
metadata_variable_name="metadata",
|
|
)
|
|
synthetic_metadata["user_api_key_metadata"] = copy.deepcopy(user_api_key_auth.metadata)
|
|
synthetic_metadata["user_api_key_team_metadata"] = copy.deepcopy(user_api_key_auth.team_metadata)
|
|
merged_guardrails: Final = (
|
|
*TypeAdapter(tuple[object, ...]).validate_python(synthetic_metadata.get("guardrails") or ()),
|
|
*TypeAdapter(tuple[object, ...]).validate_python(parent_metadata.get("guardrails") or ()),
|
|
)
|
|
synthetic_metadata["guardrails"] = [ # mutable-ok: existing guardrail selection and policy hooks require a list
|
|
selection for index, selection in enumerate(merged_guardrails) if selection not in merged_guardrails[:index]
|
|
]
|
|
return synthetic_data
|
|
|
|
def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None:
|
|
"""
|
|
Convert LLM guardrail result back to MCP response format.
|
|
"""
|
|
from litellm.types.mcp import MCPPreCallResponseObject
|
|
|
|
# If result is an exception, it means the guardrail blocked the request
|
|
if isinstance(llm_result, Exception):
|
|
return MCPPreCallResponseObject(
|
|
should_proceed=False,
|
|
error_message=str(llm_result),
|
|
modified_arguments=None,
|
|
)
|
|
|
|
# If result is a dict with modified messages, check for content filtering
|
|
if isinstance(llm_result, dict):
|
|
modified_messages: Final = llm_result.get("messages")
|
|
if modified_messages:
|
|
# Check if content was blocked/modified
|
|
original_content: Final = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
|
new_content: Final = modified_messages[0].get("content", "") if modified_messages else ""
|
|
|
|
if new_content != original_content:
|
|
# Content was modified - could be masking, redaction, or blocking
|
|
if not new_content or "blocked" in new_content.lower() or "violation" in new_content.lower():
|
|
# Content was blocked completely
|
|
return MCPPreCallResponseObject(
|
|
should_proceed=False,
|
|
error_message="Content blocked by guardrail",
|
|
modified_arguments=None,
|
|
)
|
|
else:
|
|
# Content was masked/redacted - extract the modified arguments
|
|
try:
|
|
# Try to parse the modified arguments from the masked content
|
|
modified_args = self._extract_modified_arguments_from_content(new_content, request_obj)
|
|
if modified_args is not None:
|
|
# Return the masked/redacted arguments for the MCP call to use
|
|
return MCPPreCallResponseObject(
|
|
should_proceed=True,
|
|
error_message=None,
|
|
modified_arguments=modified_args,
|
|
)
|
|
else:
|
|
# Could not parse modified arguments, allow original call but warn
|
|
verbose_proxy_logger.warning(
|
|
"Could not parse modified arguments from guardrail response: %s", new_content
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Error parsing modified arguments: %s", e)
|
|
# Fallback: allow original call
|
|
return None
|
|
|
|
# If result is a string, it's likely an error message
|
|
if isinstance(llm_result, str):
|
|
return MCPPreCallResponseObject(should_proceed=False, error_message=llm_result, modified_arguments=None)
|
|
|
|
return None
|
|
|
|
def _extract_modified_arguments_from_content(self, masked_content: str, request_obj) -> dict | None:
|
|
"""
|
|
Extract modified/masked arguments from the guardrail response content.
|
|
"""
|
|
import json
|
|
|
|
verbose_proxy_logger.debug("Extracting modified args from content: %s", masked_content)
|
|
|
|
try:
|
|
# The format should be: "Tool: <tool_name>\nArguments: <json_arguments>"
|
|
# Parse the arguments section
|
|
lines: Final = masked_content.strip().split("\n")
|
|
for i, line in enumerate(lines):
|
|
if line.startswith("Arguments:"):
|
|
# Get the arguments part - everything after "Arguments: "
|
|
args_text = line[len("Arguments:") :].strip()
|
|
|
|
verbose_proxy_logger.debug("Found arguments text: %s", args_text)
|
|
|
|
# Try to parse as JSON first
|
|
try:
|
|
modified_args = json.loads(args_text)
|
|
verbose_proxy_logger.debug("Successfully parsed JSON args: %s", modified_args)
|
|
return modified_args
|
|
except json.JSONDecodeError as e:
|
|
# If JSON parsing fails, try to extract key-value pairs manually
|
|
verbose_proxy_logger.debug("Failed to parse JSON arguments: %s, error: %s", args_text, e)
|
|
return self._parse_arguments_manually(args_text, request_obj.arguments)
|
|
|
|
# If we can't find the Arguments: line, return None
|
|
verbose_proxy_logger.warning("Could not find 'Arguments:' line in masked content")
|
|
return None
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Error extracting modified arguments: %s", e)
|
|
return None
|
|
|
|
def _parse_arguments_manually(self, args_text: str, original_args: dict) -> dict | None:
|
|
"""
|
|
Try to manually parse arguments when JSON parsing fails.
|
|
This is a fallback for cases where the guardrail modifies the format.
|
|
"""
|
|
import re
|
|
|
|
try:
|
|
# Start with original arguments and try to apply modifications
|
|
modified_args: Final = original_args.copy()
|
|
|
|
# Look for simple key-value patterns
|
|
# This is a basic implementation - can be enhanced based on specific guardrail formats
|
|
for key, original_value in original_args.items():
|
|
if isinstance(original_value, str):
|
|
# Look for the key in the masked content and try to extract its value
|
|
pattern = rf"['\"]?{re.escape(key)}['\"]?\s*:\s*['\"]?([^,'\"]*)['\"]?"
|
|
match = re.search(pattern, args_text, re.IGNORECASE)
|
|
if match:
|
|
new_value = match.group(1).strip()
|
|
if new_value:
|
|
modified_args[key] = new_value
|
|
|
|
return modified_args
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Error in manual argument parsing: %s", e)
|
|
return None
|
|
|
|
def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> MCPDuringCallResponseObject | None:
|
|
"""
|
|
Convert LLM guardrail result back to MCP during call response format.
|
|
"""
|
|
# If result is an exception, it means the guardrail wants to stop execution
|
|
if isinstance(llm_result, Exception):
|
|
return MCPDuringCallResponseObject(should_continue=False, error_message=str(llm_result))
|
|
|
|
# If result is a dict with modified messages, check for content filtering
|
|
if isinstance(llm_result, dict):
|
|
modified_messages: Final = llm_result.get("messages")
|
|
if modified_messages:
|
|
# Check if content was blocked/modified
|
|
original_content: Final = f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
|
new_content: Final = modified_messages[0].get("content", "") if modified_messages else ""
|
|
|
|
if new_content != original_content:
|
|
# Content was modified, could be masking or blocking
|
|
if not new_content or "blocked" in new_content.lower():
|
|
# Content was blocked
|
|
return MCPDuringCallResponseObject(
|
|
should_continue=False,
|
|
error_message="Content blocked by guardrail during execution",
|
|
)
|
|
else:
|
|
# Content was masked/modified - for now, stop execution
|
|
return MCPDuringCallResponseObject(
|
|
should_continue=False,
|
|
error_message="Content modified by guardrail during execution",
|
|
)
|
|
|
|
# If result is a string, it's likely an error message
|
|
if isinstance(llm_result, str):
|
|
return MCPDuringCallResponseObject(should_continue=False, error_message=llm_result)
|
|
|
|
return None
|
|
|
|
def get_combined_callback_list(self, dynamic_success_callbacks: list | None, global_callbacks: list) -> list:
|
|
if dynamic_success_callbacks is None:
|
|
return list(global_callbacks)
|
|
return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks))
|
|
|
|
def _parse_pre_mcp_call_hook_response(
|
|
self,
|
|
response: MCPPreCallResponseObject,
|
|
original_request: MCPPreCallRequestObject,
|
|
) -> Mapping[str, object]:
|
|
"""
|
|
Parse the response from the pre_mcp_tool_call_hook
|
|
|
|
1. Check if the call should proceed
|
|
2. Apply any argument modifications
|
|
3. Handle validation errors
|
|
"""
|
|
result: Final = {
|
|
"should_proceed": response.should_proceed,
|
|
"modified_arguments": response.modified_arguments or original_request.arguments,
|
|
"error_message": response.error_message,
|
|
"hidden_params": response.hidden_params,
|
|
}
|
|
return result
|
|
|
|
def _create_mcp_request_object_from_kwargs(self, kwargs: dict) -> "MCPPreCallRequestObject":
|
|
"""
|
|
Helper function to create MCPPreCallRequestObject from kwargs for standard pre_call_hook.
|
|
"""
|
|
from litellm.types.llms.base import HiddenParams
|
|
from litellm.types.mcp import MCPPreCallRequestObject
|
|
|
|
user_api_key_auth_dict: Final = self._convert_user_api_key_auth_to_dict(kwargs.get("user_api_key_auth"))
|
|
|
|
return MCPPreCallRequestObject(
|
|
tool_name=kwargs.get("name", ""),
|
|
arguments=kwargs.get("arguments", {}),
|
|
server_name=kwargs.get("server_name"),
|
|
user_api_key_auth=user_api_key_auth_dict,
|
|
hidden_params=HiddenParams(),
|
|
)
|
|
|
|
def _convert_mcp_hook_response_to_kwargs(self, response_data: dict | None, original_kwargs: dict) -> dict:
|
|
"""
|
|
Helper function to convert pre_call_hook response back to kwargs for MCP usage.
|
|
|
|
Supports:
|
|
- modified_arguments: Override tool call arguments
|
|
- extra_headers: Inject custom headers into the outbound MCP request
|
|
"""
|
|
if not response_data:
|
|
return original_kwargs
|
|
|
|
modified_kwargs: Final = original_kwargs.copy()
|
|
|
|
if response_data.get("modified_arguments"):
|
|
modified_kwargs["arguments"] = response_data["modified_arguments"]
|
|
|
|
if response_data.get("extra_headers"):
|
|
# Merge rather than replace — a prior guardrail in the chain may have
|
|
# already injected headers (e.g. tracing IDs). Later guardrails win on
|
|
# key collisions so that the most-specific guardrail (e.g. JWT signer)
|
|
# takes precedence over earlier ones.
|
|
existing: Final = modified_kwargs.get("extra_headers") or {}
|
|
modified_kwargs["extra_headers"] = {
|
|
**existing,
|
|
**response_data["extra_headers"],
|
|
}
|
|
|
|
return modified_kwargs
|
|
|
|
async def process_pre_call_hook_response(self, response, data, call_type):
|
|
if isinstance(response, Exception):
|
|
raise response
|
|
if isinstance(response, dict):
|
|
return response
|
|
if isinstance(response, str):
|
|
if call_type in ["completion", "text_completion"]:
|
|
raise RejectedRequestError(
|
|
message=response,
|
|
model=data.get("model", ""),
|
|
llm_provider="",
|
|
request_data=data,
|
|
)
|
|
else:
|
|
raise HTTPException(status_code=400, detail={"error": response})
|
|
return data
|
|
|
|
def _should_use_guardrail_load_balancing(
|
|
self,
|
|
guardrail_name: str,
|
|
) -> bool:
|
|
"""
|
|
Check if load balancing should be used for this guardrail.
|
|
|
|
Returns True if the router has multiple deployments for this guardrail name.
|
|
"""
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
if llm_router is None or not hasattr(llm_router, "guardrail_list"):
|
|
return False
|
|
|
|
matching: Final = [g for g in llm_router.guardrail_list if g.get("guardrail_name") == guardrail_name]
|
|
return len(matching) > 1
|
|
|
|
async def _execute_guardrail_hook(
|
|
self,
|
|
callback: "CustomGuardrail",
|
|
hook_type: str,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth | None,
|
|
call_type: CallTypesLiteral,
|
|
response: LLMResponseTypes | None = None,
|
|
) -> object:
|
|
"""
|
|
Execute a single guardrail's hook.
|
|
|
|
Args:
|
|
callback: The guardrail callback to execute
|
|
hook_type: One of "pre_call", "during_call", "post_call"
|
|
data: Request data
|
|
user_api_key_dict: User API key auth
|
|
call_type: Type of call
|
|
response: Response object (for post_call hooks)
|
|
|
|
Returns:
|
|
Result from the guardrail execution
|
|
"""
|
|
# Use unified_guardrail if callback has apply_guardrail method
|
|
has_apply_guardrail: Final = "apply_guardrail" in type(callback).__dict__ and not getattr(
|
|
callback, "use_native_lifecycle_hooks", False
|
|
)
|
|
use_unified: Final = has_apply_guardrail and not (
|
|
hook_type == "during_call" and getattr(callback, "use_native_during_call_hook", False)
|
|
)
|
|
if use_unified:
|
|
data["guardrail_to_apply"] = callback
|
|
|
|
target: Final = unified_guardrail if use_unified else callback
|
|
|
|
if hook_type == "pre_call":
|
|
return await target.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=self.call_details["user_api_key_cache"],
|
|
data=data,
|
|
call_type=call_type,
|
|
)
|
|
elif hook_type == "during_call":
|
|
return await target.async_moderation_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
)
|
|
elif hook_type == "post_call":
|
|
return await target.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown hook_type: {hook_type}")
|
|
|
|
async def _execute_guardrail_with_load_balancing(
|
|
self,
|
|
guardrail_name: str,
|
|
hook_type: str,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth | None,
|
|
call_type: CallTypesLiteral,
|
|
response: LLMResponseTypes | None = None,
|
|
) -> object:
|
|
"""
|
|
Execute a guardrail using the router's load balancing.
|
|
|
|
Args:
|
|
guardrail_name: Name of the guardrail
|
|
hook_type: One of "pre_call", "during_call", "post_call"
|
|
data: Request data
|
|
user_api_key_dict: User API key auth
|
|
call_type: Type of call
|
|
response: Response object (for post_call hooks)
|
|
|
|
Returns:
|
|
Result from the guardrail execution
|
|
"""
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
if llm_router is None:
|
|
raise ValueError("Router not initialized")
|
|
|
|
# Select guardrail using router's load balancing
|
|
selected_guardrail: Final = llm_router.get_available_guardrail(guardrail_name=guardrail_name)
|
|
|
|
callback: Final[CustomGuardrail | None] = selected_guardrail.get("callback")
|
|
if callback is None:
|
|
raise ValueError(f"No callback found for guardrail: {guardrail_name}")
|
|
|
|
return await self._execute_guardrail_hook(
|
|
callback=callback,
|
|
hook_type=hook_type,
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
response=response,
|
|
)
|
|
|
|
async def _process_guardrail_callback(
|
|
self,
|
|
callback: CustomGuardrail,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth | None,
|
|
call_type: CallTypesLiteral,
|
|
event_type: GuardrailEventHooks,
|
|
) -> dict | None:
|
|
"""
|
|
Process a guardrail callback during pre-call hook.
|
|
|
|
Supports load balancing when multiple guardrail deployments exist.
|
|
|
|
Args:
|
|
callback: The CustomGuardrail callback to process
|
|
data: The request data dictionary
|
|
user_api_key_dict: User API key authentication details
|
|
call_type: The type of API call being made
|
|
|
|
Returns:
|
|
Updated data dictionary if guardrail passes, None if guardrail should be skipped
|
|
"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
# Determine the event type based on call type
|
|
if event_type is GuardrailEventHooks.pre_call and call_type == CallTypes.call_mcp_tool.value:
|
|
event_type = GuardrailEventHooks.pre_mcp_call
|
|
|
|
# Check if the guardrail should run for this request
|
|
if callback.should_run_guardrail(data=data, event_type=event_type) is not True:
|
|
return None
|
|
|
|
guardrail_name: Final = callback.guardrail_name
|
|
|
|
# Track timing and errors for prometheus metrics
|
|
# Use time.perf_counter() for more accurate duration measurements
|
|
guardrail_start_time: Final = time.perf_counter()
|
|
status = "success"
|
|
error_type = None
|
|
|
|
try:
|
|
# Check if load balancing should be used
|
|
if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name):
|
|
response = await self._execute_guardrail_with_load_balancing(
|
|
guardrail_name=guardrail_name,
|
|
hook_type="pre_call",
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
)
|
|
else:
|
|
# Single guardrail - execute directly
|
|
response = await self._execute_guardrail_hook(
|
|
callback=callback,
|
|
hook_type="pre_call",
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
)
|
|
|
|
# Process the response if one was returned
|
|
if response is not None:
|
|
data = await self.process_pre_call_hook_response(response=response, data=data, call_type=call_type)
|
|
|
|
callback.mark_pre_call_hook_ran(data)
|
|
|
|
except SensitiveDataRouteException:
|
|
status = "intervened"
|
|
raise
|
|
except Exception as e:
|
|
status = "error"
|
|
error_type = type(e).__name__
|
|
_enrich_http_exception_with_guardrail_context(e, callback)
|
|
# Re-raise the exception to maintain existing behavior
|
|
raise
|
|
finally:
|
|
# Record prometheus metrics
|
|
guardrail_end_time: Final = time.perf_counter()
|
|
latency_seconds: Final = guardrail_end_time - guardrail_start_time
|
|
|
|
# Get guardrail name for metrics (fallback if not set)
|
|
metrics_guardrail_name: Final = (
|
|
guardrail_name or getattr(callback, "guardrail_name", callback.__class__.__name__) or "unknown"
|
|
)
|
|
|
|
self._emit_guardrail_metrics(
|
|
guardrail_name=metrics_guardrail_name,
|
|
latency_seconds=latency_seconds,
|
|
status=status,
|
|
error_type=error_type,
|
|
hook_type="pre_call",
|
|
)
|
|
|
|
return data
|
|
|
|
async def _run_sequential_guardrail_callback(
|
|
self,
|
|
callback: CustomGuardrail,
|
|
data: dict, # mutable-ok: matches _process_guardrail_callback's own request-payload typing
|
|
raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
call_type: CallTypesLiteral,
|
|
) -> dict: # mutable-ok: callers reassign the loop's own data from this return value
|
|
"""
|
|
Run one guardrail from the sequential pre_call loop and return what the
|
|
rest of the loop should carry forward.
|
|
|
|
A guardrail opted into ``scan_raw_request`` always evaluates a fresh
|
|
copy of ``raw_request_snapshot`` (taken before any guardrail in this
|
|
hook ran) instead of ``data`` (the live, possibly already-mutated
|
|
payload), so its block/pass decision can never depend on where it's
|
|
declared relative to a guardrail that masks or rewrites content. It's
|
|
declared block-only, same contract as ``run_in_parallel``: any data it
|
|
returns is discarded, since applying its view on top of a stale
|
|
snapshot would silently undo whatever a later guardrail already did to
|
|
the live request. A guardrail that mutates content (e.g. PII masking)
|
|
should never set this flag -- if one does anyway, its returned
|
|
mutation is discarded and a warning is logged so the misconfiguration
|
|
is visible instead of silently forwarding unredacted content.
|
|
"""
|
|
scans_raw_request: Final = callback.scan_raw_request
|
|
should_use_raw_snapshot: Final = scans_raw_request and raw_request_snapshot is not None
|
|
input_data: Final = ( # mutable-ok: same request-payload shape as data
|
|
independent_snapshot(raw_request_snapshot) if should_use_raw_snapshot else data
|
|
)
|
|
# _process_guardrail_callback always calls mark_pre_call_hook_ran on a
|
|
# successful run, which unconditionally stamps bookkeeping metadata onto
|
|
# the dict regardless of whether the guardrail's own hook mutated
|
|
# anything -- so comparing `result` straight against `input_data` would
|
|
# warn on every single scan_raw_request call. Apply that same stamp to a
|
|
# throwaway, guaranteed-independent copy first (never the live request or
|
|
# raw_request_snapshot itself) so the comparison isolates the guardrail's
|
|
# own content mutation from this bookkeeping noise without risking a
|
|
# premature marker write into shared state.
|
|
expected_if_unmutated: Final[dict | None] = ( # mutable-ok: same request-payload shape as data
|
|
independent_snapshot(input_data) if scans_raw_request else None
|
|
)
|
|
if expected_if_unmutated is not None:
|
|
callback.mark_pre_call_hook_ran(expected_if_unmutated)
|
|
try:
|
|
result: Final = await self._process_guardrail_callback(
|
|
callback=callback,
|
|
data=input_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
except SensitiveDataRouteException:
|
|
raise
|
|
except Exception:
|
|
_record_raising_guardrail(data, callback)
|
|
raise
|
|
if (
|
|
scans_raw_request
|
|
and expected_if_unmutated is not None
|
|
and result is not None
|
|
and result != expected_if_unmutated
|
|
):
|
|
verbose_proxy_logger.warning(
|
|
"Guardrail '%s' has scan_raw_request=True but returned a modified payload; "
|
|
"scan_raw_request is for block-only guardrails and this mutation is being "
|
|
"discarded. Remove scan_raw_request from this guardrail's config if it needs "
|
|
"to mask/rewrite content.",
|
|
callback.guardrail_name or callback.__class__.__name__,
|
|
)
|
|
if scans_raw_request:
|
|
if result is not None:
|
|
# _process_guardrail_callback only stamped input_data (a throwaway
|
|
# snapshot copy), never the live data returned here -- without this,
|
|
# a deployment-level guardrail sharing this name would see no marker
|
|
# via _pre_call_hook_already_ran and re-run the same guardrail a
|
|
# second time on live kwargs.
|
|
callback.mark_pre_call_hook_ran(data)
|
|
return data
|
|
if result is None:
|
|
return data
|
|
return result
|
|
|
|
async def _process_prompt_template(
|
|
self,
|
|
data: dict,
|
|
litellm_logging_obj: "LiteLLMLoggingObj",
|
|
prompt_id: str,
|
|
prompt_version: int | None,
|
|
call_type: CallTypesLiteral,
|
|
) -> None:
|
|
"""Process prompt template if applicable."""
|
|
|
|
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
from litellm.utils import get_non_default_completion_params
|
|
|
|
raw_prompt_environment: Final = data.get("prompt_environment", None)
|
|
prompt_environment: Final = raw_prompt_environment if isinstance(raw_prompt_environment, str) else None
|
|
prompt_spec: Final = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec(
|
|
prompt_id,
|
|
version=prompt_version,
|
|
environment=prompt_environment,
|
|
)
|
|
custom_logger: Final = (
|
|
IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=prompt_spec)
|
|
if prompt_spec is not None
|
|
else None
|
|
)
|
|
litellm_prompt_id: str | None = None
|
|
if prompt_spec is not None:
|
|
litellm_prompt_id = prompt_spec.litellm_params.prompt_id
|
|
data.pop("prompt_id", None)
|
|
data.pop("prompt_environment", None)
|
|
|
|
if custom_logger and prompt_spec is not None:
|
|
is_responses_call: Final = call_type == "aresponses"
|
|
original_responses_input: Final = data.get("input", "") if is_responses_call else ""
|
|
client_messages: Final = (
|
|
ResponsesAPIRequestUtils.responses_input_to_chat_messages(original_responses_input)
|
|
if is_responses_call
|
|
else data.get("messages", [])
|
|
)
|
|
(
|
|
model,
|
|
messages,
|
|
optional_params,
|
|
) = await litellm_logging_obj.async_get_chat_completion_prompt(
|
|
model=data.get("model", ""),
|
|
messages=client_messages,
|
|
non_default_params=get_non_default_completion_params(kwargs=data) or {},
|
|
prompt_id=litellm_prompt_id,
|
|
prompt_spec=prompt_spec,
|
|
prompt_management_logger=custom_logger,
|
|
prompt_variables=data.pop("prompt_variables", None) or {},
|
|
prompt_label=data.pop("prompt_label", None) or {},
|
|
prompt_version=data.pop("prompt_version", None) or {},
|
|
request_kwargs=data,
|
|
injected_for_every_deployment=True,
|
|
)
|
|
|
|
data.update(optional_params)
|
|
data["model"] = model
|
|
if is_responses_call:
|
|
data["input"] = ResponsesAPIRequestUtils.merge_prompt_management_input(
|
|
original_input=original_responses_input,
|
|
client_input=client_messages,
|
|
merged_input=messages,
|
|
)
|
|
else:
|
|
data["messages"] = messages
|
|
# prevent re-processing the prompt template
|
|
data.pop("prompt_id", None)
|
|
data.pop("prompt_variables", None)
|
|
data.pop("prompt_label", None)
|
|
data.pop("prompt_version", None)
|
|
data.pop("prompt_environment", None)
|
|
|
|
def _process_guardrail_metadata(self, data: dict) -> None:
|
|
"""Process guardrails from metadata and add to applied_guardrails."""
|
|
from litellm.proxy.common_utils.callback_utils import (
|
|
add_guardrail_to_applied_guardrails_header,
|
|
)
|
|
|
|
metadata_standard: Final = data.get("metadata") or {}
|
|
metadata_litellm: Final = data.get("litellm_metadata") or {}
|
|
|
|
guardrails_in_metadata = []
|
|
if isinstance(metadata_standard, dict) and "guardrails" in metadata_standard:
|
|
guardrails_in_metadata = metadata_standard.get("guardrails", [])
|
|
elif isinstance(metadata_litellm, dict) and "guardrails" in metadata_litellm:
|
|
guardrails_in_metadata = metadata_litellm.get("guardrails", [])
|
|
|
|
if guardrails_in_metadata and isinstance(guardrails_in_metadata, list):
|
|
applied_guardrails = []
|
|
if isinstance(metadata_standard, dict) and "applied_guardrails" in metadata_standard:
|
|
applied_guardrails = metadata_standard.get("applied_guardrails", [])
|
|
elif isinstance(metadata_litellm, dict) and "applied_guardrails" in metadata_litellm:
|
|
applied_guardrails = metadata_litellm.get("applied_guardrails", [])
|
|
|
|
if not isinstance(applied_guardrails, list):
|
|
applied_guardrails = []
|
|
|
|
for guardrail_name in guardrails_in_metadata:
|
|
if isinstance(guardrail_name, str) and guardrail_name not in applied_guardrails:
|
|
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name)
|
|
|
|
async def _maybe_execute_pipelines(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
call_type: str,
|
|
event_hook: str,
|
|
raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data
|
|
response: LLMResponseTypes | None = None,
|
|
) -> tuple[dict, LLMResponseTypes | None]: # mutable-ok: returns the request-payload dict onward
|
|
"""
|
|
Execute guardrail pipelines if any are configured for this request.
|
|
|
|
Checks metadata for pipelines resolved by the policy engine
|
|
and executes them. Handles the result (allow/block/modify_response).
|
|
|
|
``raw_request_snapshot`` (taken before any guardrail or pipeline ran)
|
|
is forwarded so a pipeline step whose guardrail opted into
|
|
``scan_raw_request`` evaluates the pristine request, not whatever an
|
|
earlier ``pass_data`` step in the same pipeline already rewrote.
|
|
|
|
Returns the (possibly modified) data dict, plus the replacement
|
|
response when a post_call pipeline step returned one (None when the
|
|
response is unchanged), matching the flat callback-loop contract.
|
|
"""
|
|
pipelines: Final = _policy_pipelines(data)
|
|
if not pipelines:
|
|
return data, None
|
|
|
|
current_response = response # rebind-ok: chains each pipeline's replacement response into the next
|
|
for policy_name, pipeline in pipelines:
|
|
if pipeline.mode != event_hook:
|
|
continue
|
|
|
|
step_input: dict = (
|
|
{**data, "response": current_response} if current_response is not None else data
|
|
) # mutable-ok: same request-payload shape as data
|
|
|
|
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
|
|
steps=pipeline.steps,
|
|
mode=pipeline.mode,
|
|
data=step_input,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
policy_name=policy_name,
|
|
raw_request_snapshot=raw_request_snapshot,
|
|
)
|
|
|
|
data = self._handle_pipeline_result(
|
|
result=result,
|
|
data=data,
|
|
policy_name=policy_name,
|
|
original_response=current_response,
|
|
)
|
|
|
|
if current_response is not None and result.modified_data is not None:
|
|
current_response = result.modified_data.get("response", current_response)
|
|
|
|
return data, current_response if current_response is not response else None
|
|
|
|
@staticmethod
|
|
def _handle_pipeline_result(
|
|
result: PipelineExecutionResult,
|
|
data: dict,
|
|
policy_name: str,
|
|
original_response: "LLMResponseTypes | Sequence[object] | None" = None,
|
|
) -> dict:
|
|
"""
|
|
Handle a PipelineExecutionResult — allow, block, or modify_response.
|
|
|
|
Returns data dict if allowed, raises on block/modify_response.
|
|
``original_response`` is set on the post_call path, where the request
|
|
payload (already sent upstream) must stay untouched; a replacement
|
|
response carried in ``modified_data`` is adopted by the caller, and
|
|
metadata-bucket writes (applied guardrails, guardrail logging info)
|
|
are merged back so headers and spend logs still see them, on block
|
|
and modify_response too, so failure spend records keep guardrail
|
|
cost and status. On the
|
|
streaming path it is the buffered chunk list, carried into
|
|
``ModifyResponseException.original_response`` for usage reporting.
|
|
"""
|
|
if result.terminal_action == "allow":
|
|
if result.modified_data is not None:
|
|
if original_response is None:
|
|
data.update(result.modified_data)
|
|
else:
|
|
_merge_pipeline_metadata_writes(data, result.modified_data)
|
|
return data
|
|
|
|
if result.modified_data is not None:
|
|
_merge_pipeline_metadata_writes(data, result.modified_data)
|
|
|
|
if result.terminal_action == "block":
|
|
blocking_step: Final = result.step_results[-1] if result.step_results else None
|
|
callback: Final = (
|
|
PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name)
|
|
if blocking_step is not None
|
|
else None
|
|
)
|
|
if callback is not None:
|
|
_record_raising_guardrail(data, callback)
|
|
original_exception: Final = result.original_exception
|
|
if original_exception is not None and not _exception_changes_request_flow(original_exception):
|
|
if callback is not None:
|
|
_enrich_http_exception_with_guardrail_context(original_exception, callback)
|
|
raise original_exception
|
|
|
|
step_results_serializable: Final = [
|
|
{
|
|
"guardrail": sr.guardrail_name,
|
|
"outcome": sr.outcome,
|
|
"action": sr.action_taken,
|
|
}
|
|
for sr in result.step_results
|
|
]
|
|
error_detail: Final = {
|
|
"error": {
|
|
"message": f"Content blocked by guardrail pipeline '{policy_name}'",
|
|
"type": "guardrail_pipeline_error",
|
|
"pipeline_context": {
|
|
"policy": policy_name,
|
|
"step_results": step_results_serializable,
|
|
},
|
|
}
|
|
}
|
|
raise HTTPException(status_code=400, detail=error_detail)
|
|
|
|
if result.terminal_action == "modify_response":
|
|
raise ModifyResponseException(
|
|
message=result.modify_response_message or "Response modified by pipeline",
|
|
model=data.get("model", "unknown"),
|
|
request_data=data,
|
|
guardrail_name=f"pipeline:{policy_name}",
|
|
detection_info=None,
|
|
original_response=original_response,
|
|
)
|
|
|
|
return data
|
|
|
|
def has_pre_call_guardrails(self, request_metadata: Mapping[str, object]) -> bool:
|
|
"""
|
|
Whether anything configured would inspect the content of a request carrying this metadata.
|
|
|
|
Evaluated with the same predicate the pre-call loop uses, so a proxy configured only with
|
|
post-call guardrails answers False. Callers that must pay a real cost to build the hook's
|
|
input, such as streaming a batch input file off disk, use this to skip that work.
|
|
|
|
A content-enforcing ``CustomLogger`` counts too. It is not a guardrail and has no event
|
|
hook to consult, but it judges the payload the same way, so a proxy configured only with
|
|
one of those still has something to say about every record.
|
|
"""
|
|
if request_metadata.get("_guardrail_pipelines"):
|
|
return True
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
if caps.has_content_enforcer:
|
|
return True
|
|
probe: Final = {"metadata": dict(request_metadata)} # mutable-ok: should_run_guardrail takes a dict
|
|
return any(
|
|
isinstance(callback, CustomGuardrail)
|
|
and callback.should_run_guardrail(data=probe, event_type=GuardrailEventHooks.pre_call)
|
|
for callback in caps.resolved_callbacks
|
|
)
|
|
|
|
# The actual implementation of the function
|
|
@overload
|
|
async def pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
data: None,
|
|
call_type: CallTypesLiteral,
|
|
guardrails_only: bool = False,
|
|
) -> None:
|
|
pass
|
|
|
|
@overload
|
|
async def pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
data: dict,
|
|
call_type: CallTypesLiteral,
|
|
guardrails_only: bool = False,
|
|
) -> dict:
|
|
pass
|
|
|
|
async def pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
data: dict | None,
|
|
call_type: CallTypesLiteral,
|
|
guardrails_only: bool = False,
|
|
) -> dict | None:
|
|
"""
|
|
Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body.
|
|
|
|
Covers:
|
|
1. /chat/completions
|
|
2. /embeddings
|
|
3. /image/generation
|
|
|
|
With ``guardrails_only`` the walk is limited to guardrails and guardrail pipelines: rate
|
|
limiting, budget accounting, prompt templates and hanging-request alerting are skipped.
|
|
Use it to scan a payload that is not itself a request, such as one record of a batch file.
|
|
"""
|
|
verbose_proxy_logger.debug("Inside Proxy Logging Pre-call hook!")
|
|
|
|
if not guardrails_only:
|
|
self._init_response_taking_too_long_task(data=data)
|
|
|
|
if data is None:
|
|
return None
|
|
|
|
litellm_logging_obj: Final = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None))
|
|
prompt_id: Final[str | None] = data.get("prompt_id", None)
|
|
|
|
## PROMPT TEMPLATE CHECK ##
|
|
|
|
if (
|
|
not guardrails_only
|
|
and litellm_logging_obj is not None
|
|
and prompt_id is not None
|
|
and (call_type == "completion" or call_type == "acompletion" or call_type == "aresponses")
|
|
):
|
|
from litellm.proxy.prompts.prompt_registry import parse_prompt_version
|
|
|
|
await self._process_prompt_template(
|
|
data=data,
|
|
litellm_logging_obj=litellm_logging_obj,
|
|
prompt_id=prompt_id,
|
|
prompt_version=parse_prompt_version(data.get("prompt_version", None)),
|
|
call_type=call_type,
|
|
)
|
|
|
|
# Snapshotted here, before _maybe_execute_pipelines or any guardrail in
|
|
# this hook has run, so a scan_raw_request guardrail's block/pass
|
|
# decision never depends on its position in the guardrails list or on
|
|
# a pipeline that runs ahead of it: an earlier guardrail (pipelined or
|
|
# not) that masks/rewrites content can't hide a violation from a later
|
|
# one that opted into scanning the original request. Only computed
|
|
# when at least one registered guardrail actually opted in, and via
|
|
# independent_snapshot (not safe_deep_copy) since this isolation
|
|
# guarantee must hold even under litellm.safe_memory_mode, which
|
|
# otherwise makes deep copies return the original object.
|
|
needs_raw_request_snapshot: Final = any(
|
|
isinstance(cb, CustomGuardrail) and cb.scan_raw_request
|
|
for cb in ProxyLogging._callback_capabilities().resolved_callbacks
|
|
)
|
|
raw_request_snapshot: Final[dict | None] = ( # mutable-ok: same request-payload shape as data
|
|
independent_snapshot(data) if needs_raw_request_snapshot else None
|
|
)
|
|
|
|
try:
|
|
# Execute guardrail pipelines before the normal callback loop
|
|
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
event_hook="pre_call",
|
|
raw_request_snapshot=raw_request_snapshot,
|
|
)
|
|
|
|
# Get pipeline-managed guardrails to skip in normal loop
|
|
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "pre_call")
|
|
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
# Skip the per-request callback walk entirely when nothing in
|
|
# ``litellm.callbacks`` overrides ``async_pre_call_hook`` and no
|
|
# CustomGuardrail is configured. Saves the loop overhead +
|
|
# ``time.time()`` x2 per registered callback for the common
|
|
# "callbacks=[]" case on small / dev deployments.
|
|
if (
|
|
not caps.has_guardrail
|
|
and not caps.has_content_enforcer
|
|
and (guardrails_only or not caps.has_pre_call_override)
|
|
):
|
|
if data is not None:
|
|
self._process_guardrail_metadata(data)
|
|
return data
|
|
|
|
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
|
|
cb
|
|
for cb in caps.resolved_callbacks
|
|
if isinstance(cb, CustomGuardrail)
|
|
and getattr(cb, "run_in_parallel", False)
|
|
and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed)
|
|
)
|
|
|
|
deferred_route_exc: SensitiveDataRouteException | None = None
|
|
for _callback in caps.resolved_callbacks:
|
|
start_time = time.time()
|
|
try:
|
|
if isinstance(_callback, CustomGuardrail) and data is not None:
|
|
# Skip guardrails managed by a pipeline
|
|
if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed:
|
|
continue
|
|
|
|
if getattr(_callback, "run_in_parallel", False):
|
|
continue
|
|
|
|
data = await self._run_sequential_guardrail_callback(
|
|
callback=_callback,
|
|
data=data,
|
|
raw_request_snapshot=raw_request_snapshot,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
)
|
|
|
|
elif (
|
|
_callback is not None
|
|
and isinstance(_callback, CustomLogger)
|
|
and (not guardrails_only or _callback.enforces_request_content)
|
|
and "async_pre_call_hook" in vars(_callback.__class__)
|
|
and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook
|
|
):
|
|
if call_type == "call_mcp_tool" and user_api_key_dict is None:
|
|
continue
|
|
|
|
response: Exception | str | Mapping[str, object] | None = await _callback.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=self.call_details["user_api_key_cache"],
|
|
data=data,
|
|
call_type=call_type,
|
|
)
|
|
if response is not None:
|
|
data = await self.process_pre_call_hook_response(
|
|
response=response, data=data, call_type=call_type
|
|
)
|
|
except SensitiveDataRouteException as e:
|
|
# Defer the reroute until remaining guardrails have run so later
|
|
# security checks are not skipped; the first reroute wins and a
|
|
# later guardrail that blocks still propagates. Fall through to the
|
|
# service-span recording below so the triggering guardrail is still
|
|
# timed like every other callback.
|
|
if deferred_route_exc is None:
|
|
deferred_route_exc = e
|
|
|
|
end_time = time.time()
|
|
duration = end_time - start_time
|
|
if (
|
|
hasattr(self, "service_logging_obj") and duration > 0.01
|
|
): # only if duration is non-negligible - don't spam the logs
|
|
await self.service_logging_obj.async_service_success_hook(
|
|
service=ServiceTypes.PROXY_PRE_CALL,
|
|
duration=duration,
|
|
call_type=f"{_callback.__class__.__name__}",
|
|
parent_otel_span=user_api_key_dict.parent_otel_span,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
|
|
if deferred_route_exc is not None and data is not None:
|
|
data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict)
|
|
|
|
if parallel_guardrails and data is not None:
|
|
await self._run_parallel_pre_call_guardrails(
|
|
guardrails=parallel_guardrails,
|
|
data=data,
|
|
raw_request_snapshot=raw_request_snapshot,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
)
|
|
|
|
if data is not None:
|
|
self._process_guardrail_metadata(data)
|
|
|
|
return data
|
|
except SensitiveDataRouteException as e:
|
|
data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict)
|
|
if data is not None:
|
|
self._process_guardrail_metadata(data)
|
|
return data
|
|
except Exception:
|
|
if data is not None:
|
|
self._process_guardrail_metadata(data)
|
|
raise
|
|
|
|
async def _run_parallel_pre_call_guardrails(
|
|
self,
|
|
guardrails: tuple[CustomGuardrail, ...],
|
|
data: dict,
|
|
raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
call_type: CallTypesLiteral,
|
|
) -> None:
|
|
"""
|
|
Run opted-in pre_call guardrails concurrently against one shared payload
|
|
snapshot. These guardrails are declared block-only, so any modified data
|
|
they return is discarded; they run for their blocking side effect (raising
|
|
to reject the request before it reaches the LLM). Every guardrail is
|
|
awaited to completion (``return_exceptions=True``) so a raise by one never
|
|
leaves the others running as unobserved background tasks. A guardrail that
|
|
blocks (any exception other than a reroute or passthrough) takes precedence
|
|
over one that only changes the request flow, so a fast reroute can never
|
|
let a slower block be bypassed; the request is rejected before it reaches
|
|
the LLM, preserving the pre-call barrier that ``during_call`` guardrails
|
|
cannot provide. Per-guardrail latency is recorded by
|
|
``_process_guardrail_callback``'s own metrics.
|
|
|
|
A guardrail that also opted into ``scan_raw_request`` evaluates
|
|
``raw_request_snapshot`` (taken before the sequential loop ran) instead
|
|
of ``data`` (the sequential loop's output), for the same reason the
|
|
sequential branch does: its block decision must not depend on what a
|
|
sequential guardrail already masked or rewrote.
|
|
"""
|
|
|
|
def _input_for(callback: CustomGuardrail) -> dict: # mutable-ok: same request-payload shape as data
|
|
if not callback.scan_raw_request or raw_request_snapshot is None:
|
|
return data
|
|
return independent_snapshot(raw_request_snapshot)
|
|
|
|
results: Final = await asyncio.gather(
|
|
*(
|
|
self._process_guardrail_callback(
|
|
callback=callback,
|
|
data=_input_for(callback),
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
for callback in guardrails
|
|
),
|
|
return_exceptions=True,
|
|
)
|
|
for callback, result in zip(guardrails, results, strict=True):
|
|
# _process_guardrail_callback stamped mark_pre_call_hook_ran on
|
|
# _input_for's throwaway snapshot copy for a scan_raw_request
|
|
# guardrail, never on the live, shared `data` -- without this, a
|
|
# deployment-level guardrail sharing this name would see no marker
|
|
# via _pre_call_hook_already_ran and re-run it a second time on
|
|
# live kwargs.
|
|
if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None:
|
|
callback.mark_pre_call_hook_ran(data)
|
|
if isinstance(result, BaseException) and not isinstance(result, SensitiveDataRouteException):
|
|
_record_raising_guardrail(data, callback)
|
|
raised: Final = tuple(result for result in results if isinstance(result, BaseException))
|
|
blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
|
|
if blocking is not None:
|
|
raise blocking
|
|
if raised:
|
|
raise raised[0]
|
|
|
|
async def _handle_sensitive_data_route_exception(
|
|
self,
|
|
exc: SensitiveDataRouteException,
|
|
data: dict | None,
|
|
user_api_key_dict: UserAPIKeyAuth | None,
|
|
) -> dict | None:
|
|
"""
|
|
Handle SensitiveDataRouteException by rerouting the current request to
|
|
the target model and, when sticky_session_routing is enabled, persisting
|
|
the session override so subsequent requests reuse the same model.
|
|
"""
|
|
if data is None:
|
|
return None
|
|
|
|
verbose_proxy_logger.info(
|
|
"SensitiveDataRouteException caught: session_id=%s route_to_model=%s guardrail=%s sticky=%s",
|
|
exc.session_id,
|
|
exc.route_to_model,
|
|
exc.guardrail_name,
|
|
exc.sticky_session_routing,
|
|
)
|
|
|
|
if exc.sticky_session_routing:
|
|
sensitive_routing_hook: Final = self.get_proxy_hook("sensitive_data_routing")
|
|
if isinstance(sensitive_routing_hook, _PROXY_SensitiveDataRoutingHandler):
|
|
await sensitive_routing_hook.set_session_routing(
|
|
session_id=exc.session_id,
|
|
model=exc.route_to_model,
|
|
user_api_key_dict=user_api_key_dict,
|
|
guardrail_name=exc.guardrail_name,
|
|
)
|
|
else:
|
|
verbose_proxy_logger.warning(
|
|
"SensitiveDataRouteException requested sticky routing for session_id=%s "
|
|
"but the 'sensitive_data_routing' hook is not registered. Only this request "
|
|
"will be rerouted; subsequent requests will not be sticky.",
|
|
exc.session_id,
|
|
)
|
|
|
|
original_model: Final = data.get("model")
|
|
data["model"] = exc.route_to_model
|
|
|
|
metadata: Final = data.get("metadata") or {}
|
|
metadata["sensitive_data_routing_applied"] = True
|
|
metadata["sensitive_data_routing_original_model"] = original_model
|
|
metadata["sensitive_data_routing_guardrail"] = exc.guardrail_name
|
|
metadata["sensitive_data_routing_detection_info"] = exc.detection_info
|
|
data["metadata"] = metadata
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def _emit_guardrail_metrics(
|
|
guardrail_name: str,
|
|
latency_seconds: float,
|
|
status: str,
|
|
error_type: str | None,
|
|
hook_type: str,
|
|
) -> None:
|
|
for prom_callback in litellm.callbacks:
|
|
if isinstance(prom_callback, PrometheusLogger):
|
|
prom_callback._record_guardrail_metrics(
|
|
guardrail_name=guardrail_name,
|
|
latency_seconds=latency_seconds,
|
|
status=status,
|
|
error_type=error_type,
|
|
hook_type=hook_type,
|
|
)
|
|
break
|
|
|
|
@staticmethod
|
|
async def _run_guardrail_with_metrics(
|
|
callback: object,
|
|
coro: Awaitable[_T],
|
|
hook_type: str,
|
|
request_data: Mapping[str, object],
|
|
) -> _T:
|
|
"""
|
|
Await `coro`, recording its latency and status to the
|
|
`litellm_guardrail_latency_seconds` metric under `hook_type`, and
|
|
enriching any raised HTTPException with the originating callback's
|
|
`guardrail_name`/`guardrail_mode` before re-raising.
|
|
"""
|
|
guardrail_name: Final = getattr(callback, "guardrail_name", None) or type(callback).__name__
|
|
start_time: Final = time.perf_counter()
|
|
status = "success"
|
|
error_type: str | None = None
|
|
try:
|
|
return await coro
|
|
except SensitiveDataRouteException:
|
|
status = "intervened"
|
|
raise
|
|
except Exception as e:
|
|
status = "error"
|
|
error_type = type(e).__name__
|
|
_enrich_http_exception_with_guardrail_context(e, callback)
|
|
_record_raising_guardrail(request_data, callback)
|
|
raise
|
|
finally:
|
|
ProxyLogging._emit_guardrail_metrics(
|
|
guardrail_name=guardrail_name,
|
|
latency_seconds=time.perf_counter() - start_time,
|
|
status=status,
|
|
error_type=error_type,
|
|
hook_type=hook_type,
|
|
)
|
|
|
|
@staticmethod
|
|
async def _wrap_streaming_iterator_with_enrichment(
|
|
callback: object,
|
|
response: AsyncIterable[_T],
|
|
hook: _StreamIteratorHook[_T],
|
|
request_data: Mapping[str, object],
|
|
) -> AsyncGenerator[_T, None]:
|
|
upstream: Final = _UpstreamStreamBoundary(response)
|
|
try:
|
|
async for chunk in hook(response=upstream):
|
|
yield chunk
|
|
except Exception as e:
|
|
if e is not upstream.failure:
|
|
_enrich_http_exception_with_guardrail_context(e, callback)
|
|
_record_raising_guardrail(request_data, callback)
|
|
raise
|
|
|
|
# Cache for callback-capability detection. Keyed on a signature of
|
|
# litellm.callbacks (length + each item's id) so we recompute when the
|
|
# callback list mutates (add/remove) without iterating every request.
|
|
_callback_capabilities_cache: ClassVar[dict[tuple[int, tuple[int, ...]], "_CallbackCapabilities"]] = {}
|
|
|
|
@staticmethod
|
|
def _callback_capabilities() -> "_CallbackCapabilities":
|
|
"""
|
|
Inspect ``litellm.callbacks`` once and answer the per-hook capability
|
|
questions used to short-circuit no-op work on the chat-completions hot
|
|
path. Per-request callers iterated ``litellm.callbacks`` and called
|
|
``get_custom_logger_compatible_class`` for every string entry — that
|
|
scanning cost dominated the proxy overhead on low-config deployments.
|
|
|
|
Cache invalidates whenever the list length or member identities change.
|
|
"""
|
|
callbacks: Final = litellm.callbacks
|
|
sig: Final = (len(callbacks), tuple(id(c) for c in callbacks))
|
|
cache: Final = ProxyLogging._callback_capabilities_cache
|
|
cached: Final = cache.get(sig)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
has_post_call_response_headers = False
|
|
has_iterator_override = False
|
|
has_streaming_chunk_override = False
|
|
has_guardrail = False
|
|
has_pre_call_override = False
|
|
has_content_enforcer = False
|
|
iterator_overrides: Final[list[tuple[Any, str]]] = [] # (callback, kind)
|
|
resolved_callbacks: Final[list[CustomLogger]] = []
|
|
|
|
for callback in callbacks:
|
|
if isinstance(callback, str):
|
|
resolved = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback)
|
|
)
|
|
else:
|
|
resolved = callback
|
|
if resolved is None or not isinstance(resolved, CustomLogger):
|
|
continue
|
|
resolved_callbacks.append(resolved)
|
|
cls = type(resolved)
|
|
if cls is CustomLogger:
|
|
continue
|
|
if isinstance(resolved, CustomGuardrail):
|
|
has_guardrail = True
|
|
# Use the same leaf-class ``__dict__`` check as the other hook
|
|
# capabilities: only callbacks that actually override the hook
|
|
# contribute to the flag. Setting this for every ``CustomLogger``
|
|
# instance (the prior behaviour) forced the full
|
|
# ``post_call_response_headers_hook`` body to run on every request
|
|
# even when no registered callback customized response headers.
|
|
cls_attrs = cls.__dict__
|
|
if "async_post_call_response_headers_hook" in cls_attrs:
|
|
has_post_call_response_headers = True
|
|
if "async_post_call_streaming_iterator_hook" in cls_attrs:
|
|
has_iterator_override = True
|
|
iterator_overrides.append((resolved, "override"))
|
|
elif "apply_guardrail" in cls_attrs and not getattr(resolved, "use_native_lifecycle_hooks", False):
|
|
iterator_overrides.append((resolved, "apply_guardrail"))
|
|
# Walk the MRO for ``async_post_call_streaming_hook`` rather than
|
|
# using the leaf-class ``__dict__`` check used by the other flags:
|
|
# before this PR the hook was unconditionally invoked, so a
|
|
# callback that inherits an override from an intermediate parent
|
|
# (e.g. a vendor base class providing the override, with the
|
|
# registered class adding nothing else) MUST still be detected.
|
|
# A leaf-class miss here would silently drop the inherited hook.
|
|
base_streaming_hook = CustomLogger.async_post_call_streaming_hook
|
|
cls_streaming_hook = getattr(
|
|
cls,
|
|
"async_post_call_streaming_hook",
|
|
base_streaming_hook,
|
|
)
|
|
if getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr(
|
|
base_streaming_hook, "__func__", base_streaming_hook
|
|
):
|
|
has_streaming_chunk_override = True
|
|
if "async_pre_call_hook" in cls_attrs:
|
|
has_pre_call_override = True
|
|
if resolved.enforces_request_content is True:
|
|
has_content_enforcer = True
|
|
|
|
caps: Final = _CallbackCapabilities(
|
|
has_post_call_response_headers=has_post_call_response_headers,
|
|
has_iterator_override=has_iterator_override
|
|
or any(kind == "apply_guardrail" for _, kind in iterator_overrides),
|
|
has_streaming_chunk_override=has_streaming_chunk_override,
|
|
has_guardrail=has_guardrail,
|
|
has_pre_call_override=has_pre_call_override,
|
|
has_content_enforcer=has_content_enforcer,
|
|
iterator_overrides=tuple(iterator_overrides),
|
|
resolved_callbacks=tuple(resolved_callbacks),
|
|
)
|
|
# Limit cache to handle test churn without leaking; production
|
|
# callback lists are stable so this rarely grows past 1 entry.
|
|
if len(cache) >= 32:
|
|
cache.clear()
|
|
cache[sig] = caps
|
|
return caps
|
|
|
|
@staticmethod
|
|
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
|
route: Final = user_api_key_dict.request_route
|
|
if not route:
|
|
return False
|
|
call_types: Final = get_call_types_for_route(route)
|
|
if not call_types:
|
|
return False
|
|
return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES
|
|
|
|
@staticmethod
|
|
def has_post_call_response_headers_callbacks() -> bool:
|
|
return ProxyLogging._callback_capabilities().has_post_call_response_headers
|
|
|
|
@staticmethod
|
|
def has_streaming_callbacks() -> bool:
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
return caps.has_iterator_override or caps.has_streaming_chunk_override or caps.has_guardrail
|
|
|
|
@staticmethod
|
|
def has_streaming_chunk_hook_overrides() -> bool:
|
|
"""True iff any callback overrides ``async_post_call_streaming_hook``
|
|
(the per-chunk hook, distinct from the iterator wrapper)."""
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
return caps.has_streaming_chunk_override or caps.has_guardrail
|
|
|
|
def needs_iterator_wrap(self) -> bool:
|
|
"""Whether ``async_data_generator`` needs to wrap the upstream stream
|
|
through ``async_post_call_streaming_iterator_hook``. Instance method
|
|
so tests can override the gate via ``MagicMock(spec=ProxyLogging)``.
|
|
"""
|
|
return ProxyLogging._callback_capabilities().has_iterator_override
|
|
|
|
def needs_per_chunk_streaming_hook(self) -> bool:
|
|
"""Whether ``async_data_generator`` needs to call the per-chunk
|
|
``_apply_streaming_chunk_hooks`` for every emitted chunk. Instance
|
|
method for the same reason as :py:meth:`needs_iterator_wrap`.
|
|
"""
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
return caps.has_streaming_chunk_override or caps.has_guardrail
|
|
|
|
@staticmethod
|
|
def has_during_call_guardrails() -> bool:
|
|
return ProxyLogging._callback_capabilities().has_guardrail
|
|
|
|
async def during_call_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth | None,
|
|
call_type: CallTypesLiteral,
|
|
):
|
|
"""
|
|
Runs the CustomGuardrail's async_moderation_hook() in parallel
|
|
"""
|
|
# Fast path: skip the entire guardrail scan when no CustomGuardrail
|
|
# callbacks are registered. Saves per-request iteration over
|
|
# ``litellm.callbacks`` plus an ``asyncio.gather([])`` round trip on
|
|
# deployments with no guardrails configured.
|
|
if not ProxyLogging._callback_capabilities().has_guardrail:
|
|
return data
|
|
# Step 1: Collect all guardrail tasks to run in parallel
|
|
guardrail_tasks: Final = []
|
|
|
|
for callback in litellm.callbacks:
|
|
if isinstance(callback, CustomGuardrail):
|
|
################################################################
|
|
# Check if guardrail should be run for GuardrailEventHooks.during_call hook
|
|
################################################################
|
|
|
|
# V1 implementation - backwards compatibility
|
|
if callback.event_hook is None and hasattr(callback, "moderation_check"):
|
|
if callback.moderation_check == "pre_call":
|
|
return
|
|
else:
|
|
# Main - V2 Guardrails implementation
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
event_type = GuardrailEventHooks.during_call
|
|
if call_type == CallTypes.call_mcp_tool.value:
|
|
event_type = GuardrailEventHooks.during_mcp_call
|
|
|
|
if callback.should_run_guardrail(data=data, event_type=event_type) is not True:
|
|
continue
|
|
# Convert user_api_key_dict to proper format for async_moderation_hook
|
|
if call_type == CallTypes.call_mcp_tool.value:
|
|
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict)
|
|
else:
|
|
user_api_key_auth_dict = user_api_key_dict
|
|
guardrail_tasks.append(
|
|
self._run_during_call_guardrail(
|
|
callback=callback,
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
user_api_key_auth_dict=user_api_key_auth_dict,
|
|
call_type=call_type,
|
|
)
|
|
)
|
|
|
|
# Step 2: Run all guardrail tasks in parallel
|
|
if guardrail_tasks:
|
|
try:
|
|
await asyncio.gather(*guardrail_tasks)
|
|
except Exception as e:
|
|
# If any guardrail raises an exception, it will propagate here
|
|
raise e
|
|
|
|
return data
|
|
|
|
async def _run_during_call_guardrail(
|
|
self,
|
|
callback: CustomGuardrail,
|
|
data: dict[str, object], # mutable-ok: request payload dict, guardrail_to_apply is written in place
|
|
user_api_key_dict: UserAPIKeyAuth | None,
|
|
user_api_key_auth_dict: UserAPIKeyAuth | dict[str, object] | None,
|
|
call_type: CallTypesLiteral,
|
|
) -> None:
|
|
if (
|
|
"apply_guardrail" in type(callback).__dict__
|
|
and not callback.use_native_lifecycle_hooks
|
|
and user_api_key_dict is not None
|
|
and not callback.use_native_during_call_hook
|
|
):
|
|
data["guardrail_to_apply"] = callback
|
|
await self._run_guardrail_with_metrics(
|
|
callback,
|
|
unified_guardrail.async_moderation_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
call_type=call_type,
|
|
),
|
|
"during_call",
|
|
request_data=data,
|
|
)
|
|
return
|
|
await self._run_guardrail_with_metrics(
|
|
callback,
|
|
callback.async_moderation_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_auth_dict,
|
|
call_type=call_type,
|
|
),
|
|
"during_call",
|
|
request_data=data,
|
|
)
|
|
|
|
async def failed_tracking_alert(
|
|
self,
|
|
error_message: str,
|
|
failing_model: str,
|
|
):
|
|
if self.alerting is None:
|
|
return
|
|
|
|
if self.slack_alerting_instance:
|
|
await self.slack_alerting_instance.failed_tracking_alert(
|
|
error_message=error_message,
|
|
failing_model=failing_model,
|
|
)
|
|
|
|
async def budget_alerts(
|
|
self,
|
|
type: Literal[
|
|
"token_budget",
|
|
"user_budget",
|
|
"soft_budget",
|
|
"max_budget_alert",
|
|
"team_budget",
|
|
"organization_budget",
|
|
"proxy_budget",
|
|
"projected_limit_exceeded",
|
|
"project_budget",
|
|
],
|
|
user_info: CallInfo,
|
|
):
|
|
# For soft_budget alerts with alert_emails set, allow email sending even if alerting is None
|
|
# This enables team-specific soft budget email alerts via metadata.soft_budget_alerting_emails
|
|
# Note: user_info is a CallInfo that can represent user/team/org level info. For team budgets,
|
|
# alert_emails is populated from team_object.metadata.soft_budget_alerting_emails (see auth_checks.py)
|
|
is_soft_budget_with_alert_emails: Final = (
|
|
type == "soft_budget" and user_info.alert_emails is not None and len(user_info.alert_emails) > 0
|
|
)
|
|
|
|
if self.alerting is None and not is_soft_budget_with_alert_emails:
|
|
# do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails)
|
|
return
|
|
|
|
if self.alerting is not None and (
|
|
"slack" in self.alerting or "ms_teams" in self.alerting or "webhook" in self.alerting
|
|
):
|
|
if self.slack_alerting_instance is not None:
|
|
await self.slack_alerting_instance.budget_alerts(
|
|
type=type,
|
|
user_info=user_info,
|
|
)
|
|
|
|
# Call email_logging_instance if:
|
|
# 1. "email" is in alerting config, OR
|
|
# 2. It's a soft_budget alert with team-specific alert_emails (bypasses global alerting config)
|
|
should_send_email = (self.alerting is not None and "email" in self.alerting) or is_soft_budget_with_alert_emails
|
|
|
|
if should_send_email and self.email_logging_instance is not None:
|
|
await self.email_logging_instance.budget_alerts(
|
|
type=type,
|
|
user_info=user_info,
|
|
)
|
|
|
|
async def alerting_handler(
|
|
self,
|
|
message: str,
|
|
level: Literal["Low", "Medium", "High"],
|
|
alert_type: AlertType,
|
|
request_data: dict | None = None,
|
|
):
|
|
"""
|
|
Alerting based on thresholds: - https://github.com/BerriAI/litellm/issues/1298
|
|
|
|
- Responses taking too long
|
|
- Requests are hanging
|
|
- Calls are failing
|
|
- DB Read/Writes are failing
|
|
- Proxy Close to max budget
|
|
- Key Close to max budget
|
|
|
|
Parameters:
|
|
level: str - Low|Medium|High - if calls might fail (Medium) or are failing (High); Currently, no alerts would be 'Low'.
|
|
message: str - what is the alert about
|
|
"""
|
|
if self.alerting is None:
|
|
return
|
|
|
|
from datetime import datetime
|
|
|
|
# Get the current timestamp
|
|
current_time: Final = datetime.now().strftime("%H:%M:%S")
|
|
_proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None)
|
|
formatted_message = f"Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
|
|
if _proxy_base_url is not None:
|
|
formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`"
|
|
|
|
extra_kwargs: Final = {}
|
|
alerting_metadata = {}
|
|
if request_data is not None:
|
|
_url: Final = await _add_langfuse_trace_id_to_alert(request_data=request_data)
|
|
|
|
if _url is not None:
|
|
extra_kwargs["🪢 Langfuse Trace"] = _url
|
|
formatted_message += f"\n\n🪢 Langfuse Trace: {_url}"
|
|
if (
|
|
"metadata" in request_data
|
|
and request_data["metadata"].get("alerting_metadata", None) is not None
|
|
and isinstance(request_data["metadata"]["alerting_metadata"], dict)
|
|
):
|
|
alerting_metadata = request_data["metadata"]["alerting_metadata"]
|
|
if "slack" in self.alerting or "ms_teams" in self.alerting:
|
|
await self.slack_alerting_instance.send_alert(
|
|
message=message,
|
|
level=level,
|
|
alert_type=alert_type,
|
|
user_info=None,
|
|
alerting_metadata=alerting_metadata,
|
|
**extra_kwargs,
|
|
)
|
|
for client in self.alerting:
|
|
if client == "sentry":
|
|
if litellm.utils.sentry_sdk_instance is not None:
|
|
litellm.utils.sentry_sdk_instance.capture_message(formatted_message)
|
|
else:
|
|
raise Exception("Missing SENTRY_DSN from environment")
|
|
|
|
async def failure_handler(self, original_exception, duration: float, call_type: str, traceback_str=""):
|
|
"""
|
|
Log failed db read/writes
|
|
|
|
Currently only logs exceptions to sentry
|
|
"""
|
|
### ALERTING ###
|
|
if AlertType.db_exceptions not in self.alert_types:
|
|
return
|
|
if isinstance(original_exception, HTTPException):
|
|
if isinstance(original_exception.detail, str):
|
|
error_message = original_exception.detail
|
|
elif isinstance(original_exception.detail, dict):
|
|
error_message = json.dumps(original_exception.detail)
|
|
else:
|
|
error_message = str(original_exception)
|
|
else:
|
|
error_message = str(original_exception)
|
|
if isinstance(traceback_str, str):
|
|
error_message += traceback_str[:1000]
|
|
error_message = _redact_string(error_message)
|
|
asyncio.create_task(
|
|
self.alerting_handler(
|
|
message=f"DB read/write call failed: {error_message}",
|
|
level="High",
|
|
alert_type=AlertType.db_exceptions,
|
|
request_data={},
|
|
)
|
|
)
|
|
|
|
if hasattr(self, "service_logging_obj"):
|
|
await self.service_logging_obj.async_service_failure_hook(
|
|
service=ServiceTypes.DB,
|
|
duration=duration,
|
|
error=error_message,
|
|
call_type=call_type,
|
|
)
|
|
|
|
if litellm.utils.capture_exception:
|
|
litellm.utils.capture_exception(error=original_exception)
|
|
|
|
async def post_call_failure_hook(
|
|
self,
|
|
request_data: dict,
|
|
original_exception: Exception,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
error_type: ProxyErrorTypes | None = None,
|
|
route: str | None = None,
|
|
traceback_str: str | None = None,
|
|
) -> HTTPException | None:
|
|
"""
|
|
Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body.
|
|
Callbacks can return or raise HTTPException to transform error responses sent to clients.
|
|
|
|
Covers:
|
|
1. /chat/completions
|
|
2. /embeddings
|
|
3. /image/generation
|
|
|
|
Args:
|
|
- request_data: dict - The request data.
|
|
- original_exception: Exception - The original exception.
|
|
- user_api_key_dict: UserAPIKeyAuth - The user api key dict.
|
|
- error_type: Optional[ProxyErrorTypes] - The error type.
|
|
- route: Optional[str] - The route.
|
|
- traceback_str: Optional[str] - The traceback string, sometimes upstream endpoints might need to send the upstream traceback. In which case we use this
|
|
|
|
Returns:
|
|
- Optional[HTTPException]: If any callback returns or raises an HTTPException, the first one found is returned.
|
|
Otherwise, returns None and the original exception is used.
|
|
"""
|
|
|
|
### ALERTING ###
|
|
await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail")
|
|
if AlertType.llm_exceptions in self.alert_types and not _is_client_error_exception(original_exception):
|
|
"""
|
|
Just alert on LLM API exceptions. Do not alert on user errors
|
|
|
|
Related issue - https://github.com/BerriAI/litellm/issues/3395
|
|
"""
|
|
litellm_debug_info: Final[str | None] = getattr(original_exception, "litellm_debug_info", None)
|
|
exception_str = str(original_exception)
|
|
if litellm_debug_info is not None:
|
|
exception_str += litellm_debug_info
|
|
|
|
asyncio.create_task(
|
|
self.alerting_handler(
|
|
message=_redact_string(f"LLM API call failed: `{exception_str}`"),
|
|
level="High",
|
|
alert_type=AlertType.llm_exceptions,
|
|
request_data=request_data,
|
|
)
|
|
)
|
|
|
|
# Auth and pass-through failure bodies are unstripped client input, and
|
|
# the logging handler below flattens body keys into model_call_details,
|
|
# so drop the key before it can masquerade as the built payload.
|
|
request_data.pop("standard_logging_object", None)
|
|
|
|
### LOGGING ###
|
|
if self._is_proxy_only_llm_api_error(
|
|
original_exception=original_exception,
|
|
error_type=error_type,
|
|
route=user_api_key_dict.request_route,
|
|
):
|
|
await self._handle_logging_proxy_only_error(
|
|
request_data=request_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
route=route,
|
|
original_exception=original_exception,
|
|
)
|
|
|
|
request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data))
|
|
|
|
# Remove before callbacks iterate — not serialisable
|
|
request_data.pop("litellm_logging_obj", None)
|
|
|
|
redacted_traceback_str: Final = _redact_string(traceback_str) if traceback_str is not None else None
|
|
|
|
# Track the first HTTPException returned or raised by any callback
|
|
transformed_exception: HTTPException | None = None
|
|
|
|
for callback in litellm.callbacks:
|
|
try:
|
|
_callback: CustomLogger | None = None
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback)
|
|
)
|
|
else:
|
|
_callback = callback
|
|
if _callback is not None and isinstance(_callback, CustomLogger):
|
|
try:
|
|
hook_result = await _callback.async_post_call_failure_hook(
|
|
request_data=request_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
original_exception=original_exception,
|
|
traceback_str=redacted_traceback_str,
|
|
)
|
|
# If callback returned an HTTPException, use it (first one wins)
|
|
if isinstance(hook_result, HTTPException) and transformed_exception is None:
|
|
transformed_exception = hook_result
|
|
except HTTPException as e:
|
|
# If callback raised an HTTPException, use it (first one wins)
|
|
if transformed_exception is None:
|
|
transformed_exception = e
|
|
except Exception as e:
|
|
# Log non-HTTPException errors from callbacks but don't break the flow
|
|
verbose_proxy_logger.exception(
|
|
"[Non-Blocking] Error in async_post_call_failure_hook callback: %s", e
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception("[Non-Blocking] Error setting up post_call_failure_hook callback: %s", e)
|
|
|
|
return transformed_exception
|
|
|
|
def _is_proxy_only_llm_api_error(
|
|
self,
|
|
original_exception: Exception,
|
|
error_type: ProxyErrorTypes | None = None,
|
|
route: str | None = None,
|
|
) -> bool:
|
|
"""
|
|
Return True if the error is a Proxy Only LLM API Error
|
|
|
|
Prevents double logging of LLM API exceptions
|
|
|
|
e.g should only return True for:
|
|
- Authentication Errors from user_api_key_auth
|
|
- HTTP HTTPException (rate limit errors)
|
|
- ProxyException (guardrail blocks, budget / rate-limit errors)
|
|
- GuardrailRaisedException (guardrail blocks / guardrail failures)
|
|
"""
|
|
|
|
#########################################################
|
|
# Only log LLM API and info route errors for proxy level hooks
|
|
# eg. Authentication errors, rate limit errors, etc.
|
|
# Note: This fixes a security issue where we
|
|
# would log temporary keys/auth info
|
|
# from management endpoints
|
|
#########################################################
|
|
if route is None:
|
|
return False
|
|
if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)):
|
|
return False
|
|
|
|
return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error)
|
|
|
|
async def _handle_logging_proxy_only_error(
|
|
self,
|
|
request_data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
route: str | None = None,
|
|
original_exception: Exception | None = None,
|
|
):
|
|
"""
|
|
Handle logging for proxy only errors by calling `litellm_logging_obj.async_failure_handler`
|
|
|
|
Is triggered when self._is_proxy_only_error() returns True
|
|
"""
|
|
litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None)
|
|
if litellm_logging_obj is None:
|
|
from litellm._uuid import uuid
|
|
|
|
request_data.setdefault("litellm_call_id", str(uuid.uuid4()))
|
|
user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
|
|
user_api_key_dict=user_api_key_dict
|
|
)
|
|
|
|
litellm_logging_obj, data = litellm.utils.function_setup(
|
|
original_function=route or "IGNORE_THIS",
|
|
rules_obj=litellm.utils.Rules(),
|
|
start_time=datetime.now(),
|
|
**request_data,
|
|
)
|
|
request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller
|
|
if "metadata" not in request_data:
|
|
request_data["metadata"] = {}
|
|
request_data["metadata"].update(user_api_key_logged_metadata)
|
|
|
|
if litellm_logging_obj is not None:
|
|
## UPDATE LOGGING INPUT
|
|
_optional_params: Final = {}
|
|
_litellm_params: Final = {}
|
|
|
|
litellm_param_keys: Final = LoggedLiteLLMParams.__annotations__.keys()
|
|
for k, v in request_data.items():
|
|
if k in litellm_param_keys:
|
|
_litellm_params[k] = v
|
|
elif k not in ("model", "user", "litellm_logging_obj"):
|
|
_optional_params[k] = v
|
|
|
|
litellm_logging_obj.update_environment_variables(
|
|
model=request_data.get("model", ""),
|
|
user=request_data.get("user", ""),
|
|
optional_params=_optional_params,
|
|
litellm_params=_litellm_params,
|
|
)
|
|
|
|
input: list | str | dict = ""
|
|
body_shape_call_type: str | None = None
|
|
if "messages" in request_data and isinstance(request_data["messages"], list):
|
|
input = request_data["messages"]
|
|
litellm_logging_obj.model_call_details["messages"] = input
|
|
body_shape_call_type = CallTypes.acompletion.value
|
|
elif "prompt" in request_data and isinstance(request_data["prompt"], str):
|
|
input = request_data["prompt"]
|
|
litellm_logging_obj.model_call_details["prompt"] = input
|
|
body_shape_call_type = CallTypes.atext_completion.value
|
|
elif "input" in request_data and isinstance(request_data["input"], list):
|
|
input = request_data["input"]
|
|
litellm_logging_obj.model_call_details["input"] = input
|
|
body_shape_call_type = CallTypes.aembedding.value
|
|
resolved_call_type: Final = _call_type_for_route(route) or body_shape_call_type
|
|
if resolved_call_type is not None and litellm_logging_obj.call_type != CallTypes.pass_through.value:
|
|
litellm_logging_obj.call_type = resolved_call_type
|
|
litellm_logging_obj.model_call_details["call_type"] = resolved_call_type
|
|
# Pass-through endpoints are logged via the callback loop's
|
|
# async_post_call_failure_hook — skip pre_call and failure handlers.
|
|
if litellm_logging_obj.call_type == CallTypes.pass_through.value:
|
|
return
|
|
# This is a proxy-gate error (auth/rate-limit) for a request that never
|
|
# reached a provider. ``pre_call`` below still fires every callback's
|
|
# input hook so the failure is logged — but tracing callbacks must not
|
|
# fabricate an LLM-call span for a call that did not happen (and, since
|
|
# this runs inside the live ``auth`` phase span, would otherwise nest it
|
|
# under auth). The marker tells them to skip span creation.
|
|
litellm_logging_obj.model_call_details[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True
|
|
litellm_logging_obj.pre_call(
|
|
input=input,
|
|
api_key="",
|
|
)
|
|
|
|
await self._dispatch_proxy_only_failure_handlers(
|
|
litellm_logging_obj=litellm_logging_obj,
|
|
original_exception=original_exception,
|
|
)
|
|
|
|
@staticmethod
|
|
async def _dispatch_proxy_only_failure_handlers(
|
|
litellm_logging_obj: Logging,
|
|
original_exception: Exception | None,
|
|
) -> None:
|
|
"""Runs the async failure handler plus the threaded sync handler. Expected
|
|
client (4xx) errors skip traceback formatting unless
|
|
litellm.log_client_error_tracebacks is set."""
|
|
include_traceback: Final = litellm.log_client_error_tracebacks or not is_expected_client_error(
|
|
original_exception
|
|
)
|
|
traceback_str: Final = traceback.format_exc() if include_traceback else ""
|
|
await litellm_logging_obj.async_failure_handler(
|
|
exception=original_exception,
|
|
traceback_exception=traceback_str,
|
|
)
|
|
|
|
threading.Thread(
|
|
target=litellm_logging_obj.failure_handler,
|
|
args=(
|
|
original_exception,
|
|
traceback_str,
|
|
),
|
|
daemon=True,
|
|
).start()
|
|
|
|
async def _run_post_call_pipelines(
|
|
self,
|
|
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: LLMResponseTypes,
|
|
) -> LLMResponseTypes | None:
|
|
if _is_pending_background_response(response):
|
|
_defer_post_call_pipelines(data, response)
|
|
return None
|
|
_, pipeline_response = await self._maybe_execute_pipelines(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
|
|
event_hook="post_call",
|
|
response=response,
|
|
)
|
|
return pipeline_response
|
|
|
|
async def post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
response: LLMResponseTypes,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
):
|
|
"""
|
|
Allow user to modify outgoing data
|
|
|
|
Covers:
|
|
1. /chat/completions
|
|
2. /embeddings
|
|
3. /image/generation
|
|
4. /files
|
|
"""
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
pipeline_response: Final = await self._run_post_call_pipelines(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=response,
|
|
)
|
|
if pipeline_response is not None:
|
|
response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below
|
|
|
|
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
|
|
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
|
|
try:
|
|
# Merge model-level guardrails before checking which guardrails to run
|
|
guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
|
|
|
|
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple(
|
|
callback
|
|
for callback in guardrail_callbacks
|
|
if getattr(callback, "run_in_parallel", False)
|
|
and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed)
|
|
)
|
|
|
|
for callback in guardrail_callbacks:
|
|
# Main - V2 Guardrails implementation
|
|
|
|
if callback.guardrail_name and callback.guardrail_name in pipeline_managed:
|
|
continue
|
|
|
|
if getattr(callback, "run_in_parallel", False):
|
|
continue
|
|
|
|
if (
|
|
callback.should_run_guardrail(
|
|
data=guardrail_data,
|
|
event_type=GuardrailEventHooks.post_call,
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
|
|
guardrail_response: Any | None = None
|
|
|
|
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
|
|
data["guardrail_to_apply"] = callback
|
|
guardrail_response = await self._run_guardrail_with_metrics(
|
|
callback,
|
|
unified_guardrail.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
),
|
|
"post_call",
|
|
request_data=data,
|
|
)
|
|
else:
|
|
guardrail_response = await self._run_guardrail_with_metrics(
|
|
callback,
|
|
callback.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
),
|
|
"post_call",
|
|
request_data=data,
|
|
)
|
|
|
|
if guardrail_response is not None:
|
|
response = guardrail_response
|
|
|
|
if parallel_guardrails:
|
|
await self._run_parallel_post_call_guardrails(
|
|
guardrails=parallel_guardrails,
|
|
data=data,
|
|
guardrail_data=guardrail_data,
|
|
response=response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
############ Handle CustomLogger ###############################
|
|
#################################################################
|
|
|
|
for callback in other_callbacks:
|
|
callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict, data=data, response=response
|
|
)
|
|
if callback_response is not None:
|
|
response = callback_response
|
|
except Exception as e:
|
|
raise e
|
|
return response
|
|
|
|
async def _run_parallel_post_call_guardrails(
|
|
self,
|
|
guardrails: tuple[CustomGuardrail, ...],
|
|
data: dict,
|
|
guardrail_data: dict,
|
|
response: LLMResponseTypes,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
) -> None:
|
|
"""
|
|
Run opted-in post_call guardrails concurrently against the response
|
|
produced by the sequential guardrails. These guardrails are declared
|
|
block-only, so any modified response they return is discarded; they run
|
|
for their blocking side effect (raising to reject the response before it
|
|
reaches the client). Every guardrail is awaited to completion
|
|
(``return_exceptions=True``) so a raise by one never leaves the others
|
|
running as unobserved background tasks. A guardrail that blocks (any
|
|
exception other than a passthrough) takes precedence over one that only
|
|
changes the response flow, so a fast passthrough can never let a slower
|
|
block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply``
|
|
immediately before awaiting, and the unified hook pops it before its first
|
|
suspension point, so concurrent guardrails never race on that key.
|
|
"""
|
|
|
|
async def _run_one(callback: CustomGuardrail) -> None:
|
|
if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True:
|
|
return
|
|
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
|
|
data["guardrail_to_apply"] = callback
|
|
await self._run_guardrail_with_metrics(
|
|
callback,
|
|
unified_guardrail.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
),
|
|
"post_call",
|
|
request_data=data,
|
|
)
|
|
else:
|
|
await self._run_guardrail_with_metrics(
|
|
callback,
|
|
callback.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
),
|
|
"post_call",
|
|
request_data=data,
|
|
)
|
|
|
|
results: Final = await asyncio.gather(
|
|
*(_run_one(callback) for callback in guardrails),
|
|
return_exceptions=True,
|
|
)
|
|
raised: Final = tuple(result for result in results if isinstance(result, BaseException))
|
|
blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
|
|
if blocking is not None:
|
|
raise blocking
|
|
if raised:
|
|
raise raised[0]
|
|
|
|
async def post_mcp_call_hook(
|
|
self,
|
|
response: "CallToolResult",
|
|
request_data: Mapping[str, Any],
|
|
user_api_key_dict: UserAPIKeyAuth | None = None,
|
|
) -> "CallToolResult":
|
|
"""
|
|
Run guardrails configured for ``post_mcp_call`` against an MCP tool result.
|
|
|
|
The MCP counterpart of ``post_call_success_hook``: guardrails that
|
|
implement ``apply_guardrail`` see the tool result's text through the
|
|
unified guardrail seam (``MCPGuardrailTranslationHandler``), so a text
|
|
guardrail can mask sensitive values in the result without any MCP-specific
|
|
code of its own. Guardrails that instead implement
|
|
``async_post_mcp_tool_call_hook`` are dispatched by
|
|
``Logging.async_post_mcp_tool_call_hook`` and are not run here.
|
|
|
|
A guardrail that rejects the result raises, and the exception propagates
|
|
(matching the inbound ``pre_mcp_call`` behavior) rather than being
|
|
swallowed into an unguarded result.
|
|
"""
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
if not caps.has_guardrail:
|
|
return response
|
|
|
|
handler_cls: Final = load_guardrail_translation_mappings().get(CallTypes.call_mcp_tool)
|
|
if handler_cls is None:
|
|
verbose_proxy_logger.debug("MCP guardrail translation handler unavailable; skipping post_mcp_call hook")
|
|
return response
|
|
|
|
for callback in caps.resolved_callbacks:
|
|
if not isinstance(callback, CustomGuardrail):
|
|
continue
|
|
if "apply_guardrail" not in type(callback).__dict__ or callback.use_native_lifecycle_hooks:
|
|
continue
|
|
if (
|
|
callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call)
|
|
is not True
|
|
):
|
|
continue
|
|
response = await self._run_guardrail_with_metrics(
|
|
callback,
|
|
handler_cls().process_output_response(
|
|
response=response,
|
|
guardrail_to_apply=callback,
|
|
litellm_logging_obj=request_data.get("litellm_logging_obj"),
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data=request_data,
|
|
),
|
|
"post_mcp_call",
|
|
request_data=request_data,
|
|
)
|
|
return response
|
|
|
|
async def post_call_response_headers_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: object,
|
|
request_headers: dict[str, str] | None = None,
|
|
) -> dict[str, str]:
|
|
"""
|
|
Calls async_post_call_response_headers_hook on all CustomLogger callbacks.
|
|
Merges all returned header dicts (later callbacks override earlier ones).
|
|
|
|
Returns:
|
|
Dict[str, str]: Merged headers from all callbacks.
|
|
"""
|
|
merged_headers: Final[dict[str, str]] = {}
|
|
# Outer call sites in common_request_processing.py already gate this
|
|
# call with ``has_post_call_response_headers_callbacks()``. The
|
|
# cached detection makes the redundant interior guard cheap, but the
|
|
# guard would still iterate every code path through this function so
|
|
# keep it cheap and rely on the cached capability lookup.
|
|
if not ProxyLogging._callback_capabilities().has_post_call_response_headers:
|
|
return merged_headers
|
|
|
|
try:
|
|
# Build litellm_call_info — normalized routing metadata for callbacks
|
|
litellm_call_info: Final = self._build_litellm_call_info(data=data, response=response)
|
|
|
|
for callback in litellm.callbacks:
|
|
_callback: CustomLogger | None = None
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback)
|
|
)
|
|
else:
|
|
_callback = callback
|
|
|
|
if _callback is not None and isinstance(_callback, CustomLogger):
|
|
if _accepts_litellm_call_info(_callback):
|
|
result = await _callback.async_post_call_response_headers_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=response,
|
|
request_headers=request_headers,
|
|
litellm_call_info=litellm_call_info,
|
|
)
|
|
else:
|
|
# Backwards compat: callback doesn't accept litellm_call_info
|
|
result = await _callback.async_post_call_response_headers_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=response,
|
|
request_headers=request_headers,
|
|
)
|
|
if result is not None:
|
|
merged_headers.update(result)
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception("Error in post_call_response_headers_hook: %s", str(e))
|
|
return merged_headers
|
|
|
|
@staticmethod
|
|
def _build_litellm_call_info(data: dict, response: object) -> dict[str, object]:
|
|
"""
|
|
Build a normalized dict of routing metadata from response._hidden_params
|
|
and data, abstracting away the metadata vs litellm_metadata split.
|
|
"""
|
|
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
|
|
|
|
# model_info: check both metadata keys (chat uses "metadata", responses uses "litellm_metadata")
|
|
model_info: Final = (
|
|
(data.get("metadata") or {}).get("model_info")
|
|
or (data.get("litellm_metadata") or {}).get("model_info")
|
|
or {}
|
|
)
|
|
|
|
return {
|
|
"custom_llm_provider": hidden_params.get("custom_llm_provider")
|
|
or getattr(response, "custom_llm_provider", None),
|
|
"model_info": model_info,
|
|
"api_base": hidden_params.get("api_base"),
|
|
"model_id": hidden_params.get("model_id"),
|
|
}
|
|
|
|
def is_a2a_streaming_response(self, response: dict) -> bool:
|
|
expected_keys: Final = ["jsonrpc", "id", "result"]
|
|
return all(key in response for key in expected_keys)
|
|
|
|
async def async_post_call_streaming_hook(
|
|
self,
|
|
data: dict,
|
|
response: ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
str_so_far: str | None = None,
|
|
):
|
|
"""
|
|
Allow user to modify outgoing streaming data -> per chunk
|
|
|
|
Covers:
|
|
1. /chat/completions
|
|
"""
|
|
# Per-chunk fast path: skip the response-string materialization and
|
|
# callback scan when no configured callback overrides
|
|
# ``async_post_call_streaming_hook`` AND no CustomGuardrail is
|
|
# active. ``get_response_string`` walks every choice/delta on the
|
|
# chunk so paying it per chunk for no-op callbacks dominated stream
|
|
# CPU time even after the iterator-chain fix.
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
if not caps.has_streaming_chunk_override and not caps.has_guardrail:
|
|
return response
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
response_str: str | None = None
|
|
if isinstance(response, (ModelResponse, ModelResponseStream)):
|
|
response_str = litellm.get_response_string(response_obj=response)
|
|
elif isinstance(response, dict) and self.is_a2a_streaming_response(response):
|
|
from litellm.llms.a2a.common_utils import extract_text_from_a2a_response
|
|
|
|
response_str = extract_text_from_a2a_response(response)
|
|
if response_str is not None:
|
|
# Cache model-level guardrails check per-request to avoid repeated
|
|
# dict lookups + llm_router.get_deployment() per callback per chunk.
|
|
_cached_guardrail_data: dict | None = None
|
|
_guardrail_data_computed = False
|
|
pipeline_gated: Final = (
|
|
stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
|
|
)
|
|
|
|
for callback in litellm.callbacks:
|
|
try:
|
|
_callback: CustomLogger | None = None
|
|
if isinstance(callback, CustomGuardrail):
|
|
if callback.guardrail_name in pipeline_gated:
|
|
continue
|
|
# Main - V2 Guardrails implementation
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
## CHECK FOR MODEL-LEVEL GUARDRAILS (cached per-request)
|
|
if not _guardrail_data_computed:
|
|
_cached_guardrail_data = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=llm_router
|
|
)
|
|
_guardrail_data_computed = True
|
|
|
|
if (
|
|
callback.should_run_guardrail(
|
|
data=_cached_guardrail_data,
|
|
event_type=GuardrailEventHooks.post_call,
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback)
|
|
)
|
|
else:
|
|
_callback = callback
|
|
if _callback is not None and isinstance(_callback, CustomLogger):
|
|
if str_so_far is not None:
|
|
complete_response = str_so_far + response_str
|
|
else:
|
|
complete_response = response_str
|
|
callback_response: (
|
|
ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None
|
|
)
|
|
callback_response = await _callback.async_post_call_streaming_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=complete_response,
|
|
)
|
|
if callback_response is not None:
|
|
response = callback_response
|
|
except Exception as e:
|
|
raise e
|
|
return response
|
|
|
|
async def async_post_call_streaming_iterator_hook(
|
|
self,
|
|
response,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
request_data: dict,
|
|
):
|
|
"""
|
|
Allow user to modify outgoing streaming data -> Given a whole response iterator.
|
|
This hook is best used when you need to modify multiple chunks of the response at once.
|
|
|
|
Covers:
|
|
1. /chat/completions
|
|
"""
|
|
caps: Final = ProxyLogging._callback_capabilities()
|
|
post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
|
|
# Fast path: no real overrides. Internal proxy CustomLogger callbacks
|
|
# (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default
|
|
# ``async for chunk: yield chunk`` body, so wrapping the iterator
|
|
# through each of them adds N pass-through trampolines per chunk for
|
|
# zero behavior change. Skip the chain entirely and stream through.
|
|
if not caps.iterator_overrides and not post_call_pipelines:
|
|
try:
|
|
async for chunk in response:
|
|
yield chunk
|
|
except (GeneratorExit, asyncio.CancelledError):
|
|
raise
|
|
except Exception as e:
|
|
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
|
|
ProxyLogging._fire_deferred_stream_logging(request_data)
|
|
raise
|
|
ProxyLogging._fire_deferred_stream_logging(request_data)
|
|
return
|
|
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
# Merge model-level guardrails before checking which guardrails to run
|
|
request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router)
|
|
|
|
current_response = response
|
|
stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
|
|
|
|
pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines)
|
|
for resolved_callback, kind in caps.iterator_overrides:
|
|
if isinstance(resolved_callback, CustomGuardrail):
|
|
if resolved_callback.guardrail_name in pipeline_gated_names:
|
|
continue
|
|
if (
|
|
resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call)
|
|
is not True
|
|
):
|
|
continue
|
|
effective_kind = (
|
|
"apply_guardrail"
|
|
if (
|
|
kind == "override"
|
|
and stream_needs_translation
|
|
and isinstance(resolved_callback, CustomGuardrail)
|
|
and resolved_callback.uses_apply_guardrail_interface()
|
|
and getattr(resolved_callback, "use_native_lifecycle_hooks", False) is not True
|
|
and not resolved_callback.mask_response_content
|
|
)
|
|
else kind
|
|
)
|
|
hook: _StreamIteratorHook[object] = (
|
|
partial(
|
|
resolved_callback.async_post_call_streaming_iterator_hook,
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data=request_data,
|
|
)
|
|
if effective_kind == "override"
|
|
else partial(
|
|
unified_guardrail.async_post_call_streaming_iterator_hook,
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data=request_data,
|
|
guardrail_to_apply=resolved_callback,
|
|
buffer_until_moderated_default=(kind == "override"),
|
|
)
|
|
)
|
|
current_response = self._wrap_streaming_iterator_with_enrichment(
|
|
resolved_callback,
|
|
current_response,
|
|
hook,
|
|
request_data=request_data,
|
|
)
|
|
|
|
pipeline_translation: Final = (
|
|
resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None
|
|
)
|
|
if pipeline_translation is not None:
|
|
current_response = self._pipeline_gated_stream(
|
|
response=current_response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
request_data=request_data,
|
|
pipelines=post_call_pipelines,
|
|
translation=pipeline_translation,
|
|
)
|
|
|
|
try:
|
|
async for chunk in current_response:
|
|
yield chunk
|
|
except (GeneratorExit, asyncio.CancelledError):
|
|
raise
|
|
except Exception as e:
|
|
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
|
|
ProxyLogging._fire_deferred_stream_logging(request_data)
|
|
raise
|
|
|
|
# Fire deferred logging AFTER all guardrail end-of-stream blocks
|
|
# completed. unified_guardrail writes guardrail_information during
|
|
# its end-of-stream block (inside current_response), so by the time
|
|
# we reach this point the metadata is fully populated.
|
|
ProxyLogging._fire_deferred_stream_logging(request_data)
|
|
|
|
async def _pipeline_gated_stream(
|
|
self,
|
|
response: "AsyncGenerator[object, None]",
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
request_data: dict, # mutable-ok: same request-payload shape the hooks mutate
|
|
pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
|
|
translation: "tuple[str, BaseTranslation]",
|
|
) -> "AsyncGenerator[Any, None]":
|
|
"""
|
|
Execute post_call policy pipelines against a streamed response.
|
|
|
|
Buffers the whole stream (nothing reaches the client until every
|
|
pipeline allows it), then runs each pipeline's steps against the
|
|
assembled output through the endpoint guardrail translation, the same
|
|
machinery flat post_call guardrails use at end of stream. An allow
|
|
releases the buffered chunks: verbatim when no guardrail rewrote the
|
|
output, rewritten in place when one rewrote text or a tool call and the
|
|
translation delivers ended-stream rewrites (later steps then re-scan the
|
|
rewritten chunks, so rewrites chain). A rewrite the translation cannot
|
|
deliver yet (one on a route without write-back, or a shape the route
|
|
refuses) is discarded by the executor and the original chunks are
|
|
released; a block or modify_response terminates with the translation's
|
|
block chunks or the raised error.
|
|
"""
|
|
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
|
|
async for item in response:
|
|
buffered.append(item)
|
|
if not buffered:
|
|
return
|
|
|
|
call_type, endpoint_translation = translation
|
|
|
|
for policy_name, pipeline in pipelines:
|
|
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
|
|
steps=pipeline.steps,
|
|
mode="post_call",
|
|
data=request_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
policy_name=policy_name,
|
|
streaming_chunks=buffered,
|
|
endpoint_translation=endpoint_translation,
|
|
)
|
|
try:
|
|
ProxyLogging._handle_pipeline_result(
|
|
result, data=request_data, policy_name=policy_name, original_response=buffered
|
|
)
|
|
except ModifyResponseException as e:
|
|
if e.original_response is None:
|
|
e.original_response = buffered
|
|
async for block_chunk in unified_guardrail.handle_streaming_block(
|
|
e, endpoint_translation, stream_started=False, responses_so_far=()
|
|
):
|
|
yield block_chunk
|
|
return
|
|
except HTTPException as e:
|
|
async for error_chunk in unified_guardrail.emit_streaming_http_error(
|
|
e, call_type, buffered, request_data
|
|
):
|
|
yield error_chunk
|
|
return
|
|
|
|
for buffered_item in buffered:
|
|
yield buffered_item
|
|
|
|
@staticmethod
|
|
def _fire_deferred_stream_logging(request_data: dict) -> None:
|
|
"""
|
|
Fire the deferred streaming logging callback after the full streaming
|
|
pipeline (including guardrail end-of-stream blocks) has completed.
|
|
|
|
CSW.__anext__ stores the callback and args on logging_obj instead of
|
|
scheduling via create_task (which would race with unified_guardrail's
|
|
end-of-stream block). This method retrieves and fires them.
|
|
"""
|
|
logging_obj: Final = request_data.get("litellm_logging_obj")
|
|
if logging_obj is None:
|
|
return
|
|
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
|
|
logging_obj, "_on_deferred_stream_complete", None
|
|
)
|
|
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
|
|
if _deferred_cb is not None and _args is not None:
|
|
logging_obj._on_deferred_stream_complete = None
|
|
logging_obj._deferred_stream_complete_args = None
|
|
asyncio.create_task(_deferred_cb(*_args))
|
|
|
|
@staticmethod
|
|
def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool:
|
|
"""Drop the parked success dispatch for an assembled chat stream that ends in an error
|
|
``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead.
|
|
Returns False when the parked dispatch should still be flushed by the caller."""
|
|
logging_obj: Final = request_data.get("litellm_logging_obj")
|
|
if not isinstance(logging_obj, Logging):
|
|
return False
|
|
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
|
|
assembled: Final = _args[0] if _args else None
|
|
if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse):
|
|
return False
|
|
logging_obj._on_deferred_stream_complete = None
|
|
logging_obj._deferred_stream_complete_args = None
|
|
logging_obj.record_assembled_response_for_failure(assembled)
|
|
return True
|
|
|
|
async def _arelease_max_parallel_requests_on_disconnect(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
) -> None:
|
|
"""
|
|
Release the api-key max_parallel_requests slot when a streaming
|
|
response is cancelled mid-flight (client disconnect) and no logging
|
|
callback fired for it. Neither the success nor failure callback runs on
|
|
the resulting CancelledError / GeneratorExit, so the pre-call +1 would
|
|
otherwise leak.
|
|
|
|
Awaited from the shielded streaming cleanup rather than scheduled
|
|
fire-and-forget, so the caller can make it the single owner of the
|
|
release: when a disconnect-time success event does fire (partial-spend
|
|
billing or a deferred-guardrail flush), that event's own limiter
|
|
callback releases the slot and this is not called at all. Two
|
|
concurrent releases of the same acquisition would otherwise race and
|
|
double-decrement under the limiter's in-memory fallback.
|
|
"""
|
|
limiter: Final = self.get_proxy_hook("parallel_request_limiter")
|
|
if not isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3):
|
|
return
|
|
await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
|
|
|
|
def _init_response_taking_too_long_task(self, data: dict | None = None):
|
|
"""
|
|
Initialize the response taking too long task if user is using slack alerting
|
|
|
|
Only run task if user is using slack alerting
|
|
|
|
This handles checking for if a request is hanging for too long
|
|
"""
|
|
## ALERTING ###
|
|
if self.slack_alerting_instance and self.slack_alerting_instance.alerting is not None:
|
|
asyncio.create_task(self.slack_alerting_instance.response_taking_too_long(request_data=data))
|
|
|
|
|
|
### DB CONNECTOR ###
|
|
# Define the retry decorator with backoff strategy
|
|
# Function to be called whenever a retry is about to happen
|
|
def on_backoff(details):
|
|
# The 'tries' key in the details dictionary contains the number of completed tries
|
|
print_verbose(f"Backing off... this was attempt #{details['tries']}")
|
|
|
|
|
|
def jsonify_object(data: dict) -> dict:
|
|
db_data: Final = copy.deepcopy(data)
|
|
|
|
for k, v in db_data.items():
|
|
if isinstance(v, dict):
|
|
try:
|
|
db_data[k] = json.dumps(v)
|
|
except Exception:
|
|
# This avoids Prisma retrying this 5 times, and making 5 clients
|
|
db_data[k] = "failed-to-serialize-json"
|
|
return db_data
|
|
|
|
|
|
# In-memory cache for deprecated key lookups:
|
|
# maps old_token_hash -> (active_token_id, cache_expires_at_ts, revoke_at_ts).
|
|
# Avoids a DB query on every auth request for non-deprecated keys.
|
|
# Bounded to prevent memory leaks from accumulated rotations.
|
|
_deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(max_size=1000)
|
|
_DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60
|
|
_PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5)
|
|
|
|
|
|
async def _lookup_deprecated_key(
|
|
db: PrismaWrapper | RoutingPrismaWrapper,
|
|
hashed_token: str,
|
|
) -> str | None:
|
|
"""
|
|
Check if a token exists in the deprecated keys table and is still within its grace period.
|
|
|
|
Returns the active_token_id if found and valid, otherwise None.
|
|
Uses an in-memory cache to avoid DB queries on every auth request.
|
|
"""
|
|
now: Final = datetime.now(timezone.utc)
|
|
now_ts: Final = now.timestamp()
|
|
|
|
# Check cache first
|
|
cached: Final = _deprecated_key_cache.get(hashed_token)
|
|
if cached is not None:
|
|
active_token_id, cache_expires_at_ts, revoke_at_ts = cached
|
|
if now_ts < cache_expires_at_ts and now_ts < revoke_at_ts:
|
|
return active_token_id
|
|
_deprecated_key_cache.pop(hashed_token, None)
|
|
|
|
try:
|
|
deprecated_keys_table: Final[
|
|
LiteLLM_DeprecatedVerificationTokenActions[LiteLLM_DeprecatedVerificationToken]
|
|
] = db.litellm_deprecatedverificationtoken
|
|
deprecated_row: Final = await deprecated_keys_table.find_first(
|
|
where={
|
|
"token": hashed_token,
|
|
"revoke_at": {"gt": now},
|
|
}
|
|
)
|
|
if deprecated_row and deprecated_row.active_token_id:
|
|
revoke_at: Final = deprecated_row.revoke_at
|
|
_deprecated_key_cache[hashed_token] = (
|
|
deprecated_row.active_token_id,
|
|
now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS,
|
|
revoke_at.timestamp(),
|
|
)
|
|
return deprecated_row.active_token_id
|
|
# Only cache positive results; negative lookups are fast on indexed columns
|
|
# and caching them risks evicting real deprecated key entries.
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug("Deprecated key lookup skipped: %s", e)
|
|
|
|
return None
|
|
|
|
|
|
# DualCache for LiteLLM_Config param_name reads.
|
|
# Redis layer is attached in proxy_server._init_cache.
|
|
LITELLM_CONFIG_CACHE_TTL_SECONDS: Final[int] = int(os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60"))
|
|
_CONFIG_CACHE_MISS: Final[str] = "__litellm_config_param_miss__"
|
|
|
|
litellm_config_cache: Final[DualCache] = DualCache(
|
|
default_in_memory_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS,
|
|
default_redis_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS,
|
|
)
|
|
|
|
|
|
class _ConfigRow:
|
|
"""Mimics the Prisma litellm_config row shape for cached entries."""
|
|
|
|
__slots__ = ("param_name", "param_value")
|
|
|
|
def __init__(self, param_name: str, param_value: Any) -> None:
|
|
self.param_name = param_name
|
|
self.param_value = param_value
|
|
|
|
|
|
def _config_cache_key(param_name: str) -> str:
|
|
return f"litellm_config:param:{param_name}"
|
|
|
|
|
|
def _pack_config_row(row: Any) -> dict[str, object]:
|
|
return {"param_name": row.param_name, "param_value": row.param_value}
|
|
|
|
|
|
def _unpack_config_row(cached: Any) -> _ConfigRow | None:
|
|
if cached is None or cached == _CONFIG_CACHE_MISS:
|
|
return None
|
|
if isinstance(cached, dict):
|
|
return _ConfigRow(cached["param_name"], cached["param_value"])
|
|
return None
|
|
|
|
|
|
async def get_config_param(prisma_client: "PrismaClient", param_name: str) -> Any | None:
|
|
"""Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None."""
|
|
cache_key: Final = _config_cache_key(param_name)
|
|
cached: Final = await litellm_config_cache.async_get_cache(cache_key)
|
|
if cached is not None:
|
|
return _unpack_config_row(cached)
|
|
|
|
row: Final = await prisma_client.get_generic_data(key="param_name", value=param_name, table_name="config")
|
|
cache_value: Final[Mapping[str, object] | str] = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS
|
|
await litellm_config_cache.async_set_cache(cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS)
|
|
return row
|
|
|
|
|
|
async def evict_config_param(param_name: str) -> None:
|
|
await litellm_config_cache.async_delete_cache(_config_cache_key(param_name))
|
|
|
|
|
|
async def invalidate_config_param(param_name: str) -> None:
|
|
"""Evict from both cache layers; call after every LiteLLM_Config write."""
|
|
await evict_config_param(param_name)
|
|
await publish_config_param_change(param_name)
|
|
|
|
|
|
async def prefetch_config_params(prisma_client: "PrismaClient | None", param_names: list[str]) -> None:
|
|
"""Batch-load LiteLLM_Config rows into the cache with one find_many."""
|
|
if not param_names:
|
|
return
|
|
try:
|
|
config_table: Final = cast( # cast-ok: ConfigRepository.table is prisma's litellm_config actions object
|
|
"TableActions[prisma_models.LiteLLM_Config]", ConfigRepository(prisma_client).table
|
|
)
|
|
rows: Final = await config_table.find_many(where={"param_name": {"in": param_names}})
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
"prefetch_config_params failed, falling through to per-param queries: %s",
|
|
e,
|
|
)
|
|
return
|
|
by_name: Final = {row.param_name: row for row in rows}
|
|
for name in param_names:
|
|
row = by_name.get(name)
|
|
cache_value: Mapping[str, object] | str = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS
|
|
await litellm_config_cache.async_set_cache(
|
|
_config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS
|
|
)
|
|
|
|
|
|
_WRITER_WRITABILITY_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only"
|
|
_WRITER_WRITABILITY_PROBE_ROWS: Final = TypeAdapter(list[dict[str, object]])
|
|
_READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS: Final = 600
|
|
|
|
|
|
class _ForcedRecreateDeclined(Exception):
|
|
"""A forced recreate was declined by the engine-generation guard.
|
|
|
|
Distinct from a reconnect *failure*: the machinery worked, it just found
|
|
that another path had already replaced the writer, so it left the engines
|
|
alone. The caller's engine may still be poisoned, so the cycle must not
|
|
report success, but it must not count as a failure either, or the record
|
|
of what could not be repaired would gate the retry that recovers.
|
|
"""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _StaleReadEngine:
|
|
"""The read engine a query observed, identified rather than only counted.
|
|
|
|
`PrismaClient.read_db` resolves to the reader while it is available and to
|
|
the writer once it is not, and the two carry independent generation
|
|
counters that both start at zero and advance on the same reconnect
|
|
cadence. A bare generation compared across that switch would silently pit
|
|
one engine's counter against another's, so the wrapper is carried with the
|
|
number and a switch counts as the engine having moved.
|
|
|
|
Holding the wrapper itself rather than its `id()` is load-bearing, not
|
|
incidental: the strong reference keeps the wrapper alive, so its address
|
|
cannot be recycled under a stored observation and match an unrelated
|
|
engine later. It is only free because writer and reader both live as long
|
|
as the client does; a replaceable reader would make this a retention leak.
|
|
"""
|
|
|
|
wrapper: PrismaWrapper
|
|
generation: int
|
|
|
|
@classmethod
|
|
def observe(cls, wrapper: PrismaWrapper) -> "_StaleReadEngine":
|
|
return cls(wrapper=wrapper, generation=wrapper.engine_generation)
|
|
|
|
def is_still_live(self, current: PrismaWrapper) -> bool:
|
|
"""Whether this exact engine is still serving reads, unreplaced.
|
|
|
|
A True answer must never be the only thing standing between a poisoned
|
|
engine and its repair. The generation moves only after a replacement
|
|
connects, and a recreate whose connect raises leaves it unmoved until
|
|
some later recreate succeeds, so this can report an engine as live
|
|
after it has stopped working. What bounds that is the failed-repair
|
|
record in `_cooldown_applies`, written by a repair attempt that fails
|
|
rather than by whatever broke the engine: the two need not be the same
|
|
recreate, since the synchronous token-refresh fallback in
|
|
`PrismaWrapper.__getattr__` recreates outside the reconnect machinery
|
|
and records nothing. The record is written only for callers that named
|
|
an engine, and it collapses the rest of the burst for up to one
|
|
cooldown window rather than guaranteeing a repair, since the cooldown
|
|
conjunct underneath it still expires and lets a later caller retry.
|
|
"""
|
|
return self.wrapper is current and self.generation == current.engine_generation
|
|
|
|
|
|
class PrismaClient:
|
|
spend_log_transactions: list = []
|
|
_spend_log_transactions_lock = asyncio.Lock()
|
|
spend_log_flush_requested: "asyncio.Event | None" = None
|
|
spend_log_queue_bytes: ClassVar[int] = 0
|
|
spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None
|
|
tool_usage_transactions: list["ToolUsageTransaction"] = []
|
|
_tool_usage_transactions_lock = asyncio.Lock()
|
|
autorouter_turn_transactions: ClassVar[
|
|
list["AutoRouterTurnTransaction"]
|
|
] = [] # mutable-ok: drained queue, mirrors tool_usage_transactions
|
|
_autorouter_turn_transactions_lock = asyncio.Lock()
|
|
|
|
# How long a health probe failure waits for an in-flight planned engine
|
|
# replacement to settle before deciding whether to report itself. Generous
|
|
# against a replacement that takes well under a second, and far short of the
|
|
# reconnect budget an outage-hung `connect()` runs under, so a real outage
|
|
# is never waited out.
|
|
PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS: ClassVar[float] = 5.0
|
|
|
|
def __init__(
|
|
self,
|
|
database_url: str,
|
|
proxy_logging_obj: ProxyLogging,
|
|
http_client: "HttpConfig | None" = None,
|
|
):
|
|
## init logging object
|
|
self.proxy_logging_obj = proxy_logging_obj
|
|
self.token_auth: DatabaseTokenAuth | None = resolve_database_token_auth()
|
|
verbose_proxy_logger.debug("Creating Prisma Client..")
|
|
try:
|
|
from prisma import Prisma
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Failed to import Prisma client: %s", e)
|
|
verbose_proxy_logger.error("This usually means 'prisma generate' hasn't been run yet.")
|
|
verbose_proxy_logger.error("Please run 'prisma generate' to generate the Prisma client.")
|
|
raise Exception("Unable to find Prisma binaries. Please run 'prisma generate' first.")
|
|
token_auth: Final = self.token_auth
|
|
writer_token_auth: Final = None if database_url_is_pooled() else token_auth
|
|
# When read-replica routing is on, tag log lines with [writer]/[reader]
|
|
# so the two wrappers' interleaved token refresh logs can be told apart.
|
|
# Single-DB deployments get an empty prefix (logs unchanged).
|
|
read_replica_url = os.getenv("DATABASE_URL_READ_REPLICA")
|
|
writer_log_prefix: Final = "[writer]" if read_replica_url else ""
|
|
if http_client is not None:
|
|
writer_wrapper = PrismaWrapper(
|
|
original_prisma=Prisma(http=http_client),
|
|
token_auth=writer_token_auth,
|
|
log_prefix=writer_log_prefix,
|
|
)
|
|
else:
|
|
writer_wrapper = PrismaWrapper(
|
|
original_prisma=Prisma(),
|
|
token_auth=writer_token_auth,
|
|
log_prefix=writer_log_prefix,
|
|
)
|
|
|
|
# Optional read-replica routing. When DATABASE_URL_READ_REPLICA is set,
|
|
# reads (find_*, count, group_by, query_raw/_first) are routed to the
|
|
# reader endpoint and writes stay on the writer. Falls back to the
|
|
# writer-only wrapper when the env var is unset, preserving existing
|
|
# single-DB deployments.
|
|
self.db: PrismaWrapper | RoutingPrismaWrapper
|
|
if read_replica_url:
|
|
try:
|
|
# If token auth is enabled, the reader refreshes its own token on
|
|
# the same cadence as the writer. We parse the static endpoint
|
|
# pieces (host/port/user/db) once from the reader URL — only
|
|
# the token rotates after that.
|
|
reader_iam_endpoint: Final = (
|
|
parse_iam_endpoint_from_url(read_replica_url) if token_auth is not None else None
|
|
)
|
|
# Mint a fresh token for the reader BEFORE constructing the
|
|
# Prisma client. Mirrors what `proxy_cli.py` already does for
|
|
# the writer — without this, the reader Prisma is built with
|
|
# whatever placeholder URL the user supplied (no real token),
|
|
# and the first query falls through to the synchronous fallback
|
|
# path in `PrismaWrapper.__getattr__`, which deadlocks the event
|
|
# loop and times out after 30s.
|
|
if token_auth is not None and reader_iam_endpoint is not None:
|
|
reader_token: Final = mint_database_token(token_auth, reader_iam_endpoint)
|
|
read_replica_url = add_missing_query_params(
|
|
reader_iam_endpoint.build_url(reader_token),
|
|
token_refresh_params_from_url(read_replica_url),
|
|
)
|
|
os.environ["DATABASE_URL_READ_REPLICA"] = read_replica_url
|
|
reader_kwargs: Final[dict[str, Any]] = {"datasource": {"url": read_replica_url}}
|
|
if http_client is not None:
|
|
reader_prisma = Prisma(http=http_client, **reader_kwargs)
|
|
else:
|
|
reader_prisma = Prisma(**reader_kwargs)
|
|
reader_wrapper: Final = PrismaWrapper(
|
|
original_prisma=reader_prisma,
|
|
token_auth=token_auth,
|
|
db_url_env_var="DATABASE_URL_READ_REPLICA",
|
|
iam_endpoint=reader_iam_endpoint,
|
|
recreate_uses_datasource=True,
|
|
log_prefix="[reader]",
|
|
)
|
|
self.db = RoutingPrismaWrapper(writer=writer_wrapper, reader=reader_wrapper)
|
|
verbose_proxy_logger.info(
|
|
"PrismaClient: read-replica routing enabled via DATABASE_URL_READ_REPLICA"
|
|
+ (f" (with {token_auth.label} auto-refresh)" if token_auth is not None else "")
|
|
)
|
|
except Exception as e:
|
|
# Reader is opt-in; never let its construction fail proxy
|
|
# startup. Mirrors the runtime contract from
|
|
# `RoutingPrismaWrapper.connect`: reader-side failures are
|
|
# logged and we keep serving traffic via the writer alone.
|
|
# This recovers from transient credential-provider hiccups
|
|
# during the reader token mint, malformed DATABASE_URL_READ_REPLICA,
|
|
# and Prisma construction errors. Operator restart is required
|
|
# to retry read-routing once the underlying issue is resolved.
|
|
verbose_proxy_logger.warning(
|
|
"Failed to initialize read replica Prisma client: %s. "
|
|
"Falling back to writer-only mode (no read routing) until proxy restart.",
|
|
e,
|
|
)
|
|
self.db = writer_wrapper
|
|
else:
|
|
self.db = writer_wrapper # Client to connect to Prisma db
|
|
self._db_reconnect_lock = asyncio.Lock()
|
|
self._db_health_watchdog_task: asyncio.Task | None = None
|
|
self._db_last_reconnect_attempt_ts: float = 0.0
|
|
self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15")))
|
|
self._db_read_only_recreate_ts: float = 0.0
|
|
self._db_read_only_recreate_streak: int = 0
|
|
self._db_health_watchdog_interval_seconds: int = max(
|
|
5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30"))
|
|
)
|
|
self._db_health_watchdog_enabled: bool = (
|
|
str_to_bool(os.getenv("PRISMA_HEALTH_WATCHDOG_ENABLED", "true")) is True
|
|
)
|
|
self._db_health_watchdog_probe_timeout_seconds: float = max(
|
|
0.5,
|
|
float(os.getenv("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", "5.0")),
|
|
)
|
|
self._db_watchdog_reconnect_timeout_seconds: float = max(
|
|
1.0, float(os.getenv("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", "30.0"))
|
|
)
|
|
self._db_auth_reconnect_timeout_seconds: float = max(
|
|
0.5, float(os.getenv("PRISMA_AUTH_RECONNECT_TIMEOUT_SECONDS", "2.0"))
|
|
)
|
|
self._db_auth_reconnect_lock_timeout_seconds: float = max(
|
|
0.0,
|
|
float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")),
|
|
)
|
|
self._consecutive_reconnect_failures: int = 0
|
|
# Last generation of each read engine whose repair was attempted and
|
|
# failed. Scoped to the engine rather than counted globally so an
|
|
# unrelated reconnect failure cannot suppress a stale reader's
|
|
# recovery, and keyed per wrapper rather than held in one slot so a
|
|
# writer failure cannot evict the reader's record and hand the waiver
|
|
# back to a caller whose engine is still unrepaired. Bounded at two
|
|
# entries: a client has one writer and at most one reader.
|
|
self._failed_recreate_generations: Mapping[PrismaWrapper, int] = MappingProxyType({})
|
|
self._reconnect_escalation_threshold: int = max(1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3")))
|
|
self._engine_pidfd: int = -1
|
|
self._engine_pid: int = 0
|
|
self._watching_engine: bool = False
|
|
self._engine_confirmed_dead: bool = False
|
|
self._engine_wait_thread: threading.Thread | None = None
|
|
verbose_proxy_logger.debug("Success - Created Prisma Client")
|
|
|
|
@property
|
|
def writer_db(self) -> PrismaWrapper:
|
|
"""Underlying writer Prisma wrapper, regardless of read-replica routing."""
|
|
if isinstance(self.db, RoutingPrismaWrapper):
|
|
return self.db.writer
|
|
return self.db
|
|
|
|
@property
|
|
def read_db(self) -> PrismaWrapper:
|
|
"""Underlying wrapper that top-level reads are dispatched to.
|
|
|
|
Identical to `writer_db` without a read replica. With one configured
|
|
it is the reader, which is the engine `query_first` actually runs on,
|
|
so anything reasoning about the state of the connection that served a
|
|
read has to consult this rather than the writer.
|
|
"""
|
|
if isinstance(self.db, RoutingPrismaWrapper):
|
|
return self.db.read_target
|
|
return self.db
|
|
|
|
def tx(self, *, timeout: timedelta = _PRISMA_DEFAULT_TX_TIMEOUT) -> "TransactionManager":
|
|
"""Open an interactive transaction on the writer.
|
|
|
|
Callers go through this instead of reaching into ``self.db`` so writer
|
|
selection and read-replica routing stay encapsulated in the wrapper.
|
|
"""
|
|
return cast("TransactionManager", self.db.tx(timeout=timeout)) # cast-ok: untyped __getattr__ delegate
|
|
|
|
def get_request_status(self, payload: dict | SpendLogsPayload) -> Literal["success", "failure"]:
|
|
"""
|
|
Determine if a request was successful or failed based on payload metadata.
|
|
|
|
Args:
|
|
payload (Union[dict, SpendLogsPayload]): Request payload containing metadata
|
|
|
|
Returns:
|
|
Literal["success", "failure"]: Request status
|
|
"""
|
|
try:
|
|
# Get metadata and convert to dict if it's a JSON string
|
|
payload_metadata: Final[dict | SpendLogsMetadata | str] = payload.get("metadata", {})
|
|
if isinstance(payload_metadata, str):
|
|
payload_metadata_json: dict | SpendLogsMetadata = cast(dict, json.loads(payload_metadata))
|
|
else:
|
|
payload_metadata_json = payload_metadata
|
|
|
|
# Check status in metadata dict
|
|
return "failure" if payload_metadata_json.get("status") == "failure" else "success"
|
|
|
|
except (json.JSONDecodeError, AttributeError):
|
|
# Default to success if metadata parsing fails
|
|
return "success"
|
|
|
|
def hash_token(self, token: str):
|
|
# Hash the string using SHA-256
|
|
hashed_token: Final = hashlib.sha256(token.encode()).hexdigest()
|
|
|
|
return hashed_token
|
|
|
|
def jsonify_object(self, data: Mapping[str, object]) -> dict[str, object]:
|
|
db_data: Final[dict[str, object]] = copy.deepcopy(dict(data))
|
|
|
|
for k, v in db_data.items():
|
|
if isinstance(v, dict):
|
|
try:
|
|
db_data[k] = json.dumps(v)
|
|
except Exception:
|
|
# This avoids Prisma retrying this 5 times, and making 5 clients
|
|
db_data[k] = "failed-to-serialize-json"
|
|
return db_data
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def check_view_exists(self):
|
|
"""
|
|
Checks if the LiteLLM_VerificationTokenView and MonthlyGlobalSpend exists in the user's db.
|
|
|
|
LiteLLM_VerificationTokenView: This view is used for getting the token + team data in user_api_key_auth
|
|
|
|
MonthlyGlobalSpend: This view is used for the admin view to see global spend for this month
|
|
|
|
If the view doesn't exist, one will be created.
|
|
"""
|
|
|
|
# Check to see if all of the necessary views exist and if they do, simply return
|
|
# This is more efficient because it lets us check for all views in one
|
|
# query instead of multiple queries.
|
|
try:
|
|
expected_views: Final = [
|
|
"LiteLLM_VerificationTokenView",
|
|
"MonthlyGlobalSpend",
|
|
"Last30dKeysBySpend",
|
|
"Last30dModelsBySpend",
|
|
"MonthlyGlobalSpendPerKey",
|
|
"MonthlyGlobalSpendPerUserPerKey",
|
|
"Last30dTopEndUsersSpend",
|
|
"DailyTagSpend",
|
|
]
|
|
required_view: Final = "LiteLLM_VerificationTokenView"
|
|
expected_views_str: Final = ", ".join(f"'{view}'" for view in expected_views)
|
|
pg_schema: Final = os.getenv("DATABASE_SCHEMA", "public")
|
|
ret: Final[Sequence[_ViewCountRow]] = await self.db.query_raw(f"""
|
|
WITH existing_views AS (
|
|
SELECT viewname
|
|
FROM pg_views
|
|
WHERE schemaname = '{pg_schema}' AND viewname IN (
|
|
{expected_views_str}
|
|
)
|
|
)
|
|
SELECT
|
|
(SELECT COUNT(*) FROM existing_views) AS view_count,
|
|
ARRAY_AGG(viewname) AS view_names
|
|
FROM existing_views
|
|
""")
|
|
expected_total_views: Final = len(expected_views)
|
|
if ret[0]["view_count"] == expected_total_views:
|
|
verbose_proxy_logger.info("All necessary views exist!")
|
|
return
|
|
else:
|
|
## check if required view exists ##
|
|
if ret[0]["view_names"] and required_view not in ret[0]["view_names"]:
|
|
await self.health_check() # make sure we can connect to db
|
|
await create_view_tolerating_race(
|
|
self.db,
|
|
"LiteLLM_VerificationTokenView",
|
|
"""
|
|
CREATE VIEW "LiteLLM_VerificationTokenView" AS
|
|
SELECT
|
|
v.*,
|
|
t.spend AS team_spend,
|
|
t.max_budget AS team_max_budget,
|
|
t.model_max_budget AS team_model_max_budget,
|
|
t.tpm_limit AS team_tpm_limit,
|
|
t.rpm_limit AS team_rpm_limit,
|
|
t.tpd_limit AS team_tpd_limit
|
|
FROM "LiteLLM_VerificationToken" v
|
|
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
|
|
""",
|
|
)
|
|
else:
|
|
should_create_views: Final = await should_create_missing_views(db=self.db)
|
|
if should_create_views:
|
|
await create_missing_views(db=self.db)
|
|
else:
|
|
# don't block execution if these views are missing
|
|
# Convert lists to sets for efficient difference calculation
|
|
ret_view_names_set: Final = set(ret[0]["view_names"]) if ret[0]["view_names"] else set()
|
|
expected_views_set: Final = set(expected_views)
|
|
# Find missing views
|
|
missing_views: Final = expected_views_set - ret_view_names_set
|
|
|
|
verbose_proxy_logger.warning(
|
|
"\n\n\x1b[93mNot all views exist in db, needed for UI 'Usage' tab. Missing=%s.\nRun 'create_views.py' from https://github.com/BerriAI/litellm/tree/main/db_scripts to create missing views.\x1b[0m\n",
|
|
missing_views,
|
|
)
|
|
|
|
except Exception:
|
|
raise
|
|
return
|
|
|
|
@log_db_metrics
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=1, # maximum number of retries
|
|
max_time=2, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def get_generic_data(
|
|
self,
|
|
key: str,
|
|
value: object,
|
|
table_name: Literal["users", "keys", "config", "spend"],
|
|
):
|
|
"""
|
|
Generic implementation of get data.
|
|
|
|
Self-heals across a single transient transport blip via
|
|
`call_with_db_reconnect_retry`: on `httpx.ReadError` /
|
|
`ClientNotConnectedError` / similar, attempt one DB reconnect and
|
|
retry once before surfacing the failure. Restores the 1.82.6 behavior
|
|
that was lost in 1.83.x — see issue #25143.
|
|
"""
|
|
start_time: Final = time.time()
|
|
|
|
async def _do_query():
|
|
if table_name == "users":
|
|
return await UserRepository(self).table.find_first(where={key: value})
|
|
elif table_name == "keys":
|
|
return await VerificationTokenRepository(self).table.find_first(where={key: value})
|
|
elif table_name == "config":
|
|
config_table: Final = cast( # cast-ok: ConfigRepository.table is prisma's litellm_config actions object
|
|
"TableActions[prisma_models.LiteLLM_Config]", ConfigRepository(self).table
|
|
)
|
|
return await config_table.find_first(where={key: value})
|
|
elif table_name == "spend":
|
|
return await self.db.l.find_first(where={key: value})
|
|
return None
|
|
|
|
try:
|
|
return await call_with_db_reconnect_retry(
|
|
self,
|
|
_do_query,
|
|
reason=f"prisma_get_generic_data_{table_name}_lookup_failure",
|
|
)
|
|
except Exception as e:
|
|
error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {e}"
|
|
verbose_proxy_logger.error(error_msg)
|
|
error_msg = error_msg + f"\nException Type: {type(e)}"
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
traceback_str=error_traceback,
|
|
call_type="get_generic_data",
|
|
)
|
|
)
|
|
|
|
raise e
|
|
|
|
async def _query_first_with_cached_plan_fallback(self, sql_query: str, *args) -> dict | None:
|
|
"""
|
|
Execute a query, recovering once from PostgreSQL's "cached plan must not
|
|
change result type" error.
|
|
|
|
That error surfaces during rolling deployments when a schema change
|
|
invalidates the prepared-statement plans that pooled connections still
|
|
hold. Clearing only the server-side plans with DEALLOCATE ALL makes
|
|
things worse: Prisma's query engine keeps a per-connection client-side
|
|
cache of prepared-statement names, so once the server drops a plan the
|
|
engine re-sends a name PostgreSQL no longer recognizes and the
|
|
connection breaks with `prepared statement "sN" does not exist`. With a
|
|
small pool that connection stays poisoned and every auth lookup fails.
|
|
|
|
Recreating the Prisma client kills the engine subprocess and drops the
|
|
server-side plans and the engine's client-side name cache together, so
|
|
the retried query is prepared fresh. We reconnect through
|
|
`attempt_db_reconnect`, which is singleflight: when a schema change
|
|
poisons every pooled connection at once, the first cached-plan error
|
|
recreates the client and the concurrent waiters reuse that single
|
|
recreate instead of racing to kill each other's fresh engine. We pass
|
|
`force_recreate` so the reconnect skips its `SELECT 1` liveness probe:
|
|
the connection is healthy here, it is the prepared statements on it
|
|
that are stale, so a passing probe would otherwise skip the recreate
|
|
and leave the retry to hit the same error. We then retry the identical
|
|
query exactly once.
|
|
|
|
The retry reuses the original query byte-for-byte. Mutating the SQL
|
|
(e.g. injecting a unique comment) would defeat PostgreSQL's plan cache,
|
|
forcing a fresh plan on every request and pegging the database CPU.
|
|
|
|
The reconnect cooldown must not gate the engine this query itself saw
|
|
as stale, or a migration landing within the cooldown of an earlier
|
|
reconnect leaves auth failing until it elapses. The engine observed
|
|
before the query names it, so the reconnect bypasses the cooldown only
|
|
while that same engine is still the live one.
|
|
|
|
It is observed from `read_db`, not `writer_db`: `query_first` is a
|
|
top-level read, so with a read replica configured it runs on the reader
|
|
and it is the reader's prepared statements that went stale. Naming the
|
|
writer here would let an unrelated writer reconnect re-arm the cooldown
|
|
while the reader stayed poisoned.
|
|
"""
|
|
stale_read_engine: Final = _StaleReadEngine.observe(self.read_db)
|
|
try:
|
|
return await self.db.query_first(sql_query, *args)
|
|
except Exception as e:
|
|
if "cached plan must not change result type" not in str(e):
|
|
raise
|
|
verbose_proxy_logger.warning(
|
|
"PostgreSQL cached plan error detected for token lookup; "
|
|
"recreating the database connection and retrying with the same "
|
|
"query. This may occur during rolling deployments when schema "
|
|
"changes are applied."
|
|
)
|
|
await self.attempt_db_reconnect(
|
|
reason="postgres_cached_plan_error",
|
|
force_recreate=True,
|
|
stale_read_engine=stale_read_engine,
|
|
)
|
|
return await self.db.query_first(sql_query, *args)
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
@log_db_metrics
|
|
async def get_data(
|
|
self,
|
|
token: str | list | None = None,
|
|
user_id: str | None = None,
|
|
user_id_list: Sequence[str] | None = None,
|
|
team_id: str | None = None,
|
|
team_id_list: Sequence[str] | None = None,
|
|
key_val: dict | None = None,
|
|
table_name: Literal[
|
|
"user", "key", "config", "spend", "enduser", "budget", "team", "user_notification", "combined_view"
|
|
]
|
|
| None = None,
|
|
query_type: Literal["find_unique", "find_all"] = "find_unique",
|
|
expires: datetime | None = None,
|
|
reset_at: datetime | None = None,
|
|
offset: int | None = None, # pagination, what row number to start from
|
|
limit: int | None = None, # pagination, number of rows to getch when find_all==True
|
|
parent_otel_span: Span | None = None,
|
|
proxy_logging_obj: ProxyLogging | None = None,
|
|
budget_id_list: list[str] | None = None,
|
|
check_deprecated: bool = True,
|
|
):
|
|
args_passed_in: Final = locals()
|
|
start_time: Final = time.time()
|
|
hashed_token: str | None = None
|
|
try:
|
|
response: Any = None
|
|
if (token is not None and table_name is None) or (table_name is not None and table_name == "key"):
|
|
# check if plain text or hash
|
|
if token is not None:
|
|
if isinstance(token, str):
|
|
hashed_token = _hash_token_if_needed(token=token)
|
|
verbose_proxy_logger.debug("PrismaClient: find_unique for token: %s", hashed_token)
|
|
if query_type == "find_unique" and hashed_token is not None:
|
|
if token is None:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={"error": f"No token passed in. Token={token}"},
|
|
)
|
|
response = await VerificationTokenRepository(self).table.find_unique(
|
|
where={"token": hashed_token},
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
if response is not None:
|
|
# for prisma we need to cast the expires time to str
|
|
if response.expires is not None and isinstance(response.expires, datetime):
|
|
response.expires = response.expires.isoformat()
|
|
else:
|
|
# Token does not exist.
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail=f"Authentication Error: invalid user key - user key does not exist in db. User Key={token}",
|
|
)
|
|
elif query_type == "find_all" and user_id is not None:
|
|
response = await VerificationTokenRepository(self).table.find_many(
|
|
where={"user_id": user_id},
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
if response is not None and len(response) > 0:
|
|
for r in response:
|
|
if isinstance(r.expires, datetime):
|
|
r.expires = r.expires.isoformat()
|
|
elif query_type == "find_all" and team_id is not None:
|
|
response = await VerificationTokenRepository(self).table.find_many(
|
|
take=limit,
|
|
where={"team_id": team_id},
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
if response is not None and len(response) > 0:
|
|
for r in response:
|
|
if isinstance(r.expires, datetime):
|
|
r.expires = r.expires.isoformat()
|
|
elif query_type == "find_all" and expires is not None and reset_at is not None:
|
|
response = await VerificationTokenRepository(self).table.find_many(
|
|
take=limit,
|
|
where={
|
|
"OR": [
|
|
{"expires": None},
|
|
{"expires": {"gt": expires}},
|
|
],
|
|
"budget_reset_at": {"lt": reset_at},
|
|
"NOT": {"budget_duration": None},
|
|
},
|
|
)
|
|
if response is not None and len(response) > 0:
|
|
for r in response:
|
|
if isinstance(r.expires, datetime):
|
|
r.expires = r.expires.isoformat()
|
|
elif query_type == "find_all":
|
|
where_filter: Final[dict[str, dict[str, Sequence[str]]]] = {}
|
|
if token is not None:
|
|
where_filter["token"] = {}
|
|
if isinstance(token, str):
|
|
token = _hash_token_if_needed(token=token)
|
|
where_filter["token"]["in"] = [token]
|
|
elif isinstance(token, list):
|
|
hashed_tokens: Final[list[str]] = []
|
|
for t in token:
|
|
assert isinstance(t, str)
|
|
if t.startswith("sk-"):
|
|
new_token = self.hash_token(token=t)
|
|
hashed_tokens.append(new_token)
|
|
else:
|
|
hashed_tokens.append(t)
|
|
where_filter["token"]["in"] = hashed_tokens
|
|
response = await VerificationTokenRepository(self).table.find_many(
|
|
order={"spend": "desc"},
|
|
where=where_filter,
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
if response is not None:
|
|
return response
|
|
else:
|
|
# Token does not exist.
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Authentication Error: invalid user key - token does not exist",
|
|
)
|
|
elif (user_id is not None and table_name is None) or (table_name is not None and table_name == "user"):
|
|
if query_type == "find_unique":
|
|
if key_val is None:
|
|
key_val = {"user_id": user_id}
|
|
|
|
response = await UserRepository(self).table.find_unique(
|
|
where=key_val,
|
|
include={"organization_memberships": True},
|
|
)
|
|
|
|
elif query_type == "find_all" and key_val is not None:
|
|
response = await UserRepository(self).table.find_many(where=key_val)
|
|
elif query_type == "find_all" and reset_at is not None:
|
|
response = await UserRepository(self).table.find_many(
|
|
take=limit,
|
|
where={
|
|
# A user seeded from default_internal_user_params
|
|
# (or created via /user/new without an explicit
|
|
# budget_reset_at) has budget_duration set but
|
|
# budget_reset_at = NULL. `{"lt": reset_at}` never
|
|
# matches NULL, so such users would never be reset
|
|
# and their spend would accumulate for the lifetime
|
|
# of the row, silently exceeding max_budget. Treat a
|
|
# NULL budget_reset_at with a non-NULL budget_duration
|
|
# as due, matching the budget-table query below.
|
|
"NOT": {"budget_duration": None},
|
|
"OR": [
|
|
{"budget_reset_at": None},
|
|
{"budget_reset_at": {"lt": reset_at}},
|
|
],
|
|
},
|
|
)
|
|
elif query_type == "find_all" and user_id_list is not None:
|
|
response = await UserRepository(self).table.find_many(where={"user_id": {"in": user_id_list}})
|
|
elif query_type == "find_all":
|
|
if expires is not None:
|
|
response = await UserRepository(self).table.find_many(
|
|
order={"spend": "desc"},
|
|
where={
|
|
"OR": [
|
|
{"expires": None},
|
|
{"expires": {"gt": expires}},
|
|
],
|
|
},
|
|
)
|
|
else:
|
|
# return all users in the table, get their key aliases ordered by spend
|
|
sql_query = """
|
|
SELECT
|
|
u.*,
|
|
json_agg(v.key_alias) AS key_aliases
|
|
FROM
|
|
"LiteLLM_UserTable" u
|
|
LEFT JOIN "LiteLLM_VerificationToken" v ON u.user_id = v.user_id
|
|
GROUP BY
|
|
u.user_id
|
|
ORDER BY u.spend DESC
|
|
LIMIT $1
|
|
OFFSET $2
|
|
"""
|
|
response = await self.db.query_raw(sql_query, limit, offset)
|
|
return response
|
|
elif table_name == "spend":
|
|
verbose_proxy_logger.debug("PrismaClient: get_data: table_name == 'spend'")
|
|
if key_val is not None:
|
|
if query_type == "find_unique":
|
|
response = await SpendLogsRepository(self).table.find_unique(
|
|
where={
|
|
key_val["key"]: key_val["value"],
|
|
}
|
|
)
|
|
elif query_type == "find_all":
|
|
response = await SpendLogsRepository(self).table.find_many(
|
|
where={
|
|
key_val["key"]: key_val["value"],
|
|
}
|
|
)
|
|
return response
|
|
else:
|
|
response = await SpendLogsRepository(self).table.find_many(
|
|
order={"startTime": "desc"},
|
|
)
|
|
return response
|
|
elif table_name == "budget" and reset_at is not None:
|
|
if query_type == "find_all":
|
|
response = await BudgetRepository(self).table.find_many(
|
|
take=limit,
|
|
where={
|
|
"NOT": {"budget_duration": None},
|
|
"OR": [
|
|
{"budget_reset_at": None},
|
|
{"budget_reset_at": {"lt": reset_at}},
|
|
],
|
|
},
|
|
)
|
|
return response
|
|
|
|
elif table_name == "enduser" and budget_id_list is not None:
|
|
if query_type == "find_all":
|
|
response = await EndUserRepository(self).table.find_many(
|
|
where={"budget_id": {"in": budget_id_list}}
|
|
)
|
|
return response
|
|
elif table_name == "team":
|
|
if query_type == "find_unique":
|
|
response = await TeamRepository(self).table.find_unique(
|
|
where={"team_id": team_id},
|
|
include={"litellm_model_table": True},
|
|
)
|
|
elif query_type == "find_all" and reset_at is not None:
|
|
response = await TeamRepository(self).table.find_many(
|
|
take=limit,
|
|
where={
|
|
# Same NULL budget_reset_at gap as the user query
|
|
# above: a team with a budget_duration but no
|
|
# initialized budget_reset_at would never be reset.
|
|
"NOT": {"budget_duration": None},
|
|
"OR": [
|
|
{"budget_reset_at": None},
|
|
{"budget_reset_at": {"lt": reset_at}},
|
|
],
|
|
},
|
|
)
|
|
elif query_type == "find_all" and user_id is not None:
|
|
response = await TeamRepository(self).table.find_many(
|
|
where={
|
|
"members": {"has": user_id},
|
|
},
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
elif query_type == "find_all" and team_id_list is not None:
|
|
response = await TeamRepository(self).table.find_many(where={"team_id": {"in": team_id_list}})
|
|
elif query_type == "find_all" and team_id_list is None:
|
|
response = await TeamRepository(self).table.find_many(take=MAX_TEAM_LIST_LIMIT)
|
|
return response
|
|
elif table_name == "user_notification":
|
|
if query_type == "find_unique":
|
|
response = await UserNotificationsRepository(self).table.find_unique(where={"user_id": user_id})
|
|
elif query_type == "find_all":
|
|
response = await UserNotificationsRepository(self).table.find_many()
|
|
return response
|
|
elif table_name == "combined_view":
|
|
# check if plain text or hash
|
|
if token is not None:
|
|
if isinstance(token, str):
|
|
hashed_token = _hash_token_if_needed(token=token)
|
|
verbose_proxy_logger.debug("PrismaClient: find_unique for token: %s", hashed_token)
|
|
if query_type == "find_unique":
|
|
if token is None:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={"error": f"No token passed in. Token={token}"},
|
|
)
|
|
|
|
sql_query = """
|
|
SELECT
|
|
v.*,
|
|
t.spend AS team_spend,
|
|
t.max_budget AS team_max_budget,
|
|
t.soft_budget AS team_soft_budget,
|
|
t.model_max_budget AS team_model_max_budget,
|
|
t.tpm_limit AS team_tpm_limit,
|
|
t.rpm_limit AS team_rpm_limit,
|
|
t.tpd_limit AS team_tpd_limit,
|
|
t.models AS team_models,
|
|
t.metadata AS team_metadata,
|
|
t.blocked AS team_blocked,
|
|
t.team_alias AS team_alias,
|
|
t.metadata AS team_metadata,
|
|
t.members_with_roles AS team_members_with_roles,
|
|
t.object_permission_id AS team_object_permission_id,
|
|
t.organization_id as org_id,
|
|
p.project_alias AS project_alias,
|
|
tm.spend AS team_member_spend,
|
|
b_tm.tpm_limit AS team_member_tpm_limit,
|
|
b_tm.rpm_limit AS team_member_rpm_limit,
|
|
m.aliases AS team_model_aliases,
|
|
-- Added comma to separate b.* columns
|
|
b.max_budget AS litellm_budget_table_max_budget,
|
|
b.tpm_limit AS litellm_budget_table_tpm_limit,
|
|
b.rpm_limit AS litellm_budget_table_rpm_limit,
|
|
b.tpd_limit AS litellm_budget_table_tpd_limit,
|
|
b.model_max_budget as litellm_budget_table_model_max_budget,
|
|
b.soft_budget as litellm_budget_table_soft_budget,
|
|
o.metadata as organization_metadata,
|
|
o.organization_alias as organization_alias,
|
|
b2.max_budget as organization_max_budget,
|
|
b2.tpm_limit as organization_tpm_limit,
|
|
b2.rpm_limit as organization_rpm_limit
|
|
FROM "LiteLLM_VerificationToken" AS v
|
|
LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id
|
|
LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id
|
|
LEFT JOIN "LiteLLM_BudgetTable" AS b_tm ON tm.budget_id = b_tm.budget_id
|
|
LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id
|
|
LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id
|
|
LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id
|
|
LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id
|
|
LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id
|
|
WHERE v.token = $1
|
|
"""
|
|
|
|
response = await self._query_first_with_cached_plan_fallback(sql_query, hashed_token)
|
|
|
|
# If not found in main table, check deprecated keys (grace period)
|
|
# check_deprecated=False on the recursive call prevents unbounded chaining
|
|
if response is None and hashed_token is not None and check_deprecated:
|
|
active_token_id: Final = await _lookup_deprecated_key(db=self.db, hashed_token=hashed_token)
|
|
if active_token_id:
|
|
# The recursive call returns a finished
|
|
# LiteLLM_VerificationTokenView; the dict
|
|
# normalization below would crash subscripting it.
|
|
deprecated_response: Final = await self.get_data(
|
|
token=active_token_id,
|
|
table_name="combined_view",
|
|
query_type="find_unique",
|
|
parent_otel_span=parent_otel_span,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
check_deprecated=False,
|
|
)
|
|
if deprecated_response is not None:
|
|
verbose_proxy_logger.debug("Deprecated key used during grace period")
|
|
return deprecated_response
|
|
|
|
if response is not None:
|
|
if response["team_models"] is None:
|
|
response["team_models"] = []
|
|
if response["team_blocked"] is None:
|
|
response["team_blocked"] = False
|
|
|
|
team_member: Member | None = None
|
|
if response["team_members_with_roles"] is not None and response["user_id"] is not None:
|
|
## find the team member corresponding to user id
|
|
"""
|
|
[
|
|
{
|
|
"role": "admin",
|
|
"user_id": "default_user_id",
|
|
"user_email": null
|
|
},
|
|
{
|
|
"role": "user",
|
|
"user_id": null,
|
|
"user_email": "test@email.com"
|
|
}
|
|
]
|
|
"""
|
|
for tm in response["team_members_with_roles"]:
|
|
if tm.get("user_id") is not None and response["user_id"] == tm.get("user_id"):
|
|
team_member = Member(**tm)
|
|
response["team_member"] = team_member
|
|
response = LiteLLM_VerificationTokenView(**response, last_refreshed_at=time.time())
|
|
# for prisma we need to cast the expires time to str
|
|
if response.expires is not None and isinstance(response.expires, datetime):
|
|
response.expires = response.expires.isoformat()
|
|
return response
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
prisma_query_info: Final = (
|
|
f"LiteLLM Prisma Client Exception: Error with `get_data`. Args passed in: {args_passed_in}"
|
|
)
|
|
error_msg: Final = prisma_query_info + str(e)
|
|
print_verbose(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
verbose_proxy_logger.debug(error_traceback)
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="get_data",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
def jsonify_team_object(self, db_data: Mapping[str, object]) -> dict[str, object]:
|
|
db_data = self.jsonify_object(data=db_data)
|
|
if db_data.get("members_with_roles", None) is not None and isinstance(db_data["members_with_roles"], list):
|
|
db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"])
|
|
if db_data.get("budget_limits", None) is not None and isinstance(db_data["budget_limits"], list):
|
|
db_data["budget_limits"] = json.dumps(db_data["budget_limits"])
|
|
return db_data
|
|
|
|
# Define a retrying strategy with exponential backoff
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def insert_data(
|
|
self,
|
|
data: Mapping[str, object],
|
|
table_name: Literal["user", "key", "config", "spend", "team", "user_notification"],
|
|
):
|
|
"""
|
|
Add a key to the database. If it already exists, do nothing.
|
|
"""
|
|
start_time: Final = time.time()
|
|
try:
|
|
verbose_proxy_logger.debug(
|
|
"PrismaClient: insert_data: %s",
|
|
{**data, "token": self.hash_token(token=cast("str", data["token"]))} # cast-ok: a key token is a str
|
|
if data.get("token") is not None
|
|
else data,
|
|
)
|
|
if table_name == "key":
|
|
token: Final = cast("str", data["token"]) # cast-ok: the key table's token column is a str
|
|
hashed_token: Final = self.hash_token(token=token)
|
|
db_data = self.jsonify_object(data=data)
|
|
db_data["token"] = hashed_token
|
|
# Prisma rejects nullable JSON fields set to None (no default).
|
|
# Strip them so the DB stores NULL via the column's nullable constraint.
|
|
if db_data.get("budget_limits") is None:
|
|
db_data.pop("budget_limits", None)
|
|
print_verbose("PrismaClient: Before upsert into litellm_verificationtoken")
|
|
new_verification_token: Final = await VerificationTokenRepository(self).table.upsert(
|
|
where={
|
|
"token": hashed_token,
|
|
},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {}, # don't do anything if it already exists
|
|
},
|
|
include={"litellm_budget_table": True},
|
|
)
|
|
verbose_proxy_logger.info("Data Inserted into Keys Table")
|
|
return new_verification_token
|
|
elif table_name == "user":
|
|
db_data = self.jsonify_object(data=data)
|
|
try:
|
|
new_user_row: Final = await UserRepository(self).table.upsert(
|
|
where={"user_id": data["user_id"]},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {}, # don't do anything if it already exists
|
|
},
|
|
)
|
|
except Exception as e:
|
|
if (
|
|
"Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"
|
|
in str(e)
|
|
):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail={
|
|
"error": f"Foreign Key Constraint failed. Organization ID={db_data['organization_id']} does not exist in LiteLLM_OrganizationTable. Create via `/organization/new`."
|
|
},
|
|
)
|
|
raise e
|
|
verbose_proxy_logger.info("Data Inserted into User Table")
|
|
return new_user_row
|
|
elif table_name == "team":
|
|
db_data = self.jsonify_team_object(db_data=data)
|
|
new_team_row: Final = await TeamRepository(self).table.upsert(
|
|
where={"team_id": data["team_id"]},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {}, # don't do anything if it already exists
|
|
},
|
|
)
|
|
verbose_proxy_logger.info("Data Inserted into Team Table")
|
|
return new_team_row
|
|
elif table_name == "config":
|
|
"""
|
|
For each param,
|
|
get the existing table values
|
|
|
|
Add the new values
|
|
|
|
Update DB
|
|
"""
|
|
tasks: Final = []
|
|
for k, v in data.items():
|
|
updated_data = v
|
|
updated_data = json.dumps(updated_data)
|
|
updated_table_row = ConfigRepository(self).table.upsert(
|
|
where={"param_name": k},
|
|
data={
|
|
"create": {"param_name": k, "param_value": updated_data},
|
|
"update": {"param_value": updated_data},
|
|
},
|
|
)
|
|
|
|
tasks.append(updated_table_row)
|
|
await asyncio.gather(*tasks)
|
|
# invalidate cache so other pods see writes from save_config
|
|
for k in data:
|
|
await invalidate_config_param(k)
|
|
verbose_proxy_logger.info("Data Inserted into Config Table")
|
|
elif table_name == "spend":
|
|
db_data = self.jsonify_object(data=data)
|
|
new_spend_row: Final = await SpendLogsRepository(self).table.upsert(
|
|
where={"request_id": data["request_id"]},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {}, # don't do anything if it already exists
|
|
},
|
|
)
|
|
verbose_proxy_logger.info("Data Inserted into Spend Table")
|
|
return new_spend_row
|
|
elif table_name == "user_notification":
|
|
db_data = self.jsonify_object(data=data)
|
|
new_user_notification_row: Final = await UserNotificationsRepository(self).table.upsert(
|
|
where={"request_id": data["request_id"]},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {}, # don't do anything if it already exists
|
|
},
|
|
)
|
|
verbose_proxy_logger.info("Data Inserted into Model Request Table")
|
|
return new_user_notification_row
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg: Final = f"LiteLLM Prisma Client Exception in insert_data: {e}"
|
|
print_verbose(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="insert_data",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
# Define a retrying strategy with exponential backoff
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def update_data(
|
|
self,
|
|
token: str | None = None,
|
|
data: Mapping[str, object] = {},
|
|
data_list: list | None = None,
|
|
user_id: str | None = None,
|
|
team_id: str | None = None,
|
|
query_type: Literal["update", "update_many"] = "update",
|
|
table_name: Literal["user", "key", "config", "spend", "team", "enduser", "budget"] | None = None,
|
|
update_key_values: dict[str, object] | None = None,
|
|
update_key_values_custom_query: dict[str, object] | None = None,
|
|
):
|
|
"""
|
|
Update existing data
|
|
"""
|
|
verbose_proxy_logger.debug("PrismaClient: update_data, table_name: %s", table_name)
|
|
start_time: Final = time.time()
|
|
try:
|
|
db_data: Final = self.jsonify_object(data=data)
|
|
if update_key_values is not None:
|
|
update_key_values = self.jsonify_object(data=update_key_values)
|
|
if token is not None:
|
|
print_verbose(f"token: [set={token is not None}]")
|
|
# check if plain text or hash
|
|
token = _hash_token_if_needed(token=token)
|
|
db_data["token"] = token
|
|
response: Final = await VerificationTokenRepository(self).table.update(
|
|
where={"token": token},
|
|
data=with_settings_updated_at(db_data),
|
|
)
|
|
verbose_proxy_logger.debug("\033[91m" + f"DB Token Table update succeeded {response}" + "\033[0m")
|
|
_data: dict = {}
|
|
if response is not None:
|
|
try:
|
|
_data = response.model_dump()
|
|
except Exception:
|
|
_data = response.dict() # pyright: ignore[reportDeprecated] # pydantic-v1 row fallback
|
|
return {"token": token, "data": _data}
|
|
elif user_id is not None or (table_name is not None and table_name == "user") and query_type == "update":
|
|
"""
|
|
If data['spend'] + data['user'], update the user table with spend info as well
|
|
"""
|
|
if user_id is None:
|
|
user_id = cast("str", db_data["user_id"]) # cast-ok: the user table's user_id column is a str
|
|
if update_key_values is None:
|
|
if update_key_values_custom_query is not None:
|
|
update_key_values = update_key_values_custom_query
|
|
else:
|
|
update_key_values = db_data
|
|
update_user_row: Final = await UserRepository(self).table.upsert(
|
|
where={"user_id": user_id},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {**update_key_values}, # just update user-specified values, if it already exists
|
|
},
|
|
)
|
|
verbose_proxy_logger.info(
|
|
"\033[91m" + f"DB User Table - update succeeded {update_user_row}" + "\033[0m"
|
|
)
|
|
return {"user_id": user_id, "data": update_user_row}
|
|
elif team_id is not None or (table_name is not None and table_name == "team") and query_type == "update":
|
|
"""
|
|
If data['spend'] + data['user'], update the user table with spend info as well
|
|
"""
|
|
if team_id is None:
|
|
team_id = cast("str | None", db_data["team_id"]) # cast-ok: team_id column is a nullable str
|
|
if update_key_values is None:
|
|
update_key_values = db_data
|
|
if "team_id" not in db_data and team_id is not None:
|
|
db_data["team_id"] = team_id
|
|
if "members_with_roles" in db_data and isinstance(db_data["members_with_roles"], list):
|
|
db_data["members_with_roles"] = json.dumps(db_data["members_with_roles"])
|
|
if "members_with_roles" in update_key_values and isinstance(
|
|
update_key_values["members_with_roles"], list
|
|
):
|
|
update_key_values["members_with_roles"] = json.dumps(update_key_values["members_with_roles"])
|
|
update_team_row: Final = await TeamRepository(self).table.upsert(
|
|
where={"team_id": team_id},
|
|
data={
|
|
"create": {**db_data},
|
|
"update": {**update_key_values}, # just update user-specified values, if it already exists
|
|
},
|
|
)
|
|
verbose_proxy_logger.info(
|
|
"\033[91m" + f"DB Team Table - update succeeded {update_team_row}" + "\033[0m"
|
|
)
|
|
return {"team_id": team_id, "data": update_team_row}
|
|
elif (
|
|
table_name is not None
|
|
and table_name == "key"
|
|
and query_type == "update_many"
|
|
and data_list is not None
|
|
and isinstance(data_list, list)
|
|
):
|
|
"""
|
|
Batch write update queries
|
|
"""
|
|
batcher = self.db.batch_()
|
|
for idx, t in enumerate(data_list):
|
|
# check if plain text or hash
|
|
if t.token.startswith("sk-"):
|
|
t.token = self.hash_token(token=t.token)
|
|
try:
|
|
data_json = self.jsonify_object(data=t.model_dump(exclude_none=True))
|
|
except Exception:
|
|
data_json = self.jsonify_object(data=t.dict(exclude_none=True))
|
|
batcher.litellm_verificationtoken.update(
|
|
where={"token": t.token},
|
|
data={**data_json},
|
|
)
|
|
await batcher.commit()
|
|
print_verbose("\033[91m" + "DB Token Table update succeeded" + "\033[0m")
|
|
elif (
|
|
table_name is not None
|
|
and table_name == "user"
|
|
and query_type == "update_many"
|
|
and data_list is not None
|
|
and isinstance(data_list, list)
|
|
):
|
|
"""
|
|
Batch write update queries
|
|
"""
|
|
batcher = self.db.batch_()
|
|
for idx, user in enumerate(data_list):
|
|
try:
|
|
data_json = self.jsonify_object(data=user.model_dump(exclude_none=True))
|
|
except Exception:
|
|
data_json = self.jsonify_object(data=user.dict())
|
|
batcher.litellm_usertable.upsert(
|
|
where={"user_id": user.user_id},
|
|
data={
|
|
"create": {**data_json},
|
|
"update": {**data_json}, # just update user-specified values, if it already exists
|
|
},
|
|
)
|
|
await batcher.commit()
|
|
verbose_proxy_logger.info("\033[91m" + "DB User Table Batch update succeeded" + "\033[0m")
|
|
elif (
|
|
table_name is not None
|
|
and table_name == "enduser"
|
|
and query_type == "update_many"
|
|
and data_list is not None
|
|
and isinstance(data_list, list)
|
|
):
|
|
"""
|
|
Batch write update queries
|
|
"""
|
|
batcher = self.db.batch_()
|
|
for enduser in data_list:
|
|
try:
|
|
data_json = self.jsonify_object(data=enduser.model_dump(exclude_none=True))
|
|
except Exception:
|
|
data_json = self.jsonify_object(data=enduser.dict())
|
|
batcher.litellm_endusertable.upsert(
|
|
where={"user_id": enduser.user_id},
|
|
data={
|
|
"create": {**data_json},
|
|
"update": {**data_json}, # just update end-user-specified values, if it already exists
|
|
},
|
|
)
|
|
await batcher.commit()
|
|
verbose_proxy_logger.info("\033[91m" + "DB End User Table Batch update succeeded" + "\033[0m")
|
|
elif (
|
|
table_name is not None
|
|
and table_name == "budget"
|
|
and query_type == "update_many"
|
|
and data_list is not None
|
|
and isinstance(data_list, list)
|
|
):
|
|
"""
|
|
Batch write update queries
|
|
"""
|
|
batcher = self.db.batch_()
|
|
for budget in data_list:
|
|
try:
|
|
data_json = self.jsonify_object(data=budget.model_dump(exclude_none=True))
|
|
except Exception:
|
|
data_json = self.jsonify_object(data=budget.dict())
|
|
batcher.litellm_budgettable.upsert(
|
|
where={"budget_id": budget.budget_id},
|
|
data={
|
|
"create": {**data_json},
|
|
"update": {**data_json}, # just update end-user-specified values, if it already exists
|
|
},
|
|
)
|
|
await batcher.commit()
|
|
verbose_proxy_logger.info("\033[91m" + "DB Budget Table Batch update succeeded" + "\033[0m")
|
|
elif (
|
|
table_name is not None
|
|
and table_name == "team"
|
|
and query_type == "update_many"
|
|
and data_list is not None
|
|
and isinstance(data_list, list)
|
|
):
|
|
# Batch write update queries
|
|
batcher = self.db.batch_()
|
|
for idx, team in enumerate(data_list):
|
|
try:
|
|
data_json = self.jsonify_team_object(db_data=team.model_dump(exclude_none=True))
|
|
except Exception:
|
|
data_json = self.jsonify_object(data=team.dict(exclude_none=True))
|
|
batcher.litellm_teamtable.upsert(
|
|
where={"team_id": team.team_id},
|
|
data={
|
|
"create": {**data_json},
|
|
"update": {**data_json}, # just update user-specified values, if it already exists
|
|
},
|
|
)
|
|
await batcher.commit()
|
|
verbose_proxy_logger.info("\033[91m" + "DB Team Table Batch update succeeded" + "\033[0m")
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg: Final = f"LiteLLM Prisma Client Exception - update_data: {e}"
|
|
print_verbose(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="update_data",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
# Define a retrying strategy with exponential backoff
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def delete_data(
|
|
self,
|
|
tokens: Sequence[str | None] | None = None,
|
|
team_id_list: Sequence[str] | None = None,
|
|
table_name: Literal["user", "key", "config", "spend", "team"] | None = None,
|
|
user_id: str | None = None,
|
|
):
|
|
"""
|
|
Allow user to delete a key(s)
|
|
|
|
Ensure user owns that key, unless admin.
|
|
"""
|
|
start_time: Final = time.time()
|
|
try:
|
|
if tokens is not None and isinstance(tokens, list):
|
|
hashed_tokens: Final[list[str | None]] = []
|
|
for token in tokens:
|
|
if isinstance(token, str) and token.startswith("sk-"):
|
|
hashed_token = self.hash_token(token=token)
|
|
else:
|
|
hashed_token = token
|
|
hashed_tokens.append(hashed_token)
|
|
filter_query: dict[str, object] = {}
|
|
if user_id is not None:
|
|
filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]}
|
|
else:
|
|
filter_query = {"token": {"in": hashed_tokens}}
|
|
|
|
deleted_tokens: Final[int] = await VerificationTokenRepository(self).table.delete_many(
|
|
where=filter_query
|
|
)
|
|
verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens)
|
|
return {"deleted_keys": deleted_tokens}
|
|
elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list):
|
|
# admin only endpoint -> `/team/delete`
|
|
await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}})
|
|
return {"deleted_teams": team_id_list}
|
|
elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list):
|
|
# admin only endpoint -> `/team/delete`
|
|
await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}})
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg: Final = f"LiteLLM Prisma Client Exception - delete_data: {e}"
|
|
print_verbose(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="delete_data",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
# Define a retrying strategy with exponential backoff
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def connect(self):
|
|
start_time: Final = time.time()
|
|
try:
|
|
verbose_proxy_logger.debug("PrismaClient: connect() called Attempting to Connect to DB")
|
|
if self.db.is_connected() is False:
|
|
verbose_proxy_logger.debug("PrismaClient: DB not connected, Attempting to Connect to DB")
|
|
await self.db.connect()
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg: Final = f"LiteLLM Prisma Client Exception connect(): {e}"
|
|
verbose_proxy_logger.warning(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="connect",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
# Define a retrying strategy with exponential backoff
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception, # base exception to catch for the backoff
|
|
max_tries=3, # maximum number of retries
|
|
max_time=10, # maximum total time to retry for
|
|
on_backoff=on_backoff, # specifying the function to call on backoff
|
|
)
|
|
async def disconnect(self):
|
|
start_time: Final = time.time()
|
|
try:
|
|
await self.db.disconnect()
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg: Final = f"LiteLLM Prisma Client Exception disconnect(): {e}"
|
|
print_verbose(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="disconnect",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
def _get_engine_pid(self) -> int:
|
|
"""Get the PID of the writer's engine subprocess, or 0 if unavailable.
|
|
|
|
Must never raise: prisma's ``_engine`` property raises
|
|
``ClientNotConnectedError`` on a disconnected client, and an exception
|
|
escaping from the reconnect path would leave it unable to recover.
|
|
"""
|
|
try:
|
|
prisma_obj: Final = self.writer_db._original_prisma
|
|
if prisma_obj.is_connected() is not True:
|
|
return 0
|
|
engine: Final = prisma_obj._engine
|
|
process: Final = getattr(engine, "process", None) if engine is not None else None
|
|
if process is not None:
|
|
pid: Final[object] = process.pid
|
|
if isinstance(pid, int):
|
|
return pid
|
|
except (AttributeError, TypeError):
|
|
pass
|
|
return 0
|
|
|
|
def _is_engine_alive(self) -> bool:
|
|
if self._engine_pid <= 0:
|
|
return True
|
|
try:
|
|
os.kill(self._engine_pid, 0)
|
|
return True
|
|
except ProcessLookupError:
|
|
return False
|
|
except (PermissionError, OSError):
|
|
return True
|
|
|
|
@staticmethod
|
|
def _reap_all_zombies() -> set:
|
|
"""Reap ALL zombie child processes via waitpid(-1, WNOHANG).
|
|
|
|
Returns a set of reaped PIDs. As PID 1 in Docker (or any
|
|
process that spawns children), we must reap ALL terminated
|
|
children to prevent zombie accumulation.
|
|
|
|
No-op on Windows: os.waitpid and os.WNOHANG are Unix-only.
|
|
"""
|
|
if sys.platform == "win32":
|
|
return set()
|
|
reaped: Final[set] = set()
|
|
while True:
|
|
try:
|
|
pid, _ = os.waitpid(-1, os.WNOHANG)
|
|
if pid == 0:
|
|
break
|
|
reaped.add(pid)
|
|
except ChildProcessError:
|
|
break
|
|
return reaped
|
|
|
|
def _try_waitpid_watch(self, pid: int) -> bool:
|
|
"""Watch engine PID via os.waitpid() in a dedicated thread.
|
|
|
|
The thread blocks on os.waitpid(pid, 0) which is a kernel-level
|
|
wait and with zero CPU overhead, instant detection when the process exits.
|
|
When the process dies, the thread notifies the asyncio event loop
|
|
via call_soon_threadsafe.
|
|
|
|
Returns True if the thread was started, False on failure.
|
|
On Windows, returns False immediately (os.waitpid/WNOHANG are Unix-only);
|
|
caller falls back to os.kill polling.
|
|
"""
|
|
if sys.platform == "win32":
|
|
return False
|
|
try:
|
|
probe_pid, _ = os.waitpid(pid, os.WNOHANG)
|
|
except ChildProcessError:
|
|
verbose_proxy_logger.debug(
|
|
"PID %s is not a child process; skipping waitpid watch.",
|
|
pid,
|
|
)
|
|
return False
|
|
|
|
if probe_pid == pid:
|
|
verbose_proxy_logger.warning(
|
|
"prisma-query-engine PID %s already dead at watch start.",
|
|
pid,
|
|
)
|
|
if self._consume_expected_death(pid):
|
|
verbose_proxy_logger.info(
|
|
"PID %s death was planned (engine already replaced); not reconnecting.",
|
|
pid,
|
|
)
|
|
self._cleanup_engine_watcher()
|
|
return True
|
|
self._engine_confirmed_dead = True
|
|
self._reap_all_zombies()
|
|
self._cleanup_engine_watcher()
|
|
asyncio.create_task(
|
|
self.attempt_db_reconnect(
|
|
reason="engine_process_death",
|
|
force=True,
|
|
)
|
|
)
|
|
return True
|
|
|
|
try:
|
|
loop: Final = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return False
|
|
|
|
thread: Final = threading.Thread(
|
|
target=self._waitpid_thread_func,
|
|
args=(pid, loop),
|
|
daemon=True,
|
|
name=f"prisma-engine-waitpid-{pid}",
|
|
)
|
|
thread.start()
|
|
self._engine_wait_thread = thread
|
|
return True
|
|
|
|
def _waitpid_thread_func(self, pid: int, loop: asyncio.AbstractEventLoop) -> None:
|
|
"""Thread function: block until engine PID exits, then notify event loop.
|
|
|
|
Note: uvloop/libuv may reap the child first via waitpid(-1, WNOHANG)
|
|
in its SIGCHLD handler. In that case our waitpid raises ChildProcessError.
|
|
we still notify the event loop because the engine is dead either way.
|
|
"""
|
|
try:
|
|
os.waitpid(pid, 0)
|
|
except ChildProcessError:
|
|
pass
|
|
except OSError:
|
|
pass
|
|
try:
|
|
loop.call_soon_threadsafe(self._on_engine_death_from_thread, pid)
|
|
except RuntimeError:
|
|
pass
|
|
|
|
def _consume_expected_death(self, pid: int) -> bool:
|
|
"""True iff ``pid`` was killed on purpose by a planned recreate.
|
|
|
|
`PrismaWrapper.recreate_prisma_client` records the old engine PID in
|
|
`_expected_engine_deaths` before SIGTERM-ing it (IAM token refresh,
|
|
guarded reconnect). When the watcher then sees that PID die, this lets
|
|
it recognize the death as planned and skip its own reconnect, which
|
|
would otherwise kill the engine the recreate just spawned (#29176).
|
|
|
|
Consumes (removes) the PID so a later real crash of a reused PID is
|
|
still handled. Tolerant of `self.db` stand-ins (tests / older clients)
|
|
that don't expose a real set.
|
|
"""
|
|
expected: Final = getattr(self.db, "_expected_engine_deaths", None)
|
|
if isinstance(expected, set) and pid in expected:
|
|
expected.discard(pid)
|
|
return True
|
|
return False
|
|
|
|
def _on_engine_death_from_thread(self, dead_pid: int) -> None:
|
|
"""Called on the event loop thread when the waitpid thread detects engine death."""
|
|
if self._engine_confirmed_dead:
|
|
return
|
|
if dead_pid != self._engine_pid:
|
|
return
|
|
if self._consume_expected_death(dead_pid):
|
|
verbose_proxy_logger.info(
|
|
"prisma-query-engine PID %s exited as part of a planned restart; "
|
|
"not reconnecting (engine already replaced).",
|
|
dead_pid,
|
|
)
|
|
self._cleanup_engine_watcher()
|
|
return
|
|
verbose_proxy_logger.error(
|
|
"prisma-query-engine PID %s exited (waitpid thread); triggering reconnect.",
|
|
dead_pid,
|
|
)
|
|
self._engine_confirmed_dead = True
|
|
self._reap_all_zombies()
|
|
self._cleanup_engine_watcher()
|
|
asyncio.create_task(
|
|
self.attempt_db_reconnect(
|
|
reason="engine_process_death",
|
|
force=True,
|
|
)
|
|
)
|
|
|
|
def _try_pidfd_watch(self, pid: int) -> bool:
|
|
"""
|
|
Watch engine PID via pidfd_open + asyncio event loop reader.
|
|
|
|
Returns True if pidfd watch was set up, False if unavailable or failed.
|
|
Broad OSError catch handles both ENOSYS and SECCOMP-blocked syscalls.
|
|
"""
|
|
if not hasattr(os, "pidfd_open"):
|
|
return False
|
|
fd = -1
|
|
try:
|
|
fd = os.pidfd_open(pid, 0)
|
|
asyncio.get_running_loop().add_reader(fd, self._on_pidfd_readable)
|
|
self._engine_pidfd = fd
|
|
return True
|
|
except OSError:
|
|
if fd >= 0:
|
|
os.close(fd)
|
|
return False
|
|
|
|
def _on_pidfd_readable(self) -> None:
|
|
"""pidfd became readable: engine process exited or became zombie.
|
|
|
|
Sets _engine_confirmed_dead BEFORE cleanup so _run_reconnect_cycle
|
|
takes the heavy path (recreate Prisma client + re-arm watcher).
|
|
"""
|
|
if self._engine_confirmed_dead:
|
|
# Already handled -- just clean up pidfd resources.
|
|
if self._engine_pidfd >= 0:
|
|
try:
|
|
asyncio.get_running_loop().remove_reader(self._engine_pidfd)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
os.close(self._engine_pidfd)
|
|
except OSError:
|
|
pass
|
|
self._engine_pidfd = -1
|
|
return
|
|
dead_pid: Final = self._engine_pid
|
|
if self._consume_expected_death(dead_pid):
|
|
verbose_proxy_logger.info(
|
|
"prisma-query-engine PID %s exited (pidfd event) as part of a "
|
|
"planned restart; not reconnecting (engine already replaced).",
|
|
dead_pid,
|
|
)
|
|
self._cleanup_engine_watcher()
|
|
return
|
|
verbose_proxy_logger.error(
|
|
"prisma-query-engine PID %s exited (pidfd event); triggering reconnect.",
|
|
dead_pid,
|
|
)
|
|
self._engine_confirmed_dead = True
|
|
self._reap_all_zombies()
|
|
self._cleanup_engine_watcher()
|
|
asyncio.create_task(
|
|
self.attempt_db_reconnect(
|
|
reason="engine_process_death",
|
|
force=True,
|
|
)
|
|
)
|
|
|
|
async def _poll_engine_proc(self) -> None:
|
|
"""poll via os.kill(pid, 0) every 1s.
|
|
Only used when BOTH waitpid thread and pidfd are unavailable
|
|
(e.g., PID is not our child process and pidfd_open fails)
|
|
"""
|
|
while self._watching_engine and self._engine_pid > 0:
|
|
try:
|
|
os.kill(self._engine_pid, 0)
|
|
except ProcessLookupError:
|
|
dead_pid = self._engine_pid
|
|
if self._consume_expected_death(dead_pid):
|
|
verbose_proxy_logger.info(
|
|
"prisma-query-engine PID %s gone as part of a planned "
|
|
"restart; not reconnecting (engine already replaced).",
|
|
dead_pid,
|
|
)
|
|
self._cleanup_engine_watcher()
|
|
return
|
|
verbose_proxy_logger.error(
|
|
"prisma-query-engine PID %s gone; triggering reconnect.",
|
|
dead_pid,
|
|
)
|
|
self._engine_confirmed_dead = True
|
|
self._reap_all_zombies()
|
|
self._cleanup_engine_watcher()
|
|
await self.attempt_db_reconnect(
|
|
reason="engine_process_death",
|
|
force=True,
|
|
)
|
|
return
|
|
except (PermissionError, OSError):
|
|
verbose_proxy_logger.debug(
|
|
"Cannot signal PID %s; stopping engine poll.",
|
|
self._engine_pid,
|
|
)
|
|
self._cleanup_engine_watcher()
|
|
return
|
|
await asyncio.sleep(1)
|
|
|
|
def _cleanup_engine_watcher(self) -> None:
|
|
"""Clean up pidfd reader, waitpid thread ref, or stop polling and reset state."""
|
|
self._watching_engine = False
|
|
if self._engine_pidfd >= 0:
|
|
try:
|
|
asyncio.get_running_loop().remove_reader(self._engine_pidfd)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
os.close(self._engine_pidfd)
|
|
except OSError:
|
|
pass
|
|
self._engine_pidfd = -1
|
|
self._engine_wait_thread = None
|
|
self._engine_pid = 0
|
|
|
|
async def _start_engine_watcher(self) -> None:
|
|
"""
|
|
Start watching the Prisma query engine process for death.
|
|
|
|
Detection priority:
|
|
1. os.waitpid() in a dedicated thread, works with all event loops.
|
|
2. pidfd_open kernel fd registered with asyncio.
|
|
3. os.kill(pid, 0) polling (1s), last-resort fallback when neither
|
|
waitpid thread nor pidfd are available.
|
|
|
|
"""
|
|
if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None:
|
|
return
|
|
pid: Final = self._get_engine_pid()
|
|
if pid == 0:
|
|
verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.")
|
|
return
|
|
self._engine_pid = pid
|
|
self._engine_confirmed_dead = False
|
|
verbose_proxy_logger.info("Found prisma-query-engine at PID %s.", pid)
|
|
waitpid_ok: Final = self._try_waitpid_watch(pid)
|
|
pidfd_ok: Final = False if waitpid_ok else self._try_pidfd_watch(pid)
|
|
if waitpid_ok:
|
|
verbose_proxy_logger.info(
|
|
"Watching engine PID %s via waitpid thread.",
|
|
pid,
|
|
)
|
|
elif pidfd_ok:
|
|
verbose_proxy_logger.info(
|
|
"Watching engine PID %s via pidfd.",
|
|
pid,
|
|
)
|
|
else:
|
|
verbose_proxy_logger.info(
|
|
"Watching engine PID %s via os.kill polling.",
|
|
pid,
|
|
)
|
|
self._watching_engine = True
|
|
asyncio.create_task(self._poll_engine_proc())
|
|
|
|
def _stop_engine_watcher(self) -> None:
|
|
"""Stop watching the engine process and clean up all resources."""
|
|
self._cleanup_engine_watcher()
|
|
self._engine_confirmed_dead = False
|
|
verbose_proxy_logger.debug("Stopped engine process watcher.")
|
|
|
|
def _handle_writer_engine_replaced(self) -> None:
|
|
"""Re-arm the engine watcher after a planned writer-engine restart.
|
|
|
|
Wired as `PrismaWrapper.on_engine_replaced` and invoked from inside
|
|
`recreate_prisma_client` once the new engine is connected (IAM token
|
|
refresh, guarded reconnect). The old watcher was tracking the engine
|
|
we just intentionally killed, so we tear it down and re-arm on the new
|
|
PID. Scheduling `_start_engine_watcher` as a task (rather than awaiting)
|
|
keeps us from blocking the recreate while it still holds the wrapper's
|
|
reconnection lock. Without this re-arm, a planned restart would leave
|
|
the proxy with no engine-death detection until the next reconnect.
|
|
"""
|
|
self._engine_confirmed_dead = False
|
|
self._cleanup_engine_watcher()
|
|
asyncio.create_task(self._start_engine_watcher())
|
|
|
|
async def _run_reconnect_cycle(
|
|
self,
|
|
timeout_seconds: float | None = None,
|
|
force_recreate: bool = False,
|
|
) -> None:
|
|
"""
|
|
Run a reconnect cycle with a single overall timeout budget.
|
|
|
|
Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll
|
|
handlers) to choose between heavy reconnect (engine dead -- recreate
|
|
Prisma client, re-arm watcher) and direct reconnect (network blip --
|
|
recreate Prisma client, re-arm watcher, SELECT 1). Both paths recreate
|
|
the client via the non-blocking kill-then-construct flow rather than
|
|
calling disconnect(), which blocks the event loop on the synchronous
|
|
subprocess.Popen.wait() inside prisma-client-py (see issue #26191).
|
|
|
|
`force_recreate` skips the direct path's liveness probe, for callers
|
|
whose failure lives in the session state rather than the connection
|
|
(stale prepared statements after a schema change): a reachable writer
|
|
proves nothing about those, so the probe must not veto the recreate.
|
|
"""
|
|
effective_timeout: Final = (
|
|
timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds
|
|
)
|
|
|
|
# Snapshot the writer's engine generation BEFORE any await. Both
|
|
# reconnect branches forward it to recreate_prisma_client as an
|
|
# optimistic-lock token: if a concurrent IAM token refresh replaces the
|
|
# engine after this point, the generation moves and the recreate becomes
|
|
# a no-op instead of killing the engine the refresh just spawned
|
|
# (#29176). Captured here — atomically with the dead-engine decision
|
|
# below — rather than inside the reconnect closures, because those run
|
|
# after an `asyncio.wait_for(...)` yield during which a refresh could
|
|
# otherwise slip in and bump the very generation the closure then reads.
|
|
expected_generation: Final = getattr(self.writer_db, "_engine_generation", None)
|
|
|
|
engine_is_dead: Final = self._engine_confirmed_dead or (self._engine_pid > 0 and not self._is_engine_alive())
|
|
|
|
if engine_is_dead:
|
|
dead_pid: Final = self._engine_pid
|
|
verbose_proxy_logger.warning(
|
|
"prisma-query-engine PID %s is dead; reconnecting.",
|
|
dead_pid,
|
|
)
|
|
self._reap_all_zombies()
|
|
self._cleanup_engine_watcher()
|
|
|
|
async def _do_heavy_reconnect() -> None:
|
|
db_url: Final = os.getenv("DATABASE_URL", "")
|
|
if not db_url:
|
|
verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.")
|
|
raise RuntimeError("DATABASE_URL not set")
|
|
# Forward the entry-snapshot generation. The engine was
|
|
# confirmed dead, but a concurrent IAM refresh may have already
|
|
# respawned it; the guard makes this recreate a no-op in that
|
|
# case rather than killing the fresh engine (#29176). Unlike the
|
|
# direct path there is no SELECT 1 probe here, so the generation
|
|
# guard is the only thing standing between a crash-reconnect and
|
|
# a refresh that raced it.
|
|
recreated: Final = await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation)
|
|
await self._start_engine_watcher()
|
|
# Same contract as the direct path below: a forced caller asked
|
|
# for its engine to be replaced, so a decline is not a success.
|
|
# Reachable here because the escalation threshold flips
|
|
# `_engine_confirmed_dead`, which routes the next cycle, forced
|
|
# callers included, down this branch.
|
|
if force_recreate is True and recreated is False:
|
|
# Clear the dead-engine flag first, restoring the policy the
|
|
# non-forced path already has: a decline does not raise for
|
|
# it, so it falls through to the clear below. Only the
|
|
# forced branch would strand the flag, and stranding it
|
|
# routes the next cycle back down this probe-free branch,
|
|
# where the refreshed generation now matches and the
|
|
# recreate kills the healthy engine a refresh just spawned
|
|
# (#29176). This has to stay AFTER `_start_engine_watcher`
|
|
# above: clearing the flag while the watcher is still torn
|
|
# down would be worse than either alone.
|
|
self._engine_confirmed_dead = False
|
|
raise _ForcedRecreateDeclined(
|
|
"Forced Prisma recreate declined by the generation guard; "
|
|
"the engine that failed was not replaced"
|
|
)
|
|
|
|
await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout)
|
|
# Only clear the "dead engine" flag after the heavy reconnect
|
|
# actually completed. If `_do_heavy_reconnect()` raises (timeout,
|
|
# missing DATABASE_URL, recreate failure), the flag stays True so
|
|
# the next attempt re-enters the heavy branch instead of silently
|
|
# demoting to the lightweight path.
|
|
self._engine_confirmed_dead = False
|
|
else:
|
|
verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).")
|
|
|
|
async def _do_direct_reconnect() -> None:
|
|
db_url: Final = os.getenv("DATABASE_URL", "")
|
|
if not db_url:
|
|
verbose_proxy_logger.error("DATABASE_URL not set; cannot reconnect Prisma client.")
|
|
raise RuntimeError("DATABASE_URL not set")
|
|
# Probe the writer BEFORE recreating. A concurrent IAM token
|
|
# refresh may have just replaced the engine (issue #29176); if
|
|
# the writer answers SELECT 1 the connection is already healthy
|
|
# and recreating would needlessly kill that fresh engine. If we
|
|
# do recreate, the entry-snapshot generation lets the wrapper
|
|
# detect a refresh that landed since cycle entry and skip the
|
|
# redundant restart.
|
|
writer: Final = self.writer_db
|
|
if force_recreate is False:
|
|
try:
|
|
if await self._writer_is_read_only(writer):
|
|
verbose_proxy_logger.warning(
|
|
"Writer answers the probe but its session is read-only "
|
|
"(writes fail with SQLSTATE 25006); recreating Prisma client."
|
|
)
|
|
else:
|
|
verbose_proxy_logger.info(
|
|
"Writer healthy on probe; skipping recreate (engine "
|
|
"likely already replaced by a token refresh)."
|
|
)
|
|
if isinstance(self.db, RoutingPrismaWrapper):
|
|
self.db.mark_writer_recovered()
|
|
await self._start_engine_watcher()
|
|
return
|
|
except Exception as probe_err:
|
|
verbose_proxy_logger.warning(
|
|
"Writer probe failed (%s); recreating Prisma client.",
|
|
probe_err,
|
|
)
|
|
# Fresh Prisma client + new engine subprocess. The previous
|
|
# "lightweight" path called `disconnect()` which blocks the
|
|
# event loop on `subprocess.Popen.wait()`; since that call
|
|
# ends up killing the engine anyway, we do it non-blockingly
|
|
# via `_kill_engine_process` inside `recreate_prisma_client`.
|
|
self._cleanup_engine_watcher()
|
|
recreated: Final = await self.db.recreate_prisma_client(db_url, expected_generation=expected_generation)
|
|
await self._start_engine_watcher()
|
|
# Smoke-test the writer specifically; query_raw on the routing
|
|
# wrapper sends to the reader, which would not validate the
|
|
# newly-recreated writer engine. The reader is left to the
|
|
# caller's own retried query, a stronger check than SELECT 1,
|
|
# and a reader that fails to come back sets `_reader_unavailable`
|
|
# so reads fall through to the writer just recreated here.
|
|
await self.writer_db.query_raw("SELECT 1")
|
|
# A recreate can decline: the optimistic-lock guard no-ops when
|
|
# the writer generation moved since cycle entry, and the routing
|
|
# wrapper then leaves the reader untouched as well. Callers that
|
|
# merely suspect a transport blip are happy either way, but a
|
|
# forced caller asked for this engine to be replaced because its
|
|
# session state is poisoned, and it was not. Do not report that
|
|
# as a success: it would reset the consecutive-failure count and
|
|
# log a repair that never happened.
|
|
if force_recreate is True and recreated is False:
|
|
raise _ForcedRecreateDeclined(
|
|
"Forced Prisma recreate declined by the generation guard; "
|
|
"the engine that failed was not replaced"
|
|
)
|
|
|
|
await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout)
|
|
|
|
def _cooldown_applies(self, stale_read_engine: "_StaleReadEngine | None") -> bool:
|
|
"""
|
|
Whether the reconnect cooldown should still gate this caller.
|
|
|
|
The cooldown collapses a burst of callers onto one recreate, so it
|
|
keeps gating a caller whose named engine has already been replaced:
|
|
that recreate is the one it was waiting for. While that engine is still
|
|
the live one the damage is still being served, so deferring to an
|
|
unrelated reconnect's cooldown would leave it broken until the cooldown
|
|
elapses.
|
|
|
|
A named engine always describes the one that served the failing read
|
|
(see `_query_first_with_cached_plan_fallback`), so it is compared
|
|
against `read_db`, identity included: `read_db` can resolve to a
|
|
different wrapper than it did at observation time.
|
|
|
|
The waiver is withdrawn once a repair of this same engine has been
|
|
tried and failed. A failed recreate leaves the generation where it was,
|
|
so without this every queued caller would still see its own engine live
|
|
and run its own full recreate serially instead of collapsing onto one
|
|
attempt, which is what the cooldown is for. The record is scoped to the
|
|
engine rather than to a global failure count: an unrelated reconnect
|
|
failing somewhere else says nothing about whether this engine can be
|
|
repaired, and gating on it would suppress the recovery this method
|
|
exists to allow.
|
|
|
|
The record is never cleared, and does not need to be. Generations are
|
|
monotonic per wrapper, so once the engine is repaired every later
|
|
caller names a higher one and the entry can never match again. And this
|
|
method is only ever the first half of the gate: the cooldown window
|
|
itself still expires, so an engine that can never be repaired degrades
|
|
to the plain cooldown rather than being suppressed forever.
|
|
"""
|
|
if stale_read_engine is None:
|
|
return True
|
|
if self._failed_recreate_generations.get(stale_read_engine.wrapper) == stale_read_engine.generation:
|
|
return True
|
|
return not stale_read_engine.is_still_live(self.read_db)
|
|
|
|
async def _attempt_reconnect_inside_lock(
|
|
self,
|
|
force: bool,
|
|
reason: str,
|
|
timeout_seconds: float | None,
|
|
force_recreate: bool = False,
|
|
stale_read_engine: "_StaleReadEngine | None" = None,
|
|
) -> bool:
|
|
now: Final = time.time()
|
|
if (
|
|
force is False
|
|
and self._cooldown_applies(stale_read_engine)
|
|
and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds
|
|
):
|
|
verbose_proxy_logger.debug(
|
|
"Skipping DB reconnect attempt inside lock due to cooldown. reason=%s",
|
|
reason,
|
|
)
|
|
return False
|
|
|
|
# Escalate to heavy reconnect after consecutive lightweight failures.
|
|
# When the Prisma engine process is alive but not accepting connections
|
|
# (e.g., startup race condition), lightweight reconnects (disconnect +
|
|
# connect) will never succeed. Force a full Prisma client recreation
|
|
# to recover from this state.
|
|
if self._consecutive_reconnect_failures >= self._reconnect_escalation_threshold:
|
|
verbose_proxy_logger.warning(
|
|
"Escalating to heavy reconnect after %d consecutive failures. reason=%s",
|
|
self._consecutive_reconnect_failures,
|
|
reason,
|
|
)
|
|
self._engine_confirmed_dead = True
|
|
|
|
verbose_proxy_logger.warning("Attempting Prisma DB reconnect. reason=%s", reason)
|
|
|
|
reconnect_succeeded = False
|
|
try:
|
|
await self._run_reconnect_cycle(timeout_seconds=timeout_seconds, force_recreate=force_recreate)
|
|
reconnect_succeeded = True
|
|
self._consecutive_reconnect_failures = 0
|
|
verbose_proxy_logger.info("Prisma DB reconnect succeeded. reason=%s", reason)
|
|
except _ForcedRecreateDeclined as declined:
|
|
# A decline is raised only when the recreate returns False, which
|
|
# happens only at the generation guard, and the generation moves
|
|
# only after a replacement has connected. So a decline is proof
|
|
# that a replacement SUCCEEDED, and zeroing a consecutive-failure
|
|
# count on that proof is right by definition rather than by
|
|
# analogy to what a reported success used to do. Note what it
|
|
# proves is that the WRITER was replaced, not that this caller's
|
|
# engine was repaired: on a read replica the reader can still be
|
|
# poisoned, since the wrapper returns before touching it. Leaving
|
|
# the count at the threshold would let the escalation check above
|
|
# re-arm the dead-engine flag on the very next attempt and send a
|
|
# healthy replacement back down the probe-free heavy path.
|
|
self._consecutive_reconnect_failures = 0
|
|
verbose_proxy_logger.warning("Prisma DB reconnect declined. reason=%s detail=%s", reason, declined)
|
|
except Exception as reconnect_err:
|
|
self._consecutive_reconnect_failures += 1
|
|
# Remember WHICH engine could not be repaired, so the rest of this
|
|
# caller's burst collapses onto the cooldown instead of each
|
|
# retrying the recreate that just failed. Recorded only for a
|
|
# caller that named a generation: a watchdog or transport-error
|
|
# reconnect failing here is unrelated to any stale read engine and
|
|
# must not suppress its waiver.
|
|
if stale_read_engine is not None:
|
|
# Key off the wrapper the CALLER named, never a freshly resolved
|
|
# `read_db`. A failed reader recreate is itself what marks the
|
|
# reader unavailable, so re-resolving here would file the
|
|
# reader's failure under the writer: the poisoned reader would
|
|
# lose its record and the healthy writer would gain a spurious
|
|
# one, wrong in both directions at once.
|
|
self._failed_recreate_generations = MappingProxyType(
|
|
{**self._failed_recreate_generations, stale_read_engine.wrapper: stale_read_engine.generation}
|
|
)
|
|
verbose_proxy_logger.error(
|
|
"Prisma DB reconnect failed (%d consecutive). reason=%s error=%s",
|
|
self._consecutive_reconnect_failures,
|
|
reason,
|
|
reconnect_err,
|
|
)
|
|
finally:
|
|
self._db_last_reconnect_attempt_ts = time.time()
|
|
|
|
return reconnect_succeeded
|
|
|
|
async def attempt_db_reconnect(
|
|
self,
|
|
reason: str,
|
|
force: bool = False,
|
|
timeout_seconds: float | None = None,
|
|
lock_timeout_seconds: float | None = None,
|
|
force_recreate: bool = False,
|
|
stale_read_engine: "_StaleReadEngine | None" = None,
|
|
) -> bool:
|
|
"""
|
|
Attempt to reconnect the Prisma client in a singleflight manner.
|
|
|
|
`force` bypasses the cooldown unconditionally; `force_recreate`
|
|
bypasses the liveness probe that would otherwise skip recreating a
|
|
reachable engine; `stale_read_engine` bypasses the cooldown only while
|
|
the engine that produced the caller's failure is still the live one
|
|
(see `_cooldown_applies`).
|
|
|
|
A `force_recreate` caller can also get False for a third reason: the
|
|
generation guard declined because another path had already replaced
|
|
the engine, which is a successful outcome reported as False. Callers
|
|
that branch on the return value (`exception_handler` raises on False,
|
|
`auth_checks` retries only on True) would misread that as a dead end,
|
|
and are safe today only because neither passes `force_recreate`. Do
|
|
not add it to one of them without revisiting how it reads the result.
|
|
|
|
Returns:
|
|
bool: True if reconnection succeeded, else False.
|
|
"""
|
|
now: Final = time.time()
|
|
if (
|
|
force is False
|
|
and self._cooldown_applies(stale_read_engine)
|
|
and now - self._db_last_reconnect_attempt_ts < self._db_reconnect_cooldown_seconds
|
|
):
|
|
verbose_proxy_logger.debug(
|
|
"Skipping DB reconnect attempt due to cooldown. reason=%s",
|
|
reason,
|
|
)
|
|
return False
|
|
|
|
if lock_timeout_seconds is None:
|
|
async with self._db_reconnect_lock:
|
|
return await self._attempt_reconnect_inside_lock(
|
|
force, reason, timeout_seconds, force_recreate, stale_read_engine
|
|
)
|
|
|
|
lock_acquired_by_timeout_task = False
|
|
|
|
async def _acquire_reconnect_lock() -> bool:
|
|
nonlocal lock_acquired_by_timeout_task
|
|
await self._db_reconnect_lock.acquire()
|
|
lock_acquired_by_timeout_task = True
|
|
return True
|
|
|
|
acquire_task: Final = asyncio.create_task(_acquire_reconnect_lock())
|
|
|
|
async def _abandon_acquire_task() -> None:
|
|
acquire_task.cancel()
|
|
try:
|
|
await acquire_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
# Defensive cleanup for timeout/cancel race on Python 3.9-3.11.
|
|
if lock_acquired_by_timeout_task:
|
|
try:
|
|
self._db_reconnect_lock.release()
|
|
except RuntimeError:
|
|
pass
|
|
|
|
try:
|
|
done, _pending = await asyncio.wait(
|
|
{acquire_task},
|
|
timeout=lock_timeout_seconds,
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
except asyncio.CancelledError:
|
|
await asyncio.shield(_abandon_acquire_task())
|
|
raise
|
|
if acquire_task not in done:
|
|
await _abandon_acquire_task()
|
|
verbose_proxy_logger.debug(
|
|
"Skipping DB reconnect attempt due to lock acquisition timeout. reason=%s timeout=%ss",
|
|
reason,
|
|
lock_timeout_seconds,
|
|
)
|
|
return False
|
|
|
|
try:
|
|
acquire_task.result()
|
|
except Exception as lock_acquire_err:
|
|
verbose_proxy_logger.debug(
|
|
"Skipping DB reconnect attempt due to lock acquisition error. reason=%s error=%s",
|
|
reason,
|
|
lock_acquire_err,
|
|
)
|
|
return False
|
|
|
|
try:
|
|
return await self._attempt_reconnect_inside_lock(
|
|
force, reason, timeout_seconds, force_recreate, stale_read_engine
|
|
)
|
|
finally:
|
|
self._db_reconnect_lock.release()
|
|
|
|
async def start_db_health_watchdog_task(self) -> None:
|
|
"""Start background tasks that monitor DB health:
|
|
- A periodic SELECT 1 probe that triggers reconnect on network/connection failure.
|
|
- A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.
|
|
"""
|
|
if self._db_health_watchdog_enabled is not True:
|
|
verbose_proxy_logger.debug("Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED")
|
|
return
|
|
if self._db_health_watchdog_task is not None:
|
|
return
|
|
# Let planned writer-engine restarts (IAM token refresh, guarded
|
|
# reconnect) re-arm the watcher on the new PID instead of being
|
|
# mistaken for a crash (issue #29176). Set on the writer wrapper since
|
|
# the watcher tracks the writer engine.
|
|
self.writer_db.on_engine_replaced = self._handle_writer_engine_replaced
|
|
self._db_health_watchdog_task = asyncio.create_task(self._db_health_watchdog_loop())
|
|
verbose_proxy_logger.info(
|
|
"Started Prisma DB health watchdog (interval=%ss, reconnect_cooldown=%ss, probe_timeout=%ss, reconnect_timeout=%ss)",
|
|
self._db_health_watchdog_interval_seconds,
|
|
self._db_reconnect_cooldown_seconds,
|
|
self._db_health_watchdog_probe_timeout_seconds,
|
|
self._db_watchdog_reconnect_timeout_seconds,
|
|
)
|
|
await self._start_engine_watcher()
|
|
|
|
async def stop_db_health_watchdog_task(self) -> None:
|
|
"""Stop DB health watchdog task and engine watcher gracefully."""
|
|
self._stop_engine_watcher()
|
|
if self._db_health_watchdog_task is None:
|
|
return
|
|
self._db_health_watchdog_task.cancel()
|
|
try:
|
|
await self._db_health_watchdog_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._db_health_watchdog_task = None
|
|
verbose_proxy_logger.info("Stopped Prisma DB health watchdog")
|
|
|
|
async def _db_health_watchdog_loop(self) -> None:
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(self._db_health_watchdog_interval_seconds)
|
|
await asyncio.wait_for(
|
|
self.db.query_raw("SELECT 1"),
|
|
timeout=self._db_health_watchdog_probe_timeout_seconds,
|
|
)
|
|
if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable:
|
|
await self.attempt_db_reconnect(
|
|
reason="db_health_watchdog_writer_unavailable",
|
|
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
|
|
)
|
|
continue
|
|
if await asyncio.wait_for(
|
|
self._writer_is_read_only(self.writer_db),
|
|
timeout=self._db_health_watchdog_probe_timeout_seconds,
|
|
):
|
|
await self.recreate_read_only_writer(
|
|
reason="db_health_watchdog_writer_read_only",
|
|
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
|
|
)
|
|
continue
|
|
self._db_read_only_recreate_streak = 0
|
|
self._db_read_only_recreate_ts = 0.0
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
if isinstance(e, asyncio.TimeoutError) or PrismaDBExceptionHandler.is_database_infrastructure_error(e):
|
|
await self.attempt_db_reconnect(
|
|
reason="db_health_watchdog_connection_error",
|
|
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
|
|
)
|
|
else:
|
|
verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e)
|
|
|
|
async def recreate_read_only_writer(self, reason: str, timeout_seconds: float | None = None) -> bool:
|
|
"""Force-recreate the client behind a writer session that rejects writes
|
|
(SQLSTATE 25006). Each recreate doubles the wait before the next one
|
|
until the watchdog sees a writable session again, so a database that is
|
|
read-only as a whole (replica, failover in progress) does not get its
|
|
engine killed on every watchdog cycle or failed write."""
|
|
backoff_seconds: Final = min(
|
|
self._db_reconnect_cooldown_seconds * 2 ** min(self._db_read_only_recreate_streak, 10),
|
|
_READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS,
|
|
)
|
|
if time.time() - self._db_read_only_recreate_ts < backoff_seconds:
|
|
verbose_proxy_logger.debug(
|
|
"Writer session still read-only after %s recreate(s); backing off %ss. reason=%s",
|
|
self._db_read_only_recreate_streak,
|
|
backoff_seconds,
|
|
reason,
|
|
)
|
|
return False
|
|
verbose_proxy_logger.warning(
|
|
"Writer session is read-only (writes fail with SQLSTATE 25006); recreating Prisma client. reason=%s",
|
|
reason,
|
|
)
|
|
self._db_read_only_recreate_ts = time.time()
|
|
self._db_read_only_recreate_streak += 1
|
|
return await self.attempt_db_reconnect(reason=reason, timeout_seconds=timeout_seconds, force_recreate=True)
|
|
|
|
async def _writer_is_read_only(self, writer: PrismaWrapper) -> bool:
|
|
"""True iff the pooled writer session answers reads but rejects writes (SQLSTATE 25006)."""
|
|
rows: Final = _WRITER_WRITABILITY_PROBE_ROWS.validate_python(
|
|
await writer.query_raw(_WRITER_WRITABILITY_PROBE_SQL)
|
|
)
|
|
return any(row.get("transaction_read_only") == "on" for row in rows)
|
|
|
|
def _probe_target_wrapper(self) -> PrismaWrapper:
|
|
"""The Prisma wrapper a `SELECT 1` health probe actually reaches.
|
|
|
|
`health_check()` issues `query_raw`, which `RoutingPrismaWrapper` sends
|
|
to the reader unless the reader is degraded. The writer's engine state
|
|
therefore says nothing about a probe that failed against the reader, so
|
|
the gate has to follow the same routing rule the probe did.
|
|
"""
|
|
if isinstance(self.db, RoutingPrismaWrapper):
|
|
return self.db.writer if self.db.reader_unavailable else self.db.reader
|
|
return self.db
|
|
|
|
async def _run_health_probe(self, wrapper: PrismaWrapper) -> object:
|
|
"""Issue the `SELECT 1` a health check is made of, against `wrapper`.
|
|
|
|
Takes the wrapper rather than re-reading `self.db`, because routing is
|
|
re-resolved on every attribute access: a reader that recovers between
|
|
the caller picking its target and the query going out would send the
|
|
probe to a different engine than the one whose generation the caller is
|
|
about to check, and attribute the failure to the wrong replacement.
|
|
"""
|
|
sql_query: Final = "SELECT 1"
|
|
response: Final[object] = await wrapper.query_raw(sql_query)
|
|
return response
|
|
|
|
async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool:
|
|
try:
|
|
await self._run_health_probe(wrapper)
|
|
except Exception as probe_error: # noqa: BLE001 # any failure means the database is not answering
|
|
verbose_proxy_logger.debug("Prisma health_check() confirmation probe failed: %s", probe_error)
|
|
return False
|
|
return True
|
|
|
|
async def _planned_engine_replacement_absorbed(
|
|
self,
|
|
e: Exception,
|
|
wrapper: PrismaWrapper,
|
|
generation_before: int,
|
|
) -> bool:
|
|
"""True iff `e` is a connection-class probe failure that a completed
|
|
planned query-engine replacement explains.
|
|
|
|
Planned replacements (RDS IAM token refresh, guarded reconnect) kill the
|
|
running query engine and spawn a new one. A `SELECT 1` probe that races
|
|
that sub-second window fails with a transport error against the engine's
|
|
local HTTP port even though nothing is wrong with the database, and
|
|
reporting it drives a false-positive `db_exceptions` alert on every
|
|
replacement.
|
|
|
|
Two things must both hold, because neither is sufficient alone. The
|
|
engine generation must have moved, which says a replacement completed
|
|
rather than merely being attempted: reconnect attempts during a real
|
|
outage hold the same lock for tens of seconds, so gating on an in-flight
|
|
replacement would swallow most of an outage's alerts. And a fresh probe
|
|
must succeed, because `Prisma.connect()` polls the query engine's own
|
|
`/status` endpoint rather than round-tripping to the database, so a
|
|
future engine that binds before it validates its connection pool would
|
|
let the generation advance with the database still unreachable.
|
|
|
|
Waiting for an in-flight replacement to settle is what makes the
|
|
generation check meaningful, since the generation has not moved yet at
|
|
the instant the probe fails. The wait is generous against a replacement
|
|
that takes well under a second and short enough that an outage-hung
|
|
reconnect is not waited out; a replacement that has not settled by then
|
|
reports rather than stays silent.
|
|
"""
|
|
if not PrismaDBExceptionHandler.is_database_connection_error(e):
|
|
return False
|
|
await wrapper.wait_for_planned_engine_replacement(self.PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS)
|
|
if wrapper.engine_generation == generation_before:
|
|
return False
|
|
return await self._probe_answers_now(wrapper)
|
|
|
|
async def _report_health_check_failure(
|
|
self,
|
|
e: Exception,
|
|
duration: float,
|
|
traceback_str: str,
|
|
wrapper: PrismaWrapper,
|
|
generation_before: int,
|
|
) -> None:
|
|
if await self._planned_engine_replacement_absorbed(e, wrapper, generation_before):
|
|
verbose_proxy_logger.info(
|
|
"Prisma health_check() connection error raced a planned query-engine replacement; "
|
|
"not reporting it as a DB exception: %s",
|
|
e,
|
|
)
|
|
return
|
|
await self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=duration,
|
|
call_type="health_check",
|
|
traceback_str=traceback_str,
|
|
)
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception,
|
|
max_tries=3,
|
|
max_time=10,
|
|
on_backoff=on_backoff,
|
|
)
|
|
async def health_check(self):
|
|
"""
|
|
Health check endpoint for the prisma client
|
|
"""
|
|
start_time: Final = time.time()
|
|
probe_wrapper: Final = self._probe_target_wrapper()
|
|
generation_before: Final = probe_wrapper.engine_generation
|
|
try:
|
|
return await self._run_health_probe(probe_wrapper)
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg: Final = f"LiteLLM Prisma Client Exception health_check(): {e}"
|
|
verbose_proxy_logger.warning(error_msg)
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
self._report_health_check_failure(
|
|
e=e,
|
|
duration=_duration,
|
|
traceback_str=error_traceback,
|
|
wrapper=probe_wrapper,
|
|
generation_before=generation_before,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
async def _get_spend_logs_row_count(self) -> int:
|
|
"""
|
|
Get the row count from LiteLLM_SpendLogs table using PostgreSQL system statistics.
|
|
"""
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception,
|
|
max_tries=3,
|
|
max_time=10,
|
|
on_backoff=on_backoff,
|
|
)
|
|
async def _fetch_row_count() -> int:
|
|
sql_query: Final = """
|
|
SELECT reltuples::BIGINT
|
|
FROM pg_class
|
|
WHERE oid = '"LiteLLM_SpendLogs"'::regclass;
|
|
"""
|
|
result: Final[Sequence[_RelTuplesRow]] = await self.db.query_raw(query=sql_query)
|
|
return result[0]["reltuples"]
|
|
|
|
try:
|
|
return await _fetch_row_count()
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Error getting LiteLLM_SpendLogs row count: %s", e)
|
|
return 0
|
|
|
|
@backoff.on_exception(
|
|
backoff.expo,
|
|
Exception,
|
|
max_tries=3,
|
|
max_time=10,
|
|
on_backoff=on_backoff,
|
|
)
|
|
async def _set_spend_logs_row_count_in_proxy_state(self) -> None:
|
|
"""
|
|
Set the `LiteLLM_SpendLogs`row count in proxy state.
|
|
|
|
This is used later to determine if we should run expensive UI Usage queries.
|
|
"""
|
|
from litellm.proxy.proxy_server import proxy_state
|
|
|
|
_num_spend_logs_rows: Final = await self._get_spend_logs_row_count()
|
|
proxy_state.set_proxy_state_variable(
|
|
variable_name="spend_logs_row_count",
|
|
value=_num_spend_logs_rows,
|
|
)
|
|
|
|
# Health Check Database Methods
|
|
def _validate_response_time(self, response_time_ms: float | None) -> float | None:
|
|
"""Validate and clean response time value"""
|
|
if response_time_ms is None:
|
|
return None
|
|
try:
|
|
value: Final = float(response_time_ms)
|
|
return value if math.isfinite(value) else None
|
|
except (ValueError, TypeError):
|
|
verbose_proxy_logger.warning("Invalid response_time_ms value: %s", response_time_ms)
|
|
return None
|
|
|
|
def _clean_details(self, details: dict | None) -> dict | None:
|
|
"""Clean and validate details JSON"""
|
|
if not isinstance(details, dict):
|
|
return None
|
|
try:
|
|
return safe_json_loads(safe_dumps(details))
|
|
except Exception as e:
|
|
verbose_proxy_logger.warning("Failed to clean details JSON: %s", e)
|
|
return None
|
|
|
|
async def save_health_check_result(
|
|
self,
|
|
model_name: str,
|
|
status: str,
|
|
healthy_count: int = 0,
|
|
unhealthy_count: int = 0,
|
|
error_message: str | None = None,
|
|
response_time_ms: float | None = None,
|
|
details: dict | None = None,
|
|
checked_by: str | None = None,
|
|
model_id: str | None = None,
|
|
):
|
|
"""Save health check result to database"""
|
|
try:
|
|
# Build base data with required fields
|
|
health_check_data: Final = {
|
|
"model_name": str(model_name),
|
|
"status": str(status),
|
|
"healthy_count": int(healthy_count),
|
|
"unhealthy_count": int(unhealthy_count),
|
|
}
|
|
|
|
# Add optional fields using dict comprehension and helper methods
|
|
optional_fields: Final = {
|
|
"error_message": str(error_message)[:500] if error_message else None,
|
|
"response_time_ms": self._validate_response_time(response_time_ms),
|
|
"details": self._clean_details(details),
|
|
"checked_by": str(checked_by) if checked_by else None,
|
|
"model_id": str(model_id) if model_id else None,
|
|
}
|
|
|
|
# Add only non-None optional fields
|
|
health_check_data.update({k: v for k, v in optional_fields.items() if v is not None})
|
|
|
|
verbose_proxy_logger.debug("Saving health check data: %s", health_check_data)
|
|
return await HealthCheckRepository(self).table.create(data=health_check_data)
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Error saving health check result for model %s: %s", model_name, e)
|
|
return None
|
|
|
|
async def get_health_check_history(
|
|
self,
|
|
model_name: str | None = None,
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
status_filter: str | None = None,
|
|
) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
|
|
"""
|
|
Get health check history with optional filtering
|
|
"""
|
|
try:
|
|
where_clause: Final[dict[str, str]] = {}
|
|
if model_name:
|
|
where_clause["model_name"] = model_name
|
|
if status_filter:
|
|
where_clause["status"] = status_filter
|
|
|
|
results: Final = await HealthCheckRepository(self).table.find_many(
|
|
where=where_clause,
|
|
order={"checked_at": "desc"},
|
|
take=limit,
|
|
skip=offset,
|
|
)
|
|
return results
|
|
except Exception as e:
|
|
verbose_proxy_logger.error("Error getting health check history: %s", e)
|
|
return []
|
|
|
|
async def get_all_latest_health_checks(self) -> tuple[LatestHealthCheckRow, ...]:
|
|
"""Latest health check per (model_id, model_name), deduplicated in Postgres."""
|
|
return await fetch_latest_health_checks(self)
|
|
|
|
async def get_latest_health_checks_for_models(self, model_names: Sequence[str]) -> tuple[LatestHealthCheckRow, ...]:
|
|
"""Same as ``get_all_latest_health_checks``, bounded to the named models."""
|
|
return await fetch_latest_health_checks_for_models(self, model_names)
|
|
|
|
|
|
### HELPER FUNCTIONS ###
|
|
|
|
|
|
async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient):
|
|
"""
|
|
Check if a user_id exists in cache,
|
|
if not retrieve it.
|
|
"""
|
|
cache_key: Final = f"{user_id}_user_api_key_user_id"
|
|
response: Final = cache.get_cache(key=cache_key)
|
|
if response is None: # Cache miss
|
|
user_row: Final = await db.get_data(user_id=user_id)
|
|
if user_row is not None:
|
|
print_verbose(f"User Row: {user_row}, type = {type(user_row)}")
|
|
if hasattr(user_row, "model_dump_json") and callable(getattr(user_row, "model_dump_json")):
|
|
cache_value: Final[str] = user_row.model_dump_json()
|
|
cache.set_cache(key=cache_key, value=cache_value, ttl=600) # store for 10 minutes
|
|
|
|
|
|
def _should_use_smtp_ssl(smtp_port: int) -> bool:
|
|
"""
|
|
Port 465 expects an immediate TLS handshake (implicit SSL), so a plain
|
|
smtplib.SMTP connection hangs waiting for an SMTP banner. Use SMTP_SSL
|
|
there, or when SMTP_USE_SSL is explicitly enabled.
|
|
"""
|
|
return os.getenv("SMTP_USE_SSL", "False") == "True" or smtp_port == 465
|
|
|
|
|
|
def _create_smtp_connection(smtp_host: str, smtp_port: int, timeout: float) -> smtplib.SMTP:
|
|
if _should_use_smtp_ssl(smtp_port=smtp_port):
|
|
return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context(), timeout=timeout)
|
|
return smtplib.SMTP(host=smtp_host, port=smtp_port, timeout=timeout)
|
|
|
|
|
|
def _send_smtp_message(
|
|
email_message: MIMEMultipart,
|
|
smtp_host: str,
|
|
smtp_port: int,
|
|
smtp_username: str | None,
|
|
smtp_password: str | None,
|
|
sender_email: str,
|
|
receiver_email: str,
|
|
timeout: float,
|
|
) -> None:
|
|
using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port)
|
|
with _create_smtp_connection(
|
|
smtp_host=smtp_host,
|
|
smtp_port=smtp_port,
|
|
timeout=timeout,
|
|
) as server:
|
|
if not using_ssl and os.getenv("SMTP_TLS", "True") != "False":
|
|
server.starttls(context=ssl.create_default_context())
|
|
|
|
if smtp_username and smtp_password:
|
|
server.login(
|
|
user=smtp_username,
|
|
password=smtp_password,
|
|
)
|
|
|
|
server.send_message(
|
|
msg=email_message,
|
|
from_addr=sender_email,
|
|
to_addrs=receiver_email,
|
|
)
|
|
|
|
|
|
async def send_email(
|
|
receiver_email: str | None = None,
|
|
subject: str | None = None,
|
|
html: str | None = None,
|
|
):
|
|
"""
|
|
smtp_host,
|
|
smtp_port,
|
|
smtp_username,
|
|
smtp_password,
|
|
sender_name,
|
|
sender_email,
|
|
"""
|
|
## SERVER SETUP ##
|
|
|
|
smtp_host: Final = os.getenv("SMTP_HOST")
|
|
smtp_port: Final = int(os.getenv("SMTP_PORT", "587")) # default to port 587
|
|
smtp_username: Final = os.getenv("SMTP_USERNAME")
|
|
smtp_password: Final = os.getenv("SMTP_PASSWORD")
|
|
sender_email: Final = os.getenv("SMTP_SENDER_EMAIL", None)
|
|
if sender_email is None:
|
|
raise ValueError("Trying to use SMTP, but SMTP_SENDER_EMAIL is not set")
|
|
if receiver_email is None:
|
|
raise ValueError(f"No receiver email provided for SMTP email. {receiver_email}")
|
|
if subject is None:
|
|
raise ValueError(f"No subject provided for SMTP email. {subject}")
|
|
if html is None:
|
|
raise ValueError(f"No HTML body provided for SMTP email. {html}")
|
|
|
|
## EMAIL SETUP ##
|
|
email_message: Final = MIMEMultipart()
|
|
email_message["From"] = sender_email
|
|
email_message["To"] = receiver_email
|
|
email_message["Subject"] = subject
|
|
verbose_proxy_logger.debug("sending email from %s to %s", sender_email, receiver_email)
|
|
|
|
if smtp_host is None:
|
|
raise ValueError("Trying to use SMTP, but SMTP_HOST is not set")
|
|
|
|
# Attach the body to the email
|
|
email_message.attach(MIMEText(html, "html"))
|
|
|
|
try:
|
|
smtp_timeout: Final = float(os.getenv("SMTP_TIMEOUT", "30"))
|
|
await asyncio.to_thread(
|
|
_send_smtp_message,
|
|
email_message=email_message,
|
|
smtp_host=smtp_host,
|
|
smtp_port=smtp_port,
|
|
smtp_username=smtp_username,
|
|
smtp_password=smtp_password,
|
|
sender_email=sender_email,
|
|
receiver_email=receiver_email,
|
|
timeout=smtp_timeout,
|
|
)
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e))
|
|
|
|
|
|
def hash_token(token: str):
|
|
import hashlib
|
|
|
|
# Hash the string using SHA-256
|
|
hashed_token: Final = hashlib.sha256(token.encode()).hexdigest()
|
|
|
|
return hashed_token
|
|
|
|
|
|
def hash_password(password: str) -> str:
|
|
"""Hash a password using scrypt with a random salt."""
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
|
|
salt: Final = os.urandom(16)
|
|
dk: Final = hashlib.scrypt(password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32)
|
|
return "scrypt:" + base64.b64encode(salt + dk).decode()
|
|
|
|
|
|
def verify_password(password: str, stored: str) -> bool:
|
|
"""Verify a password against a stored hash. Supports scrypt and SHA256."""
|
|
import base64
|
|
import hashlib
|
|
import secrets
|
|
|
|
if stored.startswith("scrypt:"):
|
|
try:
|
|
raw: Final = base64.b64decode(stored[7:])
|
|
salt, dk = raw[:16], raw[16:]
|
|
dk2: Final = hashlib.scrypt(password.encode(), salt=salt, n=16384, r=8, p=1, dklen=32)
|
|
return secrets.compare_digest(dk, dk2)
|
|
except Exception:
|
|
return False
|
|
# SHA256 fallback (not vulnerable to pass-the-hash: checks sha256(input) == stored)
|
|
if len(stored) == 64 and all(c in "0123456789abcdef" for c in stored):
|
|
return secrets.compare_digest(hashlib.sha256(password.encode()).hexdigest().encode(), stored.encode())
|
|
return False
|
|
|
|
|
|
async def migrate_passwords_to_scrypt_async(prisma_client) -> str:
|
|
"""
|
|
Migrate plaintext passwords in the DB to scrypt. SHA256 passwords
|
|
are left alone (they migrate on next login via the SHA256 fallback).
|
|
Skips quickly if no plaintext passwords exist.
|
|
"""
|
|
all_with_pw: Final = await UserRepository(prisma_client).table.find_many(
|
|
where={"password": {"not": None}},
|
|
)
|
|
|
|
def _is_sha256_hex(s: str) -> bool:
|
|
return len(s) == 64 and all(c in "0123456789abcdef" for c in s)
|
|
|
|
plaintext_users: Final = [
|
|
(u.user_id, u.password)
|
|
for u in all_with_pw
|
|
if u.password and not u.password.startswith("scrypt:") and not _is_sha256_hex(u.password)
|
|
]
|
|
if not plaintext_users:
|
|
return "No plaintext passwords found"
|
|
|
|
for user_id, plaintext_password in plaintext_users:
|
|
await UserRepository(prisma_client).table.update(
|
|
where={"user_id": user_id},
|
|
data={"password": hash_password(plaintext_password)},
|
|
)
|
|
return f"Migrated {len(plaintext_users)} plaintext passwords to scrypt"
|
|
|
|
|
|
def _hash_token_if_needed(token: str) -> str:
|
|
"""
|
|
Hash the token if it's a string and starts with "sk-"
|
|
|
|
Else return the token as is
|
|
"""
|
|
if token.startswith("sk-"):
|
|
return hash_token(token=token)
|
|
else:
|
|
return token
|
|
|
|
|
|
async def enqueue_spend_logs(
|
|
prisma_client: PrismaClient,
|
|
logs: Sequence[Mapping[str, object]],
|
|
*,
|
|
at_head: bool = False,
|
|
max_bytes: int = SPEND_LOG_QUEUE_MAX_BYTES,
|
|
) -> None:
|
|
"""Queue spend logs for the next flush, held under ``SPEND_LOG_QUEUE_MAX_BYTES``.
|
|
|
|
``at_head`` replays a batch the DB refused, so it flushes before the logs
|
|
that piled up during the outage. Past the budget the oldest logs are
|
|
dropped, which keeps a long outage from growing the queue until the pod
|
|
dies.
|
|
"""
|
|
added: Final = sum(spend_log_row_bytes(row) for row in logs)
|
|
async with prisma_client._spend_log_transactions_lock:
|
|
queued: Final = (
|
|
tuple(logs) + tuple(prisma_client.spend_log_transactions)
|
|
if at_head
|
|
else tuple(prisma_client.spend_log_transactions) + tuple(logs)
|
|
)
|
|
kept, kept_bytes = spend_log_queue_within_budget(queued, PrismaClient.spend_log_queue_bytes + added, max_bytes)
|
|
prisma_client.spend_log_transactions[:] = kept
|
|
PrismaClient.spend_log_queue_bytes = kept_bytes
|
|
if len(kept) < len(queued):
|
|
verbose_proxy_logger.error(
|
|
"Spend tracking - spend log queue is at its %d byte budget; dropped the %d oldest spend logs",
|
|
max_bytes,
|
|
len(queued) - len(kept),
|
|
)
|
|
|
|
|
|
def request_spend_log_flush(prisma_client: PrismaClient) -> None:
|
|
"""Wake this client's queue monitor now rather than leaving the rows for its next poll.
|
|
|
|
The Responses API hands the client an id it can chain from straight away, and that
|
|
lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval.
|
|
Repeated requests coalesce into the monitor's next pass, so the batching holds.
|
|
A request made before the monitor is running is dropped, and loses nothing: the
|
|
monitor reads the queue on its first pass, before it ever waits on a request.
|
|
"""
|
|
flush_requested: Final = prisma_client.spend_log_flush_requested
|
|
if flush_requested is not None:
|
|
flush_requested.set()
|
|
|
|
|
|
async def _wait_for_spend_log_flush_request(flush_requested: asyncio.Event, interval: float) -> bool:
|
|
"""Wait out ``interval``, returning early and True when a flush was requested."""
|
|
try:
|
|
await asyncio.wait_for(flush_requested.wait(), timeout=interval)
|
|
except asyncio.TimeoutError:
|
|
return False
|
|
flush_requested.clear()
|
|
return True
|
|
|
|
|
|
async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]:
|
|
"""Take up to ``limit`` of the oldest queued spend logs off the queue.
|
|
|
|
Every enqueue and dequeue goes through this pair so the byte total the
|
|
queue is bounded by stays in step with what the queue actually holds.
|
|
"""
|
|
async with prisma_client._spend_log_transactions_lock:
|
|
popped: Final = prisma_client.spend_log_transactions[:limit]
|
|
prisma_client.spend_log_transactions[:] = prisma_client.spend_log_transactions[limit:]
|
|
PrismaClient.spend_log_queue_bytes = max(
|
|
0, PrismaClient.spend_log_queue_bytes - sum(spend_log_row_bytes(row) for row in popped)
|
|
)
|
|
return popped
|
|
|
|
|
|
class ProxyUpdateSpend:
|
|
@staticmethod
|
|
async def update_end_user_spend(
|
|
n_retry_times: int,
|
|
prisma_client: PrismaClient,
|
|
proxy_logging_obj: ProxyLogging,
|
|
end_user_list_transactions: dict[str, float],
|
|
):
|
|
for i in range(n_retry_times + 1):
|
|
start_time = time.time()
|
|
try:
|
|
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
|
|
batcher: _EndUserSpendBatch
|
|
async with transaction.batch_() as batcher:
|
|
# Sort by end_user_id for consistent lock ordering across pods to prevent deadlocks.
|
|
for end_user_id, response_cost in sorted(end_user_list_transactions.items()):
|
|
if litellm.max_end_user_budget is not None:
|
|
pass
|
|
batcher.litellm_endusertable.upsert(
|
|
where={"user_id": end_user_id},
|
|
data={
|
|
"create": {
|
|
"user_id": end_user_id,
|
|
"spend": response_cost,
|
|
"blocked": False,
|
|
},
|
|
"update": {"spend": {"increment": response_cost}},
|
|
},
|
|
)
|
|
|
|
break
|
|
except Exception as e:
|
|
await DBSpendUpdateWriter._handle_spend_update_failure(
|
|
e=e,
|
|
attempt=i,
|
|
n_retry_times=n_retry_times,
|
|
start_time=start_time,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
@staticmethod
|
|
async def update_spend_logs(
|
|
n_retry_times: int,
|
|
prisma_client: PrismaClient,
|
|
db_writer_client: AsyncHTTPHandler | None,
|
|
proxy_logging_obj: ProxyLogging,
|
|
logs_to_process: list[dict[str, object]] | None = None,
|
|
):
|
|
BATCH_SIZE: Final = 1000 # Preferred size of each batch to write to the database
|
|
MAX_LOGS_PER_INTERVAL: Final = 10000 # Maximum number of logs to flush in a single interval
|
|
popped_batch = False
|
|
if logs_to_process is None:
|
|
logs_to_process = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL)
|
|
popped_batch = True
|
|
if len(logs_to_process) > 0:
|
|
verbose_proxy_logger.info(
|
|
"Spend tracking - processing %d spend logs for DB write",
|
|
len(logs_to_process),
|
|
)
|
|
start_time: Final = time.time()
|
|
try:
|
|
for i in range(n_retry_times + 1):
|
|
try:
|
|
base_url = os.getenv("SPEND_LOGS_URL", None)
|
|
if len(logs_to_process) > 0 and base_url is not None and db_writer_client is not None:
|
|
if not base_url.endswith("/"):
|
|
base_url += "/"
|
|
verbose_proxy_logger.debug("base_url: %s", base_url)
|
|
json_data = json.dumps(logs_to_process)
|
|
response = await db_writer_client.post(
|
|
url=base_url + "spend/update",
|
|
data=json_data,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
del json_data
|
|
if response.status_code == 200:
|
|
# Items already removed from queue at start of function
|
|
pass
|
|
else:
|
|
for j in range(0, len(logs_to_process), BATCH_SIZE):
|
|
batch = logs_to_process[j : j + BATCH_SIZE]
|
|
batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch]
|
|
isolation_budget = MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH
|
|
for statement_rows in spend_log_write_batches(
|
|
batch_with_dates,
|
|
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
|
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
|
|
):
|
|
isolation_budget = await _create_spend_logs_with_poison_isolation(
|
|
SpendLogsRepository(prisma_client),
|
|
statement_rows,
|
|
isolation_budget,
|
|
)
|
|
verbose_proxy_logger.debug("Flushed %s logs to the DB.", len(batch))
|
|
# Explicitly clear batch memory
|
|
del batch, batch_with_dates
|
|
|
|
# Items already removed from queue at start of function
|
|
async with prisma_client._spend_log_transactions_lock:
|
|
remaining_count = len(prisma_client.spend_log_transactions)
|
|
verbose_proxy_logger.debug(
|
|
"%s logs processed. Remaining in queue: %s", len(logs_to_process), remaining_count
|
|
)
|
|
break
|
|
except Exception as e:
|
|
if not _is_transient_spend_log_write_error(e):
|
|
if PrismaDBExceptionHandler.is_prisma_error(e):
|
|
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
|
verbose_proxy_logger.warning(
|
|
"Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s",
|
|
len(logs_to_process),
|
|
str(e),
|
|
)
|
|
raise
|
|
verbose_proxy_logger.warning(
|
|
"Spend tracking - transient DB error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
|
|
i + 1,
|
|
n_retry_times,
|
|
len(logs_to_process),
|
|
str(e),
|
|
)
|
|
if i >= n_retry_times:
|
|
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
|
raise
|
|
await asyncio.sleep(2**i)
|
|
except Exception as e:
|
|
_raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj)
|
|
finally:
|
|
# Clean up logs_to_process only if we popped it (caller-owned otherwise)
|
|
if popped_batch:
|
|
del logs_to_process
|
|
|
|
@staticmethod
|
|
def disable_spend_updates() -> bool:
|
|
"""
|
|
returns True if should not update spend in db
|
|
Skips writing spend logs and updates to key, team, user spend to DB
|
|
"""
|
|
from litellm.proxy.proxy_server import general_settings
|
|
|
|
if general_settings.get("disable_spend_updates") is True:
|
|
return True
|
|
return False
|
|
|
|
|
|
async def update_spend(
|
|
prisma_client: PrismaClient,
|
|
db_writer_client: AsyncHTTPHandler | None,
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
"""
|
|
Batch write updates to db.
|
|
|
|
Triggered every minute.
|
|
|
|
NOTE: This job now skips tag spend updates, which are handled by a separate
|
|
scheduler job (update_daily_tag_spend) at a longer interval to reduce contention.
|
|
|
|
Requires:
|
|
user_id_list: dict,
|
|
keys_list: list,
|
|
team_list: list,
|
|
spend_logs: list,
|
|
"""
|
|
n_retry_times: Final = 3
|
|
await proxy_logging_obj.db_spend_update_writer.db_update_spend_transaction_handler(
|
|
prisma_client=prisma_client,
|
|
n_retry_times=n_retry_times,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
### UPDATE SPEND LOGS ###
|
|
# Check queue size with lock protection
|
|
queue_size: Final = await _total_queued_spend_transactions(prisma_client)
|
|
verbose_proxy_logger.debug("Spend Logs transactions: %s", queue_size)
|
|
|
|
# Process spend log transactions when called directly.
|
|
# This keeps backwards compatibility with the old behavior.
|
|
# See update_spend_logs_job and _monitor_spend_logs_queue for the new behavior.
|
|
# Safe to keep: under high concurrency this can take up to ~30s to run,
|
|
# so it's unlikely to overlap with monitor_spend_logs_queue.
|
|
if queue_size > 0:
|
|
await update_spend_logs_job(
|
|
prisma_client=prisma_client,
|
|
db_writer_client=db_writer_client,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
|
|
async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int:
|
|
"""Pending entries across every request-time spend queue, sized under each queue's
|
|
lock. Every drain trigger reads this one owner, so a queue added later joins the
|
|
direct path, the batch job's emptiness check and the monitor at once."""
|
|
async with prisma_client._spend_log_transactions_lock:
|
|
spend_queue_size: Final = len(prisma_client.spend_log_transactions)
|
|
async with prisma_client._tool_usage_transactions_lock:
|
|
tool_queue_size: Final = len(prisma_client.tool_usage_transactions)
|
|
async with prisma_client._autorouter_turn_transactions_lock:
|
|
autorouter_queue_size: Final = len(prisma_client.autorouter_turn_transactions)
|
|
from litellm.proxy.db.shadow_eval_funnel import pending_shadow_eval_funnel_events
|
|
|
|
return spend_queue_size + tool_queue_size + autorouter_queue_size + pending_shadow_eval_funnel_events()
|
|
|
|
|
|
async def update_daily_tag_spend(
|
|
prisma_client: PrismaClient,
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
"""
|
|
Separate scheduler job to commit daily tag spend updates.
|
|
|
|
Runs at a longer interval (2.3x default) than the main update_spend job
|
|
to reduce query contention for DailyTagSpend table.
|
|
|
|
This is called by a dedicated scheduler job and does NOT process:
|
|
- Regular spend updates (user, key, team, org)
|
|
- End-user spend
|
|
- Agent spend
|
|
- Spend logs
|
|
|
|
Only processes tag spend transactions from the daily_tag_spend_update_queue.
|
|
|
|
Args:
|
|
prisma_client: PrismaClient instance
|
|
proxy_logging_obj: ProxyLogging instance for error handling
|
|
"""
|
|
n_retry_times: Final = 3
|
|
try:
|
|
if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis():
|
|
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis(
|
|
prisma_client=prisma_client,
|
|
n_retry_times=n_retry_times,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
else:
|
|
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db(
|
|
prisma_client=prisma_client,
|
|
n_retry_times=n_retry_times,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
except Exception as e:
|
|
# NOTE: keep this as a plain ``error`` (no traceback) to match the
|
|
# historical behavior of this site. ``spend_log_error`` would attach
|
|
# the active exception's traceback whenever the suppression env var
|
|
# is unset, which would be a regression for operators who never saw
|
|
# one here before.
|
|
verbose_proxy_logger.error("Error updating daily tag spend: %s", e)
|
|
|
|
|
|
async def update_spend_logs_job(
|
|
prisma_client: PrismaClient,
|
|
db_writer_client: AsyncHTTPHandler | None,
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
"""
|
|
Job to process spend_log_transactions queue.
|
|
|
|
This job is triggered based on queue size rather than time.
|
|
Pops the batch once, writes spend logs, then runs guardrail usage tracking.
|
|
"""
|
|
n_retry_times: Final = 3
|
|
MAX_LOGS_PER_INTERVAL: Final = 10000
|
|
|
|
# Atomically pop batch from queue. The tool usage queue counts toward the
|
|
# emptiness check: a spend-log write failure aborts a run before the tool
|
|
# drain below, and those entries must not strand once the spend queue drains.
|
|
if await _total_queued_spend_transactions(prisma_client) == 0:
|
|
return
|
|
|
|
logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL)
|
|
|
|
try:
|
|
await ProxyUpdateSpend.update_spend_logs(
|
|
n_retry_times=n_retry_times,
|
|
prisma_client=prisma_client,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
db_writer_client=db_writer_client,
|
|
logs_to_process=logs_to_process,
|
|
)
|
|
except asyncio.CancelledError:
|
|
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
|
verbose_proxy_logger.warning(
|
|
"Spend tracking - spend log write cancelled, requeued %d rows for the next flush",
|
|
len(logs_to_process),
|
|
)
|
|
raise
|
|
|
|
# Guardrail/policy usage tracking (same batch, outside spend-logs update)
|
|
try:
|
|
from litellm.proxy.guardrails.usage_tracking import (
|
|
process_spend_logs_guardrail_usage,
|
|
)
|
|
|
|
await process_spend_logs_guardrail_usage(
|
|
prisma_client=prisma_client,
|
|
logs_to_process=logs_to_process,
|
|
)
|
|
except Exception as guardrail_tracking_err:
|
|
verbose_proxy_logger.warning(
|
|
"Spend tracking - guardrail usage tracking failed (non-fatal): %s",
|
|
guardrail_tracking_err,
|
|
)
|
|
|
|
# Tool usage tracking: drain the request-time queue into the tool index and the
|
|
# LiteLLM_DailyToolSpend rollup. Never retried; a dropped batch is permanently
|
|
# absent from the rollup, so failures log at error.
|
|
async with prisma_client._tool_usage_transactions_lock:
|
|
tool_usage_to_process: Final = prisma_client.tool_usage_transactions[:MAX_LOGS_PER_INTERVAL]
|
|
prisma_client.tool_usage_transactions = prisma_client.tool_usage_transactions[len(tool_usage_to_process) :]
|
|
try:
|
|
from litellm.proxy.db.spend_log_tool_index import flush_tool_usage_transactions
|
|
|
|
await flush_tool_usage_transactions(
|
|
prisma_client=prisma_client,
|
|
transactions=tool_usage_to_process,
|
|
)
|
|
except Exception as tool_tracking_err:
|
|
verbose_proxy_logger.error(
|
|
"Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s",
|
|
len(tool_usage_to_process),
|
|
tool_tracking_err,
|
|
)
|
|
|
|
async with prisma_client._autorouter_turn_transactions_lock:
|
|
autorouter_turns_to_process: Final = prisma_client.autorouter_turn_transactions[:MAX_LOGS_PER_INTERVAL]
|
|
remaining_autorouter_turns: Final = prisma_client.autorouter_turn_transactions[
|
|
len(autorouter_turns_to_process) :
|
|
]
|
|
prisma_client.autorouter_turn_transactions = remaining_autorouter_turns # rebind-ok: drain under lock
|
|
try:
|
|
from litellm.proxy.db.autorouter_session_rollup import flush_autorouter_turn_transactions
|
|
|
|
await flush_autorouter_turn_transactions(
|
|
prisma_client=prisma_client,
|
|
transactions=autorouter_turns_to_process,
|
|
)
|
|
except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job
|
|
verbose_proxy_logger.error(
|
|
"Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s",
|
|
len(autorouter_turns_to_process),
|
|
autorouter_tracking_err,
|
|
)
|
|
|
|
try:
|
|
from litellm.proxy.db.shadow_eval_funnel import flush_shadow_eval_funnel
|
|
|
|
await flush_shadow_eval_funnel(prisma_client)
|
|
except Exception as funnel_err: # noqa: BLE001 # a drain bug must not abort the spend job
|
|
verbose_proxy_logger.error("Spend tracking - shadow eval funnel drain failed: %s", funnel_err)
|
|
|
|
|
|
MAX_SPEND_LOG_DRAIN_ITERATIONS: Final = 20
|
|
|
|
|
|
async def drain_spend_logs_queue(
|
|
prisma_client: PrismaClient,
|
|
db_writer_client: "AsyncHTTPHandler | None",
|
|
proxy_logging_obj: ProxyLogging,
|
|
) -> None:
|
|
monitor_task: Final = prisma_client.spend_logs_queue_monitor_task
|
|
if monitor_task is not None:
|
|
monitor_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await monitor_task
|
|
prisma_client.spend_logs_queue_monitor_task = None # rebind-ok: the client owns its monitor handle
|
|
|
|
for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS):
|
|
if await _total_queued_spend_transactions(prisma_client) == 0:
|
|
return
|
|
await update_spend_logs_job(
|
|
prisma_client=prisma_client,
|
|
db_writer_client=db_writer_client,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
remaining: Final = await _total_queued_spend_transactions(prisma_client)
|
|
if remaining > 0:
|
|
spend_log_error(
|
|
"Spend tracking - %d spend log rows still queued after %d drain passes",
|
|
remaining,
|
|
MAX_SPEND_LOG_DRAIN_ITERATIONS,
|
|
)
|
|
|
|
|
|
async def _monitor_spend_logs_queue(
|
|
prisma_client: PrismaClient,
|
|
db_writer_client: AsyncHTTPHandler | None,
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
"""
|
|
Background task that monitors the spend_log_transactions queue size
|
|
and triggers processing when the threshold is reached.
|
|
|
|
Args:
|
|
prisma_client: Prisma client instance
|
|
db_writer_client: Optional HTTP handler for external spend logs endpoint
|
|
proxy_logging_obj: Proxy logging object
|
|
"""
|
|
from litellm.constants import (
|
|
SPEND_LOG_QUEUE_POLL_INTERVAL,
|
|
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
|
|
)
|
|
|
|
threshold: Final = SPEND_LOG_QUEUE_SIZE_THRESHOLD
|
|
base_interval: Final = SPEND_LOG_QUEUE_POLL_INTERVAL
|
|
max_backoff: Final = 30.0 # Maximum backoff interval in seconds
|
|
backoff_multiplier: Final = 1.5 # Exponential backoff multiplier
|
|
current_interval = base_interval
|
|
flush_requested: Final = asyncio.Event()
|
|
prisma_client.spend_log_flush_requested = flush_requested # rebind-ok: the client owns its monitor's flush signal
|
|
|
|
verbose_proxy_logger.info(
|
|
"Starting spend logs queue monitor (threshold: %s, poll_interval: %ss)", threshold, base_interval
|
|
)
|
|
|
|
while True:
|
|
try:
|
|
# Check queue sizes with lock protection; the tool usage queue keeps
|
|
# the monitor firing when a prior failed run left it nonempty.
|
|
queue_size = await _total_queued_spend_transactions(prisma_client)
|
|
|
|
if queue_size > 0:
|
|
if queue_size >= threshold:
|
|
verbose_proxy_logger.debug(
|
|
"Spend logs queue size (%s) reached threshold (%s), triggering processing",
|
|
queue_size,
|
|
threshold,
|
|
)
|
|
# Reset to base interval when threshold is reached
|
|
current_interval = base_interval
|
|
else:
|
|
verbose_proxy_logger.debug(
|
|
"Spend logs queue size (%s) below threshold (%s), processing with backoff",
|
|
queue_size,
|
|
threshold,
|
|
)
|
|
# Exponential backoff when below threshold but still processing
|
|
current_interval = min(current_interval * backoff_multiplier, max_backoff)
|
|
|
|
await update_spend_logs_job(
|
|
prisma_client=prisma_client,
|
|
db_writer_client=db_writer_client,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
else:
|
|
# Exponential backoff when no logs to process
|
|
current_interval = min(current_interval * backoff_multiplier, max_backoff)
|
|
|
|
if await _wait_for_spend_log_flush_request(flush_requested, current_interval):
|
|
current_interval = base_interval
|
|
except Exception as e:
|
|
spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e)
|
|
# Continue monitoring even if there's an error, with exponential backoff
|
|
current_interval = min(current_interval * backoff_multiplier, max_backoff)
|
|
await asyncio.sleep(current_interval)
|
|
|
|
|
|
MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
|
|
|
|
|
|
def _is_transient_spend_log_write_error(e: Exception) -> bool:
|
|
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
|
|
|
|
|
|
async def _create_spend_logs_with_poison_isolation(
|
|
repo: SpendLogsRepository,
|
|
rows: Sequence[Mapping[str, object]],
|
|
failure_budget: int,
|
|
) -> int:
|
|
"""Write spend-log rows, isolating any row Postgres rejects on its data.
|
|
|
|
``create_many`` writes the whole batch in a single statement, so one row
|
|
carrying bytes Postgres refuses (a residual NUL byte is the canonical case)
|
|
fails the entire insert and drops every good row alongside it. On a genuine
|
|
data-layer rejection the batch is bisected so the good rows still persist
|
|
and only the offending row is dropped and logged. Transport failures,
|
|
including the "can't reach database server" outage that prisma mislabels as
|
|
a ``DataError``, are re-raised unchanged so the caller's connection-retry
|
|
path still runs.
|
|
|
|
``failure_budget`` caps the *failed* inserts the isolation may issue, which
|
|
is the work an authenticated caller flooding poisoned rows can amplify. The
|
|
one insert a statement needs when nothing is poisoned is not charged, so a
|
|
caller can thread a single budget through every statement of a flush and
|
|
bound the whole flush's failed inserts and log lines by the initial value,
|
|
without a large healthy flush ever running out and losing rows. When the
|
|
budget is spent the still-failing remainder is dropped wholesale (the
|
|
pre-existing drop-the-batch behavior) under one log line, and a statement
|
|
reached afterwards is still attempted, so clean rows behind a poison flood
|
|
persist. Returns the budget left after this subtree.
|
|
"""
|
|
try:
|
|
await repo.table.create_many(data=rows, skip_duplicates=True)
|
|
return failure_budget
|
|
except Exception as e:
|
|
if not PrismaDBExceptionHandler.is_prisma_data_error(e):
|
|
raise
|
|
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
|
raise
|
|
if PrismaDBExceptionHandler.is_deadlock_error(e):
|
|
raise
|
|
budget_left: Final = max(failure_budget - 1, 0)
|
|
if len(rows) == 1:
|
|
request_id: Final = rows[0].get("request_id")
|
|
spend_log_error(
|
|
"Spend tracking - dropping spend log row Postgres rejected. request_id=%s error=%s",
|
|
request_id,
|
|
str(e),
|
|
exc=e,
|
|
)
|
|
return budget_left
|
|
if budget_left <= 0:
|
|
spend_log_error(
|
|
"Spend tracking - dropping %d spend log rows without per-row isolation; "
|
|
"isolation failure budget exhausted for this flush",
|
|
len(rows),
|
|
)
|
|
return 0
|
|
mid: Final = len(rows) // 2
|
|
remaining: Final = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], budget_left)
|
|
if remaining <= 0:
|
|
spend_log_error(
|
|
"Spend tracking - dropping %d spend log rows without per-row isolation; "
|
|
"isolation failure budget exhausted for this flush",
|
|
len(rows) - mid,
|
|
)
|
|
return 0
|
|
return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining)
|
|
|
|
|
|
def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_logging_obj: ProxyLogging):
|
|
"""
|
|
Raise an exception for failed update spend logs
|
|
|
|
- Calls proxy_logging_obj.failure_handler to log the error
|
|
- Ensures error messages says "Non-Blocking"
|
|
"""
|
|
import traceback
|
|
|
|
error_msg: Final = f"[Non-Blocking]LiteLLM Prisma Client Exception - update spend logs: {e}"
|
|
error_traceback: Final = error_msg + "\n" + traceback.format_exc()
|
|
end_time: Final = time.time()
|
|
_duration: Final = end_time - start_time
|
|
asyncio.create_task(
|
|
proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="update_spend",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
|
|
def _get_month_end_date(today: date) -> date:
|
|
if today.month == 12:
|
|
return date(today.year + 1, 1, 1) - timedelta(days=1)
|
|
return date(today.year, today.month + 1, 1) - timedelta(days=1)
|
|
|
|
|
|
def _is_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None):
|
|
if soft_budget_limit is None:
|
|
# If there's no limit, we can't exceed it.
|
|
return False
|
|
|
|
today: Final = date.today()
|
|
|
|
# Finding the first day of the next month, then subtracting one day to get the end of the current month.
|
|
end_month: Final = _get_month_end_date(today)
|
|
|
|
remaining_days: Final = (end_month - today).days
|
|
|
|
# Check for the start of the month to avoid division by zero
|
|
if today.day == 1:
|
|
daily_spend_estimate = current_spend
|
|
else:
|
|
daily_spend_estimate = current_spend / (today.day - 1)
|
|
|
|
# Total projected spend for the month
|
|
projected_spend: Final = current_spend + (daily_spend_estimate * remaining_days)
|
|
|
|
if projected_spend > soft_budget_limit:
|
|
print_verbose("Projected spend exceeds soft budget limit!")
|
|
return True
|
|
return False
|
|
|
|
|
|
def _get_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None) -> tuple | None:
|
|
if soft_budget_limit is None:
|
|
return None
|
|
|
|
today: Final = date.today()
|
|
end_month: Final = _get_month_end_date(today)
|
|
remaining_days: Final = (end_month - today).days
|
|
|
|
# assuming the current spend till today (not including today)
|
|
if today.day == 1:
|
|
daily_spend = current_spend
|
|
else:
|
|
daily_spend = current_spend / (today.day - 1)
|
|
projected_spend: Final = current_spend + (daily_spend * remaining_days)
|
|
|
|
if projected_spend > soft_budget_limit:
|
|
if daily_spend <= 0:
|
|
limit_exceed_date = today
|
|
else:
|
|
remaining_budget: Final = soft_budget_limit - current_spend
|
|
if remaining_budget <= 0:
|
|
limit_exceed_date = today
|
|
else:
|
|
approx_days: Final = remaining_budget / daily_spend
|
|
limit_exceed_date = today + timedelta(days=approx_days)
|
|
|
|
# return the projected spend and the date it will exceeded
|
|
return projected_spend, limit_exceed_date
|
|
|
|
return None
|
|
|
|
|
|
def _is_valid_team_configs(team_id=None, team_config=None, request_data=None):
|
|
if team_id is None or team_config is None or request_data is None:
|
|
return
|
|
# check if valid model called for team
|
|
if "models" in team_config:
|
|
valid_models: Final = team_config.pop("models")
|
|
model_in_request: Final = request_data["model"]
|
|
if model_in_request not in valid_models:
|
|
raise Exception(
|
|
f"Invalid model for team {team_id}: {model_in_request}. Valid models for team are: {valid_models}\n"
|
|
)
|
|
return
|
|
|
|
|
|
def _to_ns(dt):
|
|
return int(dt.timestamp() * 1e9)
|
|
|
|
|
|
def _check_and_merge_model_level_guardrails(
|
|
data: dict,
|
|
llm_router: Router | None,
|
|
trust_client_model_info: bool = True,
|
|
model_alias: str | None = None,
|
|
) -> dict:
|
|
"""
|
|
Check if the model has guardrails defined and merge them with existing guardrails in the request data.
|
|
|
|
Args:
|
|
data: The request data dict
|
|
llm_router: The LLM router instance to get deployment info from
|
|
model_alias: Resolve guardrails for this model group instead of data["model"]
|
|
trust_client_model_info: If False, ignore metadata.model_info.id and
|
|
resolve guardrails by alias-union only. Set to False on the
|
|
pre_call path because add_litellm_data_to_request preserves
|
|
client-supplied model_info when allow_client_pricing_override is
|
|
set, so a caller could spoof an unguarded model_info.id while
|
|
requesting a guarded alias and bypass guardrails (veria-ai HIGH
|
|
on #29654). Defaults to True for post_call paths where the
|
|
router has populated model_info.id itself.
|
|
|
|
Returns:
|
|
Modified data dict with merged guardrails (if any model-level guardrails exist)
|
|
"""
|
|
if llm_router is None:
|
|
return data
|
|
|
|
metadata: Final = data.get("metadata") or {}
|
|
litellm_metadata: Final = data.get("litellm_metadata") or {}
|
|
model_info: Final = metadata.get("model_info") or {}
|
|
model_id: Final = model_info.get("id") if trust_client_model_info else None
|
|
# route_request resolves team-scoped public model names with the
|
|
# server-populated team id; pre_call lookup must do the same so
|
|
# team-scoped guardrails are not silently skipped (greptile/veria-ai
|
|
# Medium on #29654).
|
|
team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id")
|
|
|
|
model_level_guardrails: list[object] | None = None
|
|
if model_id is not None:
|
|
deployment: Final = llm_router.get_deployment(model_id=model_id)
|
|
if deployment is None:
|
|
return data
|
|
deployment_guardrails: Final = deployment.litellm_params.get("guardrails")
|
|
# Bare-string guardrail names were truthy-accepted before; preserve
|
|
# that contract so post_call callers don't silently lose them.
|
|
if isinstance(deployment_guardrails, list):
|
|
model_level_guardrails = deployment_guardrails
|
|
elif deployment_guardrails:
|
|
model_level_guardrails = [deployment_guardrails]
|
|
else:
|
|
# Pre_call paths run before route_request picks a deployment, so we
|
|
# don't know which deployment's litellm_params.guardrails will apply.
|
|
# Take the UNION across all deployments in the group so a guardrail
|
|
# set on ANY eligible deployment still fires (#29652; addresses
|
|
# veria-ai HIGH on the single-deployment fallback that would skip
|
|
# non-first deployments).
|
|
alias: Final = model_alias if model_alias is not None else data.get("model")
|
|
if not isinstance(alias, str) or not alias:
|
|
return data
|
|
# Pass team_id so team-scoped public model names resolve the same way
|
|
# route_request resolves them; otherwise team-scoped deployments are
|
|
# invisible to this lookup and their guardrails are silently dropped.
|
|
deployments: Final = llm_router.get_model_list(model_name=alias, team_id=team_id) or []
|
|
seen: Final[set] = set()
|
|
union: Final[list] = []
|
|
for dep in deployments:
|
|
litellm_params_dep = dep.get("litellm_params") or {}
|
|
guardrails = litellm_params_dep.get("guardrails")
|
|
if isinstance(guardrails, str):
|
|
guardrails = [guardrails]
|
|
elif not isinstance(guardrails, list):
|
|
continue
|
|
for g in guardrails:
|
|
key = g if isinstance(g, str) else repr(g)
|
|
if key not in seen:
|
|
seen.add(key)
|
|
union.append(g)
|
|
model_level_guardrails = union or None
|
|
|
|
if model_level_guardrails is None:
|
|
return data
|
|
|
|
# Merge model-level guardrails with existing ones
|
|
return _merge_guardrails_with_existing(data, model_level_guardrails)
|
|
|
|
|
|
def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) -> dict:
|
|
"""
|
|
Merge model-level guardrails with any existing guardrails in the request data.
|
|
|
|
Args:
|
|
data: The request data dict
|
|
model_level_guardrails: Guardrails defined at the model level
|
|
|
|
Returns:
|
|
Modified data dict with merged guardrails in metadata
|
|
"""
|
|
modified_data: Final = data.copy()
|
|
metadata: Final = modified_data.setdefault("metadata", {})
|
|
existing_guardrails = metadata.get("guardrails", [])
|
|
|
|
# Ensure existing_guardrails is a list
|
|
if not isinstance(existing_guardrails, list):
|
|
existing_guardrails = [existing_guardrails] if existing_guardrails else []
|
|
|
|
# Ensure model_level_guardrails is a list
|
|
if not isinstance(model_level_guardrails, list):
|
|
model_level_guardrails = [model_level_guardrails] if model_level_guardrails else []
|
|
|
|
# Combine existing and model-level guardrails
|
|
metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails))
|
|
return modified_data
|
|
|
|
|
|
def get_error_message_str(e: Exception) -> str:
|
|
error_message = ""
|
|
if isinstance(e, HTTPException):
|
|
if isinstance(e.detail, str):
|
|
error_message = e.detail
|
|
elif isinstance(e.detail, dict):
|
|
error_message = json.dumps(e.detail)
|
|
elif hasattr(e, "message"):
|
|
_error: Final = getattr(e, "message", None)
|
|
if isinstance(_error, str):
|
|
error_message = _error
|
|
elif isinstance(_error, dict):
|
|
error_message = json.dumps(_error)
|
|
else:
|
|
error_message = str(e)
|
|
else:
|
|
error_message = str(e)
|
|
return error_message
|
|
|
|
|
|
def _get_redoc_url() -> str | None:
|
|
"""
|
|
Get the Redoc URL from the environment variables.
|
|
|
|
- If REDOC_URL is set, return it.
|
|
- If NO_REDOC is True, return None.
|
|
- Otherwise, default to "/redoc".
|
|
"""
|
|
if redoc_url := os.getenv("REDOC_URL"):
|
|
return redoc_url
|
|
|
|
if str_to_bool(os.getenv("NO_REDOC")) is True:
|
|
return None
|
|
|
|
return "/redoc"
|
|
|
|
|
|
def _get_docs_url() -> str | None:
|
|
"""
|
|
Get the docs (Swagger UI) URL from the environment variables.
|
|
|
|
- If DOCS_URL is set, return it.
|
|
- If NO_DOCS is True, return None.
|
|
- Otherwise, default to "/".
|
|
"""
|
|
if docs_url := os.getenv("DOCS_URL"):
|
|
return docs_url
|
|
|
|
if str_to_bool(os.getenv("NO_DOCS")) is True:
|
|
return None
|
|
|
|
return "/"
|
|
|
|
|
|
def _get_openapi_url() -> str | None:
|
|
"""
|
|
Get the OpenAPI JSON URL from the environment variables.
|
|
|
|
- If OPENAPI_URL is set, return it.
|
|
- If NO_OPENAPI is True, return None.
|
|
- Otherwise, default to "/openapi.json".
|
|
"""
|
|
if openapi_url := os.getenv("OPENAPI_URL"):
|
|
return openapi_url
|
|
|
|
if str_to_bool(os.getenv("NO_OPENAPI")) is True:
|
|
return None
|
|
|
|
return "/openapi.json"
|
|
|
|
|
|
def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | None") -> None:
|
|
if prisma_client is None:
|
|
return
|
|
asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction"))
|
|
|
|
|
|
def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException:
|
|
"""
|
|
Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible
|
|
"""
|
|
from fastapi import status
|
|
|
|
verbose_proxy_logger.exception("Exception: %s", e)
|
|
if PrismaDBExceptionHandler.is_read_only_transaction_error(e):
|
|
from litellm.proxy.proxy_server import prisma_client
|
|
|
|
_recreate_writer_on_read_only_transaction(prisma_client)
|
|
|
|
headers: Final = litellm_call_id_headers(litellm_call_id)
|
|
if isinstance(e, HTTPException):
|
|
return ProxyException(
|
|
message=getattr(e, "detail", f"error({e})"),
|
|
type=ProxyErrorTypes.internal_server_error,
|
|
param=openai_error_param(e),
|
|
headers=headers,
|
|
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
return with_litellm_call_id(e, litellm_call_id)
|
|
_status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
|
|
return ProxyException(
|
|
message=str(e),
|
|
type=ProxyErrorTypes.internal_server_error,
|
|
param=openai_error_param(e),
|
|
headers=headers,
|
|
code=_status_code,
|
|
)
|
|
|
|
|
|
def _premium_user_check(feature: str | None = None):
|
|
"""
|
|
Raises an HTTPException if the user is not a premium user
|
|
"""
|
|
from litellm.proxy.proxy_server import premium_user
|
|
|
|
if feature:
|
|
detail_msg = f"This feature is only available for LiteLLM Enterprise users: {feature}. {CommonProxyErrors.not_premium_user.value}"
|
|
else:
|
|
detail_msg = (
|
|
f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}"
|
|
)
|
|
|
|
if not premium_user:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail={"error": detail_msg},
|
|
)
|
|
|
|
|
|
def is_known_model(model: str | None, llm_router: Router | None) -> bool:
|
|
"""
|
|
Returns True if the model is in the llm_router model names
|
|
"""
|
|
if model is None or llm_router is None:
|
|
return False
|
|
model_names: Final = llm_router.get_model_names()
|
|
|
|
model_names_set: Final = set(model_names)
|
|
|
|
is_in_list = False
|
|
if model in model_names_set:
|
|
is_in_list = True
|
|
|
|
return is_in_list
|
|
|
|
|
|
def is_known_vector_store_index(index_name: str) -> bool:
|
|
"""
|
|
Returns True if the vector store index is in the llm_router vector store indexes
|
|
"""
|
|
|
|
if litellm.vector_store_index_registry is None:
|
|
return False
|
|
return index_name in litellm.vector_store_index_registry.get_vector_store_indexes()
|
|
|
|
|
|
def join_paths(base_path: str, route: str) -> str:
|
|
# Remove trailing slashes from base_path and leading slashes from route
|
|
base_path = base_path.rstrip("/")
|
|
route = route.lstrip("/")
|
|
|
|
# If base_path is empty, return route with leading slash
|
|
if not base_path:
|
|
return f"/{route}" if route else "/"
|
|
|
|
# If route is empty, return just base_path
|
|
if not route:
|
|
return base_path
|
|
|
|
# Check if base_path already ends with the route to avoid duplication
|
|
if base_path.endswith(f"/{route}"):
|
|
final_path = base_path
|
|
else:
|
|
# Join with single slash
|
|
final_path = f"{base_path}/{route}"
|
|
|
|
return final_path
|
|
|
|
|
|
def get_custom_url(request_base_url: str, route: str | None = None) -> str:
|
|
# Use environment variable value, otherwise use URL from request
|
|
server_base_url: Final = get_proxy_base_url()
|
|
if server_base_url is not None:
|
|
base_url = server_base_url
|
|
else:
|
|
base_url = request_base_url
|
|
|
|
# get_request_root_path() returns the prefix the router is actually
|
|
# resolving this request under: the matched SERVER_ROOT_PATHS entry when
|
|
# PerRequestRootPathMiddleware ran, otherwise the SERVER_ROOT_PATH scalar.
|
|
# This keeps the emitted URL under one prefix — the one the client called —
|
|
# instead of stacking the scalar onto a request already living under a
|
|
# dynamic prefix (which would produce /tenant-a/legacy/... — a path that
|
|
# doesn't exist). join_paths()'s tail-dedup then collapses the append when
|
|
# base_url (i.e. request.base_url) already ends in the same prefix.
|
|
from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports utils
|
|
get_request_root_path,
|
|
)
|
|
|
|
server_root_path: Final = get_request_root_path()
|
|
if route is not None:
|
|
if server_root_path != "":
|
|
# First join base_url with server_root_path, then with route
|
|
intermediate_url: Final = join_paths(base_url, server_root_path)
|
|
return join_paths(intermediate_url, route)
|
|
else:
|
|
return join_paths(base_url, route)
|
|
else:
|
|
return join_paths(base_url, server_root_path)
|
|
|
|
|
|
def get_proxy_base_url() -> str | None:
|
|
"""
|
|
Get the proxy base url from the environment variables.
|
|
"""
|
|
return os.getenv("PROXY_BASE_URL")
|
|
|
|
|
|
def get_server_root_path() -> str:
|
|
"""
|
|
Get the server root path from the environment variables.
|
|
|
|
- If SERVER_ROOT_PATH is set, return it.
|
|
- Otherwise, default to "/".
|
|
"""
|
|
return os.getenv("SERVER_ROOT_PATH", "")
|
|
|
|
|
|
def normalize_route_for_root_path(route: str) -> str | None:
|
|
"""Strip SERVER_ROOT_PATH prefix. Returns de-prefixed route, or None if route is not under root path."""
|
|
root_path: Final = get_server_root_path()
|
|
if root_path and root_path != "/":
|
|
if route.startswith(root_path + "/"):
|
|
return route[len(root_path) :]
|
|
return None
|
|
return route
|
|
|
|
|
|
def get_prisma_client_or_throw(message: str):
|
|
from litellm.proxy.proxy_server import prisma_client
|
|
|
|
if prisma_client is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
detail={"error": message},
|
|
)
|
|
return prisma_client
|
|
|
|
|
|
def is_valid_api_key(key: str) -> bool:
|
|
"""
|
|
Validates API key format:
|
|
- sk- keys: must match ^sk-[A-Za-z0-9_-]+$
|
|
- hashed keys: must match ^[a-fA-F0-9]{64}$
|
|
- Length between 20 and 100 characters
|
|
"""
|
|
import re
|
|
|
|
if not isinstance(key, str):
|
|
return False
|
|
if 3 <= len(key) <= 100:
|
|
if re.match(r"^sk-[A-Za-z0-9_-]+$", key):
|
|
return True
|
|
if re.match(r"^[a-fA-F0-9]{64}$", key):
|
|
return True
|
|
return False
|
|
|
|
|
|
def construct_database_url_from_env_vars() -> str | None:
|
|
"""
|
|
Construct a DATABASE_URL from individual environment variables.
|
|
Returns:
|
|
Optional[str]: The constructed DATABASE_URL or None if required variables are missing
|
|
"""
|
|
import urllib.parse
|
|
|
|
# Check if all required variables are provided
|
|
database_host: Final = os.getenv("DATABASE_HOST")
|
|
database_username: Final = os.getenv("DATABASE_USERNAME")
|
|
database_password: Final = os.getenv("DATABASE_PASSWORD")
|
|
database_name: Final = os.getenv("DATABASE_NAME")
|
|
database_schema: Final = os.getenv("DATABASE_SCHEMA")
|
|
|
|
if database_host and database_username and database_name:
|
|
# Handle the problem of special character escaping in the database URL
|
|
database_username_enc: Final = urllib.parse.quote_plus(database_username)
|
|
database_password_enc: Final = urllib.parse.quote_plus(database_password) if database_password else ""
|
|
database_name_enc: Final = urllib.parse.quote_plus(database_name)
|
|
|
|
# Construct DATABASE_URL from the provided variables
|
|
if database_password:
|
|
database_url = (
|
|
f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}"
|
|
)
|
|
else:
|
|
database_url = f"postgresql://{database_username_enc}@{database_host}/{database_name_enc}"
|
|
|
|
if database_schema:
|
|
database_url += f"?schema={database_schema}"
|
|
|
|
return add_missing_query_params(database_url, DatabaseURLSettings.from_env().tls_params())
|
|
|
|
return None
|
|
|
|
|
|
async def _get_validated_team_object(
|
|
user_api_key_dict: "UserAPIKeyAuth",
|
|
team_id: str,
|
|
prisma_client: "PrismaClient",
|
|
user_api_key_cache: "UserApiKeyCache",
|
|
proxy_logging_obj: "ProxyLogging",
|
|
) -> "LiteLLM_TeamTableCachedObj":
|
|
from litellm.proxy.auth.auth_checks import get_team_object
|
|
from litellm.proxy.management_endpoints.team_endpoints import validate_membership
|
|
|
|
team_object: Final = await get_team_object(
|
|
team_id=team_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object)
|
|
return team_object
|
|
|
|
|
|
async def _get_team_object_for_access_groups(
|
|
team_id: str | None,
|
|
prisma_client: Optional["PrismaClient"],
|
|
user_api_key_cache: Optional["UserApiKeyCache"],
|
|
proxy_logging_obj: Optional["ProxyLogging"],
|
|
) -> Optional["LiteLLM_TeamTableCachedObj"]:
|
|
from litellm.proxy.auth.auth_checks import get_team_object
|
|
|
|
if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None:
|
|
return None
|
|
try:
|
|
return await get_team_object(
|
|
team_id=team_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
except HTTPException:
|
|
verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id)
|
|
return None
|
|
|
|
|
|
async def _get_access_group_models(
|
|
user_api_key_dict: "UserAPIKeyAuth",
|
|
team_object: Optional["LiteLLM_TeamTableCachedObj"],
|
|
prisma_client: Optional["PrismaClient"],
|
|
user_api_key_cache: Optional["UserApiKeyCache"],
|
|
proxy_logging_obj: Optional["ProxyLogging"],
|
|
) -> tuple[str, ...]:
|
|
from litellm.proxy.auth.auth_checks import (
|
|
_get_models_from_access_groups,
|
|
get_authorized_resources_from_key_access_groups,
|
|
)
|
|
|
|
team_group_models: Final = await _get_models_from_access_groups(
|
|
access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (),
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
key_group_models: Final = await get_authorized_resources_from_key_access_groups(
|
|
valid_token=user_api_key_dict,
|
|
team_object=team_object,
|
|
resource_field="access_model_names",
|
|
)
|
|
return tuple(dict.fromkeys((*team_group_models, *key_group_models)))
|
|
|
|
|
|
async def get_available_models_for_user(
|
|
user_api_key_dict: "UserAPIKeyAuth",
|
|
llm_router: Optional["Router"],
|
|
general_settings: dict,
|
|
user_model: str | None,
|
|
prisma_client: Optional["PrismaClient"] = None,
|
|
proxy_logging_obj: Optional["ProxyLogging"] = None,
|
|
team_id: str | None = None,
|
|
include_model_access_groups: bool = False,
|
|
only_model_access_groups: bool = False,
|
|
return_wildcard_routes: bool = False,
|
|
user_api_key_cache: Optional["UserApiKeyCache"] = None,
|
|
) -> list[str]:
|
|
"""
|
|
Get the list of models available to a user based on their API key and team permissions.
|
|
|
|
Args:
|
|
user_api_key_dict: User API key authentication object
|
|
llm_router: LiteLLM router instance
|
|
general_settings: General settings from config
|
|
user_model: User-specific model
|
|
prisma_client: Prisma client for database operations
|
|
proxy_logging_obj: Proxy logging object
|
|
team_id: Specific team ID to check (optional)
|
|
include_model_access_groups: Whether to include model access groups
|
|
only_model_access_groups: Whether to only return model access groups
|
|
return_wildcard_routes: Whether to return wildcard routes
|
|
|
|
Returns:
|
|
List of model names available to the user
|
|
"""
|
|
from litellm.proxy.auth.model_checks import (
|
|
get_complete_model_list,
|
|
get_key_models,
|
|
get_team_models,
|
|
)
|
|
|
|
# Get proxy model list and access groups
|
|
if llm_router is None:
|
|
proxy_model_list = []
|
|
model_access_groups = {}
|
|
else:
|
|
proxy_model_list = llm_router.get_model_names()
|
|
model_access_groups = llm_router.get_model_access_groups()
|
|
|
|
requested_team_object: Final = (
|
|
await _get_validated_team_object(
|
|
user_api_key_dict=user_api_key_dict,
|
|
team_id=team_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
if team_id and prisma_client and proxy_logging_obj and user_api_key_cache
|
|
else None
|
|
)
|
|
|
|
key_models: Final[Sequence[str]] = (
|
|
()
|
|
if requested_team_object is not None
|
|
else get_key_models(
|
|
user_api_key_dict=user_api_key_dict,
|
|
proxy_model_list=proxy_model_list,
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=include_model_access_groups,
|
|
)
|
|
)
|
|
|
|
team_models: Final = get_team_models(
|
|
team_models=(
|
|
requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models
|
|
),
|
|
proxy_model_list=proxy_model_list,
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=include_model_access_groups,
|
|
)
|
|
|
|
effective_team_id: Final = team_id or user_api_key_dict.team_id
|
|
|
|
access_group_models: Final = (
|
|
await _get_access_group_models(
|
|
user_api_key_dict=user_api_key_dict,
|
|
team_object=requested_team_object
|
|
or await _get_team_object_for_access_groups(
|
|
team_id=effective_team_id,
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
),
|
|
prisma_client=prisma_client,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
if key_models or team_models
|
|
else ()
|
|
)
|
|
|
|
granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models
|
|
granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models
|
|
|
|
# Get complete model list
|
|
all_models: Final = get_complete_model_list(
|
|
key_models=granted_key_models,
|
|
team_models=granted_team_models,
|
|
proxy_model_list=proxy_model_list,
|
|
user_model=user_model,
|
|
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
|
|
return_wildcard_routes=return_wildcard_routes,
|
|
llm_router=llm_router,
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=include_model_access_groups,
|
|
only_model_access_groups=only_model_access_groups,
|
|
team_id=effective_team_id,
|
|
)
|
|
|
|
return all_models
|
|
|
|
|
|
def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None:
|
|
try:
|
|
return get_model_info(model)
|
|
except Exception as e:
|
|
verbose_proxy_logger.debug(
|
|
"create_model_info_response: cost map lookup failed for %s: %s",
|
|
model,
|
|
e,
|
|
)
|
|
return None
|
|
|
|
|
|
def _resolve_listing_model_info(
|
|
deployment_model: str | None,
|
|
listed_model: str,
|
|
listed_info: ModelInfo | None,
|
|
get_model_info: Callable[[str], ModelInfo],
|
|
) -> tuple[ModelInfo, ...]:
|
|
"""
|
|
Cost-map entries describing one deployment behind a listed model, best source first.
|
|
|
|
The name a model is listed under is an arbitrary public alias, so it often misses the
|
|
cost map and lands on a fallback-generalization rule that answers with a conservative
|
|
family baseline instead of the real model's limits; the deployment's underlying model
|
|
is what the request actually reaches. Both names are kept because either can
|
|
generalize, and because a deployment's own model is registered into the cost map as a
|
|
stub that carries no limits of its own. Exact entries are consulted before generalized
|
|
ones, and each field is then taken from the first entry that has it.
|
|
|
|
``listed_info`` is resolved once by the caller, since a group with several distinct
|
|
underlying models resolves the same alias for each of them.
|
|
"""
|
|
# Fast path, and the only one a wildcard-expanded name takes: with a single name
|
|
# there is nothing to order, so skip the generalization test entirely. This keeps
|
|
# the per-model cost of the listing on the hot path #33721 exists to protect.
|
|
if deployment_model is None or deployment_model == listed_model:
|
|
return () if listed_info is None else (listed_info,)
|
|
|
|
deployment_info: Final = _safe_get_model_info(deployment_model, get_model_info)
|
|
if deployment_info is None:
|
|
return () if listed_info is None else (listed_info,)
|
|
if listed_info is None:
|
|
return (deployment_info,)
|
|
|
|
from litellm.utils import is_generalized_model_info
|
|
|
|
# Both names resolved: the deployment's model leads unless it only generalized
|
|
# while the listed name is an exact cost-map entry.
|
|
if is_generalized_model_info(deployment_info) and not is_generalized_model_info(listed_info):
|
|
return (listed_info, deployment_info)
|
|
return (deployment_info, listed_info)
|
|
|
|
|
|
def _first_token_limit(candidates: tuple[ModelInfo, ...], field: str) -> int | None:
|
|
return next(
|
|
(limit for limit in (coerce_token_limit(info.get(field)) for info in candidates) if limit is not None),
|
|
None,
|
|
)
|
|
|
|
|
|
def _group_token_limit(candidate_sets: tuple[tuple[ModelInfo, ...], ...], field: str) -> int | None:
|
|
"""The widest limit any deployment behind the listed name declares for ``field``.
|
|
|
|
A model group is normally one model behind several interchangeable deployments, so
|
|
there is a single value to report and the choice of aggregate does not arise.
|
|
|
|
When a group genuinely mixes models no single number is right, and the widest is the
|
|
deliberate pick over the narrowest for two reasons. It is what ``/model_group/info``
|
|
has long reported to the Admin UI, so the two surfaces agree; disagreeing is the very
|
|
complaint this resolution path exists to fix. And of the two ways to be wrong,
|
|
under-advertising is worse: a client that trusts a narrowed window silently refuses
|
|
prompts the group would have served, while an over-long prompt that reaches a smaller
|
|
deployment comes back as a legible context-length error -- and does not reach one at
|
|
all when ``enable_pre_call_checks`` is set, which filters deployments the prompt does
|
|
not fit.
|
|
"""
|
|
limits: Final = tuple(
|
|
limit for limit in (_first_token_limit(candidates, field) for candidates in candidate_sets) if limit is not None
|
|
)
|
|
return max(limits) if limits else None
|
|
|
|
|
|
def create_model_info_response(
|
|
model_id: str,
|
|
provider: str,
|
|
include_metadata: bool = False,
|
|
fallback_type: str | None = None,
|
|
llm_router: Optional["Router"] = None,
|
|
get_model_info: Callable[[str], ModelInfo] = litellm.get_model_info,
|
|
) -> ModelInfoResponse:
|
|
"""
|
|
Create a standardized OpenAI-compatible model object.
|
|
|
|
When include_metadata is true, attaches the model's configured fallbacks
|
|
(resolved via the router under fallback_type, defaulting to "general").
|
|
Raises HTTPException(400) for an unknown fallback_type.
|
|
"""
|
|
from litellm.proxy.auth.model_checks import get_all_fallbacks
|
|
|
|
base: Final[ModelInfoResponse] = {
|
|
"id": model_id,
|
|
"object": "model",
|
|
"created": DEFAULT_MODEL_CREATED_AT_TIME,
|
|
"owned_by": provider,
|
|
}
|
|
|
|
alias_target: Final = (
|
|
resolve_model_group_alias(llm_router.model_group_alias, model_id) if llm_router is not None else None
|
|
)
|
|
lookup_model: Final = alias_target if alias_target is not None else model_id
|
|
|
|
listing_info: Final = llm_router.get_model_listing_info(lookup_model) if llm_router is not None else None
|
|
|
|
# One entry per distinct model behind the listed name; (None,) when the router knows
|
|
# nothing about it, so the listed name is resolved on its own as before.
|
|
deployment_models: Final[tuple[str | None, ...]] = (
|
|
listing_info.cost_map_keys if listing_info is not None and listing_info.cost_map_keys else (None,)
|
|
)
|
|
listed_info: Final = _safe_get_model_info(lookup_model, get_model_info)
|
|
candidate_sets: Final = tuple(
|
|
_resolve_listing_model_info(
|
|
deployment_model=deployment_model,
|
|
listed_model=lookup_model,
|
|
listed_info=listed_info,
|
|
get_model_info=get_model_info,
|
|
)
|
|
for deployment_model in deployment_models
|
|
)
|
|
|
|
max_input_tokens: int | None = _group_token_limit(candidate_sets, "max_input_tokens")
|
|
max_output_tokens: int | None = _group_token_limit(candidate_sets, "max_output_tokens")
|
|
mode: Final = next(
|
|
(
|
|
m
|
|
for m in (
|
|
cast("Mapping[str, object]", info).get("mode") # cast-ok: an entry need not carry "mode"
|
|
for candidates in candidate_sets
|
|
for info in candidates
|
|
)
|
|
if isinstance(m, str)
|
|
),
|
|
None,
|
|
)
|
|
if mode is not None:
|
|
base["mode"] = mode
|
|
|
|
if listing_info is not None:
|
|
if listing_info.max_input_tokens is not None:
|
|
max_input_tokens = listing_info.max_input_tokens
|
|
if listing_info.max_output_tokens is not None:
|
|
max_output_tokens = listing_info.max_output_tokens
|
|
|
|
if llm_router is not None:
|
|
configured_mode: Final = llm_router.get_configured_mode(lookup_model)
|
|
if isinstance(configured_mode, str):
|
|
base["mode"] = configured_mode
|
|
|
|
if max_input_tokens is not None:
|
|
base["max_input_tokens"] = max_input_tokens
|
|
if max_output_tokens is not None:
|
|
base["max_output_tokens"] = max_output_tokens
|
|
|
|
if not include_metadata:
|
|
return base
|
|
|
|
effective_fallback_type: Final = fallback_type if fallback_type is not None else "general"
|
|
|
|
valid_fallback_types: Final = ["general", "context_window", "content_policy"]
|
|
if effective_fallback_type not in valid_fallback_types:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"Invalid fallback_type. Must be one of: {valid_fallback_types}",
|
|
)
|
|
|
|
fallbacks: Final = get_all_fallbacks(
|
|
model=model_id,
|
|
llm_router=llm_router,
|
|
fallback_type=effective_fallback_type,
|
|
)
|
|
return {**base, "metadata": {"fallbacks": fallbacks}}
|
|
|
|
|
|
def validate_model_access(
|
|
model_id: str,
|
|
available_models: list[str],
|
|
) -> None:
|
|
"""
|
|
Validate that a model is accessible to the user.
|
|
Supports batch requests with comma-separated model IDs.
|
|
|
|
Args:
|
|
model_id: The model ID to validate (can be comma-separated for batch requests)
|
|
available_models: List of models available to the user
|
|
|
|
Raises:
|
|
HTTPException: If the model is not accessible
|
|
"""
|
|
# Handle batch requests with comma-separated models
|
|
if "," in model_id:
|
|
models: Final = [m.strip() for m in model_id.split(",")]
|
|
inaccessible_models: Final = [m for m in models if m not in available_models]
|
|
if inaccessible_models:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="The following model(s) do not exist or are not accessible: {}".format(
|
|
", ".join(inaccessible_models)
|
|
),
|
|
)
|
|
else:
|
|
# Single model validation
|
|
if model_id not in available_models:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"The model `{model_id}` does not exist or is not accessible",
|
|
)
|
|
|
|
|
|
_PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [
|
|
("message", "content"), # null when tool_calls present (issue #6677)
|
|
("message", "role"), # always required by OpenAI spec
|
|
("delta", "content"), # null in streaming chunks
|
|
]
|
|
|
|
|
|
def model_dump_with_preserved_fields(
|
|
obj: Any,
|
|
preserve_fields: list[str] | None = None,
|
|
exclude_unset: bool = True,
|
|
) -> dict[str, object]:
|
|
"""
|
|
Serialize a Pydantic model to a dictionary while preserving specific fields
|
|
even if they are None.
|
|
|
|
Fields listed in _PRESERVED_NONE_FIELDS are restored after
|
|
model_dump(exclude_none=True) strips them.
|
|
|
|
Args:
|
|
obj: The Pydantic BaseModel instance to serialize
|
|
preserve_fields: Deprecated, kept for backward compatibility.
|
|
exclude_unset: Whether to exclude fields that were not explicitly set
|
|
|
|
Returns:
|
|
Dictionary representation with None values excluded except for preserved fields
|
|
"""
|
|
result: Final = obj.model_dump(exclude_none=True, exclude_unset=exclude_unset)
|
|
|
|
choices: Final = result.get("choices")
|
|
if not choices:
|
|
return result
|
|
|
|
obj_choices: Final = obj.choices
|
|
for choice_obj, choice_dict in zip(obj_choices, choices):
|
|
for sub_object, field_name in _PRESERVED_NONE_FIELDS:
|
|
sub_dict = choice_dict.get(sub_object)
|
|
if sub_dict is None:
|
|
continue
|
|
if field_name not in sub_dict:
|
|
sub_obj = getattr(choice_obj, sub_object, None)
|
|
if sub_obj is not None and hasattr(sub_obj, field_name):
|
|
sub_dict[field_name] = getattr(sub_obj, field_name)
|
|
|
|
return result
|