Repository: run-llama/llama_index
Pull request: Fix AgentRunner AgentRunStepStartEvent dispatch
# Description

Fixing an issue where in many cases, the `AgentRunStepStartEvent` dispatched in `AgentRunner._run_step` is not assigned the current `TaskStep`. I added an extensive test that verifies all agent-related events being dispatched.

Fixes #14822 

## New Package?

Did I fill in the `tool.llamahub` section in the `pyproject.toml` and provide a detailed README.md for my new integration or package?

- [ ] Yes
- [x] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [ ] Yes
- [x] No

## Type of Change

Please delete options that are not relevant.

- [x] Bug fix (non-breaking change which fixes an issue)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [x] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [x] I stared at the code and made sure it makes sense

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Bug]: AgentRunStepStartEvent may be dispatched without associated TaskStep
### Bug Description

At the very beginning of `AgentRunner._run_step`, an `AgentRunStepStartEvent` is dispatched. In the current implementation, the value of the `TaskStep` object attached to `AgentRunStepStartEvent` will be`None` if the step is popped from `step_queue` (as opposed to being passed explicitly as an argument to `AgentRunner._run_step`). This means that in many cases, the `TaskStep` associated with the current step being executed is not attached to the corresponding `AgentRunStepStartEvent`.

> 

### Version

0.10.55

### Steps to Reproduce

1. Attach an event handler to the root dispatcher to record / inspect events
2. Create an agent that will run one or more steps
3. Run the agent
4. Inspect all `AgentRunStepStartEvent` dispatched, and notice that their `step` field is `None` despite the fact that one or more steps were executed

### Relevant Logs/Tracbacks

_No response_

Hints from issue discussion:

Commit message: Fix AgentRunner AgentRunStepStartEvent dispatch

Patch to review:
diff --git a/llama-index-core/llama_index/core/agent/runner/base.py b/llama-index-core/llama_index/core/agent/runner/base.py
index 0d677d948e..bd9301bbde 100644
--- a/llama-index-core/llama_index/core/agent/runner/base.py
+++ b/llama-index-core/llama_index/core/agent/runner/base.py
@@ -392,15 +392,16 @@ def _run_step(
         **kwargs: Any,
     ) -> TaskStepOutput:
         """Execute step."""
-        dispatcher.event(
-            AgentRunStepStartEvent(task_id=task_id, step=step, input=input)
-        )
         task = self.state.get_task(task_id)
         step_queue = self.state.get_step_queue(task_id)
         step = step or step_queue.popleft()
         if input is not None:
             step.input = input
 
+        dispatcher.event(
+            AgentRunStepStartEvent(task_id=task_id, step=step, input=input)
+        )
+
         if self.verbose:
             print(f"> Running step {step.step_id}. Step input: {step.input}")
 
diff --git a/llama-index-core/tests/agent/runner/test_base.py b/llama-index-core/tests/agent/runner/test_base.py
index 284bd96cee..6bf9f08620 100644
--- a/llama-index-core/tests/agent/runner/test_base.py
+++ b/llama-index-core/tests/agent/runner/test_base.py
@@ -2,7 +2,7 @@
 
 import uuid
 from typing import Any, cast
-
+import llama_index.core.instrumentation as instrument
 from llama_index.core.agent.runner.base import AgentRunner
 from llama_index.core.agent.runner.parallel import ParallelAgentRunner
 from llama_index.core.agent.types import (
@@ -13,6 +13,27 @@
 )
 from llama_index.core.base.llms.types import ChatMessage, MessageRole
 from llama_index.core.chat_engine.types import AgentChatResponse
+from llama_index.core.instrumentation.event_handlers.base import BaseEventHandler
+from llama_index.core.instrumentation.events.agent import (
+    AgentChatWithStepEndEvent,
+    AgentChatWithStepStartEvent,
+    AgentRunStepEndEvent,
+    AgentRunStepStartEvent,
+)
+from llama_index.core.instrumentation.events.base import BaseEvent
+
+dispatcher = instrument.get_dispatcher()
+
+
+class _TestEventHandler(BaseEventHandler):
+    events = []
+
+    @classmethod
+    def class_name(cls):
+        return "_TestEventHandler"
+
+    def handle(self, e: BaseEvent):
+        self.events.append(e)
 
 
 # define mock agent worker
