feat: workspace persistence across conversational exchanges via content-parts

pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "uipath-runtime"
-version = "0.12.3"
+version = "0.12.4"
 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"
 dependencies = [
-  "uipath-core>=0.5.28, <0.6.0",
+  "uipath-core>=0.5.31, <0.6.0",
   "vaderSentiment>=3.3.2, <4.0",
   "chardet>=5.2.0, <8.0",
 ]

src/uipath/runtime/__init__.py

@@ -45,6 +45,7 @@
 from uipath.runtime.storage import UiPathRuntimeStorageProtocol
 from uipath.runtime.workspace import (
     AttachmentRegistryEntry,
+    ConversationalWorkspaceRuntime,
     HydrationPolicy,
     HydrationRuntime,
     Workspace,
@@ -82,6 +83,7 @@
     "UiPathChatProtocol",
     "UiPathChatRuntime",
     "AttachmentRegistryEntry",
+    "ConversationalWorkspaceRuntime",
     "HydrationPolicy",
     "HydrationRuntime",
     "Workspace",

src/uipath/runtime/workspace/__init__.py

@@ -1,5 +1,6 @@
 """Workspace persistence primitives for runtime implementations."""
 
+from uipath.runtime.workspace.conversational import ConversationalWorkspaceRuntime
 from uipath.runtime.workspace.hydration import (
     HydrationPolicy,
     HydrationRuntime,
@@ -13,6 +14,7 @@
 
 __all__ = [
     "AttachmentRegistryEntry",
+    "ConversationalWorkspaceRuntime",
     "HydrationPolicy",
     "HydrationRuntime",
     "Workspace",

src/uipath/runtime/workspace/cas_uri.py

@@ -0,0 +1,22 @@
+"""CAS file URI helpers for workspace attachments."""
+
+import uuid
+
+CAS_FILE_URI_PREFIX = "urn:uipath:cas:file:orchestrator:"
+
+
+def build_cas_uri(attachment_key: str) -> str:
+    """Build a CAS file URI for an Orchestrator attachment key."""
+    return f"{CAS_FILE_URI_PREFIX}{attachment_key}"
+
+
+def parse_attachment_id(uri: str) -> str | None:
+    """Return the normalized attachment ID from an Orchestrator CAS URI."""
+    if not uri.startswith(CAS_FILE_URI_PREFIX):
+        return None
+
+    attachment_id = uri.removeprefix(CAS_FILE_URI_PREFIX)
+    try:
+        return str(uuid.UUID(attachment_id))
+    except ValueError:
+        return None

src/uipath/runtime/workspace/conversational.py

@@ -0,0 +1,296 @@
+"""Attachment-backed workspace persistence for conversational agents."""
+
+import logging
+import mimetypes
+from collections.abc import Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import Any, AsyncGenerator
+
+from uipath.core.chat import (
+    UiPathConversationContentPartEndEvent,
+    UiPathConversationContentPartEvent,
+    UiPathConversationContentPartStartEvent,
+    UiPathConversationMessage,
+    UiPathConversationMessageEvent,
+    UiPathExternalValue,
+)
+
+from uipath.runtime.base import (
+    UiPathExecuteOptions,
+    UiPathRuntimeProtocol,
+    UiPathStreamOptions,
+)
+from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeMessageEvent
+from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus
+from uipath.runtime.schema import UiPathRuntimeSchema
+from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id
+from uipath.runtime.workspace.hydrator import WorkspaceHydrator
+from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore
+
+logger = logging.getLogger(__name__)
+
+WORKSPACE_FILE_KIND = "workspace"
+DEFAULT_MIME_TYPE = "application/octet-stream"
+
+_EXTRA_MIME_TYPES = {
+    ".md": "text/markdown",
+    ".yaml": "application/yaml",
+    ".yml": "application/yaml",
+}
+
+
+def _guess_mime_type(virtual_path: str) -> str:
+    suffix = Path(virtual_path).suffix.lower()
+    if suffix in _EXTRA_MIME_TYPES:
+        return _EXTRA_MIME_TYPES[suffix]
+    return mimetypes.guess_type(virtual_path)[0] or DEFAULT_MIME_TYPE
+
+
+def _last_assistant_message(
+    messages: Sequence[object],
+) -> UiPathConversationMessage | None:
+    for message in reversed(messages):
+        if isinstance(message, UiPathConversationMessage):
+            if message.role == "assistant":
+                return message
+        elif isinstance(message, Mapping) and message.get("role") == "assistant":
+            return UiPathConversationMessage.model_validate(message)
+    return None
+
+
+def _attachment_keys_from_history(
+    input: Mapping[str, object] | None,
+) -> dict[str, str]:
+    """Collect workspace attachment keys from the last assistant message."""
+    messages = (input or {}).get("messages")
+    if not isinstance(messages, list):
+        return {}
+
+    last_assistant = _last_assistant_message(messages)
+    if last_assistant is None:
+        return {}
+
+    attachment_keys_by_path: dict[str, str] = {}
+    for part in last_assistant.content_parts:
+        if not isinstance(part.data, UiPathExternalValue):
+            continue
+        is_workspace_file = (
+            part.metadata.get("fileKind") == WORKSPACE_FILE_KIND
+            if part.metadata is not None
+            else part.content_part_id.startswith("ws-")
+        )
+        if not is_workspace_file:
+            continue
+        attachment_id = parse_attachment_id(part.data.uri)
+        if attachment_id is None or not part.name:
+            continue
+        attachment_keys_by_path[part.name] = attachment_id
+    return attachment_keys_by_path
+
+
+def _conversation_message_from_event(
+    event: UiPathRuntimeEvent,
+) -> tuple[UiPathRuntimeMessageEvent, UiPathConversationMessageEvent] | None:
+    if not isinstance(event, UiPathRuntimeMessageEvent):
+        return None
+    if not isinstance(event.payload, UiPathConversationMessageEvent):
+        return None
+    return event, event.payload
+
+
+class _AssistantMessageEndBuffer:
+    """Delay assistant end events, including end-only events emitted after resume."""
+
+    def __init__(self) -> None:
+        self._open_message_roles: dict[str, str] = {}
+        self._pending_end: UiPathRuntimeMessageEvent | None = None
+
+    def process(self, event: UiPathRuntimeEvent) -> Iterator[UiPathRuntimeEvent]:
+        conversation_event = _conversation_message_from_event(event)
+        if conversation_event is None:
+            yield event
+            return
+        message_event, message = conversation_event
+
+        if self._pending_end is not None:
+            pending_end = self._pending_end
+            self._pending_end = None
+            yield pending_end
+
+        if message.start is not None:
+            self._open_message_roles[message.message_id] = message.start.role
+
+        message_role = self._open_message_roles.get(message.message_id)
+        is_assistant_end = message.end is not None and (
+            message_role is None or message_role == "assistant"
+        )
+        if message.end is not None:
+            self._open_message_roles.pop(message.message_id, None)
+        if is_assistant_end:
+            self._pending_end = message_event
+        else:
+            yield message_event
+
+    def get_pending_end(self) -> UiPathRuntimeMessageEvent | None:
+        pending_end = self._pending_end
+        self._pending_end = None
+        return pending_end
+
+
+class ConversationalWorkspaceRuntime:
+    """Persists a workspace in marked file parts on assistant messages."""
+
+    def __init__(
+        self,
+        delegate: UiPathRuntimeProtocol,
+        *,
+        hydrator: WorkspaceHydrator,
+        registry_store: WorkspaceRegistryStore | None = None,
+    ):
+        """Initialize the wrapper with its delegate and hydrator."""
+        self.delegate = delegate
+        self.hydrator = hydrator
+        self.registry_store = registry_store
+        self._registry: dict[str, dict[str, Any]] = {}
+        self._hydrated = False
+
+    async def execute(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathExecuteOptions | None = None,
+    ) -> UiPathRuntimeResult:
+        """Execute by draining the stream."""
+        result: UiPathRuntimeResult | None = None
+        stream_options = (
+            UiPathStreamOptions.model_validate(options.model_dump())
+            if options is not None
+            else None
+        )
+        async for event in self.stream(input, options=stream_options):
+            if isinstance(event, UiPathRuntimeResult):
+                result = event
+        if result is None:
+            raise RuntimeError("Delegate stream completed without a runtime result")
+        return result
+
+    async def stream(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathStreamOptions | None = None,
+    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
+        """Hydrate, stream delegate events, and emit workspace content-parts."""
+        await self._hydrate(input)
+
+        end_buffer = _AssistantMessageEndBuffer()
+        final_result: UiPathRuntimeResult | None = None
+
+        async for event in self.delegate.stream(input, options=options):
+            if isinstance(event, UiPathRuntimeResult):
+                final_result = event
+                continue
+
+            for event_to_emit in end_buffer.process(event):
+                yield event_to_emit
+
+        pending_end = end_buffer.get_pending_end()
+
+        if (
+            final_result is not None
+            and final_result.status == UiPathRuntimeStatus.SUCCESSFUL
+        ):
+            if pending_end is None:
+                logger.error(
+                    "Skipping workspace persistence because the delegate did not "
+                    "complete an assistant message"
+                )
+            else:
+                for file_event in await self._dehydrate(pending_end.payload.message_id):
+                    yield file_event
+
+        if pending_end is not None:
+            yield pending_end
+        if final_result is not None:
+            yield final_result
+
+    async def get_schema(self) -> UiPathRuntimeSchema:
+        """Passthrough schema from delegate runtime."""
+        return await self.delegate.get_schema()
+
+    async def dispose(self) -> None:
+        """Release resources owned by this wrapper."""
+
+    async def _hydrate(self, input: Mapping[str, object] | None) -> None:
+        """Rebuild the workspace from the last assistant message's file parts."""
+        if self._hydrated:
+            return
+
+        persisted_registry = (
+            await self.registry_store.try_load() if self.registry_store else None
+        )
+        attachment_keys_by_path = (
+            _attachment_keys_from_history(input) if persisted_registry is None else {}
+        )
+        logger.info(
+            "Conversational workspace hydrate: %d file part(s) from history",
+            len(attachment_keys_by_path),
+        )
+        if persisted_registry is not None:
+            self._registry = await self.hydrator.hydrate_from_registry(
+                persisted_registry
+            )
+        elif attachment_keys_by_path:
+            self._registry = await self.hydrator.hydrate_from_attachments(
+                attachment_keys_by_path
+            )
+        self._hydrated = True
+
+    async def _dehydrate(self, message_id: str) -> list[UiPathRuntimeMessageEvent]:
+        """Upload/relink workspace files and build their content-part events."""
+        self._registry = await self.hydrator.dehydrate(self._registry)
+        if self.registry_store is not None:
+            await self.registry_store.save(self._registry)
+        logger.info(
+            "Conversational workspace dehydrate: emitting %d workspace file(s) on message %s",
+            len(self._registry),
+            message_id,
+        )
+
+        events: list[UiPathRuntimeMessageEvent] = []
+        for index, (virtual_path, entry) in enumerate(sorted(self._registry.items())):
+            content_part_id = f"ws-{message_id}-{index}"
+            mime_type = _guess_mime_type(virtual_path)
+            metadata = {
+                "fileKind": WORKSPACE_FILE_KIND,
+                "sha256": entry["sha256"],
+            }
+            events.append(
+                UiPathRuntimeMessageEvent(
+                    payload=UiPathConversationMessageEvent(
+                        message_id=message_id,
+                        content_part=UiPathConversationContentPartEvent(
+                            content_part_id=content_part_id,
+                            start=UiPathConversationContentPartStartEvent(
+                                mime_type=mime_type,
+                                name=virtual_path,
+                                external_value=UiPathExternalValue(
+                                    uri=build_cas_uri(entry["attachment_key"]),
+                                    byte_count=entry["size"],
+                                ),
+                                metadata=metadata,
+                            ),
+                        ),
+                    )
+                )
+            )
+            events.append(
+                UiPathRuntimeMessageEvent(
+                    payload=UiPathConversationMessageEvent(
+                        message_id=message_id,
+                        content_part=UiPathConversationContentPartEvent(
+                            content_part_id=content_part_id,
+                            end=UiPathConversationContentPartEndEvent(),
+                        ),
+                    )
+                )
+            )
+        return events

src/uipath/runtime/workspace/hydration.py

@@ -96,7 +96,7 @@ async def dispose(self) -> None:
 
     async def _hydrate(self) -> None:
         registry = await self.registry_store.load()
-        hydrated = await self.hydrator.hydrate(registry)
+        hydrated = await self.hydrator.hydrate_from_registry(registry)
         if hydrated != registry:
             await self.registry_store.save(hydrated)