Add guardrails for gemini integration.

docker-compose.local.yml

@@ -10,6 +10,7 @@ services:
     environment:
       - DEV_MODE=true
       - GUARDRAILS_FILE_PATH=${GUARDRAILS_FILE_PATH:+/srv/resources/guardrails.py}
+      - ${INVARIANT_API_KEY:+INVARIANT_API_KEY=${INVARIANT_API_KEY}}
     volumes:
       - type: bind
         source: ./gateway

gateway/common/config_manager.py

@@ -1,41 +1,51 @@
 """Common Configurations for the Gateway Server."""
 
+import asyncio
 import os
 import threading
 
-from invariant.analyzer import Policy
+from integrations.guardails import _preload
+
+from httpx import HTTPStatusError
 
 
 class GatewayConfig:
     """Common configurations for the Gateway Server."""
 
     def __init__(self):
-        self.guardrails = self._load_guardrails()
+        self.guardrails = self._load_guardrails_from_file()
 
-    def _load_guardrails(self) -> str:
+    def _load_guardrails_from_file(self) -> str:
         """
-        Loads and validates guardrails from the file specified in GUARDRAILS_FILE_PATH.
-        Returns the guardrails file content as a string if valid; otherwise, raises an error.
+        Loads the guardrails from the file specified in GUARDRAILS_FILE_PATH.
+        Returns the guardrails file content as a string.
         """
         guardrails_file = os.getenv("GUARDRAILS_FILE_PATH", "")
-
         if not guardrails_file:
             print("[warning: GUARDRAILS_FILE_PATH is not set. Using empty guardrails]")
             return ""
 
+        invariant_api_key = os.getenv("INVARIANT_API_KEY", "")
+        if not invariant_api_key:
+            raise ValueError(
+                "Error: INVARIANT_API_KEY is not set."
+                "It is required to validate guardrails file content."
+            )
+
         try:
             with open(guardrails_file, "r", encoding="utf-8") as f:
                 guardrails_file_content = f.read()
-            _ = Policy.from_string(guardrails_file_content)
-            return guardrails_file_content
+                asyncio.run(
+                    _preload(guardrails_file_content, "Bearer " + invariant_api_key)
+                )
+                return guardrails_file_content
 
         except (FileNotFoundError, PermissionError, OSError) as e:
             raise ValueError(
-                f"Error: Unable to read guardrails file ({guardrails_file}): {e}"
+                f"Unable to read guardrails file ({guardrails_file}): {e}"
             ) from e
-
-        except Exception as e:
-            raise ValueError(f"Invalid policy content in {guardrails_file}: {e}") from e
+        except HTTPStatusError as e:
+            raise ValueError(f"Cannot load guardrails, {e}, {e.response.text}") from e
 
     def __repr__(self) -> str:
         return f"GatewayConfig(guardrails={repr(self.guardrails)})"

gateway/integrations/explorer.py

@@ -15,9 +15,32 @@ def create_annotations_from_guardrails_errors(
 ) -> List[AnnotationCreate]:
     """Create Explorer annotations from the guardrails errors."""
     annotations = []
+
+    def _remove_prefixes(ranges: list[str]) -> list[str]:
+        """
+        Remove prefixes from the list of ranges.
+
+        If the ranges are ['messages.2', 'messages.2.content:25-30', 'messages.2.content']
+        then this returns ['messages.2.content:25-30'].
+        """
+        ranges = sorted(ranges, key=len)
+        result = []
+
+        for i, s in enumerate(ranges):
+            is_prefix = False
+            for t in ranges[i + 1 :]:
+                if t.startswith(s) and t != s:
+                    is_prefix = True
+                    break
+            if not is_prefix:
+                result.append(s)
+
+        return result
+
     for error in guardrails_errors:
         content = error.get("args")[0]
