CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-08 20:07:54 +02:00
Code Issues Packages Projects Releases Wiki Activity

Small changes related to constants and sorting order of imports.

Browse Source
This commit is contained in:
Hemang
2025-06-05 09:55:02 +02:00
committed by Hemang Sarkar
parent 24d47c4585
commit cd6d6a50b0
12 changed files with 74 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/common/config_manager.py
-1
View File
@@ -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.
gateway/common/constants.py
+3
View File
@@ -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"
gateway/common/guardrails.py
+1 -3
View File
@@ -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."""
gateway/common/request_context.py
+4 -4
View File
@@ -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.<key> 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]:
gateway/integrations/explorer.py
+2 -3
View File
@@ -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],
gateway/integrations/guardrails.py
+7 -4
View File
@@ -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 {
gateway/mcp/log.py
-1
View File
@@ -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)
gateway/mcp/sse.py
+2 -2
View File
@@ -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},
)
gateway/mcp/streamable.py
+8 -6
View File
@@ -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,
)
gateway/routes/anthropic.py
+27 -10
View File
@@ -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
)
gateway/routes/gemini.py
+12 -6
View File
@@ -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),
)
)
gateway/routes/open_ai.py
+8 -5
View File
@@ -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,
),
)