feat: add chat bridge protocol

Author: GabrielVasilescu04

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-runtime"
-version = "0.2.0"
+version = "0.2.1"
 description = "Runtime abstractions and interfaces for building agents and automation scripts in the UiPath ecosystem"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/runtime/__init__.py

@@ -7,6 +7,8 @@
     UiPathStreamNotSupportedError,
     UiPathStreamOptions,
 )
+from uipath.runtime.chat.bridge import UiPathChatBridgeProtocol
+from uipath.runtime.chat.runtime import UiPathChatRuntime
 from uipath.runtime.context import UiPathRuntimeContext
 from uipath.runtime.debug.breakpoint import UiPathBreakpointResult
 from uipath.runtime.debug.bridge import UiPathDebugBridgeProtocol
@@ -66,4 +68,6 @@
     "UiPathBreakpointResult",
     "UiPathStreamNotSupportedError",
     "UiPathResumeTriggerName",
+    "UiPathChatBridgeProtocol",
+    "UiPathChatRuntime",
 ]

src/uipath/runtime/chat/__init__.py

@@ -0,0 +1,6 @@
+"""Chat bridge protocol and runtime for conversational agents."""
+
+from uipath.runtime.chat.bridge import UiPathChatBridgeProtocol
+from uipath.runtime.chat.runtime import UiPathChatRuntime
+
+__all__ = ["UiPathChatBridgeProtocol", "UiPathChatRuntime"]

src/uipath/runtime/chat/bridge.py

@@ -0,0 +1,28 @@
+"""Abstract conversation bridge interface."""
+
+from typing import Any, Protocol
+
+from uipath.core.chat import UiPathConversationMessageEvent
+
+
+class UiPathChatBridgeProtocol(Protocol):
+    """Abstract interface for chat communication.
+
+    Implementations: WebSocket, etc.
+    """
+
+    async def connect(self) -> None:
+        """Establish connection to chat service."""
+        ...
+
+    async def disconnect(self) -> None:
+        """Close connection and send exchange end event."""
+        ...
+
+    async def emit_message_event(self, message_event: UiPathConversationMessageEvent) -> None:
+        """Wrap and send a message event.
+
+        Args:
+            message_event: UiPathConversationMessageEvent to wrap and send
+        """
+        ...

src/uipath/runtime/chat/runtime.py

