Added new default configs for llm whisperer (#78)

* Added default config for add_line_nos and output_json

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

* Write metadata to file if output_file_path is provided

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

* Fixed metadata file path

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

* Added docstring

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

* Resolved review comment

---------

Signed-off-by: Deepak <[email protected]>
Signed-off-by: Deepak K <[email protected]>

Author: Deepak-Kesavan

src/unstract/sdk/adapters/x2text/llm_whisperer/src/constants.py

@@ -61,6 +61,8 @@ class WhispererConfig:
     HORIZONTAL_STRETCH_FACTOR = "horizontal_stretch_factor"
     PAGES_TO_EXTRACT = "pages_to_extract"
     STORE_METADATA_FOR_HIGHLIGHTING = "store_metadata_for_highlighting"
+    ADD_LINE_NOS = "add_line_nos"
+    OUTPUT_JSON = "output_json"
     PAGE_SEPARATOR = "page_seperator"
 
 
@@ -87,4 +89,6 @@ class WhispererDefaults:
     POLL_INTERVAL = int(os.getenv(WhispererEnv.POLL_INTERVAL, 30))
     MAX_POLLS = int(os.getenv(WhispererEnv.MAX_POLLS, 30))
     PAGES_TO_EXTRACT = ""
+    ADD_LINE_NOS = True
+    OUTPUT_JSON = True
     PAGE_SEPARATOR = "<<< >>>"

src/unstract/sdk/adapters/x2text/llm_whisperer/src/llm_whisperer.py

@@ -1,6 +1,8 @@
+import json
 import logging
 import os
 import time
+from pathlib import Path
 from typing import Any, Optional
 
 import requests
@@ -162,6 +164,8 @@ def _get_whisper_params(self, enable_highlight: bool = False) -> dict[str, Any]:
                 WhispererConfig.PAGES_TO_EXTRACT,
                 WhispererDefaults.PAGES_TO_EXTRACT,
             ),
+            WhispererConfig.ADD_LINE_NOS: WhispererDefaults.ADD_LINE_NOS,
+            WhispererConfig.OUTPUT_JSON: WhispererDefaults.OUTPUT_JSON,
             WhispererConfig.PAGE_SEPARATOR: self.config.get(
                 WhispererConfig.PAGE_SEPARATOR,
                 WhispererDefaults.PAGE_SEPARATOR,
@@ -264,7 +268,10 @@ def _extract_async(self, whisper_hash: str) -> str:
         logger.info(f"Extracting async for whisper hash: {whisper_hash}")
 
         headers: dict[str, Any] = self._get_request_headers()
-        params = {WhisperStatus.WHISPER_HASH: whisper_hash}
+        params = {
+            WhisperStatus.WHISPER_HASH: whisper_hash,
+            WhispererConfig.OUTPUT_JSON: WhispererDefaults.OUTPUT_JSON,
+        }
 
         # Polls in fixed intervals and checks status
         self._check_status_until_ready(
@@ -278,7 +285,7 @@ def _extract_async(self, whisper_hash: str) -> str:
             params=params,
         )
         if retrieve_response.status_code == 200:
-            return retrieve_response.content.decode("utf-8")
+            return retrieve_response.json()
         else:
             raise ExtractorError(
                 "Error retrieving from LLMWhisperer: "
@@ -310,25 +317,61 @@ def _send_whisper_request(
     def _extract_text_from_response(
         self, output_file_path: Optional[str], response: requests.Response
     ) -> str:
-
-        output = ""
+        output_json = {}
         if response.status_code == 200:
-            output = response.content.decode("utf-8")
+            output_json = response.json()
         elif response.status_code == 202:
             whisper_hash = response.json().get(WhisperStatus.WHISPER_HASH)
-            output = self._extract_async(whisper_hash=whisper_hash)
+            output_json = self._extract_async(whisper_hash=whisper_hash)
         else:
             raise ExtractorError("Couldn't extract text from file")
+        if output_file_path:
+            self._write_output_to_file(
+                output_json=output_json,
+                output_file_path=Path(output_file_path),
+            )
+        return output_json.get("text", "")
+
+    def _write_output_to_file(self, output_json: dict, output_file_path: Path) -> None:
+        """Writes the extracted text and metadata to the specified output file
+        and metadata file.
+
+        Args:
+            output_json (dict): The dictionary containing the extracted data,
+                with "text" as the key for the main content.
+            output_file_path (Path): The file path where the extracted text
+                should be written.
 
+        Raises:
+            ExtractorError: If there is an error while writing the output file.
+        """
         try:
-            # Write output to a file
-            if output_file_path:
-                with open(output_file_path, "w", encoding="utf-8") as f:
-                    f.write(output)
-        except OSError as e:
-            logger.error(f"OS error while writing {output_file_path}: {e} ")
+            text_output = output_json.get("text", "")
+            logger.info(f"Writing output to {output_file_path}")
+            output_file_path.write_text(text_output, encoding="utf-8")
+            try:
+                # Define the directory of the output file and metadata paths
+                output_dir = output_file_path.parent
+                metadata_dir = output_dir / "metadata"
+                metadata_file_name = output_file_path.with_suffix(".json").name
+                metadata_file_path = metadata_dir / metadata_file_name
+                # Ensure the metadata directory exists
+                metadata_dir.mkdir(parents=True, exist_ok=True)
+                # Remove the "text" key from the metadata
+                metadata = {
+                    key: value for key, value in output_json.items() if key != "text"
+                }
+                metadata_json = json.dumps(metadata, ensure_ascii=False, indent=4)
+                logger.info(f"Writing metadata to {metadata_file_path}")
+                metadata_file_path.write_text(metadata_json, encoding="utf-8")
+            except Exception as e:
+                logger.error(
+                    f"Error while writing metadata to {metadata_file_path}: {e}"
+                )
+
+        except Exception as e:
+            logger.error(f"Error while writing {output_file_path}: {e}")
             raise ExtractorError(str(e))
-        return output
 
     def process(
         self,