Extract guardrails from header if provided (#33)

* guardrails from header

* use in-file guardrails in client.py

* support case without request

* remove client script

* tests: guardrailing rule passed in header

* include checked guardrails in annotation extra metadata

* include guardrailing action

* update guardrail metadata

gateway/common/authorization.py

@@ -9,10 +9,10 @@ API_KEYS_SEPARATOR = ";invariant-auth="
 
 def extract_authorization_from_headers(
     request: Request,
-    dataset_name: Optional[str],
-    llm_provider_api_key_header: str,
-    llm_provider_fallback_api_key_headers: list[str] | None = None
-) -> Tuple[str, str]:
+    dataset_name: Optional[str] = None,
+    llm_provider_api_key_header: Optional[str] = None,
+    llm_provider_fallback_api_key_headers: Optional[list[str]] = None,
+) -> Tuple[Optional[str], Optional[str]]:
     """
     Extracts the Invariant authorization and LLM Provider API key from the request headers.
 
@@ -22,7 +22,7 @@ def extract_authorization_from_headers(
     "invariant-authorization": "Bearer <Invariant API Key>"
     {llm_provider_api_key_header} contains the LLM Provider API Key as
     {llm_provider_api_key_header}: "<API Key>"
-    
+
     If {llm_provider_api_key_header} is not among headers, we look for
     any header among {llm_provider_fallback_api_key_headers}.
 
@@ -32,18 +32,23 @@ def extract_authorization_from_headers(
     The header in that case becomes:
     {llm_provider_api_key_header}: "<API Key>;invariant-auth=<Invariant API Key>"
     """
+    # invariant api key
     invariant_authorization = request.headers.get(INVARIANT_AUTHORIZATION_HEADER)
-    llm_provider_api_key = request.headers.get(llm_provider_api_key_header)
 
+    # llm provider api key (also check fallbacks for clients like litellm)
+    if llm_provider_api_key_header is not None:
+        llm_provider_api_key = request.headers.get(llm_provider_api_key_header)
 
-    if llm_provider_api_key is None and llm_provider_fallback_api_key_headers:
-        for header in llm_provider_fallback_api_key_headers:
-            llm_provider_api_key = request.headers.get(header)
-            if llm_provider_api_key:
-                llm_provider_api_key_header = header
-                break
-
+        if llm_provider_api_key is None and llm_provider_fallback_api_key_headers:
+            for header in llm_provider_fallback_api_key_headers:
+                llm_provider_api_key = request.headers.get(header)
+                if llm_provider_api_key:
+                    llm_provider_api_key_header = header
+                    break
+    else:
+        llm_provider_api_key = None
 
+    # if the dataset name is not None, we need to check if the invariant api key is present
     if dataset_name:
         if invariant_authorization is None:
             if llm_provider_api_key is None:
@@ -59,9 +64,7 @@ def extract_authorization_from_headers(
                 API_KEYS_SEPARATOR
             )
             if len(api_keys) != 2 or not api_keys[1].strip():
-                raise HTTPException(
-                    status_code=400, detail="Invalid API Key format"
-                )
+                raise HTTPException(status_code=400, detail="Invalid API Key format")
 
             invariant_authorization = f"Bearer {api_keys[1].strip()}"
             llm_provider_api_key = f"{api_keys[0].strip()}"

gateway/common/config_manager.py

@@ -3,9 +3,31 @@
 import asyncio
 import os
 import threading
+from typing import Optional
 
+import fastapi
 from httpx import HTTPStatusError
 
+from common.guardrails import Guardrail, GuardrailAction, GuardrailRuleSet
+from common.authorization import extract_authorization_from_headers
+
+
+def extract_policy_from_headers(request: Optional[fastapi.Request]) -> Optional[str]:
+    """
+    Extracts the guardrailing policy from the request headers if present.
+
+    Returns 'None' if no such header is present.
+    """
+    if request is None:
+        return None
+
+    policy = request.headers.get("Invariant-Guardrails")
+    # undo unicode_escape
+    if policy:
+        # interpret as bytes then decode
+        policy = policy.encode("utf-8").decode("unicode_escape")
+    return policy
+
 
 class GatewayConfig:
     """Common configurations for the Gateway Server."""
