diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index b252f31..047c714 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -3,6 +3,8 @@ from dataclasses import dataclass, field from typing import Any, Dict, Optional +import fastapi + from common.config_manager import GatewayConfig from common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction @@ -36,6 +38,7 @@ class RequestContext: invariant_authorization: Optional[str] = None, guardrails: Optional[GuardrailRuleSet] = None, config: Optional[GatewayConfig] = None, + request: fastapi.Request = None, ) -> "RequestContext": """Creates a new RequestContext instance, applying default guardrails if needed.""" diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index ec40651..70965ad 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -82,6 +82,10 @@ 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. @@ -96,12 +100,18 @@ async def preload_guardrails(context: "RequestContext") -> None: # Move these calls to a batch preload/validate API. for blocking_guardrail in context.guardrails.blocking_guardrails: task = asyncio.create_task( - _preload(blocking_guardrail.content, context.invariant_authorization) + _preload( + blocking_guardrail.content, + get_guardrails_invariant_authorization(context), + ) ) asyncio.shield(task) for logging_guadrail in context.guardrails.logging_guardrails: task = asyncio.create_task( - _preload(logging_guadrail.content, context.invariant_authorization) + _preload( + logging_guadrail.content, + get_guardrails_invariant_authorization(context), + ) ) asyncio.shield(task) except Exception as e: @@ -332,7 +342,7 @@ class InstrumentedResponse(InstrumentedStreamingResponse): async def check_guardrails( messages: List[Dict[str, Any]], guardrails: List[Guardrail], - invariant_authorization: str, + context: RequestContext, ) -> Dict[str, Any]: """ Checks guardrails on the list of messages. @@ -357,7 +367,7 @@ async def check_guardrails( "policies": [g.content for g in guardrails], }, headers={ - "Authorization": invariant_authorization, + "Authorization": get_guardrails_invariant_authorization(context), "Accept": "application/json", }, ) diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 4fd0744..9dfff01 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -157,7 +157,7 @@ async def get_guardrails_check_result( guardrails_execution_result = await check_guardrails( messages=converted_messages, guardrails=guardrails, - invariant_authorization=context.invariant_authorization, + context=context, ) return guardrails_execution_result diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 1399c33..bb26411 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -64,11 +64,16 @@ async def gemini_generate_content_gateway( status_code=400, ) headers = { - k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS + [GEMINI_AUTHORIZATION_FALLBACK_HEADER] + k: v + for k, v in request.headers.items() + if k.lower() not in IGNORED_HEADERS + [GEMINI_AUTHORIZATION_FALLBACK_HEADER] } headers["accept-encoding"] = "identity" invariant_authorization, gemini_api_key = extract_authorization_from_headers( - request, dataset_name, GEMINI_AUTHORIZATION_HEADER, [GEMINI_AUTHORIZATION_FALLBACK_HEADER] + request, + dataset_name, + GEMINI_AUTHORIZATION_HEADER, + [GEMINI_AUTHORIZATION_FALLBACK_HEADER], ) headers[GEMINI_AUTHORIZATION_HEADER] = gemini_api_key @@ -394,7 +399,7 @@ async def get_guardrails_check_result( guardrails_execution_result = await check_guardrails( messages=converted_requests + converted_responses, guardrails=guardrails, - invariant_authorization=context.invariant_authorization, + context=context, ) return guardrails_execution_result diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 8466504..5695c03 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -149,6 +149,7 @@ async def openai_chat_completions_gateway( invariant_authorization=invariant_authorization, guardrails=header_guardrails or dataset_guardrails, config=config, + request=request, ) if request_json.get("stream", False): return await handle_stream_response( @@ -546,7 +547,7 @@ async def get_guardrails_check_result( guardrails_execution_result = await check_guardrails( messages=messages, guardrails=guardrails, - invariant_authorization=context.invariant_authorization, + context=context, ) return guardrails_execution_result