Remove unnecessary methods.

gateway/common/constants.py

@@ -4,7 +4,6 @@ DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
 
 IGNORED_HEADERS = [
     "accept-encoding",
-    "content-length",
     "host",
     "invariant-authorization",
     "x-forwarded-for",

gateway/integrations/guardrails.py

@@ -9,7 +9,7 @@ from typing import Any
 import httpx
 from fastapi import HTTPException
 
-from gateway.common.constants import DEFAULT_API_URL
+from gateway.common.constants import CONTENT_TYPE_JSON, DEFAULT_API_URL
 from gateway.common.request_context import RequestContext
 from gateway.common.authorization import (
     INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER,
@@ -79,7 +79,7 @@ async def _preload(guardrails: str, invariant_authorization: str) -> None:
             json={"policy": guardrails},
             headers={
                 "Authorization": invariant_authorization,
-                "Accept": "application/json",
+                "Accept": CONTENT_TYPE_JSON,
             },
         )
         result.raise_for_status()
@@ -146,7 +146,7 @@ async def check_guardrails(
                 },
                 headers={
                     "Authorization": context.get_guardrailing_authorization(),
-                    "Accept": "application/json",
+                    "Accept": CONTENT_TYPE_JSON,
                 },
                 timeout=5,
             )

gateway/routes/anthropic.py

@@ -25,7 +25,7 @@ from gateway.converters.anthropic_to_invariant import (
     convert_anthropic_to_invariant_message_format,
 )
 from gateway.integrations.explorer import fetch_guardrails_from_explorer
-from gateway.routes.base_provider import BaseProvider, Replacement
+from gateway.routes.base_provider import BaseProvider, ExtraItem, Replacement
 from gateway.routes.instrumentation import (
     InstrumentedResponse,
     InstrumentedStreamingResponse,
@@ -241,7 +241,7 @@ class AnthropicProvider(BaseProvider):
         self,
         guardrails_execution_result: dict[str, Any],
         location: Literal["request", "response"] = "response",
-    ) -> bytes:
+    ) -> ExtraItem:
         """Anthropic streaming error format (SSE)"""
         error_chunk = json.dumps(
             {
@@ -251,7 +251,7 @@ class AnthropicProvider(BaseProvider):
                 }
             }
         )
-        return f"event: error\ndata: {error_chunk}\n\n".encode()
+        return ExtraItem(f"event: error\ndata: {error_chunk}\n\n".encode(), end_of_stream=True)
 
     def should_push_trace(self, _1: dict[str, Any], _2: bool) -> bool:
         """Anthropic always pushes traces"""
@@ -260,7 +260,34 @@ class AnthropicProvider(BaseProvider):
     def process_streaming_chunk(
         self, chunk: bytes, merged_response: dict[str, Any], chunk_state: dict[str, Any]
     ) -> None:
-        """Anthropic streaming chunk processing"""
+        """
+        Process the chunk and update the merged_response.
+        Each chunk may contain multiple events, separated by double newlines.
+        Each event has type and data fields, separated by a newline.
+        It is possible that a chunk contains some incomplete events.
+
+        Example:
+
+        b'event: message_start\ndata: {"type":"message_start","message":
+        {"id":"msg_01LkayzAaw7b7QkUAw91psyx","type":"message","role":"assistant"
+        ,"model":"claude-3-5-sonnet-20241022","content":[],"stop_reason":null,
+        "stop_sequence":null,"usage":{"input_tokens":20,"cache_creation_input_to'
+
+        and
+
+        b'kens":0,"cache_read_input_tokens":0,"output_tokens":1}}}\n\nevent: content_block_start
+        \ndata: {"type":"content_block_start","index":0,"content_block"
+        :{"type":"text","text":""} }\n\nevent: ping
+        \ndata: {"type": "ping"}\n\nevent: content_block_delta
+        \ndata: {"type":"content_block_delta","index":0,"delta":{"type":
+        "text_delta","text":"Originally"} }\n\n'
+
+        In this case the first chunk ends with 'cache_creation_input_to' which is
+        continued in the next chunk.
+
+        in this case we need to maintain a buffer of the incomplete events.
+        We filter out the ping events and update a merged_response.
+        """
         decoded_chunk = chunk.decode("utf-8", errors="replace")
         chunk_state["sse_buffer"] = chunk_state.get("sse_buffer", "") + decoded_chunk
 
@@ -282,7 +309,7 @@ class AnthropicProvider(BaseProvider):
                     elif line.startswith("data:"):
                         event_data = line[5:].strip()
 
-                if event_data and event_type != "ping":
+                if event_data and event_type != "ping": # Ignore ping events
                     try:
                         event_json = json.loads(event_data)
                         update_merged_response(event_json, merged_response)
@@ -320,7 +347,3 @@ class AnthropicProvider(BaseProvider):
     def initialize_streaming_state(self) -> dict[str, Any]:
         """Anthropic streaming state"""
         return {"sse_buffer": ""}
-
-    def streaming_error_should_end_stream(self) -> bool:
-        """Anthropic continues stream on error"""
-        return True

gateway/routes/base_provider.py

@@ -65,7 +65,7 @@ class BaseProvider(ABC):
         guardrails_execution_result: dict[str, Any],
         location: Literal["request", "response"] = "response",
         status_code: int = 400,
-    ) -> ExtraItem:
+    ) -> Replacement:
         """Create provider-specific error response for non-streaming"""
 
     @abstractmethod
@@ -73,7 +73,7 @@ class BaseProvider(ABC):
         self,
         guardrails_execution_result: dict[str, Any],
         location: Literal["request", "response"] = "response",
-    ) -> bytes:
+    ) -> ExtraItem:
         """Create provider-specific error chunk for streaming"""
 
     @abstractmethod
@@ -105,10 +105,6 @@ class BaseProvider(ABC):
     def initialize_streaming_state(self) -> dict[str, Any]:
         """Initialize provider-specific state for streaming (e.g., OpenAI's mappings)"""
 
-    @abstractmethod
-    def streaming_error_should_end_stream(self) -> bool:
-        """Whether streaming errors should end the stream"""
-
     def check_error_in_non_streaming_response(self, response: httpx.Response) -> None:
         """Check response status and parse JSON for non-streaming requests"""
         try:

gateway/routes/gemini.py

@@ -26,7 +26,7 @@ from gateway.converters.gemini_to_invariant import (
     convert_response,
 )
 from gateway.integrations.explorer import fetch_guardrails_from_explorer
-from gateway.routes.base_provider import BaseProvider, Replacement
+from gateway.routes.base_provider import BaseProvider, ExtraItem, Replacement
 from gateway.routes.instrumentation import (
     InstrumentedResponse,
     InstrumentedStreamingResponse,
@@ -288,9 +288,9 @@ class GeminiProvider(BaseProvider):
         self,
         guardrails_execution_result: dict[str, Any],
         location: Literal["request", "response"] = "response",
-    ) -> bytes:
+    ) -> ExtraItem:
         """Gemini streaming error format"""
