Remove feature flag related code changes for remote storage (#157)

* Remove feature flag related code changes

* Fix summary issue

* Read entire length for mime type detection

* Read entire length for mime type detection

* Add MINIO as permanent file storage proivder

* SDK version

* Fix SDK for OSS use cases

* Promote to GA

---------

Signed-off-by: Gayathri <[email protected]>

Author: gaya3-zipstack

src/unstract/sdk/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.57.0rc5"
+__version__ = "0.58.0"
 
 
 def get_sdk_version():

src/unstract/sdk/adapters/x2text/llm_whisperer_v2/src/helper.py

@@ -225,7 +225,7 @@ def send_whisper_request(
 
         response: requests.Response
         try:
-            input_file_data = BytesIO(fs.read(input_file_path, "rb"))
+            input_file_data = BytesIO(fs.read(path=input_file_path, mode="rb"))
             response = LLMWhispererHelper.make_request(
                 config=config,
                 params=params,
@@ -276,7 +276,10 @@ def write_output_to_file(
             text_output = output_json.get("result_text", "")
             logger.info(f"Writing output to {output_file_path}")
             fs.write(
-                path=str(output_file_path), mode="w", data=text_output, encoding="utf-8"
+                path=str(output_file_path),
+                mode="w",
+                data=text_output,
+                encoding="utf-8",
             )
         except Exception as e:
             logger.error(f"Error while writing {output_file_path}: {e}")

src/unstract/sdk/constants.py

@@ -11,15 +11,13 @@ class ToolEnv:
         PLATFORM_API_KEY (str): Platform service API key.
         PLATFORM_HOST (str): Platform service host.
         PLATFORM_PORT (str): Platform service port.
-        DATA_DIR (str): The environment variable for the tool data directory.
         EXECUTION_BY_TOOL (str): Implicitly set to 1 by the SDK if its executed
             by a tool.
     """
 
     PLATFORM_API_KEY = "PLATFORM_SERVICE_API_KEY"
     PLATFORM_HOST = "PLATFORM_SERVICE_HOST"
     PLATFORM_PORT = "PLATFORM_SERVICE_PORT"
-    DATA_DIR = "TOOL_DATA_DIR"
     EXECUTION_BY_TOOL = "EXECUTION_BY_TOOL"
     EXECUTION_DATA_DIR = "EXECUTION_DATA_DIR"
     WORKFLOW_EXECUTION_FILE_STORAGE_CREDENTIALS = (

src/unstract/sdk/file_storage/constants.py

@@ -3,7 +3,6 @@
 
 class FileOperationParams:
     READ_ENTIRE_LENGTH = -1
-    MIME_TYPE_DEFAULT_READ_LENGTH = 100
     EXTENSION_DEFAULT_READ_LENGTH = 100
     DEFAULT_ENCODING = "utf-8"
 

src/unstract/sdk/file_storage/impl.py

@@ -41,7 +41,7 @@ def read(
 
         Args:
             path (str): Path to the file to be opened
-            mode (str): Mode in whicg the file is to be opened. Usual options
+            mode (str): Mode in which the file is to be opened. Usual options
                         include r, rb, w and wb
             encoding (str): Encoding type like utf-8 or utf-16
             seek_position (int): Position to start reading from
@@ -212,19 +212,18 @@ def modification_time(self, path: str) -> datetime:
             file_mtime = datetime.fromtimestamp(file_mtime)
         return file_mtime
 
-    @skip_local_cache
     def mime_type(
         self,
         path: str,
-        read_length: int = FileOperationParams.MIME_TYPE_DEFAULT_READ_LENGTH,
+        read_length: int = FileOperationParams.READ_ENTIRE_LENGTH,
     ) -> str:
         """Gets the file MIME type for an input file. Uses libmagic to perform
         the same.
 
         Args:
             path (str): Path of the input file
             read_length (int): Length(bytes) to be read from the file for in
-            order to identify the mime type
+            order to identify the mime type. Defaults to read the entire length.
 
         Returns:
             str: MIME type of the file
@@ -364,7 +363,6 @@ def yaml_load(
             data: dict[str, Any] = yaml.safe_load(f)
             return data
 
-    @skip_local_cache
     def guess_extension(self, path: str) -> str:
         """Returns the extension of the file passed.
 

src/unstract/sdk/file_storage/interface.py

@@ -77,7 +77,7 @@ def modification_time(self, path: str) -> datetime:
     def mime_type(
         self,
         path: str,
-        read_length: int = FileOperationParams.MIME_TYPE_DEFAULT_READ_LENGTH,
+        read_length: int = FileOperationParams.READ_ENTIRE_LENGTH,
     ) -> str:
         pass
 

src/unstract/sdk/file_storage/permanent.py

@@ -1,6 +1,9 @@
 import logging
 from typing import Any, Optional, Union
 
+import filetype
+import magic
+
 from unstract.sdk.exceptions import FileOperationError, FileStorageError
 from unstract.sdk.file_storage.constants import FileOperationParams
 from unstract.sdk.file_storage.impl import FileStorage
@@ -13,6 +16,7 @@ class PermanentFileStorage(FileStorage):
     SUPPORTED_FILE_STORAGE_TYPES = [
         FileStorageProvider.GCS.value,
         FileStorageProvider.S3.value,
+        FileStorageProvider.MINIO.value,
         FileStorageProvider.LOCAL.value,
     ]
 
@@ -27,7 +31,11 @@ def __init__(
                 f"supported in Permanent mode. "
                 f"Supported providers: {self.SUPPORTED_FILE_STORAGE_TYPES}"
             )
-        if provider == FileStorageProvider.GCS or provider == FileStorageProvider.LOCAL:
+        if (
+            provider == FileStorageProvider.GCS
+            or provider == FileStorageProvider.LOCAL
+            or provider == FileStorageProvider.MINIO
+        ):
             super().__init__(provider, **storage_config)
         else:
             raise NotImplementedError(f"Provider {provider.value} is not implemented")
@@ -53,6 +61,11 @@ def _copy_on_read(self, path: str, legacy_storage_path: str):
             # to remote storage
             if local_file_storage.exists(local_file_path):
                 self.upload(local_file_path, path)
+                logger.info(
+                    f"Uploading {local_file_path} from "
+                    f"{local_file_storage.provider} to remote "
+                    f"storage {self.provider} in the path {path}"
+                )
 
     def read(
         self,
@@ -87,3 +100,54 @@ def read(
             if isinstance(e, FileNotFoundError) or isinstance(e, FileOperationError):
                 raise e
             raise FileOperationError(str(e)) from e
+
+    def mime_type(
+        self,
+        path: str,
+        read_length: int = FileOperationParams.READ_ENTIRE_LENGTH,
+        legacy_storage_path: Optional[str] = None,
+    ) -> str:
+        """Gets the file MIME type for an input file. Uses libmagic to perform
+        the same.
+
+        Args:
+            path (str): Path of the input file
+            read_length (int): Length(bytes) to be read from the file for in
+            order to identify the mime type. Defaults to read the entire length.
+            legacy_storage_path (str):  Legacy path to the same file
+
+        Returns:
+            str: MIME type of the file
+        """
+        sample_contents = self.read(
+            path=path,
+            mode="rb",
+            length=read_length,
+            legacy_storage_path=legacy_storage_path,
+        )
+        mime_type = magic.from_buffer(sample_contents, mime=True)
+        return mime_type
+
+    def guess_extension(
+        self, path: str, legacy_storage_path: Optional[str] = None
+    ) -> str:
+        """Returns the extension of the file passed.
+
+        Args:
+            path (str): String holding the path
+            legacy_storage_path (str):  Legacy path to the same file
+
+        Returns:
+            str: File extension
+        """
+        file_extension = ""
+        sample_contents = self.read(
+            path=path,
+            mode="rb",
+            length=FileOperationParams.EXTENSION_DEFAULT_READ_LENGTH,
+            legacy_storage_path=legacy_storage_path,
+        )
+        if sample_contents:
+            file_type = filetype.guess(sample_contents)
+            file_extension = file_type.EXTENSION
+        return file_extension

src/unstract/sdk/index.py

@@ -297,8 +297,7 @@ def index(
 
             if doc_id_found and not reindex:
                 self.tool.stream_log(f"File was indexed already under {doc_id}")
-
-                if not fs.exists(output_file_path):
+                if output_file_path and not fs.exists(output_file_path):
                     # Added this as a workaround to handle extraction
                     # for documents uploaded twice in different projects.
                     # to be reconsidered after permanent fixes.

src/unstract/sdk/tool/base.py

@@ -1,8 +1,7 @@
 import datetime
-import json
 import os
 from abc import ABC, abstractmethod
-from json import JSONDecodeError, loads
+from json import JSONDecodeError
 from pathlib import Path
 from typing import Any, Union
 
@@ -127,34 +126,14 @@ def handle_static_command(self, command: str) -> None:
         else:
             raise ValueError(f"Unknown command {command}")
 
-    def _get_data_dir(self) -> Path:
-        """Gets the TOOL_DATA_DIR that houses the input and output files of
-        tool execution.
-
-        Returns:
-            Path: Path object of the TOOL_DATA_DIR that's configured.
-        """
-        data_dir = self.get_env_or_die(ToolEnv.DATA_DIR)
-        base_path = Path(data_dir)
-        if not base_path.exists():
-            self.stream_error_and_exit(f"{data_dir} does not exist")
-        if not base_path.is_dir():
-            self.stream_error_and_exit(f"{data_dir} is not a directory")
-        return base_path.absolute()
-
     def _get_file_from_data_dir(self, file_to_get: str, raise_err: bool = False) -> str:
-        if self.workflow_filestorage:
-            base_path = self.execution_dir
-            file_path = base_path / file_to_get
-            if raise_err and not self.workflow_filestorage.exists(path=file_path):
-                self.stream_error_and_exit(
-                    f"{file_to_get} is missing in EXECUTION_DATA_DIR"
-                )
-        else:
-            base_path: Path = self._get_data_dir()
-            file_path = base_path / file_to_get
-            if raise_err and not file_path.exists():
-                self.stream_error_and_exit(f"{file_to_get} is missing in TOOL_DATA_DIR")
+        base_path = self.execution_dir
+        file_path = base_path / file_to_get
+        if raise_err and not self.workflow_filestorage.exists(path=file_path):
+            self.stream_error_and_exit(
+                f"{file_to_get} is missing in EXECUTION_DATA_DIR"
+            )
+
         return str(file_path)
 
     def get_source_file(self) -> str:
@@ -183,10 +162,7 @@ def get_output_dir(self) -> str:
         Returns:
             str: Absolute path to the output directory.
         """
-        if self.workflow_filestorage:
-            base_path = self.execution_dir
-        else:
-            base_path: Path = self._get_data_dir()
+        base_path = self.execution_dir
         return str(base_path / ToolExecKey.OUTPUT_DIR)
 
     @property
@@ -205,20 +181,13 @@ def _get_exec_metadata(self) -> dict[str, Any]:
         Returns:
             dict[str, Any]: Contents of METADATA.json
         """
-        if self.workflow_filestorage:
-            base_path = self.execution_dir
-        else:
-            base_path: Path = self._get_data_dir()
+        base_path = self.execution_dir
         metadata_path = base_path / ToolExecKey.METADATA_FILE
         metadata_json = {}
         try:
-            if self.workflow_filestorage:
-                metadata_json = ToolUtils.load_json(
-                    file_to_load=metadata_path, fs=self.workflow_filestorage
-                )
-            else:
-                with open(metadata_path, encoding="utf-8") as f:
-                    metadata_json = loads(f.read())
+            metadata_json = ToolUtils.load_json(
+                file_to_load=metadata_path, fs=self.workflow_filestorage
+            )
         except JSONDecodeError as e:
             self.stream_error_and_exit(f"JSON decode error for {metadata_path}: {e}")
         except FileNotFoundError:
@@ -233,19 +202,13 @@ def _write_exec_metadata(self, metadata: dict[str, Any]) -> None:
         Args:
             metadata (dict[str, Any]): Metadata to write
         """
-        if self.workflow_filestorage:
-            base_path = self.execution_dir
-            metadata_path = base_path / ToolExecKey.METADATA_FILE
-            ToolUtils.dump_json(
-                file_to_dump=metadata_path,
-                json_to_dump=metadata,
-                fs=self.workflow_filestorage,
-            )
-        else:
-            base_path: Path = self._get_data_dir()
-            metadata_path = base_path / ToolExecKey.METADATA_FILE
-            with metadata_path.open("w", encoding="utf-8") as f:
-                f.write(ToolUtils.json_to_str(metadata))
+        base_path = self.execution_dir
+        metadata_path = base_path / ToolExecKey.METADATA_FILE
+        ToolUtils.dump_json(
+            file_to_dump=metadata_path,
+            json_to_dump=metadata,
+            fs=self.workflow_filestorage,
+        )
 
     def _update_exec_metadata(self) -> None:
         """Updates the execution metadata after a tool executes.
@@ -314,18 +277,13 @@ def write_tool_result(self, data: Union[str, dict[str, Any]]) -> None:
         self.stream_result(result)
 
         self._update_exec_metadata()
-        json_data = json.dumps(data)
         # INFILE is overwritten for next tool to run
         input_file_path: Path = Path(self.get_input_file())
-        if self.workflow_filestorage:
-            ToolUtils.dump_json(
-                file_to_dump=input_file_path,
-                json_to_dump=data,
-                fs=self.workflow_filestorage,
-            )
-        else:
-            with input_file_path.open("w", encoding="utf-8") as f:
-                f.write(json_data)
+        ToolUtils.dump_json(
+            file_to_dump=input_file_path,
+            json_to_dump=data,
+            fs=self.workflow_filestorage,
+        )
 
     def validate(self, input_file: str, settings: dict[str, Any]) -> None:
         """Override to implement custom validation for the tool.

src/unstract/sdk/tool/executor.py

@@ -57,8 +57,6 @@ def execute_run(self, args: argparse.Namespace) -> None:
             f"Execution ID: {self.tool.execution_id}, "
             f"SDK Version: {get_sdk_version()}"
         )
-        if not self.tool.workflow_filestorage:
-            self._setup_for_run()
         validator = ToolValidator(self.tool)
         settings = validator.validate_pre_execution(settings=settings)