@@ -58,7 +80,7 @@ class GatewayConfigManager:
     _lock = threading.Lock()
 
     @classmethod
-    def get_config(cls):
+    def get_config(cls, request: fastapi.Request = None) -> GatewayConfig:
         """Initializes and returns the gateway configuration using double-checked locking."""
         local_config = cls._config_instance
 
@@ -68,4 +90,27 @@ class GatewayConfigManager:
                 if local_config is None:
                     local_config = GatewayConfig()
                     cls._config_instance = local_config
+
         return local_config
+
+
+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=guardrails,
+            logging_guardrails=[],
+        )

gateway/common/guardrails.py

@@ -24,7 +24,7 @@ class Guardrail:
 
 
 @dataclass(frozen=True)
-class DatasetGuardrails:
+class GuardrailRuleSet:
     """Grouped guardrail rules separated by their action."""
 
     blocking_guardrails: List[Guardrail]

gateway/common/request_context.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
 from common.config_manager import GatewayConfig
-from common.guardrails import DatasetGuardrails, Guardrail, GuardrailAction
+from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction
 
 
 @dataclass(frozen=True)
@@ -14,7 +14,8 @@ class RequestContext:
     request_json: Dict[str, Any]
     dataset_name: Optional[str] = None
     invariant_authorization: Optional[str] = None
-    dataset_guardrails: Optional[DatasetGuardrails] = None
+    # the set of guardrails to enforce for this request
+    guardrails: Optional[GuardrailRuleSet] = None
     config: Dict[str, Any] = None
 
     _created_via_factory: bool = field(
@@ -33,7 +34,7 @@ class RequestContext:
         request_json: Dict[str, Any],
         dataset_name: Optional[str] = None,
         invariant_authorization: Optional[str] = None,
-        dataset_guardrails: Optional[DatasetGuardrails] = None,
+        guardrails: Optional[GuardrailRuleSet] = None,
         config: Optional[GatewayConfig] = None,
     ) -> "RequestContext":
         """Creates a new RequestContext instance, applying default guardrails if needed."""
@@ -45,26 +46,24 @@ 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 = dataset_guardrails
+        # If no guardrails are configured and the config specifies
+        # guardrails_from_file, use those instead.
         if (
             (
-                not dataset_guardrails
+                not guardrails
                 or (
-                    not dataset_guardrails.blocking_guardrails
-                    and not dataset_guardrails.logging_guardrails
+                    not guardrails.blocking_guardrails
+                    and not guardrails.logging_guardrails
                 )
             )
             and config
             and config.guardrails_from_file
         ):
-            # TODO: Support logging guardrails via file.
-            guardrails = DatasetGuardrails(
+            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,
                     )
@@ -76,7 +75,7 @@ class RequestContext:
             request_json=request_json,
             dataset_name=dataset_name,
             invariant_authorization=invariant_authorization,
-            dataset_guardrails=guardrails,
+            guardrails=guardrails,
             config=context_config,
             _created_via_factory=True,
         )
@@ -86,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"dataset_guardrails={self.dataset_guardrails}, "
+            f"invariant_authorization=inv-*****{self.invariant_authorization[-4:]}, "
+            f"guardrails={self.guardrails}, "
             f"config={self.config})"
         )

gateway/integrations/explorer.py

@@ -3,7 +3,7 @@
 import os
 from typing import Any, Dict, List
 
-from common.guardrails import DatasetGuardrails, Guardrail, GuardrailAction
+from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction
 from invariant_sdk.async_client import AsyncClient
 from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesResponse
 from invariant_sdk.types.annotations import AnnotationCreate
