diff --git a/gateway/common/constants.py b/gateway/common/constants.py index 02cea2c..20ab1d8 100644 --- a/gateway/common/constants.py +++ b/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", diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index 7938a70..add8839 100644 --- a/gateway/integrations/guardrails.py +++ b/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, ) diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 846f74e..a4915f0 100644 --- a/gateway/routes/anthropic.py +++ b/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 diff --git a/gateway/routes/base_provider.py b/gateway/routes/base_provider.py index 4589ffb..2d97d8a 100644 --- a/gateway/routes/base_provider.py +++ b/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: diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 49e9dea..d43bd59 100644 --- a/gateway/routes/gemini.py +++ b/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 diff --git a/gateway/routes/instrumentation.py b/gateway/routes/instrumentation.py index 4974cd0..f99d266 100644 --- a/gateway/routes/instrumentation.py +++ b/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, ) diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 36bb83b..461f727 100644 --- a/gateway/routes/open_ai.py +++ b/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,