@@ -266,3 +287,61 @@ def test_agent_from_llm() -> None:
     llm = MockLLM()
     agent_runner = AgentRunner.from_llm(llm=llm)
     assert isinstance(agent_runner, ReActAgent)
+
+
+def test_agent_dispatches_events() -> None:
+    event_handler = _TestEventHandler()
+    dispatcher.add_event_handler(event_handler)
+
+    num_steps = 3
+    chat_message_input = "hello world"
+    agent_runner = AgentRunner(agent_worker=MockAgentWorker(limit=num_steps))
+    response = agent_runner.chat(chat_message_input)
+
+    # Expect start / end event pair for agent chat, then a pair for each step run
+    assert len(event_handler.events) == (2 + (2 * num_steps))
+
+    # Check AgentChatWithStepStartEvent
+    assert isinstance(event_handler.events[0], AgentChatWithStepStartEvent)
+    assert event_handler.events[0].user_msg == chat_message_input
+
+    # Check all AgentRunStepStartEvent
+    for i in range(num_steps):
+        run_step_start_idx = 2 * i + 1
+        assert isinstance(
+            event_handler.events[run_step_start_idx], AgentRunStepStartEvent
+        )
+        assert event_handler.events[run_step_start_idx].step is not None
+        assert (
+            event_handler.events[run_step_start_idx].step.task_id
+            == event_handler.events[run_step_start_idx].task_id
+        )
+        assert event_handler.events[run_step_start_idx].input is None
+
+    # Check all AgentRunStepEndEvent
+    for i in range(num_steps):
+        run_step_start_idx = 2 * i + 2
+        assert isinstance(
+            event_handler.events[run_step_start_idx], AgentRunStepEndEvent
+        )
+        assert (
+            event_handler.events[run_step_start_idx].step_output.output.response
+            == f"counter: {i+1}"
+        )
+        assert (
+            event_handler.events[run_step_start_idx].step_output.task_step
+            == event_handler.events[run_step_start_idx - 1].step
+        )
+        assert len(event_handler.events[run_step_start_idx].step_output.next_steps) == 1
+
+        # NOTE: MockAgentWorker generates a next step for the last step, so don't test this
+        is_last_step = i == (num_steps - 1)
+        if not is_last_step:
+            assert (
+                event_handler.events[run_step_start_idx].step_output.next_steps[0]
+                == event_handler.events[run_step_start_idx + 1].step
+            )
+
+    # Check AgentChatWithStepEndEvent
+    assert isinstance(event_handler.events[-1], AgentChatWithStepEndEvent)
+    assert event_handler.events[-1].response == response


Code review comments:
diff_hunk: @@ -266,3 +287,61 @@ def test_agent_from_llm() -> None:
     llm = MockLLM()
     agent_runner = AgentRunner.from_llm(llm=llm)
     assert isinstance(agent_runner, ReActAgent)
+
+
+def test_agent_dispatches_events() -> None:
+    event_handler = _TestEventHandler()
+    dispatcher.add_event_handler(event_handler)
+
+    num_steps = 3
+    chat_message_input = "hello world"
+    agent_runner = AgentRunner(agent_worker=MockAgentWorker(limit=num_steps))
+    response = agent_runner.chat(chat_message_input)
+
+    # Expect start / end event pair for agent chat, then a pair for each step run
+    assert len(event_handler.events) == (2 + (2 * num_steps))
+
+    # Check AgentChatWithStepStartEvent
+    assert isinstance(event_handler.events[0], AgentChatWithStepStartEvent)
+    assert event_handler.events[0].user_msg == chat_message_input
+
+    # Check all AgentRunStepStartEvent
+    for i in range(num_steps):
+        run_step_start_idx = 2 * i + 1
+        assert isinstance(
+            event_handler.events[run_step_start_idx], AgentRunStepStartEvent
+        )
+        assert event_handler.events[run_step_start_idx].step is not None
+        assert (
+            event_handler.events[run_step_start_idx].step.task_id
+            == event_handler.events[run_step_start_idx].task_id
+        )
+        assert event_handler.events[run_step_start_idx].input is None
+
+    # Check all AgentRunStepEndEvent
+    for i in range(num_steps):
+        run_step_start_idx = 2 * i + 2
path: llama-index-core/tests/agent/runner/test_base.py
text: @user1:
nit: should this be `run_step_end_idx`?

@author:
Oops, good catch, fixed!

@user1:
lgtm! thanks :)