[FIX] Fixes for Prompt studio Indexing and tool runs

backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py

@@ -381,18 +381,36 @@ def index_document(
         process_text = None
         if text_processor:
             process_text = text_processor.process
-        doc_id = PromptStudioHelper.dynamic_indexer(
-            profile_manager=default_profile,
-            tool_id=tool_id,
-            file_path=file_path,
-            org_id=org_id,
-            document_id=document_id,
-            is_summary=is_summary,
-            reindex=True,
-            run_id=run_id,
-            user_id=user_id,
-            process_text=process_text,
-        )
+        if not check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+            doc_id = PromptStudioHelper.dynamic_indexer(
+                profile_manager=default_profile,
+                tool_id=tool_id,
+                file_path=file_path,
+                org_id=org_id,
+                document_id=document_id,
+                is_summary=is_summary,
+                reindex=True,
+                run_id=run_id,
+                user_id=user_id,
+                process_text=process_text,
+            )
+        else:
+            fs_instance = FileStorageHelper.initialize_file_storage(
+                type=FileStorageType.PERMANENT
+            )
+            doc_id = PromptStudioHelper.dynamic_indexer(
+                profile_manager=default_profile,
+                tool_id=tool_id,
+                file_path=file_path,
+                org_id=org_id,
+                document_id=document_id,
+                is_summary=is_summary,
+                reindex=True,
+                run_id=run_id,
+                user_id=user_id,
+                process_text=process_text,
+                fs=fs_instance,
+            )
 
         elapsed_time = time.time() - start_time
         logger.info(
@@ -870,7 +888,7 @@ def _fetch_response(
         if not check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
             file_hash = ToolUtils.get_hash_from_file(file_path=doc_path)
         else:
-            file_hash = ToolUtils.get_hash_from_file(file_path=doc_path, fs=fs_instance)
+            file_hash = fs_instance.get_hash_from_file(path=doc_path)
 
         payload = {
             TSPKeys.TOOL_SETTINGS: tool_settings,
@@ -1123,17 +1141,34 @@ def _fetch_single_pass_response(
         if not default_profile:
             raise DefaultProfileError()
 
-        index_result = PromptStudioHelper.dynamic_indexer(
-            profile_manager=default_profile,
-            file_path=file_path,
-            tool_id=tool_id,
-            org_id=org_id,
-            is_summary=tool.summarize_as_source,
-            document_id=document_id,
-            run_id=run_id,
-            user_id=user_id,
-            process_text=process_text,
-        )
+        if not check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+            index_result = PromptStudioHelper.dynamic_indexer(
+                profile_manager=default_profile,
+                file_path=file_path,
+                tool_id=tool_id,
+                org_id=org_id,
+                is_summary=tool.summarize_as_source,
+                document_id=document_id,
+                run_id=run_id,
+                user_id=user_id,
+                process_text=process_text,
+            )
+        else:
+            fs_instance = FileStorageHelper.initialize_file_storage(
+                type=FileStorageType.PERMANENT
+            )
+            index_result = PromptStudioHelper.dynamic_indexer(
+                profile_manager=default_profile,
+                file_path=file_path,
+                tool_id=tool_id,
+                org_id=org_id,
+                is_summary=tool.summarize_as_source,
+                document_id=document_id,
+                run_id=run_id,
+                user_id=user_id,
+                process_text=process_text,
+                fs=fs_instance,
+            )
         if index_result.get("status") == IndexingStatus.PENDING_STATUS.value:
             return {
                 "status": IndexingStatus.PENDING_STATUS.value,
@@ -1174,7 +1209,10 @@ def _fetch_single_pass_response(
         if tool.summarize_as_source:
             path = Path(file_path)
             file_path = str(path.parent / TSPKeys.SUMMARIZE / (path.stem + ".txt"))
-        file_hash = ToolUtils.get_hash_from_file(file_path=file_path)
+        if not check_feature_flag_status(FeatureFlag.REMOTE_FILE_STORAGE):
+            file_hash = ToolUtils.get_hash_from_file(file_path=file_path)
+        else:
+            file_hash = fs_instance.get_hash_from_file(path=file_path)
 
         payload = {
             TSPKeys.TOOL_SETTINGS: tool_settings,

backend/sample.env

@@ -160,12 +160,10 @@ API_EXECUTION_DIR_PREFIX="unstract/api"
 
 # Storage Provider for Workflow Execution
 # Valid options: MINIO, S3, etc..
-WORKFLOW_EXECUTION_FS_PROVIDER="MINIO"
-WORKFLOW_EXECUTION_FS_CREDENTIAL='{"endpoint_url": "", "key": "", "secret": ""}'
+WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS='{"provider":"minio","credentials": {"endpoint_url":"http://unstract-minio:9000","key":"XXX","secret":"XXX"}}'
 
 # Storage Provider for API Execution
-API_STORAGE_FS_PROVIDER="MINIO"
-API_STORAGE_FS_CREDENTIAL='{"endpoint_url": "", "key": "", "secret": ""}'
+API_FILE_STORAGE_CREDENTIALS='{"provider": "minio", "credentials": {"endpoint_url": "http://unstract-minio:9000", "key": "XXX", "secret": "XXX"}}'
 
 # Optional: Legacy storage path (if applicable)
 LEGACY_STORAGE_PATH="/path/to/legacy/storage"

tools/classifier/sample.env

@@ -11,5 +11,4 @@ X2TEXT_PORT=3004
 # (e.g., bucket/execution/org_id/workflow_id/execution_id)
 EXECUTION_DATA_DIR=<execution_dir_path_with_bucket>
 # Storage provider for Workflow Execution (e.g., minio, S3)
-WORKFLOW_EXECUTION_FS_PROVIDER="minio"
-WORKFLOW_EXECUTION_FS_CREDENTIAL={"endpoint_url":"http://localhost:9000","key":"","secret":""}
+WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS='{"provider":"minio","credentials"={"endpoint_url":"http://localhost:9000","key":"XXX","secret":"XXX"}}'

tools/classifier/src/helper.py

@@ -145,7 +145,7 @@ def _extract_from_adapter(self, file: str, adapter_id: str) -> Optional[str]:
 
         try:
             extraction_result: TextExtractionResult = x2text.process(
-                input_file_path=file
+                input_file_path=file, fs=self.tool.workflow_filestorage
             )
             extracted_text: str = extraction_result.extracted_text
             return extracted_text

tools/structure/sample.env

@@ -15,5 +15,4 @@ X2TEXT_PORT=3004
 # (e.g., bucket/execution/org_id/workflow_id/execution_id)
 EXECUTION_DATA_DIR=<execution_dir_path_with_bucket>
 # Storage provider for Workflow Execution (e.g., minio, S3)
-WORKFLOW_EXECUTION_FS_PROVIDER="minio"
-WORKFLOW_EXECUTION_FS_CREDENTIAL={"endpoint_url":"http://localhost:9000","key":"","secret":""}
+WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS='{"provider":"minio","credentials"={"endpoint_url":"http://localhost:9000","key":"","secret":""}}'

tools/structure/src/main.py

@@ -193,6 +193,11 @@ def run(
                             reindex=reindex,
                             usage_kwargs=usage_kwargs,
                             process_text=process_text,
+                            **(
+                                {"fs": self.workflow_filestorage}
+                                if self.workflow_filestorage is not None
+                                else {}
+                            ),
                         )
                         index_metrics[output[SettingsKeys.NAME]] = {
                             SettingsKeys.INDEXING: index.get_metrics()
@@ -379,9 +384,14 @@ def _summarize_and_index(
                     f.write(summarized_context)
 
         self.stream_log("Indexing summarized context")
-        summarize_file_hash: str = ToolUtils.get_hash_from_file(
-            file_path=summarize_file_path
-        )
+        if self.workflow_filestorage:
+            summarize_file_hash: str = self.workflow_filestorage.get_hash_from_file(
+                file_path=summarize_file_path
+            )
+        else:
+            summarize_file_hash: str = ToolUtils.get_hash_from_file(
+                file_path=summarize_file_path
+            )
         index.index(
             tool_id=tool_id,
             embedding_instance_id=embedding_instance_id,

tools/text_extractor/sample.env

@@ -12,5 +12,4 @@ X2TEXT_PORT=
 # (e.g., bucket/execution/org_id/workflow_id/execution_id)
 EXECUTION_DATA_DIR=<execution_dir_path_with_bucket>
 # Storage provider for Workflow Execution (e.g., minio, S3)
-WORKFLOW_EXECUTION_FS_PROVIDER="minio"
-WORKFLOW_EXECUTION_FS_CREDENTIAL={"endpoint_url":"http://localhost:9000","key":"","secret":""}
+WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS='{"provider":"minio","credentials"={"endpoint_url":"http://localhost:9000","key":"XXX","secret":"XXX"}}'

tools/text_extractor/src/main.py

@@ -58,9 +58,14 @@ def run(
             tool=self, adapter_instance_id=text_extraction_adapter_id
         )
         self.stream_log("Text extraction adapter has been created successfully.")
-        extraction_result: TextExtractionResult = text_extraction_adapter.process(
-            input_file_path=input_file
-        )
+        if self.workflow_filestorage:
+            extraction_result: TextExtractionResult = text_extraction_adapter.process(
+                input_file_path=input_file, fs=self.workflow_filestorage
+            )
+        else:
+            extraction_result: TextExtractionResult = text_extraction_adapter.process(
+                input_file_path=input_file
+            )
         extracted_text = self.convert_to_actual_string(extraction_result.extracted_text)
 
         self.stream_log("Text has been extracted successfully.")

unstract/filesystem/src/unstract/filesystem/file_storage_config.py

@@ -1,8 +1,7 @@
-import json
 import logging
 import os
 
-from unstract.sdk.file_storage import FileStorageProvider, SharedTemporaryFileStorage
+from unstract.sdk.file_storage import FileStorageProvider, StorageType
 
 from .exceptions import ProviderNotFound
 from .file_storage_types import FileStorageType  # Import the shared enum
@@ -29,22 +28,11 @@ def get_provider(var_name: str, default: str = "minio") -> FileStorageProvider:
 
 # Mappings for storage types, credentials, and providers
 STORAGE_MAPPING = {
-    FileStorageType.WORKFLOW_EXECUTION: SharedTemporaryFileStorage,
-    FileStorageType.API_EXECUTION: SharedTemporaryFileStorage,
+    FileStorageType.WORKFLOW_EXECUTION: StorageType.SHARED_TEMPORARY,
+    FileStorageType.API_EXECUTION: StorageType.SHARED_TEMPORARY,
 }
 
-FILE_STORAGE_CREDENTIALS_MAPPING = {
-    FileStorageType.WORKFLOW_EXECUTION: json.loads(
-        os.environ.get("WORKFLOW_EXECUTION_FS_CREDENTIAL", "{}")
-    ),
-    FileStorageType.API_EXECUTION: json.loads(
-        os.environ.get("API_STORAGE_FS_CREDENTIAL", "{}")
-    ),
-}
-
-FILE_STORAGE_PROVIDER_MAPPING = {
-    FileStorageType.WORKFLOW_EXECUTION: get_provider(
-        "WORKFLOW_EXECUTION_FS_PROVIDER", "minio"
-    ),
-    FileStorageType.API_EXECUTION: get_provider("API_STORAGE_FS_PROVIDER", "minio"),
+FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING = {
+    FileStorageType.WORKFLOW_EXECUTION: "WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS",
+    FileStorageType.API_EXECUTION: "API_FILE_STORAGE_CREDENTIALS",
 }

unstract/filesystem/src/unstract/filesystem/filesystem.py

@@ -1,18 +1,11 @@
 import logging
-import os
-from typing import Any
 
-from unstract.sdk.file_storage import (
-    FileStorage,
-    FileStorageProvider,
-    PermanentFileStorage,
-)
+from unstract.sdk.file_storage import EnvHelper, FileStorage
 
 from unstract.filesystem.file_storage_types import FileStorageType
 
 from .file_storage_config import (
-    FILE_STORAGE_CREDENTIALS_MAPPING,
-    FILE_STORAGE_PROVIDER_MAPPING,
+    FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING,
     STORAGE_MAPPING,
 )
 
@@ -23,31 +16,12 @@ class FileSystem:
     def __init__(self, file_storage_type: FileStorageType):
         self.file_storage_type = file_storage_type
 
-    def _get_credentials(self) -> dict[str, Any]:
-        """Retrieve storage credentials based on the storage type."""
-        return FILE_STORAGE_CREDENTIALS_MAPPING.get(self.file_storage_type, {})
-
-    def _get_provider(self) -> FileStorageProvider:
-        """Retrieve provider based on the storage type."""
-        provider_name = FILE_STORAGE_PROVIDER_MAPPING.get(self.file_storage_type)
-        if not provider_name:
-            raise ValueError(f"No provider found for {self.file_storage_type}.")
-        return provider_name
-
     def get_file_storage(self) -> FileStorage:
         """Initialize and return the appropriate file storage instance."""
         # Get the storage class based on the FileStorageType
-        storage_class = STORAGE_MAPPING.get(self.file_storage_type)
-
-        if not storage_class:
-            raise ValueError(f"Unsupported FileStorageType: {self.file_storage_type}")
-
-        # Base configuration for all storage classes
-        credentials = self._get_credentials()
-
-        # Add specific parameters based on the storage class type
-        if storage_class == PermanentFileStorage:
-            credentials["legacy_storage_path"] = os.getenv("LEGACY_STORAGE_PATH", None)
-
-        # Return an instance of the storage class with the appropriate configuration
-        return storage_class(provider=self._get_provider(), **credentials)
+        storage_type = STORAGE_MAPPING.get(self.file_storage_type)
+        env_name = FILE_STORAGE_CREDENTIALS_TO_ENV_NAME_MAPPING.get(
+            self.file_storage_type
+        )
+        file_storage = EnvHelper.get_storage(storage_type, env_name)
+        return file_storage

worker/src/unstract/worker/constants.py

@@ -30,8 +30,9 @@ class Env:
     LOG_LEVEL = "LOG_LEVEL"
     REMOVE_CONTAINER_ON_EXIT = "REMOVE_CONTAINER_ON_EXIT"
     WORKFLOW_EXECUTION_DIR_PREFIX = "WORKFLOW_EXECUTION_DIR_PREFIX"
-    WORKFLOW_EXECUTION_FS_PROVIDER = "WORKFLOW_EXECUTION_FS_PROVIDER"
-    WORKFLOW_EXECUTION_FS_CREDENTIAL = "WORKFLOW_EXECUTION_FS_CREDENTIAL"
+    WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS = (
+        "WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS"
+    )
     EXECUTION_DATA_DIR = "EXECUTION_DATA_DIR"
     FLIPT_SERVICE_AVAILABLE = "FLIPT_SERVICE_AVAILABLE"
 

worker/src/unstract/worker/worker.py

@@ -194,11 +194,8 @@ def run_container(
                 workflow_id,
                 execution_id,
             )
-            envs[Env.WORKFLOW_EXECUTION_FS_PROVIDER] = os.getenv(
-                Env.WORKFLOW_EXECUTION_FS_PROVIDER, ""
-            )
-            envs[Env.WORKFLOW_EXECUTION_FS_CREDENTIAL] = os.getenv(
-                Env.WORKFLOW_EXECUTION_FS_CREDENTIAL, "{}"
+            envs[Env.WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS] = os.getenv(
+                Env.WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS, "{}"
             )
         else:
             envs[Env.TOOL_DATA_DIR] = tool_data_dir