Fix/adding whisper hash in metadata (#64)

* Update metadat with whisper hash

* Update metadat with whisper hash

* Update metadata with whisper hash

* Removed unwanted comments

* trying editable install

* reverted the pytoml dependency

* Updated the upstream function

* Enabling condition in the upstream

* Updated the handling using data class

* reverting the PDM lock file

---------

Signed-off-by: Rahul Johny <[email protected]>
Co-authored-by: Gayathri <[email protected]>

Author: johnyrahul

src/unstract/sdk/index.py

@@ -12,6 +12,9 @@
 )
 from typing_extensions import deprecated
 from unstract.adapters.exceptions import AdapterError
+from unstract.adapters.x2text.constants import X2TextConstants
+from unstract.adapters.x2text.dto import TextExtractionResult
+from unstract.adapters.x2text.llm_whisperer.src import LLMWhisperer
 
 from unstract.sdk.adapters import ToolAdapter
 from unstract.sdk.constants import LogLevel
@@ -131,6 +134,7 @@ def index(
         reindex: bool = False,
         file_hash: Optional[str] = None,
         output_file_path: Optional[str] = None,
+        enable_highlight: bool = False,
         usage_kwargs: dict[Any, Any] = {},
     ) -> str:
         """Indexes an individual file using the passed arguments.
@@ -246,9 +250,29 @@ def index(
                     x2text = X2Text(
                         tool=self.tool, adapter_instance_id=x2text_instance_id
                     )
-                    extracted_text = x2text.process(
-                        input_file_path=file_path, output_file_path=output_file_path
-                    )
+                    if enable_highlight and isinstance(
+                        x2text._x2text_instance, LLMWhisperer
+                    ):
+                        process_response: TextExtractionResult = x2text.process(
+                            input_file_path=file_path,
+                            output_file_path=output_file_path,
+                            enable_highlight=enable_highlight,
+                        )
+                        whisper_hash_value = (
+                            process_response.extraction_metadata.whisper_hash
+                        )
+
+                        metadata = {X2TextConstants.WHISPER_HASH: whisper_hash_value}
+
+                        self.tool.update_exec_metadata(metadata)
+
+                    else:
+                        process_response: TextExtractionResult = x2text.process(
+                            input_file_path=file_path,
+                            output_file_path=output_file_path,
+                        )
+
+                    extracted_text = process_response.extracted_text
             except AdapterError as e:
                 # Wrapping AdapterErrors with SdkError
                 raise IndexingError(str(e)) from e

src/unstract/sdk/tool/base.py

@@ -56,9 +56,7 @@ def from_tool_args(cls, args: list[str]) -> "BaseTool":
         if parsed_args.command not in Command.static_commands():
             tool._exec_metadata = tool._get_exec_metadata()
             tool.workflow_id = tool._exec_metadata.get(MetadataKey.WORKFLOW_ID)
-            tool.execution_id = tool._exec_metadata.get(
-                MetadataKey.EXECUTION_ID
-            )
+            tool.execution_id = tool._exec_metadata.get(MetadataKey.EXECUTION_ID)
             tool.org_id = tool._exec_metadata.get(MetadataKey.ORG_ID)
         return tool
 
@@ -108,15 +106,11 @@ def _get_data_dir(self) -> Path:
             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:
+    def _get_file_from_data_dir(self, file_to_get: str, raise_err: bool = False) -> str:
         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"
-            )
+            self.stream_error_and_exit(f"{file_to_get} is missing in TOOL_DATA_DIR")
         return str(file_path)
 
     def get_source_file(self) -> str:
@@ -171,17 +165,11 @@ def _get_exec_metadata(self) -> dict[str, Any]:
             with open(metadata_path, encoding="utf-8") as f:
                 metadata_json = loads(f.read())
         except JSONDecodeError as e:
-            self.stream_error_and_exit(
-                f"JSON decode error for {metadata_path}: {e}"
-            )
+            self.stream_error_and_exit(f"JSON decode error for {metadata_path}: {e}")
         except FileNotFoundError:
-            self.stream_error_and_exit(
-                f"Metadata file not found at {metadata_path}"
-            )
+            self.stream_error_and_exit(f"Metadata file not found at {metadata_path}")
         except OSError as e:
-            self.stream_error_and_exit(
-                f"OS Error while opening {metadata_path}: {e}"
-            )
+            self.stream_error_and_exit(f"OS Error while opening {metadata_path}: {e}")
         return metadata_json
 
     def _write_exec_metadata(self, metadata: dict[str, Any]) -> None:
@@ -204,18 +192,12 @@ def _update_exec_metadata(self) -> None:
         tool_metadata = {
             MetadataKey.TOOL_NAME: self.properties[PropKey.FUNCTION_NAME],
             MetadataKey.ELAPSED_TIME: self.elapsed_time(),
-            MetadataKey.OUTPUT_TYPE: self.properties[PropKey.RESULT][
-                PropKey.TYPE
-            ],
+            MetadataKey.OUTPUT_TYPE: self.properties[PropKey.RESULT][PropKey.TYPE],
         }
         if MetadataKey.TOTAL_ELA_TIME not in self._exec_metadata:
-            self._exec_metadata[
-                MetadataKey.TOTAL_ELA_TIME
-            ] = self.elapsed_time()
+            self._exec_metadata[MetadataKey.TOTAL_ELA_TIME] = self.elapsed_time()
         else:
-            self._exec_metadata[
-                MetadataKey.TOTAL_ELA_TIME
-            ] += self.elapsed_time()
+            self._exec_metadata[MetadataKey.TOTAL_ELA_TIME] += self.elapsed_time()
 
         if MetadataKey.TOOL_META not in self._exec_metadata:
             self._exec_metadata[MetadataKey.TOOL_META] = [tool_metadata]
@@ -224,16 +206,35 @@ def _update_exec_metadata(self) -> None:
 
         self._write_exec_metadata(metadata=self._exec_metadata)
 
+    def update_exec_metadata(self, metadata: dict[str, Any]) -> None:
+        """Helps update the execution metadata with the provided metadata
+        dictionary.
+
+        This method iterates over the key-value pairs in the input metadata dictionary
+        and updates the internal `_exec_metadata` dictionary of the tool instance
+        accordingly. It then writes the updated metadata to the `METADATA.json`
+        file in the tool's data directory.
+
+        Args:
+            metadata (dict[str, Any]): A dictionary containing the metadata
+            key-value pairs to update in the execution metadata.
+
+        Returns:
+            None
+        """
+        for key, value in metadata.items():
+            self._exec_metadata[key] = value
+
+        self._write_exec_metadata(metadata=self._exec_metadata)
+
     def write_tool_result(self, data: Union[str, dict[str, Any]]) -> None:
         """Helps write contents of the tool result into TOOL_DATA_DIR.
 
         Args:
             data (Union[str, dict[str, Any]]): Data to be written
         """
         output_type = self.properties[PropKey.RESULT][PropKey.TYPE]
-        if output_type is PropKey.OutputType.JSON and not isinstance(
-            data, dict
-        ):
+        if output_type is PropKey.OutputType.JSON and not isinstance(data, dict):
             # TODO: Validate JSON type output with output schema as well
             self.stream_error_and_exit(
                 f"Expected result to have type {PropKey.OutputType.JSON} "

src/unstract/sdk/x2txt.py

@@ -5,6 +5,7 @@
 from unstract.adapters.constants import Common
 from unstract.adapters.x2text import adapters
 from unstract.adapters.x2text.constants import X2TextConstants
+from unstract.adapters.x2text.dto import TextExtractionResult
 from unstract.adapters.x2text.x2text_adapter import X2TextAdapter
 
 from unstract.sdk.adapters import ToolAdapter
@@ -67,7 +68,7 @@ def process(
         input_file_path: str,
         output_file_path: Optional[str] = None,
         **kwargs: dict[Any, Any],
-    ) -> str:
+    ) -> TextExtractionResult:
         return self._x2text_instance.process(
             input_file_path, output_file_path, **kwargs
         )