[Feature Request] Ability to cancel workflow execution

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

@@ -29,6 +29,7 @@ def __init__(self, workflow: "Workflow", stepwise: bool = False) -> None:
         self._queues: Dict[str, asyncio.Queue] = {}
         self._tasks: Set[asyncio.Task] = set()
         self._broker_log: List[Event] = []
+        self._cancel_flag: asyncio.Event = asyncio.Event()
         self._step_flags: Dict[str, asyncio.Event] = {}
         self._step_event_holding: Optional[Event] = None
         self._step_lock: asyncio.Lock = asyncio.Lock()

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

@@ -12,3 +12,7 @@ class WorkflowRuntimeError(Exception):
 
 class WorkflowDone(Exception):
     pass
+
+
+class WorkflowCancelledByUser(Exception):
+    pass

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

@@ -104,3 +104,9 @@ async def run_step(self) -> Optional[Event]:
             raise ValueError("Context is not set!")
 
         return retval
+
+    async def cancel_run(self) -> None:
+        """Method to cancel a Workflow execution."""
+        if self.ctx:
+            self.ctx._cancel_flag.set()
+            await asyncio.sleep(0)

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

@@ -8,7 +8,13 @@
 
 from .decorators import StepConfig, step
 from .context import Context
-from .events import InputRequiredEvent, HumanResponseEvent, Event, StartEvent, StopEvent
+from .events import (
+    InputRequiredEvent,
+    HumanResponseEvent,
+    Event,
+    StartEvent,
+    StopEvent,
+)
 from .errors import *
 from .service import ServiceManager
 from .utils import (
@@ -162,6 +168,7 @@ def _start(self, stepwise: bool = False, ctx: Optional[Context] = None) -> Conte
             ctx._step_flags = {}
             ctx._retval = None
             ctx._step_event_holding = None
+            ctx._cancel_flag.clear()
 
         for name, step_func in self._get_steps().items():
             ctx._queues[name] = asyncio.Queue()
@@ -275,6 +282,17 @@ async def _task(
                     )
                 )
 
+            # add dedicated cancel task
+            async def _cancel_workflow_task() -> None:
+                await ctx._cancel_flag.wait()
+                raise WorkflowCancelledByUser
+
+            ctx._tasks.add(
+                asyncio.create_task(
+                    _cancel_workflow_task(), name="cancel_workflow_task"
+                )
+            )
+
         return ctx
 
     def send_event(self, message: Event, step: Optional[str] = None) -> None:

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

@@ -19,6 +19,7 @@
     WorkflowTimeoutError,
     WorkflowValidationError,
     WorkflowRuntimeError,
+    WorkflowCancelledByUser,
 )
 
 from .conftest import AnotherTestEvent, LastEvent, OneTestEvent
@@ -64,6 +65,21 @@ async def test_workflow_run_step(workflow):
     assert result == "Workflow completed"
 
 
+@pytest.mark.asyncio()
+async def test_workflow_cancelled_by_user(workflow):
+    handler = workflow.run(stepwise=True)
+
+    event = await handler.run_step()
+    assert isinstance(event, OneTestEvent)
+    assert not handler.is_done()
+    handler.ctx.send_event(event)
+
+    await handler.cancel_run()
+    await asyncio.sleep(0.1)  # let workflow get cancelled
+    assert handler.is_done()
+    assert type(handler.exception()) == WorkflowCancelledByUser
+
+
 @pytest.mark.asyncio()
 async def test_workflow_run_step_continue_context():
     class DummyWorkflow(Workflow):