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

Refactor the LLM provider routes to move common functionalities into a BaseInstrumentedResponse class and move provier specific implementations in the corresponding BaseProvider implementations.

Browse Source
This commit is contained in:
Hemang
2025-06-11 15:20:10 +02:00
committed by Hemang Sarkar
parent 42a9c1cc30
commit 9f564a0401
9 changed files with 1155 additions and 1655 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/common/constants.py
+1
View File
@@ -4,6 +4,7 @@ DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
IGNORED_HEADERS = [
"accept-encoding",
"content-length",
"host",
"invariant-authorization",
"x-forwarded-for",
gateway/integrations/guardrails.py
+4 -226
View File
@@ -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"<ExtraItem value={self.value} end_of_stream={self.end_of_stream}>"
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"<Replacement value={self.value}>"
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],
gateway/routes/anthropic.py
+190 -506
View File
@@ -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
gateway/routes/base_provider.py
+144
View File
@@ -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"<ExtraItem value={self.value} end_of_stream={self.end_of_stream}>"
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"<Replacement value={self.value}>"
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)
gateway/routes/gemini.py
+195 -472
View File
@@ -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/<version>/models/<model-name>:generateContent \
/api/v1/gateway/<dataset-name>/gemini/<version>models/<model-name>:generateContent \
/api/v1/gateway/gemini/<version>/models/<model-name>:streamGenerateContent or \
/api/v1/gateway/<dataset-name>/gemini/<version>models/<model-name>: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
gateway/routes/instrumentation.py
+468
View File
@@ -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]
gateway/routes/open_ai.py
+151 -448
View File
@@ -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()
tests/integration/anthropic/test_anthropic_with_tool_call.py
+1 -1
View File
@@ -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"
tests/integration/guardrails/test_header_guardrails.py
+1 -2
View File
@@ -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"