-        for r in error.get("ranges", []):
+        filtered_ranges = _remove_prefixes(list(error.get("ranges", [])))
+        for r in filtered_ranges:
             annotations.append(
                 AnnotationCreate(
                     content=content,

gateway/integrations/guardails.py

@@ -7,7 +7,6 @@ from typing import Any, Dict, List
 from functools import wraps
 
 import httpx
-from common.request_context_data import RequestContextData
 
 DEFAULT_API_URL = "https://guardrail.invariantnet.com"
 
@@ -70,20 +69,18 @@ async def _preload(guardrails: str, invariant_authorization: str) -> None:
     """
     async with httpx.AsyncClient() as client:
         url = os.getenv("GUADRAILS_API_URL", DEFAULT_API_URL).rstrip("/")
-        try:
-            await client.post(
-                f"{url}/api/v1/policy/load",
-                json={"policy": guardrails},
-                headers={
-                    "Authorization": invariant_authorization,
-                    "Accept": "application/json",
-                },
-            )
-        except Exception as e:
-            print(f"Failed to load guardrails: {e}")
+        result = await client.post(
+            f"{url}/api/v1/policy/load",
+            json={"policy": guardrails},
+            headers={
+                "Authorization": invariant_authorization,
+                "Accept": "application/json",
+            },
+        )
+        result.raise_for_status()
 
 
-async def preload_guardrails(context: RequestContextData) -> None:
+async def preload_guardrails(context: "RequestContextData") -> None:
     """
     Preloads the guardrails for faster checking later.
 

gateway/routes/gemini.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import json
-from typing import Any
+from typing import Any, Optional
 
 import httpx
 from common.config_manager import GatewayConfig, GatewayConfigManager
@@ -15,8 +15,8 @@ from common.constants import (
 from common.authorization import extract_authorization_from_headers
 from common.request_context_data import RequestContextData
 from converters.gemini_to_invariant import convert_request, convert_response
-from integrations.explorer import push_trace
-from integrations.guardails import preload_guardrails
+from integrations.explorer import create_annotations_from_guardrails_errors, push_trace
+from integrations.guardails import check_guardrails, preload_guardrails
 
 gateway = APIRouter()
 
@@ -117,12 +117,45 @@ async def stream_response(
             if not chunk_text:
                 continue
 
-            # Yield chunk immediately to the client
-            yield chunk
-
             # Parse and update merged_response incrementally
             process_chunk_text(merged_response, chunk_text)
 
+            if (
+                merged_response.get("candidates", [])
+                and merged_response.get("candidates")[0].get("finishReason", "")
+                and context.config
+                and context.config.guardrails
+            ):
+                # Block on the guardrails check
+                guardrails_execution_result = await get_guardrails_check_result(
+                    context, merged_response
+                )
+                if guardrails_execution_result.get("errors", []):
+                    error_chunk = json.dumps(
+                        {
+                            "error": {
+                                "code": 400,
+                                "message": "[Invariant] The response did not pass the guardrails",
+                                "details": guardrails_execution_result,
+                                "status": "INVARIANT_GUARDRAILS_VIOLATION",
+                            },
+                        }
+                    )
+                    # Push annotated trace to the explorer - don't block on its response
+                    if context.dataset_name:
+                        asyncio.create_task(
+                            push_to_explorer(
+                                context,
+                                merged_response,
+                                guardrails_execution_result,
+                            )
+                        )
+                    yield f"data: {error_chunk}\n\n".encode()
+                    return
+
+            # Yield chunk immediately to the client
+            yield chunk
+
         if context.dataset_name:
             # Push to Explorer - don't block on the response
             asyncio.create_task(
@@ -209,18 +242,42 @@ def create_metadata(
     return metadata
 
 
+async def get_guardrails_check_result(
+    context: RequestContextData, response_json: dict[str, Any]
+) -> dict[str, Any]:
+    """Get the guardrails check result"""
+    converted_requests = convert_request(context.request_json)
+    converted_responses = convert_response(response_json)
+
+    # Block on the guardrails check
+    guardrails_execution_result = await check_guardrails(
+        messages=converted_requests + converted_responses,
+        guardrails=context.config.guardrails,
+        invariant_authorization=context.invariant_authorization,
+    )
+    return guardrails_execution_result
+
+
 async def push_to_explorer(
     context: RequestContextData,
     response_json: dict[str, Any],
+    guardrails_execution_result: Optional[dict] = None,
 ) -> None:
     """Pushes the full trace to the Invariant Explorer"""
+    guardrails_execution_result = guardrails_execution_result or {}
+    annotations = create_annotations_from_guardrails_errors(
+        guardrails_execution_result.get("errors", [])
+    )
+
     converted_requests = convert_request(context.request_json)
     converted_responses = convert_response(response_json)
+
     _ = await push_trace(
         dataset_name=context.dataset_name,
         messages=[converted_requests + converted_responses],
         invariant_authorization=context.invariant_authorization,
         metadata=[create_metadata(context, response_json)],
+        annotations=[annotations] if annotations else None,
     )
 
 
@@ -241,13 +298,36 @@ async def handle_non_streaming_response(
             status_code=response.status_code,
             detail=response_json.get("error", "Unknown error from Gemini API"),
         )
+    guardrails_execution_result = {}
+    response_string = json.dumps(response_json)
+    response_code = response.status_code
+
+    if context.config and context.config.guardrails:
+        # Block on the guardrails check
+        guardrails_execution_result = await get_guardrails_check_result(
+            context, response_json
+        )
+        if guardrails_execution_result.get("errors", []):
+            response_string = json.dumps(
+                {
+                    "error": {
+                        "code": 400,
+                        "message": "[Invariant] The response did not pass the guardrails",
+                        "details": guardrails_execution_result,
+                        "status": "INVARIANT_GUARDRAILS_VIOLATION",
+                    },
+                }
+            )
+            response_code = 400
     if context.dataset_name:
-        # Push to Explorer - don't block on the response
-        asyncio.create_task(push_to_explorer(context, response_json))
+        # Push to Explorer - don't block on its response
+        asyncio.create_task(
+            push_to_explorer(context, response_json, guardrails_execution_result)
+        )
 
     return Response(
-        content=json.dumps(response_json),
-        status_code=response.status_code,
+        content=response_string,
+        status_code=response_code,
         media_type="application/json",
         headers=dict(response.headers),
     )

gateway/validate_config.py

@@ -9,5 +9,5 @@ try:
     print("[gateway config validated successfully]")
     sys.exit(0)
 except Exception as e:
-    print(f"Error loading GatewayConfig error: {e}")
+    print(f"Error loading GatewayConfig: {e}")
     sys.exit(1)

run.sh

@@ -23,14 +23,21 @@ up() {
   if [[ -n "$GUARDRAILS_FILE_PATH" ]]; then
     if [[ -f "$GUARDRAILS_FILE_PATH" ]]; then
       GUARDRAILS_FILE_PATH=$(realpath "$GUARDRAILS_FILE_PATH")
+      export GUARDRAILS_FILE_PATH="$GUARDRAILS_FILE_PATH"
     else
       echo "Error: Specified guardrails file does not exist: $GUARDRAILS_FILE_PATH"
       exit 1
     fi
+
+    # If GUARDRAILS_FILE_PATH is set, then INVARIANT_API_KEY **must** be set
+    if [[ -z "$INVARIANT_API_KEY" ]]; then
+      echo "Error: A guardrails file is specified, but INVARIANT_API_KEY env var is not set. This is required to validate guardrails."
+      exit 1
+    fi
   fi
 
   # Start Docker Compose with the correct environment variable
-  GUARDRAILS_FILE_PATH="$GUARDRAILS_FILE_PATH" docker compose -f docker-compose.local.yml up -d
+  docker compose -f docker-compose.local.yml up -d
 
   # Get the status of the container
   sleep 2
@@ -46,6 +53,7 @@ up() {
   if [ -n "$GUARDRAILS_FILE_PATH" ]; then
     echo "Using Guardrails File: $GUARDRAILS_FILE_PATH"
   fi
+  unset GUARDRAILS_FILE_PATH
 }
 
 build() {