feat: add event_type parameter to add_standard_logging_guardrail_information_to_request_data

This commit is contained in:
Yuta Saito 2025-12-27 07:16:30 +09:00
parent 6888d34ea5
commit 8bf0b9e19b
8 changed files with 84 additions and 10 deletions

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:

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(

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]]:
"""
@ -172,6 +173,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 +194,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
raise
@ -199,6 +202,7 @@ class IBMGuardrailDetector(CustomGuardrail):
async def _call_orchestrator(
self,
content: str,
event_type: GuardrailEventHooks,
request_data: Optional[dict] = None,
) -> List[IBMDetectorDetection]:
"""
@ -258,6 +262,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 +283,7 @@ class IBMGuardrailDetector(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
raise
@ -472,6 +478,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 +507,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 +565,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 +594,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 +683,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 +713,7 @@ class IBMGuardrailDetector(CustomGuardrail):
orchestrator_result = await self._call_orchestrator(
content=content,
request_data=data,
event_type=GuardrailEventHooks.post_call,
)
verbose_proxy_logger.debug(

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 = ""

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,
)
#########################################################

View file

@ -327,6 +327,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 +352,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
duration=duration,
start_time=start_time,
end_time=end_time,
event_type=event_type,
)
return response

View file

@ -163,6 +163,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 +214,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 +244,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 +296,7 @@ class NomaGuardrail(CustomGuardrail):
start_time=start_time.timestamp(),
end_time=end_time.timestamp(),
duration=duration,
event_type=event_type,
)
if self.monitor_mode:
@ -602,7 +606,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 +625,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 +657,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 +676,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 +710,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 +729,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 +741,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 +757,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 +875,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

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