passing tags to x2text adapter in kwargs (#150)

* passing tags to x2text adapter in kwargs

* re difined the DTO for params

* Update src/unstract/sdk/adapters/x2text/llm_whisperer_v2/src/dto.py

Signed-off-by: Chandrasekharan M <[email protected]>

---------

Signed-off-by: Chandrasekharan M <[email protected]>
Co-authored-by: Chandrasekharan M <[email protected]>

Author: muhammad-ali-e

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

@@ -3,6 +3,7 @@ class X2TextConstants:
     X2TEXT_HOST = "X2TEXT_HOST"
     X2TEXT_PORT = "X2TEXT_PORT"
     ENABLE_HIGHLIGHT = "enable_highlight"
+    TAGS = "tags"
     EXTRACTED_TEXT = "extracted_text"
     WHISPER_HASH = "whisper-hash"
     WHISPER_HASH_V2 = "whisper_hash"

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

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class WhispererRequestParams:
+    """DTO for LLM Whisperer API request parameters.
+
+    Args:
+        tag: Tag value. Can be initialized with List[str], str, or None.
+             Will be converted to str | None after initialization.
+    """
+
+    # TODO: Extend this DTO to include all Whisperer API parameters
+    tag: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        # TODO: Allow list of tags once its supported in LLMW v2
+        if isinstance(self.tag, list):
+            self.tag = self.tag[0] if self.tag else None

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

@@ -20,6 +20,7 @@
     WhispererHeader,
     WhisperStatus,
 )
+from unstract.sdk.adapters.x2text.llm_whisperer_v2.src.dto import WhispererRequestParams
 from unstract.sdk.constants import MimeType
 from unstract.sdk.file_storage import FileStorage, FileStorageProvider
 