@@ -0,0 +1,86 @@
+"""Chat runtime implementation."""
+
+import logging
+from typing import Any, AsyncGenerator, cast
+
+from uipath.runtime.base import (
+    UiPathExecuteOptions,
+    UiPathRuntimeProtocol,
+    UiPathStreamOptions,
+)
+from uipath.runtime.chat.bridge import UiPathChatBridgeProtocol
+from uipath.runtime.events import (
+    UiPathRuntimeEvent,
+    UiPathRuntimeMessageEvent,
+)
+from uipath.runtime.result import (
+    UiPathRuntimeResult,
+    UiPathRuntimeStatus,
+)
+from uipath.runtime.schema import UiPathRuntimeSchema
+
+logger = logging.getLogger(__name__)
+
+
+class UiPathChatRuntime:
+    """Specialized runtime for chat mode that streams message events to a chat bridge."""
+
+    def __init__(
+        self,
+        delegate: UiPathRuntimeProtocol,
+        chat_bridge: UiPathChatBridgeProtocol,
+    ):
+        """Initialize the UiPathChatRuntime.
+
+        Args:
+            delegate: The underlying runtime to wrap
+            chat_bridge: Bridge for chat event communication
+        """
+        super().__init__()
+        self.delegate = delegate
+        self.chat_bridge = chat_bridge
+
+    async def execute(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathExecuteOptions | None = None,
+    ) -> UiPathRuntimeResult:
+        """Execute the workflow with chat support."""
+        result: UiPathRuntimeResult | None = None
+        async for event in self.stream(input, cast(UiPathStreamOptions, options)):
+            if isinstance(event, UiPathRuntimeResult):
+                result = event
+
+        return (
+            result
+            if result
+            else UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL)
+        )
+
+    async def stream(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathStreamOptions | None = None,
+    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
+        """Stream execution events with chat support."""
+        await self.chat_bridge.connect()
+
+        async for event in self.delegate.stream(input, options=options):
+            if isinstance(event, UiPathRuntimeMessageEvent):
+                if event.payload:
+                    await self.chat_bridge.emit_message_event(event.payload)
+
+            yield event
+
+        await self.chat_bridge.disconnect()
+
+    async def get_schema(self) -> UiPathRuntimeSchema:
+        """Get schema from the delegate runtime."""
+        return await self.delegate.get_schema()
+
+    async def dispose(self) -> None:
+        """Cleanup runtime resources."""
+        try:
+            await self.chat_bridge.disconnect()
+        except Exception as e:
+            logger.warning(f"Error disconnecting chat bridge: {e}")

tests/test_chat_runtime.py

@@ -0,0 +1,220 @@
+"""Tests for UiPathChatRuntime with mocked runtime and chat bridge."""
+
+from __future__ import annotations
+
+from typing import Any, AsyncGenerator, Sequence, cast
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from uipath.core.chat import (
+    UiPathConversationMessageEvent,
+    UiPathConversationMessageStartEvent,
+)
+
+from uipath.runtime import (
+    UiPathExecuteOptions,
+    UiPathRuntimeResult,
+    UiPathRuntimeStatus,
+    UiPathStreamOptions,
+)
+from uipath.runtime.chat import (
+    UiPathChatBridgeProtocol,
+    UiPathChatRuntime,
+)
+from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeMessageEvent
+from uipath.runtime.schema import UiPathRuntimeSchema
+
+
+def make_chat_bridge_mock() -> UiPathChatBridgeProtocol:
+    """Create a chat bridge mock with all methods that UiPathChatRuntime uses.
+    """
+    bridge_mock: Mock = Mock(spec=UiPathChatBridgeProtocol)
+
+    bridge_mock.connect = AsyncMock()
+    bridge_mock.disconnect = AsyncMock()
+    bridge_mock.emit_message_event = AsyncMock()
+
+    return cast(UiPathChatBridgeProtocol, bridge_mock)
+
+
+class StreamingMockRuntime:
+    """Mock runtime that streams message events and a final result."""
+
+    def __init__(
+        self,
+        messages: Sequence[str],
+        *,
+        error_in_stream: bool = False,
+    ) -> None:
+        super().__init__()
+        self.messages: list[str] = list(messages)
+        self.error_in_stream: bool = error_in_stream
+        self.execute_called: bool = False
+
+    async def dispose(self) -> None:
+        pass
+
+    async def execute(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathExecuteOptions | None = None,
+    ) -> UiPathRuntimeResult:
+        """Fallback execute path."""
+        self.execute_called = True
+        return UiPathRuntimeResult(
+            status=UiPathRuntimeStatus.SUCCESSFUL,
+            output={"mode": "execute"},
+        )
+
+    async def stream(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathStreamOptions | None = None,
+    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
+        """Async generator yielding message events and final result."""
+        if self.error_in_stream:
+            raise RuntimeError("Stream blew up")
+
+        for idx, message_text in enumerate(self.messages):
+            message_event = UiPathConversationMessageEvent(
+                message_id=f"msg-{idx}",
+                start=UiPathConversationMessageStartEvent(
+                    role="assistant",
+                    timestamp="2025-01-01T00:00:00.000Z",
+                ),
+            )
+            yield UiPathRuntimeMessageEvent(payload=message_event)
+
+        # Final result at the end of streaming
+        yield UiPathRuntimeResult(
+            status=UiPathRuntimeStatus.SUCCESSFUL,
+            output={"messages": self.messages},
+        )
+
+    async def get_schema(self) -> UiPathRuntimeSchema:
+        raise NotImplementedError()
+
+
+@pytest.mark.asyncio
+async def test_chat_runtime_streams_and_emits_messages():
+    """UiPathChatRuntime should stream events and emit message events to bridge."""
+
+    runtime_impl = StreamingMockRuntime(
+        messages=["Hello", "How are you?", "Goodbye"],
+    )
+    bridge = make_chat_bridge_mock()
+
+    chat_runtime = UiPathChatRuntime(
+        delegate=runtime_impl,
+        chat_bridge=bridge,
+    )
+
+    result = await chat_runtime.execute({})
+
+    # Result propagation
+    assert isinstance(result, UiPathRuntimeResult)
+    assert result.status == UiPathRuntimeStatus.SUCCESSFUL
+    assert result.output == {"messages": ["Hello", "How are you?", "Goodbye"]}
+
+    # Bridge lifecycle
+    cast(AsyncMock, bridge.connect).assert_awaited_once()
+    cast(AsyncMock, bridge.disconnect).assert_awaited_once()
+
+    assert cast(AsyncMock, bridge.emit_message_event).await_count == 3
+
+    # Verify message events were passed as UiPathConversationMessageEvent objects
+    calls = cast(AsyncMock, bridge.emit_message_event).await_args_list
+    assert isinstance(calls[0][0][0], UiPathConversationMessageEvent)
+    assert calls[0][0][0].message_id == "msg-0"
+    assert isinstance(calls[1][0][0], UiPathConversationMessageEvent)
+    assert calls[1][0][0].message_id == "msg-1"
+    assert isinstance(calls[2][0][0], UiPathConversationMessageEvent)
+    assert calls[2][0][0].message_id == "msg-2"
+
+
+@pytest.mark.asyncio
+async def test_chat_runtime_stream_yields_all_events():
+    """UiPathChatRuntime.stream() should yield all events from delegate."""
+
+    runtime_impl = StreamingMockRuntime(
+        messages=["Message 1", "Message 2"],
+    )
+    bridge = make_chat_bridge_mock()
+
+    chat_runtime = UiPathChatRuntime(
+        delegate=runtime_impl,
+        chat_bridge=bridge,
+    )
+
+    events = []
+    async for event in chat_runtime.stream({}):
+        events.append(event)
+
+    # Should have 2 message events + 1 final result
+    assert len(events) == 3
+    assert isinstance(events[0], UiPathRuntimeMessageEvent)
+    assert isinstance(events[1], UiPathRuntimeMessageEvent)
+    assert isinstance(events[2], UiPathRuntimeResult)
+
+    # Bridge methods called
+    cast(AsyncMock, bridge.connect).assert_awaited_once()
+    cast(AsyncMock, bridge.disconnect).assert_awaited_once()
+    assert cast(AsyncMock, bridge.emit_message_event).await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_chat_runtime_handles_errors():
+    """On unexpected errors, UiPathChatRuntime should propagate them."""
+
+    runtime_impl = StreamingMockRuntime(
+        messages=["Message"],
+        error_in_stream=True,
+    )
+    bridge = make_chat_bridge_mock()
+
+    chat_runtime = UiPathChatRuntime(
+        delegate=runtime_impl,
+        chat_bridge=bridge,
+    )
+
+    # Error should propagate
+    with pytest.raises(RuntimeError, match="Stream blew up"):
+        await chat_runtime.execute({})
+
+    cast(AsyncMock, bridge.connect).assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_chat_runtime_dispose_calls_disconnect():
+    """dispose() should call chat bridge disconnect."""
+
+    runtime_impl = StreamingMockRuntime(messages=["Message"])
+    bridge = make_chat_bridge_mock()
+
+    chat_runtime = UiPathChatRuntime(
+        delegate=runtime_impl,
+        chat_bridge=bridge,
+    )
+
+    await chat_runtime.dispose()
+
+    # Bridge disconnect should be called
+    cast(AsyncMock, bridge.disconnect).assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_chat_runtime_dispose_suppresses_disconnect_errors():
+    """Errors from chat_bridge.disconnect should be suppressed."""
+
+    runtime_impl = StreamingMockRuntime(messages=["Message"])
+    bridge = make_chat_bridge_mock()
+    cast(AsyncMock, bridge.disconnect).side_effect = RuntimeError("disconnect failed")
+
+    chat_runtime = UiPathChatRuntime(
+        delegate=runtime_impl,
+        chat_bridge=bridge,
+    )
+
+    await chat_runtime.dispose()
+
+    cast(AsyncMock, bridge.disconnect).assert_awaited_once()