allow different guardrailing api key

gateway/common/authorization.py

@@ -4,9 +4,22 @@ from typing import Tuple, Optional
 from fastapi import HTTPException, Request
 
 INVARIANT_AUTHORIZATION_HEADER = "invariant-authorization"
+INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER = "invariant-guardrails-authorization"
 API_KEYS_SEPARATOR = ";invariant-auth="
 
 
+def extract_guardrail_service_authorization_from_headers(
+    request: Request,
+) -> Tuple[Optional[str], Optional[str]]:
+    """
+    Extracts the optional Invariant-Guardrail-Service authorization header from the request.
+
+    This header can be specifified to use a different API key for guardrailing compared to
+    Explorer interactions.
+    """
+    return request.headers.get(INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER)
+
+
 def extract_authorization_from_headers(
     request: Request,
     dataset_name: Optional[str] = None,

gateway/common/request_context.py

@@ -7,6 +7,9 @@ import fastapi
 
 from common.config_manager import GatewayConfig
 from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction
+from gateway.common.authorization import (
+    extract_guardrail_service_authorization_from_headers,
+)
 
 
 @dataclass(frozen=True)
@@ -15,7 +18,10 @@ class RequestContext:
 
     request_json: Dict[str, Any]
     dataset_name: Optional[str] = None
+    # authorization to use for invariant service like explorer
     invariant_authorization: Optional[str] = None
+    # authorization to use for invariant guardrailing specifically
+    guardrail_authorization: Optional[str] = None
     # the set of guardrails to enforce for this request
     guardrails: Optional[GuardrailRuleSet] = None
     config: Dict[str, Any] = None
@@ -74,15 +80,37 @@ class RequestContext:
                 logging_guardrails=[],
             )
 
+        # if additionally provided, extract separate API key to use with guardrailing service
+        guardrail_service_authorization = None
+        if (
+            guardrail_authorization
+            := extract_guardrail_service_authorization_from_headers(request)
+        ):
+            guardrail_service_authorization = guardrail_authorization
+
         return cls(
             request_json=request_json,
             dataset_name=dataset_name,
             invariant_authorization=invariant_authorization,
+            guardrail_service_authorization=guardrail_service_authorization,
             guardrails=guardrails,
             config=context_config,
             _created_via_factory=True,
+            guardrail_authorization=guardrail_service_authorization,
         )
 
+    def get_guardrailing_authorization(self) -> Optional[str]:
+        """
+        Returns the authorization to use for the guardrailing service.
+
+        This can be different from the invariant authorization, but falls back
+        "to be the same if not explicitly set via header.
+
+        See also extract_guardrail_service_authorization_from_headers(...)
+        """
+
+        return self.guardrail_authorization or self.invariant_authorization
+
     def __repr__(self) -> str:
         return (
             f"RequestContext("

gateway/integrations/guardrails.py

@@ -82,10 +82,6 @@ async def _preload(guardrails: str, invariant_authorization: str) -> None:
         result.raise_for_status()
 
 
-def get_guardrails_invariant_authorization(context: "RequestContext") -> str:
-    return context.invariant_authorization
-
-
 async def preload_guardrails(context: "RequestContext") -> None:
     """
     Preloads the guardrails for faster checking later.
@@ -101,8 +97,7 @@ async def preload_guardrails(context: "RequestContext") -> None:
         for blocking_guardrail in context.guardrails.blocking_guardrails:
             task = asyncio.create_task(
                 _preload(
-                    blocking_guardrail.content,
-                    get_guardrails_invariant_authorization(context),
+                    blocking_guardrail.content, context.get_guardrailing_authorization()
                 )
             )
             asyncio.shield(task)
@@ -110,7 +105,7 @@ async def preload_guardrails(context: "RequestContext") -> None:
             task = asyncio.create_task(
                 _preload(
                     logging_guadrail.content,
-                    get_guardrails_invariant_authorization(context),
+                    context.get_guardrailing_authorization(),
                 )
             )
             asyncio.shield(task)
@@ -367,7 +362,7 @@ async def check_guardrails(
                     "policies": [g.content for g in guardrails],
                 },
                 headers={
-                    "Authorization": get_guardrails_invariant_authorization(context),
+                    "Authorization": context.get_guardrailing_authorization(),
                     "Accept": "application/json",
                 },
             )

gateway/routes/anthropic.py

@@ -105,6 +105,7 @@ async def anthropic_v1_messages_gateway(
         invariant_authorization=invariant_authorization,
         guardrails=header_guardrails or dataset_guardrails,
         config=config,
+        request=request,
     )
     if request_json.get("stream"):
         return await handle_streaming_response(context, client, anthropic_request)

gateway/routes/gemini.py

@@ -103,6 +103,7 @@ async def gemini_generate_content_gateway(
         invariant_authorization=invariant_authorization,
         guardrails=header_guardrails or dataset_guardrails,
         config=config,
+        request=request,
     )
     if alt == "sse" or endpoint == "streamGenerateContent":
         return await stream_response(