✨ (Feat) Log Guardrails run, guardrail response on logging integratio…

…ns (#7445) * add guardrail_information to SLP * use standard_logging_guardrail_information * track StandardLoggingGuardrailInformation * use log_guardrail_information * use log_guardrail_information * docs guardrails * docs guardrails * update quick start * fix presidio logging for sync functions * update Guardrail type * enforce add_standard_logging_guardrail_information_to_request_data * update gd docs
BerriAI · Dec 27, 2024 · 62753ee · 62753ee
1 parent d1686ee
commit 62753ee
Show file tree

Hide file tree

Showing 14 changed files with 223 additions and 29 deletions.
diff --git a/docs/my-website/docs/proxy/guardrails/quick_start.md b/docs/my-website/docs/proxy/guardrails/quick_start.md
@@ -112,14 +112,49 @@ curl -i http://localhost:4000/v1/chat/completions \
 
 </Tabs>
 
+## **Using Guardrails client side**
+
+
+### ✨ View available guardrails (/guardrails/list)
+
+Show available guardrails on the proxy server. This makes it easier for developers to know what guardrails are available / can be used.
+
+```shell
+curl -X GET 'http://0.0.0.0:4000/guardrails/list'
+```
+
+Expected response
+
+```json
+{
+    "guardrails": [
+        {
+        "guardrail_name": "bedrock-pre-guard",
+        "guardrail_info": {
+            "params": [
+            {
+                "name": "toxicity_score",
+                "type": "float",
+                "description": "Score between 0-1 indicating content toxicity level"
+            },
+            {
+                "name": "pii_detection",
+                "type": "boolean"
+            }
+            ]
+        }
+        }
+    ]
+}
+```
+
 
-## Advanced
 
 ### ✨ Pass additional parameters to guardrail
 
 :::info
 
-✨ This is an Enterprise only feature [Contact us to get a free trial](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat)
+✨ This is an Enterprise only feature [Get a free trial](https://www.litellm.ai/#trial)
 
 :::
 
@@ -196,11 +231,40 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 </Tabs>
 
 
+## **Proxy Admin Controls**
+
+### ✨ Monitoring Guardrails
+
+Monitor which guardrails were executed and whether they passed or failed. e.g. guardrail going rogue and failing requests we don't intend to fail
+
+:::info
+
+✨ This is an Enterprise only feature [Get a free trial](https://www.litellm.ai/#trial)
+
+:::
+
+### Setup
+
+1. Connect LiteLLM to a [supported logging provider](../logging)
+2. Make a request with a `guardrails` parameter
+3. Check your logging provider for the guardrail trace
+
+#### Traced Guardrail Success
+
+<Image img={require('../../../img/gd_success.png')} />
+
+#### Traced Guardrail Failure
+
+<Image img={require('../../../img/gd_fail.png')} />
+
+
+
+
 ### ✨ Control Guardrails per Project (API Key)
 
 :::info
 
-✨ This is an Enterprise only feature [Contact us to get a free trial](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat)
+✨ This is an Enterprise only feature [Get a free trial](https://www.litellm.ai/#trial)
 
 :::
 
@@ -262,7 +326,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 
 :::info
 
-✨ This is an Enterprise only feature [Contact us to get a free trial](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat)
+✨ This is an Enterprise only feature [Get a free trial](https://www.litellm.ai/#trial)
 
 :::
 
@@ -320,22 +384,6 @@ The `pii_masking` guardrail ran on this request because api key=sk-jNm1Zar7XfNdZ
 
 
 
-### ✨ List guardrails
-
-Show available guardrails on the proxy server. This makes it easier for developers to know what guardrails are available / can be used.
-
-```shell
-curl -X GET 'http://0.0.0.0:4000/guardrails/list'
-```
-
-Expected response
-
-```json
-{
-    "guardrails": ["aporia-pre-guard", "aporia-post-guard"]
-}
-```
-
 ## Spec: `guardrails` Parameter
 
 The `guardrails` parameter can be passed to any LiteLLM Proxy endpoint (`/chat/completions`, `/completions`, `/embeddings`).

diff --git a/docs/my-website/img/gd_fail.png b/docs/my-website/img/gd_fail.png
diff --git a/docs/my-website/img/gd_success.png b/docs/my-website/img/gd_success.png
diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py
@@ -1,8 +1,9 @@
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.types.guardrails import DynamicGuardrailParams, GuardrailEventHooks
+from litellm.types.utils import StandardLoggingGuardrailInformation
 
 
 class CustomGuardrail(CustomLogger):
@@ -119,3 +120,101 @@ def _validate_premium_user(self) -> bool:
             )
             return False
         return True
+
+    def add_standard_logging_guardrail_information_to_request_data(
+        self,
+        guardrail_json_response: Union[Exception, str, dict],
+        request_data: dict,
+        guardrail_status: Literal["success", "failure"],
+    ) -> None:
+        """
+        Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc.
+        """
+        from litellm.proxy.proxy_server import premium_user
+
+        if premium_user is not True:
+            verbose_logger.warning(
+                f"Guardrail Tracing is only available for premium users. Skipping guardrail logging for guardrail={self.guardrail_name} event_hook={self.event_hook}"
+            )
+            return
+        if isinstance(guardrail_json_response, Exception):
+            guardrail_json_response = str(guardrail_json_response)
+        slg = StandardLoggingGuardrailInformation(
+            guardrail_name=self.guardrail_name,
+            guardrail_mode=self.event_hook,
+            guardrail_response=guardrail_json_response,
+            guardrail_status=guardrail_status,
+        )
+        if "metadata" in request_data:
+            request_data["metadata"]["standard_logging_guardrail_information"] = slg
+        elif "litellm_metadata" in request_data:
+            request_data["litellm_metadata"][
+                "standard_logging_guardrail_information"
+            ] = slg
+        else:
+            verbose_logger.warning(
+                "unable to log guardrail information. No metadata found in request_data"
+            )
+
+
+def log_guardrail_information(func):
+    """
+    Decorator to add standard logging guardrail information to any function
+
+    Add this decorator to ensure your guardrail response is logged to DataDog, OTEL, s3, GCS etc.
+
+    Logs for:
+        - pre_call
+        - during_call
+        - TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run
+    """
+    import asyncio
+    import functools
+
+    def process_response(self, response, request_data):
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response=response,
+            request_data=request_data,
+            guardrail_status="success",
+        )
+        return response
+
+    def process_error(self, e, request_data):
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response=e,
+            request_data=request_data,
+            guardrail_status="failure",
+        )
+        raise e
+
+    @functools.wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        self: CustomGuardrail = args[0]
+        request_data: Optional[dict] = (
+            kwargs.get("data") or kwargs.get("request_data") or {}
+        )
+        try:
+            response = await func(*args, **kwargs)
+            return process_response(self, response, request_data)
+        except Exception as e:
+            return process_error(self, e, request_data)
+
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        self: CustomGuardrail = args[0]
+        request_data: Optional[dict] = (
+            kwargs.get("data") or kwargs.get("request_data") or {}
+        )
+        try:
+            response = func(*args, **kwargs)
+            return process_response(self, response, request_data)
+        except Exception as e:
+            return process_error(self, e, request_data)
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper(*args, **kwargs)
+        return sync_wrapper(*args, **kwargs)
+
+    return wrapper
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
@@ -3038,6 +3038,9 @@ def get_standard_logging_object_payload(
             response_cost_failure_debug_info=kwargs.get(
                 "response_cost_failure_debug_information"
             ),
+            guardrail_information=metadata.get(
+                "standard_logging_guardrail_information", None
+            ),
         )
 
         return payload

diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py
@@ -19,7 +19,10 @@
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,
+)
 from litellm.litellm_core_utils.logging_utils import (
     convert_litellm_response_object_to_str,
 )
@@ -142,6 +145,7 @@ async def make_aporia_api_request(
                     },
                 )
 
