Support remote storage for tool-registry + SDK roll changes for platform service

backend/sample.env

@@ -146,6 +146,9 @@ TOOL_REGISTRY_CONFIG_PATH="/data/tool_registry_config"
 # Flipt Service
 FLIPT_SERVICE_AVAILABLE=False
 
+# Remote storage config for tool registry
+TOOL_REGISTRY_STORAGE_CREDENTIALS='{"provider":"local"}'
+
 
 # File System Configuration for Workflow and API Execution
 

docker/dockerfiles/platform.Dockerfile

@@ -12,7 +12,14 @@ ENV BUILD_CONTEXT_PATH platform-service
 ENV BUILD_PACKAGES_PATH unstract
 ENV PDM_VERSION 2.16.1
 
-RUN pip install --no-cache-dir -U pip pdm~=${PDM_VERSION}; \
+RUN apt-get update; \
+    apt-get --no-install-recommends install -y  \
+    # unstract sdk
+    build-essential libmagic-dev; \
+    \
+    apt-get clean && rm -rf /var/lib/apt/lists/* /var/cache/apt/archives/*; \
+    \
+    pip install --no-cache-dir -U pip pdm~=${PDM_VERSION}; \
     \
     # Creates a non-root user with an explicit UID and adds permission to access the /app folder
     # For more info, please refer to https://aka.ms/vscode-docker-python-configure-containers

platform-service/pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
     "cryptography>=41.0.7",
     "unstract-flags @ file:///${PROJECT_ROOT}/../unstract/flags",
     "requests>=2.31.0",
-    "unstract-sdk~=0.54.0rc1",
+    "unstract-sdk==0.54.0rc2",
     "gcsfs==2024.10.0",
 ]
 requires-python = ">=3.9,<3.11.1"

platform-service/src/unstract/platform_service/helper/cost_calculation.py

@@ -30,21 +30,26 @@ def __init__(
         self.file_path = file_path
 
         if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
-            self.file_storage, self.file_path = self.__get_storage_credentials()
+            self.file_storage, self.file_path = self._get_storage_credentials()
         self.model_token_data = self._get_model_token_data()
 
     if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
 
-        def __get_storage_credentials(self) -> tuple[PermanentFileStorage, str]:
+        def _get_storage_credentials(self) -> tuple[PermanentFileStorage, str]:
             try:
                 # Not creating constants for now for the keywords below as this
                 # logic ought to change in the near future to maintain unformity
                 # across services
                 file_storage = json.loads(os.environ.get("FILE_STORAGE_CREDENTIALS"))
                 provider = FileStorageProvider(file_storage["provider"])
-                credentials = file_storage["credentials"]
+                credentials = file_storage.get("credentials", {})
                 file_path = file_storage["model_prices_file_path"]
                 return PermanentFileStorage(provider, **credentials), file_path
+            except KeyError as e:
+                app.logger.error(
+                    f"Required credentials is missing in the env: {str(e)}"
+                )
+                raise e
             except FileStorageError as e:
                 app.logger.error(
                     "Error while initialising storage: %s",
@@ -142,8 +147,6 @@ def _fetch_and_save_json(self) -> Optional[dict[str, Any]]:
             if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
                 self.file_storage.json_dump(
                     path=self.file_path,
-                    mode="w",
-                    encoding="utf-8",
                     data=json_data,
                     ensure_ascii=False,
                     indent=4,

unstract/tool-registry/pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     # Hence instead, add the dependencies without version constraints where the
     # assumption is they are added as direct dependencies in main project itself.
     "unstract-tool-sandbox",
+    "unstract-flags",
 ]
 requires-python = ">=3.9,<3.11.1"
 readme = "README.md"

unstract/tool-registry/src/unstract/tool_registry/constants.py

@@ -1,6 +1,13 @@
 from typing import Any
 
 
+class FeatureFlag:
+    """Temporary feature flags."""
+
+    # For enabling remote storage feature
+    REMOTE_FILE_STORAGE = "remote_file_storage"
+
+
 class Tools:
     TOOLS_DIRECTORY = "tools"
     IMAGE_LATEST_TAG = "latest"

unstract/tool-registry/src/unstract/tool_registry/helper.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Any, Optional
 
+from unstract.sdk.file_storage import FileStorage, FileStorageProvider
 from unstract.tool_registry.constants import PropKey
 from unstract.tool_registry.dto import Tool, ToolMeta
 from unstract.tool_registry.exceptions import (
@@ -20,6 +21,7 @@ def __init__(
         registry: str,
         private_tools_file: str,
         public_tools_file: str,
+        fs: FileStorage = FileStorage(FileStorageProvider.LOCAL),
     ) -> None:
         """Helper class for ToolRegistry.
 
@@ -32,6 +34,7 @@ def __init__(
         self.registry_file = registry
         self.private_tools_file = private_tools_file
         self.public_tools_file = public_tools_file
+        self.fs = fs
         self.tools = self._load_tools_from_registry_file()
         if self.tools:
             logger.info(f"Loaded tools from registry YAML: {self.tools}")
@@ -214,7 +217,7 @@ def save_registry(self, data: dict[str, Any]) -> None:
             RegistryNotFound: _description_
         """
         try:
-            ToolUtils.save_registry(self.registry_file, data=data)
+            ToolUtils.save_registry(self.registry_file, data=data, fs=self.fs)
         except FileNotFoundError:
             logger.error(f"File not found: {self.registry_file}")
             raise RegistryNotFound()
@@ -298,7 +301,9 @@ def get_all_tools_from_disk(self) -> dict[str, dict[str, Any]]:
         tools = {}
         for tool_file in tool_files:
             try:
-                data = ToolUtils.get_all_tools_from_disk(file_path=tool_file)
+                data = ToolUtils.get_all_tools_from_disk(
+                    file_path=tool_file, fs=self.fs
+                )
                 if not data:
                     logger.info(f"No data from {tool_file}")
                 tool_version_list = [
@@ -309,7 +314,6 @@ def get_all_tools_from_disk(self) -> dict[str, dict[str, Any]]:
                 tools.update(data)
             except FileNotFoundError:
                 logger.warning(f"Unable to find tool file to load tools: {tool_file}")
-                pass
         return tools
 
     def get_tool_data_by_id(self, tool_uid: str) -> dict[str, Any]:

unstract/tool-registry/src/unstract/tool_registry/tool_registry.py

@@ -1,21 +1,37 @@
+import json
 import logging
 import os
 from typing import Any, Optional
 
-from unstract.tool_registry.constants import PropKey, ToolJsonField, ToolKey
+from unstract.tool_registry.constants import (
+    FeatureFlag,
+    PropKey,
+    ToolJsonField,
+    ToolKey,
+)
 from unstract.tool_registry.dto import Tool
 from unstract.tool_registry.exceptions import InvalidToolURLException
 from unstract.tool_registry.helper import ToolRegistryHelper
 from unstract.tool_registry.schema_validator import JsonSchemaValidator
 from unstract.tool_registry.tool_utils import ToolUtils
 
+from unstract.flags.feature_flag import check_feature_flag_status
+
+if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+    from unstract.sdk.exceptions import FileStorageError
+    from unstract.sdk.file_storage import FileStorageProvider, PermanentFileStorage
+
 logger = logging.getLogger(__name__)
 
 
 class ToolRegistry:
     REGISTRY_FILE = "registry.yaml"
-    PRIVATE_TOOL_CONFIG_FILE = "private_tools.json"
-    PUBLIC_TOOL_CONFIG_FILE = "public_tools.json"
+    if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+        PRIVATE_TOOL_CONFIG_FILE = "private_tools_remote_storage.json"
+        PUBLIC_TOOL_CONFIG_FILE = "public_tools_remote_storage.json"
+    else:
+        PRIVATE_TOOL_CONFIG_FILE = "private_tools.json"
+        PUBLIC_TOOL_CONFIG_FILE = "public_tools.json"
 
     def __init__(
         self,
@@ -45,12 +61,43 @@ def __init__(
                 "Env 'TOOL_REGISTRY_CONFIG_PATH' is not set, please add the tool "
                 "registry JSONs and YAML to a directory and set the env."
             )
+
+        if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+            self.fs = self._get_storage_credentials()
+        else:
+            self.fs = None
         self.helper = ToolRegistryHelper(
             registry=os.path.join(directory, registry_file),
             private_tools_file=os.path.join(directory, private_tools),
             public_tools_file=os.path.join(directory, public_tools),
+            fs=self.fs,
         )
 
+    if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+
+        def _get_storage_credentials(self) -> PermanentFileStorage:
+            try:
+                # Not creating constants for now for the keywords below as this
+                # logic ought to change in the near future to maintain unformity
+                # across services
+                file_storage = json.loads(
+                    os.environ.get("TOOL_REGISTRY_STORAGE_CREDENTIALS", {})
+                )
+                provider = FileStorageProvider(file_storage["provider"])
+                credentials = file_storage.get("credentials", {})
+                return PermanentFileStorage(provider, **credentials)
+            except KeyError as e:
+                logger.error(f"Required credentials is missing in the env: {str(e)}")
+                raise e
+            except FileStorageError as e:
+                logger.error(
+                    "Error while initialising storage: %s",
+                    e,
+                    stack_info=True,
+                    exc_info=True,
+                )
+                raise e
+
     def load_all_tools_to_disk(self) -> None:
         self.helper.load_all_tools_to_disk()
 

unstract/tool-registry/src/unstract/tool_registry/tool_utils.py

@@ -5,10 +5,13 @@
 
 import yaml
 from unstract.sdk.adapters.enums import AdapterTypes
-from unstract.tool_registry.constants import AdapterPropertyKey, Tools
+from unstract.sdk.file_storage import FileStorage, FileStorageProvider
+from unstract.tool_registry.constants import AdapterPropertyKey, FeatureFlag, Tools
 from unstract.tool_registry.dto import AdapterProperties, Spec, Tool, ToolMeta
 from unstract.tool_registry.exceptions import InvalidToolURLException, RegistryNotFound
 
+from unstract.flags.feature_flag import check_feature_flag_status
+
 logger = logging.getLogger(__name__)
 
 
@@ -55,39 +58,67 @@ def get_tool_meta_from_tool_url(registry_tool: str) -> ToolMeta:
         )
 
     @staticmethod
-    def save_tools_in_to_disk(file_path: str, data: dict[str, Any]) -> None:
-        with open(file_path, "w") as json_file:
-            json.dump(data, json_file)
+    def save_tools_in_to_disk(
+        file_path: str,
+        data: dict[str, Any],
+        fs: FileStorage = FileStorage(FileStorageProvider.LOCAL),
+    ) -> None:
+        if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+            fs.json_dump(path=file_path, mode="w", encoding="utf-8", data=data)
+        else:
+            with open(file_path, "w") as json_file:
+                json.dump(data, json_file)
 
     @staticmethod
-    def get_all_tools_from_disk(file_path: str) -> dict[str, Any]:
+    def get_all_tools_from_disk(
+        file_path: str, fs: FileStorage = FileStorage(FileStorageProvider.LOCAL)
+    ) -> dict[str, Any]:
         try:
-            with open(file_path) as json_file:
-                data: dict[str, Any] = json.load(json_file)
+            if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+                data = fs.json_load(file_path)
+            else:
+                with open(file_path) as json_file:
+                    data: dict[str, Any] = json.load(json_file)
             return data
         except json.JSONDecodeError as e:
             logger.warning(f"Error loading tools from {file_path}: {e}")
             return {}
 
     @staticmethod
-    def save_registry(file_path: str, data: dict[str, Any]) -> None:
-        with open(file_path, "w") as file:
-            yaml.dump(data, file, default_flow_style=False)
+    def save_registry(
+        file_path: str,
+        data: dict[str, Any],
+        fs: FileStorage = FileStorage(FileStorageProvider.LOCAL),
+    ) -> None:
+        if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+            fs.yaml_dump(path=file_path, mode="w", encoding="utf-8", data=data)
+        else:
+            with open(file_path, "w") as file:
+                yaml.dump(data, file, default_flow_style=False)
 
     @staticmethod
-    def get_registry(file_path: str, raise_exc: bool = False) -> dict[str, Any]:
+    def get_registry(
+        file_path: str,
+        fs: FileStorage = FileStorage(FileStorageProvider.LOCAL),
+        raise_exc: bool = False,
+    ) -> dict[str, Any]:
         """Get Registry File.
 
         Args:
             file_path (str): file path of registry.yaml
+            fs (FileStorage): file storage to be used
 
         Returns:
             dict[str, Any]: yaml data
         """
         yml_data: dict[str, Any] = {}
         try:
-            with open(file_path) as file:
-                yml_data = yaml.safe_load(file)
+            logger.debug(f"Reading tool registry YAML: {file_path}")
+            if check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+                yml_data = fs.yaml_load(file_path)
+            else:
+                with open(file_path) as file:
+                    yml_data = yaml.safe_load(file)
         except FileNotFoundError:
             logger.warning(f"Could not find tool registry YAML: {str(file_path)}")
             if raise_exc: