From 278b3a8bf4838085a6b506c8b180c8acabf5d191 Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Thu, 3 Apr 2025 10:08:03 +0200 Subject: [PATCH] allow different guardrailing api key --- gateway/common/authorization.py | 13 +++++++++++++ gateway/common/request_context.py | 28 ++++++++++++++++++++++++++++ gateway/integrations/guardrails.py | 11 +++-------- gateway/routes/anthropic.py | 1 + gateway/routes/gemini.py | 1 + 5 files changed, 46 insertions(+), 8 deletions(-) diff --git a/gateway/common/authorization.py b/gateway/common/authorization.py index 8833e0c..12a68bc 100644 --- a/gateway/common/authorization.py +++ b/gateway/common/authorization.py @@ -4,9 +4,22 @@ from typing import Tuple, Optional from fastapi import HTTPException, Request INVARIANT_AUTHORIZATION_HEADER = "invariant-authorization" +INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER = "invariant-guardrails-authorization" API_KEYS_SEPARATOR = ";invariant-auth=" +def extract_guardrail_service_authorization_from_headers( + request: Request, +) -> Tuple[Optional[str], Optional[str]]: + """ + Extracts the optional Invariant-Guardrail-Service authorization header from the request. + + This header can be specifified to use a different API key for guardrailing compared to + Explorer interactions. + """ + return request.headers.get(INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER) + + def extract_authorization_from_headers( request: Request, dataset_name: Optional[str] = None, diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index 047c714..42ffc45 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -7,6 +7,9 @@ import fastapi from common.config_manager import GatewayConfig from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction +from gateway.common.authorization import ( + extract_guardrail_service_authorization_from_headers, +) @dataclass(frozen=True) @@ -15,7 +18,10 @@ class RequestContext: request_json: Dict[str, Any] dataset_name: Optional[str] = None + # authorization to use for invariant service like explorer invariant_authorization: Optional[str] = None + # authorization to use for invariant guardrailing specifically + guardrail_authorization: Optional[str] = None # the set of guardrails to enforce for this request guardrails: Optional[GuardrailRuleSet] = None config: Dict[str, Any] = None @@ -74,15 +80,37 @@ class RequestContext: logging_guardrails=[], ) + # if additionally provided, extract separate API key to use with guardrailing service + guardrail_service_authorization = None + if ( + guardrail_authorization + := extract_guardrail_service_authorization_from_headers(request) + ): + guardrail_service_authorization = guardrail_authorization + return cls( request_json=request_json, dataset_name=dataset_name, invariant_authorization=invariant_authorization, + guardrail_service_authorization=guardrail_service_authorization, guardrails=guardrails, config=context_config, _created_via_factory=True, + guardrail_authorization=guardrail_service_authorization, ) + def get_guardrailing_authorization(self) -> Optional[str]: + """ + Returns the authorization to use for the guardrailing service. + + This can be different from the invariant authorization, but falls back + "to be the same if not explicitly set via header. + + See also extract_guardrail_service_authorization_from_headers(...) + """ + + return self.guardrail_authorization or self.invariant_authorization + def __repr__(self) -> str: return ( f"RequestContext(" diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index 70965ad..4bfc9f5 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -82,10 +82,6 @@ async def _preload(guardrails: str, invariant_authorization: str) -> None: result.raise_for_status() -def get_guardrails_invariant_authorization(context: "RequestContext") -> str: - return context.invariant_authorization - - async def preload_guardrails(context: "RequestContext") -> None: """ Preloads the guardrails for faster checking later. @@ -101,8 +97,7 @@ async def preload_guardrails(context: "RequestContext") -> None: for blocking_guardrail in context.guardrails.blocking_guardrails: task = asyncio.create_task( _preload( - blocking_guardrail.content, - get_guardrails_invariant_authorization(context), + blocking_guardrail.content, context.get_guardrailing_authorization() ) ) asyncio.shield(task) @@ -110,7 +105,7 @@ async def preload_guardrails(context: "RequestContext") -> None: task = asyncio.create_task( _preload( logging_guadrail.content, - get_guardrails_invariant_authorization(context), + context.get_guardrailing_authorization(), ) ) asyncio.shield(task) @@ -367,7 +362,7 @@ async def check_guardrails( "policies": [g.content for g in guardrails], }, headers={ - "Authorization": get_guardrails_invariant_authorization(context), + "Authorization": context.get_guardrailing_authorization(), "Accept": "application/json", }, ) diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 9dfff01..13b5f20 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -105,6 +105,7 @@ async def anthropic_v1_messages_gateway( invariant_authorization=invariant_authorization, guardrails=header_guardrails or dataset_guardrails, config=config, + request=request, ) if request_json.get("stream"): return await handle_streaming_response(context, client, anthropic_request) diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index bb26411..00101f9 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -103,6 +103,7 @@ async def gemini_generate_content_gateway( invariant_authorization=invariant_authorization, guardrails=header_guardrails or dataset_guardrails, config=config, + request=request, ) if alt == "sse" or endpoint == "streamGenerateContent": return await stream_response(