feat: add storage and tests

Author: GabrielVasilescu04

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/chat/messages.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from datetime import datetime, timezone
 from typing import Any
@@ -22,23 +23,36 @@
     UiPathConversationToolCallStartEvent,
 )
 
+from uipath_llamaindex.runtime.storage import SqliteResumableStorage
+
 logger = logging.getLogger(__name__)
 
+STORAGE_NAMESPACE_EVENT_MAPPER = "chat-event-mapper"
+STORAGE_KEY_TOOL_ID_TO_MESSAGE_ID_MAP = "tool_id_map"
+
 
 class UiPathChatMessagesMapper:
     """Stateful mapper that converts LlamaIndex agent events to UiPath message events.
 
     Maintains state across events to properly track:
     - The current AI message ID (generated per agent turn, since LlamaIndex doesn't provide one)
     - Pending tool calls per message ID for correct message_end timing
+
+    When a storage backend is provided, the tool_id → message_id mapping is persisted
+    so that it survives workflow suspension and can be correctly resolved on resume.
     """
 
-    def __init__(self, runtime_id: str) -> None:
+    def __init__(
+        self,
+        runtime_id: str,
+        storage: SqliteResumableStorage | None = None,
+    ) -> None:
         self.runtime_id = runtime_id
+        self.storage = storage
         self._current_message_id: str | None = None
-        # message_id -> set of tool_ids still pending completion
+        self._storage_lock = asyncio.Lock()
+        # In-memory fallback state used when no storage is provided
         self._pending_tool_calls: dict[str, set[str]] = {}
-        # tool_id -> message_id for correlating ToolCallResult with its parent AI message
         self._tool_id_to_message_id: dict[str, str] = {}
 
     @staticmethod
