feat: workspace persistence across conversational exchanges via content-parts

pyproject.toml

@@ -1,6 +1,6 @@
 [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"

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,7 @@
 """Workspace persistence primitives for runtime implementations."""
 
+from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id
+from uipath.runtime.workspace.conversational import ConversationalWorkspaceRuntime
 from uipath.runtime.workspace.hydration import (
     HydrationPolicy,
     HydrationRuntime,
@@ -13,9 +15,12 @@
 
 __all__ = [
     "AttachmentRegistryEntry",
+    "ConversationalWorkspaceRuntime",
     "HydrationPolicy",
     "HydrationRuntime",
     "Workspace",
     "WorkspaceHydrator",
     "WorkspaceRegistryStore",
+    "build_cas_uri",
+    "parse_attachment_id",
 ]

src/uipath/runtime/workspace/cas_uri.py

@@ -0,0 +1,25 @@
+"""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 | None) -> str | None:
+    """Return the normalized attachment ID from an Orchestrator CAS URI."""
+    if not uri or not uri.startswith(CAS_FILE_URI_PREFIX):
+        return None
+
+    attachment_id = uri.removeprefix(CAS_FILE_URI_PREFIX)
+    if not attachment_id or ":" in attachment_id:
+        return None
+
+    try:
+        return str(uuid.UUID(attachment_id))
+    except (ValueError, AttributeError):
+        return None

src/uipath/runtime/workspace/conversational.py

@@ -0,0 +1,271 @@
+"""Attachment-backed workspace persistence for conversational agents."""
+
+import logging
+import mimetypes
+from pathlib import Path
+from typing import Any, AsyncGenerator
+
+from pydantic import ValidationError
+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
+from uipath.runtime.workspace.workspace import Workspace
+
+logger = logging.getLogger(__name__)
+
+WORKSPACE_FILE_KIND = "workspace"
+DEFAULT_MIME_TYPE = "application/octet-stream"
+
+# stdlib mimetypes lacks these until newer Python versions
+_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
+
+
+class ConversationalWorkspaceRuntime:
+    """Persists a workspace in marked file parts on assistant messages."""
+
+    def __init__(
+        self,
+        delegate: UiPathRuntimeProtocol,
+        *,
+        hydrator: WorkspaceHydrator,
+        workspace: Workspace,
+        runtime_id: str | None = None,
+        registry_store: WorkspaceRegistryStore | None = None,
+    ):
+        """Initialize the wrapper with the hydrator and workspace to persist."""
+        self.delegate = delegate
+        self.hydrator = hydrator
+        self.workspace = workspace
+        self.registry_store = registry_store
+        if runtime_id is not None:
+            self.runtime_id = runtime_id
+        self._registry: dict[str, dict[str, Any]] = {}
+        self._hydrated = False
+        self._assistant_message_ids: set[str] = set()
+
+    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)
+
+        # Workspace parts must precede the final assistant end event.
+        pending_end: UiPathRuntimeMessageEvent | None = None
+        final_result: UiPathRuntimeResult | None = None
+
+        async for event in self.delegate.stream(input, options=options):
+            if isinstance(event, UiPathRuntimeResult):
+                final_result = event
+                continue
+
+            if isinstance(event, UiPathRuntimeMessageEvent) and isinstance(
+                event.payload, UiPathConversationMessageEvent
+            ):
+                payload = event.payload
+                if pending_end is not None:
+                    yield pending_end
+                    pending_end = None
+                if payload.start and payload.start.role == "assistant":
+                    self._assistant_message_ids.add(payload.message_id)
+                if payload.end and payload.message_id in self._assistant_message_ids:
+                    pending_end = event
+                    continue
+
+            yield event
+
+        if final_result is None:
+            if pending_end is not None:
+                yield pending_end
+            return
+
+        if final_result.status == UiPathRuntimeStatus.SUCCESSFUL:
+            if pending_end is None:
+                raise RuntimeError(
+                    "Conversational workspace requires a completed assistant message"
+                )
+            for file_event in await self._dehydrate(pending_end.payload.message_id):
+                yield file_event
+
+        if pending_end is not None:
+            yield pending_end
+        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:
+        """Dispose delegate and workspace."""
+        try:
+            await self.delegate.dispose()
+        finally:
+            await self.workspace.dispose()
+
+    async def _hydrate(self, input: dict[str, Any] | None) -> None:
+        """Rebuild the workspace from the last assistant message's file parts."""
+        if self._hydrated:
+            return
+
+        persisted_registry = (
+            await self.registry_store.load() if self.registry_store else {}
+        )
+        files = self._files_from_history(input) if not persisted_registry else {}
+        logger.info(
+            "Conversational workspace hydrate: %d file part(s) from history",
+            len(files),
+        )
+        if persisted_registry:
+            self._registry = await self.hydrator.hydrate(persisted_registry)
+        elif files:
+            self._registry = await self.hydrator.hydrate_files(files)
+        self._hydrated = True
+
+    @staticmethod
+    def _files_from_history(input: dict[str, Any] | None) -> dict[str, str]:
+        """Collect workspace files from the last assistant message's file parts."""
+        messages = (input or {}).get("messages")
+        if not isinstance(messages, list):
+            return {}
+
+        last_assistant = ConversationalWorkspaceRuntime._last_assistant_message(
+            messages
+        )
+        if last_assistant is None:
+            return {}
+
+        files: dict[str, str] = {}
+        for part in last_assistant.content_parts:
+            if not isinstance(part.data, UiPathExternalValue):
+                continue
+            metadata = getattr(part, "metadata", None)
+            has_workspace_metadata = (
+                isinstance(metadata, dict)
+                and metadata.get("fileKind") == WORKSPACE_FILE_KIND
+            )
+            is_workspace_file = (
+                has_workspace_metadata or 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
+            files[part.name] = attachment_id
+        return files
+
+    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)
+        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
+
+    @staticmethod
+    def _last_assistant_message(
+        messages: list[Any],
+    ) -> UiPathConversationMessage | None:
+        last: UiPathConversationMessage | None = None
+        for message in messages:
+            if isinstance(message, UiPathConversationMessage):
+                parsed = message
+            elif isinstance(message, dict):
+                try:
+                    parsed = UiPathConversationMessage.model_validate(message)
+                except ValidationError:
+                    continue
+            else:
+                continue
+            if parsed.role == "assistant":
+                last = parsed
+        return last

src/uipath/runtime/workspace/hydrator.py

@@ -68,6 +68,39 @@ async def hydrate(
             )
         return self._dump_registry(normalized)
 
+    async def hydrate_files(
+        self,
+        files: dict[str, str],
+    ) -> dict[str, dict[str, Any]]:
+        """Download attachments into the workspace and build a registry from them.
+
+        Args:
+            files: Mapping of virtual path to attachment key.
+
+        Unlike ``hydrate``, no prior registry (and thus no SHA-256) is known, so
+        every file is downloaded and its digest/size are computed from the
+        downloaded bytes. The returned registry can be fed to ``dehydrate`` so
+        unchanged files are relinked instead of re-uploaded.
+        """
+        registry: dict[str, AttachmentRegistryEntry] = {}
+        for virtual_path, attachment_key in files.items():
+            target = self._resolve_workspace_path(virtual_path)
+            target.parent.mkdir(parents=True, exist_ok=True)
+            await self.attachments.download_async(
+                key=UUID(attachment_key),
+                destination_path=str(target),
+                folder_key=self.folder_key,
+                folder_path=self.folder_path,
+            )
+            registry[virtual_path] = AttachmentRegistryEntry(
+                attachment_key=attachment_key,
+                sha256=self._sha256(target),
+                size=target.stat().st_size,
+                uploaded_at=datetime.now(timezone.utc).isoformat(),
+                attachment_name=self._attachment_name_for_virtual_path(virtual_path),
+            )
+        return self._dump_registry(registry)
+
     async def dehydrate(
         self,
         registry: dict[str, dict[str, Any]],

tests/workspace/test_cas_uri.py

@@ -0,0 +1,16 @@
+import uuid
+
+from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id
+
+
+def test_parse_attachment_id_accepts_orchestrator_cas_uri() -> None:
+    attachment_id = uuid.uuid4()
+
+    assert parse_attachment_id(build_cas_uri(str(attachment_id))) == str(attachment_id)
+
+
+def test_parse_attachment_id_rejects_other_uri_namespaces() -> None:
+    attachment_id = uuid.uuid4()
+
+    assert parse_attachment_id(f"urn:other:{attachment_id}") is None
+    assert parse_attachment_id(f"https://example.test/{attachment_id}") is None