From e55df7348786074af24a9edc70c84c4932009836 Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 26 Mar 2025 15:16:05 +0100 Subject: [PATCH] Add guardrails for gemini integration. --- docker-compose.local.yml | 1 + gateway/common/config_manager.py | 34 ++++++---- gateway/integrations/explorer.py | 25 +++++++- gateway/integrations/guardails.py | 23 +++---- gateway/routes/gemini.py | 100 +++++++++++++++++++++++++++--- gateway/validate_config.py | 2 +- run.sh | 10 ++- 7 files changed, 157 insertions(+), 38 deletions(-) diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 069a926..60ea5ad 100644 --- a/docker-compose.local.yml +++ b/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 diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py index c0438d4..57a1d89 100644 --- a/gateway/common/config_manager.py +++ b/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)})" diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index ed4098d..fd0760b 100644 --- a/gateway/integrations/explorer.py +++ b/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, diff --git a/gateway/integrations/guardails.py b/gateway/integrations/guardails.py index 6944e1e..4c571fc 100644 --- a/gateway/integrations/guardails.py +++ b/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. diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 905f14f..4042b4e 100644 --- a/gateway/routes/gemini.py +++ b/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), ) diff --git a/gateway/validate_config.py b/gateway/validate_config.py index 99a24bb..6ad9392 100644 --- a/gateway/validate_config.py +++ b/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) diff --git a/run.sh b/run.sh index 6cc9931..6447ede 100755 --- a/run.sh +++ b/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() {