minor refactor for getting invariant api keys for guardrailing

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."""
 

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",
                 },
             )

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
 

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
 

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