From 9d7dbb1bd5842c1687252982e9649a1117ba9fa9 Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Tue, 8 Apr 2025 10:49:39 +0200 Subject: [PATCH] push behavior header --- README.md | 2 +- gateway/common/request_context.py | 11 +++++++++++ gateway/routes/anthropic.py | 1 + gateway/routes/gemini.py | 1 + gateway/routes/open_ai.py | 1 + 5 files changed, 15 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index c3d58f4..ace781b 100644 --- a/README.md +++ b/README.md @@ -339,7 +339,7 @@ Then, this combined key can be passed as `Authorization` header only. - `Invariant-Guardrail-Service-Authorization`: Additional header value to specify a different API key to use for Guardrails evaluation specifically. `Invariant-Authorization` will then only be used to interact with the configured Explorer instance, not for guardrailing. - `Invariant-Guardrails`: Guardrailing rules to be checked for a specific request. This list of rules will replace all rules being evaluated, i.e. no additional rules will be pulled from [Explorer](https://explorer.invariantlabs.ai). -- `Invariant-NoPush`: Disables pushing to Explorer alltogether. Gateway will still pull and evaluate the guardrailing rules stored in Explorer, but it will no longer push the resulting trace. +- `Invariant-Push`: Configures push behavior of Gateway. Valid values are `push` (default, pushes to Explorer), `skip` (does not push anything) --- diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index 68cb1d7..706796a 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -24,7 +24,10 @@ class RequestContext: guardrail_authorization: Optional[str] = None # the set of guardrails to enforce for this request guardrails: Optional[GuardrailRuleSet] = None + # configuration parameters for this request config: Dict[str, Any] = None + # push behavior + push_behavior: str = 'push' _created_via_factory: bool = field( default=False, init=True, repr=False, compare=False @@ -55,6 +58,13 @@ class RequestContext: if key != "guardrails_from_file" } + # Read potential Invariant-NoPush header + push_behavior = 'push' + if push_behavior := request.headers.get("Invariant-Push"): + if not push_behavior.lower() in ["push", "skip"]: + raise fastapi.HTTPException(status_code=400, detail="Invalid value for Invariant-Push header. Valid values are 'push' and 'skip'.") + push_behavior = push_behavior.lower() + # If no guardrails are configured and the config specifies # guardrails_from_file, use those instead. if ( @@ -95,6 +105,7 @@ class RequestContext: guardrail_authorization=guardrail_service_authorization, guardrails=guardrails, config=context_config, + push_behavior=push_behavior, _created_via_factory=True, ) diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 13b5f20..8f7d2f5 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -196,6 +196,7 @@ async def push_to_explorer( invariant_authorization=context.invariant_authorization, metadata=[create_metadata(context, merged_response)], annotations=[annotations] if annotations else None, + push_behavior=context.push_behavior, ) diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 00101f9..17e9858 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -437,6 +437,7 @@ async def push_to_explorer( invariant_authorization=context.invariant_authorization, metadata=[create_metadata(context, response_json)], annotations=[annotations] if annotations else None, + push_behavior=context.push_behavior, ) diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 5695c03..9124398 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -520,6 +520,7 @@ async def push_to_explorer( messages=[messages], annotations=[annotations], metadata=[create_metadata(context, merged_response)], + push_behavior=context.push_behavior, )