From cd6d6a50b05838e21cad1f38a587ef89bef3f151 Mon Sep 17 00:00:00 2001 From: Hemang Date: Thu, 5 Jun 2025 09:55:02 +0200 Subject: [PATCH] Small changes related to constants and sorting order of imports. --- gateway/common/config_manager.py | 1 - gateway/common/constants.py | 3 +++ gateway/common/guardrails.py | 4 +--- gateway/common/request_context.py | 8 +++---- gateway/integrations/explorer.py | 5 ++-- gateway/integrations/guardrails.py | 11 +++++---- gateway/mcp/log.py | 1 - gateway/mcp/sse.py | 4 ++-- gateway/mcp/streamable.py | 14 ++++++----- gateway/routes/anthropic.py | 37 ++++++++++++++++++++++-------- gateway/routes/gemini.py | 18 ++++++++++----- gateway/routes/open_ai.py | 13 +++++++---- 12 files changed, 74 insertions(+), 45 deletions(-) diff --git a/gateway/common/config_manager.py b/gateway/common/config_manager.py index 98d83e5..0c6095d 100644 --- a/gateway/common/config_manager.py +++ b/gateway/common/config_manager.py @@ -10,7 +10,6 @@ from httpx import HTTPStatusError from gateway.common.guardrails import Guardrail, GuardrailAction, GuardrailRuleSet - def extract_policy_from_headers(request: Optional[fastapi.Request]) -> Optional[str]: """ Extracts the guardrailing policy from the request headers if present. diff --git a/gateway/common/constants.py b/gateway/common/constants.py index a84eaea..20ab1d8 100644 --- a/gateway/common/constants.py +++ b/gateway/common/constants.py @@ -16,3 +16,6 @@ IGNORED_HEADERS = [ CLIENT_TIMEOUT = 60.0 +CONTENT_TYPE_HEADER = "content-type" +CONTENT_TYPE_JSON = "application/json" +CONTENT_TYPE_EVENT_STREAM = "text/event-stream" diff --git a/gateway/common/guardrails.py b/gateway/common/guardrails.py index e4164c1..ec7a94e 100644 --- a/gateway/common/guardrails.py +++ b/gateway/common/guardrails.py @@ -1,11 +1,9 @@ """Common guardrails data class.""" +from dataclasses import dataclass from enum import Enum from typing import List -from dataclasses import dataclass - - class GuardrailAction(str, Enum): """Enum representing the action to be taken for guardrail rules.""" diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index 11b477d..fbce8a3 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -5,11 +5,11 @@ from typing import Any, Dict, Optional import fastapi -from gateway.common.config_manager import GatewayConfig -from gateway.common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction from gateway.common.authorization import ( extract_guardrail_service_authorization_from_headers, ) +from gateway.common.config_manager import GatewayConfig +from gateway.common.guardrails import GuardrailRuleSet, Guardrail, GuardrailAction @dataclass(frozen=True) @@ -25,7 +25,7 @@ class RequestContext: # the set of guardrails to enforce for this request guardrails: Optional[GuardrailRuleSet] = None config: Dict[str, Any] = None - + # extra parameters available as input. during guardrail evaluation guardrails_parameters: Optional[Dict[str, Any]] = None @@ -100,7 +100,7 @@ class RequestContext: guardrails=guardrails, config=context_config, _created_via_factory=True, - guardrails_parameters=guardrails_parameters + guardrails_parameters=guardrails_parameters, ) def get_guardrailing_authorization(self) -> Optional[str]: diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index e547831..621f1c2 100644 --- a/gateway/integrations/explorer.py +++ b/gateway/integrations/explorer.py @@ -2,8 +2,9 @@ import os import json - from typing import Any, Dict, List + +import httpx from fastapi import HTTPException from gateway.common.constants import DEFAULT_API_URL @@ -12,8 +13,6 @@ from invariant_sdk.async_client import AsyncClient from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesResponse from invariant_sdk.types.annotations import AnnotationCreate -import httpx - def create_annotations_from_guardrails_errors( guardrails_errors: List[dict], diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index aa73dad..2b3f912 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -376,11 +376,14 @@ async def check_guardrails( if result.status_code == 401: raise HTTPException( status_code=401, - detail="The provided Invariant API key is not valid for guardrail checking. Please ensure you are using the correct API key or pass an alternative API key for guardrail checking specifically via the '{}' header.".format( - INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER + detail=( + "The provided Invariant API key is not valid for guardrail checking. " + "Please ensure you are using the correct API key or pass an " + "alternative API key for guardrail checking specifically via the " + f"'{INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER}' header." ), ) - raise Exception( + raise Exception( # pylint: disable=broad-exception-raised f"Guardrails check failed: {result.status_code} - {result.text}" ) guardrails_result = result.json() @@ -412,7 +415,7 @@ async def check_guardrails( return aggregated_errors except HTTPException as e: raise e - except Exception as e: + except Exception as e: # pylint: disable=broad-except print(f"Failed to verify guardrails: {e}") # make sure runtime errors are also visible in e.g. Explorer return { diff --git a/gateway/mcp/log.py b/gateway/mcp/log.py index 7287513..388355b 100644 --- a/gateway/mcp/log.py +++ b/gateway/mcp/log.py @@ -2,7 +2,6 @@ import os import sys - from builtins import print as builtins_print os.makedirs(os.path.join(os.path.expanduser("~"), ".invariant"), exist_ok=True) diff --git a/gateway/mcp/sse.py b/gateway/mcp/sse.py index 998d5de..6c6b971 100644 --- a/gateway/mcp/sse.py +++ b/gateway/mcp/sse.py @@ -10,7 +10,7 @@ from httpx_sse import aconnect_sse, ServerSentEvent from fastapi import APIRouter, HTTPException, Request, Response from fastapi.responses import StreamingResponse -from gateway.common.constants import CLIENT_TIMEOUT +from gateway.common.constants import CLIENT_TIMEOUT, CONTENT_TYPE_EVENT_STREAM from gateway.mcp.constants import MCP_CUSTOM_HEADER_PREFIX, UTF_8 from gateway.mcp.mcp_sessions_manager import ( McpSessionsManager, @@ -293,7 +293,7 @@ class SseTransport(McpTransportBase): return StreamingResponse( event_generator(), - media_type="text/event-stream", + media_type=CONTENT_TYPE_EVENT_STREAM, headers={"X-Proxied-By": "mcp-gateway", **response_headers}, ) diff --git a/gateway/mcp/streamable.py b/gateway/mcp/streamable.py index bf5f694..930abbd 100644 --- a/gateway/mcp/streamable.py +++ b/gateway/mcp/streamable.py @@ -8,7 +8,12 @@ from httpx_sse import aconnect_sse from fastapi import APIRouter, HTTPException, Request, Response from fastapi.responses import StreamingResponse -from gateway.common.constants import CLIENT_TIMEOUT +from gateway.common.constants import ( + CLIENT_TIMEOUT, + CONTENT_TYPE_HEADER, + CONTENT_TYPE_JSON, + CONTENT_TYPE_EVENT_STREAM, +) from gateway.mcp.constants import ( INVARIANT_SESSION_ID_PREFIX, MCP_CUSTOM_HEADER_PREFIX, @@ -23,9 +28,6 @@ from gateway.mcp.mcp_transport_base import McpTransportBase gateway = APIRouter() mcp_sessions_manager = McpSessionsManager() -CONTENT_TYPE_JSON = "application/json" -CONTENT_TYPE_SSE = "text/event-stream" -CONTENT_TYPE_HEADER = "content-type" MCP_SESSION_ID_HEADER = "mcp-session-id" MCP_SERVER_POST_AND_DELETE_HEADERS = { "connection", @@ -187,7 +189,7 @@ class StreamableTransport(McpTransportBase): return StreamingResponse( event_generator(), - media_type=CONTENT_TYPE_SSE, + media_type=CONTENT_TYPE_EVENT_STREAM, headers={"X-Proxied-By": "mcp-gateway", **response_headers}, ) @@ -382,7 +384,7 @@ class StreamableTransport(McpTransportBase): return StreamingResponse( event_generator(), - media_type=CONTENT_TYPE_SSE, + media_type=CONTENT_TYPE_EVENT_STREAM, headers=response_headers, ) diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 36c4a63..e455b11 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -14,7 +14,12 @@ from gateway.common.config_manager import ( GatewayConfigManager, extract_guardrails_from_header, ) -from gateway.common.constants import CLIENT_TIMEOUT, IGNORED_HEADERS +from gateway.common.constants import ( + CLIENT_TIMEOUT, + CONTENT_TYPE_JSON, + CONTENT_TYPE_EVENT_STREAM, + IGNORED_HEADERS, +) from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet from gateway.common.request_context import RequestContext from gateway.converters.anthropic_to_invariant import ( @@ -218,7 +223,10 @@ class InstrumentedAnthropicResponse(InstrumentedResponse): self.guardrails_execution_result = {} async def on_start(self): - """Check guardrails in a pipelined fashion, before processing the first chunk (for input guardrailing).""" + """ + Check guardrails in a pipelined fashion, before processing the first + chunk (for input guardrailing). + """ if self.context.guardrails: self.guardrails_execution_result = await get_guardrails_check_result( self.context, action=GuardrailAction.BLOCK, response_json={} @@ -249,8 +257,8 @@ class InstrumentedAnthropicResponse(InstrumentedResponse): Response( content=error_chunk, status_code=400, - media_type="application/json", - headers={"content-type": "application/json"}, + media_type=CONTENT_TYPE_JSON, + headers={"content-type": CONTENT_TYPE_JSON}, ) ) @@ -263,7 +271,10 @@ class InstrumentedAnthropicResponse(InstrumentedResponse): except json.JSONDecodeError as e: raise HTTPException( status_code=self.response.status_code, - detail=f"Invalid JSON response received from Anthropic: {self.response.text}, got error{e}", + detail=( + "Invalid JSON response received from Anthropic: " + f"{self.response.text}, got error: {e}" + ), ) from e if self.response.status_code != 200: raise HTTPException( @@ -289,12 +300,15 @@ class InstrumentedAnthropicResponse(InstrumentedResponse): return Response( content=content, status_code=status_code, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, headers=dict(updated_headers), ) async def on_end(self): - """Checks guardrails after the response is received, and asynchronously pushes to Explorer.""" + """ + Checks guardrails after the response is received, and asynchronously + pushes to Explorer. + """ # ensure the response data is available assert self.response is not None, "response is None" assert self.response_json is not None, "response_json is None" @@ -383,7 +397,10 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse): self.sse_buffer = "" # Buffer for incomplete events async def on_start(self): - """Check guardrails in a pipelined fashion, before processing the first chunk (for input guardrailing).""" + """ + Check guardrails in a pipelined fashion, before processing the + first chunk (for input guardrailing). + """ if self.context.guardrails: self.guardrails_execution_result = await get_guardrails_check_result( self.context, @@ -503,7 +520,7 @@ class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse): f"JSON parsing error in event: {e}. Event data: {event_data[:100]}...", flush=True, ) - except Exception as e: + except Exception as e: # pylint: disable=broad-except print(f"Error processing event: {e}", flush=True) # on last stream chunk, run output guardrails @@ -582,7 +599,7 @@ async def handle_streaming_response( ) return StreamingResponse( - response.instrumented_event_generator(), media_type="text/event-stream" + response.instrumented_event_generator(), media_type=CONTENT_TYPE_EVENT_STREAM ) diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 6e7fe26..14a3a63 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -16,6 +16,8 @@ from gateway.common.config_manager import ( ) from gateway.common.constants import ( CLIENT_TIMEOUT, + CONTENT_TYPE_JSON, + CONTENT_TYPE_EVENT_STREAM, IGNORED_HEADERS, ) from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet @@ -82,7 +84,11 @@ async def gemini_generate_content_gateway( request_json = json.loads(request_body_bytes) client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) - gemini_api_url = f"https://generativelanguage.googleapis.com/{api_version}/models/{model}:{endpoint}" + gemini_api_url = ( + f"https://generativelanguage.googleapis.com/" + f"{api_version}/models/" + f"{model}:{endpoint}" + ) if alt == "sse": gemini_api_url += "?alt=sse" gemini_request = client.build_request( @@ -303,7 +309,7 @@ async def stream_response( return StreamingResponse( event_generator(), - media_type="text/event-stream", + media_type=CONTENT_TYPE_EVENT_STREAM, ) @@ -511,9 +517,9 @@ class InstrumentedGeminiResponse(InstrumentedResponse): Response( content=error_chunk, status_code=400, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, headers={ - "Content-Type": "application/json", + "Content-Type": CONTENT_TYPE_JSON, }, ) ) @@ -541,7 +547,7 @@ class InstrumentedGeminiResponse(InstrumentedResponse): return Response( content=response_string, status_code=response_code, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, headers=dict(self.response.headers), ) @@ -584,7 +590,7 @@ class InstrumentedGeminiResponse(InstrumentedResponse): Response( content=response_string, status_code=response_code, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, headers=dict(self.response.headers), ) ) diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 4d25871..2d8501e 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -16,6 +16,8 @@ from gateway.common.config_manager import ( ) from gateway.common.constants import ( CLIENT_TIMEOUT, + CONTENT_TYPE_JSON, + CONTENT_TYPE_EVENT_STREAM, IGNORED_HEADERS, ) from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet @@ -273,7 +275,8 @@ class InstrumentedOpenAIStreamResponse(InstrumentedStreamingResponse): } ) - # yield an extra error chunk (without preventing the original chunk to go through after) + # yield an extra error chunk (without preventing the original + # chunk to go through after) return ExtraItem(f"data: {error_chunk}\n\n".encode()) # push will happen in on_end @@ -324,7 +327,7 @@ async def handle_stream_response( ) return StreamingResponse( - response.instrumented_event_generator(), media_type="text/event-stream" + response.instrumented_event_generator(), media_type=CONTENT_TYPE_EVENT_STREAM ) @@ -606,7 +609,7 @@ class InstrumentedOpenAIResponse(InstrumentedResponse): } ), status_code=400, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, ), end_of_stream=True, ) @@ -634,7 +637,7 @@ class InstrumentedOpenAIResponse(InstrumentedResponse): return Response( content=response_string, status_code=response_code, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, headers=dict(self.response.headers), ) @@ -686,7 +689,7 @@ class InstrumentedOpenAIResponse(InstrumentedResponse): Response( content=response_string, status_code=response_code, - media_type="application/json", + media_type=CONTENT_TYPE_JSON, ), )