feat: add chat bridge implementation and support for the runtime

GabrielVasilescu04 · GabrielVasilescu04 · commit 4a69715fb629 · 2025-12-04T17:15:37.000+02:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.2.11"
+version = "2.2.12"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
diff --git a/src/uipath/_cli/_chat/_bridge.py b/src/uipath/_cli/_chat/_bridge.py
@@ -0,0 +1,272 @@
+"""Chat bridge implementations for conversational agents."""
+
+import asyncio
+import logging
+import os
+from typing import Any
+from urllib.parse import urlparse
+
+import socketio
+from socketio import AsyncClient
+from uipath.core.chat import (
+    UiPathConversationEvent,
+    UiPathConversationExchangeEndEvent,
+    UiPathConversationExchangeEvent,
+)
+from uipath.runtime.chat import UiPathChatBridgeProtocol
+from uipath.runtime.context import UiPathRuntimeContext
+
+logger = logging.getLogger(__name__)
+
+
+class WebSocketChatBridge:
+    """WebSocket-based chat bridge for streaming conversational events to CAS.
+
+    Implements UiPathChatBridgeProtocol using python-socketio library.
+    """
+
+    def __init__(
+        self,
+        websocket_url: str,
+        conversation_id: str,
+        exchange_id: str,
+        headers: dict[str, str],
+        auth: dict[str, Any] | None = None,
+    ):
+        """Initialize the WebSocket chat bridge.
+
+        Args:
+            websocket_url: The WebSocket server URL to connect to
+            conversation_id: The conversation ID for this session
+            exchange_id: The exchange ID for this session
+            headers: HTTP headers to send during connection
+            auth: Optional authentication data to send during connection
+        """
+        self.websocket_url = websocket_url
+        self.conversation_id = conversation_id
+        self.exchange_id = exchange_id
+        self.auth = auth
+        self.headers = headers
+        self._client: AsyncClient | None = None
+        self._connected_event = asyncio.Event()
+
+    async def connect(self, timeout: float = 10.0) -> None:
+        """Establish WebSocket connection to the server.
+
+        Args:
+            timeout: Connection timeout in seconds (default: 10.0)
+
+        Raises:
+            RuntimeError: If connection fails or times out
+
+        Example:
+            ```python
+            manager = WebSocketManager("http://localhost:3000")
+            await manager.connect()
+            ```
+        """
+        if self._client is not None:
+            logger.warning("WebSocket client already connected")
+            return
+
+        # Create new SocketIO client
+        self._client = socketio.AsyncClient(
+            logger=logger,
+            engineio_logger=logger,
+        )
+
+        # Register connection event handlers
+        self._client.on("connect", self._handle_connect)
+        self._client.on("disconnect", self._handle_disconnect)
+        self._client.on("connect_error", self._handle_connect_error)
+
+        # Clear connection event
+        self._connected_event.clear()
+
+        try:
+            # Attempt to connect with timeout
+            logger.info(f"Connecting to WebSocket server: {self.websocket_url}")
+
+            await asyncio.wait_for(
+                self._client.connect(
+                    url=self.websocket_url,
+                    headers=self.headers,
+                    auth=self.auth,
+                    transports=["websocket"],
+                ),
+                timeout=timeout,
+            )
+
+            # Wait for connection confirmation
+            await asyncio.wait_for(self._connected_event.wait(), timeout=timeout)
+
+            logger.info("WebSocket connection established successfully")
+
+        except asyncio.TimeoutError as e:
+            error_message = f"Failed to connect to WebSocket server within {timeout}s timeout"
+            logger.error(error_message)
+            await self._cleanup_client()
+            raise RuntimeError(error_message) from e
+
+        except Exception as e:
+            error_message = f"Failed to connect to WebSocket server: {e}"
+            logger.error(error_message)
+            await self._cleanup_client()
+            raise RuntimeError(error_message) from e
+
+    async def disconnect(self) -> None:
+        """Close the WebSocket connection gracefully.
+
+        Sends an exchange end event before disconnecting to signal that the
+        exchange is complete. Uses stored conversation/exchange IDs.
+        """
+        if self._client is None:
+            logger.warning("WebSocket client not connected")
+            return
+
+        # Send exchange end event using stored IDs
+        if self._client and self._connected_event.is_set():
+            try:
+                end_event = UiPathConversationEvent(
+                    conversation_id=self.conversation_id,
+                    exchange=UiPathConversationExchangeEvent(
+                        exchange_id=self.exchange_id,
+                        end=UiPathConversationExchangeEndEvent(),
+                    ),
+                )
+                event_data = end_event.model_dump(
+                    mode="json", exclude_none=True, by_alias=True
+                )
+                await self._client.emit("ConversationEvent", event_data)
+                logger.info("Exchange end event sent")
+            except Exception as e:
+                logger.warning(f"Error sending exchange end event: {e}")
+
+        try:
+            logger.info("Disconnecting from WebSocket server")
+            await self._client.disconnect()
+            logger.info("WebSocket disconnected successfully")
+        except Exception as e:
+            logger.error(f"Error during WebSocket disconnect: {e}")
+        finally:
+            await self._cleanup_client()
+
+    async def emit_message_event(self, message_event: Any) -> None:
+        """Wrap and send a message event to the WebSocket server.
+
+        Args:
+            message_event: UiPathConversationMessageEvent to wrap and send
+
+        Raises:
+            RuntimeError: If client is not connected
+        """
+        if self._client is None:
+            raise RuntimeError("WebSocket client not connected. Call connect() first.")
+
+        if not self._connected_event.is_set():
+            raise RuntimeError("WebSocket client not in connected state")
+
+        try:
+            # Wrap message event with conversation/exchange IDs
+            wrapped_event = UiPathConversationEvent(
+                conversation_id=self.conversation_id,
+                exchange=UiPathConversationExchangeEvent(
+                    exchange_id=self.exchange_id,
+                    message=message_event,
+                ),
+            )
+
+            event_data = wrapped_event.model_dump(
+                mode="json", exclude_none=True, by_alias=True
+            )
+
+            logger.debug("Sending conversation event to WebSocket")
+            await self._client.emit("ConversationEvent", event_data)
+            logger.debug("Conversation event sent successfully")
+
+        except Exception as e:
+            logger.error(f"Error sending conversation event to WebSocket: {e}")
+            raise RuntimeError(f"Failed to send conversation event: {e}") from e
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the WebSocket is currently connected.
+
+        Returns:
+            True if connected, False otherwise
+        """
+        return self._client is not None and self._connected_event.is_set()
+
+    async def _handle_connect(self) -> None:
+        """Handle successful connection event."""
+        logger.info("WebSocket connection established")
+        self._connected_event.set()
+
+    async def _handle_disconnect(self) -> None:
+        """Handle disconnection event."""
+        logger.info("WebSocket connection closed")
+        self._connected_event.clear()
+
+    async def _handle_connect_error(self, data: Any) -> None:
+        """Handle connection error event."""
+        logger.error(f"WebSocket connection error: {data}")
+
+    async def _cleanup_client(self) -> None:
+        """Clean up client resources."""
+        self._connected_event.clear()
+        self._client = None
+
+
+def get_chat_bridge(
+    context: UiPathRuntimeContext,
+    conversation_id: str,
+    exchange_id: str,
+) -> UiPathChatBridgeProtocol:
+    """Factory to get WebSocket chat bridge for conversational agents.
+
+    Args:
+        context: The runtime context containing environment configuration
+        conversation_id: The conversation ID for this session
+        exchange_id: The exchange ID for this session
+
+    Returns:
+        WebSocketChatBridge instance configured for CAS
+
+    Raises:
+        RuntimeError: If UIPATH_URL is not set or invalid
+
+    Example:
+        ```python
+        bridge = get_chat_bridge(context, "conv-123", "exch-456")
+        await bridge.connect()
+        await bridge.emit_message_event(message_event)
+        await bridge.disconnect(conversation_id, exchange_id)
+        ```
+    """
+    # Extract host from UIPATH_URL
+    base_url = os.environ.get("UIPATH_URL")
+    if not base_url:
+        raise RuntimeError("UIPATH_URL environment variable required for conversational mode")
+
+    parsed = urlparse(base_url)
+    if not parsed.netloc:
+        raise RuntimeError(f"Invalid UIPATH_URL format: {base_url}")
+
+    host = parsed.netloc
+
+    # Construct WebSocket URL for CAS
+    websocket_url = f"wss://{host}/autopilotforeveryone_/websocket_/socket.io?conversationId={conversation_id}"
+
+    # Build headers from context
+    headers = {
+        "Authorization": f"Bearer {os.environ.get('UIPATH_ACCESS_TOKEN', '')}",
+        "X-UiPath-Internal-TenantId": context.tenant_id or os.environ.get("UIPATH_TENANT_ID", ""),
+        "X-UiPath-Internal-AccountId": context.org_id or os.environ.get("UIPATH_ORGANIZATION_ID", ""),
+    }
+
+    return WebSocketChatBridge(
+        websocket_url=websocket_url,
+        conversation_id=conversation_id,
+        exchange_id=exchange_id,
+        headers=headers,
+    )
diff --git a/src/uipath/_cli/cli_run.py b/src/uipath/_cli/cli_run.py
@@ -3,6 +3,7 @@
 import click
 from uipath.core.tracing import UiPathTraceManager
 from uipath.runtime import (
+    UiPathChatRuntime,
     UiPathExecuteOptions,
     UiPathRuntimeFactoryProtocol,
     UiPathRuntimeFactoryRegistry,
@@ -13,11 +14,13 @@
 from uipath.runtime.context import UiPathRuntimeContext
 from uipath.runtime.debug import UiPathDebugBridgeProtocol
 from uipath.runtime.errors import UiPathRuntimeError
-from uipath.runtime.events import UiPathRuntimeStateEvent
+from uipath.runtime.events import UiPathRuntimeMessageEvent, UiPathRuntimeStateEvent
 
+from uipath._cli._chat._bridge import get_chat_bridge
 from uipath._cli._debug._bridge import ConsoleDebugBridge
 from uipath._cli._utils._common import read_resource_overwrites_from_file
 from uipath._cli._utils._debug import setup_debugging
+from uipath._cli.models.uipath_json_schema import UiPathJsonConfig
 from uipath._utils._bindings import ResourceOverwritesContext
 from uipath.tracing import JsonLinesFileExporter, LlmOpsHttpExporter
 
@@ -132,6 +135,22 @@ async def debug_runtime(
                         await debug_bridge.emit_state_update(event)
                 return ctx.result
 
+            async def conversational_runtime(
+                ctx: UiPathRuntimeContext,
+                runtime: UiPathRuntimeProtocol,
+                conversation_id: str,
+                exchange_id: str,
+            ) -> UiPathRuntimeResult | None:
+                """Execute runtime in chat mode with WebSocket streaming."""
+                chat_bridge = get_chat_bridge(ctx, conversation_id, exchange_id)
+                chat_runtime = UiPathChatRuntime(delegate=runtime, chat_bridge=chat_bridge)
+
+                options = UiPathExecuteOptions(resume=resume)
+                ctx.result = await chat_runtime.execute(
+                    input=ctx.get_input(), options=options
+                )
+                return ctx.result
+
             async def execute() -> None:
                 trace_manager = UiPathTraceManager()
 
@@ -151,6 +170,21 @@ async def execute() -> None:
                         JsonLinesFileExporter(ctx.trace_file)
                     )
 
+                # Check if this is a conversational agent
+                config = UiPathJsonConfig.load_from_file("uipath.json")
+                is_conversational = config.runtime_options.is_conversational
+
+                # Extract conversation_id and exchange_id
+                if is_conversational:
+                    input_data = ctx.get_input()
+                    conversation_id = input_data.get("conversationId")
+                    exchange_id = input_data.get("exchangeId")
+
+                    if not conversation_id or not exchange_id:
+                        raise RuntimeError(
+                            "Conversational agents require 'conversationId' and 'exchangeId' in input"
+                        )
+
                 async with ResourceOverwritesContext(
                     lambda: read_resource_overwrites_from_file(ctx.runtime_dir)
                 ):
@@ -162,7 +196,12 @@ async def execute() -> None:
                             runtime = await factory.new_runtime(
                                 entrypoint, ctx.job_id or "default"
                             )
-                            if ctx.job_id:
+
+                            if is_conversational:
+                                ctx.result = await conversational_runtime(
+                                    ctx, runtime, conversation_id, exchange_id
+                                )
+                            elif ctx.job_id:
                                 trace_manager.add_span_exporter(LlmOpsHttpExporter())
                                 ctx.result = await execute_runtime(ctx, runtime)
                             else:
diff --git a/uv.lock b/uv.lock