@@ -50,7 +50,12 @@ def create_annotations_from_guardrails_errors(
                     address=r,
                     extra_metadata={
                         "source": "guardrails-error",
-                        "guardrail-action": action,
+                        # if included in error, also include information about guardrail source
+                        **(
+                            {"guardrail": error.get("guardrail")}
+                            if error.get("guardrail")
+                            else {}
+                        ),
                     },
                 )
             )
@@ -101,11 +106,11 @@ async def push_trace(
 
 async def fetch_guardrails_from_explorer(
     dataset_name: str, invariant_authorization: str
-) -> DatasetGuardrails:
+) -> GuardrailRuleSet:
     """Get the guardrails for the dataset.
 
     Returns:
-        DatasetGuardrails: The guardrails for the dataset grouped by their action.
+        GuardrailRuleSet: The guardrails for the dataset grouped by their action.
     """
 
     # TODO: Implement a single API in explorer backend which can return
@@ -134,7 +139,7 @@ async def fetch_guardrails_from_explorer(
     if policies_response.status_code != 200:
         if policies_response.status_code == 404:
             # If the dataset does not exist, return empty guardrails.
-            return DatasetGuardrails(
+            return GuardrailRuleSet(
                 blocking_guardrails=[],
                 logging_guardrails=[],
             )
@@ -169,7 +174,7 @@ async def fetch_guardrails_from_explorer(
         else:
             logging_guardrails.append(guardrail)
 
-    return DatasetGuardrails(
+    return GuardrailRuleSet(
         blocking_guardrails=blocking_guardrails,
         logging_guardrails=logging_guardrails,
     )

gateway/integrations/guardrails.py

@@ -89,17 +89,17 @@ async def preload_guardrails(context: "RequestContext") -> None:
     Args:
         context: RequestContext object.
     """
-    if not context.dataset_guardrails:
+    if not context.guardrails:
         return
 
     try:
         # Move these calls to a batch preload/validate API.
-        for blocking_guardrail in context.dataset_guardrails.blocking_guardrails:
+        for blocking_guardrail in context.guardrails.blocking_guardrails:
             task = asyncio.create_task(
                 _preload(blocking_guardrail.content, context.invariant_authorization)
             )
             asyncio.shield(task)
-        for logging_guadrail in context.dataset_guardrails.logging_guardrails:
+        for logging_guadrail in context.guardrails.logging_guardrails:
             task = asyncio.create_task(
                 _preload(logging_guadrail.content, context.invariant_authorization)
             )
@@ -365,13 +365,42 @@ async def check_guardrails(
                 raise Exception(
                     f"Guardrails check failed: {result.status_code} - {result.text}"
                 )
-            print(f"Guardrail check response: {result.json()}")
-
             guardrails_result = result.json()
+
             aggregated_errors = {"errors": []}
-            for res in guardrails_result.get("result", []):
-                aggregated_errors["errors"].extend(res.get("errors", []))
+            for res, guardrail in zip(guardrails_result.get("result", []), guardrails):
+                for error in 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": {
+                                "id": guardrail.id,
+                                "name": guardrail.name,
+                                "content": guardrail.content,
+                                "action": guardrail.action,
+                            },
+                        }
+                    )
+
+                # check for any error_message
+                if error_message := res.get("error_message"):
+                    return {
+                        "errors": [
+                            {"args": [error_message], "kwargs": {}, "ranges": []}
+                        ]
+                    }
             return aggregated_errors
         except Exception as e:
             print(f"Failed to verify guardrails: {e}")
-            return {"error": str(e)}
+            # make sure runtime errors are also visible in e.g. Explorer
+            return {
+                "errors": [
+                    {
+                        "args": ["Gateway: " + str(e)],
+                        "kwargs": {},
+                        "ranges": ["messages[0].content:L0"],
+                    }
+                ]
+            }

gateway/routes/anthropic.py

@@ -9,12 +9,16 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
 from starlette.responses import StreamingResponse
 
 from common.authorization import extract_authorization_from_headers
-from common.config_manager import GatewayConfig, GatewayConfigManager
+from common.config_manager import (
+    GatewayConfig,
+    GatewayConfigManager,
+    GuardrailsInHeader,
+)
 from common.constants import (
     CLIENT_TIMEOUT,
     IGNORED_HEADERS,
 )
-from common.guardrails import GuardrailAction
+from common.guardrails import GuardrailAction, GuardrailRuleSet
 from common.request_context import RequestContext
 from converters.anthropic_to_invariant import (
     convert_anthropic_to_invariant_message_format,
@@ -66,6 +70,7 @@ async def anthropic_v1_messages_gateway(
     request: Request,
     dataset_name: str = None,  # This is None if the client doesn't want to push to Explorer
     config: GatewayConfig = Depends(GatewayConfigManager.get_config),  # pylint: disable=unused-argument
+    header_guardrails: GuardrailRuleSet = Depends(GuardrailsInHeader),
 ):
     """Proxy calls to the Anthropic APIs"""
     headers = {
@@ -98,11 +103,9 @@ async def anthropic_v1_messages_gateway(
         request_json=request_json,
         dataset_name=dataset_name,
         invariant_authorization=invariant_authorization,
-        dataset_guardrails=dataset_guardrails,
+        guardrails=header_guardrails or dataset_guardrails,
         config=config,
     )
-    asyncio.create_task(preload_guardrails(context))
-
     if request_json.get("stream"):
         return await handle_streaming_response(context, client, anthropic_request)
     return await handle_non_streaming_response(context, client, anthropic_request)
@@ -140,9 +143,9 @@ async def get_guardrails_check_result(
     """Get the guardrails check result"""
     # Determine which guardrails to apply based on the action
     guardrails = (
-        context.dataset_guardrails.logging_guardrails
+        context.guardrails.logging_guardrails
         if action == GuardrailAction.LOG
-        else context.dataset_guardrails.blocking_guardrails
+        else context.guardrails.blocking_guardrails
     )
     if not guardrails:
         return {}
@@ -219,7 +222,7 @@ class InstrumentedAnthropicResponse(InstrumentedResponse):
 
     async def on_start(self):
         """Check guardrails in a pipelined fashion, before processing the first chunk (for input guardrailing)."""
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context, action=GuardrailAction.BLOCK, response_json={}
             )
@@ -300,7 +303,7 @@ class InstrumentedAnthropicResponse(InstrumentedResponse):
         assert self.response_json is not None, "response_json is None"
         assert self.response_string is not None, "response_string is None"
 
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             # Block on the guardrails check
             guardrails_execution_result = await get_guardrails_check_result(
                 self.context,
@@ -382,7 +385,7 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse):
 
     async def on_start(self):
         """Check guardrails in a pipelined fashion, before processing the first chunk (for input guardrailing)."""
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context,
                 action=GuardrailAction.BLOCK,
@@ -443,7 +446,7 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse):
         process_chunk(decoded_chunk, self.merged_response)
 
         # on last stream chunk, run output guardrails
-        if "event: message_stop" in decoded_chunk and self.context.dataset_guardrails:
+        if "event: message_stop" in decoded_chunk and self.context.guardrails:
             # Block on the guardrails check
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context,

gateway/routes/gemini.py

@@ -9,12 +9,16 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
 from fastapi.responses import StreamingResponse
 
 from common.authorization import extract_authorization_from_headers
-from common.config_manager import GatewayConfig, GatewayConfigManager
+from common.config_manager import (
+    GatewayConfig,
+    GatewayConfigManager,
+    GuardrailsInHeader,
+)
 from common.constants import (
     CLIENT_TIMEOUT,
     IGNORED_HEADERS,
 )
-from common.guardrails import GuardrailAction
+from common.guardrails import GuardrailAction, GuardrailRuleSet
 from common.request_context import RequestContext
 from converters.gemini_to_invariant import convert_request, convert_response
 from integrations.explorer import (
@@ -49,6 +53,7 @@ async def gemini_generate_content_gateway(
         None, title="Response Format", description="Set to 'sse' for streaming"
     ),
     config: GatewayConfig = Depends(GatewayConfigManager.get_config),  # pylint: disable=unused-argument
+    header_guardrails: GuardrailRuleSet = Depends(GuardrailsInHeader),
 ) -> Response:
     """Proxy calls to the Gemini GenerateContent API"""
     if endpoint not in ["generateContent", "streamGenerateContent"]:
@@ -91,11 +96,9 @@ async def gemini_generate_content_gateway(
         request_json=request_json,
         dataset_name=dataset_name,
         invariant_authorization=invariant_authorization,
-        dataset_guardrails=dataset_guardrails,
+        guardrails=header_guardrails or dataset_guardrails,
         config=config,
     )
-    asyncio.create_task(preload_guardrails(context))
-
     if alt == "sse" or endpoint == "streamGenerateContent":
         return await stream_response(
             context,
@@ -176,7 +179,7 @@ class InstrumentedStreamingGeminiResponse(InstrumentedStreamingResponse):
         Check guardrails in a pipelined fashion, before processing the first chunk
         (for input guardrailing).
         """
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context, action=GuardrailAction.BLOCK, response_json={}
             )
@@ -230,7 +233,7 @@ class InstrumentedStreamingGeminiResponse(InstrumentedStreamingResponse):
         if (
             self.merged_response.get("candidates", [])
             and self.merged_response.get("candidates")[0].get("finishReason", "")
-            and self.context.dataset_guardrails
+            and self.context.guardrails
         ):
             # Block on the guardrails check
             self.guardrails_execution_result = await get_guardrails_check_result(
@@ -377,9 +380,9 @@ async def get_guardrails_check_result(
     """Get the guardrails check result"""
     # Determine which guardrails to apply based on the action
     guardrails = (
-        context.dataset_guardrails.logging_guardrails
+        context.guardrails.logging_guardrails
         if action == GuardrailAction.LOG
-        else context.dataset_guardrails.blocking_guardrails
+        else context.guardrails.blocking_guardrails
     )
     if not guardrails:
         return {}
@@ -459,7 +462,7 @@ class InstrumentedGeminiResponse(InstrumentedResponse):
         Check guardrails in a pipelined fashion, before processing the first chunk
         (for input guardrailing).
         """
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context, action=GuardrailAction.BLOCK, response_json={}
             )
@@ -540,7 +543,7 @@ class InstrumentedGeminiResponse(InstrumentedResponse):
         response_string = json.dumps(self.response_json)
         response_code = self.response.status_code
 
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             # Block on the guardrails check
             guardrails_execution_result = await get_guardrails_check_result(
                 self.context,

gateway/routes/open_ai.py

@@ -9,12 +9,16 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
 from fastapi.responses import StreamingResponse
 
 from common.authorization import extract_authorization_from_headers
-from common.config_manager import GatewayConfig, GatewayConfigManager
+from common.config_manager import (
+    GatewayConfig,
+    GatewayConfigManager,
+    GuardrailsInHeader,
+)
 from common.constants import (
     CLIENT_TIMEOUT,
     IGNORED_HEADERS,
 )
-from common.guardrails import GuardrailAction
+from common.guardrails import GuardrailAction, GuardrailRuleSet
 from common.request_context import RequestContext
 from integrations.explorer import (
     create_annotations_from_guardrails_errors,
@@ -109,6 +113,7 @@ async def openai_chat_completions_gateway(
     request: Request,
     dataset_name: str = None,  # This is None if the client doesn't want to push to Explorer
     config: GatewayConfig = Depends(GatewayConfigManager.get_config),  # pylint: disable=unused-argument
+    header_guardrails: GuardrailRuleSet = Depends(GuardrailsInHeader),
 ) -> Response:
     """Proxy calls to the OpenAI APIs"""
     headers = {
@@ -142,11 +147,9 @@ async def openai_chat_completions_gateway(
         request_json=request_json,
         dataset_name=dataset_name,
         invariant_authorization=invariant_authorization,
-        dataset_guardrails=dataset_guardrails,
+        guardrails=header_guardrails or dataset_guardrails,
         config=config,
     )
-    asyncio.create_task(preload_guardrails(context))
-
     if request_json.get("stream", False):
         return await handle_stream_response(
             context,
@@ -203,7 +206,7 @@ class InstrumentedOpenAIStreamResponse(InstrumentedStreamingResponse):
         Check guardrails in a pipelined fashion, before processing the first chunk
         (for input guardrailing).
         """
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context,
                 action=GuardrailAction.BLOCK,
@@ -253,7 +256,7 @@ class InstrumentedOpenAIStreamResponse(InstrumentedStreamingResponse):
         )
 
         # check guardrails at the end of the stream (on the '[DONE]' SSE chunk.)
-        if "data: [DONE]" in chunk_text and self.context.dataset_guardrails:
+        if "data: [DONE]" in chunk_text and self.context.guardrails:
             # Block on the guardrails check
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context,
@@ -527,10 +530,11 @@ async def get_guardrails_check_result(
     """Get the guardrails check result"""
     # Determine which guardrails to apply based on the action
     guardrails = (
-        context.dataset_guardrails.logging_guardrails
+        context.guardrails.logging_guardrails
         if action == GuardrailAction.LOG
-        else context.dataset_guardrails.blocking_guardrails
+        else context.guardrails.blocking_guardrails
     )
+
     if not guardrails:
         return {}
 
@@ -578,7 +582,7 @@ class InstrumentedOpenAIResponse(InstrumentedResponse):
         Checks guardrails in a pipelined fashion, before processing
         the first chunk (for input guardrailing)
         """
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             # block on the guardrails check
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context, action=GuardrailAction.BLOCK
@@ -653,7 +657,7 @@ class InstrumentedOpenAIResponse(InstrumentedResponse):
         response_code = self.response.status_code
 
         # if we have guardrails, check the response
-        if self.context.dataset_guardrails:
+        if self.context.guardrails:
             # run guardrails again, this time on request + response
             self.guardrails_execution_result = await get_guardrails_check_result(
                 self.context,

tests/integration/guardrails/test_header_guardrails.py

@@ -0,0 +1,185 @@
+"""Test the guardrails from header with the OpenAI route."""
+
+import os
+import sys
+import uuid
+import time
+
+# Add integration folder (parent) to sys.path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import pytest
+import requests
+from httpx import Client
+from openai import OpenAI, BadRequestError, APIError
+
+# Pytest plugins
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+@pytest.mark.parametrize(
+    "do_stream, push_to_explorer",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_input_guardrail_in_header(
+    explorer_api_url, gateway_url, do_stream, push_to_explorer
+):
+    """Test the message content guardrail."""
+    if not os.getenv("INVARIANT_API_KEY"):
+        pytest.fail("No INVARIANT_API_KEY set, failing")
+
+    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
+
+    policy = """
+# For input guardrailing specifically
+raise "Users must not mention the magic phrase 'Abracadabra'" if:
+    (msg: Message)
+    msg.role == "user"
+    "Abracadabra" in msg.content
+"""
+
+    client = OpenAI(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}",
+                "Invariant-Guardrails": policy.encode("unicode-escape"),
+            },
+        ),
+        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway/openai",
+    )
+
+    request = {
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "Tell me more about Abracadabra."}],
+    }
+
+    if not do_stream:
+        with pytest.raises(BadRequestError) as exc_info:
+            chat_response = client.chat.completions.create(
+                **request,
+                stream=False,
+            )
+
+        assert exc_info.value.status_code == 400
+        assert "[Invariant] The request did not pass the guardrails" in str(
+            exc_info.value
+        )
+        assert "Users must not mention the magic phrase 'Abracadabra'" in str(
+            exc_info.value
+        )
+
+    else:
+        with pytest.raises(APIError) as exc_info:
+            chat_response = client.chat.completions.create(
+                **request,
+                stream=True,
+            )
+
+            for _ in chat_response:
+                pass
+        assert (
+            "[Invariant] The request did not pass the guardrails"
+            in exc_info.value.message
+        )
+        assert "Users must not mention the magic phrase 'Abracadabra'" in str(
+            exc_info.value.body
+        )
+
+    if push_to_explorer:
+        # Wait for the trace to be saved
+        # This is needed because the trace is saved asynchronously
+        time.sleep(2)
+
+        # Fetch the trace ids for the dataset
+        traces_response = requests.get(
+            f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces",
+            timeout=5,
+        )
+        traces = traces_response.json()
+        assert len(traces) == 1
+        trace_id = traces[0]["id"]
+
+        # Fetch the trace
+        trace_response = requests.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}",
+            timeout=5,
+        )
+        trace = trace_response.json()
+
+        # in case of input guardrailing, the pushed trace will not contain a response
+        assert len(trace["messages"]) == 1
+        assert trace["messages"][0] == {
+            "role": "user",
+            "content": "Tell me more about Abracadabra.",
+        }
+
+        # Fetch annotations
+        annotations_response = requests.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}/annotations",
+            timeout=5,
+        )
+        annotations = annotations_response.json()
+
+        assert len(annotations) == 1
+        assert (
+            annotations[0]["content"]
+            == "Users must not mention the magic phrase 'Abracadabra'"
+            and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        )
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+@pytest.mark.parametrize(
+    "do_stream, push_to_explorer",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_invalid_guardrail_in_header(gateway_url, do_stream, push_to_explorer):
+    """Test the message content guardrail."""
+    if not os.getenv("INVARIANT_API_KEY"):
+        pytest.fail("No INVARIANT_API_KEY set, failing")
+
+    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
+
+    policy = """
+# For input guardrailing specifically
+raise "Users must not mention the magic phrase 'Abracadabra'" if:
+    (msg: Message)
+    msg.role == "user"
+    "Abracadabra" in msg.content
+    illegal statement
+"""
+
+    client = OpenAI(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}",
+                "Invariant-Guardrails": policy.encode("unicode-escape"),
+            },
+        ),
+        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway/openai",
+    )
+
+    request = {
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "Tell me more about Abracadabra."}],
+    }
+
+    if not do_stream:
+        with pytest.raises(BadRequestError) as exc_info:
+            chat_response = client.chat.completions.create(
+                **request,
+                stream=False,
+            )
+
+        print(exc_info.value.message, flush=True)
+        assert "Failed to create policy from policy source." in str(
+            exc_info.value
+        ), "guardrails check fails because of an invalid guardrailing rule"
+        assert "illegal statement" in str(
+            exc_info.value
+        ), "error points to illegal statement in the rule definition"

tests/test-client.py

@@ -0,0 +1,33 @@
+"""
+Simple (non-streaming) test client for the Gateway (uses OpenAI integration).
+"""
+
+from openai import OpenAI
+from httpx import Client
+import os
+
+# unicode escape everything
+guardrails = """
+raise "Rule 1: Do not talk about Fight Club" if:
+    (msg: Message)
+    "fight club" in msg.content
+""".encode("unicode_escape")
+
+openai_client = OpenAI(
+    default_headers={
+        "Invariant-Authorization": "Bearer " + os.getenv("INVARIANT_API_KEY"),
+        "Invariant-Guardrails": guardrails,
+    },
+    base_url="http://localhost:9999/api/v1/gateway/non-streaming/openai",
+)
+
+response = openai_client.chat.completions.create(
+    model="gpt-4",
+    messages=[
+        {
+            "role": "user",
+            "content": "What can you tell me about fight club?",
+        }
+    ],
+)
+print("Response: ", response.choices[0].message.content)