ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revert "fix: fix ruff errors"

Browse Source 
This reverts commit eef864360e.
This commit is contained in:
Ishaan Jaffer 2025-11-08 14:33:38 -08:00
parent bce8a5d6b7
commit 0a1fc0eeb2
21 changed files with 632 additions and 355 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
enterprise/enterprise_hooks/aporia_ai.py
View File

@ -8,8 +8,6 @@
import os
import sys
from litellm.types.utils import CallTypesLiteral
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
@ -168,7 +166,16 @@ class AporiaGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,

14
enterprise/enterprise_hooks/google_text_moderation.py
View File

@ -6,13 +6,14 @@
# +-----------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
from typing import Literal
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
class _ENTERPRISE_GoogleTextModeration(CustomLogger):
@ -88,7 +89,16 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
- Calls Google's Text Moderation API

13
enterprise/enterprise_hooks/openai_moderation.py
View File

@ -12,6 +12,7 @@ sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import sys
from typing import Literal
from fastapi import HTTPException
@ -19,7 +20,6 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
class _ENTERPRISE_OpenAI_Moderation(CustomLogger):
@ -35,7 +35,16 @@ class _ENTERPRISE_OpenAI_Moderation(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
text = ""
if "messages" in data and isinstance(data["messages"], list):

13
enterprise/litellm_enterprise/enterprise_callbacks/llama_guard.py
View File

@ -23,7 +23,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral, Choices, ModelResponse
from litellm.types.utils import Choices, ModelResponse
class _ENTERPRISE_LlamaGuard(CustomLogger):
@ -98,7 +98,16 @@ class _ENTERPRISE_LlamaGuard(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
- Calls the Llama Guard Endpoint

12
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View File

@ -17,7 +17,6 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_formatted_prompt
@ -121,7 +120,16 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
- Calls the LLM Guard Endpoint

14
enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py
View File

@ -31,7 +31,6 @@ from litellm.types.integrations.pagerduty import (
PagerDutyRequestBody,
)
from litellm.types.utils import (
CallTypesLiteral,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -143,7 +142,18 @@ class PagerDutyAlerting(SlackAlerting):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, dict]]:
"""
Example of detecting hanging requests by waiting a given threshold.

24
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View File

@ -36,7 +36,6 @@ from litellm.types.llms.openai import (
OpenAIFilesPurpose,
)
from litellm.types.utils import (
CallTypesLiteral,
LiteLLMBatch,
LiteLLMFineTuningJob,
LLMResponseTypes,
@ -273,7 +272,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: Dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"acreate_batch",
"aretrieve_batch",
"acreate_file",
"afile_list",
"afile_delete",
"afile_content",
"acreate_fine_tuning_job",
"aretrieve_fine_tuning_job",
"alist_fine_tuning_jobs",
"acancel_fine_tuning_job",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, Dict, None]:
"""
- Detect litellm_proxy/ file_id

25
litellm/proxy/example_config_yaml/custom_callbacks1.py
View File

@ -3,7 +3,6 @@ from typing import Literal, Optional
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.proxy_server import DualCache, UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
# This file includes the custom callbacks for LiteLLM Proxy
@ -22,7 +21,18 @@ class MyCustomHandler(
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
):
return data
@ -48,7 +58,16 @@ class MyCustomHandler(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
pass

25
litellm/proxy/example_config_yaml/custom_guardrail.py
View File

@ -6,7 +6,6 @@ from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.utils import CallTypesLiteral
class myCustomGuardrail(CustomGuardrail):
@ -24,7 +23,18 @@ class myCustomGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, dict]]:
"""
Runs before the LLM API call
@ -52,7 +62,16 @@ class myCustomGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
Runs in parallel to LLM API call

27
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
View File

@ -7,7 +7,7 @@
import asyncio
import json
import os
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Type, Union
from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal, Optional, Type, Union
from fastapi import HTTPException
from pydantic import BaseModel
@ -23,7 +23,6 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import (
CallTypesLiteral,
Choices,
EmbeddingResponse,
ImageResponse,
@ -68,7 +67,18 @@ class AimGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Pre-Call Hook")
return await self.call_aim_guardrail(
@ -79,7 +89,16 @@ class AimGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug("Inside AIM Moderation Hook")

31
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View File

@ -42,7 +42,6 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
)
from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
Choices,
GuardrailStatus,
ModelResponse,
@ -609,7 +608,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
verbose_proxy_logger.debug(
"Inside Bedrock Pre-Call Hook for call_type: %s", call_type
@ -675,7 +685,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@ -1144,7 +1163,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
This method allows users to test Bedrock guardrails without making actual LLM calls.
It creates a mock request and response to test the guardrail functionality.
Args:
text: The text to analyze
language: Optional language parameter (not used by Bedrock)
@ -1156,11 +1175,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
mock_messages: List[AllMessageValues] = [
ChatCompletionUserMessage(role="user", content=text)
]
# Use provided request_data or create a mock one for testing
if request_data is None:
request_data = {"messages": mock_messages}
bedrock_response = await self.make_bedrock_api_request(
source="INPUT",
messages=mock_messages,

179
litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
View File

@ -7,7 +7,16 @@
import os
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional, Type, Union
from typing import (
Any,
AsyncGenerator,
Dict,
List,
Literal,
Optional,
Type,
Union,
)
import httpx
@ -27,7 +36,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import (
DynamoAIRequest,
DynamoAIResponse,
)
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
from litellm.types.utils import GuardrailStatus, ModelResponseStream
GUARDRAIL_NAME = "dynamoai"
@ -35,10 +44,9 @@ GUARDRAIL_NAME = "dynamoai"
class DynamoAIGuardrails(CustomGuardrail):
"""
DynamoAI Guardrails integration for LiteLLM.
Provides content moderation and policy enforcement using DynamoAI's guardrail API.
"""
def __init__(
self,
guardrail_name: str = "litellm_test",
@ -51,30 +59,28 @@ class DynamoAIGuardrails(CustomGuardrail):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
# Set API configuration
self.api_key = api_key or os.getenv("DYNAMOAI_API_KEY")
if not self.api_key:
raise ValueError(
"DynamoAI API key is required. Set DYNAMOAI_API_KEY environment variable or pass api_key parameter."
)
self.api_base = api_base or os.getenv(
"DYNAMOAI_API_BASE", "https://api.dynamo.ai"
)
self.api_url = f"{self.api_base}/v1/moderation/analyze/"
# Model ID for tracking/logging purposes
self.model_id = model_id or os.getenv("DYNAMOAI_MODEL_ID", "")
# Policy IDs - get from parameter, env var, or use empty list
env_policy_ids = os.getenv("DYNAMOAI_POLICY_IDS", "")
self.policy_ids = policy_ids or (
env_policy_ids.split(",") if env_policy_ids else []
)
self.policy_ids = policy_ids or (env_policy_ids.split(",") if env_policy_ids else [])
self.guardrail_name = guardrail_name
self.guardrail_provider = "dynamoai"
# store kwargs as optional_params
self.optional_params = kwargs
@ -102,38 +108,38 @@ class DynamoAIGuardrails(CustomGuardrail):
) -> DynamoAIResponse:
"""
Call DynamoAI Guardrails API to analyze messages for policy violations.
Args:
messages: List of messages to analyze
text_type: Type of text being analyzed ("input" or "output")
request_data: Optional request data for logging purposes
Returns:
DynamoAIResponse: Response from the DynamoAI Guardrails API
"""
start_time = datetime.now()
payload: DynamoAIRequest = {
"messages": messages,
}
# Add optional fields if provided
if self.policy_ids:
payload["policyIds"] = self.policy_ids
if self.model_id:
payload["modelId"] = self.model_id
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
"Authorization": f"Bearer {self.api_key}"
}
verbose_proxy_logger.debug(
"DynamoAI request to %s with payload=%s",
self.api_url,
payload,
)
try:
response = await self.async_handler.post(
url=self.api_url,
@ -142,10 +148,10 @@ class DynamoAIGuardrails(CustomGuardrail):
)
response.raise_for_status()
response_json = response.json()
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Add guardrail information to request trace
if request_data:
guardrail_status = self._determine_guardrail_status(response_json)
@ -158,15 +164,17 @@ class DynamoAIGuardrails(CustomGuardrail):
end_time=end_time.timestamp(),
duration=duration,
)
return response_json
except httpx.HTTPError as e:
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
verbose_proxy_logger.error("DynamoAI API request failed: %s", str(e))
verbose_proxy_logger.error(
"DynamoAI API request failed: %s", str(e)
)
# Add guardrail information with failure status
if request_data:
self.add_standard_logging_guardrail_information_to_request_data(
@ -178,7 +186,7 @@ class DynamoAIGuardrails(CustomGuardrail):
end_time=end_time.timestamp(),
duration=duration,
)
raise
def _process_dynamoai_guardrails_response(
@ -186,35 +194,35 @@ class DynamoAIGuardrails(CustomGuardrail):
) -> DynamoAIProcessedResult:
"""
Process the response from the DynamoAI Guardrails API
Args:
response: The response from the API with 'finalAction' and 'appliedPolicies' keys
Returns:
DynamoAIProcessedResult: Processed response with detected violations
"""
final_action = response.get("finalAction", "NONE")
applied_policies = response.get("appliedPolicies", [])
violations_detected: List[str] = []
violation_details: Dict[str, Any] = {}
# For now, only handle BLOCK action
if final_action == "BLOCK":
for applied_policy in applied_policies:
policy_info = applied_policy.get("policy", {})
policy_outputs = applied_policy.get("outputs", {})
# Get policy name and action
policy_name = policy_info.get("name", "unknown")
# Check for action in multiple places
policy_action = (
applied_policy.get("action")
or (policy_outputs.get("action") if policy_outputs else None)
or "NONE"
applied_policy.get("action") or
(policy_outputs.get("action") if policy_outputs else None) or
"NONE"
)
# Only include policies with BLOCK action
if policy_action == "BLOCK":
violations_detected.append(policy_name)
@ -223,14 +231,12 @@ class DynamoAIGuardrails(CustomGuardrail):
"action": policy_action,
"method": policy_info.get("method"),
"description": policy_info.get("description"),
"message": (
policy_outputs.get("message") if policy_outputs else None
),
"message": policy_outputs.get("message") if policy_outputs else None,
}
return {
"violations_detected": violations_detected,
"violation_details": violation_details,
"violation_details": violation_details
}
def _determine_guardrail_status(
@ -238,7 +244,7 @@ class DynamoAIGuardrails(CustomGuardrail):
) -> GuardrailStatus:
"""
Determine the guardrail status based on DynamoAI API response.
Returns:
"success": Content allowed through with no violations (finalAction is NONE)
"guardrail_intervened": Content blocked (finalAction is BLOCK)
@ -247,21 +253,21 @@ class DynamoAIGuardrails(CustomGuardrail):
try:
if not isinstance(response_json, dict):
return "guardrail_failed_to_respond"
# Check for error in response
if response_json.get("error"):
return "guardrail_failed_to_respond"
final_action = response_json.get("finalAction", "NONE")
if final_action == "NONE":
return "success"
elif final_action == "BLOCK":
return "guardrail_intervened"
# For now, treat other actions as success (WARN, REDACT, SANITIZE not implemented yet)
return "success"
except Exception as e:
verbose_proxy_logger.error(
"Error determining DynamoAI guardrail status: %s", str(e)
@ -271,24 +277,22 @@ class DynamoAIGuardrails(CustomGuardrail):
def _create_error_message(self, processed_result: DynamoAIProcessedResult) -> str:
"""
Create a detailed error message from processed guardrail results.
Args:
processed_result: Processed response with detected violations
Returns:
Formatted error message string
"""
violations_detected = processed_result["violations_detected"]
violation_details = processed_result["violation_details"]
error_message = (
f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n"
)
error_message = f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n"
for policy_name in violations_detected:
error_message += f"- {policy_name.upper()}:\n"
details = violation_details.get(policy_name, {})
# Format violation details
if details.get("action"):
error_message += f" Action: {details['action']}\n"
@ -301,7 +305,7 @@ class DynamoAIGuardrails(CustomGuardrail):
if details.get("policyId"):
error_message += f" Policy ID: {details['policyId']}\n"
error_message += "\n"
return error_message.strip()
async def async_pre_call_hook(
@ -309,7 +313,18 @@ class DynamoAIGuardrails(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
@ -334,9 +349,7 @@ class DynamoAIGuardrails(CustomGuardrail):
request_data=data,
)
verbose_proxy_logger.debug(
"Guardrails async_pre_call_hook result=%s", result
)
verbose_proxy_logger.debug("Guardrails async_pre_call_hook result=%s", result)
# Process the guardrails response
processed_result = self._process_dynamoai_guardrails_response(result)
@ -351,14 +364,23 @@ class DynamoAIGuardrails(CustomGuardrail):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=self.guardrail_name
)
return data
async def async_moderation_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
Runs in parallel to LLM API call
@ -382,9 +404,7 @@ class DynamoAIGuardrails(CustomGuardrail):
request_data=data,
)
verbose_proxy_logger.debug(
"Guardrails async_moderation_hook result=%s", result
)
verbose_proxy_logger.debug("Guardrails async_moderation_hook result=%s", result)
# Process the guardrails response
processed_result = self._process_dynamoai_guardrails_response(result)
@ -429,26 +449,22 @@ class DynamoAIGuardrails(CustomGuardrail):
return
verbose_proxy_logger.debug("async_post_call_success_hook response=%s", response)
# Check if the ModelResponse has text content in its choices
# to avoid sending empty content to DynamoAI (e.g., during tool calls)
if isinstance(response, litellm.ModelResponse):
has_text_content = False
dynamoai_messages: List[Dict[str, Any]] = []
for choice in response.choices:
if isinstance(choice, litellm.Choices):
if choice.message.content and isinstance(
choice.message.content, str
):
if choice.message.content and isinstance(choice.message.content, str):
has_text_content = True
dynamoai_messages.append(
{
"role": choice.message.role or "assistant",
"content": choice.message.content,
}
)
dynamoai_messages.append({
"role": choice.message.role or "assistant",
"content": choice.message.content
})
if not has_text_content:
verbose_proxy_logger.warning(
"DynamoAI: not running guardrail. No output text in response"
@ -462,9 +478,7 @@ class DynamoAIGuardrails(CustomGuardrail):
request_data=data,
)
verbose_proxy_logger.debug(
"Guardrails async_post_call_success_hook result=%s", result
)
verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result=%s", result)
# Process the guardrails response
processed_result = self._process_dynamoai_guardrails_response(result)
@ -503,3 +517,4 @@ class DynamoAIGuardrails(CustomGuardrail):
)
return DynamoAIGuardrailConfigModel

