From 0157aad27b3a44d0bb3aafbb90b93478617698ad Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Thu, 27 Mar 2025 15:00:58 +0100 Subject: [PATCH] handle non-streaming case --- gateway/integrations/guardails.py | 5 ++++- gateway/routes/open_ai.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/gateway/integrations/guardails.py b/gateway/integrations/guardails.py index 5217e1c..d77386b 100644 --- a/gateway/integrations/guardails.py +++ b/gateway/integrations/guardails.py @@ -171,6 +171,10 @@ class StreamInstrumentor: """ def decorator(func): + assert asyncio.iscoroutinefunction( + func + ), "Listener must be an async function" + if event == "chunk": if self.on_chunk_listeners is None: self.on_chunk_listeners = [] @@ -296,7 +300,6 @@ class StreamInstrumentor: # yield item yield item - # finally, execute on complete listeners on_complete_tasks = [ asyncio.create_task(listener(), name="instrumentor:end") for listener in self.on_complete_listeners diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index 4d88bc9..48696d8 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -225,7 +225,7 @@ async def stream_response( raise YieldException(f"data: {error_chunk}\n\n".encode()) @instrumentor.on("end") - def send_to_explorer() -> None: + async def send_to_explorer() -> None: # Send full merged response to the explorer # Don't block on the response from explorer if context.dataset_name: