From f1111f9a1bc93ea9c21b0b986947259631320dc7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 4 Sep 2024 13:57:04 -0700
Subject: [PATCH] handle logging_only logic for guardrails

---
 litellm/litellm_core_utils/litellm_logging.py | 19 ++++++++++++++++++-
 litellm/types/guardrails.py                   |  5 +++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 537ca15a477..eb77a0a198c 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -25,6 +25,7 @@ from litellm import (
 )
 from litellm.caching import DualCache, InMemoryCache, S3Cache
 from litellm.cost_calculator import _select_model_name_for_cost_calc
+from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.redact_messages import (
     redact_message_input_output_from_logging,
@@ -1350,7 +1351,23 @@ class Logging:
         ## LOGGING HOOK ##
 
         for callback in callbacks:
-            if isinstance(callback, CustomLogger):
+            if isinstance(callback, CustomGuardrail):
+                from litellm.types.guardrails import GuardrailEventHooks
+
+                if (
+                    callback.should_run_guardrail(
+                        data=self.model_call_details,
+                        event_type=GuardrailEventHooks.logging_only,
+                    )
+                    is not True
+                ):
+                    self.model_call_details, result = await callback.async_logging_hook(
+                        kwargs=self.model_call_details,
+                        result=result,
+                        call_type=self.call_type,
+                    )
+                    continue
+            elif isinstance(callback, CustomLogger):
                 self.model_call_details, result = await callback.async_logging_hook(
                     kwargs=self.model_call_details,
                     result=result,
diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py
index 10f4be7e1eb..cb70de5052d 100644
--- a/litellm/types/guardrails.py
+++ b/litellm/types/guardrails.py
@@ -84,6 +84,10 @@ class LitellmParams(TypedDict, total=False):
     guardrailIdentifier: Optional[str]
     guardrailVersion: Optional[str]
 
+    # Presidio params
+    output_parse_pii: Optional[bool]
+    presidio_ad_hoc_recognizers: Optional[str]
+
 
 class Guardrail(TypedDict):
     guardrail_name: str
@@ -98,6 +102,7 @@ class GuardrailEventHooks(str, Enum):
     pre_call = "pre_call"
     post_call = "post_call"
     during_call = "during_call"
+    logging_only = "logging_only"
 
 
 class BedrockTextContent(TypedDict, total=False):