28
litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py
View File

@ -7,7 +7,7 @@
import os
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Union
import httpx
@ -25,7 +25,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
EnkryptAIProcessedResult,
EnkryptAIResponse,
)
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
from litellm.types.utils import GuardrailStatus, ModelResponseStream
GUARDRAIL_NAME = "enkryptai"
@ -284,7 +284,18 @@ class EnkryptAIGuardrails(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
@ -337,7 +348,16 @@ class EnkryptAIGuardrails(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
Runs in parallel to LLM API call

30
litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
View File

@ -7,7 +7,7 @@
import os
from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
from typing import Any, AsyncGenerator, Dict, List, Literal, Optional, Union
from urllib.parse import urlencode
import httpx
@ -26,7 +26,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.ibm import (
IBMDetectorDetection,
IBMDetectorResponseOrchestrator,
)
from litellm.types.utils import CallTypesLiteral, GuardrailStatus, ModelResponseStream
from litellm.types.utils import GuardrailStatus, ModelResponseStream
GUARDRAIL_NAME = "ibm_guardrails"
@ -47,7 +47,7 @@ class IBMGuardrailDetector(CustomGuardrail):
):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
params={"ssl_verify": verify_ssl},
params={"ssl_verify": verify_ssl}
)
# Set API configuration
@ -436,7 +436,18 @@ class IBMGuardrailDetector(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call
@ -522,7 +533,16 @@ class IBMGuardrailDetector(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
"""
Runs in parallel to LLM API call

17
litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
View File

@ -1,5 +1,5 @@
from datetime import datetime
from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Type, Union
from fastapi import HTTPException
@ -18,7 +18,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.javelin import (
JavelinGuardRequest,
JavelinGuardResponse,
)
from litellm.types.utils import CallTypesLiteral, GuardrailStatus
from litellm.types.utils import GuardrailStatus
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
@ -165,7 +165,18 @@ class JavelinGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: litellm.DualCache,
data: Dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, Dict]]:
"""
Pre-call hook for the Javelin guardrail.

32
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View File

@ -1,7 +1,7 @@
import copy
import os
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Union
from typing import Dict, List, Literal, Optional, Tuple, Union
from fastapi import HTTPException
@ -20,7 +20,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.lakera_ai_v2 import (
LakeraAIRequest,
LakeraAIResponse,
)
from litellm.types.utils import CallTypesLiteral, GuardrailStatus
from litellm.types.utils import GuardrailStatus
class LakeraAIGuardrail(CustomGuardrail):
@ -183,7 +183,18 @@ class LakeraAIGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: litellm.DualCache,
data: Dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, Dict]]:
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@ -246,7 +257,16 @@ class LakeraAIGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@ -313,7 +333,7 @@ class LakeraAIGuardrail(CustomGuardrail):
breakdown = lakera_response.get("breakdown", []) or []
if not breakdown:
return False
has_violations = False
for item in breakdown:
if item.get("detected", False):
@ -321,7 +341,7 @@ class LakeraAIGuardrail(CustomGuardrail):
detector_type = item.get("detector_type", "") or ""
if not detector_type.startswith("pii/"):
return False
# Return True only if there are violations and they are all PII
return has_violations

