[FEAT] Unified remote APIs - Prompt studio changes

backend/file_management/file_management_helper.py

@@ -23,6 +23,8 @@
 from file_management.file_management_dto import FileInformation
 from fsspec import AbstractFileSystem
 from pydrive2.files import ApiRequestError
+from utils.constants import FileStorageType
+from utils.file_storage_helper import FileStorageUtil
 
 from unstract.connectors.filesystems import connectors as fs_connectors
 from unstract.connectors.filesystems.unstract_file_system import UnstractFileSystem
@@ -31,7 +33,6 @@
 
 
 class FileManagerHelper:
-
     @staticmethod
     def get_file_system(connector: ConnectorInstance) -> UnstractFileSystem:
         """Creates the `UnstractFileSystem` for the corresponding connector."""
@@ -141,15 +142,15 @@ def upload_file(
             remote_file.write(file.read())
 
     @staticmethod
-    def fetch_file_contents(file_system: UnstractFileSystem, file_path: str) -> Any:
-        fs = file_system.get_fsspec_fs()
-        try:
-            file_info = fs.info(file_path)
-        except FileNotFoundError:
+    def fetch_file_contents(file_path: str) -> Any:
+        file_storage = FileStorageUtil.initialize_file_storage(
+            FileStorageType.PERMANENT
+        )
+        if not file_storage.exists(file_path):
             raise FileNotFound
 
-        file_content_type = file_info.get("ContentType")
-        file_type = file_info.get("type")
+        # file_content_type = file_info.get("ContentType")
+        # file_type = file_info.get("type")
         if file_type != "file":
             raise InvalidFileType
         try:

backend/utils/common_utils.py

@@ -1,4 +1,5 @@
 import json
+import os
 from enum import Enum
 
 
@@ -31,6 +32,23 @@ def is_json(string: str) -> bool:
             return False
         return True
 
+    @staticmethod
+    def get_env_or_die(env_key: str) -> str:
+        """Returns the value of an env variable.
+
+        If its empty or None, raises an error and exits
+
+        Args:
+            env_key (str): Key to retrieve
+
+        Returns:
+            str: Value of the env
+        """
+        env_value = os.environ.get(env_key)
+        if env_value is None or env_value == "":
+            raise ValueError(f"Env variable '{env_key}' is required")
+        return env_value
+
 
 class ModelEnum(Enum):
     @classmethod

backend/utils/constants.py

@@ -1,4 +1,5 @@
 import os
+from enum import Enum
 
 from django.conf import settings
 from utils.common_utils import CommonUtils
@@ -78,3 +79,14 @@ class ExecutionLogConstants:
     PERIODIC_TASK_NAME_V2 = "workflow_log_history_v2"
     TASK = "workflow_manager.workflow.execution_log_utils.consume_log_history"
     TASK_V2 = "consume_log_history"
+
+
+class FileStorageKeys:
+    FILE_STORAGE_PROVIDER = "FILE_STORAGE_PROVIDER"
+    FILE_STORAGE_CREDENTIALS = "FILE_STORAGE_CREDENTIALS"
+
+
+class FileStorageType(Enum):
+    PERMANENT = "permanent"
+    DEFAULT = "default"
+    TEMPORARY = "temporary"

backend/utils/file_storage_helper.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from unstract.sdk.file_storage.fs_impl import FileStorage
+from unstract.sdk.file_storage.fs_permanent import PermanentFileStorage
+from utils.common_utils import CommonUtils
+from utils.constants import FileStorageKeys, FileStorageType
+
+
+class FileStorageUtil:
+    @staticmethod
+    def initialize_file_storage(type: FileStorageType) -> FileStorage:
+        provider_data = FileStorageUtil.load_file_storage_envs()
+        provider = provider_data[FileStorageKeys.FILE_STORAGE_PROVIDER]
+        credentials = provider_data[FileStorageKeys.FILE_STORAGE_CREDENTIALS]
+        if type.value == FileStorageType.PERMANENT.value:
+            file_storage = PermanentFileStorage(
+                provider=provider, credentials=credentials
+            )
+        if type.value == FileStorageType.DEFAULT.value:
+            file_storage = FileStorage(provider=provider, credentials=credentials)
+        return file_storage
+
+    @staticmethod
+    def load_file_storage_envs() -> dict[str, Any]:
+        provider: str = CommonUtils.get_env_or_die(
+            env_key=FileStorageKeys.FILE_STORAGE_PROVIDER
+        )
+        credentials: str = CommonUtils.get_env_or_die(
+            env_key=FileStorageKeys.FILE_STORAGE_CREDENTIALS
+        )
+        provider_data: dict[str, Any] = {}
+        provider_data[FileStorageKeys.FILE_STORAGE_PROVIDER] = provider
+        provider_data[FileStorageKeys.FILE_STORAGE_CREDENTIALS] = credentials
+        return provider_data