diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index ae6d646b5ae..305a5c20764 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -901,8 +901,16 @@ class CustomGuardrail(CustomLogger): for key, value in (litellm_params.get("metadata") or {}).items() if key != "standard_logging_guardrail_information" } + response: Final = ( + kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result + ) + output_translation: Final = ( + get_guardrail_translation_mapping(CallTypes.acompletion)() + if isinstance(response, ModelResponse) + else translation + ) try: - await self._scan_logged_call(kwargs, result, translation, scratch_metadata) + await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata) except Exception as e: verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) recorded: Final = scratch_metadata.get("standard_logging_guardrail_information") @@ -919,12 +927,11 @@ class CustomGuardrail(CustomLogger): async def _scan_logged_call( self, kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract - result: object, + response: object | None, translation: "BaseTranslation", + output_translation: "BaseTranslation", scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: - from litellm.llms import get_guardrail_translation_mapping - optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -936,16 +943,8 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - response: Final = ( - kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result - ) if response is None: return - output_translation: Final = ( - get_guardrail_translation_mapping(CallTypes.acompletion)() - if isinstance(response, ModelResponse) - else translation - ) await output_translation.process_output_response( response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request )