Adding output recipe, cleaning up some imports

Author: ryanmeans

litellm/proxy/guardrails/guardrail_hooks/pangea.py

@@ -1,17 +1,9 @@
 # litellm/proxy/guardrails/guardrail_hooks/pangea.py
 import os
-import sys
-
-# Adds the parent directory to the system path to allow importing litellm modules
-sys.path.insert(
-    0, os.path.abspath("../../..")
-)
-import json
 from typing import Any, List, Literal, Optional, Union
 
 from fastapi import HTTPException
 
-import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
@@ -50,7 +42,8 @@ class PangeaHandler(CustomGuardrail):
     def __init__(
         self,
         guardrail_name: str,
-        pangea_recipe: str,
+        pangea_input_recipe: Optional[str] = None,
+        pangea_output_recipe: Optional[str] = None,
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
         **kwargs,
@@ -80,20 +73,22 @@ def __init__(
             or os.environ.get("PANGEA_API_BASE")
             or "https://ai-guard.aws.us.pangea.cloud"
         )
-        self.pangea_recipe = pangea_recipe
+        self.pangea_input_recipe = pangea_input_recipe
+        self.pangea_output_recipe = pangea_output_recipe
         self.guardrail_endpoint = f"{self.api_base}/v1/text/guard"
 
         # Pass relevant kwargs to the parent class
         super().__init__(guardrail_name=guardrail_name, **kwargs)
         verbose_proxy_logger.info(
-            f"Initialized Pangea Guardrail: name={guardrail_name}, recipe={pangea_recipe}, api_base={self.api_base}"
+            f"Initialized Pangea Guardrail: name={guardrail_name}, recipe={pangea_input_recipe}, api_base={self.api_base}"
         )
 
     def _prepare_payload(
         self,
         messages: Optional[List[AllMessageValues]] = None,
         text_input: Optional[str] = None,
         request_data: Optional[dict] = None,
+        recipe: Optional[str] = None,
     ) -> dict:
         """
         Prepares the payload for the Pangea AI Guard API request.
@@ -107,9 +102,12 @@ def _prepare_payload(
             dict: The payload dictionary for the API request.
         """
         payload: dict[str, Any] = {
-            "recipe": self.pangea_recipe,
             "debug": False,  # Or make this configurable if needed
         }
+
+        if recipe:
+            payload["recipe"] = recipe
+
         if messages:
             # Ensure messages are in the format Pangea expects (list of dicts with 'role' and 'content')
             payload["messages"] = [
@@ -253,7 +251,7 @@ async def async_moderation_hook(
 
         try:
             payload = self._prepare_payload(
-                messages=messages, text_input=text_input, request_data=data
+                messages=messages, text_input=text_input, request_data=data, recipe=self.pangea_input_recipe
             )
             await self._call_pangea_guard(
                 payload=payload, request_data=data, hook_name="moderation_hook"
@@ -303,7 +301,7 @@ async def async_post_call_success_hook(
         try:
             # Scan only the output text in the post-call hook
             payload = self._prepare_payload(
-                text_input=response_str, request_data=data
+                text_input=response_str, request_data=data, recipe=self.pangea_output_recipe
             )
             await self._call_pangea_guard(
                 payload=payload,
@@ -321,4 +319,4 @@ async def async_post_call_success_hook(
                     "error": f"Error preparing Pangea payload for response: {ve}",
                     "guardrail_name": self.guardrail_name,
                 },
-            )
+            )

litellm/proxy/guardrails/guardrail_initializers.py

@@ -121,9 +121,10 @@ def initialize_pangea(litellm_params, guardrail):
 
     _pangea_callback = PangeaHandler(
         guardrail_name=guardrail["guardrail_name"],
-        pangea_recipe=litellm_params["pangea_recipe"],
+        pangea_input_recipe=litellm_params["pangea_input_recipe"],
+        pangea_output_recipe=litellm_params["pangea_output_recipe"],
         api_base=litellm_params["api_base"],
         api_key=litellm_params["api_key"],
         default_on=litellm_params["default_on"],
     )
-    litellm.logging_callback_manager.add_litellm_callback(_pangea_callback)
+    litellm.logging_callback_manager.add_litellm_callback(_pangea_callback)

litellm/types/guardrails.py

@@ -115,7 +115,8 @@ class LitellmParams(TypedDict):
     ]  # will mask response content if guardrail makes any changes
 
     # pangea params
-    pangea_recipe: Optional[str]
+    pangea_input_recipe: Optional[str]
+    pangea_output_recipe: Optional[str]
 
 class Guardrail(TypedDict, total=False):
     guardrail_name: str