+    @log_guardrail_information
     async def async_post_call_success_hook(
         self,
         data: dict,
@@ -173,6 +177,7 @@ async def async_post_call_success_hook(
 
         pass
 
+    @log_guardrail_information
     async def async_moderation_hook(  ### 👈 KEY CHANGE ###
         self,
         data: dict,

diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
@@ -19,7 +19,10 @@
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,
+)
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
@@ -231,6 +234,7 @@ async def make_bedrock_api_request(
                 response.text,
             )
 
+    @log_guardrail_information
     async def async_moderation_hook(  ### 👈 KEY CHANGE ###
         self,
         data: dict,
@@ -263,6 +267,7 @@ async def async_moderation_hook(  ### 👈 KEY CHANGE ###
             )
             pass
 
+    @log_guardrail_information
     async def async_post_call_success_hook(
         self,
         data: dict,

diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py
@@ -3,7 +3,10 @@
 import litellm
 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
 
 
@@ -17,6 +20,7 @@ def __init__(
 
         super().__init__(**kwargs)
 
+    @log_guardrail_information
     async def async_pre_call_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -55,6 +59,7 @@ async def async_pre_call_hook(
 
         return data
 
+    @log_guardrail_information
     async def async_moderation_hook(
         self,
         data: dict,
@@ -84,6 +89,7 @@ async def async_moderation_hook(
                     if "litellm" in _content.lower():
                         raise ValueError("Guardrail failed words - `litellm` detected")
 
+    @log_guardrail_information
     async def async_post_call_success_hook(
         self,
         data: dict,

diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai.py
@@ -12,7 +12,10 @@
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,
+)
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_content_from_model_response,
 )
@@ -79,6 +82,7 @@ async def make_guardrails_ai_api_request(self, llm_output: str, request_data: di
             )
         return _json_response
 
+    @log_guardrail_information
     async def async_post_call_success_hook(
         self,
         data: dict,

diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
@@ -20,7 +20,10 @@
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_guardrail import (
+    CustomGuardrail,
+    log_guardrail_information,
+)
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -294,6 +297,7 @@ async def _check(  # noqa: PLR0915
             """
             self._check_response_flagged(response=response.json())
 
+    @log_guardrail_information
     async def async_pre_call_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -330,6 +334,7 @@ async def async_pre_call_hook(
             data=data, user_api_key_dict=user_api_key_dict, call_type=call_type
         )
 
+    @log_guardrail_information
     async def async_moderation_hook(  ### 👈 KEY CHANGE ###
         self,
         data: dict,