fix: return earlier

example_policy.gr

@@ -3,4 +3,5 @@ from invariant.detectors import prompt_injection
 raise "Don't say 'Hello'" if:
     (msg: Message)
     msg.role == "user"
-    prompt_injection(msg.content)
+    prompt_injection(msg.content)
+    # "Hello" in msg.content

gateway/integrations/guardails.py

@@ -219,12 +219,15 @@ class StreamInstrumentor:
                 nonlocal start_first_item_request, aiterable
 
                 r = await aiterable.__anext__()
-                self.stat_first_item_time = time.time() - start_first_item_request
+                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())
 
             # wait for all before listeners to finish
+            has_end_of_stream = False
             for before_task in before_tasks:
                 try:
                     await before_task
@@ -240,12 +243,14 @@ class StreamInstrumentor:
                             # [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
-                                )
-                        else:
-                            print(
-                                "before yields, but next item already ready", flush=True
-                            )
+                                    time.time() - start_first_item_request
+                                ) + 0.01
+                        has_end_of_stream = True
+
+            # don't wait for the first item if end_of stream is True
+            if has_end_of_stream:
+                # if end_of_stream is True, stop the stream
+                return
 
             # [STAT] capture before time stamp
             self.stat_before_time = time.time() - start
@@ -281,11 +286,12 @@ class StreamInstrumentor:
 
                 # if end_of_stream is True, stop the stream
                 if any_end_of_stream:
-                    break
+                    return
 
                 # yield item
                 yield item
-            # execute on complete listeners
+
+            # finally, execute on complete listeners
             on_complete_tasks = [
                 asyncio.create_task(listener())
                 for listener in self.on_complete_listeners
@@ -310,6 +316,7 @@ class StreamInstrumentor:
             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)})]"

gateway/routes/open_ai.py

@@ -161,6 +161,17 @@ async def stream_response(
                             }
                         }
                     )
+
+                    # Push annotated trace to the explorer - don't block on its response
+                    if context.dataset_name:
+                        asyncio.create_task(
+                            push_to_explorer(
+                                context,
+                                merged_response,
+                                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
                     raise YieldException(
@@ -215,15 +226,17 @@ async def stream_response(
                     # yield an extra error chunk (without preventing the original chunk to go through after)
                     raise YieldException(f"data: {error_chunk}\n\n".encode())
 
+        @instrumentor.on("end")
+        def send_to_explorer() -> None:
+            # Send full merged response to the explorer
+            # Don't block on the response from explorer
+            if context.dataset_name:
+                asyncio.create_task(push_to_explorer(context, merged_response))
+
         async for chunk in instrumentor.stream(request_and_stream()):
             # Yield chunk to the client
             yield chunk
 
-        # Send full merged response to the explorer
-        # Don't block on the response from explorer
-        if context.dataset_name:
-            asyncio.create_task(push_to_explorer(context, merged_response))
-
     return StreamingResponse(event_generator(), media_type="text/event-stream")