From e66232215ec9df3caac24b7655c33b850366632f Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Thu, 27 Mar 2025 15:14:22 +0100 Subject: [PATCH] fix more cases --- gateway/integrations/guardails.py | 4 ++-- gateway/routes/open_ai.py | 10 ++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/gateway/integrations/guardails.py b/gateway/integrations/guardails.py index d77386b..feabfb8 100644 --- a/gateway/integrations/guardails.py +++ b/gateway/integrations/guardails.py @@ -365,8 +365,8 @@ class RequestInstrumentor(StreamInstrumentor): # 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] + assert len(result) >= 1, "RequestInstrumentor must yield at least one item" + return result[-1] async def check_guardrails( diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 48696d8..0418e2d 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -486,6 +486,16 @@ async def handle_non_streaming_response( # block on the guardrails check guardrails_execution_result = await get_guardrails_check_result(context) if guardrails_execution_result.get("errors", []): + # Push annotated trace to the explorer - don't block on its response + if context.dataset_name: + asyncio.create_task( + push_to_explorer( + context, + {}, + guardrails_execution_result, + ) + ) + # replace the response with the error message raise YieldException( Response(