Repository: run-llama/llama_index
Pull request: waiting for events refactor
This refactor attempts to upgrade `ctx.wait_for_event()` to better handle cases where multiple waiters are waiting for events

We do this by tapping into the broker log, and letting every waiter observe the events that come in

Fixes https://github.com/run-llama/llama_index/issues/18220

Problem statement:
[Question]: Parallel Human in Loop with Agent Workflow Issues
### Question Validation

- [x] I have searched both the documentation and discord for an answer.

### Question

Searching and debugging a long time to find a solution. Thanks for any help!

When an agent workflow is prompted to execute multiple tools that contain human in the loop events at once,  the workflow get stuck. 

-For example asking my agent to create 5 'interview questions' results in  attempting to calling the following tool 5 times, as expected. 
- The tool seems to be called in parallel, 5 back to back executions. If context is `set` during these executions, other tools are not aware of it, as it happens at nearly the same time.
-  When the agent reaches the line of code to await a human response event, further executions of the tool call are lost. The agent sometimes only can call the tool 3/5 times before the function awaits the human in the loop response. The context gets serialized and saved... resulting in a loss of tool calls.
- It seems I need to either prevent this tool from being called in parallel, or figure out how to wait for all planned tool calls to execute before awaiting the human response event. 

Thanks again for any help! 

```
# Create Question Tool
async def create_interview_question(
    ctx: Context, current_user: User, interview_uuid: str, question: str
):
    # Start Tool
    # ... logic to check if is_new_question
   
    try:
        if is_new_question:
            ctx.write_event_to_stream(
                InputRequiredEvent(
                    prefix="Confirm you want to add question, respond "yes"."
                )
            )

        event = await ctx.wait_for_event(HumanResponseEvent)

        if event.response.lower() == "yes":
            create_question(...)

            if not interview_question:
                raise Exception("Failed to create interview question.")

            return "The question has been created. Proceed to make the next question. When all questions are created, inform the user!"
        else:
            return f"The user has provided additional input: {event.response}"

    except Exception as e:
        logger.error(e)
        raise e
```

Hints from issue discussion:
Hmm, I can see how this is an issue. Probably some upgrade to the framework would be needed to get around this 

BUT there is several workarounds i can think of
- change your tool to accept a list of questions, then you'll only need a single human event
- update your system so that at some point after questions are added, some human in the loop is required (this can even be done outside the workflow by inspecting the ctx)

Both are valid options to workaround this slight limitation. 
I really appreciate your response and confirmation of the issue! I am glad it's not just me missing something obvious!

These two options are already on my list of next attempts! Thanks for the suggestions here. 

It seems this might actually be something that can cause an issue on any function that has human in the loop with some sort of chat interface. For example even if I batch the questions in a list, it might still be possible that the user asks the agent to create two interviews at once each with a set of questions?

Thanks for listening to my thoughts. 

If you are able to point me in the right direction, I am more than happy to attempt a PR with a potential fix. From the digging I have done, it seems there needs to be a similar `num_workers` property that can be attached to tools? I could be very wrong here! 

Thanks again! I'll be watching for updates here! Super powerful stuff!
I think the main issue here is `ctx.wait_for_event` is a singular queue. You can technically filter incoming events based on attributes, which would ALMOST solve the problem

`ctx.wait_for_event(MyEvent(), requirements={"user_id", "123"})` -- this waits for an event with those specific values

However, this still falls apart because it is a singular queue -- thinking of making this some kind of circular queue/list instead, should have a PR out today :)

Commit message: waiting for events refactor

Patch to review:
diff --git a/llama-index-core/llama_index/core/workflow/context.py b/llama-index-core/llama_index/core/workflow/context.py
index bcc15c58bf..f57daa3ed6 100644
--- a/llama-index-core/llama_index/core/workflow/context.py
+++ b/llama-index-core/llama_index/core/workflow/context.py
@@ -65,6 +65,10 @@ def __init__(
         self._lock = asyncio.Lock()
         self._globals: Dict[str, Any] = {}
 
+        # Asyncio machinery for waiting on events
+        self._new_event_available = asyncio.Event()
+        self._last_seen_broker_idx: int = -1
+
     def _init_broker_data(self) -> None:
         self._waiter_id = str(uuid.uuid4())
         self._queues: Dict[str, asyncio.Queue] = {self._waiter_id: asyncio.Queue()}
@@ -191,6 +195,11 @@ def from_dict(
             for k, v in data["queues"].items()
         }
         context._in_progress = defaultdict(list)
+
+        # explicitly set last seen broker idx to the last index in the broker log
+        # this ensures old events are not re-processed by the wait_for_event method
+        context._last_seen_broker_idx = len(context._broker_log) - 1
+
         return context
 
     async def set(self, key: str, value: Any, make_private: bool = False) -> None:
@@ -352,6 +361,7 @@ def send_event(self, message: Event, step: Optional[str] = None) -> None:
                 )
 
         self._broker_log.append(message)
