[FEAT] Remote File Storage for Prompt Studio

backend/backend/constants.py

@@ -34,3 +34,4 @@ class FeatureFlag:
     """Temporary feature flags."""
 
     APP_DEPLOYMENT = "app_deployment"
+    REMOTE_FILE_STORAGE = "remote_file_storage"

backend/file_management/file_management_helper.py

@@ -7,6 +7,7 @@
 
 import magic
 from connector_v2.models import ConnectorInstance
+from deprecated import deprecated
 from django.conf import settings
 from django.http import StreamingHttpResponse
 from file_management.exceptions import (
@@ -141,6 +142,7 @@ def upload_file(
             remote_file.write(file.read())
 
     @staticmethod
+    @deprecated(reason="Use remote FS APIs from SDK")
     def fetch_file_contents(file_system: UnstractFileSystem, file_path: str) -> Any:
         fs = file_system.get_fsspec_fs()
         try:

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -12,7 +12,6 @@
 from adapter_processor.models import AdapterInstance
 from django.conf import settings
 from django.db.models.manager import BaseManager
-from file_management.file_management_helper import FileManagerHelper
 from prompt_studio.modifier_loader import ModifierConfig
 from prompt_studio.modifier_loader import load_plugins as load_modifier_plugins
 from prompt_studio.prompt_profile_manager.models import ProfileManager
@@ -51,6 +50,7 @@
 from unstract.sdk.index import Index
 from unstract.sdk.prompt import PromptTool
 from unstract.sdk.utils.tool_utils import ToolUtils
+from utils.file_storage.helpers.prompt_studio_file_helper import PromptStudioFileHelper
 from utils.local_context import StateStore
 
 from unstract.core.pubsub_helper import LogPublisher
@@ -331,7 +331,7 @@ def index_document(
             file_path = file_name
         else:
             default_profile = ProfileManager.get_default_llm_profile(tool)
-            file_path = FileManagerHelper.handle_sub_directory_for_tenants(
+            file_path = PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
                 org_id,
                 is_create=False,
                 user_id=user_id,
@@ -609,7 +609,7 @@ def _execute_prompts_in_single_pass(
 
     @staticmethod
     def _get_document_path(org_id, user_id, tool_id, doc_name):
-        doc_path = FileManagerHelper.handle_sub_directory_for_tenants(
+        doc_path = PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
             org_id=org_id,
             user_id=user_id,
             tool_id=tool_id,
@@ -621,7 +621,7 @@ def _get_document_path(org_id, user_id, tool_id, doc_name):
     def _get_extract_or_summary_document_path(
         org_id, user_id, tool_id, doc_name, doc_type
     ) -> str:
-        doc_path = FileManagerHelper.handle_sub_directory_for_tenants(
+        doc_path = PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
             org_id=org_id,
             user_id=user_id,
             tool_id=tool_id,

backend/prompt_studio/prompt_studio_core/views.py

@@ -52,9 +52,12 @@
 from rest_framework.versioning import URLPathVersioning
 from tool_instance.models import ToolInstance
 from unstract.sdk.utils.common_utils import CommonUtils
+from utils.file_storage.helpers.prompt_studio_file_helper import PromptStudioFileHelper
 from utils.user_session import UserSessionUtils
 
+from backend.constants import FeatureFlag
 from unstract.connectors.filesystems.local_storage.local_storage import LocalStorageFS
+from unstract.flags.feature_flag import check_feature_flag_status
 
 from .models import CustomTool
 from .serializers import (
@@ -403,31 +406,41 @@ def fetch_contents_ide(self, request: HttpRequest, pk: Any = None) -> Response:
                 f"{FileViewTypes.SUMMARIZE.lower()}/"
                 f"{filename_without_extension}.txt"
             )
-
-        file_path = file_path = FileManagerHelper.handle_sub_directory_for_tenants(
-            UserSessionUtils.get_organization_id(request),
-            is_create=True,
-            user_id=custom_tool.created_by.user_id,
-            tool_id=str(custom_tool.tool_id),
-        )
-        file_system = LocalStorageFS(settings={"path": file_path})
-        if not file_path.endswith("/"):
-            file_path += "/"
-        file_path += file_name
-        # Temporary Hack for frictionless onboarding as the user id will be empty
-        try:
-            contents = FileManagerHelper.fetch_file_contents(file_system, file_path)
-        except FileNotFound:
+        if not check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
             file_path = file_path = FileManagerHelper.handle_sub_directory_for_tenants(
                 UserSessionUtils.get_organization_id(request),
                 is_create=True,
-                user_id="",
+                user_id=custom_tool.created_by.user_id,
                 tool_id=str(custom_tool.tool_id),
             )
+            file_system = LocalStorageFS(settings={"path": file_path})
             if not file_path.endswith("/"):
                 file_path += "/"
-                file_path += file_name
-            contents = FileManagerHelper.fetch_file_contents(file_system, file_path)
+            file_path += file_name
+            # TODO : Handle this with proper fix
+            # Temporary Hack for frictionless onboarding as the user id will be empty
+            try:
+                contents = FileManagerHelper.fetch_file_contents(file_system, file_path)
+            except FileNotFound:
+                file_path = file_path = (
+                    FileManagerHelper.handle_sub_directory_for_tenants(
+                        UserSessionUtils.get_organization_id(request),
+                        is_create=True,
+                        user_id="",
+                        tool_id=str(custom_tool.tool_id),
+                    )
+                )
+                if not file_path.endswith("/"):
+                    file_path += "/"
+                    file_path += file_name
+                contents = FileManagerHelper.fetch_file_contents(file_system, file_path)
+        else:
+            PromptStudioFileHelper.fetch_file_contents(
+                file_name=file_name,
+                org_id=UserSessionUtils.get_organization_id(request),
+                user_id=custom_tool.created_by.user_id,
+                tool_id=str(custom_tool.tool_id),
+            )
 
         return Response({"data": contents}, status=status.HTTP_200_OK)
 
@@ -437,15 +450,6 @@ def upload_for_ide(self, request: HttpRequest, pk: Any = None) -> Response:
         serializer = FileUploadIdeSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
         uploaded_files: Any = serializer.validated_data.get("file")
-
-        file_path = FileManagerHelper.handle_sub_directory_for_tenants(
-            UserSessionUtils.get_organization_id(request),
-            is_create=True,
-            user_id=custom_tool.created_by.user_id,
-            tool_id=str(custom_tool.tool_id),
-        )
-        file_system = LocalStorageFS(settings={"path": file_path})
-
         documents = []
         for uploaded_file in uploaded_files:
             file_name = uploaded_file.name
@@ -464,11 +468,11 @@ def upload_for_ide(self, request: HttpRequest, pk: Any = None) -> Response:
             logger.info(
                 f"Uploading file: {file_name}" if file_name else "Uploading file"
             )
-            FileManagerHelper.upload_file(
-                file_system,
-                file_path,
-                uploaded_file,
-                file_name,
+            PromptStudioFileHelper.upload_for_ide(
+                org_id=UserSessionUtils.get_organization_id(request),
+                user_id=custom_tool.created_by.user_id,
+                tool_id=str(custom_tool.tool_id),
+                uploaded_file=uploaded_file,
             )
             documents.append(doc)
         return Response({"data": documents})
@@ -485,7 +489,7 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response:
         user_id = custom_tool.created_by.user_id
         document: DocumentManager = DocumentManager.objects.get(pk=document_id)
         file_name: str = document.document_name
-        file_path = FileManagerHelper.handle_sub_directory_for_tenants(
+        file_path = PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
             org_id=org_id,
             is_create=False,
             user_id=user_id,

backend/utils/file_storage/common_utils.py

@@ -0,0 +1,17 @@
+import os
+
+
+class FileStorageUtil:
+    @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

backend/utils/file_storage/constants.py

@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class FileStorageKeys:
+    FILE_STORAGE_PROVIDER = "FILE_STORAGE_PROVIDER"
+    FILE_STORAGE_CREDENTIALS = "FILE_STORAGE_CREDENTIALS"
+
+
+class FileStorageType(Enum):
+    PERMANENT = "permanent"
+    TEMPORARY = "temporary"
+
+
+class FileStorageConstants:
+    PROMPT_STUDIO_FILE_PATH = "PROMPT_STUDIO_FILE_PATH"

backend/utils/file_storage/exceptions.py

@@ -0,0 +1,6 @@
+from rest_framework.exceptions import APIException
+
+
+class OrgIdNotValid(APIException):
+    status_code = 400
+    default_detail = "Organization ID is not valid"

backend/utils/file_storage/helpers/common_file_helper.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+from unstract.sdk.file_storage import FileStorageProvider
+from unstract.sdk.file_storage.fs_impl import FileStorage
+from unstract.sdk.file_storage.fs_permanent import PermanentFileStorage
+from unstract.sdk.file_storage.fs_shared_temporary import SharedTemporaryFileStorage
+from utils.file_storage.common_utils import FileStorageUtil
+from utils.file_storage.constants import FileStorageKeys, FileStorageType
+
+
+class FileStorageHelper:
+    # TODO : Optimize this to a singleton class
+    @staticmethod
+    def initialize_file_storage(type: FileStorageType) -> FileStorage:
+        provider_data = FileStorageHelper.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.TEMPORARY.value:
+            file_storage = SharedTemporaryFileStorage(
+                provider=provider, credentials=credentials
+            )
+        file_storage = FileStorage(
+            provider=FileStorageProvider.Local, credentials=credentials
+        )
+        return file_storage
+
+    @staticmethod
+    def load_file_storage_envs() -> dict[str, Any]:
+        provider: str = FileStorageUtil.get_env_or_die(
+            env_key=FileStorageKeys.FILE_STORAGE_PROVIDER
+        )
+        credentials: str = FileStorageUtil.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

backend/utils/file_storage/helpers/prompt_studio_file_helper.py

@@ -0,0 +1,96 @@
+from typing import Any, Union
+
+from file_management.exceptions import OrgIdNotValid
+from utils.file_storage.common_utils import FileStorageUtil
+from utils.file_storage.constants import FileStorageConstants, FileStorageType
+from utils.file_storage.helpers.common_file_helper import FileStorageHelper
+
+
+class PromptStudioFileHelper:
+    @staticmethod
+    def handle_sub_directory_for_prompt_studio(
+        org_id: str, user_id: str, tool_id: str, is_create: bool
+    ) -> str:
+        """Resolves a directory path meant for a user running prompt studio.
+
+        Args:
+            org_id (str): Organization ID
+            user_id (str): User ID
+            tool_id (str): ID of the prompt studio tool
+            is_create (bool): Flag to create the directory
+
+        Returns:
+            str: The absolute path to the directory meant for prompt studio
+        """
+        if not org_id:
+            raise OrgIdNotValid()
+        base_path = FileStorageUtil.get_env_or_die(
+            env_key=FileStorageConstants.PROMPT_STUDIO_FILE_PATH
+        )
+        file_path = f"{base_path}/{org_id}/{user_id}/{tool_id}"
+        extract_file_path = f"{file_path}/extract"
+        summarize_file_path = f"{file_path}/summarize"
+        if is_create:
+            fs_instance = FileStorageHelper.initialize_file_storage(
+                type=FileStorageType.PERMANENT
+            )
+            fs_instance.mkdir(file_path, create_parents=True)
+            fs_instance.mkdir(extract_file_path, create_parents=True)
+            fs_instance.mkdir(summarize_file_path, create_parents=True)
+        return str(file_path)
+
+    @staticmethod
+    def upload_for_ide(
+        org_id: str, user_id: str, tool_id: str, uploaded_file: Any
+    ) -> None:
+        fs_instance = FileStorageHelper.initialize_file_storage(
+            type=FileStorageType.PERMANENT
+        )
+        file_system_path = (
+            PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
+                org_id=org_id,
+                is_create=True,
+                user_id=user_id,
+                tool_id=str(tool_id),
+            )
+        )
+        file_path = f"{file_system_path}/{uploaded_file.name}"
+        fs_instance.write(path=file_path, mode="wb", data=uploaded_file.read())
+
+    @staticmethod
+    def fetch_file_contents(
+        org_id: str, user_id: str, tool_id: str, file_name: str
+    ) -> Union[bytes, str]:
+        fs_instance = FileStorageHelper.initialize_file_storage(
+            type=FileStorageType.PERMANENT
+        )
+        file_system_path = (
+            PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
+                org_id=org_id,
+                is_create=True,
+                user_id=user_id,
+                tool_id=str(tool_id),
+            )
+        )
+        # TODO : Handle this with proper fix
+        # Temporary Hack for frictionless onboarding as the user id will be empty
+        if not fs_instance.exists(file_system_path):
+            file_system_path = (
+                PromptStudioFileHelper.handle_sub_directory_for_prompt_studio(
+                    org_id=org_id,
+                    is_create=True,
+                    user_id="",
+                    tool_id=str(tool_id),
+                )
+            )
+        file_path = f"{file_system_path}/{file_name}"
+        file_content_type = fs_instance.mime_type(file_path)
+        text_content: Union[bytes, str]
+        if file_content_type == "application/pdf":
+            # Read contents of PDF file into a string
+            text_content = fs_instance.read(path=file_path, mode="rb")
+
+        elif file_content_type == "text/plain":
+            text_content = fs_instance.read(path=file_path, mode="r")
+
+        return text_content