@@ -108,7 +109,9 @@ def make_request(
         return response
 
     @staticmethod
-    def get_whisperer_params(config: dict[str, Any]) -> dict[str, Any]:
+    def get_whisperer_params(
+        config: dict[str, Any], extra_params: WhispererRequestParams
+    ) -> dict[str, Any]:
         """Gets query params meant for /whisper endpoint.
 
         The params is filled based on the configuration passed.
@@ -152,7 +155,8 @@ def get_whisperer_params(config: dict[str, Any]) -> dict[str, Any]:
             ),
             # Not providing default value to maintain legacy compatablity
             # these are optional params and identifiers for audit
-            WhispererConfig.TAG: config.get(
+            WhispererConfig.TAG: extra_params.tag
+            or config.get(
                 WhispererConfig.TAG,
                 WhispererDefaults.TAG,
             ),
@@ -292,11 +296,14 @@ def extract_async(config: dict[str, Any], whisper_hash: str) -> dict[Any, Any]:
     def send_whisper_request(
         input_file_path: str,
         config: dict[str, Any],
+        extra_params: WhispererRequestParams,
         fs: FileStorage = FileStorage(provider=FileStorageProvider.LOCAL),
     ) -> requests.Response:
         headers = LLMWhispererHelper.get_request_headers(config)
         headers["Content-Type"] = "application/octet-stream"
-        params = LLMWhispererHelper.get_whisperer_params(config)
+        params = LLMWhispererHelper.get_whisperer_params(
+            config=config, extra_params=extra_params
+        )
 
         response: requests.Response
         try:

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

@@ -14,6 +14,7 @@
     HTTPMethod,
     WhispererEndpoint,
 )
+from unstract.sdk.adapters.x2text.llm_whisperer_v2.src.dto import WhispererRequestParams
 from unstract.sdk.adapters.x2text.llm_whisperer_v2.src.helper import LLMWhispererHelper
 from unstract.sdk.adapters.x2text.x2text_adapter import X2TextAdapter
 from unstract.sdk.file_storage import FileStorage, FileStorageProvider
@@ -76,8 +77,12 @@ def process(
             str: Extracted text
         """
 
+        extra_params = WhispererRequestParams(tag=kwargs.get(X2TextConstants.TAGS))
         response: requests.Response = LLMWhispererHelper.send_whisper_request(
-            input_file_path, self.config, fs=fs
+            input_file_path=input_file_path,
+            config=self.config,
+            fs=fs,
+            extra_params=extra_params,
         )
         response_text = response.text
         reponse_dict = json.loads(response_text)

src/unstract/sdk/constants.py

@@ -128,6 +128,7 @@ class MetadataKey:
     WORKFLOW_ID = "workflow_id"
     EXECUTION_ID = "execution_id"
     FILE_EXECUTION_ID = "file_execution_id"
+    TAGS = "tags"
     ORG_ID = "organization_id"
     TOOL_META = "tool_metadata"
     TOOL_NAME = "tool_name"

src/unstract/sdk/index.py

@@ -128,6 +128,7 @@ def extract_text(
         usage_kwargs: dict[Any, Any] = {},
         process_text: Optional[Callable[[str], str]] = None,
         fs: FileStorage = FileStorage(FileStorageProvider.LOCAL),
+        tags: Optional[list[str]] = None,
     ) -> str:
         """Extracts text from a document.
 
@@ -147,6 +148,7 @@ def extract_text(
                 Defaults to {}.
             process_text (Optional[Callable[[str], str]], optional): Optional function
                 to post-process the text. Defaults to None.
+            tags: (Optional[list[str]], optional): Tags
 
         Raises:
             IndexingError: Errors during text extraction
@@ -164,14 +166,18 @@ def extract_text(
                     input_file_path=file_path,
                     output_file_path=output_file_path,
                     enable_highlight=enable_highlight,
+                    tags=tags,
                     fs=fs,
                 )
                 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, fs=fs
+                    input_file_path=file_path,
+                    output_file_path=output_file_path,
+                    tags=tags,
+                    fs=fs,
                 )
             extracted_text = process_response.extracted_text
         # TODO: Handle prepend of context where error is raised and remove this
@@ -193,6 +199,7 @@ def extract_text(
                 )
         return extracted_text
 
+    # TODO: Reduce the number of params by some dataclass
     @log_elapsed(operation="CHECK_AND_INDEX(overall)")
     @capture_metrics
     def index(
@@ -211,6 +218,7 @@ def index(
         usage_kwargs: dict[Any, Any] = {},
         process_text: Optional[Callable[[str], str]] = None,
         fs: FileStorage = FileStorage(provider=FileStorageProvider.LOCAL),
+        tags: Optional[list[str]] = None,
     ) -> str:
         """Indexes an individual file using the passed arguments.
 
@@ -231,6 +239,8 @@ def index(
             output_file_path (Optional[str], optional): File path to write
                 the extracted contents into. Defaults to None.
             fs (FileStorage): file storage object to perfrom file operations
+            tags (Optional[list[str]], optional): List of tags to be associated with
+                the indexed document.
 
         Returns:
             str: A unique ID for the file and indexing arguments combination
@@ -300,6 +310,7 @@ def index(
                         usage_kwargs=usage_kwargs,
                         process_text=process_text,
                         fs=fs,
+                        tags=tags,
                     )
                 return doc_id
 
@@ -310,6 +321,7 @@ def index(
                 enable_highlight=enable_highlight,
                 usage_kwargs=usage_kwargs,
                 process_text=process_text,
+                tags=tags,
                 fs=fs,
             )
             if not extracted_text:

src/unstract/sdk/tool/base.py

@@ -40,6 +40,7 @@ def __init__(self, log_level: LogLevel = LogLevel.INFO) -> None:
         self.workflow_id = ""
         self.execution_id = ""
         self.file_execution_id = ""
+        self.tags = []
         self.source_file_name = ""
         self.org_id = ""
         self._exec_metadata = {}
@@ -90,6 +91,7 @@ def from_tool_args(cls, args: list[str]) -> "BaseTool":
             tool.file_execution_id = tool._exec_metadata.get(
                 MetadataKey.FILE_EXECUTION_ID, ""
             )
+            tool.tags = tool._exec_metadata.get(MetadataKey.TAGS, [])
             tool.source_file_name = tool._exec_metadata.get(MetadataKey.SOURCE_NAME, "")
             tool.org_id = tool._exec_metadata.get(MetadataKey.ORG_ID)
         return tool