+        self._new_event_available.set()
 
     async def wait_for_event(
         self,
@@ -373,18 +383,29 @@ async def wait_for_event(
             asyncio.TimeoutError: If the timeout is reached before receiving matching event
         """
         requirements = requirements or {}
+        start_time = time.time()
 
+        last_seen_broker_idx = self._last_seen_broker_idx
         while True:
-            event = await asyncio.wait_for(
-                self._queues[self._waiter_id].get(), timeout=timeout
-            )
-            if type(event) is event_type:
-                if all(
-                    event.get(k, default=None) == v for k, v in requirements.items()
-                ):
-                    return event
-                else:
-                    continue
+            await self._new_event_available.wait()
+            self._new_event_available.clear()  # Clear the event flag after being notified
+
+            # check new events in broker log
+            for i in range(last_seen_broker_idx + 1, len(self._broker_log)):
+                ev = self._broker_log[i]
+                if type(ev) is event_type:
+                    if all(
+                        ev.get(k, default=None) == v for k, v in requirements.items()
+                    ):
+                        return ev
+
+            # update last seen broker idx
+            last_seen_broker_idx = len(self._broker_log) - 1
+
+            if time.time() - start_time > timeout:
+                raise asyncio.TimeoutError(
+                    f"Timeout waiting for event type {event_type} with requirements {requirements} after {timeout} seconds"
+                )
 
     def write_event_to_stream(self, ev: Optional[Event]) -> None:
         self._streaming_queue.put_nowait(ev)
diff --git a/llama-index-core/tests/workflow/test_context.py b/llama-index-core/tests/workflow/test_context.py
index b19a57639c..ee6f0c308c 100644
--- a/llama-index-core/tests/workflow/test_context.py
+++ b/llama-index-core/tests/workflow/test_context.py
@@ -1,12 +1,18 @@
 import asyncio
 import json
-from typing import Optional, Union
+from typing import List, Optional, Union
 from unittest import mock
 
 import pytest
 from llama_index.core.workflow.decorators import StepConfig, step
 from llama_index.core.workflow.errors import WorkflowRuntimeError
-from llama_index.core.workflow.events import Event, StartEvent, StopEvent
+from llama_index.core.workflow.events import (
+    Event,
+    HumanResponseEvent,
+    InputRequiredEvent,
+    StartEvent,
+    StopEvent,
+)
 from llama_index.core.workflow.workflow import (
     Context,
     Workflow,
@@ -188,6 +194,88 @@ async def step1(self, ctx: Context, ev: StartEvent) -> StopEvent:
     assert result == "bar"
 
 
+class Waiter1(Event):
+    msg: str
+
+
+class Waiter2(Event):
+    msg: str
+
+
+class ResultEvent(Event):
+    result: str
+
+
+class WaitingWorkflow(Workflow):
+    @step
+    async def spawn_waiters(self, ctx: Context, ev: StartEvent) -> Waiter1 | Waiter2:
+        ctx.send_event(Waiter1(msg="foo"))
+        ctx.send_event(Waiter2(msg="bar"))
+
+    @step
+    async def waiter_one(self, ctx: Context, ev: Waiter1) -> ResultEvent:
+        ctx.write_event_to_stream(InputRequiredEvent(prefix="waiter_one"))
+
+        ev: HumanResponseEvent = await ctx.wait_for_event(
+            HumanResponseEvent, {"waiter_id": "waiter_one"}
+        )
+        return ResultEvent(result=ev.response)
+
+    @step
+    async def waiter_two(self, ctx: Context, ev: Waiter2) -> ResultEvent:
+        ctx.write_event_to_stream(InputRequiredEvent(prefix="waiter_two"))
+
+        ev: HumanResponseEvent = await ctx.wait_for_event(
+            HumanResponseEvent, {"waiter_id": "waiter_two"}
+        )
+        return ResultEvent(result=ev.response)
+
+    @step
+    async def collect_waiters(self, ctx: Context, ev: ResultEvent) -> StopEvent:
+        events: Optional[List[ResultEvent]] = ctx.collect_events(
+            ev, [ResultEvent, ResultEvent]
+        )
+        if events is None:
+            return None
+
+        return StopEvent(result=[e.result for e in events])
+
+
+@pytest.mark.asyncio()
+async def test_wait_for_multiple_events_in_workflow():
+    workflow = WaitingWorkflow()
+    handler = workflow.run()
+    async for ev in handler.stream_events():
+        if isinstance(ev, InputRequiredEvent) and ev.prefix == "waiter_one":
+            handler.ctx.send_event(
+                HumanResponseEvent(response="foo", waiter_id="waiter_one")
+            )
+        elif isinstance(ev, InputRequiredEvent) and ev.prefix == "waiter_two":
+            handler.ctx.send_event(
+                HumanResponseEvent(response="bar", waiter_id="waiter_two")
+            )
+
+    result = await handler
+    assert result == ["foo", "bar"]
+
+    # serialize and resume
+    ctx_dict = handler.ctx.to_dict()
+    ctx = Context.from_dict(workflow, ctx_dict)
+    handler = workflow.run(ctx=ctx)
+    async for ev in handler.stream_events():
+        if isinstance(ev, InputRequiredEvent) and ev.prefix == "waiter_one":
+            handler.ctx.send_event(
+                HumanResponseEvent(response="fizz", waiter_id="waiter_one")
+            )
+        elif isinstance(ev, InputRequiredEvent) and ev.prefix == "waiter_two":
+            handler.ctx.send_event(
+                HumanResponseEvent(response="buzz", waiter_id="waiter_two")
+            )
+
+    result = await handler
+    assert result == ["fizz", "buzz"]
+
+
 def test_get_holding_events(ctx):
     ctx._step_events_holding = None
     assert ctx.get_holding_events() == []


Code review comments:
diff_hunk: @@ -373,18 +383,29 @@ async def wait_for_event(
             asyncio.TimeoutError: If the timeout is reached before receiving matching event
         """
         requirements = requirements or {}
+        start_time = time.time()
 
+        last_seen_broker_idx = self._last_seen_broker_idx
path: llama-index-core/llama_index/core/workflow/context.py
text: @user1:
I see we only set `self._last_seen_broker_idx` at Context creation, could it be that I need to call `wait_for_event` multiple times on the same Context instance? In that case I think we would keep getting the same events from the first time this function was called.

@author:
hmm, thats true. ugh 🤔

diff_hunk: @@ -373,18 +383,29 @@ async def wait_for_event(
             asyncio.TimeoutError: If the timeout is reached before receiving matching event
         """
         requirements = requirements or {}
+        start_time = time.time()
 
+        last_seen_broker_idx = self._last_seen_broker_idx
         while True:
-            event = await asyncio.wait_for(
-                self._queues[self._waiter_id].get(), timeout=timeout
-            )
-            if type(event) is event_type:
-                if all(
-                    event.get(k, default=None) == v for k, v in requirements.items()
-                ):
-                    return event
-                else:
-                    continue
+            await self._new_event_available.wait()
+            self._new_event_available.clear()  # Clear the event flag after being notified
+
+            # check new events in broker log
+            for i in range(last_seen_broker_idx + 1, len(self._broker_log)):
+                ev = self._broker_log[i]
+                if type(ev) is event_type:
+                    if all(
+                        ev.get(k, default=None) == v for k, v in requirements.items()
+                    ):
+                        return ev
+
+            # update last seen broker idx
+            last_seen_broker_idx = len(self._broker_log) - 1
path: llama-index-core/llama_index/core/workflow/context.py
text: @user1:
Thanks for writing this today. I have spent a bit of time trying to see if it solved the issue but it did not. I might also just be doing things wrong but I hope my testing is helpful! 

1. I now call an agent to trigger multiple events. 
2. I see that all the events get registered, and the last broker IDX is updated, marking them all as seen.
3. Now my system is saving the context, but the context marks all the events as done due to this last seen event. 
4. I reload the context, and it is starting from the next idx, ignoring any of the events that I had previously been waiting on.



@author:
I'm not sure I fully understand -- can you provide a complete minimal script or colab notebook that you are using to test/reproduce this behaviour @user1 ?