Merge pull request #94 from UiPath/fix/use_sqllite_llamaindex_storage

fix: store context in sqllite

Author: cristipufu

pyproject.toml

@@ -1,15 +1,16 @@
 [project]
 name = "uipath-llamaindex"
-version = "0.1.4"
+version = "0.1.5"
 description = "UiPath LlamaIndex SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
 dependencies = [
+    "aiosqlite>=0.21.0",
     "llama-index>=0.14.8",
     "llama-index-embeddings-azure-openai>=0.4.1",
     "llama-index-llms-azure-openai>=0.4.2",
     "openinference-instrumentation-llama-index>=4.3.9",
-    "uipath>=2.2.16, <2.3.0",
+    "uipath>=2.2.26, <2.3.0",
 ]
 classifiers = [
     "Intended Audience :: Developers",

…maindex/runtime/_attribute_normalizer.py …/uipath_llamaindex/runtime/_telemetry.pysrc/uipath_llamaindex/runtime/_attribute_normalizer.py renamed to src/uipath_llamaindex/runtime/_telemetry.py src/uipath_llamaindex/runtime/_attribute_normalizer.py renamed to src/uipath_llamaindex/runtime/_telemetry.py

@@ -15,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 
-class AttributeNormalizingSpanProcessor(SpanProcessor):
+class ToolCallAttributeNormalizer(SpanProcessor):
     """Normalizes LlamaIndex tool call attributes to match other frameworks.
 
     Unwraps {"kwargs": {...}} to flat {...} format for consistency with LangChain.

src/uipath_llamaindex/runtime/factory.py

@@ -17,16 +17,16 @@
 from uipath.runtime.errors import UiPathErrorCategory
 from workflows import Workflow
 
-from uipath_llamaindex.runtime._attribute_normalizer import (
-    AttributeNormalizingSpanProcessor,
+from uipath_llamaindex.runtime._telemetry import (
+    ToolCallAttributeNormalizer,
 )
 from uipath_llamaindex.runtime.config import LlamaIndexConfig
 from uipath_llamaindex.runtime.errors import (
     UiPathLlamaIndexErrorCode,
     UiPathLlamaIndexRuntimeError,
 )
 from uipath_llamaindex.runtime.runtime import UiPathLlamaIndexRuntime
-from uipath_llamaindex.runtime.storage import PickleResumableStorage
+from uipath_llamaindex.runtime.storage import SQLiteResumableStorage
 from uipath_llamaindex.runtime.workflow import LlamaIndexWorkflowLoader
 
 
@@ -45,12 +45,14 @@ def __init__(
         """
         self.context = context
         self._config: LlamaIndexConfig | None = None
-        self._storage_path: str | None = None
 
         self._workflow_cache: dict[str, Workflow] = {}
         self._workflow_loaders: dict[str, LlamaIndexWorkflowLoader] = {}
         self._workflow_lock = asyncio.Lock()
 
+        self._storage_lock = asyncio.Lock()
+        self._storage: SQLiteResumableStorage | None = None
+
         self._setup_instrumentation(self.context.trace_manager)
 
     def _setup_instrumentation(self, trace_manager: UiPathTraceManager | None) -> None:
@@ -60,26 +62,37 @@ def _setup_instrumentation(self, trace_manager: UiPathTraceManager | None) -> No
 
         if trace_manager:
             trace_manager.tracer_provider.add_span_processor(
-                AttributeNormalizingSpanProcessor()
+                ToolCallAttributeNormalizer()
             )
 
     def _get_storage_path(self) -> str:
         """Get the storage path for workflow state."""
-        if self._storage_path is None:
-            if self.context.runtime_dir and self.context.state_file:
-                path = os.path.join(self.context.runtime_dir, self.context.state_file)
-                if not self.context.resume and self.context.job_id is None:
-                    # If not resuming and no job id, delete the previous state file
-                    if os.path.exists(path):
-                        os.remove(path)
-                os.makedirs(self.context.runtime_dir, exist_ok=True)
-                self._storage_path = path
-            else:
-                default_path = os.path.join("__uipath", "state.db")
-                os.makedirs(os.path.dirname(default_path), exist_ok=True)
-                self._storage_path = default_path
-
-        return self._storage_path
+        if self.context.runtime_dir and self.context.state_file:
+            path = os.path.join(self.context.runtime_dir, self.context.state_file)
+            if not self.context.resume and self.context.job_id is None:
+                # If not resuming and no job id, delete the previous state file
+                if os.path.exists(path):
+                    os.remove(path)
+            os.makedirs(self.context.runtime_dir, exist_ok=True)
+            return path
+
+        default_path = os.path.join("__uipath", "state.db")
+        os.makedirs(os.path.dirname(default_path), exist_ok=True)
+        return default_path
+
+    async def _get_storage(self) -> SQLiteResumableStorage:
+        """Get or create the shared storage instance."""
+        if self._storage is not None:
+            return self._storage
+
+        async with self._storage_lock:
+            if self._storage is not None:
+                return self._storage
+
+            storage_path = self._get_storage_path()
+            self._storage = SQLiteResumableStorage(storage_path)
+            await self._storage.setup()
+            return self._storage
 
     def _load_config(self) -> LlamaIndexConfig:
         """Load llama_index.json configuration."""
@@ -199,7 +212,6 @@ async def discover_runtimes(self) -> list[UiPathRuntimeProtocol]:
             List of LlamaIndexRuntime instances, one per entrypoint
         """
         entrypoints = self.discover_entrypoints()
-        storage_path = self._get_storage_path()
 
         runtimes: list[UiPathRuntimeProtocol] = []
         for entrypoint in entrypoints:
@@ -209,7 +221,6 @@ async def discover_runtimes(self) -> list[UiPathRuntimeProtocol]:
                 workflow=workflow,
                 runtime_id=entrypoint,
                 entrypoint=entrypoint,
-                storage_path=storage_path,
             )
             runtimes.append(runtime)
 
@@ -220,7 +231,6 @@ async def _create_runtime_instance(
         workflow: Workflow,
         runtime_id: str,
         entrypoint: str,
-        storage_path: str,
     ) -> UiPathRuntimeProtocol:
         """
         Create a runtime instance from a workflow.
@@ -229,12 +239,12 @@ async def _create_runtime_instance(
             workflow: The workflow
             runtime_id: Unique identifier for the runtime instance
             entrypoint: Workflow entrypoint name
-            storage_path: Path for state storage
 
         Returns:
             Configured runtime instance
         """
-        storage = PickleResumableStorage(storage_path)
+
+        storage = await self._get_storage()
 
         base_runtime = UiPathLlamaIndexRuntime(
             workflow=workflow,
@@ -265,15 +275,12 @@ async def new_runtime(
         Returns:
             Configured runtime instance with workflow
         """
-        storage_path = self._get_storage_path()
-
         workflow = await self._resolve_workflow(entrypoint)
 
         return await self._create_runtime_instance(
             workflow=workflow,
             runtime_id=runtime_id,
             entrypoint=entrypoint,
-            storage_path=storage_path,
         )
 
     async def dispose(self) -> None:

src/uipath_llamaindex/runtime/runtime.py

@@ -47,7 +47,7 @@
     UiPathLlamaIndexRuntimeError,
 )
 from uipath_llamaindex.runtime.schema import get_entrypoints_schema, get_workflow_schema
-from uipath_llamaindex.runtime.storage import PickleResumableStorage
+from uipath_llamaindex.runtime.storage import SQLiteResumableStorage
 
 from ._serialize import serialize_output
 
@@ -62,7 +62,7 @@ def __init__(
         workflow: Workflow,
         runtime_id: str | None = None,
         entrypoint: str | None = None,
-        storage: PickleResumableStorage | None = None,
+        storage: SQLiteResumableStorage | None = None,
         debug_mode: bool = False,
     ):
         """
@@ -76,7 +76,7 @@ def __init__(
         self.workflow: Workflow = workflow
         self.runtime_id: str = runtime_id or "default"
         self.entrypoint: str | None = entrypoint
-        self.storage: PickleResumableStorage | None = storage
+        self.storage: SQLiteResumableStorage | None = storage
         self.debug_mode: bool = debug_mode
         self._context: Context | None = None
 
@@ -142,7 +142,7 @@ async def _run_workflow(
         is_resuming = bool(options and options.resume)
 
         if is_resuming:
-            self._context = self._load_context()
+            self._context = await self._load_context()
         else:
             self._context = Context(self.workflow)
 
@@ -215,7 +215,7 @@ async def _run_workflow(
 
         if suspended_event is not None:
             await asyncio.sleep(0)  # Yield control to event loop
-            self._save_context()
+            await self._save_context()
             await handler.cancel_run()
             if isinstance(suspended_event, BreakpointEvent):
                 yield self._create_breakpoint_result(suspended_event)
@@ -325,12 +325,12 @@ def _create_runtime_error(self, e: Exception) -> UiPathLlamaIndexRuntimeError:
             UiPathErrorCategory.USER,
         )
 
-    def _load_context(self) -> Context:
+    async def _load_context(self) -> Context:
         """Load the workflow context from storage."""
         if not self.storage:
             return Context(self.workflow)
 
-        context_dict = self.storage.load_context()
+        context_dict = await self.storage.load_context(runtime_id=self.runtime_id)
 
         if context_dict:
             from workflows.context.serializers import JsonPickleSerializer
@@ -344,7 +344,7 @@ def _load_context(self) -> Context:
         else:
             return Context(self.workflow)
 
-    def _save_context(self) -> None:
+    async def _save_context(self) -> None:
         """Save the current workflow context to storage."""
         if not self.storage or not self._context:
             return None
@@ -354,7 +354,9 @@ def _save_context(self) -> None:
         serializer = JsonPickleSerializer()
         context_dict = self._context.to_dict(serializer=serializer)
 
-        self.storage.save_context(context_dict)
+        await self.storage.save_context(
+            runtime_id=self.runtime_id, context_dict=context_dict
+        )
 
     async def dispose(self) -> None:
         """Cleanup runtime resources."""