diff --git a/client.py b/client.py new file mode 100644 index 0000000..e4771b6 --- /dev/null +++ b/client.py @@ -0,0 +1,29 @@ +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": open("guardrails.py").read().encode("unicode_escape"), + }, + base_url="http://localhost:8000/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) diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py index 3e62b55..ad85149 100644 --- a/gateway/common/config_manager.py +++ b/gateway/common/config_manager.py @@ -3,15 +3,31 @@ import asyncio import os import threading +from typing import Optional +import fastapi from httpx import HTTPStatusError +def extract_policy_from_headers(request: fastapi.Request) -> Optional[str]: + """ + Extracts the guardrailing policy from the request headers if present. + + Returns 'None' if no such header is present. + """ + 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.""" - def __init__(self): - self.guardrails = self._load_guardrails_from_file() + def __init__(self, guardrails: Optional[str] = None): + self.guardrails = guardrails or self._load_guardrails_from_file() def _load_guardrails_from_file(self) -> str: """ @@ -50,6 +66,12 @@ class GatewayConfig: def __repr__(self) -> str: return f"GatewayConfig(guardrails={repr(self.guardrails)})" + def with_guardrails(self, guardrails: str) -> "GatewayConfig": + """ + Returns a new GatewayConfig instance with the specified guardrails. + """ + return GatewayConfig(guardrails) + class GatewayConfigManager: """Manager for Gateway Configuration.""" @@ -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,9 @@ class GatewayConfigManager: if local_config is None: local_config = GatewayConfig() cls._config_instance = local_config + + # if provided in header, use custom guardrailing policy + if guardrail_file_contents := extract_policy_from_headers(request): + local_config = local_config.with_guardrails(guardrail_file_contents) + return local_config diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index b0f3601..e3c236f 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -355,4 +355,13 @@ async def check_guardrails( return result.json() 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"], + } + ] + }