4006 lines
154 KiB
Python
4006 lines
154 KiB
Python
import asyncio
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import smtplib
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from datetime import datetime, timedelta
|
|
from email.mime.multipart import MIMEMultipart
|
|
from email.mime.text import MIMEText
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Dict,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Union,
|
|
cast,
|
|
overload,
|
|
)
|
|
|
|
from litellm.constants import MAX_TEAM_LIST_LIMIT, DEFAULT_MODEL_CREATED_AT_TIME
|
|
from litellm.proxy._types import (
|
|
DB_CONNECTION_ERROR_TYPES,
|
|
CommonProxyErrors,
|
|
ProxyErrorTypes,
|
|
ProxyException,
|
|
SpendLogsMetadata,
|
|
SpendLogsPayload,
|
|
)
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
try:
|
|
import backoff
|
|
except ImportError:
|
|
raise ImportError(
|
|
"backoff is not installed. Please install it via 'pip install backoff'"
|
|
)
|
|
|
|
from fastapi import HTTPException, status
|
|
|
|
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 verbose_proxy_logger
|
|
from litellm._service_logger import ServiceLogging, ServiceTypes
|
|
from litellm.caching.caching import DualCache, RedisCache
|
|
from litellm.exceptions import (
|
|
BlockedPiiEntityError,
|
|
GuardrailRaisedException,
|
|
RejectedRequestError,
|
|
)
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
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.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.llms.custom_httpx.httpx_handler import HTTPHandler
|
|
from litellm.proxy._types import (
|
|
AlertType,
|
|
CallInfo,
|
|
LiteLLM_VerificationTokenView,
|
|
Member,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.proxy.auth.route_checks import RouteChecks
|
|
from litellm.proxy.db.create_views import (
|
|
create_missing_views,
|
|
should_create_missing_views,
|
|
)
|
|
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
|
from litellm.proxy.db.log_db_metrics import log_db_metrics
|
|
from litellm.proxy.db.prisma_client import PrismaWrapper
|
|
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
|
|
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
|
|
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
|
|
from litellm.proxy.hooks.parallel_request_limiter import (
|
|
_PROXY_MaxParallelRequestsHandler,
|
|
)
|
|
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
|
from litellm.secret_managers.main import str_to_bool
|
|
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
|
from litellm.types.mcp import (
|
|
MCPDuringCallResponseObject,
|
|
MCPPreCallRequestObject,
|
|
MCPPreCallResponseObject,
|
|
)
|
|
from litellm.types.utils import CallTypes, LLMResponseTypes, LoggedLiteLLMParams
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
|
|
|
Span = Union[_Span, Any]
|
|
else:
|
|
Span = Any
|
|
|
|
|
|
def print_verbose(print_statement):
|
|
"""
|
|
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("{}\n{}".format(print_statement, traceback.format_exc()))
|
|
if litellm.set_verbose:
|
|
print(f"LiteLLM Proxy: {print_statement}") # noqa
|
|
|
|
|
|
class InternalUsageCache:
|
|
def __init__(self, dual_cache: DualCache):
|
|
self.dual_cache: DualCache = dual_cache
|
|
|
|
async def async_get_cache(
|
|
self,
|
|
key,
|
|
litellm_parent_otel_span: Union[Span, None],
|
|
local_only: bool = False,
|
|
**kwargs,
|
|
) -> 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,
|
|
value,
|
|
litellm_parent_otel_span: Union[Span, None],
|
|
local_only: bool = False,
|
|
**kwargs,
|
|
) -> 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,
|
|
litellm_parent_otel_span: Union[Span, None],
|
|
local_only: bool = False,
|
|
**kwargs,
|
|
) -> 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: list,
|
|
parent_otel_span: Optional[Span] = None,
|
|
local_only: bool = False,
|
|
):
|
|
return await self.dual_cache.async_batch_get_cache(
|
|
keys=keys,
|
|
parent_otel_span=parent_otel_span,
|
|
local_only=local_only,
|
|
)
|
|
|
|
async def async_increment_cache(
|
|
self,
|
|
key,
|
|
value: float,
|
|
litellm_parent_otel_span: Union[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,
|
|
value,
|
|
local_only: bool = False,
|
|
**kwargs,
|
|
) -> None:
|
|
return self.dual_cache.set_cache(
|
|
key=key,
|
|
value=value,
|
|
local_only=local_only,
|
|
**kwargs,
|
|
)
|
|
|
|
def get_cache(
|
|
self,
|
|
key,
|
|
local_only: bool = False,
|
|
**kwargs,
|
|
) -> Any:
|
|
return self.dual_cache.get_cache(
|
|
key=key,
|
|
local_only=local_only,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
### LOGGING ###
|
|
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: DualCache,
|
|
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.max_budget_limiter = _PROXY_MaxBudgetLimiter()
|
|
self.cache_control_check = _PROXY_CacheControlCheck()
|
|
self.alerting: Optional[List] = None
|
|
self.alerting_threshold: float = 300 # default to 5 min. threshold
|
|
self.alert_types: List[AlertType] = DEFAULT_ALERT_TYPES
|
|
self.alert_to_webhook_url: Optional[dict] = 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.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
|
|
|
|
def startup_event(
|
|
self,
|
|
llm_router: Optional[Router],
|
|
redis_usage_cache: Optional[RedisCache],
|
|
):
|
|
"""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
|
|
)
|
|
) # 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
|
|
|
|
def update_values(
|
|
self,
|
|
alerting: Optional[List] = None,
|
|
alerting_threshold: Optional[float] = None,
|
|
redis_cache: Optional[RedisCache] = None,
|
|
alert_types: Optional[List[AlertType]] = None,
|
|
alerting_args: Optional[dict] = None,
|
|
alert_to_webhook_url: Optional[dict] = 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 updated_slack_alerting is True:
|
|
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,
|
|
)
|
|
|
|
if self.alerting is not None and "slack" 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) # type: ignore
|
|
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: Optional[Router] = 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)
|
|
import inspect
|
|
|
|
expected_args = inspect.getfullargspec(proxy_hook).args
|
|
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) -> Optional[CustomLogger]:
|
|
"""
|
|
Get a proxy hook from the proxy_hook_mapping
|
|
"""
|
|
return self.proxy_hook_mapping.get(hook)
|
|
|
|
def _init_litellm_callbacks(self, llm_router: Optional[Router] = None):
|
|
self._add_proxy_hooks(llm_router)
|
|
litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore
|
|
for callback in litellm.callbacks:
|
|
if isinstance(callback, str):
|
|
callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore
|
|
callback,
|
|
internal_usage_cache=self.internal_usage_cache.dual_cache,
|
|
llm_router=llm_router,
|
|
)
|
|
if callback is None:
|
|
continue
|
|
if callback not in litellm.input_callback:
|
|
litellm.input_callback.append(callback) # type: ignore
|
|
if callback not in litellm.success_callback:
|
|
litellm.logging_callback_manager.add_litellm_success_callback(callback) # type: ignore
|
|
if callback not in litellm.failure_callback:
|
|
litellm.logging_callback_manager.add_litellm_failure_callback(callback) # type: ignore
|
|
if callback not in litellm._async_success_callback:
|
|
litellm.logging_callback_manager.add_litellm_async_success_callback(callback) # type: ignore
|
|
if callback not in litellm._async_failure_callback:
|
|
litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) # type: ignore
|
|
if callback not in litellm.service_callback:
|
|
litellm.service_callback.append(callback) # type: ignore
|
|
|
|
if (
|
|
len(litellm.input_callback) > 0
|
|
or len(litellm.success_callback) > 0
|
|
or len(litellm.failure_callback) > 0
|
|
):
|
|
callback_list = list(
|
|
set(
|
|
litellm.input_callback
|
|
+ litellm.success_callback
|
|
+ litellm.failure_callback
|
|
)
|
|
)
|
|
litellm.litellm_core_utils.litellm_logging.set_callbacks(
|
|
callback_list=callback_list
|
|
)
|
|
|
|
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="request_status:{}".format(litellm_call_id),
|
|
value=status,
|
|
local_only=True,
|
|
ttl=alerting_threshold,
|
|
litellm_parent_otel_span=None,
|
|
)
|
|
|
|
async def async_pre_mcp_tool_call_hook(
|
|
self,
|
|
kwargs: dict,
|
|
request_obj: Any,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
) -> Optional[Any]:
|
|
"""
|
|
Pre MCP Tool Call Hook
|
|
|
|
Use this to validate and modify MCP tool calls before execution.
|
|
Reuses existing LLM guardrail logic by converting MCP calls to message format.
|
|
"""
|
|
from litellm.types.llms.base import HiddenParams
|
|
from litellm.types.mcp import MCPPreCallRequestObject
|
|
|
|
callbacks = self.get_combined_callback_list(
|
|
dynamic_success_callbacks=getattr(self, "dynamic_success_callbacks", None),
|
|
global_callbacks=litellm.success_callback,
|
|
)
|
|
|
|
# Create the request object if it's not already one
|
|
if not isinstance(request_obj, MCPPreCallRequestObject):
|
|
# Convert UserAPIKeyAuth object to dict if needed
|
|
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(
|
|
kwargs.get("user_api_key_auth")
|
|
)
|
|
|
|
request_obj = 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(),
|
|
)
|
|
|
|
for callback in callbacks:
|
|
try:
|
|
_callback: Optional[CustomLogger] = None
|
|
if isinstance(callback, str):
|
|
from typing import cast
|
|
|
|
from litellm import _custom_logger_compatible_callbacks_literal
|
|
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback)
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
|
|
if _callback is not None and isinstance(_callback, CustomGuardrail):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
# Check if guardrail should be run for pre_call hook (reusing existing logic)
|
|
if (
|
|
_callback.should_run_guardrail(
|
|
data=kwargs, event_type=GuardrailEventHooks.pre_mcp_call
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
|
|
# Convert MCP tool call to LLM message format for existing guardrail logic
|
|
synthetic_llm_data = self._convert_mcp_to_llm_format(
|
|
request_obj, kwargs
|
|
)
|
|
# Reuse existing LLM guardrail logic
|
|
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(
|
|
kwargs.get("user_api_key_auth")
|
|
)
|
|
|
|
result = await _callback.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_auth_dict, # type: ignore
|
|
cache=self.call_details["user_api_key_cache"],
|
|
data=synthetic_llm_data,
|
|
call_type="mcp_call",
|
|
)
|
|
|
|
# Convert result back to MCP response format if blocked/modified
|
|
if result is not None:
|
|
mcp_response = self._convert_llm_result_to_mcp_response(
|
|
result, request_obj
|
|
)
|
|
if mcp_response is not None:
|
|
return self._parse_pre_mcp_call_hook_response(
|
|
response=mcp_response, original_request=request_obj
|
|
)
|
|
|
|
except (
|
|
BlockedPiiEntityError,
|
|
GuardrailRaisedException,
|
|
HTTPException,
|
|
) as e:
|
|
# Re-raise guardrail exceptions so they can be properly handled
|
|
raise e
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception(
|
|
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(
|
|
str(e)
|
|
)
|
|
)
|
|
return 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 user_api_key_auth_obj
|
|
|
|
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
|
|
|
|
# Create a synthetic message that represents the tool call
|
|
tool_call_content = (
|
|
f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
|
)
|
|
|
|
synthetic_message = ChatCompletionUserMessage(
|
|
role="user", content=tool_call_content
|
|
)
|
|
|
|
# Create synthetic LLM data that guardrails can process
|
|
synthetic_data = {
|
|
"messages": [synthetic_message],
|
|
"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
|
|
}
|
|
|
|
return synthetic_data
|
|
|
|
def _convert_llm_result_to_mcp_response(
|
|
self, llm_result, request_obj
|
|
) -> Optional[Any]:
|
|
"""
|
|
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 = llm_result.get("messages")
|
|
if modified_messages:
|
|
# Check if content was blocked/modified
|
|
original_content = (
|
|
f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
|
)
|
|
new_content = (
|
|
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(
|
|
f"Could not parse modified arguments from guardrail response: {new_content}"
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(
|
|
f"Error parsing modified arguments: {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
|
|
) -> Optional[dict]:
|
|
"""
|
|
Extract modified/masked arguments from the guardrail response content.
|
|
"""
|
|
import json
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"Extracting modified args from content: {masked_content}"
|
|
)
|
|
|
|
try:
|
|
# The format should be: "Tool: <tool_name>\nArguments: <json_arguments>"
|
|
# Parse the arguments section
|
|
lines = 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(f"Found arguments text: {args_text}")
|
|
|
|
# Try to parse as JSON first
|
|
try:
|
|
modified_args = json.loads(args_text)
|
|
verbose_proxy_logger.debug(
|
|
f"Successfully parsed JSON args: {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(
|
|
f"Failed to parse JSON arguments: {args_text}, error: {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(f"Error extracting modified arguments: {e}")
|
|
return None
|
|
|
|
def _parse_arguments_manually(
|
|
self, args_text: str, original_args: dict
|
|
) -> Optional[dict]:
|
|
"""
|
|
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 = 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(f"Error in manual argument parsing: {e}")
|
|
return None
|
|
|
|
def _convert_llm_result_to_mcp_during_response(
|
|
self, llm_result, request_obj
|
|
) -> Optional[Any]:
|
|
"""
|
|
Convert LLM guardrail result back to MCP during call response format.
|
|
"""
|
|
from litellm.types.mcp import MCPDuringCallResponseObject
|
|
|
|
# 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 = llm_result.get("messages")
|
|
if modified_messages:
|
|
# Check if content was blocked/modified
|
|
original_content = (
|
|
f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
|
|
)
|
|
new_content = (
|
|
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: Optional[List], global_callbacks: List
|
|
) -> List:
|
|
if dynamic_success_callbacks is None:
|
|
return global_callbacks
|
|
return list(set(dynamic_success_callbacks + global_callbacks))
|
|
|
|
def _parse_pre_mcp_call_hook_response(
|
|
self,
|
|
response: MCPPreCallResponseObject,
|
|
original_request: MCPPreCallRequestObject,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
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 = {
|
|
"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
|
|
|
|
async def async_during_mcp_tool_call_hook(
|
|
self,
|
|
kwargs: dict,
|
|
request_obj: Any,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
) -> Optional[Any]:
|
|
"""
|
|
During MCP Tool Call Hook
|
|
|
|
Use this for concurrent monitoring and validation during tool execution.
|
|
Reuses existing LLM guardrail logic by converting MCP calls to message format.
|
|
"""
|
|
from litellm.types.llms.base import HiddenParams
|
|
from litellm.types.mcp import MCPDuringCallRequestObject
|
|
|
|
callbacks = self.get_combined_callback_list(
|
|
dynamic_success_callbacks=getattr(self, "dynamic_success_callbacks", None),
|
|
global_callbacks=litellm.success_callback,
|
|
)
|
|
|
|
# Create the request object if it's not already one
|
|
if not isinstance(request_obj, MCPDuringCallRequestObject):
|
|
request_obj = MCPDuringCallRequestObject(
|
|
tool_name=kwargs.get("name", ""),
|
|
arguments=kwargs.get("arguments", {}),
|
|
server_name=kwargs.get("server_name"),
|
|
start_time=start_time.timestamp() if start_time else None,
|
|
hidden_params=HiddenParams(),
|
|
)
|
|
|
|
for callback in callbacks:
|
|
try:
|
|
_callback: Optional[CustomLogger] = None
|
|
if isinstance(callback, str):
|
|
from typing import cast
|
|
|
|
from litellm import _custom_logger_compatible_callbacks_literal
|
|
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
cast(_custom_logger_compatible_callbacks_literal, callback)
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
|
|
if _callback is not None and isinstance(_callback, CustomGuardrail):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
# Check if guardrail should be run for during_call hook (reusing existing logic)
|
|
if (
|
|
_callback.should_run_guardrail(
|
|
data=kwargs, event_type=GuardrailEventHooks.during_mcp_call
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
# Convert MCP tool call to LLM message format for existing guardrail logic
|
|
synthetic_llm_data = self._convert_mcp_to_llm_format(
|
|
request_obj, kwargs
|
|
)
|
|
|
|
# Reuse existing LLM guardrail logic for during call
|
|
user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(
|
|
kwargs.get("user_api_key_auth")
|
|
)
|
|
|
|
result = await _callback.async_moderation_hook(
|
|
data=synthetic_llm_data,
|
|
user_api_key_dict=user_api_key_auth_dict, # type: ignore
|
|
call_type="mcp_call",
|
|
)
|
|
# Convert result back to MCP response format if blocked/modified
|
|
if result is not None:
|
|
mcp_response = self._convert_llm_result_to_mcp_during_response(
|
|
result, request_obj
|
|
)
|
|
if mcp_response is not None:
|
|
return self._parse_during_mcp_call_hook_response(
|
|
response=mcp_response
|
|
)
|
|
|
|
except Exception as e:
|
|
raise e
|
|
verbose_proxy_logger.exception(
|
|
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {}".format(
|
|
str(e)
|
|
)
|
|
)
|
|
return None
|
|
|
|
def _parse_during_mcp_call_hook_response(
|
|
self, response: MCPDuringCallResponseObject
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Parse the response from the during_mcp_tool_call_hook
|
|
|
|
1. Check if execution should continue
|
|
2. Handle any error messages
|
|
3. Apply any hidden parameter updates
|
|
"""
|
|
result = {
|
|
"should_continue": response.should_continue,
|
|
"error_message": response.error_message,
|
|
"hidden_params": response.hidden_params,
|
|
}
|
|
return result
|
|
|
|
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
|
|
|
|
# The actual implementation of the function
|
|
@overload
|
|
async def pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
data: None,
|
|
call_type: Literal[
|
|
"completion",
|
|
"text_completion",
|
|
"embeddings",
|
|
"image_generation",
|
|
"moderation",
|
|
"audio_transcription",
|
|
"pass_through_endpoint",
|
|
"rerank",
|
|
],
|
|
) -> None:
|
|
pass
|
|
|
|
@overload
|
|
async def pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
data: dict,
|
|
call_type: Literal[
|
|
"completion",
|
|
"text_completion",
|
|
"embeddings",
|
|
"image_generation",
|
|
"moderation",
|
|
"audio_transcription",
|
|
"pass_through_endpoint",
|
|
"rerank",
|
|
],
|
|
) -> dict:
|
|
pass
|
|
|
|
async def pre_call_hook(
|
|
self,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
data: Optional[dict],
|
|
call_type: Literal[
|
|
"completion",
|
|
"text_completion",
|
|
"embeddings",
|
|
"image_generation",
|
|
"moderation",
|
|
"audio_transcription",
|
|
"pass_through_endpoint",
|
|
"rerank",
|
|
],
|
|
) -> Optional[dict]:
|
|
"""
|
|
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
|
|
"""
|
|
from litellm.utils import get_non_default_completion_params
|
|
|
|
verbose_proxy_logger.debug("Inside Proxy Logging Pre-call hook!")
|
|
|
|
self._init_response_taking_too_long_task(data=data)
|
|
|
|
if data is None:
|
|
return None
|
|
|
|
litellm_logging_obj = cast(
|
|
Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)
|
|
)
|
|
prompt_id = data.get("prompt_id", None)
|
|
|
|
## PROMPT TEMPLATE CHECK ##
|
|
if (
|
|
litellm_logging_obj is not None
|
|
and prompt_id is not None
|
|
and (call_type == "completion" or call_type == "acompletion")
|
|
):
|
|
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
|
|
|
custom_logger = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(
|
|
prompt_id
|
|
)
|
|
prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id)
|
|
litellm_prompt_id: Optional[str] = None
|
|
if prompt_spec is not None:
|
|
litellm_prompt_id = prompt_spec.litellm_params.prompt_id
|
|
|
|
if custom_logger and litellm_prompt_id is not None:
|
|
(
|
|
model,
|
|
messages,
|
|
optional_params,
|
|
) = litellm_logging_obj.get_chat_completion_prompt(
|
|
model=data.get("model", ""),
|
|
messages=data.get("messages", []),
|
|
non_default_params=get_non_default_completion_params(kwargs=data),
|
|
prompt_id=litellm_prompt_id,
|
|
prompt_management_logger=custom_logger,
|
|
prompt_variables=data.get("prompt_variables", None),
|
|
prompt_label=data.get("prompt_label", None),
|
|
prompt_version=data.get("prompt_version", None),
|
|
)
|
|
|
|
data.update(optional_params)
|
|
data["model"] = model
|
|
data["messages"] = messages
|
|
|
|
try:
|
|
for callback in litellm.callbacks:
|
|
_callback = None
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
callback
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
if _callback is not None and isinstance(_callback, CustomGuardrail):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
if (
|
|
_callback.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
|
|
response = await _callback.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=self.call_details["user_api_key_cache"],
|
|
data=data, # type: ignore
|
|
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
|
|
)
|
|
|
|
elif (
|
|
_callback is not None
|
|
and isinstance(_callback, CustomLogger)
|
|
and "async_pre_call_hook" in vars(_callback.__class__)
|
|
and _callback.__class__.async_pre_call_hook
|
|
!= CustomLogger.async_pre_call_hook
|
|
):
|
|
response = await _callback.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=self.call_details["user_api_key_cache"],
|
|
data=data, # type: ignore
|
|
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
|
|
)
|
|
|
|
return data
|
|
except Exception as e:
|
|
raise e
|
|
|
|
async def during_call_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
call_type: Literal[
|
|
"completion",
|
|
"responses",
|
|
"embeddings",
|
|
"image_generation",
|
|
"moderation",
|
|
"audio_transcription",
|
|
],
|
|
):
|
|
"""
|
|
Runs the CustomGuardrail's async_moderation_hook()
|
|
"""
|
|
for callback in litellm.callbacks:
|
|
try:
|
|
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": # type: ignore
|
|
return
|
|
else:
|
|
# Main - V2 Guardrails implementation
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
if (
|
|
callback.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.during_call
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
await callback.async_moderation_hook(
|
|
data=data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
call_type=call_type,
|
|
)
|
|
except Exception as e:
|
|
raise e
|
|
return 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",
|
|
"team_budget",
|
|
"proxy_budget",
|
|
"projected_limit_exceeded",
|
|
],
|
|
user_info: CallInfo,
|
|
):
|
|
if self.alerting is None:
|
|
# do nothing if alerting is not switched on
|
|
return
|
|
await self.slack_alerting_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: Optional[dict] = 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 = datetime.now().strftime("%H:%M:%S")
|
|
_proxy_base_url = 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 = {}
|
|
alerting_metadata = {}
|
|
if request_data is not None:
|
|
_url = await _add_langfuse_trace_id_to_alert(request_data=request_data)
|
|
|
|
if _url is not None:
|
|
extra_kwargs["🪢 Langfuse Trace"] = _url
|
|
formatted_message += "\n\n🪢 Langfuse Trace: {}".format(_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"]
|
|
for client in self.alerting:
|
|
if client == "slack":
|
|
await self.slack_alerting_instance.send_alert(
|
|
message=message,
|
|
level=level,
|
|
alert_type=alert_type,
|
|
user_info=None,
|
|
alerting_metadata=alerting_metadata,
|
|
**extra_kwargs,
|
|
)
|
|
elif 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]
|
|
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: Optional[ProxyErrorTypes] = None,
|
|
route: Optional[str] = None,
|
|
traceback_str: Optional[str] = None,
|
|
):
|
|
"""
|
|
Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body.
|
|
|
|
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
|
|
"""
|
|
|
|
### 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 isinstance(
|
|
original_exception, HTTPException
|
|
):
|
|
"""
|
|
Just alert on LLM API exceptions. Do not alert on user errors
|
|
|
|
Related issue - https://github.com/BerriAI/litellm/issues/3395
|
|
"""
|
|
litellm_debug_info = 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=f"LLM API call failed: `{exception_str}`",
|
|
level="High",
|
|
alert_type=AlertType.llm_exceptions,
|
|
request_data=request_data,
|
|
)
|
|
)
|
|
|
|
### 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,
|
|
)
|
|
|
|
for callback in litellm.callbacks:
|
|
try:
|
|
_callback: Optional[CustomLogger] = None
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
callback
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
if _callback is not None and isinstance(_callback, CustomLogger):
|
|
asyncio.create_task(
|
|
_callback.async_post_call_failure_hook(
|
|
request_data=request_data,
|
|
user_api_key_dict=user_api_key_dict,
|
|
original_exception=original_exception,
|
|
traceback_str=traceback_str,
|
|
)
|
|
)
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception(
|
|
f"[Non-Blocking] Error in post_call_failure_hook: {e}"
|
|
)
|
|
return
|
|
|
|
def _is_proxy_only_llm_api_error(
|
|
self,
|
|
original_exception: Exception,
|
|
error_type: Optional[ProxyErrorTypes] = None,
|
|
route: Optional[str] = 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)
|
|
"""
|
|
|
|
#########################################################
|
|
# Only log LLM API 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 RouteChecks.is_llm_api_route(route) is not True:
|
|
return False
|
|
|
|
return isinstance(original_exception, HTTPException) or (
|
|
error_type == ProxyErrorTypes.auth_error
|
|
)
|
|
|
|
async def _handle_logging_proxy_only_error(
|
|
self,
|
|
request_data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
route: Optional[str] = None,
|
|
original_exception: Optional[Exception] = 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: Optional[Logging] = request_data.get(
|
|
"litellm_logging_obj", None
|
|
)
|
|
if litellm_logging_obj is None:
|
|
import uuid
|
|
|
|
request_data["litellm_call_id"] = str(uuid.uuid4())
|
|
user_api_key_logged_metadata = (
|
|
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,
|
|
)
|
|
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 = {}
|
|
_litellm_params = {}
|
|
|
|
litellm_param_keys = LoggedLiteLLMParams.__annotations__.keys()
|
|
for k, v in request_data.items():
|
|
if k in litellm_param_keys:
|
|
_litellm_params[k] = v
|
|
elif k != "model" and k != "user":
|
|
_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: Union[list, str, dict] = ""
|
|
if "messages" in request_data and isinstance(
|
|
request_data["messages"], list
|
|
):
|
|
input = request_data["messages"]
|
|
litellm_logging_obj.model_call_details["messages"] = input
|
|
litellm_logging_obj.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
|
|
litellm_logging_obj.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
|
|
litellm_logging_obj.call_type = CallTypes.aembedding.value
|
|
litellm_logging_obj.pre_call(
|
|
input=input,
|
|
api_key="",
|
|
)
|
|
|
|
# log the custom exception
|
|
await litellm_logging_obj.async_failure_handler(
|
|
exception=original_exception,
|
|
traceback_exception=traceback.format_exc(),
|
|
)
|
|
|
|
threading.Thread(
|
|
target=litellm_logging_obj.failure_handler,
|
|
args=(
|
|
original_exception,
|
|
traceback.format_exc(),
|
|
),
|
|
).start()
|
|
|
|
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
|
|
"""
|
|
|
|
for callback in litellm.callbacks:
|
|
try:
|
|
_callback: Optional[CustomLogger] = None
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
callback
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
|
|
if _callback is not None:
|
|
############## Handle Guardrails ########################################
|
|
#############################################################################
|
|
if isinstance(callback, CustomGuardrail):
|
|
# Main - V2 Guardrails implementation
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
if (
|
|
callback.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.post_call
|
|
)
|
|
is not True
|
|
):
|
|
continue
|
|
|
|
await callback.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
)
|
|
|
|
############ Handle CustomLogger ###############################
|
|
#################################################################
|
|
elif isinstance(_callback, CustomLogger):
|
|
await _callback.async_post_call_success_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
data=data,
|
|
response=response,
|
|
)
|
|
except Exception as e:
|
|
raise e
|
|
return response
|
|
|
|
async def async_post_call_streaming_hook(
|
|
self,
|
|
data: dict,
|
|
response: Union[
|
|
ModelResponse, EmbeddingResponse, ImageResponse, ModelResponseStream
|
|
],
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
str_so_far: Optional[str] = None,
|
|
):
|
|
"""
|
|
Allow user to modify outgoing streaming data -> per chunk
|
|
|
|
Covers:
|
|
1. /chat/completions
|
|
"""
|
|
from litellm.proxy.proxy_server import llm_router
|
|
|
|
response_str: Optional[str] = None
|
|
if isinstance(response, (ModelResponse, ModelResponseStream)):
|
|
response_str = litellm.get_response_string(response_obj=response)
|
|
if response_str is not None:
|
|
for callback in litellm.callbacks:
|
|
try:
|
|
_callback: Optional[CustomLogger] = None
|
|
if isinstance(callback, CustomGuardrail):
|
|
# Main - V2 Guardrails implementation
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
## CHECK FOR MODEL-LEVEL GUARDRAILS
|
|
modified_data = _check_and_merge_model_level_guardrails(
|
|
data=data, llm_router=llm_router
|
|
)
|
|
|
|
if (
|
|
callback.should_run_guardrail(
|
|
data=modified_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(
|
|
callback
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
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
|
|
potential_error_response = (
|
|
await _callback.async_post_call_streaming_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=complete_response,
|
|
)
|
|
)
|
|
if isinstance(
|
|
potential_error_response, str
|
|
) and potential_error_response.startswith("data: "):
|
|
return potential_error_response
|
|
except Exception as e:
|
|
raise e
|
|
return response
|
|
|
|
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
|
|
"""
|
|
for callback in litellm.callbacks:
|
|
_callback: Optional[CustomLogger] = None
|
|
if isinstance(callback, str):
|
|
_callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
|
|
callback
|
|
)
|
|
else:
|
|
_callback = callback # type: ignore
|
|
if _callback is not None and isinstance(_callback, CustomLogger):
|
|
if not isinstance(
|
|
_callback, CustomGuardrail
|
|
) or _callback.should_run_guardrail(
|
|
data=request_data, event_type=GuardrailEventHooks.post_call
|
|
):
|
|
response = _callback.async_post_call_streaming_iterator_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
response=response,
|
|
request_data=request_data,
|
|
)
|
|
return response
|
|
|
|
def _init_response_taking_too_long_task(self, data: Optional[dict] = 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 = 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
|
|
|
|
|
|
class PrismaClient:
|
|
spend_log_transactions: List = []
|
|
|
|
def __init__(
|
|
self,
|
|
database_url: str,
|
|
proxy_logging_obj: ProxyLogging,
|
|
http_client: Optional[Any] = None,
|
|
):
|
|
## init logging object
|
|
self.proxy_logging_obj = proxy_logging_obj
|
|
self.iam_token_db_auth: Optional[bool] = str_to_bool(
|
|
os.getenv("IAM_TOKEN_DB_AUTH")
|
|
)
|
|
verbose_proxy_logger.debug("Creating Prisma Client..")
|
|
try:
|
|
from prisma import Prisma # type: ignore
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"Failed to import Prisma client: {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."
|
|
)
|
|
if http_client is not None:
|
|
self.db = PrismaWrapper(
|
|
original_prisma=Prisma(http=http_client),
|
|
iam_token_db_auth=(
|
|
self.iam_token_db_auth
|
|
if self.iam_token_db_auth is not None
|
|
else False
|
|
),
|
|
)
|
|
else:
|
|
self.db = PrismaWrapper(
|
|
original_prisma=Prisma(),
|
|
iam_token_db_auth=(
|
|
self.iam_token_db_auth
|
|
if self.iam_token_db_auth is not None
|
|
else False
|
|
),
|
|
) # Client to connect to Prisma db
|
|
verbose_proxy_logger.debug("Success - Created Prisma Client")
|
|
|
|
def get_request_status(
|
|
self, payload: Union[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: Union[Dict, SpendLogsMetadata, str] = payload.get(
|
|
"metadata", {}
|
|
)
|
|
if isinstance(payload_metadata, str):
|
|
payload_metadata_json: Union[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 = hashlib.sha256(token.encode()).hexdigest()
|
|
|
|
return hashed_token
|
|
|
|
def jsonify_object(self, data: dict) -> dict:
|
|
db_data = 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
|
|
|
|
@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 = [
|
|
"LiteLLM_VerificationTokenView",
|
|
"MonthlyGlobalSpend",
|
|
"Last30dKeysBySpend",
|
|
"Last30dModelsBySpend",
|
|
"MonthlyGlobalSpendPerKey",
|
|
"MonthlyGlobalSpendPerUserPerKey",
|
|
"Last30dTopEndUsersSpend",
|
|
"DailyTagSpend",
|
|
]
|
|
required_view = "LiteLLM_VerificationTokenView"
|
|
expected_views_str = ", ".join(f"'{view}'" for view in expected_views)
|
|
pg_schema = os.getenv("DATABASE_SCHEMA", "public")
|
|
ret = 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 = 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 self.db.execute_raw(
|
|
"""
|
|
CREATE VIEW "LiteLLM_VerificationTokenView" AS
|
|
SELECT
|
|
v.*,
|
|
t.spend AS team_spend,
|
|
t.max_budget AS team_max_budget,
|
|
t.tpm_limit AS team_tpm_limit,
|
|
t.rpm_limit AS team_rpm_limit
|
|
FROM "LiteLLM_VerificationToken" v
|
|
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
|
|
"""
|
|
)
|
|
|
|
verbose_proxy_logger.info(
|
|
"LiteLLM_VerificationTokenView Created in DB!"
|
|
)
|
|
else:
|
|
should_create_views = 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 = (
|
|
set(ret[0]["view_names"]) if ret[0]["view_names"] else set()
|
|
)
|
|
expected_views_set = set(expected_views)
|
|
# Find missing views
|
|
missing_views = expected_views_set - ret_view_names_set
|
|
|
|
verbose_proxy_logger.warning(
|
|
"\n\n\033[93mNot all views exist in db, needed for UI 'Usage' tab. Missing={}.\nRun 'create_views.py' from https://github.com/BerriAI/litellm/tree/main/db_scripts to create missing views.\033[0m\n".format(
|
|
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: Any,
|
|
table_name: Literal["users", "keys", "config", "spend"],
|
|
):
|
|
"""
|
|
Generic implementation of get data
|
|
"""
|
|
start_time = time.time()
|
|
try:
|
|
if table_name == "users":
|
|
response = await self.db.litellm_usertable.find_first(
|
|
where={key: value} # type: ignore
|
|
)
|
|
elif table_name == "keys":
|
|
response = await self.db.litellm_verificationtoken.find_first( # type: ignore
|
|
where={key: value} # type: ignore
|
|
)
|
|
elif table_name == "config":
|
|
response = await self.db.litellm_config.find_first( # type: ignore
|
|
where={key: value} # type: ignore
|
|
)
|
|
elif table_name == "spend":
|
|
response = await self.db.l.find_first( # type: ignore
|
|
where={key: value} # type: ignore
|
|
)
|
|
return response
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}"
|
|
verbose_proxy_logger.error(error_msg)
|
|
error_msg = error_msg + "\nException Type: {}".format(type(e))
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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
|
|
|
|
@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( # noqa: PLR0915
|
|
self,
|
|
token: Optional[Union[str, list]] = None,
|
|
user_id: Optional[str] = None,
|
|
user_id_list: Optional[list] = None,
|
|
team_id: Optional[str] = None,
|
|
team_id_list: Optional[list] = None,
|
|
key_val: Optional[dict] = None,
|
|
table_name: Optional[
|
|
Literal[
|
|
"user",
|
|
"key",
|
|
"config",
|
|
"spend",
|
|
"enduser",
|
|
"budget",
|
|
"team",
|
|
"user_notification",
|
|
"combined_view",
|
|
]
|
|
] = None,
|
|
query_type: Literal["find_unique", "find_all"] = "find_unique",
|
|
expires: Optional[datetime] = None,
|
|
reset_at: Optional[datetime] = None,
|
|
offset: Optional[int] = None, # pagination, what row number to start from
|
|
limit: Optional[
|
|
int
|
|
] = None, # pagination, number of rows to getch when find_all==True
|
|
parent_otel_span: Optional[Span] = None,
|
|
proxy_logging_obj: Optional[ProxyLogging] = None,
|
|
budget_id_list: Optional[List[str]] = None,
|
|
):
|
|
args_passed_in = locals()
|
|
start_time = time.time()
|
|
hashed_token: Optional[str] = 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(
|
|
f"PrismaClient: find_unique for token: {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 self.db.litellm_verificationtoken.find_unique(
|
|
where={"token": hashed_token}, # type: ignore
|
|
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 self.db.litellm_verificationtoken.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 self.db.litellm_verificationtoken.find_many(
|
|
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 self.db.litellm_verificationtoken.find_many(
|
|
where={ # type:ignore
|
|
"OR": [
|
|
{"expires": None},
|
|
{"expires": {"gt": expires}},
|
|
],
|
|
"budget_reset_at": {"lt": reset_at},
|
|
}
|
|
)
|
|
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: dict = {}
|
|
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 = []
|
|
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 self.db.litellm_verificationtoken.find_many(
|
|
order={"spend": "desc"},
|
|
where=where_filter, # type: ignore
|
|
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 self.db.litellm_usertable.find_unique( # type: ignore
|
|
where=key_val, # type: ignore
|
|
include={"organization_memberships": True},
|
|
)
|
|
|
|
elif query_type == "find_all" and key_val is not None:
|
|
response = await self.db.litellm_usertable.find_many(
|
|
where=key_val # type: ignore
|
|
) # type: ignore
|
|
elif query_type == "find_all" and reset_at is not None:
|
|
response = await self.db.litellm_usertable.find_many(
|
|
where={ # type:ignore
|
|
"budget_reset_at": {"lt": reset_at},
|
|
}
|
|
)
|
|
elif query_type == "find_all" and user_id_list is not None:
|
|
response = await self.db.litellm_usertable.find_many(
|
|
where={"user_id": {"in": user_id_list}}
|
|
)
|
|
elif query_type == "find_all":
|
|
if expires is not None:
|
|
response = await self.db.litellm_usertable.find_many( # type: ignore
|
|
order={"spend": "desc"},
|
|
where={ # type:ignore
|
|
"OR": [
|
|
{"expires": None}, # type:ignore
|
|
{"expires": {"gt": expires}}, # type:ignore
|
|
],
|
|
},
|
|
)
|
|
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 self.db.litellm_spendlogs.find_unique( # type: ignore
|
|
where={ # type: ignore
|
|
key_val["key"]: key_val["value"], # type: ignore
|
|
}
|
|
)
|
|
elif query_type == "find_all":
|
|
response = await self.db.litellm_spendlogs.find_many( # type: ignore
|
|
where={
|
|
key_val["key"]: key_val["value"], # type: ignore
|
|
}
|
|
)
|
|
return response
|
|
else:
|
|
response = await self.db.litellm_spendlogs.find_many( # type: ignore
|
|
order={"startTime": "desc"},
|
|
)
|
|
return response
|
|
elif table_name == "budget" and reset_at is not None:
|
|
if query_type == "find_all":
|
|
response = await self.db.litellm_budgettable.find_many(
|
|
where={ # type:ignore
|
|
"OR": [
|
|
{
|
|
"AND": [
|
|
{"budget_reset_at": None},
|
|
{"NOT": {"budget_duration": 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 self.db.litellm_endusertable.find_many(
|
|
where={"budget_id": {"in": budget_id_list}}
|
|
)
|
|
return response
|
|
elif table_name == "team":
|
|
if query_type == "find_unique":
|
|
response = await self.db.litellm_teamtable.find_unique(
|
|
where={"team_id": team_id}, # type: ignore
|
|
include={"litellm_model_table": True}, # type: ignore
|
|
)
|
|
elif query_type == "find_all" and reset_at is not None:
|
|
response = await self.db.litellm_teamtable.find_many(
|
|
where={ # type:ignore
|
|
"budget_reset_at": {"lt": reset_at},
|
|
}
|
|
)
|
|
elif query_type == "find_all" and user_id is not None:
|
|
response = await self.db.litellm_teamtable.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 self.db.litellm_teamtable.find_many(
|
|
where={"team_id": {"in": team_id_list}}
|
|
)
|
|
elif query_type == "find_all" and team_id_list is None:
|
|
response = await self.db.litellm_teamtable.find_many(
|
|
take=MAX_TEAM_LIST_LIMIT
|
|
)
|
|
return response
|
|
elif table_name == "user_notification":
|
|
if query_type == "find_unique":
|
|
response = await self.db.litellm_usernotifications.find_unique( # type: ignore
|
|
where={"user_id": user_id} # type: ignore
|
|
)
|
|
elif query_type == "find_all":
|
|
response = await self.db.litellm_usernotifications.find_many() # type: ignore
|
|
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(
|
|
f"PrismaClient: find_unique for token: {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 = f"""
|
|
SELECT
|
|
v.*,
|
|
t.spend AS team_spend,
|
|
t.max_budget AS team_max_budget,
|
|
t.tpm_limit AS team_tpm_limit,
|
|
t.rpm_limit AS team_rpm_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.organization_id as org_id,
|
|
tm.spend AS team_member_spend,
|
|
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.model_max_budget as litellm_budget_table_model_max_budget,
|
|
b.soft_budget as litellm_budget_table_soft_budget
|
|
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_ModelTable" m ON t.model_id = m.id
|
|
LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id
|
|
WHERE v.token = '{token}'
|
|
"""
|
|
|
|
response = await self.db.query_first(query=sql_query)
|
|
|
|
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: Optional[Member] = 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 = f"LiteLLM Prisma Client Exception: Error with `get_data`. Args passed in: {args_passed_in}"
|
|
error_msg = prisma_query_info + str(e)
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
verbose_proxy_logger.debug(error_traceback)
|
|
end_time = time.time()
|
|
_duration = 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: dict):
|
|
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"])
|
|
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( # noqa: PLR0915
|
|
self,
|
|
data: dict,
|
|
table_name: Literal[
|
|
"user", "key", "config", "spend", "team", "user_notification"
|
|
],
|
|
):
|
|
"""
|
|
Add a key to the database. If it already exists, do nothing.
|
|
"""
|
|
start_time = time.time()
|
|
try:
|
|
verbose_proxy_logger.debug("PrismaClient: insert_data: %s", data)
|
|
if table_name == "key":
|
|
token = data["token"]
|
|
hashed_token = self.hash_token(token=token)
|
|
db_data = self.jsonify_object(data=data)
|
|
db_data["token"] = hashed_token
|
|
print_verbose(
|
|
"PrismaClient: Before upsert into litellm_verificationtoken"
|
|
)
|
|
new_verification_token = await self.db.litellm_verificationtoken.upsert( # type: ignore
|
|
where={
|
|
"token": hashed_token,
|
|
},
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"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 = await self.db.litellm_usertable.upsert(
|
|
where={"user_id": data["user_id"]},
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"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 = await self.db.litellm_teamtable.upsert(
|
|
where={"team_id": data["team_id"]},
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"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 = []
|
|
for k, v in data.items():
|
|
updated_data = v
|
|
updated_data = json.dumps(updated_data)
|
|
updated_table_row = self.db.litellm_config.upsert(
|
|
where={"param_name": k}, # type: ignore
|
|
data={
|
|
"create": {"param_name": k, "param_value": updated_data}, # type: ignore
|
|
"update": {"param_value": updated_data},
|
|
},
|
|
)
|
|
|
|
tasks.append(updated_table_row)
|
|
await asyncio.gather(*tasks)
|
|
verbose_proxy_logger.info("Data Inserted into Config Table")
|
|
elif table_name == "spend":
|
|
db_data = self.jsonify_object(data=data)
|
|
new_spend_row = await self.db.litellm_spendlogs.upsert(
|
|
where={"request_id": data["request_id"]},
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"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 = (
|
|
await self.db.litellm_usernotifications.upsert( # type: ignore
|
|
where={"request_id": data["request_id"]},
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"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 = f"LiteLLM Prisma Client Exception in insert_data: {str(e)}"
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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( # noqa: PLR0915
|
|
self,
|
|
token: Optional[str] = None,
|
|
data: dict = {},
|
|
data_list: Optional[List] = None,
|
|
user_id: Optional[str] = None,
|
|
team_id: Optional[str] = None,
|
|
query_type: Literal["update", "update_many"] = "update",
|
|
table_name: Optional[
|
|
Literal["user", "key", "config", "spend", "team", "enduser", "budget"]
|
|
] = None,
|
|
update_key_values: Optional[dict] = None,
|
|
update_key_values_custom_query: Optional[dict] = None,
|
|
):
|
|
"""
|
|
Update existing data
|
|
"""
|
|
verbose_proxy_logger.debug(
|
|
f"PrismaClient: update_data, table_name: {table_name}"
|
|
)
|
|
start_time = time.time()
|
|
try:
|
|
db_data = 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: {token}")
|
|
# check if plain text or hash
|
|
token = _hash_token_if_needed(token=token)
|
|
db_data["token"] = token
|
|
response = await self.db.litellm_verificationtoken.update(
|
|
where={"token": token}, # type: ignore
|
|
data={**db_data}, # type: ignore
|
|
)
|
|
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() # type: ignore
|
|
except Exception:
|
|
_data = response.dict()
|
|
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 = db_data["user_id"]
|
|
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 = await self.db.litellm_usertable.upsert(
|
|
where={"user_id": user_id}, # type: ignore
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"update": {
|
|
**update_key_values # type: ignore
|
|
}, # 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 = db_data["team_id"]
|
|
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 = await self.db.litellm_teamtable.upsert(
|
|
where={"team_id": team_id}, # type: ignore
|
|
data={
|
|
"create": {**db_data}, # type: ignore
|
|
"update": {
|
|
**update_key_values # type: ignore
|
|
}, # 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-"): # type: ignore
|
|
t.token = self.hash_token(token=t.token) # type: ignore
|
|
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}, # type: ignore
|
|
data={**data_json}, # type: ignore
|
|
)
|
|
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}, # type: ignore
|
|
data={
|
|
"create": {**data_json}, # type: ignore
|
|
"update": {
|
|
**data_json # type: ignore
|
|
}, # 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}, # type: ignore
|
|
data={
|
|
"create": {**data_json}, # type: ignore
|
|
"update": {
|
|
**data_json # type: ignore
|
|
}, # 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}, # type: ignore
|
|
data={
|
|
"create": {**data_json}, # type: ignore
|
|
"update": {
|
|
**data_json # type: ignore
|
|
}, # 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}, # type: ignore
|
|
data={
|
|
"create": {**data_json}, # type: ignore
|
|
"update": {
|
|
**data_json # type: ignore
|
|
}, # 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 = f"LiteLLM Prisma Client Exception - update_data: {str(e)}"
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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: Optional[List] = None,
|
|
team_id_list: Optional[List] = None,
|
|
table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None,
|
|
user_id: Optional[str] = None,
|
|
):
|
|
"""
|
|
Allow user to delete a key(s)
|
|
|
|
Ensure user owns that key, unless admin.
|
|
"""
|
|
start_time = time.time()
|
|
try:
|
|
if tokens is not None and isinstance(tokens, List):
|
|
hashed_tokens = []
|
|
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 = {}
|
|
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 = await self.db.litellm_verificationtoken.delete_many(
|
|
where=filter_query # type: ignore
|
|
)
|
|
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 self.db.litellm_teamtable.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 self.db.litellm_verificationtoken.delete_many(
|
|
where={"team_id": {"in": team_id_list}}
|
|
)
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg = f"LiteLLM Prisma Client Exception - delete_data: {str(e)}"
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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 = 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 = f"LiteLLM Prisma Client Exception connect(): {str(e)}"
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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 = time.time()
|
|
try:
|
|
await self.db.disconnect()
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg = f"LiteLLM Prisma Client Exception disconnect(): {str(e)}"
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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
|
|
|
|
async def health_check(self):
|
|
"""
|
|
Health check endpoint for the prisma client
|
|
"""
|
|
start_time = time.time()
|
|
try:
|
|
sql_query = "SELECT 1"
|
|
|
|
# Execute the raw query
|
|
# The asterisk before `user_id_list` unpacks the list into separate arguments
|
|
response = await self.db.query_raw(sql_query)
|
|
return response
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg = f"LiteLLM Prisma Client Exception disconnect(): {str(e)}"
|
|
print_verbose(error_msg)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = end_time - start_time
|
|
asyncio.create_task(
|
|
self.proxy_logging_obj.failure_handler(
|
|
original_exception=e,
|
|
duration=_duration,
|
|
call_type="health_check",
|
|
traceback_str=error_traceback,
|
|
)
|
|
)
|
|
raise e
|
|
|
|
async def _get_spend_logs_row_count(self) -> int:
|
|
try:
|
|
sql_query = """
|
|
SELECT reltuples::BIGINT
|
|
FROM pg_class
|
|
WHERE oid = '"LiteLLM_SpendLogs"'::regclass;
|
|
"""
|
|
result = await self.db.query_raw(query=sql_query)
|
|
return result[0]["reltuples"]
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(
|
|
f"Error getting LiteLLM_SpendLogs row count: {e}"
|
|
)
|
|
return 0
|
|
|
|
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 = 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: Optional[float]
|
|
) -> Optional[float]:
|
|
"""Validate and clean response time value"""
|
|
if response_time_ms is None:
|
|
return None
|
|
try:
|
|
value = float(response_time_ms)
|
|
return (
|
|
value
|
|
if value == value and value not in (float("inf"), float("-inf"))
|
|
else None
|
|
)
|
|
except (ValueError, TypeError):
|
|
verbose_proxy_logger.warning(
|
|
f"Invalid response_time_ms value: {response_time_ms}"
|
|
)
|
|
return None
|
|
|
|
def _clean_details(self, details: Optional[dict]) -> Optional[dict]:
|
|
"""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(f"Failed to clean details JSON: {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: Optional[str] = None,
|
|
response_time_ms: Optional[float] = None,
|
|
details: Optional[dict] = None,
|
|
checked_by: Optional[str] = None,
|
|
model_id: Optional[str] = None,
|
|
):
|
|
"""Save health check result to database"""
|
|
try:
|
|
# Build base data with required fields
|
|
health_check_data = {
|
|
"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 = {
|
|
"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(f"Saving health check data: {health_check_data}")
|
|
return await self.db.litellm_healthchecktable.create(data=health_check_data)
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(
|
|
f"Error saving health check result for model {model_name}: {e}"
|
|
)
|
|
return None
|
|
|
|
async def get_health_check_history(
|
|
self,
|
|
model_name: Optional[str] = None,
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
status_filter: Optional[str] = None,
|
|
):
|
|
"""
|
|
Get health check history with optional filtering
|
|
"""
|
|
try:
|
|
where_clause = {}
|
|
if model_name:
|
|
where_clause["model_name"] = model_name
|
|
if status_filter:
|
|
where_clause["status"] = status_filter
|
|
|
|
results = await self.db.litellm_healthchecktable.find_many(
|
|
where=where_clause,
|
|
order={"checked_at": "desc"},
|
|
take=limit,
|
|
skip=offset,
|
|
)
|
|
return results
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"Error getting health check history: {e}")
|
|
return []
|
|
|
|
async def get_all_latest_health_checks(self):
|
|
"""
|
|
Get the latest health check for each model
|
|
"""
|
|
try:
|
|
# Get all unique model names first
|
|
all_checks = await self.db.litellm_healthchecktable.find_many(
|
|
order={"checked_at": "desc"}
|
|
)
|
|
|
|
# Group by model_name and get the latest for each
|
|
latest_checks = {}
|
|
for check in all_checks:
|
|
if check.model_name not in latest_checks:
|
|
latest_checks[check.model_name] = check
|
|
|
|
return list(latest_checks.values())
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"Error getting all latest health checks: {e}")
|
|
return []
|
|
|
|
|
|
### 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 = f"{user_id}_user_api_key_user_id"
|
|
response = cache.get_cache(key=cache_key)
|
|
if response is None: # Cache miss
|
|
user_row = 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 = user_row.model_dump_json()
|
|
cache.set_cache(
|
|
key=cache_key, value=cache_value, ttl=600
|
|
) # store for 10 minutes
|
|
return
|
|
|
|
|
|
async def send_email(
|
|
receiver_email: Optional[str] = None,
|
|
subject: Optional[str] = None,
|
|
html: Optional[str] = None,
|
|
):
|
|
"""
|
|
smtp_host,
|
|
smtp_port,
|
|
smtp_username,
|
|
smtp_password,
|
|
sender_name,
|
|
sender_email,
|
|
"""
|
|
## SERVER SETUP ##
|
|
|
|
smtp_host = os.getenv("SMTP_HOST")
|
|
smtp_port = int(os.getenv("SMTP_PORT", "587")) # default to port 587
|
|
smtp_username = os.getenv("SMTP_USERNAME")
|
|
smtp_password = os.getenv("SMTP_PASSWORD")
|
|
sender_email = 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 = 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:
|
|
# Establish a secure connection with the SMTP server
|
|
with smtplib.SMTP(
|
|
host=smtp_host,
|
|
port=smtp_port,
|
|
) as server:
|
|
if os.getenv("SMTP_TLS", "True") != "False":
|
|
server.starttls()
|
|
|
|
# Login to your email account only if smtp_username and smtp_password are provided
|
|
if smtp_username and smtp_password:
|
|
server.login(
|
|
user=smtp_username,
|
|
password=smtp_password,
|
|
)
|
|
|
|
# Send the email
|
|
server.send_message(
|
|
msg=email_message,
|
|
from_addr=sender_email,
|
|
to_addrs=receiver_email,
|
|
)
|
|
|
|
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 = hashlib.sha256(token.encode()).hexdigest()
|
|
|
|
return hashed_token
|
|
|
|
|
|
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
|
|
|
|
|
|
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:
|
|
async with transaction.batch_() as batcher:
|
|
for (
|
|
end_user_id,
|
|
response_cost,
|
|
) in 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 DB_CONNECTION_ERROR_TYPES as e:
|
|
if i >= n_retry_times: # If we've reached the maximum number of retries
|
|
_raise_failed_update_spend_exception(
|
|
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
# Optionally, sleep for a bit before retrying
|
|
await asyncio.sleep(2**i) # Exponential backoff
|
|
except Exception as e:
|
|
_raise_failed_update_spend_exception(
|
|
e=e, 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: Optional[HTTPHandler],
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
BATCH_SIZE = 100 # Preferred size of each batch to write to the database
|
|
MAX_LOGS_PER_INTERVAL = (
|
|
1000 # Maximum number of logs to flush in a single interval
|
|
)
|
|
# Get initial logs to process
|
|
logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL]
|
|
start_time = 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: {}".format(base_url))
|
|
response = await db_writer_client.post(
|
|
url=base_url + "spend/update",
|
|
data=json.dumps(logs_to_process),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
if response.status_code == 200:
|
|
prisma_client.spend_log_transactions = (
|
|
prisma_client.spend_log_transactions[
|
|
len(logs_to_process) :
|
|
]
|
|
)
|
|
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
|
|
]
|
|
await prisma_client.db.litellm_spendlogs.create_many(
|
|
data=batch_with_dates, skip_duplicates=True
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"Flushed {len(batch)} logs to the DB."
|
|
)
|
|
|
|
prisma_client.spend_log_transactions = (
|
|
prisma_client.spend_log_transactions[len(logs_to_process) :]
|
|
)
|
|
verbose_proxy_logger.debug(
|
|
f"{len(logs_to_process)} logs processed. Remaining in queue: {len(prisma_client.spend_log_transactions)}"
|
|
)
|
|
break
|
|
except DB_CONNECTION_ERROR_TYPES:
|
|
if i is None:
|
|
i = 0
|
|
if i >= n_retry_times:
|
|
raise
|
|
await asyncio.sleep(2**i)
|
|
except Exception as e:
|
|
prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[
|
|
len(logs_to_process) :
|
|
]
|
|
_raise_failed_update_spend_exception(
|
|
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
|
|
)
|
|
|
|
@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( # noqa: PLR0915
|
|
prisma_client: PrismaClient,
|
|
db_writer_client: Optional[HTTPHandler],
|
|
proxy_logging_obj: ProxyLogging,
|
|
):
|
|
"""
|
|
Batch write updates to db.
|
|
|
|
Triggered every minute.
|
|
|
|
Requires:
|
|
user_id_list: dict,
|
|
keys_list: list,
|
|
team_list: list,
|
|
spend_logs: list,
|
|
"""
|
|
n_retry_times = 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 ###
|
|
verbose_proxy_logger.debug(
|
|
"Spend Logs transactions: {}".format(len(prisma_client.spend_log_transactions))
|
|
)
|
|
|
|
if len(prisma_client.spend_log_transactions) > 0:
|
|
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,
|
|
)
|
|
|
|
|
|
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 = (
|
|
f"[Non-Blocking]LiteLLM Prisma Client Exception - update spend logs: {str(e)}"
|
|
)
|
|
error_traceback = error_msg + "\n" + traceback.format_exc()
|
|
end_time = time.time()
|
|
_duration = 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 _is_projected_spend_over_limit(
|
|
current_spend: float, soft_budget_limit: Optional[float]
|
|
):
|
|
from datetime import date
|
|
|
|
if soft_budget_limit is None:
|
|
# If there's no limit, we can't exceed it.
|
|
return False
|
|
|
|
today = date.today()
|
|
|
|
# Finding the first day of the next month, then subtracting one day to get the end of the current month.
|
|
if today.month == 12: # December edge case
|
|
end_month = date(today.year + 1, 1, 1) - timedelta(days=1)
|
|
else:
|
|
end_month = date(today.year, today.month + 1, 1) - timedelta(days=1)
|
|
|
|
remaining_days = (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 = 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: Optional[float]
|
|
) -> Optional[tuple]:
|
|
import datetime
|
|
|
|
if soft_budget_limit is None:
|
|
return None
|
|
|
|
today = datetime.date.today()
|
|
end_month = datetime.date(today.year, today.month + 1, 1) - datetime.timedelta(
|
|
days=1
|
|
)
|
|
remaining_days = (end_month - today).days
|
|
|
|
daily_spend = current_spend / (
|
|
today.day - 1
|
|
) # assuming the current spend till today (not including today)
|
|
projected_spend = daily_spend * remaining_days
|
|
|
|
if projected_spend > soft_budget_limit:
|
|
approx_days = soft_budget_limit / daily_spend
|
|
limit_exceed_date = today + datetime.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 = team_config.pop("models")
|
|
model_in_request = 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: Optional[Router]
|
|
) -> 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
|
|
|
|
Returns:
|
|
Modified data dict with merged guardrails (if any model-level guardrails exist)
|
|
"""
|
|
if llm_router is None:
|
|
return data
|
|
|
|
# Get the model ID from the data
|
|
metadata = data.get("metadata") or {}
|
|
model_info = metadata.get("model_info") or {}
|
|
model_id = model_info.get("id", None)
|
|
|
|
if model_id is None:
|
|
return data
|
|
|
|
# Check if the model has guardrails
|
|
deployment = llm_router.get_deployment(model_id=model_id)
|
|
if deployment is None:
|
|
return data
|
|
|
|
model_level_guardrails = deployment.litellm_params.get("guardrails")
|
|
|
|
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: Any) -> 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 = data.copy()
|
|
metadata = 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 = 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() -> Optional[str]:
|
|
"""
|
|
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() -> Optional[str]:
|
|
"""
|
|
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 handle_exception_on_proxy(e: Exception) -> ProxyException:
|
|
"""
|
|
Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible
|
|
"""
|
|
from fastapi import status
|
|
|
|
verbose_proxy_logger.exception(f"Exception: {e}")
|
|
|
|
if isinstance(e, HTTPException):
|
|
return ProxyException(
|
|
message=getattr(e, "detail", f"error({str(e)})"),
|
|
type=ProxyErrorTypes.internal_server_error,
|
|
param=getattr(e, "param", "None"),
|
|
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
|
)
|
|
elif isinstance(e, ProxyException):
|
|
return e
|
|
return ProxyException(
|
|
message="Internal Server Error, " + str(e),
|
|
type=ProxyErrorTypes.internal_server_error,
|
|
param=getattr(e, "param", "None"),
|
|
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
|
)
|
|
|
|
|
|
def _premium_user_check():
|
|
"""
|
|
Raises an HTTPException if the user is not a premium user
|
|
"""
|
|
from litellm.proxy.proxy_server import premium_user
|
|
|
|
if not premium_user:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail={
|
|
"error": f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}"
|
|
},
|
|
)
|
|
|
|
|
|
def is_known_model(model: Optional[str], llm_router: Optional[Router]) -> 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 = llm_router.get_model_names()
|
|
|
|
is_in_list = False
|
|
if model in model_names:
|
|
is_in_list = True
|
|
|
|
return is_in_list
|
|
|
|
|
|
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
|
|
|
|
# Join with single slash
|
|
return f"{base_path}/{route}"
|
|
|
|
|
|
def get_custom_url(request_base_url: str, route: Optional[str] = None) -> str:
|
|
# Use environment variable value, otherwise use URL from request
|
|
server_base_url = get_proxy_base_url()
|
|
if server_base_url is not None:
|
|
base_url = server_base_url
|
|
else:
|
|
base_url = request_base_url
|
|
|
|
server_root_path = get_server_root_path()
|
|
if route is not None:
|
|
if server_root_path != "":
|
|
# First join base_url with server_root_path, then with route
|
|
intermediate_url = 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() -> Optional[str]:
|
|
"""
|
|
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 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() -> Optional[str]:
|
|
"""
|
|
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 = os.getenv("DATABASE_HOST")
|
|
database_username = os.getenv("DATABASE_USERNAME")
|
|
database_password = os.getenv("DATABASE_PASSWORD")
|
|
database_name = os.getenv("DATABASE_NAME")
|
|
|
|
if database_host and database_username and database_name:
|
|
# Handle the problem of special character escaping in the database URL
|
|
database_username_enc = urllib.parse.quote_plus(database_username)
|
|
database_password_enc = (
|
|
urllib.parse.quote_plus(database_password) if database_password else ""
|
|
)
|
|
database_name_enc = 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}"
|
|
|
|
return database_url
|
|
|
|
return None
|
|
|
|
|
|
async def get_available_models_for_user(
|
|
user_api_key_dict: "UserAPIKeyAuth",
|
|
llm_router: Optional["Router"],
|
|
general_settings: dict,
|
|
user_model: Optional[str],
|
|
prisma_client: Optional["PrismaClient"] = None,
|
|
proxy_logging_obj: Optional["ProxyLogging"] = None,
|
|
team_id: Optional[str] = None,
|
|
include_model_access_groups: bool = False,
|
|
only_model_access_groups: bool = False,
|
|
return_wildcard_routes: bool = False,
|
|
user_api_key_cache: Optional["DualCache"] = 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_key_models,
|
|
get_team_models,
|
|
get_complete_model_list,
|
|
)
|
|
from litellm.proxy.auth.auth_checks import get_team_object
|
|
from litellm.proxy.management_endpoints.team_endpoints import validate_membership
|
|
|
|
# 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()
|
|
|
|
# Get key models
|
|
key_models = 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,
|
|
)
|
|
|
|
# Get team models
|
|
team_models: List[str] = user_api_key_dict.team_models
|
|
|
|
# If specific team_id is provided, validate and get team models
|
|
if team_id and prisma_client and proxy_logging_obj and user_api_key_cache:
|
|
key_models = []
|
|
team_object = 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,
|
|
)
|
|
validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object)
|
|
team_models = team_object.models
|
|
|
|
team_models = get_team_models(
|
|
team_models=team_models,
|
|
proxy_model_list=proxy_model_list,
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=include_model_access_groups,
|
|
)
|
|
|
|
# Get complete model list
|
|
all_models = get_complete_model_list(
|
|
key_models=key_models,
|
|
team_models=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,
|
|
)
|
|
|
|
return all_models
|
|
|
|
|
|
def create_model_info_response(
|
|
model_id: str,
|
|
provider: str,
|
|
include_metadata: bool = False,
|
|
fallback_type: Optional[str] = None,
|
|
llm_router: Optional["Router"] = None,
|
|
) -> dict:
|
|
"""
|
|
Create a standardized model info response.
|
|
|
|
Args:
|
|
model_id: The model ID
|
|
provider: The model provider
|
|
include_metadata: Whether to include metadata
|
|
fallback_type: Type of fallbacks to include
|
|
llm_router: LiteLLM router instance
|
|
|
|
Returns:
|
|
Dictionary containing model information
|
|
"""
|
|
from litellm.proxy.auth.model_checks import get_all_fallbacks
|
|
|
|
model_info = {
|
|
"id": model_id,
|
|
"object": "model",
|
|
"created": DEFAULT_MODEL_CREATED_AT_TIME,
|
|
"owned_by": provider,
|
|
}
|
|
|
|
# Add metadata if requested
|
|
if include_metadata:
|
|
metadata = {}
|
|
|
|
# Default fallback_type to "general" if include_metadata is true
|
|
effective_fallback_type = (
|
|
fallback_type if fallback_type is not None else "general"
|
|
)
|
|
|
|
# Validate fallback_type
|
|
valid_fallback_types = ["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 = get_all_fallbacks(
|
|
model=model_id,
|
|
llm_router=llm_router,
|
|
fallback_type=effective_fallback_type,
|
|
)
|
|
metadata["fallbacks"] = fallbacks
|
|
|
|
model_info["metadata"] = metadata
|
|
|
|
return model_info
|
|
|
|
|
|
def validate_model_access(
|
|
model_id: str,
|
|
available_models: List[str],
|
|
) -> None:
|
|
"""
|
|
Validate that a model is accessible to the user.
|
|
|
|
Args:
|
|
model_id: The model ID to validate
|
|
available_models: List of models available to the user
|
|
|
|
Raises:
|
|
HTTPException: If the model is not accessible
|
|
"""
|
|
if model_id not in available_models:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail="The model `{}` does not exist or is not accessible".format(model_id)
|
|
) |