feat: invoking events with streaming

Author: gusye1234

drive_flow/core.py

@@ -119,13 +119,26 @@ async def run_event(current_event: BaseEvent, current_event_input: Any):
                             )
                             queue.append((cand_event, build_input))
 
-        while len(queue):
-            this_batch_events = queue[:max_async_events] if max_async_events else queue
-            queue = queue[max_async_events:] if max_async_events else []
-            logger.debug(
-                f"Running a turn with {len(this_batch_events)} event tasks, left {len(queue)} event tasks in queue"
-            )
-            await asyncio.gather(
-                *[run_event(*run_event_input) for run_event_input in this_batch_events]
-            )
+        tasks = set()
+        try:
+            while len(queue) or len(tasks):
+                this_batch_events = (
+                    queue[:max_async_events] if max_async_events else queue
+                )
+                queue = queue[max_async_events:] if max_async_events else []
+                new_tasks = {
+                    asyncio.create_task(run_event(*run_event_input))
+                    for run_event_input in this_batch_events
+                }
+                tasks.update(new_tasks)
+                done, tasks = await asyncio.wait(
+                    tasks, return_when=asyncio.FIRST_COMPLETED
+                )
+                for task in done:
+                    await task  # Handle any exceptions
+        except asyncio.CancelledError:
+            for task in tasks:
+                task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
         return this_run_ctx

tests/test_run.py

@@ -1,3 +1,4 @@
+import asyncio
 import pytest
 from drive_flow import default_drive, EventInput
 from drive_flow.types import ReturnBehavior
@@ -70,6 +71,7 @@ async def a(event: EventInput, global_ctx):
 
     @default_drive.listen_group([start])
     async def b(event: EventInput, global_ctx):
+        await asyncio.sleep(0.2)
         return 2
 
     @default_drive.listen_group([a, b])