From 0df61524da3020133f585c4d801488a7f62a9516 Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Thu, 27 Mar 2025 14:55:52 +0100 Subject: [PATCH] non-streaming case --- gateway/integrations/guardails.py | 41 ++++++-- gateway/routes/open_ai.py | 159 +++++++++++++++++++++--------- 2 files changed, 147 insertions(+), 53 deletions(-) diff --git a/gateway/integrations/guardails.py b/gateway/integrations/guardails.py index 871454e..5217e1c 100644 --- a/gateway/integrations/guardails.py +++ b/gateway/integrations/guardails.py @@ -205,7 +205,8 @@ class StreamInstrumentor: # schedule all before listeners which can be run concurrently before_tasks = [ - asyncio.create_task(listener()) for listener in self.before_listeners + asyncio.create_task(listener(), name="instrumentor:start") + for listener in self.before_listeners ] # create async iterator from async_iterable @@ -224,7 +225,9 @@ class StreamInstrumentor: self.stat_first_item_time = time.time() - start_first_item_request return r - next_item_task = asyncio.create_task(wait_for_first_item()) + next_item_task = asyncio.create_task( + wait_for_first_item(), name="instrumentor:next:first" + ) # wait for all before listeners to finish has_end_of_stream = False @@ -263,7 +266,9 @@ class StreamInstrumentor: break # schedule next item - next_item_task = asyncio.create_task(aiterable.__anext__()) + next_item_task = asyncio.create_task( + aiterable.__anext__(), name="instrumentor:next" + ) # [STAT] capture token time stamp if len(self.stat_token_times) == 0: @@ -274,8 +279,8 @@ class StreamInstrumentor: ) # invoke on_chunk listeners + any_end_of_stream = False for listener in self.on_chunk_listeners: - any_end_of_stream = False try: await listener(item) except YieldException as e: @@ -293,7 +298,7 @@ class StreamInstrumentor: # finally, execute on complete listeners on_complete_tasks = [ - asyncio.create_task(listener()) + asyncio.create_task(listener(), name="instrumentor:end") for listener in self.on_complete_listeners ] for result in asyncio.as_completed(on_complete_tasks): @@ -306,7 +311,6 @@ class StreamInstrumentor: # [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: @@ -337,6 +341,31 @@ class StreamInstrumentor: print(f" [total: {time.time() - start:.2f}s]") +class RequestInstrumentor(StreamInstrumentor): + """ + Like 'StreamInstrumentor', but for non-streaming requests. + + Supports similar 'start', 'end' events, but not 'chunk', since everything is assumed + to be processed in one chunk (i.e., the request). + """ + + def on(self, event): + assert event in [ + "start", + "end", + ], "RequestInstrumentor does not support 'chunk' events" + return super().on(event) + + async def execute(self, request_task): + async def wrapped_request_task(): + yield await request_task + + # pretend the 'request_task' is an async iterable with a single item + result = [item async for item in self.stream(wrapped_request_task())] + assert len(result) == 1, "RequestInstrumentor should yield exactly one item" + return result[0] + + async def check_guardrails( messages: List[Dict[str, Any]], guardrails: str, invariant_authorization: str ) -> Dict[str, Any]: diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 3692901..4d88bc9 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -14,6 +14,7 @@ from common.constants import ( ) from integrations.explorer import create_annotations_from_guardrails_errors, push_trace from integrations.guardails import ( + RequestInstrumentor, StreamInstrumentor, YieldException, check_guardrails, @@ -84,11 +85,8 @@ async def openai_chat_completions_gateway( client, open_ai_request, ) - response = await client.send(open_ai_request) - return await handle_non_streaming_response( - context, - response, - ) + + return await handle_non_streaming_response(context, client, open_ai_request) async def stream_response( @@ -156,7 +154,7 @@ async def stream_response( error_chunk = json.dumps( { "error": { - "message": "[Invariant] The response did not pass the guardrails", + "message": "[Invariant] The request did not pass the guardrails", "details": guardrails_execution_result, } } @@ -386,7 +384,7 @@ def create_metadata( { key: value for key, value in merged_response.items() - if key in ("usage", "model") + if key in ("usage", "model") and merged_response.get(key) is not None } ) return metadata @@ -421,11 +419,13 @@ async def push_to_explorer( async def get_guardrails_check_result( - context: RequestContextData, json_response: dict[str, Any] + context: RequestContextData, json_response: dict[str, Any] | None = None ) -> dict[str, Any]: """Get the guardrails check result""" messages = list(context.request_json.get("messages", [])) - messages += [choice["message"] for choice in json_response.get("choices", [])] + + if json_response is not None: + messages += [choice["message"] for choice in json_response.get("choices", [])] # Block on the guardrails check guardrails_execution_result = await check_guardrails( @@ -437,48 +437,113 @@ async def get_guardrails_check_result( async def handle_non_streaming_response( - context: RequestContextData, response: httpx.Response + context: RequestContextData, + client: httpx.AsyncClient, + open_ai_request: httpx.Request, ) -> Response: """Handles non-streaming OpenAI responses""" - try: - json_response = response.json() - except json.JSONDecodeError as e: - raise HTTPException( - status_code=response.status_code, - detail="Invalid JSON response received from OpenAI API", - ) from e - if response.status_code != 200: - raise HTTPException( - status_code=response.status_code, - detail=json_response.get("error", "Unknown error from OpenAI API"), - ) - guardrails_execution_result = {} - response_string = json.dumps(json_response) - response_code = response.status_code + instrumentor = RequestInstrumentor() - if context.config and context.config.guardrails: - # Block on the guardrails check - guardrails_execution_result = await get_guardrails_check_result( - context, json_response - ) - if guardrails_execution_result.get("errors", []): - response_string = json.dumps( - { - "error": "[Invariant] The response did not pass the guardrails", - "details": guardrails_execution_result, - } + # respond we get and its JSON decoded version + # available once the 'send_request' function has progressed to the point of + # being able to call 'response.json()' + response = None + json_response = None + + async def send_request(): + nonlocal response, json_response + + response = await client.send(open_ai_request) + + try: + json_response = response.json() + except json.JSONDecodeError as e: + raise HTTPException( + status_code=response.status_code, + detail="Invalid JSON response received from OpenAI API", + ) from e + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail=json_response.get("error", "Unknown error from OpenAI API"), ) - response_code = 400 - if context.dataset_name: - # Push to Explorer - don't block on its response - asyncio.create_task( - push_to_explorer(context, json_response, guardrails_execution_result) + + response_string = json.dumps(json_response) + response_code = response.status_code + + return Response( + content=response_string, + status_code=response_code, + media_type="application/json", + headers=dict(response.headers), ) - return Response( - content=response_string, - status_code=response_code, - media_type="application/json", - headers=dict(response.headers), - ) + @instrumentor.on("start") + async def precheck_guardrails() -> None: + # check guardrails in a pipelined fashion, before processing the first chunk (for input guardrailing) + if context.config and context.config.guardrails: + # block on the guardrails check + guardrails_execution_result = await get_guardrails_check_result(context) + if guardrails_execution_result.get("errors", []): + # replace the response with the error message + raise YieldException( + Response( + content=json.dumps( + { + "error": "[Invariant] The response did not pass the guardrails", + "details": guardrails_execution_result, + } + ), + status_code=400, + media_type="application/json", + ), + end_of_stream=True, + ) + + @instrumentor.on("end") + async def postprocess_guardrails() -> None: + # at this point, we are guaranteed that 'send_request' has already been executed successfully + response_code = response.status_code + + # if we have guardrails, check the response + if context.config and context.config.guardrails: + # run guardrails again, this time on request + response + guardrails_execution_result = await get_guardrails_check_result( + context, json_response + ) + if guardrails_execution_result.get("errors", []): + response_string = json.dumps( + { + "error": "[Invariant] The response did not pass the guardrails", + "details": guardrails_execution_result, + } + ) + response_code = 400 + + # Push annotated trace to the explorer - don't block on its response + if context.dataset_name: + asyncio.create_task( + push_to_explorer( + context, + json_response, + guardrails_execution_result, + ) + ) + + # replace the response with the error message + raise YieldException( + Response( + content=response_string, + status_code=response_code, + media_type="application/json", + ), + ) + + # if we don't have guardrails or if the response passed the guardrails (only then, we reach this point) + if context.dataset_name: + # Push to Explorer - don't block on its response + asyncio.create_task(push_to_explorer(context, json_response)) + + # execute instrumented request + return await instrumentor.execute(send_request())