mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
feat: add event_type parameter to add_standard_logging_guardrail_information_to_request_data
This commit is contained in:
parent
6888d34ea5
commit
8bf0b9e19b
8 changed files with 84 additions and 10 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 = ""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue