tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-01 10:20:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18489 from BerriAI/litellm_feat_guardrail-log-actual-event-type

Browse Source 
Litellm feat guardrail log actual event type
This commit is contained in:
YutaSaito
2025-12-29 07:18:21 +09:00
committed by GitHub
parent 861b103dc0 c6d6ffd567
commit 61f64f4e7d
10 changed files with 371 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/custom_guardrail.py
+39 -8
View File
@@ -243,14 +243,14 @@ class CustomGuardrail(CustomLogger):
def _is_valid_response_type(self, result: Any) -> bool:
"""
Check if result is a valid LLMResponseTypes instance.
Safely handles TypedDict types which don't support isinstance checks.
For non-LiteLLM responses (like passthrough httpx.Response), returns True
to allow them through.
"""
if result is None:
return False
try:
# Try isinstance check on valid types that support it
response_types = get_args(LLMResponseTypes)
@@ -506,6 +506,7 @@ class CustomGuardrail(CustomLogger):
duration: Optional[float] = None,
masked_entity_count: Optional[Dict[str, int]] = None,
guardrail_provider: Optional[str] = None,
event_type: Optional[GuardrailEventHooks] = None,
) -> None:
"""
Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc.
@@ -514,14 +515,19 @@ class CustomGuardrail(CustomLogger):
guardrail_json_response = str(guardrail_json_response)
from litellm.types.utils import GuardrailMode
# Use event_type if provided, otherwise fall back to self.event_hook
guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]]
if event_type is not None:
guardrail_mode = event_type
elif isinstance(self.event_hook, Mode):
guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) # type: ignore[typeddict-item]
else:
guardrail_mode = self.event_hook # type: ignore[assignment]
slg = StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider=guardrail_provider,
guardrail_mode=(
GuardrailMode(**self.event_hook.model_dump()) # type: ignore
if isinstance(self.event_hook, Mode)
else self.event_hook
),
guardrail_mode=guardrail_mode,
guardrail_response=guardrail_json_response,
guardrail_status=guardrail_status,
start_time=start_time,
@@ -589,6 +595,7 @@ class CustomGuardrail(CustomLogger):
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
event_type: Optional[GuardrailEventHooks] = None,
):
"""
Add StandardLoggingGuardrailInformation to the request data
@@ -605,6 +612,7 @@ class CustomGuardrail(CustomLogger):
duration=duration,
start_time=start_time,
end_time=end_time,
event_type=event_type,
)
return response
@@ -615,6 +623,7 @@ class CustomGuardrail(CustomLogger):
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
event_type: Optional[GuardrailEventHooks] = None,
):
"""
Add StandardLoggingGuardrailInformation to the request data
@@ -628,6 +637,7 @@ class CustomGuardrail(CustomLogger):
duration=duration,
start_time=start_time,
end_time=end_time,
event_type=event_type,
)
raise e
@@ -712,16 +722,32 @@ def log_guardrail_information(func):
Logs for:
- pre_call
- during_call
- TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run
- post_call
"""
import asyncio
import functools
def _infer_event_type_from_function_name(
func_name: str,
) -> Optional[GuardrailEventHooks]:
"""Infer the actual event type from the function name"""
if func_name == "async_pre_call_hook":
return GuardrailEventHooks.pre_call
elif func_name == "async_moderation_hook":
return GuardrailEventHooks.during_call
elif func_name in (
"async_post_call_success_hook",
"async_post_call_streaming_hook",
):
return GuardrailEventHooks.post_call
return None
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = datetime.now() # Move start_time inside the wrapper
self: CustomGuardrail = args[0]
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
event_type = _infer_event_type_from_function_name(func.__name__)
try:
response = await func(*args, **kwargs)
return self._process_response(
@@ -730,6 +756,7 @@ def log_guardrail_information(func):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
except Exception as e:
return self._process_error(
@@ -738,6 +765,7 @@ def log_guardrail_information(func):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
@functools.wraps(func)
@@ -745,18 +773,21 @@ def log_guardrail_information(func):
start_time = datetime.now() # Move start_time inside the wrapper
self: CustomGuardrail = args[0]
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
event_type = _infer_event_type_from_function_name(func.__name__)
try:
response = func(*args, **kwargs)
return self._process_response(
response=response,
request_data=request_data,
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
except Exception as e:
return self._process_error(
e=e,
request_data=request_data,
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
@functools.wraps(func)
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+17 -9
View File
@@ -449,6 +449,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prepared_request.headers,
)
event_type = (
GuardrailEventHooks.pre_call
if source == "INPUT"
else GuardrailEventHooks.post_call
)
try:
httpx_response = await self.async_handler.post(
url=prepared_request.url,
@@ -469,6 +475,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
# Re-raise the exception to maintain existing behavior
raise
@@ -486,6 +493,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
#########################################################
if httpx_response.status_code == 200:
@@ -605,10 +613,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions.
If `self.mask_request_content` or `self.mask_response_content` is set to `True`,
If `self.mask_request_content` or `self.mask_response_content` is set to `True`,
then use the output from the guardrail to mask the request or response content.
However, even with masking enabled, content with action="BLOCKED" should still
However, even with masking enabled, content with action="BLOCKED" should still
raise an exception, only content with action="ANONYMIZED" should be masked.
"""
@@ -731,9 +739,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = (
None
)
bedrock_guardrail_response: Optional[
Union[BedrockGuardrailResponse, str]
] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
@@ -803,9 +811,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = (
None
)
bedrock_guardrail_response: Optional[
Union[BedrockGuardrailResponse, str]
] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
+6
View File
@@ -97,6 +97,7 @@ class DynamoAIGuardrails(CustomGuardrail):
async def _call_dynamoai_guardrails(
self,
messages: List[Dict[str, Any]],
event_type: GuardrailEventHooks,
text_type: str = "input",
request_data: Optional[dict] = None,
) -> DynamoAIResponse:
@@ -157,6 +158,7 @@ class DynamoAIGuardrails(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
return response_json
@@ -177,6 +179,7 @@ class DynamoAIGuardrails(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
raise
@@ -332,6 +335,7 @@ class DynamoAIGuardrails(CustomGuardrail):
messages=_messages,
text_type="input",
request_data=data,
event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.debug(
@@ -380,6 +384,7 @@ class DynamoAIGuardrails(CustomGuardrail):
messages=_messages,
text_type="input",
request_data=data,
event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.debug(
@@ -460,6 +465,7 @@ class DynamoAIGuardrails(CustomGuardrail):
messages=dynamoai_messages,
text_type="output",
request_data=data,
event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(
litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
+12 -1
View File
@@ -108,6 +108,7 @@ class IBMGuardrailDetector(CustomGuardrail):
async def _call_detector_server(
self,
contents: List[str],
event_type: GuardrailEventHooks,
request_data: Optional[dict] = None,
) -> List[List[IBMDetectorDetection]]:
"""
@@ -142,7 +143,6 @@ class IBMGuardrailDetector(CustomGuardrail):
)
try:
response = await self.async_handler.post(
url=self.api_url,
json=payload,
@@ -172,6 +172,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
return response_json
@@ -192,6 +193,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
raise
@@ -199,6 +201,7 @@ class IBMGuardrailDetector(CustomGuardrail):
async def _call_orchestrator(
self,
content: str,
event_type: GuardrailEventHooks,
request_data: Optional[dict] = None,
) -> List[IBMDetectorDetection]:
"""
@@ -258,6 +261,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
return response_json.get("detections", [])
@@ -278,6 +282,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
raise
@@ -472,6 +477,7 @@ class IBMGuardrailDetector(CustomGuardrail):
result = await self._call_detector_server(
contents=contents_to_check,
request_data=data,
event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.debug(
@@ -500,6 +506,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.debug(
@@ -557,6 +564,7 @@ class IBMGuardrailDetector(CustomGuardrail):
result = await self._call_detector_server(
contents=contents_to_check,
request_data=data,
event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.debug(
@@ -585,6 +593,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.debug(
@@ -673,6 +682,7 @@ class IBMGuardrailDetector(CustomGuardrail):
result = await self._call_detector_server(
contents=contents_to_check,
request_data=data,
event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(
@@ -702,6 +712,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(
litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py
+5 -1
View File
@@ -83,6 +83,7 @@ class JavelinGuardrail(CustomGuardrail):
async def call_javelin_guard(
self,
request: JavelinGuardRequest,
event_type: GuardrailEventHooks,
) -> JavelinGuardResponse:
"""
Call the Javelin guard API.
@@ -158,6 +159,7 @@ class JavelinGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
event_type=event_type,
)
async def async_pre_call_hook(
@@ -208,7 +210,9 @@ class JavelinGuardrail(CustomGuardrail):
config=self.config if self.config else {},
)
javelin_response = await self.call_javelin_guard(request=javelin_guard_request)
javelin_response = await self.call_javelin_guard(
request=javelin_guard_request, event_type=GuardrailEventHooks.pre_call
)
assessments = javelin_response.get("assessments", [])
reject_prompt = ""
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
+4
View File
@@ -70,6 +70,7 @@ class LakeraAIGuardrail(CustomGuardrail):
self,
messages: List[AllMessageValues],
request_data: Dict,
event_type: GuardrailEventHooks,
) -> Tuple[LakeraAIResponse, Dict]:
"""
Call the Lakera AI v2 guard API.
@@ -128,6 +129,7 @@ class LakeraAIGuardrail(CustomGuardrail):
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
masked_entity_count=masked_entity_count,
event_type=event_type,
)
def _mask_pii_in_messages(
@@ -214,6 +216,7 @@ class LakeraAIGuardrail(CustomGuardrail):
lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
messages=new_messages,
request_data=data,
event_type=GuardrailEventHooks.pre_call,
)
#########################################################
@@ -279,6 +282,7 @@ class LakeraAIGuardrail(CustomGuardrail):
lakera_guardrail_response, masked_entity_count = await self.call_v2_guard(
messages=new_messages,
request_data=data,
event_type=GuardrailEventHooks.during_call,
)
#########################################################
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
+5 -1
View File
@@ -295,7 +295,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
filters = (
list(filter_results.values())
if isinstance(filter_results, dict)
else filter_results if isinstance(filter_results, list) else []
else filter_results
if isinstance(filter_results, list)
else []
)
# Prefer sanitized text from deidentifyResult if present
@@ -327,6 +329,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
event_type: Optional[GuardrailEventHooks] = None,
):
"""
Override to store only the Model Armor API response, not the entire data dict.
@@ -351,6 +354,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
duration=duration,
start_time=start_time,
end_time=end_time,
event_type=event_type,
)
return response
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
+27 -10
View File
@@ -119,9 +119,7 @@ class NomaGuardrail(CustomGuardrail):
self.api_base = api_base or os.environ.get(
"NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE
)
self.application_id = application_id or os.environ.get(
"NOMA_APPLICATION_ID"
)
self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID")
self.default_application_id = "litellm"
if monitor_mode is None:
@@ -163,6 +161,7 @@ class NomaGuardrail(CustomGuardrail):
self,
request_data: dict,
user_auth: UserAPIKeyAuth,
event_type: Optional[GuardrailEventHooks] = None,
) -> Optional[str]:
"""Shared logic for processing user message checks"""
start_time = datetime.now()
@@ -213,6 +212,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
if self.monitor_mode:
@@ -242,6 +242,7 @@ class NomaGuardrail(CustomGuardrail):
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
event_type: Optional[GuardrailEventHooks] = None,
) -> Optional[str]:
"""Shared logic for processing LLM response checks"""
@@ -293,6 +294,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
if self.monitor_mode:
@@ -578,7 +580,6 @@ class NomaGuardrail(CustomGuardrail):
data: dict,
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
verbose_proxy_logger.debug("Running Noma pre-call hook")
if (
@@ -602,7 +603,9 @@ class NomaGuardrail(CustomGuardrail):
return data
try:
return await self._check_user_message(data, user_api_key_dict)
return await self._check_user_message(
data, user_api_key_dict, GuardrailEventHooks.pre_call
)
except NomaBlockedMessage:
# Blocked requests were already logged in _process_user_message_check with "blocked" status
raise
@@ -619,6 +622,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
event_type=GuardrailEventHooks.pre_call,
)
verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}")
@@ -650,7 +654,9 @@ class NomaGuardrail(CustomGuardrail):
return data
try:
return await self._check_user_message(data, user_api_key_dict)
return await self._check_user_message(
data, user_api_key_dict, GuardrailEventHooks.during_call
)
except NomaBlockedMessage:
# Blocked requests were already logged in _process_user_message_check with "blocked" status
raise
@@ -667,6 +673,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
event_type=GuardrailEventHooks.during_call,
)
verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}")
@@ -700,7 +707,9 @@ class NomaGuardrail(CustomGuardrail):
return response
try:
return await self._check_llm_response(data, response, user_api_key_dict)
return await self._check_llm_response(
data, response, user_api_key_dict, GuardrailEventHooks.post_call
)
except NomaBlockedMessage:
# Blocked requests were already logged in _process_llm_response_check with "blocked" status
raise
@@ -717,6 +726,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=start_time.timestamp(),
duration=0.0,
event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}")
@@ -728,9 +738,12 @@ class NomaGuardrail(CustomGuardrail):
self,
request_data: dict,
user_auth: UserAPIKeyAuth,
event_type: Optional[GuardrailEventHooks] = None,
) -> Union[Exception, str, dict, None]:
"""Check user message for policy violations"""
user_message = await self._process_user_message_check(request_data, user_auth)
user_message = await self._process_user_message_check(
request_data, user_auth, event_type
)
if not user_message:
return request_data
@@ -741,10 +754,11 @@ class NomaGuardrail(CustomGuardrail):
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
event_type: Optional[GuardrailEventHooks] = None,
) -> Any:
"""Check LLM response for policy violations"""
content = await self._process_llm_response_check(
request_data, response, user_auth
request_data, response, user_auth, event_type
)
if not content:
return response
@@ -858,7 +872,10 @@ class NomaGuardrail(CustomGuardrail):
if isinstance(assembled_model_response, ModelResponse):
try:
processed_response = await self._check_llm_response(
request_data, assembled_model_response, user_api_key_dict
request_data,
assembled_model_response,
user_api_key_dict,
GuardrailEventHooks.post_call,
)
except NomaBlockedMessage:
raise
litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py
+21 -3
View File
@@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypesLiteral, ModelResponse
if TYPE_CHECKING:
@@ -523,6 +524,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
scan_result: Dict[str, Any],
data: Dict[str, Any],
start_time: datetime,
event_type: GuardrailEventHooks,
is_response: bool = False,
) -> Optional[Dict[str, Any]]:
"""Handle API errors with fail-open/fail-closed logic."""
@@ -542,6 +544,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
if scan_result.get("_always_block"):
@@ -735,7 +738,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if scan_result.get("_is_transient") or scan_result.get("_always_block"):
return self._handle_api_error_with_logging(
scan_result, data, start_time, is_response=False
scan_result,
data,
start_time,
is_response=False,
event_type=GuardrailEventHooks.pre_call,
)
end_time = datetime.now()
@@ -749,6 +756,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.pre_call,
)
action = scan_result.get("action", "block")
@@ -872,7 +880,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if scan_result.get("_is_transient") or scan_result.get("_always_block"):
self._handle_api_error_with_logging(
scan_result, data, start_time, is_response=True
scan_result,
data,
start_time,
is_response=True,
event_type=GuardrailEventHooks.post_call,
)
return response
@@ -887,6 +899,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
event_type=GuardrailEventHooks.post_call,
)
action = scan_result.get("action", "block")
@@ -1066,7 +1079,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
if scan_result.get("_is_transient") or scan_result.get("_always_block"):
self._handle_api_error_with_logging(
scan_result, request_data, start_time, is_response=True
scan_result,
request_data,
start_time,
is_response=True,
event_type=EventHooks.post_call,
)
for chunk in all_chunks:
yield chunk
@@ -1083,6 +1100,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=(end_time - start_time).total_seconds(),
event_type=EventHooks.post_call,
)
# Add guardrail to applied guardrails header for observability
tests/test_litellm/integrations/test_custom_guardrail.py
+235 -3
View File
@@ -498,7 +498,7 @@ class TestPassthroughCallTypeHandling:
def test_get_pre_call_type_with_allm_passthrough_route(self):
"""
Test that _get_pre_call_type correctly maps allm_passthrough_route.
This tests Fix #1: allm_passthrough_route was not being handled, causing call_type to be None.
"""
from litellm.proxy.common_request_processing import (
@@ -509,14 +509,14 @@ class TestPassthroughCallTypeHandling:
result = ProxyBaseLLMRequestProcessing._get_pre_call_type(
route_type="allm_passthrough_route"
)
# Should return allm_passthrough_route, not None
assert result == "allm_passthrough_route"
def test_get_pre_call_type_preserves_standard_mappings(self):
"""
Test that _get_pre_call_type still correctly maps standard route types.
Ensures Fix #1 didn't break existing functionality.
"""
from litellm.proxy.common_request_processing import (
@@ -536,3 +536,235 @@ class TestPassthroughCallTypeHandling:
ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses")
== "responses"
)
class TestEventTypeLogging:
"""Tests for event_type logging in guardrail information."""
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call
from async_pre_call_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_pre_call_hook(self, data: dict, **kwargs):
return {"result": "pre_call_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_pre_call_hook(data=request_data)
# Check that the guardrail_mode was set to pre_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
from async_post_call_success_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_post_call_success_hook(self, data: dict, **kwargs):
return {"result": "post_call_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_post_call_success_hook(data=request_data)
# Check that the guardrail_mode was set to post_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call
from async_moderation_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.during_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_moderation_hook(self, data: dict, **kwargs):
return {"result": "moderation_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_moderation_hook(data=request_data)
# Check that the guardrail_mode was set to during_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
from async_post_call_streaming_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_post_call_streaming_hook(self, data: dict, **kwargs):
return {"result": "streaming_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_post_call_streaming_hook(data=request_data)
# Check that the guardrail_mode was set to post_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
@pytest.mark.asyncio
async def test_log_guardrail_information_returns_none_for_unknown_function_name(
self,
):
"""
Test that log_guardrail_information decorator returns None for event_type
when function name doesn't match known patterns, and falls back to self.event_hook.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def some_other_hook(self, data: dict, **kwargs):
return {"result": "other_hook_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.some_other_hook(data=request_data)
# Check that the guardrail_mode falls back to self.event_hook
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
def test_add_standard_logging_uses_event_type_over_event_hook(self):
"""
Test that add_standard_logging_guardrail_information_to_request_data
prioritizes event_type parameter over self.event_hook.
"""
from litellm.types.guardrails import GuardrailEventHooks
guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
)
request_data = {"metadata": {}}
# Call with explicit event_type
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
event_type=GuardrailEventHooks.post_call,
)
# Should use the provided event_type (post_call), not the full event_hook list
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
self,
):
"""
Test that add_standard_logging_guardrail_information_to_request_data
falls back to self.event_hook when event_type is None.
"""
from litellm.types.guardrails import GuardrailEventHooks
guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
request_data = {"metadata": {}}
# Call with event_type=None
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
event_type=None,
)
# Should fall back to self.event_hook
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call