diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index e9e711f8023..979d04ac748 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -16,7 +16,10 @@ import tempfile from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import walk_user_text from litellm.types.utils import GenericGuardrailAPIInputs @@ -470,6 +473,7 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): ) return text + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, @@ -491,6 +495,7 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): return True #### CALL HOOKS - proxy only #### + @log_guardrail_information async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index da697a83b48..ee072846c27 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -1,6 +1,9 @@ import pytest from litellm_enterprise.enterprise_callbacks.secret_detection import _ENTERPRISE_SecretDetection +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth + @pytest.mark.asyncio async def test_apply_guardrail_redacts_secrets(): @@ -21,3 +24,31 @@ async def test_apply_guardrail_redacts_secrets(): assert "sk-1234998222" not in texts[0] assert "[REDACTED]" in texts[0] assert texts[1] == "this text has no secrets" + + +@pytest.mark.asyncio +async def test_async_pre_call_hook_records_guardrail_information(): + guard = _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets") + data = { + "messages": [ + { + "role": "user", + "content": "my key openai_api_key=sk-1234998222", + } + ], + "metadata": {}, + } + + await guard.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + guardrail_information = data["metadata"]["standard_logging_guardrail_information"] + assert isinstance(guardrail_information, list) + assert len(guardrail_information) == 1 + assert guardrail_information[0]["guardrail_status"] == "success" + assert "sk-1234998222" not in data["messages"][0]["content"] + assert "[REDACTED]" in data["messages"][0]["content"]