175
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
View File

@ -6,22 +6,11 @@
# +-------------------------------------------------------------+
import asyncio
import json
import os
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Dict,
Final,
List,
Literal,
Optional,
Type,
Union,
)
from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Optional, Type, Union
from urllib.parse import urljoin
import json
from fastapi import HTTPException
@ -29,22 +18,25 @@ import litellm
from litellm import DualCache, ModelResponse
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.main import stream_chunk_builder
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import EmbeddingResponse, GuardrailStatus, ImageResponse
from litellm.types.utils import (
CallTypesLiteral,
EmbeddingResponse,
GuardrailStatus,
ImageResponse,
ModelResponseStream,
TextCompletionResponse,
)
from typing import (
List,
AsyncGenerator
)
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder
from litellm.types.utils import TextCompletionResponse
# Constants
USER_ROLE: Final[Literal["user"]] = "user"
@ -56,9 +48,9 @@ MessageRole = Literal["user", "assistant"]
LLMResponse = Union[Any, ModelResponse, EmbeddingResponse, ImageResponse]
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
class NomaBlockedMessage(HTTPException):
"""Exception raised when Noma guardrail blocks a message"""
@ -168,7 +160,13 @@ class NomaGuardrail(CustomGuardrail):
return None
payload = {
"input": [{"type": "message", "role": "user", "content": user_message}]
"input": [
{
"type": "message",
"role": "user",
"content": user_message
}
]
}
response_json = await self._call_noma_api(
payload=payload,
@ -177,13 +175,13 @@ class NomaGuardrail(CustomGuardrail):
user_auth=user_auth,
extra_data=extra_data,
)
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Determine guardrail status based on response
guardrail_status = self._determine_guardrail_status(response_json)
# Always log guardrail information for consistency
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@ -224,7 +222,7 @@ class NomaGuardrail(CustomGuardrail):
user_auth: UserAPIKeyAuth,
) -> Optional[str]:
"""Shared logic for processing LLM response checks"""
start_time = datetime.now()
extra_data = self.get_guardrail_dynamic_request_body_params(request_data)
@ -245,7 +243,12 @@ class NomaGuardrail(CustomGuardrail):
{
"type": "message",
"role": "assistant",
"content": [{"type": "input_text", "text": content}],
"content": [
{
"type": "input_text",
"text": content
}
]
}
]
}
@ -257,13 +260,13 @@ class NomaGuardrail(CustomGuardrail):
user_auth=user_auth,
extra_data=extra_data,
)
end_time = datetime.now()
duration = (end_time - start_time).total_seconds()
# Determine guardrail status based on response
guardrail_status = self._determine_guardrail_status(response_json)
# Always log guardrail information for consistency
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@ -300,10 +303,10 @@ class NomaGuardrail(CustomGuardrail):
def _determine_guardrail_status(self, response_json: dict) -> GuardrailStatus:
"""
Determine the guardrail status based on NOMA API response.
Args:
response_json: Response from NOMA API
Returns:
"success": Content allowed through with no violations
"guardrail_intervened": Content blocked due to policy violations
@ -313,26 +316,24 @@ class NomaGuardrail(CustomGuardrail):
# Check if we got a valid response structure
if not isinstance(response_json, dict):
return "guardrail_failed_to_respond"
# Get the aggregatedScanResult from the response
# aggregatedScanResult=True means unsafe (block), False means safe (allow)
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
# If aggregatedScanResult is False, content is safe/allowed
if aggregated_scan_result is False:
return "success"
# If aggregatedScanResult is True, content is blocked/flagged
if aggregated_scan_result is True:
return "guardrail_intervened"
# If aggregatedScanResult is missing or invalid, treat as failure
return "guardrail_failed_to_respond"
except Exception as e:
verbose_proxy_logger.error(
f"Error determining NOMA guardrail status: {str(e)}"
)
verbose_proxy_logger.error(f"Error determining NOMA guardrail status: {str(e)}")
return "guardrail_failed_to_respond"
def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool:
@ -391,16 +392,12 @@ class NomaGuardrail(CustomGuardrail):
scan_result = response_json.get("scanResult", [])
if not scan_result:
return None
# Find the scan result matching the message type (role)
for result_item in scan_result:
if result_item.get("role") == message_type:
return (
result_item.get("results", {})
.get("anonymizedContent", {})
.get("anonymized", "")
)
return result_item.get("results", {}).get("anonymizedContent", {}).get("anonymized", "")
return None
def _should_anonymize(self, response_json: dict, message_type: MessageRole) -> bool:
@ -426,7 +423,7 @@ class NomaGuardrail(CustomGuardrail):
# aggregatedScanResult=False means safe, True means unsafe
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
# If aggregatedScanResult is False, content is safe - anonymize if available
if not aggregated_scan_result:
return True
@ -435,15 +432,13 @@ class NomaGuardrail(CustomGuardrail):
scan_result = response_json.get("scanResult", [])
if not scan_result:
return False
if not isinstance(scan_result, list) or len(scan_result) == 0:
return False
for result_item in scan_result:
if result_item.get("role") == message_type:
return self._should_only_sensitive_data_failed(
result_item.get("results", {})
)
return self._should_only_sensitive_data_failed(result_item.get("results", {}))
return False
@ -539,7 +534,7 @@ class NomaGuardrail(CustomGuardrail):
try:
# aggregatedScanResult=True means blocked, False means allowed
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
if aggregated_scan_result: # True = unsafe
msg = f"Noma guardrail blocked {type} message: {message}"
verbose_proxy_logger.warning(msg)
@ -556,9 +551,20 @@ class NomaGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, dict]]:
verbose_proxy_logger.debug("Running Noma pre-call hook")
if (
@ -589,7 +595,6 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@ -600,7 +605,7 @@ class NomaGuardrail(CustomGuardrail):
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}")
if self.block_failures:
@ -611,7 +616,16 @@ class NomaGuardrail(CustomGuardrail):
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
@ -637,7 +651,6 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@ -648,7 +661,7 @@ class NomaGuardrail(CustomGuardrail):
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}")
if self.block_failures:
@ -687,7 +700,6 @@ class NomaGuardrail(CustomGuardrail):
except Exception as e:
# Log technical failures
from datetime import datetime
start_time = datetime.now()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider="noma",
@ -698,7 +710,7 @@ class NomaGuardrail(CustomGuardrail):
end_time=start_time.timestamp(),
duration=0.0,
)
verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}")
if self.block_failures:
raise
@ -744,31 +756,37 @@ class NomaGuardrail(CustomGuardrail):
last_user_message = user_messages[-1].get("content", "")
if isinstance(last_user_message, str):
return [{"type": "input_text", "text": last_user_message}]
return [{
"type": "input_text",
"text": last_user_message
}]
elif isinstance(last_user_message, list):
converted_messages = []
for message in last_user_message:
converted_message = self._convert_single_user_message_to_payload(
message
)
converted_message = self._convert_single_user_message_to_payload(message)
if converted_message is not None:
converted_messages.append(converted_message)
return converted_messages
else:
return None
def _convert_single_user_message_to_payload(
self, user_message: Any
) -> Optional[dict]:
def _convert_single_user_message_to_payload(self, user_message: Any) -> Optional[dict]:
if isinstance(user_message, str):
return {"type": "input_text", "text": user_message}
return {
"type": "input_text",
"text": user_message
}
elif user_message.get("type", "") == "image_url":
return {
"type": "input_image",
"image_url": user_message.get("image_url", {}).get("url", ""),
"image_url": user_message.get("image_url", {}).get("url", "")
}
elif user_message.get("type", "") == "text":
return {"type": "input_text", "text": user_message.get("text", "")}
return {
"type": "input_text",
"text": user_message.get("text", "")
}
else:
return None
@ -798,16 +816,13 @@ class NomaGuardrail(CustomGuardrail):
"applicationId": extra_data.get("application_id")
or request_data.get("metadata", {})
.get("headers", {})
.get("x-noma-application-id")
or self.application_id,
.get("x-noma-application-id") or self.application_id,
"ipAddress": request_data.get("metadata", {}).get(
"requester_ip_address", None
),
"userId": (
user_auth.user_email
if user_auth.user_email
else user_auth.user_id
),
"userId": user_auth.user_email
if user_auth.user_email
else user_auth.user_id,
"sessionId": call_id,
"requestId": llm_request_id,
},
@ -829,7 +844,7 @@ class NomaGuardrail(CustomGuardrail):
"""
# aggregatedScanResult=True means blocked, False means allowed
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
if aggregated_scan_result: # True = unsafe, block it
msg = f"Noma guardrail blocked {type} message: {message}"

