diff --git a/gateway/common/authorization.py b/gateway/common/authorization.py index 74ba70b..8833e0c 100644 --- a/gateway/common/authorization.py +++ b/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 " {llm_provider_api_key_header} contains the LLM Provider API Key as {llm_provider_api_key_header}: "" - + 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}: ";invariant-auth=" """ + # 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()}" diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py index ac51c89..951e41b 100644 --- a/gateway/common/config_manager.py +++ b/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=[], + ) diff --git a/gateway/common/guardrails.py b/gateway/common/guardrails.py index cb7ef1e..e4164c1 100644 --- a/gateway/common/guardrails.py +++ b/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] diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index f2c8c3e..b252f31 100644 --- a/gateway/common/request_context.py +++ b/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})" ) diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index f65f5c6..dd15235 100644 --- a/gateway/integrations/explorer.py +++ b/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, ) diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index b7377c3..ec40651 100644 --- a/gateway/integrations/guardrails.py +++ b/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"], + } + ] + } diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 24cf097..4fd0744 100644 --- a/gateway/routes/anthropic.py +++ b/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, diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 07aa36a..1399c33 100644 --- a/gateway/routes/gemini.py +++ b/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, diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index ce6fefa..8466504 100644 --- a/gateway/routes/open_ai.py +++ b/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, diff --git a/tests/integration/guardrails/test_header_guardrails.py b/tests/integration/guardrails/test_header_guardrails.py new file mode 100644 index 0000000..e03847d --- /dev/null +++ b/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" diff --git a/tests/test-client.py b/tests/test-client.py new file mode 100644 index 0000000..388e453 --- /dev/null +++ b/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)