Repository: run-llama/llama_deploy

Issue:
Add `send_event` method to `SessionClient` (to enable a user to do the equivalent of `ctx.send_event(ev)`)



Discussion:

Patch:
diff --git a/llama_deploy/client/async_client.py b/llama_deploy/client/async_client.py
index cd521d11..bf2c2d3d 100644
--- a/llama_deploy/client/async_client.py
+++ b/llama_deploy/client/async_client.py
@@ -7,11 +7,14 @@
 
 from llama_deploy.control_plane.server import ControlPlaneConfig
 from llama_deploy.types import (
+    EventDefinition,
     ServiceDefinition,
     SessionDefinition,
     TaskDefinition,
     TaskResult,
 )
+from llama_index.core.workflow import Event
+from llama_index.core.workflow.context_serializers import JsonSerializer
 
 DEFAULT_TIMEOUT = 120.0
 DEFAULT_POLL_INTERVAL = 0.5
@@ -149,6 +152,26 @@ async def get_task_result_stream(
                         f"Task result not available after waiting for {self.timeout} seconds"
                     )
 
+    async def send_event(self, service_name: str, task_id: str, ev: Event) -> None:
+        """Send event to a Workflow service.
+
+        Args:
+            event (Event): The event to be submitted to the workflow.
+
+        Returns:
+            None
+        """
+        serializer = JsonSerializer()
+        event_def = EventDefinition(
+            event_obj_str=serializer.serialize(ev), agent_id=service_name
+        )
+
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            await client.post(
+                f"{self.control_plane_url}/sessions/{self.session_id}/tasks/{task_id}/send_event",
+                json=event_def.model_dump(),
+            )
+
 
 class AsyncLlamaDeployClient:
     def __init__(
diff --git a/llama_deploy/client/sync_client.py b/llama_deploy/client/sync_client.py
index 76d17395..741ef51c 100644
--- a/llama_deploy/client/sync_client.py
+++ b/llama_deploy/client/sync_client.py
@@ -5,11 +5,15 @@
 
 from llama_deploy.control_plane.server import ControlPlaneConfig
 from llama_deploy.types import (
+    EventDefinition,
     TaskDefinition,
     ServiceDefinition,
     TaskResult,
     SessionDefinition,
 )
+from llama_index.core.workflow import Event
+from llama_index.core.workflow.context_serializers import JsonSerializer
+
 
 DEFAULT_TIMEOUT = 120.0
 DEFAULT_POLL_INTERVAL = 0.5
@@ -151,6 +155,26 @@ def get_task_result_stream(
                             f"Task result not available after waiting for {self.timeout} seconds"
                         )
 
+    def send_event(self, service_name: str, task_id: str, ev: Event) -> None:
+        """Send event to a Workflow service.
+
+        Args:
+            event (Event): The event to be submitted to the workflow.
+
+        Returns:
+            None
+        """
+        serializer = JsonSerializer()
+        event_def = EventDefinition(
+            event_obj_str=serializer.serialize(ev), agent_id=service_name
+        )
+
+        with httpx.Client(timeout=self.timeout) as client:
+            client.post(
+                f"{self.control_plane_url}/sessions/{self.session_id}/tasks/{task_id}/send_event",
+                json=event_def.model_dump(),
+            )
+
 
 class LlamaDeployClient:
     def __init__(
diff --git a/llama_deploy/control_plane/server.py b/llama_deploy/control_plane/server.py
index 7ecec4ec..ba7cd85e 100644
--- a/llama_deploy/control_plane/server.py
+++ b/llama_deploy/control_plane/server.py
@@ -25,6 +25,7 @@
 from llama_deploy.orchestrators.utils import get_result_key, get_stream_key
 from llama_deploy.types import (
     ActionTypes,
+    EventDefinition,
     ServiceDefinition,
     SessionDefinition,
     TaskDefinition,
@@ -237,6 +238,12 @@ def __init__(
             methods=["GET"],
             tags=["Sessions"],
         )
+        self.app.add_api_route(
+            "/sessions/{session_id}/tasks/{task_id}/send_event",
+            self.send_event,
+            methods=["POST"],
+            tags=["Sessions"],
+        )
         self.app.add_api_route(
             "/sessions/{session_id}/state",
             self.get_session_state,
@@ -584,6 +591,25 @@ async def event_generator(
             media_type="application/x-ndjson",
         )
 
+    async def send_event(
+        self,
+        session_id: str,
+        task_id: str,
+        event_def: EventDefinition,
+    ) -> None:
+        task_def = TaskDefinition(
+            task_id=task_id,
+            session_id=session_id,
+            input=event_def.event_obj_str,
+            agent_id=event_def.agent_id,
+        )
+        message = QueueMessage(
+            type=event_def.agent_id,
+            action=ActionTypes.SEND_EVENT,
+            data=task_def.model_dump(),
+        )
+        await self.publish(message)
+
     async def get_session_state(self, session_id: str) -> Dict[str, Any]:
         session = await self.get_session(session_id)
         if session.task_ids is None:
diff --git a/llama_deploy/services/workflow.py b/llama_deploy/services/workflow.py
index 2dfade12..df969d02 100644
--- a/llama_deploy/services/workflow.py
+++ b/llama_deploy/services/workflow.py
@@ -3,6 +3,7 @@
 import os
 import uuid
 import uvicorn
+from collections import defaultdict
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
 from logging import getLogger
@@ -11,7 +12,11 @@
 from typing import AsyncGenerator, Dict, Optional, Any
 
 from llama_index.core.workflow import Context, Workflow
-from llama_index.core.workflow.context_serializers import JsonPickleSerializer
+from llama_index.core.workflow.handler import WorkflowHandler
+from llama_index.core.workflow.context_serializers import (
+    JsonPickleSerializer,
+    JsonSerializer,
+)
 
 from llama_deploy.message_consumers.base import BaseMessageQueueConsumer
 from llama_deploy.message_consumers.callable import CallableMessageConsumer
@@ -127,6 +132,7 @@ class WorkflowService(BaseService):
     _publish_callback: Optional[PublishCallback] = PrivateAttr()
     _lock: asyncio.Lock = PrivateAttr()
     _outstanding_calls: Dict[str, WorkflowState] = PrivateAttr()
+    _events_buffer: Dict[str, asyncio.Queue] = PrivateAttr()
 
     def __init__(
         self,
@@ -165,6 +171,7 @@ def __init__(
 
         self._outstanding_calls: Dict[str, WorkflowState] = {}
         self._ongoing_tasks: Dict[str, asyncio.Task] = {}
+        self._events_buffer: Dict[str, asyncio.Queue] = defaultdict(asyncio.Queue)
 
         self._app = FastAPI(lifespan=self.lifespan)
 
@@ -285,10 +292,29 @@ async def process_call(self, current_call: WorkflowState) -> None:
             # run the workflow
             handler = self.workflow.run(ctx=ctx, **current_call.run_kwargs)
 
+            # create send_event background task
+            close_send_events = asyncio.Event()
+
+            async def send_events(
+                handler: WorkflowHandler, close_event: asyncio.Event
+            ) -> None:
+                if handler.ctx is None:
+                    raise ValueError("handler does not have a valid Context.")
+
+                while not close_event.is_set():
+                    try:
+                        event = self._events_buffer[current_call.task_id].get_nowait()
+                        handler.ctx.send_event(event)
+                    except asyncio.QueueEmpty:
+                        pass
+                    await asyncio.sleep(self.step_interval)
+
+            _ = asyncio.create_task(send_events(handler, close_send_events))
+
             index = 0
             async for ev in handler.stream_events():
+                # send the event to control plane for client / api server streaming
                 logger.debug(f"Publishing event: {ev}")
-
                 await self.message_queue.publish(
                     QueueMessage(
                         type=CONTROL_PLANE_NAME,
@@ -344,6 +370,7 @@ async def process_call(self, current_call: WorkflowState) -> None:
             )
         finally:
             # clean up
+            close_send_events.set()
             async with self.lock:
                 self._outstanding_calls.pop(current_call.task_id, None)
             self._ongoing_tasks.pop(current_call.task_id, None)
@@ -406,6 +433,7 @@ async def process_message(self, message: QueueMessage) -> None:
         """Process a message received from the message queue."""
         if message.action == ActionTypes.NEW_TASK:
             task_def = TaskDefinition(**message.data or {})
+
             run_kwargs = json.loads(task_def.input)
             workflow_state = WorkflowState(
                 session_id=task_def.session_id,
@@ -415,6 +443,14 @@ async def process_message(self, message: QueueMessage) -> None:
 
             async with self.lock:
                 self._outstanding_calls[task_def.task_id] = workflow_state
+        elif message.action == ActionTypes.SEND_EVENT:
+            serializer = JsonSerializer()
+
+            task_def = TaskDefinition(**message.data or {})
+            event = serializer.deserialize(task_def.input)
+            async with self.lock:
+                self._events_buffer[task_def.task_id].put_nowait(event)
+
         else:
             raise ValueError(f"Unhandled action: {message.action}")
 
diff --git a/llama_deploy/types.py b/llama_deploy/types.py
index 1e6c7afa..2ae4923c 100644
--- a/llama_deploy/types.py
+++ b/llama_deploy/types.py
@@ -78,6 +78,7 @@ class ActionTypes(str, Enum):
     NEW_TOOL_CALL = "new_tool_call"
     COMPLETED_TOOL_CALL = "completed_tool_call"
     TASK_STREAM = "task_stream"
+    SEND_EVENT = "send_event"
 
 
 class TaskDefinition(BaseModel):
@@ -127,6 +128,20 @@ def current_task_id(self) -> Optional[str]:
         return self.task_ids[-1]
 
 
+class EventDefinition(BaseModel):
+    """The definition of event.
+
+    To be used as payloads for service endpoints when wanting to send serialized
+    Events.
+
+    Attributes:
+        event_object_str (str): serialized string of event.
+    """
+
+    agent_id: str
+    event_obj_str: str
+
+
 class TaskResult(BaseModel):
     """
     The result of a task.


Test patch:
diff --git a/e2e_tests/basic_hitl/__init__.py b/e2e_tests/basic_hitl/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/e2e_tests/basic_hitl/conftest.py b/e2e_tests/basic_hitl/conftest.py
new file mode 100644
index 00000000..d74efa06
--- /dev/null
+++ b/e2e_tests/basic_hitl/conftest.py
@@ -0,0 +1,55 @@
+import asyncio
+import multiprocessing
+import time
+
+import pytest
+
+from llama_deploy import (
+    ControlPlaneConfig,
+    SimpleMessageQueueConfig,
+    WorkflowServiceConfig,
+    deploy_core,
+    deploy_workflow,
+)
+
+from .workflow import HumanInTheLoopWorkflow
+
+
+def run_async_core():
+    asyncio.run(deploy_core(ControlPlaneConfig(), SimpleMessageQueueConfig()))
+
+
+@pytest.fixture(scope="package")
+def core():
+    p = multiprocessing.Process(target=run_async_core)
+    p.start()
+    time.sleep(5)
+
+    yield
+
+    p.kill()
+
+
+def run_async_workflow():
+    asyncio.run(
+        deploy_workflow(
+            HumanInTheLoopWorkflow(timeout=60),
+            WorkflowServiceConfig(
+                host="127.0.0.1",
+                port=8002,
+                service_name="hitl_workflow",
+            ),
+            ControlPlaneConfig(),
+        )
+    )
+
+
+@pytest.fixture(scope="package")
+def services(core):
+    p = multiprocessing.Process(target=run_async_workflow)
+    p.start()
+    time.sleep(5)
+
+    yield
+
+    p.kill()
diff --git a/e2e_tests/basic_hitl/test_run_client.py b/e2e_tests/basic_hitl/test_run_client.py
new file mode 100644
index 00000000..6244636a
--- /dev/null
+++ b/e2e_tests/basic_hitl/test_run_client.py
@@ -0,0 +1,75 @@
+import asyncio
+import pytest
+import time
+
+from llama_deploy import AsyncLlamaDeployClient, ControlPlaneConfig, LlamaDeployClient
+from llama_index.core.workflow.events import HumanResponseEvent
+
+
+@pytest.mark.e2ehitl
+def test_run_client(services):
+    client = LlamaDeployClient(ControlPlaneConfig(), timeout=10)
+
+    # sanity check
+    sessions = client.list_sessions()
+    assert len(sessions) == 0, "Sessions list is not empty"
+
+    # create a session
+    session = client.create_session()
+
+    # kick off run
+    task_id = session.run_nowait("hitl_workflow")
+
+    # send event
+    session.send_event(
+        ev=HumanResponseEvent(response="42"),
+        service_name="hitl_workflow",
+        task_id=task_id,
+    )
+
+    # get final result, polling to wait for workflow to finish after send event
+    final_result = None
+    while final_result is None:
+        final_result = session.get_task_result(task_id)
+        time.sleep(0.1)
+    assert final_result.result == "42", "The human's response is not consistent."
+
+    # delete the session
+    client.delete_session(session.session_id)
+    sessions = client.list_sessions()
+    assert len(sessions) == 0, "Sessions list is not empty"
+
+
+@pytest.mark.e2ehitl
+@pytest.mark.asyncio
+async def test_run_client_async(services):
+    client = AsyncLlamaDeployClient(ControlPlaneConfig(), timeout=10)
+
+    # sanity check
+    sessions = await client.list_sessions()
+    assert len(sessions) == 0, "Sessions list is not empty"
+
+    # create a session
+    session = await client.create_session()
+
+    # kick off run
+    task_id = await session.run_nowait("hitl_workflow")
+
+    # send event
+    await session.send_event(
+        ev=HumanResponseEvent(response="42"),
+        service_name="hitl_workflow",
+        task_id=task_id,
+    )
+
+    # get final result, polling to wait for workflow to finish after send event
+    final_result = None
+    while final_result is None:
+        final_result = await session.get_task_result(task_id)
+        asyncio.sleep(0.1)
+    assert final_result.result == "42", "The human's response is not consistent."
+
+    # delete the session
+    await client.delete_session(session.session_id)
+    sessions = await client.list_sessions()
+    assert len(sessions) == 0, "Sessions list is not empty"
diff --git a/e2e_tests/basic_hitl/workflow.py b/e2e_tests/basic_hitl/workflow.py
new file mode 100644
index 00000000..b7c93d4a
--- /dev/null
+++ b/e2e_tests/basic_hitl/workflow.py
@@ -0,0 +1,20 @@
+from llama_index.core.workflow import (
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
+from llama_index.core.workflow.events import (
+    HumanResponseEvent,
+    InputRequiredEvent,
+)
+
+
+class HumanInTheLoopWorkflow(Workflow):
+    @step
+    async def step1(self, ev: StartEvent) -> InputRequiredEvent:
+        return InputRequiredEvent(prefix="Enter a number: ")
+
+    @step
+    async def step2(self, ev: HumanResponseEvent) -> StopEvent:
+        return StopEvent(result=ev.response)
diff --git a/tests/services/test_workflow_service.py b/tests/services/test_workflow_service.py
index 01acddb1..cc1535d3 100644
--- a/tests/services/test_workflow_service.py
+++ b/tests/services/test_workflow_service.py
@@ -4,6 +4,8 @@
 from pydantic import PrivateAttr
 from typing import Any, List
 from llama_index.core.workflow import Workflow, StartEvent, StopEvent, step
+from llama_index.core.workflow.events import HumanResponseEvent, InputRequiredEvent
+from llama_index.core.workflow.context_serializers import JsonSerializer
 
 from llama_deploy.messages import QueueMessage
 from llama_deploy.message_consumers import BaseMessageQueueConsumer
@@ -40,6 +42,20 @@ async def run_step(self, ev: StartEvent) -> StopEvent:
     return TestWorklow()
 
 
+@pytest.fixture()
+def test_hitl_workflow() -> Workflow:
+    class TestHumanInTheLoopWorklow(Workflow):
+        @step
+        async def step1(self, ev: StartEvent) -> InputRequiredEvent:
+            return InputRequiredEvent(prefix="Enter a number: ")
+
+        @step
+        async def step2(self, ev: HumanResponseEvent) -> StopEvent:
+            return StopEvent(result=ev.response)
+
+    return TestHumanInTheLoopWorklow()
+
+
 @pytest.mark.asyncio
 async def test_workflow_service(
     test_workflow: Workflow, human_output_consumer: MockMessageConsumer
@@ -83,3 +99,66 @@ async def test_workflow_service(
     result = human_output_consumer.processed_messages[-1]
     assert result.action == ActionTypes.COMPLETED_TASK
     assert result.data["result"] == "test_arg1_result"
+
+
+@pytest.mark.asyncio()
+async def test_hitl_workflow_service(
+    test_hitl_workflow: Workflow,
+    human_output_consumer: MockMessageConsumer,
+) -> None:
+    # arrange
+    message_queue = SimpleMessageQueue()
+    _ = await message_queue.register_consumer(human_output_consumer)
+
+    # create the service
+    workflow_service = WorkflowService(
+        test_hitl_workflow,
+        message_queue,
+        service_name="test_workflow",
+        description="Test Workflow Service",
+        host="localhost",
+        port=8001,
+    )
+
+    # launch it
+    mq_task = await message_queue.launch_local()
+    server_task = await workflow_service.launch_local()
+
+    # process run task
+    task = TaskDefinition(
+        task_id="1",
+        input=json.dumps({}),
+        session_id="test_session_id",
+    )
+
+    await workflow_service.process_message(
+        QueueMessage(
+            action=ActionTypes.NEW_TASK,
+            data=task.model_dump(),
+        )
+    )
+
+    # process human response event task
+    serializer = JsonSerializer()
+    ev = HumanResponseEvent(response="42")
+    task = TaskDefinition(
+        task_id="1",
+        session_id="test_session_id",
+        input=serializer.serialize(ev),
+    )
+    await workflow_service.process_message(
+        QueueMessage(
+            action=ActionTypes.SEND_EVENT,
+            data=task.model_dump(),
+        )
+    )
+
+    # give time to process and shutdown afterwards
+    await asyncio.sleep(1)
+    mq_task.cancel()
+    server_task.cancel()
+
+    # assert
+    result = human_output_consumer.processed_messages[-1]
+    assert result.action == ActionTypes.COMPLETED_TASK
+    assert result.data["result"] == "42"