-        return json.dumps(make_refusal(location, guardrails_execution_result))
+        return ExtraItem(json.dumps(make_refusal(location, guardrails_execution_result)), end_of_stream=True)
 
     def should_push_trace(
         self, merged_response: dict[str, Any], has_errors: bool
@@ -337,7 +337,3 @@ class GeminiProvider(BaseProvider):
     def initialize_streaming_state(self) -> dict[str, Any]:
         """Gemini has no additional state"""
         return {}
-
-    def streaming_error_should_end_stream(self) -> bool:
-        """Gemini ENDS stream on error"""
-        return True

gateway/routes/instrumentation.py

@@ -10,6 +10,7 @@ import httpx
 from fastapi import Response
 
 from gateway.routes.base_provider import BaseProvider, ExtraItem
+from gateway.common.constants import CONTENT_TYPE_JSON
 from gateway.common.guardrails import GuardrailAction
 from gateway.common.request_context import RequestContext
 from gateway.integrations.explorer import (
@@ -157,7 +158,10 @@ class BaseInstrumentedResponse(ABC):
 
         if self.guardrails_execution_result.get("errors", []):
             if self.context.dataset_name:
-                print("Pushing to explorer from inside handle_input_guardrails", flush=True)
+                print(
+                    "Pushing to explorer from inside handle_input_guardrails",
+                    flush=True,
+                )
                 asyncio.create_task(
                     self.push_to_explorer(
                         response_data, self.guardrails_execution_result
@@ -165,12 +169,9 @@ class BaseInstrumentedResponse(ABC):
                 )
 
             if self.is_streaming:
-                return ExtraItem(
-                    self.provider.create_error_chunk(
+                return self.provider.create_error_chunk(
                         self.guardrails_execution_result, location="request"
-                    ),
-                    end_of_stream=True,
-                )
+                    )
             return self.provider.create_non_streaming_error_response(
                 self.guardrails_execution_result, location="request"
             )
@@ -190,7 +191,11 @@ class BaseInstrumentedResponse(ABC):
         if self.guardrails_execution_result.get("errors", []):
             # Push to explorer
             if self.context.dataset_name:
-                print("Pushing to explorer from inside handle_output_guardrails", self.guardrails_execution_result.get("errors", []), flush=True)
+                print(
+                    "Pushing to explorer from inside handle_output_guardrails",
+                    self.guardrails_execution_result.get("errors", []),
+                    flush=True,
+                )
                 asyncio.create_task(
                     self.push_to_explorer(
                         response_data, self.guardrails_execution_result
@@ -198,14 +203,10 @@ class BaseInstrumentedResponse(ABC):
                 )
 
             if self.is_streaming:
-                error_chunk = self.provider.create_error_chunk(
+                return self.provider.create_error_chunk(
                     self.guardrails_execution_result,
                     location="response",
                 )
-                return ExtraItem(
-                    error_chunk,
-                    end_of_stream=self.provider.streaming_error_should_end_stream(),
-                )
             return self.provider.create_non_streaming_error_response(
                 self.guardrails_execution_result,
                 location="response",
@@ -452,7 +453,7 @@ class InstrumentedResponse(BaseInstrumentedResponse):
         yield Response(
             content=response_string,
             status_code=self.response.status_code,
-            media_type="application/json",
+            media_type=CONTENT_TYPE_JSON,
             headers=updated_headers,
         )
 

gateway/routes/open_ai.py

@@ -26,7 +26,7 @@ from gateway.routes.instrumentation import (
     InstrumentedResponse,
     InstrumentedStreamingResponse,
 )
-from gateway.routes.base_provider import BaseProvider, ExtraItem
+from gateway.routes.base_provider import BaseProvider, ExtraItem, Replacement
 
 gateway = APIRouter()
 
@@ -222,9 +222,9 @@ class OpenAIProvider(BaseProvider):
         guardrails_execution_result: dict[str, Any],
         location: Literal["request", "response"] = "response",
         status_code: int = 400,
-    ) -> ExtraItem:
-        """OpenAI non-streaming error format"""
-        return ExtraItem(
+    ) -> Replacement:
+        """OpenAI non-streaming error format, replace the response with the error message"""
+        return Replacement(
             Response(
                 content=json.dumps(
                     {
@@ -235,14 +235,13 @@ class OpenAIProvider(BaseProvider):
                 status_code=status_code,
                 media_type=CONTENT_TYPE_JSON,
             ),
-            end_of_stream=True,
         )
 
     def create_error_chunk(
         self,
         guardrails_execution_result: dict[str, Any],
         location: Literal["request", "response"] = "response",
-    ) -> bytes:
+    ) -> ExtraItem:
         """OpenAI streaming error format"""
         error_chunk = error_chunk = json.dumps(
             {
@@ -252,7 +251,9 @@ class OpenAIProvider(BaseProvider):
                 }
             }
         )
-        return f"data: {error_chunk}\n\n".encode()
+        # return an extra error chunk (without preventing the original
+        # chunk to go through after)
+        return ExtraItem(f"data: {error_chunk}\n\n".encode(), end_of_stream=True)
 
     def should_push_trace(
         self, merged_response: dict[str, Any], has_errors: bool
@@ -299,10 +300,6 @@ class OpenAIProvider(BaseProvider):
         """OpenAI streaming state"""
         return {"choice_mapping_by_index": {}, "tool_call_mapping_by_index": {}}
 
-    def streaming_error_should_end_stream(self) -> bool:
-        """OpenAI continues stream on error"""
-        return True
-
 
 def process_chunk_text(
     chunk_text: str,