From eefd66496c54662718ecff295cd51a9e29826c29 Mon Sep 17 00:00:00 2001
From: Luca Beurer-Kellner <lucabeurerkellner@gmail.com>
Date: Thu, 3 Apr 2025 08:44:36 +0200
Subject: [PATCH] update guardrail metadata

---
 gateway/common/config_manager.py   | 22 ++++++++++++++--------
 gateway/common/request_context.py  | 12 +++++-------
 gateway/integrations/explorer.py   | 12 +++---------
 gateway/integrations/guardrails.py | 11 ++++++++---
 4 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py
index c303ee8..951e41b 100644
--- a/gateway/common/config_manager.py
+++ b/gateway/common/config_manager.py
@@ -95,16 +95,22 @@ class GatewayConfigManager:
 
 
 async def GuardrailsInHeader(request: fastapi.Request) -> Optional[GuardrailRuleSet]:
+    """
+    Extracts Invariant-Guardrails from the request header if provided, and returns a corresponding
+    GuardrailRuleSet. If no guardrails are provided, returns None.
+    """
     # if provided in header, use custom guardrailing policy
     if guardrails := extract_policy_from_headers(request):
+        guardrails = [
+            Guardrail(
+                id="guardrails-from-header",
+                name="guardrails from request header",
+                content=guardrails,
+                action=GuardrailAction.BLOCK,
+            )
+        ]
+
         return GuardrailRuleSet(
-            blocking_guardrails=[
-                Guardrail(
-                    id="guardrail-from-header",
-                    name="guardrails from request header",
-                    content=guardrails,
-                    action=GuardrailAction.BLOCK,
-                )
-            ],
+            blocking_guardrails=guardrails,
             logging_guardrails=[],
         )
diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py
index 0c4a31b..b252f31 100644
--- a/gateway/common/request_context.py
+++ b/gateway/common/request_context.py
@@ -46,9 +46,8 @@ class RequestContext:
             if key != "guardrails_from_file"
         }
 
-        # If no guardrails are configured for the dataset on Explorer,
-        # and the config specifies guardrails_from_file, use that.
-        guardrails = guardrails
+        # If no guardrails are configured and the config specifies
+        # guardrails_from_file, use those instead.
         if (
             (
                 not guardrails
@@ -60,12 +59,11 @@ class RequestContext:
             and config
             and config.guardrails_from_file
         ):
-            # TODO: Support logging guardrails via file.
             guardrails = GuardrailRuleSet(
                 blocking_guardrails=[
                     Guardrail(
-                        id="default",
-                        name="default",
+                        id="guardrails-from-gateway-config-file",
+                        name="guardrails from gateway configuration file",
                         content=config.guardrails_from_file,
                         action=GuardrailAction.BLOCK,
                     )
@@ -87,7 +85,7 @@ class RequestContext:
             f"RequestContext("
             f"request_json={self.request_json}, "
             f"dataset_name={self.dataset_name}, "
-            f"invariant_authorization={self.invariant_authorization}, "
+            f"invariant_authorization=inv-*****{self.invariant_authorization[-4:]}, "
             f"guardrails={self.guardrails}, "
             f"config={self.config})"
         )
diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py
index 3ba9f86..d48c1eb 100644
--- a/gateway/integrations/explorer.py
+++ b/gateway/integrations/explorer.py
@@ -50,16 +50,10 @@ def create_annotations_from_guardrails_errors(
                     address=r,
                     extra_metadata={
                         "source": "guardrails-error",
-                        # if included in error, also include the source of the guardrail
+                        # if included in error, also include information about guardrail source
                         **(
-                            {"guardrail_content": error.get("guardrail_content")}
-                            if error.get("guardrail_content")
-                            else {}
-                        ),
-                        # same for guardrail-action
-                        **(
-                            {"guardrail_action": error.get("guardrail_action")}
-                            if error.get("guardrail_action")
+                            {"guardrail": error.get("guardrail")}
+                            if error.get("guardrail")
                             else {}
                         ),
                     },
diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py
index ae6cf31..ec40651 100644
--- a/gateway/integrations/guardrails.py
+++ b/gateway/integrations/guardrails.py
@@ -370,12 +370,17 @@ async def check_guardrails(
             aggregated_errors = {"errors": []}
             for res, guardrail in zip(guardrails_result.get("result", []), guardrails):
                 for error in res.get("errors", []):
-                    # aggregated_errors["errors"].extend(res.get("errors", []))
+                    # add each error to the aggregated errors but keep track
+                    # of which guardrail it belongs to
                     aggregated_errors["errors"].append(
                         {
                             **error,
-                            "guardrail_content": guardrail.content,
-                            "guardrail_action": guardrail.action,
+                            "guardrail": {
+                                "id": guardrail.id,
+                                "name": guardrail.name,
+                                "content": guardrail.content,
+                                "action": guardrail.action,
+                            },
                         }
                     )