@@ -96,14 +110,14 @@ async def map_event(
             return self._map_agent_stream(event)
 
         if isinstance(event, AgentOutput):
-            return self._map_agent_output(event)
+            return await self._map_agent_output(event)
 
         # ToolCall start is handled via AgentOutput to have the message_id available
         if isinstance(event, ToolCall):
             return None
 
         if isinstance(event, ToolCallResult):
-            return self._map_tool_call_result(event)
+            return await self._map_tool_call_result(event)
 
         return None
 
@@ -124,7 +138,7 @@ def _map_agent_stream(
 
         return events if events else None
 
-    def _map_agent_output(
+    async def _map_agent_output(
         self, event: AgentOutput
     ) -> list[UiPathConversationMessageEvent] | None:
         message_id = self._current_message_id
@@ -137,31 +151,58 @@ def _map_agent_output(
         events: list[UiPathConversationMessageEvent] = []
 
         if event.tool_calls:
-            # Emit a tool_call_start event for each tool call and track them as pending
-            pending: set[str] = set()
-            for tool_call in event.tool_calls:
-                self._tool_id_to_message_id[tool_call.tool_id] = message_id
-                pending.add(tool_call.tool_id)
-                events.append(
-                    self._create_tool_call_start_event(
-                        message_id=message_id,
-                        tool_call_id=tool_call.tool_id,
-                        tool_name=tool_call.tool_name,
-                        input=tool_call.tool_kwargs,
+            if self.storage is not None:
+                async with self._storage_lock:
+                    existing: dict[str, str] | None = await self.storage.get_value(
+                        self.runtime_id,
+                        STORAGE_NAMESPACE_EVENT_MAPPER,
+                        STORAGE_KEY_TOOL_ID_TO_MESSAGE_ID_MAP,
                     )
-                )
-            self._pending_tool_calls[message_id] = pending
+                    tool_id_to_message_id: dict[str, str] = existing or {}
+
+                    for tool_call in event.tool_calls:
+                        tool_id_to_message_id[tool_call.tool_id] = message_id
+                        events.append(
+                            self._create_tool_call_start_event(
+                                message_id=message_id,
+                                tool_call_id=tool_call.tool_id,
+                                tool_name=tool_call.tool_name,
+                                input=tool_call.tool_kwargs,
+                            )
+                        )
+
+                    await self.storage.set_value(
+                        self.runtime_id,
+                        STORAGE_NAMESPACE_EVENT_MAPPER,
+                        STORAGE_KEY_TOOL_ID_TO_MESSAGE_ID_MAP,
+                        tool_id_to_message_id,
+                    )
+            else:
+                # In-memory fallback (no suspend/resume support)
+                pending: set[str] = set()
+                for tool_call in event.tool_calls:
+                    self._tool_id_to_message_id[tool_call.tool_id] = message_id
+                    pending.add(tool_call.tool_id)
+                    events.append(
+                        self._create_tool_call_start_event(
+                            message_id=message_id,
+                            tool_call_id=tool_call.tool_id,
+                            tool_name=tool_call.tool_name,
+                            input=tool_call.tool_kwargs,
+                        )
+                    )
+                self._pending_tool_calls[message_id] = pending
             # message_end will be emitted once the last ToolCallResult comes in
         else:
             # No tool calls: this is the final text response, close the message now
             events.append(self._create_message_end_event(message_id))
 
         return events if events else None
 
-    def _map_tool_call_result(
+    async def _map_tool_call_result(
         self, event: ToolCallResult
     ) -> list[UiPathConversationMessageEvent] | None:
-        message_id = self._tool_id_to_message_id.pop(event.tool_id, None)
+        message_id, is_last = await self._get_message_id_for_tool_call(event.tool_id)
         if message_id is None:
             logger.warning(
                 "ToolCallResult received for unknown tool_id '%s' — skipping.",
@@ -180,14 +221,70 @@ def _map_tool_call_result(
         ]
 
         # Close the message once all tool calls for it have completed
+        if is_last:
+            events.append(self._create_message_end_event(message_id))
+
+        return events
+
+    async def _get_message_id_for_tool_call(
+        self, tool_id: str
+    ) -> tuple[str | None, bool]:
+        """Look up the message_id for a tool_id and remove it from the map.
+
+        Returns (message_id, is_last) where is_last is True when no other
+        pending tool calls remain for the same message.
+        """
+        if self.storage is not None:
+            async with self._storage_lock:
+                tool_id_to_message_id: dict[str, str] | None = (
+                    await self.storage.get_value(
+                        self.runtime_id,
+                        STORAGE_NAMESPACE_EVENT_MAPPER,
+                        STORAGE_KEY_TOOL_ID_TO_MESSAGE_ID_MAP,
+                    )
+                )
+
+                if tool_id_to_message_id is None:
+                    logger.error(
+                        "attempt to lookup tool_id %s when no map present in storage",
+                        tool_id,
+                    )
+                    return None, False
+
+                message_id = tool_id_to_message_id.get(tool_id)
+                if message_id is None:
+                    logger.error(
+                        "tool_id to message map does not contain tool_id %s",
+                        tool_id,
+                    )
+                    return None, False
+
+                del tool_id_to_message_id[tool_id]
+
+                await self.storage.set_value(
+                    self.runtime_id,
+                    STORAGE_NAMESPACE_EVENT_MAPPER,
+                    STORAGE_KEY_TOOL_ID_TO_MESSAGE_ID_MAP,
+                    tool_id_to_message_id,
+                )
+
+                is_last = message_id not in tool_id_to_message_id.values()
+
+            return message_id, is_last
+
+        # In-memory fallback
+        message_id = self._tool_id_to_message_id.pop(tool_id, None)
+        if message_id is None:
+            return None, False
+
         pending = self._pending_tool_calls.get(message_id)
         if pending is not None:
-            pending.discard(event.tool_id)
+            pending.discard(tool_id)
             if not pending:
                 del self._pending_tool_calls[message_id]
-                events.append(self._create_message_end_event(message_id))
+                return message_id, True
 
-        return events
+        return message_id, False
 
     # ── Factory helpers ────────────────────────────────────────────────────────
 

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/runtime.py

@@ -169,7 +169,7 @@ async def _run_workflow(
 
         event_stream = handler.stream_events(expose_internal=True)
         suspended_event: InputRequiredEvent | None = None
-        chat = UiPathChatMessagesMapper(runtime_id=self.runtime_id)
+        chat = UiPathChatMessagesMapper(runtime_id=self.runtime_id, storage=self.storage)
 
         is_resumed: bool = False
         async for event in event_stream: