fix: make ctx._events_buffer json-serializable (#17676)

llama-index-core/llama_index/core/workflow/context.py

@@ -1,14 +1,14 @@
 import asyncio
 import json
-import warnings
 import uuid
+import warnings
 from collections import defaultdict
-from typing import Dict, Any, Optional, List, Type, TYPE_CHECKING, Set, Tuple, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Type, TypeVar
 
 from .context_serializers import BaseSerializer, JsonSerializer
 from .decorators import StepConfig
-from .events import Event
 from .errors import WorkflowRuntimeError
+from .events import Event
 
 if TYPE_CHECKING:  # pragma: no cover
     from .workflow import Workflow
@@ -60,7 +60,7 @@ def __init__(
         self._lock = asyncio.Lock()
         self._globals: Dict[str, Any] = {}
         # Step-specific instance
-        self._events_buffer: Dict[Type[Event], List[Event]] = defaultdict(list)
+        self._events_buffer: Dict[str, List[Event]] = defaultdict(list)
 
     def _serialize_queue(self, queue: asyncio.Queue, serializer: BaseSerializer) -> str:
         queue_items = list(queue._queue)  # type: ignore
@@ -240,14 +240,17 @@ def session(self) -> "Context":
         warnings.warn(msg, DeprecationWarning)
         return self
 
+    def _get_full_path(self, ev_type: Type[Event]) -> str:
+        return f"{ev_type.__module__}.{ev_type.__name__}"
+
     def collect_events(
         self, ev: Event, expected: List[Type[Event]]
     ) -> Optional[List[Event]]:
-        self._events_buffer[type(ev)].append(ev)
+        self._events_buffer[self._get_full_path(type(ev))].append(ev)
 
         retval: List[Event] = []
         for e_type in expected:
-            e_instance_list = self._events_buffer.get(e_type)
+            e_instance_list = self._events_buffer.get(self._get_full_path(e_type))
             if e_instance_list:
                 retval.append(e_instance_list.pop(0))
 
@@ -256,7 +259,7 @@ def collect_events(
 
         # put back the events if unable to collect all
         for ev in retval:
-            self._events_buffer[type(ev)].append(ev)
+            self._events_buffer[self._get_full_path(type(ev))].append(ev)
 
         return None
 

llama-index-core/tests/workflow/test_context.py

@@ -1,18 +1,18 @@
 import asyncio
+import json
+from typing import Optional, Union
 from unittest import mock
-from typing import Union, Optional
 
 import pytest
+from llama_index.core.workflow.decorators import step
+from llama_index.core.workflow.errors import WorkflowRuntimeError
+from llama_index.core.workflow.events import Event, StartEvent, StopEvent
 from llama_index.core.workflow.workflow import (
-    Workflow,
     Context,
+    Workflow,
 )
-from llama_index.core.workflow.decorators import step
-from llama_index.core.workflow.errors import WorkflowRuntimeError
-from llama_index.core.workflow.events import StartEvent, StopEvent, Event
-from llama_index.core.workflow.workflow import Workflow
 
-from .conftest import OneTestEvent, AnotherTestEvent
+from .conftest import AnotherTestEvent, OneTestEvent
 
 
 @pytest.mark.asyncio()
@@ -115,6 +115,11 @@ def test_get_result(ctx):
     assert ctx.get_result() == 42
 
 
+def test_to_dict_with_events_buffer(ctx):
+    ctx.collect_events(OneTestEvent(), [OneTestEvent, AnotherTestEvent])
+    assert json.dumps(ctx.to_dict())
+
+
 @pytest.mark.asyncio()
 async def test_deprecated_params(ctx):
     with pytest.warns(

llama-index-core/tests/workflow/test_workflow.py

@@ -220,6 +220,9 @@ async def original_step(
             ctx.session.send_event(OneTestEvent(test_param="test2"))
             ctx.session.send_event(OneTestEvent(test_param="test3"))
 
+            # send one extra event
+            ctx.session.send_event(AnotherTestEvent(another_test_param="test4"))
+
             return LastEvent()
 
         @step(num_workers=3)
@@ -241,11 +244,21 @@ async def final_step(
     workflow = NumWorkersWorkflow()
 
     start_time = time.time()
-    result = await workflow.run()
+    handler = workflow.run()
+    result = await handler
     end_time = time.time()
 
     assert workflow.is_done()
-    assert set(result) == {"test1", "test2", "test3"}
+    assert set(result) == {"test1", "test2", "test4"}
+
+    # ctx should have 1 extra event
+    assert (
+        len(handler.ctx._events_buffer["tests.workflow.conftest.AnotherTestEvent"]) == 1
+    )
+
+    # ensure ctx is serializable
+    ctx = handler.ctx
+    ctx.to_dict()
 
     # Check if the execution time is close to 1 second (with some tolerance)
     execution_time = end_time - start_time