From 9f564a0401ea2ad886035d673d0c2e884b631619 Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 11 Jun 2025 15:20:10 +0200 Subject: [PATCH] Refactor the LLM provider routes to move common functionalities into a BaseInstrumentedResponse class and move provier specific implementations in the corresponding BaseProvider implementations. --- gateway/common/constants.py | 1 + gateway/integrations/guardrails.py | 230 +----- gateway/routes/anthropic.py | 696 +++++------------- gateway/routes/base_provider.py | 144 ++++ gateway/routes/gemini.py | 667 +++++------------ gateway/routes/instrumentation.py | 468 ++++++++++++ gateway/routes/open_ai.py | 599 ++++----------- .../test_anthropic_with_tool_call.py | 2 +- .../guardrails/test_header_guardrails.py | 3 +- 9 files changed, 1155 insertions(+), 1655 deletions(-) create mode 100644 gateway/routes/base_provider.py create mode 100644 gateway/routes/instrumentation.py diff --git a/gateway/common/constants.py b/gateway/common/constants.py index 20ab1d8..02cea2c 100644 --- a/gateway/common/constants.py +++ b/gateway/common/constants.py @@ -4,6 +4,7 @@ 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 7b2a28f..7938a70 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -3,19 +3,18 @@ import asyncio import os import time -from typing import Any from functools import wraps +from typing import Any -from fastapi import HTTPException import httpx +from fastapi import HTTPException from gateway.common.constants import DEFAULT_API_URL -from gateway.common.guardrails import Guardrail from gateway.common.request_context import RequestContext from gateway.common.authorization import ( INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER, ) - +from gateway.common.guardrails import Guardrail # Timestamps of last API calls per guardrails string _guardrails_cache = {} @@ -113,231 +112,10 @@ async def preload_guardrails(context: "RequestContext") -> None: ) ) asyncio.shield(task) - except Exception as e: + except Exception as e: # pylint: disable=broad-exception-caught print(f"Error scheduling preload_guardrails task: {e}") -class ExtraItem: - """ - Return this class in a instrumented stream callback, to yield an extra item in the resulting stream. - """ - - def __init__(self, value, end_of_stream=False): - self.value = value - self.end_of_stream = end_of_stream - - def __str__(self): - return f"" - - -class Replacement(ExtraItem): - """ - Like ExtraItem, but used to replace the full request result in case of 'InstrumentedResponse'. - """ - - def __init__(self, value): - super().__init__(value, end_of_stream=True) - - def __str__(self): - return f"" - - -class InstrumentedStreamingResponse: - def __init__(self): - # request statistics - self.stat_token_times = [] - self.stat_before_time = None - self.stat_after_time = None - - self.stat_first_item_time = None - - async def on_chunk(self, chunk: Any) -> ExtraItem | None: - """ - This called will be called on every chunk (async). - """ - pass - - async def on_start(self) -> ExtraItem | None: - """ - Decorator to register a listener for start events. - """ - pass - - async def on_end(self) -> ExtraItem | None: - """ - Decorator to register a listener for end events. - """ - pass - - async def event_generator(self): - """ - Streams the async iterable and invokes all instrumented hooks. - - Args: - async_iterable: An async iterable to stream. - - Yields: - The streamed data. - """ - raise NotImplementedError("This method should be implemented in a subclass.") - - async def instrumented_event_generator(self): - """ - Streams the async iterable and invokes all instrumented hooks. - - Args: - async_iterable: An async iterable to stream. - - Yields: - The streamed data. - """ - try: - start = time.time() - - # schedule on_start which can be run concurrently - start_task = asyncio.create_task(self.on_start(), name="instrumentor:start") - - # create async iterator from async_iterable - aiterable = aiter(self.event_generator()) - - # [STAT] capture start time of first item - start_first_item_request = time.time() - - # waits for first item of the iterable - async def wait_for_first_item(): - nonlocal start_first_item_request, aiterable - - r = await aiterable.__anext__() - if self.stat_first_item_time is None: - # [STAT] capture time to first item - self.stat_first_item_time = time.time() - start_first_item_request - return r - - next_item_task = asyncio.create_task( - wait_for_first_item(), name="instrumentor:next:first" - ) - - # check if 'start_task' yields an extra item - if extra_item := await start_task: - # yield extra value before any real items - yield extra_item.value - # stop the stream if end_of_stream is True - if extra_item.end_of_stream: - # if first item is already available - if not next_item_task.done(): - # cancel the task - next_item_task.cancel() - # [STAT] capture time to first item to be now +0.01 - if self.stat_first_item_time is None: - self.stat_first_item_time = ( - time.time() - start_first_item_request - ) + 0.01 - # don't wait for the first item if end_of stream is True - return - - # [STAT] capture before time stamp - self.stat_before_time = time.time() - start - - while True: - # wait for first item - try: - item = await next_item_task - except StopAsyncIteration: - break - - # schedule next item - next_item_task = asyncio.create_task( - aiterable.__anext__(), name="instrumentor:next" - ) - - # [STAT] capture token time stamp - if len(self.stat_token_times) == 0: - self.stat_token_times.append(time.time() - start) - else: - self.stat_token_times.append( - time.time() - start - sum(self.stat_token_times) - ) - - if extra_item := await self.on_chunk(item): - yield extra_item.value - # if end_of_stream is True, stop the stream - if extra_item.end_of_stream: - # cancel next task - next_item_task.cancel() - return - - # yield item - yield item - - # run on_end, before closing the stream (may yield an extra value) - if extra_item := await self.on_end(): - # yield extra value before any real items - yield extra_item.value - # we ignore end_of_stream here, because we are already at the end - - # [STAT] capture after time stamp - self.stat_after_time = time.time() - start - finally: - # [STAT] end all open intervals if not already closed - if self.stat_after_time is None: - self.stat_before_time = time.time() - start - if self.stat_after_time is None: - self.stat_after_time = 0 - if self.stat_first_item_time is None: - self.stat_first_item_time = 0 - - # print statistics - token_times_5_decimale = str([f"{x:.5f}" for x in self.stat_token_times]) - print( - f"[STATS]\n [token times: {token_times_5_decimale} ({len(self.stat_token_times)})]" - ) - print(f" [before: {self.stat_before_time:.2f}s] ") - print(f" [time-to-first-item: {self.stat_first_item_time:.2f}s]") - print( - f" [zero-latency: {' TRUE' if self.stat_before_time < self.stat_first_item_time else 'FALSE'}]" - ) - print( - f" [extra-latency: {self.stat_before_time - self.stat_first_item_time:.2f}s]" - ) - print(f" [after: {self.stat_after_time:.2f}s]") - if len(self.stat_token_times) > 0: - print( - f" [average token time: {sum(self.stat_token_times) / len(self.stat_token_times):.2f}s]" - ) - print(f" [total: {time.time() - start:.2f}s]") - - -class InstrumentedResponse(InstrumentedStreamingResponse): - """ - A class to instrument an async request with hooks for concurrent - pre-processing and post-processing (input and output guardrailing). - """ - - async def event_generator(self): - """ - We implement the 'event_generator' as a single item stream, - where the item is the full result of the request. - """ - yield await self.request() - - async def request(self): - """ - This method should be implemented in a subclass to perform the actual request. - """ - raise NotImplementedError("This method should be implemented in a subclass.") - - async def instrumented_request(self): - """ - Returns the 'Response' object of the request, after applying all instrumented hooks. - """ - results = [r async for r in self.instrumented_event_generator()] - assert len(results) >= 1, "InstrumentedResponse must yield at least one item" - - # we return the last item, in case the end callback yields an extra item. Then, - # don't return the actual result but the 'end' result, e.g. for output guardrailing. - return results[-1] - - async def check_guardrails( messages: list[dict[str, Any]], guardrails: list[Guardrail], diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 0fcbb1f..846f74e 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -1,12 +1,11 @@ """Gateway service to forward requests to the Anthropic APIs""" -import asyncio import json -from typing import Any, Optional +from typing import Any, Literal import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response -from starlette.responses import StreamingResponse +from fastapi.responses import StreamingResponse from gateway.common.authorization import extract_authorization_from_headers from gateway.common.config_manager import ( @@ -16,26 +15,20 @@ from gateway.common.config_manager import ( ) from gateway.common.constants import ( CLIENT_TIMEOUT, - CONTENT_TYPE_JSON, CONTENT_TYPE_EVENT_STREAM, + CONTENT_TYPE_JSON, IGNORED_HEADERS, ) -from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet +from gateway.common.guardrails import GuardrailRuleSet from gateway.common.request_context import RequestContext from gateway.converters.anthropic_to_invariant import ( convert_anthropic_to_invariant_message_format, ) -from gateway.integrations.explorer import ( - create_annotations_from_guardrails_errors, - fetch_guardrails_from_explorer, - push_trace, -) -from gateway.integrations.guardrails import ( - ExtraItem, +from gateway.integrations.explorer import fetch_guardrails_from_explorer +from gateway.routes.base_provider import BaseProvider, Replacement +from gateway.routes.instrumentation import ( InstrumentedResponse, InstrumentedStreamingResponse, - Replacement, - check_guardrails, ) gateway = APIRouter() @@ -69,23 +62,30 @@ def validate_headers(x_api_key: str = Header(None)): ) async def anthropic_v1_messages_gateway( request: Request, - dataset_name: str | None = None, # This is None if the client doesn't want to push to Explorer - config: GatewayConfig = Depends(GatewayConfigManager.get_config), # pylint: disable=unused-argument + dataset_name: str | None = None, + config: GatewayConfig = Depends(GatewayConfigManager.get_config), header_guardrails: GuardrailRuleSet = Depends(extract_guardrails_from_header), ): - """Proxy calls to the Anthropic APIs""" + """ + Proxy calls to the Anthropic APIs + + All Anthropic-specific cases (SSE, message conversion) handled by provider + """ + + # Standard Anthropic request setup headers = { k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS } headers["accept-encoding"] = "identity" - invariant_authorization, anthopic_api_key = extract_authorization_from_headers( + invariant_authorization, anthropic_api_key = extract_authorization_from_headers( request, dataset_name, ANTHROPIC_AUTHORIZATION_HEADER ) - headers[ANTHROPIC_AUTHORIZATION_HEADER] = anthopic_api_key + headers[ANTHROPIC_AUTHORIZATION_HEADER] = anthropic_api_key request_body = await request.body() request_json = json.loads(request_body) + client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) anthropic_request = client.build_request( "POST", @@ -94,12 +94,14 @@ async def anthropic_v1_messages_gateway( data=request_body, ) + # Fetch dataset guardrails dataset_guardrails = None if dataset_name: - # Get the guardrails for the dataset from explorer. dataset_guardrails = await fetch_guardrails_from_explorer( dataset_name, invariant_authorization ) + + # Create request context context = RequestContext.create( request_json=request_json, dataset_name=dataset_name, @@ -108,501 +110,32 @@ async def anthropic_v1_messages_gateway( config=config, request=request, ) + + # Create Anthropic provider + provider = AnthropicProvider() + + # Handle streaming vs non-streaming if request_json.get("stream"): - return await handle_streaming_response(context, client, anthropic_request) - return await handle_non_streaming_response(context, client, anthropic_request) - - -def create_metadata( - context: RequestContext, response_json: dict[str, Any] -) -> dict[str, Any]: - """Creates metadata for the trace""" - metadata = {k: v for k, v in context.request_json.items() if k != "messages"} - metadata["via_gateway"] = True - if response_json.get("usage"): - metadata["usage"] = response_json.get("usage") - return metadata - - -def combine_request_and_response_messages( - context: RequestContext, response_json: dict[str, Any] -): - """Combine the request and response messages""" - messages = [] - if "system" in context.request_json: - messages.append( - {"role": "system", "content": context.request_json.get("system")} + # Use the base class directly - it handles SSE processing via the provider + response = InstrumentedStreamingResponse( + context=context, + client=client, + provider_request=anthropic_request, + provider=provider, ) - messages.extend(context.request_json.get("messages", [])) - if len(response_json) > 0: - messages.append(response_json) - return messages - - -async def get_guardrails_check_result( - context: RequestContext, action: GuardrailAction, response_json: dict[str, Any] -) -> dict[str, Any]: - """Get the guardrails check result""" - # Determine which guardrails to apply based on the action - guardrails = ( - context.guardrails.logging_guardrails - if action == GuardrailAction.LOG - else context.guardrails.blocking_guardrails - ) - if not guardrails: - return {} - - messages = combine_request_and_response_messages(context, response_json) - converted_messages = convert_anthropic_to_invariant_message_format(messages) - - # Block on the guardrails check - guardrails_execution_result = await check_guardrails( - messages=converted_messages, - guardrails=guardrails, - context=context, - ) - return guardrails_execution_result - - -async def push_to_explorer( - context: RequestContext, - merged_response: dict[str, Any], - guardrails_execution_result: dict | None = None, -) -> None: - """Pushes the full trace to the Invariant Explorer""" - guardrails_execution_result = guardrails_execution_result or {} - annotations = create_annotations_from_guardrails_errors( - guardrails_execution_result.get("errors", []) - ) - - # Execute the logging guardrails before pushing to Explorer - logging_guardrails_execution_result = await get_guardrails_check_result( - context, - action=GuardrailAction.LOG, - response_json=merged_response, - ) - logging_annotations = create_annotations_from_guardrails_errors( - logging_guardrails_execution_result.get("errors", []) - ) - # Update the annotations with the logging guardrails - annotations.extend(logging_annotations) - - # Combine the messages from the request body and Anthropic response - messages = combine_request_and_response_messages(context, merged_response) - converted_messages = convert_anthropic_to_invariant_message_format(messages) - - _ = await push_trace( - dataset_name=context.dataset_name, - messages=[converted_messages], - invariant_authorization=context.invariant_authorization, - metadata=[create_metadata(context, merged_response)], - annotations=[annotations] if annotations else None, - ) - - -class InstrumentedAnthropicResponse(InstrumentedResponse): - """Instrumented response for Anthropic API""" - - def __init__( - self, - context: RequestContext, - client: httpx.AsyncClient, - anthropic_request: httpx.Request, - ): - super().__init__() - self.context: RequestContext = context - self.client: httpx.AsyncClient = client - self.anthropic_request: httpx.Request = anthropic_request - - # response data - self.response: httpx.Response | None = None - self.response_string: str | None = None - self.response_json: dict[str, Any] | None = None - - # guardrailing response (if any) - self.guardrails_execution_result = {} - - async def on_start(self): - """ - 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={} - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - { - "error": { - "message": "[Invariant] The request did not pass the guardrails", - "details": self.guardrails_execution_result, - } - } - ) - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - {}, - self.guardrails_execution_result, - ) - ) - - # if we find something, we prevent the request from going through - # and return an error instead - return Replacement( - Response( - content=error_chunk, - status_code=400, - media_type=CONTENT_TYPE_JSON, - headers={"content-type": CONTENT_TYPE_JSON}, - ) - ) - - async def request(self): - """Make the request to the Anthropic API.""" - self.response = await self.client.send(self.anthropic_request) - - try: - response_json = self.response.json() - except json.JSONDecodeError as e: - raise HTTPException( - status_code=self.response.status_code, - detail=( - "Invalid JSON response received from Anthropic: " - f"{self.response.text}, got error: {e}" - ), - ) from e - if self.response.status_code != 200: - raise HTTPException( - status_code=self.response.status_code, - detail=response_json.get("error", "Unknown error from Anthropic"), - ) - - self.response_json = response_json - self.response_string = json.dumps(response_json) - - return self._make_response( - content=self.response_string, - status_code=self.response.status_code, + return StreamingResponse( + response.instrumented_event_generator(), + media_type=CONTENT_TYPE_EVENT_STREAM, ) - - def _make_response(self, content: str, status_code: int): - """Creates a new Response object with the correct headers and content""" - assert self.response is not None, "response is None" - - updated_headers = self.response.headers.copy() - updated_headers.pop("Content-Length", None) - - return Response( - content=content, - status_code=status_code, - 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. - """ - # 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" - assert self.response_string is not None, "response_string is None" - - if self.context.guardrails: - # Block on the guardrails check - guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.response_json, - ) - if guardrails_execution_result.get("errors", []): - guardrail_response_string = json.dumps( - { - "error": "[Invariant] The response did not pass the guardrails", - "details": guardrails_execution_result, - } - ) - - # push to explorer (if configured) - if self.context.dataset_name: - # Push to Explorer - don't block on its response - asyncio.create_task( - push_to_explorer( - self.context, - self.response_json, - guardrails_execution_result, - ) - ) - - return Replacement( - self._make_response( - content=guardrail_response_string, - status_code=400, - ) - ) - - # push to explorer (if configured) - if self.context.dataset_name: - # Push to Explorer - don't block on its response - asyncio.create_task( - push_to_explorer( - self.context, self.response_json, guardrails_execution_result - ) - ) - - -async def handle_non_streaming_response( - context: RequestContext, - client: httpx.AsyncClient, - anthropic_request: httpx.Request, -) -> Response: - """Handles non-streaming Anthropic responses""" - response = InstrumentedAnthropicResponse( + response = InstrumentedResponse( context=context, client=client, - anthropic_request=anthropic_request, + provider_request=anthropic_request, + provider=provider, ) - return await response.instrumented_request() -class InstrumentedAnthropicStreamingResponse(InstrumentedStreamingResponse): - """Instrumented streaming response for Anthropic API""" - - def __init__( - self, - context: RequestContext, - client: httpx.AsyncClient, - anthropic_request: httpx.Request, - ): - super().__init__() - - # request parameters - self.context: RequestContext = context - self.client: httpx.AsyncClient = client - self.anthropic_request: httpx.Request = anthropic_request - - # response data - self.merged_response = {} - - # guardrailing response (if any) - self.guardrails_execution_result = {} - - 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). - """ - if self.context.guardrails: - self.guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.merged_response, - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - { - "error": { - "message": "[Invariant] The request did not pass the guardrails", - "details": self.guardrails_execution_result, - } - } - ) - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.merged_response, - self.guardrails_execution_result, - ) - ) - - # if we find something, we end the stream prematurely (end_of_stream=True) - # and yield an error chunk instead of actually beginning the stream - return ExtraItem( - f"event: error\ndata: {error_chunk}\n\n".encode(), - end_of_stream=True, - ) - - async def event_generator(self): - """Actual streaming response generator""" - response = await self.client.send(self.anthropic_request, stream=True) - if response.status_code != 200: - error_content = await response.aread() - try: - error_json = json.loads(error_content) - error_detail = error_json.get("error", "Unknown error from Anthropic") - except json.JSONDecodeError: - error_detail = { - "error": "Failed to decode error response from Anthropic" - } - raise HTTPException(status_code=response.status_code, detail=error_detail) - - # iterate over the response stream - async for chunk in response.aiter_bytes(): - yield chunk - - async def on_chunk(self, chunk): - """ - 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. - """ - # Decode the chunk and add to buffer - decoded_chunk = chunk.decode("utf-8", errors="replace") - self.sse_buffer += decoded_chunk - - # Process complete events from buffer - complete_events, incomplete_events = self.process_complete_events( - self.sse_buffer - ) - self.sse_buffer = incomplete_events - - # Check if we've received message_stop in any events - message_stop_received = False - - # Update the merged_response based on complete events - for event in complete_events: - try: - if "event: message_stop" in event: - message_stop_received = True - - # Extract event data - lines = event.split("\n") - event_type = None - event_data = None - - for line in lines: - if line.startswith("event:"): - event_type = line[6:].strip() - elif line.startswith("data:"): - event_data = line[5:].strip() - - if event_data and event_type != "ping": # Skip ping events - try: - event_json = json.loads(event_data) - update_merged_response(event_json, self.merged_response) - except json.JSONDecodeError as e: - print( - f"JSON parsing error in event: {e}. Event data: {event_data[:100]}...", - flush=True, - ) - except Exception as e: # pylint: disable=broad-except - print(f"Error processing event: {e}", flush=True) - - # on last stream chunk, run output guardrails - if message_stop_received and self.context.guardrails: - # Block on the guardrails check - self.guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.merged_response, - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - { - "type": "error", - "error": { - "message": "[Invariant] The response did not pass the guardrails", - "details": self.guardrails_execution_result, - }, - } - ) - - # yield an extra error chunk (without preventing the original chunk - # to go through after, - # so client gets the proper message_stop event still) - return ExtraItem( - value=f"event: error\ndata: {error_chunk}\n\n".encode() - ) - - def process_complete_events(self, buffer): - """Process the buffer and extract complete SSE events. - - Returns: - tuple[list[str], str]: A tuple containing a list of - complete events and the remaining buffer with incomplete events. - """ - # Split on double newlines which separate SSE events - if not buffer: - return [], "" - events = [] - remaining = buffer - - # Process events that are complete (ending with \n\n) - while "\n\n" in remaining: - pos = remaining.find("\n\n") - if pos >= 0: - event = remaining[: pos + 2] - remaining = remaining[pos + 2 :] - if event.strip(): # Skip empty events - events.append(event) - - return events, remaining - - async def on_end(self): - """on_end: send full merged response to the explorer (if configured)""" - # don't block on the response from explorer (.create_task) - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.merged_response, - self.guardrails_execution_result, - ) - ) - - -async def handle_streaming_response( - context: RequestContext, - client: httpx.AsyncClient, - anthropic_request: httpx.Request, -) -> StreamingResponse: - """Handles streaming Anthropic responses""" - response = InstrumentedAnthropicStreamingResponse( - context=context, - client=client, - anthropic_request=anthropic_request, - ) - - return StreamingResponse( - response.instrumented_event_generator(), media_type=CONTENT_TYPE_EVENT_STREAM - ) - - def update_merged_response( event: dict[str, Any], merged_response: dict[str, Any] ) -> None: @@ -640,3 +173,154 @@ def update_merged_response( merged_response.get("content")[index]["input"] += delta.get("partial_json") elif event_type == MESSAGE_DELTA: merged_response["usage"].update(**event.get("usage")) + + +class AnthropicProvider(BaseProvider): + """Complete Anthropic provider covering all cases""" + + def get_provider_name(self) -> str: + return "anthropic" + + def combine_messages( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> list[dict[str, Any]]: + """Anthropic message combination with format conversion""" + messages = [] + + # Add system message if present (Anthropic-specific) + if "system" in request_json: + messages.append({"role": "system", "content": request_json.get("system")}) + + messages.extend(request_json.get("messages", [])) + + if response_json: + messages.append(response_json) + + return convert_anthropic_to_invariant_message_format(messages) + + def create_metadata( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> dict[str, Any]: + """Anthropic metadata creation""" + metadata = { + k: v + for k, v in request_json.items() + if k not in ["messages", "system"] and v is not None + } + metadata["via_gateway"] = True + + if response_json and response_json.get("usage"): + metadata["usage"] = response_json["usage"] + return metadata + + def create_non_streaming_error_response( + self, + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + status_code: int = 400, + ) -> Replacement: + """Anthropic non-streaming error format""" + error_chunk = json.dumps( + { + "error": { + "message": f"[Invariant] The {location} did not pass the guardrails", + "details": guardrails_execution_result, + } + } + ) + return Replacement( + Response( + content=error_chunk, + status_code=status_code, + media_type=CONTENT_TYPE_JSON, + headers={"content-type": CONTENT_TYPE_JSON}, + ) + ) + + def create_error_chunk( + self, + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + ) -> bytes: + """Anthropic streaming error format (SSE)""" + error_chunk = json.dumps( + { + "error": { + "message": f"[Invariant] The {location} did not pass the guardrails", + "details": guardrails_execution_result, + } + } + ) + return f"event: error\ndata: {error_chunk}\n\n".encode() + + def should_push_trace(self, _1: dict[str, Any], _2: bool) -> bool: + """Anthropic always pushes traces""" + return True + + def process_streaming_chunk( + self, chunk: bytes, merged_response: dict[str, Any], chunk_state: dict[str, Any] + ) -> None: + """Anthropic streaming chunk processing""" + decoded_chunk = chunk.decode("utf-8", errors="replace") + chunk_state["sse_buffer"] = chunk_state.get("sse_buffer", "") + decoded_chunk + + complete_events, incomplete_events = self._process_complete_events( + chunk_state["sse_buffer"] + ) + chunk_state["sse_buffer"] = incomplete_events + + # Update merged response with events + for event in complete_events: + try: + lines = event.split("\n") + event_type = None + event_data = None + + for line in lines: + if line.startswith("event:"): + event_type = line[6:].strip() + elif line.startswith("data:"): + event_data = line[5:].strip() + + if event_data and event_type != "ping": + try: + event_json = json.loads(event_data) + update_merged_response(event_json, merged_response) + except json.JSONDecodeError: + pass + except Exception: # pylint: disable=broad-except + pass + + def _process_complete_events(self, buffer: str) -> tuple[list[str], str]: + """Streaming buffer processing""" + if not buffer: + return [], "" + + events = [] + remaining = buffer + + while "\n\n" in remaining: + pos = remaining.find("\n\n") + if pos >= 0: + event = remaining[: pos + 2] + remaining = remaining[pos + 2 :] + if event.strip(): + events.append(event) + + return events, remaining + + def is_streaming_complete(self, _: dict[str, Any], chunk_text: str = "") -> bool: + """Anthropic completion detection""" + return "message_stop" in chunk_text + + def initialize_streaming_response(self) -> dict[str, Any]: + """Anthropic starts with empty response""" + return {} + + 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 new file mode 100644 index 0000000..4589ffb --- /dev/null +++ b/gateway/routes/base_provider.py @@ -0,0 +1,144 @@ +"""Base LLM Provider Class for Invariant Gateway""" + +import json +from typing import Any, Literal +from abc import ABC, abstractmethod + +import httpx +from fastapi import HTTPException + + +class ExtraItem: + """ + Return this class in a instrumented stream callback, to yield an extra item + in the resulting stream. + """ + + def __init__(self, value, end_of_stream=False): + self.value = value + self.end_of_stream = end_of_stream + + def __str__(self): + return f"" + + +class Replacement(ExtraItem): + """ + Like ExtraItem, but used to replace the full request result in case of 'InstrumentedResponse'. + """ + + def __init__(self, value): + super().__init__(value, end_of_stream=True) + + def __str__(self): + return f"" + + +class BaseProvider(ABC): + """ + Base Provider class that defines the protocol for all providers + (e.g., OpenAI, Anthropic, Gemini). + """ + + @abstractmethod + def get_provider_name(self) -> str: + """Return provider name (e.g., 'openai', 'anthropic', 'gemini')""" + + @abstractmethod + def combine_messages( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> list[dict[str, Any]]: + """ + Combine request and response messages in provider-specific way + Handles message format conversion (e.g., Anthropic/Gemini converters) + """ + + @abstractmethod + def create_metadata( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> dict[str, Any]: + """Create provider-specific metadata""" + + @abstractmethod + def create_non_streaming_error_response( + self, + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + status_code: int = 400, + ) -> ExtraItem: + """Create provider-specific error response for non-streaming""" + + @abstractmethod + def create_error_chunk( + self, + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + ) -> bytes: + """Create provider-specific error chunk for streaming""" + + @abstractmethod + def should_push_trace( + self, merged_response: dict[str, Any], has_errors: bool + ) -> bool: + """Provider-specific logic for when to push traces""" + + @abstractmethod + def process_streaming_chunk( + self, chunk: bytes, merged_response: dict[str, Any], chunk_state: dict[str, Any] + ) -> None: + """ + Process a streaming chunk and update merged_response + chunk_state can hold provider-specific state (e.g., OpenAI's choice_mapping) + """ + + @abstractmethod + def is_streaming_complete( + self, merged_response: dict[str, Any], chunk_text: str = "" + ) -> bool: + """Determine if streaming is complete""" + + @abstractmethod + def initialize_streaming_response(self) -> dict[str, Any]: + """Initialize the merged response structure for streaming""" + + @abstractmethod + 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: + response_json = response.json() + except json.JSONDecodeError as e: + raise HTTPException( + status_code=response.status_code, + detail=f"Invalid JSON response received from {self.get_provider_name()}: " + "{response.text}, error: {e}", + ) from e + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail=response_json.get( + "error", f"Unknown error from {self.get_provider_name()}" + ), + ) + + async def check_error_in_streaming_response(self, response: httpx.Response) -> None: + """Check response status and parse JSON for streaming requests""" + if response.status_code != 200: + error_content = await response.aread() + try: + error_json = json.loads(error_content.decode("utf-8")) + error_detail = error_json.get( + "error", f"Unknown error from {self.get_provider_name()}" + ) + except json.JSONDecodeError: + error_detail = { + "error": f"Failed to parse {self.get_provider_name()} error response" + } + + raise HTTPException(status_code=response.status_code, detail=error_detail) diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index d386929..49e9dea 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -1,11 +1,10 @@ """Gateway service to forward requests to the Gemini APIs""" -import asyncio import json from typing import Any, Literal import httpx -from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from fastapi import APIRouter, Depends, Query, Request, Response from fastapi.responses import StreamingResponse from gateway.common.authorization import extract_authorization_from_headers @@ -16,24 +15,21 @@ from gateway.common.config_manager import ( ) from gateway.common.constants import ( CLIENT_TIMEOUT, - CONTENT_TYPE_JSON, CONTENT_TYPE_EVENT_STREAM, + CONTENT_TYPE_JSON, IGNORED_HEADERS, ) -from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet +from gateway.common.guardrails import GuardrailRuleSet from gateway.common.request_context import RequestContext -from gateway.converters.gemini_to_invariant import convert_request, convert_response -from gateway.integrations.explorer import ( - create_annotations_from_guardrails_errors, - fetch_guardrails_from_explorer, - push_trace, +from gateway.converters.gemini_to_invariant import ( + convert_request, + convert_response, ) -from gateway.integrations.guardrails import ( - ExtraItem, +from gateway.integrations.explorer import fetch_guardrails_from_explorer +from gateway.routes.base_provider import BaseProvider, Replacement +from gateway.routes.instrumentation import ( InstrumentedResponse, InstrumentedStreamingResponse, - Replacement, - check_guardrails, ) gateway = APIRouter() @@ -49,29 +45,34 @@ async def gemini_generate_content_gateway( api_version: str, model: str, endpoint: str, - dataset_name: str | None = None, # This is None if the client doesn't want to push to Explorer + dataset_name: str | None = None, alt: str = Query( None, title="Response Format", description="Set to 'sse' for streaming" ), - config: GatewayConfig = Depends(GatewayConfigManager.get_config), # pylint: disable=unused-argument + config: GatewayConfig = Depends(GatewayConfigManager.get_config), header_guardrails: GuardrailRuleSet = Depends(extract_guardrails_from_header), ) -> Response: - """Proxy calls to the Gemini GenerateContent API""" + """ + Proxy calls to the Gemini APIs + + All Gemini-specific cases (message conversion, end_of_stream behavior) handled by provider + """ + + # Gemini endpoint validation if endpoint not in ["generateContent", "streamGenerateContent"]: return Response( - content="Invalid endpoint - the only endpoints supported are: \ - /api/v1/gateway/gemini//models/:generateContent \ - /api/v1/gateway//gemini/models/:generateContent \ - /api/v1/gateway/gemini//models/:streamGenerateContent or \ - /api/v1/gateway//gemini/models/:streamGenerateContent", + content="Invalid endpoint - only generateContent and streamGenerateContent supported", status_code=400, ) + + # Standard Gemini request setup headers = { k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS + [GEMINI_AUTHORIZATION_FALLBACK_HEADER] } headers["accept-encoding"] = "identity" + invariant_authorization, gemini_api_key = extract_authorization_from_headers( request, dataset_name, @@ -91,6 +92,7 @@ async def gemini_generate_content_gateway( ) if alt == "sse": gemini_api_url += "?alt=sse" + gemini_request = client.build_request( "POST", gemini_api_url, @@ -98,12 +100,14 @@ async def gemini_generate_content_gateway( headers=headers, ) + # Fetch dataset guardrails dataset_guardrails = None if dataset_name: - # Get the guardrails for the dataset dataset_guardrails = await fetch_guardrails_from_explorer( dataset_name, invariant_authorization ) + + # Create request context context = RequestContext.create( request_json=request_json, dataset_name=dataset_name, @@ -112,226 +116,30 @@ async def gemini_generate_content_gateway( config=config, request=request, ) + + # Create Gemini provider + provider = GeminiProvider() + + # Handle streaming and non-streaming if alt == "sse" or endpoint == "streamGenerateContent": - return await stream_response( - context, - client, - gemini_request, + # Use the base class directly - it handles Gemini streaming via the provider + response = InstrumentedStreamingResponse( + context=context, + client=client, + provider_request=gemini_request, + provider=provider, ) - return await handle_non_streaming_response( - context, - client, - gemini_request, - ) - - -class InstrumentedStreamingGeminiResponse(InstrumentedStreamingResponse): - """Instrumented streaming response for Gemini API""" - - def __init__( - self, - context: RequestContext, - client: httpx.AsyncClient, - gemini_request: httpx.Request, - ): - super().__init__() - - # request data - self.context: RequestContext = context - self.client: httpx.AsyncClient = client - self.gemini_request: httpx.Request = gemini_request - - # Store the progressively merged response - self.merged_response = { - "candidates": [{"content": {"parts": []}, "finishReason": None}] - } - - # guardrailing execution result (if any) - self.guardrails_execution_result: dict[str, Any] | None = None - - def make_refusal( - self, - location: Literal["request", "response"], - guardrails_execution_result: dict[str, Any], - ) -> dict: - """Create a refusal response for the given request or response""" - return { - "candidates": [ - { - "content": { - "parts": [ - { - "text": f"[Invariant] The {location} did not pass the guardrails", - } - ], - } - } - ], - "error": { - "code": 400, - "message": f"[Invariant] The {location} did not pass the guardrails", - "details": guardrails_execution_result, - "status": "INVARIANT_GUARDRAILS_VIOLATION", - }, - "promptFeedback": { - "blockReason": "SAFETY", - "block_reason_message": f"[Invariant] The {location} did not pass the guardrails: " - + json.dumps(guardrails_execution_result), - "safetyRatings": [ - { - "category": "HARM_CATEGORY_UNSPECIFIED", - "probability": "HIGH", - "blocked": True, - } - ], - }, - } - - async def on_start(self): - """ - 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={} - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - self.make_refusal("request", self.guardrails_execution_result) - ) - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - {}, - self.guardrails_execution_result, - ) - ) - - # if we find something, we end the stream prematurely (end_of_stream=True) - # and yield an error chunk instead of actually beginning the stream - return ExtraItem( - f"data: {error_chunk}\r\n\r\n".encode(), end_of_stream=True - ) - - async def event_generator(self): - """Event generator for streaming responses""" - response = await self.client.send(self.gemini_request, stream=True) - - if response.status_code != 200: - error_content = await response.aread() - try: - error_json = json.loads(error_content.decode("utf-8")) - error_detail = error_json.get("error", "Unknown error from Gemini API") - except json.JSONDecodeError: - error_detail = {"error": "Failed to parse Gemini error response"} - raise HTTPException(status_code=response.status_code, detail=error_detail) - - async for chunk in response.aiter_bytes(): - yield chunk - - async def on_chunk(self, chunk): - """Processes each chunk of the streaming response""" - chunk_text = chunk.decode().strip() - if not chunk_text: - return - - # Parse and update merged_response incrementally - process_chunk_text(self.merged_response, chunk_text) - - # runs on the last stream item - if ( - self.merged_response.get("candidates", []) - and self.merged_response.get("candidates")[0].get("finishReason", "") - and self.context.guardrails - ): - # Block on the guardrails check - self.guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.merged_response, - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - self.make_refusal("response", self.guardrails_execution_result) - ) - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.merged_response, - self.guardrails_execution_result, - ) - ) - - return ExtraItem( - value=f"data: {error_chunk}\r\n\r\n".encode(), - # for Gemini we have to end the stream prematurely, as the client SDK - # will not stop streaming when it encounters an error - end_of_stream=True, - ) - - async def on_end(self): - """Runs when the stream ends.""" - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.merged_response, - self.guardrails_execution_result, - ) - ) - - -async def stream_response( - context: RequestContext, - client: httpx.AsyncClient, - gemini_request: httpx.Request, -) -> Response: - """Handles streaming the Gemini response to the client""" - - response = InstrumentedStreamingGeminiResponse( + return StreamingResponse( + response.instrumented_event_generator(), + media_type=CONTENT_TYPE_EVENT_STREAM, + ) + response = InstrumentedResponse( context=context, client=client, - gemini_request=gemini_request, + provider_request=gemini_request, + provider=provider, ) - - async def event_generator(): - async for chunk in response.instrumented_event_generator(): - yield chunk - - return StreamingResponse( - event_generator(), - media_type=CONTENT_TYPE_EVENT_STREAM, - ) - - -def process_chunk_text( - merged_response: dict[str, Any], - chunk_text: str, -) -> None: - """Processes the chunk text and updates the merged_response to be sent to the explorer""" - # Split the chunk text into individual JSON strings - # A single chunk can contain multiple "data: " sections - for json_string in chunk_text.split("data: "): - json_string = json_string.replace("data: ", "").strip() - - if not json_string: - continue - - try: - json_chunk = json.loads(json_string) - except json.JSONDecodeError: - print("Warning: Could not parse chunk:", json_string) - - update_merged_response(merged_response, json_chunk) + return await response.instrumented_request() def update_merged_response(merged_response: dict[str, Any], chunk_json: dict) -> None: @@ -367,254 +175,169 @@ def update_merged_response(merged_response: dict[str, Any], chunk_json: dict) -> merged_response["modelVersion"] = chunk_json["modelVersion"] -def create_metadata( - context: RequestContext, response_json: dict[str, Any] -) -> dict[str, Any]: - """Creates metadata for the trace""" - metadata = { - k: v - for k, v in context.request_json.items() - if k not in ("systemInstruction", "contents") +def make_refusal( + location: Literal["request", "response"], + guardrails_execution_result: dict[str, Any], +) -> dict: + """Create a refusal response for the given request or response""" + return { + "candidates": [ + { + "content": { + "parts": [ + { + "text": f"[Invariant] The {location} did not pass the guardrails", + } + ], + } + } + ], + "error": { + "code": 400, + "message": f"[Invariant] The {location} did not pass the guardrails", + "details": guardrails_execution_result, + "status": "INVARIANT_GUARDRAILS_VIOLATION", + }, + "promptFeedback": { + "blockReason": "SAFETY", + "block_reason_message": f"[Invariant] The {location} did not pass the guardrails: " + + json.dumps(guardrails_execution_result), + "safetyRatings": [ + { + "category": "HARM_CATEGORY_UNSPECIFIED", + "probability": "HIGH", + "blocked": True, + } + ], + }, } - metadata["via_gateway"] = True - metadata.update( - { - key: value - for key, value in response_json.items() - if key in ("usageMetadata", "modelVersion") + + +class GeminiProvider(BaseProvider): + """Complete Gemini provider covering all cases""" + + def get_provider_name(self) -> str: + return "gemini" + + def combine_messages( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> list[dict[str, Any]]: + """Gemini message combination with format conversion""" + converted_requests = convert_request(request_json) + converted_responses = convert_response(response_json) if response_json else [] + + return converted_requests + converted_responses + + def create_metadata( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> dict[str, Any]: + """Gemini metadata creation""" + metadata = { + k: v + for k, v in request_json.items() + if k not in ["systemInstruction", "contents"] and v is not None } - ) - return metadata + metadata["via_gateway"] = True + if response_json: + if response_json.get("usageMetadata"): + metadata["usage"] = response_json["usageMetadata"] + if response_json.get("modelVersion"): + metadata["modelVersion"] = response_json["modelVersion"] + return metadata -async def get_guardrails_check_result( - context: RequestContext, action: GuardrailAction, response_json: dict[str, Any] -) -> dict[str, Any]: - """Get the guardrails check result""" - # Determine which guardrails to apply based on the action - guardrails = ( - context.guardrails.logging_guardrails - if action == GuardrailAction.LOG - else context.guardrails.blocking_guardrails - ) - if not guardrails: - return {} - - converted_requests = convert_request(context.request_json) - converted_responses = convert_response(response_json) - - # Block on the guardrails check - guardrails_execution_result = await check_guardrails( - messages=converted_requests + converted_responses, - guardrails=guardrails, - context=context, - ) - return guardrails_execution_result - - -async def push_to_explorer( - context: RequestContext, - response_json: dict[str, Any], - guardrails_execution_result: dict | None = None, -) -> None: - """Pushes the full trace to the Invariant Explorer""" - guardrails_execution_result = guardrails_execution_result or {} - annotations = create_annotations_from_guardrails_errors( - guardrails_execution_result.get("errors", []) - ) - - # Execute the logging guardrails before pushing to Explorer - logging_guardrails_execution_result = await get_guardrails_check_result( - context, - action=GuardrailAction.LOG, - response_json=response_json, - ) - logging_annotations = create_annotations_from_guardrails_errors( - logging_guardrails_execution_result.get("errors", []) - ) - # Update the annotations with the logging guardrails - annotations.extend(logging_annotations) - - converted_requests = convert_request(context.request_json) - converted_responses = convert_response(response_json) - - _ = await push_trace( - dataset_name=context.dataset_name, - messages=[converted_requests + converted_responses], - invariant_authorization=context.invariant_authorization, - metadata=[create_metadata(context, response_json)], - annotations=[annotations] if annotations else None, - ) - - -class InstrumentedGeminiResponse(InstrumentedResponse): - """Instrumented response for Gemini API""" - - def __init__( + def create_non_streaming_error_response( self, - context: RequestContext, - client: httpx.AsyncClient, - gemini_request: httpx.Request, - ): - super().__init__() - - # request data - self.context: RequestContext = context - self.client: httpx.AsyncClient = client - self.gemini_request: httpx.Request = gemini_request - - # response data - self.response: httpx.Response | None = None - self.response_json: dict[str, Any] | None = None - - # guardrails execution result (if any) - self.guardrails_execution_result: dict[str, Any] | None = None - - async def on_start(self): - """ - 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={} + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + status_code: int = 400, + ) -> Replacement: + """Gemini non-streaming error format""" + error_chunk = json.dumps( + { + "error": { + "code": status_code, + "message": f"[Invariant] The {location} did not pass the guardrails", + "details": guardrails_execution_result, + "status": "INVARIANT_GUARDRAILS_VIOLATION", + }, + "prompt_feedback": { + "blockReason": "SAFETY", + "safetyRatings": [ + { + "category": "HARM_CATEGORY_UNSPECIFIED", + "probability": 0.0, + "blocked": True, + } + ], + }, + } + ) + return Replacement( + Response( + content=error_chunk, + status_code=400, + media_type=CONTENT_TYPE_JSON, + headers={ + "Content-Type": CONTENT_TYPE_JSON, + }, ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - { - "error": { - "code": 400, - "message": "[Invariant] The request did not pass the guardrails", - "details": self.guardrails_execution_result, - "status": "INVARIANT_GUARDRAILS_VIOLATION", - }, - "prompt_feedback": { - "blockReason": "SAFETY", - "safetyRatings": [ - { - "category": "HARM_CATEGORY_UNSPECIFIED", - "probability": 0.0, - "blocked": True, - } - ], - }, - } - ) - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - {}, - self.guardrails_execution_result, - ) - ) - - # if we find something, we end the stream prematurely (end_of_stream=True) - # and yield an error chunk instead of actually beginning the stream - return Replacement( - Response( - content=error_chunk, - status_code=400, - media_type=CONTENT_TYPE_JSON, - headers={ - "Content-Type": CONTENT_TYPE_JSON, - }, - ) - ) - - async def request(self): - """Makes the request to the Gemini API and return the response""" - self.response = await self.client.send(self.gemini_request) - - response_string = self.response.text - response_code = self.response.status_code - - try: - self.response_json = self.response.json() - except json.JSONDecodeError as e: - raise HTTPException( - status_code=self.response.status_code, - detail="Invalid JSON response received from Gemini API", - ) from e - if self.response.status_code != 200: - raise HTTPException( - status_code=self.response.status_code, - detail=self.response_json.get("error", "Unknown error from Gemini API"), - ) - - return Response( - content=response_string, - status_code=response_code, - media_type=CONTENT_TYPE_JSON, - headers=dict(self.response.headers), ) - async def on_end(self): - """Runs when the request ends.""" - response_string = json.dumps(self.response_json) - response_code = self.response.status_code + def create_error_chunk( + self, + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + ) -> bytes: + """Gemini streaming error format""" + return json.dumps(make_refusal(location, guardrails_execution_result)) - if self.context.guardrails: - # Block on the guardrails check - guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.response_json, - ) - if guardrails_execution_result.get("errors", []): - response_string = json.dumps( - { - "error": { - "code": 400, - "message": "[Invariant] The response did not pass the guardrails", - "details": guardrails_execution_result, - "status": "INVARIANT_GUARDRAILS_VIOLATION", - }, - } - ) - response_code = 400 + def should_push_trace( + self, merged_response: dict[str, Any], has_errors: bool + ) -> bool: + """Gemini push criteria""" + return has_errors or ( + merged_response.get("candidates", []) + and merged_response["candidates"][0].get("finishReason") is not None + ) - if self.context.dataset_name: - # Push to Explorer - don't block on its response - asyncio.create_task( - push_to_explorer( - self.context, - self.response_json, - guardrails_execution_result, - ) - ) + def process_streaming_chunk( + self, chunk: bytes, merged_response: dict[str, Any], _: dict[str, Any] + ) -> None: + """Gemini streaming hunk processing""" + chunk_text = chunk.decode().strip() + if not chunk_text: + return - return Replacement( - Response( - content=response_string, - status_code=response_code, - media_type=CONTENT_TYPE_JSON, - headers=dict(self.response.headers), - ) - ) + for json_string in chunk_text.split("data: "): + json_string = json_string.replace("data: ", "").strip() - # Otherwise, also push to Explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, self.response_json, guardrails_execution_result - ) - ) + if not json_string: + continue + try: + json_chunk = json.loads(json_string) + update_merged_response(merged_response, json_chunk) + except json.JSONDecodeError: + continue -async def handle_non_streaming_response( - context: RequestContext, - client: httpx.AsyncClient, - gemini_request: httpx.Request, -) -> Response: - """Handles non-streaming Gemini responses""" + def is_streaming_complete( + self, merged_response: dict[str, Any], _: str = "" + ) -> bool: + """Gemini completion detection""" + return ( + merged_response.get("candidates", []) + and merged_response["candidates"][0].get("finishReason", "") != "" + ) - response = InstrumentedGeminiResponse( - context=context, - client=client, - gemini_request=gemini_request, - ) + def initialize_streaming_response(self) -> dict[str, Any]: + """Gemini streaming response structure""" + return {"candidates": [{"content": {"parts": []}, "finishReason": None}]} - return await response.instrumented_request() + 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 new file mode 100644 index 0000000..4974cd0 --- /dev/null +++ b/gateway/routes/instrumentation.py @@ -0,0 +1,468 @@ +"""Instrumentation module for LLM provider routes.""" + +import asyncio +import json +import time +from abc import ABC, abstractmethod +from typing import Any + +import httpx +from fastapi import Response + +from gateway.routes.base_provider import BaseProvider, ExtraItem +from gateway.common.guardrails import GuardrailAction +from gateway.common.request_context import RequestContext +from gateway.integrations.explorer import ( + push_trace, + create_annotations_from_guardrails_errors, +) +from gateway.integrations.guardrails import check_guardrails, preload_guardrails + + +class BaseInstrumentedResponse(ABC): + """ + Base class for instrumented responses that provides common functionality + for both streaming and non-streaming responses. + """ + + def __init__( + self, + context: RequestContext, + client: Any, + provider_request: Any, + provider: BaseProvider, + is_streaming: bool, + ): + """Configure the instrumented response for a specific provider""" + self.context = context + self.client = client + self.provider_request = provider_request + self.provider = provider + self.is_streaming = is_streaming + + # Response tracking + self.response = None + self.response_json = None + self.guardrails_execution_result = {} + + # For streaming: initialize provider-specific response and state + if is_streaming: + self.merged_response = provider.initialize_streaming_response() + self.streaming_state = provider.initialize_streaming_state() + + # request statistics + self.stat_token_times = [] + self.stat_before_time = None + self.stat_after_time = None + self.stat_first_item_time = None + + @abstractmethod + async def event_generator(self): + """ + An async iterable that yields events (e.g., chunks of data). + This method should be implemented by subclasses to provide the actual data source. + """ + + @abstractmethod + async def on_start(self): + """ + Pre-processing hook. + This can be used for input guardrails or other pre-processing tasks. + """ + + @abstractmethod + async def on_end(self): + """ + Post-processing hook. + This can be used for output guardrails or other post-processing tasks. + """ + + @abstractmethod + async def on_chunk(self, chunk: Any): + """ + Process a chunk of data. + This can be used for streaming responses to handle each chunk as it arrives. + """ + + async def check_guardrails_common( + self, messages: list[dict[str, Any]], action: GuardrailAction + ) -> dict[str, Any]: + """Common guardrails checking""" + + guardrails = ( + self.context.guardrails.logging_guardrails + if action == GuardrailAction.LOG + else self.context.guardrails.blocking_guardrails + ) + + if not guardrails: + return {} + + return await check_guardrails( + messages=messages, guardrails=guardrails, context=self.context + ) + + async def push_to_explorer( + self, response_json: dict[str, Any], guardrails_result: dict[str, Any] = None + ) -> None: + """Common explorer integration""" + guardrails_result = guardrails_result or {} + + # Create annotations from blocking guardrails errors + blocking_annotations = create_annotations_from_guardrails_errors( + guardrails_result.get("errors", []) + ) + + # Execute logging guardrails - provider handles message conversion + messages = self.provider.combine_messages( + self.context.request_json, response_json + ) + logging_result = await self.check_guardrails_common( + messages, GuardrailAction.LOG + ) + logging_annotations = create_annotations_from_guardrails_errors( + logging_result.get("errors", []) + ) + + # Combine all annotations + all_annotations = blocking_annotations + logging_annotations + + # Create provider-specific metadata + metadata = self.provider.create_metadata( + self.context.request_json, response_json + ) + + # Push to explorer + await push_trace( + dataset_name=self.context.dataset_name, + messages=[messages], + invariant_authorization=self.context.invariant_authorization, + metadata=[metadata], + annotations=[all_annotations] if all_annotations else None, + ) + + async def handle_input_guardrails(self) -> Any: + """Handle input guardrails""" + if not self.context or not self.context.guardrails: + return None + + asyncio.create_task(preload_guardrails(self.context)) + response_data = getattr(self, "merged_response", {}) + messages = self.provider.combine_messages( + self.context.request_json, response_data + ) + self.guardrails_execution_result = await self.check_guardrails_common( + messages, GuardrailAction.BLOCK + ) + + if self.guardrails_execution_result.get("errors", []): + if self.context.dataset_name: + print("Pushing to explorer from inside handle_input_guardrails", flush=True) + asyncio.create_task( + self.push_to_explorer( + response_data, self.guardrails_execution_result + ) + ) + + if self.is_streaming: + return ExtraItem( + 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" + ) + + async def handle_output_guardrails(self, response_data: dict[str, Any]) -> Any: + """Handle output guardrails""" + if not self.context or not self.context.guardrails: + return None + + messages = self.provider.combine_messages( + self.context.request_json, response_data + ) + self.guardrails_execution_result = await self.check_guardrails_common( + messages, GuardrailAction.BLOCK + ) + + 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) + asyncio.create_task( + self.push_to_explorer( + response_data, self.guardrails_execution_result + ) + ) + + if self.is_streaming: + error_chunk = 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", + ) + + async def push_successful_trace(self, response_data: dict[str, Any]) -> None: + """Push successful trace""" + if self.context.dataset_name: + should_push = self.provider.should_push_trace( + response_data, + bool(self.guardrails_execution_result.get("errors", [])), + ) + if not should_push: + return + + print("Pushing to explorer from push_successful_trace", flush=True) + asyncio.create_task( + self.push_to_explorer(response_data, self.guardrails_execution_result) + ) + + async def instrumented_event_generator(self): + """ + Streams the async iterable and invokes all instrumented hooks. + Common functionality for both streaming and non-streaming responses. + + Args: + async_iterable: An async iterable to stream. + + Yields: + The streamed data. + """ + try: + start = time.time() + + # schedule on_start which can be run concurrently + start_task = asyncio.create_task(self.on_start(), name="instrumentor:start") + + # create async iterator from async_iterable + aiterable = aiter(self.event_generator()) + + # [STAT] capture start time of first item + start_first_item_request = time.time() + + # waits for first item of the iterable + async def wait_for_first_item(): + nonlocal start_first_item_request, aiterable + + r = await aiterable.__anext__() + if self.stat_first_item_time is None: + # [STAT] capture time to first item + self.stat_first_item_time = time.time() - start_first_item_request + return r + + next_item_task = asyncio.create_task( + wait_for_first_item(), name="instrumentor:next:first" + ) + + # check if 'start_task' yields an extra item + if extra_item := await start_task: + # yield extra value before any real items + yield extra_item.value + # stop the stream if end_of_stream is True + if extra_item.end_of_stream: + # if first item is already available + if not next_item_task.done(): + # cancel the task + next_item_task.cancel() + # [STAT] capture time to first item to be now +0.01 + if self.stat_first_item_time is None: + self.stat_first_item_time = ( + time.time() - start_first_item_request + ) + 0.01 + # don't wait for the first item if end_of stream is True + return + + # [STAT] capture before time stamp + self.stat_before_time = time.time() - start + + while True: + # wait for first item + try: + item = await next_item_task + except StopAsyncIteration: + break + + # schedule next item + next_item_task = asyncio.create_task( + aiterable.__anext__(), name="instrumentor:next" + ) + + # [STAT] capture token time stamp + if len(self.stat_token_times) == 0: + self.stat_token_times.append(time.time() - start) + else: + self.stat_token_times.append( + time.time() - start - sum(self.stat_token_times) + ) + + if extra_item := await self.on_chunk(item): + yield extra_item.value + # if end_of_stream is True, stop the stream + if extra_item.end_of_stream: + # cancel next task + next_item_task.cancel() + return + + # yield item + yield item + + # run on_end, before closing the stream (may yield an extra value) + if extra_item := await self.on_end(): + # yield extra value before any real items + yield extra_item.value + # we ignore end_of_stream here, because we are already at the end + + # [STAT] capture after time stamp + self.stat_after_time = time.time() - start + finally: + # [STAT] end all open intervals if not already closed + if self.stat_after_time is None: + self.stat_before_time = time.time() - start + if self.stat_after_time is None: + self.stat_after_time = 0 + if self.stat_first_item_time is None: + self.stat_first_item_time = 0 + + # print statistics + token_times_5_decimale = str([f"{x:.5f}" for x in self.stat_token_times]) + print( + f"[STATS]\n [token times: {token_times_5_decimale} ({len(self.stat_token_times)})]" + ) + print(f" [before: {self.stat_before_time:.2f}s] ") + print(f" [time-to-first-item: {self.stat_first_item_time:.2f}s]") + print( + f" [zero-latency: {' TRUE' if self.stat_before_time < self.stat_first_item_time else 'FALSE'}]" + ) + print( + f" [extra-latency: {self.stat_before_time - self.stat_first_item_time:.2f}s]" + ) + print(f" [after: {self.stat_after_time:.2f}s]") + if len(self.stat_token_times) > 0: + print( + f" [average token time: {sum(self.stat_token_times) / len(self.stat_token_times):.2f}s]" + ) + print(f" [total: {time.time() - start:.2f}s]") + + +class InstrumentedStreamingResponse(BaseInstrumentedResponse): + """A class to instrument streaming for LLM provider responses with guardrailing.""" + + def __init__( + self, + context: RequestContext, + client: httpx.AsyncClient, + provider_request: httpx.Request, + provider: BaseProvider, + ): + super().__init__(context, client, provider_request, provider, is_streaming=True) + + async def on_chunk(self, chunk: Any) -> ExtraItem | None: + """Process a chunk of streaming data and handle guardrails.""" + # Use provider-specific chunk processing + self.provider.process_streaming_chunk( + chunk, self.merged_response, self.streaming_state + ) + + # Check if streaming is complete using provider-specific logic + chunk_text = chunk.decode("utf-8", errors="replace") + if ( + self.provider.is_streaming_complete(self.merged_response, chunk_text) + and self.context.guardrails + ): + return await self.handle_output_guardrails(self.merged_response) + + async def on_start(self) -> ExtraItem | None: + """Run pre-processing before starting the streaming response.""" + return await self.handle_input_guardrails() + + async def on_end(self) -> ExtraItem | None: + """Run post-processing after the streaming response ends.""" + print("Reached on_end: ", self.merged_response, flush=True) + await self.push_successful_trace(self.merged_response) + + async def event_generator(self): + """Generic event generator using provider protocol""" + response = await self.client.send(self.provider_request, stream=True) + await self.provider.check_error_in_streaming_response(response) + + async for chunk in response.aiter_bytes(): + yield chunk + + +class InstrumentedResponse(BaseInstrumentedResponse): + """ + A class to instrument an async request with hooks for concurrent + pre-processing and post-processing (input and output guardrailing). + """ + + def __init__( + self, + context: RequestContext, + client: httpx.AsyncClient, + provider_request: httpx.Request, + provider: BaseProvider, + ): + super().__init__( + context, client, provider_request, provider, is_streaming=False + ) + self.response: httpx.Response | None = None + self.response_json: dict | None = None + + async def on_start(self): + """Input guardrails""" + return await self.handle_input_guardrails() + + async def on_chunk(self, _: Any): + """No-op for non-streaming responses""" + return None + + async def on_end(self): + """Output guardrails and explorer integration""" + if self.response is not None and self.response_json is not None: + # Check output guardrails + result = await self.handle_output_guardrails(self.response_json) + if result: # If guardrails failed + return result + + # Push successful trace + await self.push_successful_trace(self.response_json) + + async def event_generator(self): + """ + We implement the 'event_generator' as a single item stream, + where the item is the full result of the request. + """ + self.response = await self.client.send(self.provider_request) + self.provider.check_error_in_non_streaming_response(self.response) + self.response_json = self.response.json() + + response_string = json.dumps(self.response_json) + updated_headers = dict(self.response.headers) + updated_headers.pop("content-length", None) + + yield Response( + content=response_string, + status_code=self.response.status_code, + media_type="application/json", + headers=updated_headers, + ) + + async def instrumented_request(self): + """ + Returns the 'Response' object of the request, after applying all instrumented hooks. + """ + results = [r async for r in self.instrumented_event_generator()] + assert len(results) >= 1, "InstrumentedResponse must yield at least one item" + + # we return the last item, in case the end callback yields an extra item. Then, + # don't return the actual result but the 'end' result, e.g. for output guardrailing. + return results[-1] diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 1c8e1d9..36bb83b 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -1,8 +1,7 @@ """Gateway service to forward requests to the OpenAI APIs""" -import asyncio import json -from typing import Any +from typing import Any, Literal import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response @@ -16,23 +15,18 @@ from gateway.common.config_manager import ( ) from gateway.common.constants import ( CLIENT_TIMEOUT, - CONTENT_TYPE_JSON, CONTENT_TYPE_EVENT_STREAM, + CONTENT_TYPE_JSON, IGNORED_HEADERS, ) -from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet +from gateway.common.guardrails import GuardrailRuleSet from gateway.common.request_context import RequestContext -from gateway.integrations.explorer import ( - create_annotations_from_guardrails_errors, - fetch_guardrails_from_explorer, - push_trace, -) -from gateway.integrations.guardrails import ( - ExtraItem, +from gateway.integrations.explorer import fetch_guardrails_from_explorer +from gateway.routes.instrumentation import ( InstrumentedResponse, InstrumentedStreamingResponse, - check_guardrails, ) +from gateway.routes.base_provider import BaseProvider, ExtraItem gateway = APIRouter() @@ -78,7 +72,7 @@ async def openai_models_options(request: Request): @gateway.get("/openai/models") async def openai_models_gateway( request: Request, - dataset_name: str | None = None, # This is None if the client doesn't want to push to Explorer + dataset_name: str | None = None, ): """Proxy request to OpenAI /models endpoint""" headers = { @@ -88,6 +82,7 @@ async def openai_models_gateway( request, dataset_name, OPENAI_AUTHORIZATION_HEADER ) headers[OPENAI_AUTHORIZATION_HEADER] = "Bearer " + openai_api_key + async with httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) as client: open_ai_request = client.build_request( "GET", @@ -112,11 +107,17 @@ async def openai_models_gateway( ) async def openai_chat_completions_gateway( request: Request, - dataset_name: str | None = None, # This is None if the client doesn't want to push to Explorer - config: GatewayConfig = Depends(GatewayConfigManager.get_config), # pylint: disable=unused-argument + dataset_name: str | None = None, + config: GatewayConfig = Depends(GatewayConfigManager.get_config), header_guardrails: GuardrailRuleSet = Depends(extract_guardrails_from_header), ) -> Response: - """Proxy calls to the OpenAI APIs""" + """ + Proxy calls to the OpenAI APIs + + All OpenAI-specific cases handled by the provider and base classes + """ + + # Standard OpenAI request setup headers = { k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS } @@ -138,12 +139,14 @@ async def openai_chat_completions_gateway( headers=headers, ) + # Fetch dataset guardrails dataset_guardrails = None if dataset_name: - # Get the guardrails for the dataset dataset_guardrails = await fetch_guardrails_from_explorer( dataset_name, invariant_authorization ) + + # Create request context context = RequestContext.create( request_json=request_json, dataset_name=dataset_name, @@ -152,41 +155,138 @@ async def openai_chat_completions_gateway( config=config, request=request, ) + + # Create OpenAI provider + provider = OpenAIProvider() + + # Handle streaming vs non-streaming if request_json.get("stream", False): - return await handle_stream_response( - context, - client, - open_ai_request, + # Use the base class directly - it handles everything via the provider + response = InstrumentedStreamingResponse( + context=context, + client=client, + provider_request=open_ai_request, + provider=provider, + ) + return StreamingResponse( + response.instrumented_event_generator(), + media_type=CONTENT_TYPE_EVENT_STREAM, + ) + response = InstrumentedResponse( + context=context, + client=client, + provider_request=open_ai_request, + provider=provider, + ) + return await response.instrumented_request() + + +class OpenAIProvider(BaseProvider): + """Complete OpenAI provider covering all cases""" + + def get_provider_name(self) -> str: + return "openai" + + def combine_messages( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> list[dict[str, Any]]: + """Combine request and response messages in OpenAI format""" + messages = list(request_json.get("messages", [])) + if response_json: + messages += [ + choice["message"] for choice in response_json.get("choices", []) + ] + return messages + + def create_metadata( + self, request_json: dict[str, Any], response_json: dict[str, Any] + ) -> dict[str, Any]: + """OpenAI metadata creation""" + metadata = { + k: v for k, v in request_json.items() if k != "messages" and v is not None + } + metadata["via_gateway"] = True + + if response_json: + metadata.update( + { + key: value + for key, value in response_json.items() + if key in ("usage", "model") and value is not None + } + ) + return metadata + + def create_non_streaming_error_response( + self, + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + status_code: int = 400, + ) -> ExtraItem: + """OpenAI non-streaming error format""" + return ExtraItem( + Response( + content=json.dumps( + { + "error": f"[Invariant] The {location} did not pass the guardrails", + "details": guardrails_execution_result, + } + ), + status_code=status_code, + media_type=CONTENT_TYPE_JSON, + ), + end_of_stream=True, ) - return await handle_non_stream_response(context, client, open_ai_request) - - -class InstrumentedOpenAIStreamResponse(InstrumentedStreamingResponse): - """ - Does a streaming OpenAI completion request at the core, but also checks guardrails - before (concurrent) and after the request. - """ - - def __init__( + def create_error_chunk( self, - context: RequestContext, - client: httpx.AsyncClient, - open_ai_request: httpx.Request, - ): - super().__init__() + guardrails_execution_result: dict[str, Any], + location: Literal["request", "response"] = "response", + ) -> bytes: + """OpenAI streaming error format""" + error_chunk = error_chunk = json.dumps( + { + "error": { + "message": f"[Invariant] The {location} did not pass the guardrails", + "details": guardrails_execution_result, + } + } + ) + return f"data: {error_chunk}\n\n".encode() - # request parameters - self.context: RequestContext = context - self.client: httpx.AsyncClient = client - self.open_ai_request: httpx.Request = open_ai_request + def should_push_trace( + self, merged_response: dict[str, Any], has_errors: bool + ) -> bool: + """OpenAI-specific push criteria""" - # guardrailing output (if any) - self.guardrails_execution_result: dict | None = None + return has_errors or not ( + merged_response.get("choices") + and merged_response["choices"][0].get("finish_reason") + not in FINISH_REASON_TO_PUSH_TRACE + ) - # merged_response will be updated with the data from the chunks in the stream - # At the end of the stream, this will be sent to the explorer - self.merged_response = { + def process_streaming_chunk( + self, chunk: bytes, merged_response: dict[str, Any], chunk_state: dict[str, Any] + ) -> None: + """OpenAI streaming chunk processing""" + chunk_text = chunk.decode().strip() + if not chunk_text: + return + + process_chunk_text( + chunk_text, + merged_response, + chunk_state.get("choice_mapping_by_index", {}), + chunk_state.get("tool_call_mapping_by_index", {}), + ) + + def is_streaming_complete(self, _: dict[str, Any], chunk_text: str = "") -> bool: + """OpenAI completion detection""" + return "data: [DONE]" in chunk_text + + def initialize_streaming_response(self) -> dict[str, Any]: + """OpenAI streaming response structure""" + return { "id": None, "object": "chat.completion", "created": None, @@ -195,152 +295,13 @@ class InstrumentedOpenAIStreamResponse(InstrumentedStreamingResponse): "usage": None, } - # Each chunk in the stream contains a list called "choices" each entry in the list - # has an index. - # A choice has a field called "delta" which may contain a list called "tool_calls". - # Maps the choice index in the stream to the index in the merged_response["choices"] list - self.choice_mapping_by_index = {} - # Combines the choice index and tool call index to uniquely identify a tool call - self.tool_call_mapping_by_index = {} + def initialize_streaming_state(self) -> dict[str, Any]: + """OpenAI streaming state""" + return {"choice_mapping_by_index": {}, "tool_call_mapping_by_index": {}} - async def on_start(self): - """ - 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=self.merged_response, - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - { - "error": { - "message": "[Invariant] The request did not pass the guardrails", - "details": self.guardrails_execution_result, - } - } - ) - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.merged_response, - self.guardrails_execution_result, - ) - ) - - # if we find something, we end the stream prematurely (end_of_stream=True) - # and yield an error chunk instead of actually beginning the stream - return ExtraItem( - f"data: {error_chunk}\n\n".encode(), - end_of_stream=True, - ) - - async def on_chunk(self, chunk): - """Processes each chunk of the stream and checks guardrails at the end of the stream""" - # process and check each chunk - chunk_text = chunk.decode().strip() - if not chunk_text: - return - - # Process the chunk - # This will update merged_response with the data from the chunk - process_chunk_text( - chunk_text, - self.merged_response, - self.choice_mapping_by_index, - self.tool_call_mapping_by_index, - ) - - # check guardrails at the end of the stream (on the '[DONE]' SSE chunk.) - if "data: [DONE]" in chunk_text and self.context.guardrails: - # Block on the guardrails check - self.guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.merged_response, - ) - if self.guardrails_execution_result.get("errors", []): - error_chunk = json.dumps( - { - "error": { - "message": "[Invariant] The response did not pass the guardrails", - "details": self.guardrails_execution_result, - } - } - ) - - # 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 - - async def on_end(self): - """Sends full merged response to the explorer.""" - # don't block on the response from explorer (.create_task) - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, self.merged_response, self.guardrails_execution_result - ) - ) - - async def event_generator(self): - """Actual OpenAI stream response.""" - response = await self.client.send(self.open_ai_request, stream=True) - if response.status_code != 200: - error_content = await response.aread() - try: - error_json = json.loads(error_content.decode("utf-8")) - error_detail = error_json.get("error", "Unknown error from OpenAI API") - except json.JSONDecodeError: - error_detail = {"error": "Failed to parse OpenAI error response"} - raise HTTPException(status_code=response.status_code, detail=error_detail) - - # stream out chunks - async for chunk in response.aiter_bytes(): - yield chunk - - -async def handle_stream_response( - context: RequestContext, - client: httpx.AsyncClient, - open_ai_request: httpx.Request, -) -> Response: - """ - Handles streaming the OpenAI response to the client while building a merged_response - The chunks are returned to the caller immediately - The merged_response is built from the chunks as they are received - It is sent to the Invariant Explorer at the end of the stream - """ - - response = InstrumentedOpenAIStreamResponse( - context, - client, - open_ai_request, - ) - - return StreamingResponse( - response.instrumented_event_generator(), media_type=CONTENT_TYPE_EVENT_STREAM - ) - - -def initialize_merged_response() -> dict[str, Any]: - """Initializes the full response dictionary""" - return { - "id": None, - "object": "chat.completion", - "created": None, - "model": None, - "choices": [], - "usage": None, - } + def streaming_error_should_end_stream(self) -> bool: + """OpenAI continues stream on error""" + return True def process_chunk_text( @@ -461,261 +422,3 @@ def update_existing_choice_with_delta( finish_reason = delta.get("finish_reason") if finish_reason is not None: existing_choice["finish_reason"] = finish_reason - - -def create_metadata( - context: RequestContext, merged_response: dict[str, Any] -) -> dict[str, Any]: - """Creates metadata for the trace""" - metadata = { - k: v - for k, v in context.request_json.items() - if k != "messages" and v is not None - } - metadata["via_gateway"] = True - metadata.update( - { - key: value - for key, value in merged_response.items() - if key in ("usage", "model") and merged_response.get(key) is not None - } - ) - return metadata - - -async def push_to_explorer( - context: RequestContext, - merged_response: dict[str, Any], - guardrails_execution_result: dict | None = None, -) -> None: - """Pushes the merged response to the Invariant Explorer""" - # Only push the trace to explorer if the message is an end turn message - # or if the guardrails check returned errors. - guardrails_execution_result = guardrails_execution_result or {} - guardrails_errors = guardrails_execution_result.get("errors", []) - annotations = create_annotations_from_guardrails_errors(guardrails_errors) - # Execute the logging guardrails before pushing to Explorer - logging_guardrails_execution_result = await get_guardrails_check_result( - context, - action=GuardrailAction.LOG, - response_json=merged_response, - ) - logging_annotations = create_annotations_from_guardrails_errors( - logging_guardrails_execution_result.get("errors", []) - ) - # Update the annotations with the logging guardrails - annotations.extend(logging_annotations) - - if annotations or not ( - merged_response.get("choices") - and merged_response["choices"][0].get("finish_reason") - not in FINISH_REASON_TO_PUSH_TRACE - ): - # Combine the messages from the request body and the choices from the OpenAI response - messages = list(context.request_json.get("messages", [])) - messages += [choice["message"] for choice in merged_response.get("choices", [])] - _ = await push_trace( - dataset_name=context.dataset_name, - invariant_authorization=context.invariant_authorization, - messages=[messages], - annotations=[annotations], - metadata=[create_metadata(context, merged_response)], - ) - - -async def get_guardrails_check_result( - context: RequestContext, - action: GuardrailAction, - response_json: dict[str, Any] | None = None, -) -> dict[str, Any]: - """Get the guardrails check result""" - # Determine which guardrails to apply based on the action - guardrails = ( - context.guardrails.logging_guardrails - if action == GuardrailAction.LOG - else context.guardrails.blocking_guardrails - ) - - if not guardrails: - return {} - - messages = list(context.request_json.get("messages", [])) - if response_json is not None: - messages += [choice["message"] for choice in response_json.get("choices", [])] - - # Block on the guardrails check - guardrails_execution_result = await check_guardrails( - messages=messages, - guardrails=guardrails, - context=context, - ) - return guardrails_execution_result - - -class InstrumentedOpenAIResponse(InstrumentedResponse): - """ - Does an OpenAI completion request at the core, but also checks guardrails - before (concurrent) and after the request. - """ - - def __init__( - self, - context: RequestContext, - client: httpx.AsyncClient, - open_ai_request: httpx.Request, - ): - super().__init__() - - # request parameters - self.context: RequestContext = context - self.client: httpx.AsyncClient = client - self.open_ai_request: httpx.Request = open_ai_request - - # request outputs - self.response: httpx.Response | None = None - self.response_json: dict[str, Any] | None = None - - # guardrailing output (if any) - self.guardrails_execution_result: dict | None = None - - async def on_start(self): - """ - Checks guardrails in a pipelined fashion, before processing - the first chunk (for input guardrailing) - """ - if self.context.guardrails: - # block on the guardrails check - self.guardrails_execution_result = await get_guardrails_check_result( - self.context, action=GuardrailAction.BLOCK - ) - if self.guardrails_execution_result.get("errors", []): - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - {}, - self.guardrails_execution_result, - ) - ) - - # replace the response with the error message - return ExtraItem( - Response( - content=json.dumps( - { - "error": "[Invariant] The request did not pass the guardrails", - "details": self.guardrails_execution_result, - } - ), - status_code=400, - media_type=CONTENT_TYPE_JSON, - ), - end_of_stream=True, - ) - - async def request(self): - """Actual OpenAI request.""" - self.response = await self.client.send(self.open_ai_request) - - try: - self.response_json = self.response.json() - except json.JSONDecodeError as e: - raise HTTPException( - status_code=self.response.status_code, - detail="Invalid JSON response received from OpenAI API", - ) from e - if self.response.status_code != 200: - raise HTTPException( - status_code=self.response.status_code, - detail=self.response_json.get("error", "Unknown error from OpenAI API"), - ) - - response_string = json.dumps(self.response_json) - response_code = self.response.status_code - - return Response( - content=response_string, - status_code=response_code, - media_type=CONTENT_TYPE_JSON, - headers=dict(self.response.headers), - ) - - async def on_end(self): - """Postprocesses the OpenAI response and potentially replace it with a guardrails error.""" - - # these two request outputs are guaranteed to be available by the time we reach - # this point (after self.request() was executed) - # nevertheless, we check for them to avoid any potential issues - assert ( - self.response is not None - ), "on_end called before 'self.response' was available" - assert ( - self.response_json is not None - ), "on_end called before 'self.response_json' was available" - - # extract original response status code - response_code = self.response.status_code - - # if we have guardrails, check the response - if self.context.guardrails: - # run guardrails again, this time on request + response - self.guardrails_execution_result = await get_guardrails_check_result( - self.context, - action=GuardrailAction.BLOCK, - response_json=self.response_json, - ) - if self.guardrails_execution_result.get("errors", []): - response_string = json.dumps( - { - "error": "[Invariant] The response did not pass the guardrails", - "details": self.guardrails_execution_result, - } - ) - response_code = 400 - - # Push annotated trace to the explorer - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.response_json, - self.guardrails_execution_result, - ) - ) - - # replace the response with the error message - return ExtraItem( - Response( - content=response_string, - status_code=response_code, - media_type=CONTENT_TYPE_JSON, - ), - ) - - # Push annotated trace to the explorer in any case - don't block on its response - if self.context.dataset_name: - asyncio.create_task( - push_to_explorer( - self.context, - self.response_json, - # include any guardrailing errors if available - self.guardrails_execution_result, - ) - ) - - -async def handle_non_stream_response( - context: RequestContext, - client: httpx.AsyncClient, - open_ai_request: httpx.Request, -) -> Response: - """Handles non-streaming OpenAI responses""" - - response = InstrumentedOpenAIResponse( - context, - client, - open_ai_request, - ) - - return await response.instrumented_request() diff --git a/tests/integration/anthropic/test_anthropic_with_tool_call.py b/tests/integration/anthropic/test_anthropic_with_tool_call.py index 365af35..0af1002 100644 --- a/tests/integration/anthropic/test_anthropic_with_tool_call.py +++ b/tests/integration/anthropic/test_anthropic_with_tool_call.py @@ -259,7 +259,7 @@ async def test_streaming_response_with_tool_call( elif len(response) == 1: # expected output in this case is something like this: # [[TextBlock(text="I'll help you check the weather in New York using the get_weather function.", type='text', citations=None), ToolUseBlock(id='toolu_019VZsmxuUhShou2EpPBxvpe', input={'location': 'New York, NY', 'unit': 'celsius'}, name='get_weather', type='tool_use')]] - + assert response is not None assert response[0][0].type == "text" assert response[0][1].type == "tool_use" diff --git a/tests/integration/guardrails/test_header_guardrails.py b/tests/integration/guardrails/test_header_guardrails.py index e03847d..f51c136 100644 --- a/tests/integration/guardrails/test_header_guardrails.py +++ b/tests/integration/guardrails/test_header_guardrails.py @@ -171,12 +171,11 @@ raise "Users must not mention the magic phrase 'Abracadabra'" if: if not do_stream: with pytest.raises(BadRequestError) as exc_info: - chat_response = client.chat.completions.create( + _ = client.chat.completions.create( **request, stream=False, ) - print(exc_info.value.message, flush=True) assert "Failed to create policy from policy source." in str( exc_info.value ), "guardrails check fails because of an invalid guardrailing rule"