17
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View File

@ -6,7 +6,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint
3. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_streaming_iterator_hook
"""
from typing import Any, AsyncGenerator, Union
from typing import Any, AsyncGenerator, Literal, Union
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -16,7 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.llms import load_guardrail_translation_mappings
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypes, CallTypesLiteral, ModelResponseStream
from litellm.types.utils import CallTypes, ModelResponseStream
GUARDRAIL_NAME = "unified_llm_guardrails"
endpoint_guardrail_translation_mappings = None
@ -43,7 +43,18 @@ class UnifiedLLMGuardrails(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
"""
Runs before the LLM API call

44
litellm/proxy/hooks/dynamic_rate_limiter.py
View File

@ -4,7 +4,7 @@
import asyncio
import os
from typing import List, Optional, Tuple, Union
from typing import List, Literal, Optional, Tuple, Union
from fastapi import HTTPException
@ -15,7 +15,6 @@ from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
from litellm.utils import get_utc_datetime
from .rate_limiter_utils import convert_priority_to_percent
@ -103,10 +102,10 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
"""
try:
# Get model info first for conversion
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
)
model_group_info: Optional[
ModelGroupInfo
] = self.llm_router.get_model_group_info(model_group=model)
weight: float = 1
if (
litellm.priority_reservation is None
@ -194,7 +193,18 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[
Union[Exception, str, dict]
]: # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm
@ -277,16 +287,16 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
) = await self.check_available_usage(
model=model_info["model_name"], priority=key_priority
)
response._hidden_params["additional_headers"] = (
{ # Add additional response headers - easier debugging
"x-litellm-model_group": model_info["model_name"],
"x-ratelimit-remaining-litellm-project-tokens": available_tpm,
"x-ratelimit-remaining-litellm-project-requests": available_rpm,
"x-ratelimit-remaining-model-tokens": model_tpm,
"x-ratelimit-remaining-model-requests": model_rpm,
"x-ratelimit-current-active-projects": active_projects,
}
)
response._hidden_params[
"additional_headers"
] = { # Add additional response headers - easier debugging
"x-litellm-model_group": model_info["model_name"],
"x-ratelimit-remaining-litellm-project-tokens": available_tpm,
"x-ratelimit-remaining-litellm-project-requests": available_rpm,
"x-ratelimit-remaining-model-tokens": model_tpm,
"x-ratelimit-remaining-model-requests": model_rpm,
"x-ratelimit-current-active-projects": active_projects,
}
return response
return await super().async_post_call_success_hook(

224
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View File

@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
import os
from datetime import datetime
from typing import Callable, Dict, List, Optional, Union
from typing import Callable, Dict, List, Literal, Optional, Union
from fastapi import HTTPException
@ -22,20 +22,19 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.hooks.rate_limiter_utils import convert_priority_to_percent
from litellm.proxy.utils import InternalUsageCache
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"""
Saturation-aware priority-based rate limiter using v3 infrastructure.
Key features:
1. Model capacity ALWAYS enforced at 100% (prevents over-allocation)
2. Priority usage tracked from first request (accurate accounting)
3. Priority limits only enforced when saturated >= threshold
4. Three-phase checking prevents partial counter increments
5. Reuses v3 limiter's Redis-based tracking (multi-instance safe)
How it works:
- Phase 1: Read-only check of ALL limits (no increments)
- Phase 2: Decide enforcement based on saturation
@ -44,7 +43,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
- When saturated: strict priority-based limits enforced (fair)
- Uses v3 limiter's atomic Lua scripts for race-free increments
"""
def __init__(
self,
internal_usage_cache: DualCache,
@ -58,9 +56,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
def update_variables(self, llm_router: Router):
self.llm_router = llm_router
def _get_priority_weight(
self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None
) -> float:
def _get_priority_weight(self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None) -> float:
"""Get the weight for a given priority from litellm.priority_reservation"""
weight: float = litellm.priority_reservation_settings.default_priority
if (
@ -80,32 +76,30 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
weight = convert_priority_to_percent(value, model_info)
return weight
def _normalize_priority_weights(
self, model_info: ModelGroupInfo
) -> Dict[str, float]:
def _normalize_priority_weights(self, model_info: ModelGroupInfo) -> Dict[str, float]:
"""
Normalize priority weights if they sum to > 1.0
Handles over-allocation: {key_a: 0.60, key_b: 0.80} -> {key_a: 0.43, key_b: 0.57}
Converts absolute rpm/tpm values to percentages based on model capacity.
"""
if litellm.priority_reservation is None:
return {}
# Convert all values to percentages first
weights: Dict[str, float] = {}
for k, v in litellm.priority_reservation.items():
weights[k] = convert_priority_to_percent(v, model_info)
total_weight = sum(weights.values())
if total_weight > 1.0:
normalized = {k: v / total_weight for k, v in weights.items()}
verbose_proxy_logger.debug(
f"Normalized over-allocated priorities: {weights} -> {normalized}"
)
return normalized
return weights
def _get_priority_allocation(
@ -117,31 +111,29 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> tuple[float, str]:
"""
Get priority weight and pool key for a given priority.
For explicit priorities: returns specific allocation and unique pool key
For default priority: returns default allocation and shared pool key
Args:
model: Model name
priority: Priority level (None for default)
normalized_weights: Pre-computed normalized weights
model_info: Model configuration (optional, for fallback conversion)
Returns:
tuple: (priority_weight, priority_key)
"""
# Check if this key has an explicit priority in litellm.priority_reservation
has_explicit_priority = (
priority is not None
and litellm.priority_reservation is not None
priority is not None
and litellm.priority_reservation is not None
and priority in litellm.priority_reservation
)
if has_explicit_priority and priority is not None:
# Explicit priority: get its specific allocation
priority_weight = normalized_weights.get(
priority, self._get_priority_weight(priority, model_info)
)
priority_weight = normalized_weights.get(priority, self._get_priority_weight(priority, model_info))
# Use unique key per priority level
priority_key = f"{model}:{priority}"
else:
@ -149,7 +141,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
priority_weight = litellm.priority_reservation_settings.default_priority
# Use shared key for all default-priority requests
priority_key = f"{model}:default_pool"
return priority_weight, priority_key
async def _check_model_saturation(
@ -159,16 +151,16 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> float:
"""
Check current saturation by directly querying v3 limiter's cache keys.
Reuses v3 limiter's Redis-based tracking (works across multiple instances).
Reads counters WITHOUT incrementing them.
Returns:
float: Saturation ratio (0.0 = empty, 1.0 = at capacity, >1.0 = over)
"""
try:
max_saturation = 0.0
# Query RPM saturation
if model_group_info.rpm is not None and model_group_info.rpm > 0:
# Use v3 limiter's key format: {key:value}:rate_limit_type
@ -177,24 +169,24 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
value=model,
rate_limit_type="requests",
)
# Query cache for current counter value
counter_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False, # Check Redis too
)
if counter_value is not None:
current_requests = int(counter_value)
rpm_saturation = current_requests / model_group_info.rpm
max_saturation = max(max_saturation, rpm_saturation)
verbose_proxy_logger.debug(
f"Model {model} RPM: {current_requests}/{model_group_info.rpm} "
f"({rpm_saturation:.1%})"
)
# Query TPM saturation
if model_group_info.tpm is not None and model_group_info.tpm > 0:
counter_key = self.v3_limiter.create_rate_limit_keys(
@ -202,29 +194,29 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
value=model,
rate_limit_type="tokens",
)
counter_value = await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False,
)
if counter_value is not None:
current_tokens = float(counter_value)
tpm_saturation = current_tokens / model_group_info.tpm
max_saturation = max(max_saturation, tpm_saturation)
verbose_proxy_logger.debug(
f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} "
f"({tpm_saturation:.1%})"
)
verbose_proxy_logger.debug(
f"Model {model} overall saturation: {max_saturation:.1%}"
)
return max_saturation
except Exception as e:
verbose_proxy_logger.error(
f"Error checking saturation for {model}: {str(e)}"
@ -240,17 +232,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> List[RateLimitDescriptor]:
"""
Create rate limit descriptors with normalized priority weights.
Uses normalized weights to handle over-allocation scenarios.
For explicit priorities: each priority gets its own pool (e.g., prod gets 75%)
For default priority: ALL keys without explicit priority share ONE pool (e.g., all share 25%)
"""
descriptors: List[RateLimitDescriptor] = []
# Get model group info
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(
model_group=model
)
if model_group_info is None:
return descriptors
@ -263,21 +255,21 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
normalized_weights=normalized_weights,
model_info=model_group_info,
)
rate_limit_config: RateLimitDescriptorRateLimitObject = {}
# Apply priority weight to model limits
if model_group_info.tpm is not None:
reserved_tpm = int(model_group_info.tpm * priority_weight)
rate_limit_config["tokens_per_unit"] = reserved_tpm
if model_group_info.rpm is not None:
reserved_rpm = int(model_group_info.rpm * priority_weight)
rate_limit_config["requests_per_unit"] = reserved_rpm
if rate_limit_config:
rate_limit_config["window_size"] = self.v3_limiter.window_size
descriptors.append(
RateLimitDescriptor(
key="priority_model",
@ -296,12 +288,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> RateLimitDescriptor:
"""
Create a descriptor for tracking model-wide usage.
Args:
model: Model name
model_group_info: Model configuration with RPM/TPM limits
high_limit_multiplier: Multiplier for limits (use >1 for tracking-only)
Returns:
Rate limit descriptor for model-wide tracking
"""
@ -310,19 +302,18 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
value=model,
rate_limit={
"requests_per_unit": (
model_group_info.rpm * high_limit_multiplier
if model_group_info.rpm
else None
model_group_info.rpm * high_limit_multiplier
if model_group_info.rpm else None
),
"tokens_per_unit": (
model_group_info.tpm * high_limit_multiplier
if model_group_info.tpm
else None
model_group_info.tpm * high_limit_multiplier
if model_group_info.tpm else None
),
"window_size": self.v3_limiter.window_size,
},
)
async def _check_rate_limits(
self,
model: str,
@ -334,23 +325,23 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> None:
"""
Check rate limits using THREE-PHASE approach to prevent partial increments.
Phase 1: Read-only check of ALL limits (no increments)
Phase 2: Decide which limits to enforce based on saturation
Phase 3: Increment ALL counters atomically (model + priority)
This prevents the bug where:
- Model counter increments in stage 1
- Priority check fails in stage 2
- Request blocked but model counter already incremented
Key behaviors:
- All checks performed first (read-only)
- Only increment counters if request will be allowed
- Model capacity: Always enforced at 100%
- Priority limits: Only enforced when saturated >= threshold
- Both counters tracked from first request (accurate accounting)
Args:
model: Model name
model_group_info: Model configuration
@ -358,20 +349,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
key_priority: User's priority level
saturation: Current saturation level
data: Request data dictionary
Raises:
HTTPException: If any limit is exceeded
"""
import json
saturation_threshold = (
litellm.priority_reservation_settings.saturation_threshold
)
saturation_threshold = litellm.priority_reservation_settings.saturation_threshold
should_enforce_priority = saturation >= saturation_threshold
# Build ALL descriptors upfront
descriptors_to_check: List[RateLimitDescriptor] = []
# Model-wide descriptor (always enforce)
model_wide_descriptor = self._create_model_tracking_descriptor(
model=model,
@ -379,7 +367,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
high_limit_multiplier=1,
)
descriptors_to_check.append(model_wide_descriptor)
# Priority descriptors (always track, conditionally enforce)
priority_descriptors = self._create_priority_based_descriptors(
model=model,
@ -388,33 +376,31 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
if priority_descriptors:
descriptors_to_check.extend(priority_descriptors)
# PHASE 1: Read-only check of ALL limits (no increments)
check_response = await self.v3_limiter.should_rate_limit(
descriptors=descriptors_to_check,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True, # CRITICAL: Don't increment counters yet
)
verbose_proxy_logger.debug(
f"Read-only check: {json.dumps(check_response, indent=2)}"
)
verbose_proxy_logger.debug(f"Read-only check: {json.dumps(check_response, indent=2)}")
# PHASE 2: Decide which limits to enforce
if check_response["overall_code"] == "OVER_LIMIT":
for status in check_response["statuses"]:
if status["code"] == "OVER_LIMIT":
descriptor_key = status["descriptor_key"]
# Model-wide limit exceeded (ALWAYS enforce)
if descriptor_key == "model_saturation_check":
raise HTTPException(
status_code=429,
detail={
"error": f"Model capacity reached for {model}. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
@ -422,7 +408,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"x-litellm-priority": key_priority or "default",
},
)
# Priority limit exceeded (ONLY enforce when saturated)
elif descriptor_key == "priority_model" and should_enforce_priority:
verbose_proxy_logger.debug(
@ -433,10 +419,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
f"Priority: {key_priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},
headers={
"retry-after": str(self.v3_limiter.window_size),
@ -445,12 +431,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"x-litellm-saturation": f"{saturation:.2%}",
},
)
# PHASE 3: Increment counters separately to avoid early-exit issues
# Model counter must ALWAYS increment, but priority counter might be over limit
# If we increment them together, v3_limiter's in-memory check will exit early
# and skip incrementing the model counter
# Step 3a: Increment model-wide counter (always)
model_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=[model_wide_descriptor],
@ -465,12 +451,11 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Combine responses for post-call hook
combined_response = {
"overall_code": model_increment_response["overall_code"],
"statuses": model_increment_response["statuses"]
+ priority_increment_response["statuses"],
"statuses": model_increment_response["statuses"] + priority_increment_response["statuses"]
}
data["litellm_proxy_rate_limit_response"] = combined_response
else:
@ -481,39 +466,50 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, dict]]:
"""
Saturation-aware pre-call hook for priority-based rate limiting.
Flow:
1. Check current saturation level
2. THREE-PHASE rate limit check:
- PHASE 1: Read-only check of ALL limits (no increments)
- PHASE 2: Decide which limits to enforce based on saturation
- PHASE 3: Increment ALL counters atomically if request allowed
This three-phase approach ensures:
- Model capacity is NEVER exceeded (always enforced at 100%)
- Priority usage tracked from first request (accurate metrics)
- Counters only increment when request will be allowed (prevents phantom usage)
- When under-saturated: priorities can borrow unused capacity (generous)
- When saturated: fair allocation based on normalized priority weights (strict)
Example with 100 RPM model, 60% priority allocation, 80% threshold:
- Saturation < 80%: Priority can use up to 100 RPM (model limit enforced only)
- Saturation >= 80%: Priority limited to 60 RPM (both limits enforced)
Prevents bugs where:
- Model counter increments but priority check fails → model over-capacity
- Priority counter increments but not enforced → inaccurate metrics
Args:
user_api_key_dict: User authentication and metadata
cache: Dual cache instance
data: Request data containing model name
call_type: Type of API call being made
Returns:
None if request is allowed, otherwise raises HTTPException
"""
@ -522,29 +518,26 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
model = data["model"]
key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None)
# Get model configuration
model_group_info: Optional[ModelGroupInfo] = (
self.llm_router.get_model_group_info(model_group=model)
model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info(
model_group=model
)
if model_group_info is None:
verbose_proxy_logger.debug(
f"No model group info for {model}, allowing request"
)
verbose_proxy_logger.debug(f"No model group info for {model}, allowing request")
return None
try:
# STEP 1: Check current saturation level
saturation = await self._check_model_saturation(model, model_group_info)
saturation_threshold = (
litellm.priority_reservation_settings.saturation_threshold
)
saturation_threshold = litellm.priority_reservation_settings.saturation_threshold
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={key_priority}"
)
# STEP 2: Check rate limits in THREE phases
# Phase 1: Read-only check of ALL limits (no increments)
@ -559,7 +552,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
saturation=saturation,
data=data,
)
except HTTPException:
raise
except Exception as e:
@ -586,22 +579,15 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Add additional priority-specific headers
if isinstance(response, ModelResponse):
key_priority: Optional[str] = user_api_key_dict.metadata.get(
"priority", None
)
key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None)
# Get existing additional headers
additional_headers = (
getattr(response, "_hidden_params", {}).get(
"additional_headers", {}
)
or {}
)
additional_headers = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {}
# Add priority information
additional_headers["x-litellm-priority"] = key_priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
# Update response
if not hasattr(response, "_hidden_params"):
response._hidden_params = {}

30
litellm/proxy/hooks/responses_id_security.py
View File

@ -5,7 +5,16 @@ This hook uses the DBSpendUpdateWriter to batch-write response IDs to the databa
instead of writing immediately on each request.
"""
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Tuple, Union, cast
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Literal,
Optional,
Tuple,
Union,
cast,
)
from fastapi import HTTPException
@ -20,7 +29,7 @@ from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
ResponsesAPIResponse,
)
from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, SpecialEnums
from litellm.types.utils import LLMResponseTypes, SpecialEnums
if TYPE_CHECKING:
from litellm.caching.caching import DualCache
@ -36,7 +45,18 @@ class ResponsesIDSecurity(CustomLogger):
user_api_key_dict: "UserAPIKeyAuth",
cache: "DualCache",
data: dict,
call_type: CallTypesLiteral,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Optional[Union[Exception, str, dict]]:
# MAP all the responses api response ids to the encrypted response ids
responses_api_call_types = {
@ -116,7 +136,7 @@ class ResponsesIDSecurity(CustomLogger):
split_result = response_id.split("resp_")
if len(split_result) < 2:
return False
remaining_string = split_result[1]
decrypted_value = decrypt_value_helper(
value=remaining_string, key="response_id", return_original_value=True
@ -141,7 +161,7 @@ class ResponsesIDSecurity(CustomLogger):
split_result = response_id.split("resp_")
if len(split_result) < 2:
return response_id, None, None
remaining_string = split_result[1]
decrypted_value = decrypt_value_helper(
value=remaining_string, key="response_id", return_original_value=True