fix: Added root logger to tools (#123)

* File Storage interface and implementation (#112)

* contextSizeChanges

* contextSizeChanges

* Version roll and test folder check in

* Fix enum values

* Fix test cases, address review comments

* Address review comments

* Update pyproject.toml

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

* Address mypy issues

* Change class design and implementation

* Remove unused definitions

* Add cp() and function refactoring

* Check-in sample env

* Default value of dict changed to None

* Add size()

* Refctor for using FileStorage

* Refactor to use FileStorage

* Fix issues

* Add mim_type, download functions

* change comments

* Refactor het_hash_from_file

* Add return types

* Remove permanent file storage from sdk

* Fix SDK functional issues

* Support minio

* Test cases for Minio

* Bring file variants back to sdk

* Fix copy_on_write

* Add new test cases for uploadd/download

* Add new functions to support platform-service

* Change modififcation_time return type to datetime

* Refactor env pick-up logic

* Sample env

* contextSizeChanges

* Remove commented code and some improvisations

* contextSizeChanges

* Add right JSON formatted string

* Update src/unstract/sdk/file_storage/fs_permanent.py

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

* Address review comments

* Address review comments

* Update src/unstract/sdk/file_storage/fs_shared_temporary.py

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

* Refactor for change in enum value

* Add return type

---------

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

* Added root logger for tools

* Avoid adding multiple handlers to rootLogger to prevent duplicate logs

* Revert "File Storage interface and implementation (#112)"

This reverts commit ef00f81.

---------

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

Author: 3 people

src/unstract/sdk/prompt.py

@@ -146,29 +146,19 @@ def get_exported_tool(
         """
         platform_host = tool.get_env_or_die(ToolEnv.PLATFORM_HOST)
         platform_port = tool.get_env_or_die(ToolEnv.PLATFORM_PORT)
-
-        tool.stream_log("Connecting to DB and getting exported tool metadata")
         base_url = SdkHelper.get_platform_base_url(platform_host, platform_port)
         bearer_token = tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY)
-
         url = f"{base_url}/custom_tool_instance"
         query_params = {PromptStudioKeys.PROMPT_REGISTRY_ID: prompt_registry_id}
         headers = {"Authorization": f"Bearer {bearer_token}"}
         response = requests.get(url, headers=headers, params=query_params)
         if response.status_code == 200:
-            adapter_data: dict[str, Any] = response.json()
-            tool.stream_log(
-                "Successfully retrieved metadata for the exported "
-                f"tool: {prompt_registry_id}"
-            )
-            return adapter_data
-
+            return response.json()
         elif response.status_code == 404:
             tool.stream_error_and_exit(
-                f"Exported tool {prompt_registry_id} is not found"
+                f"Exported tool '{prompt_registry_id}' is not found"
             )
             return None
-
         else:
             tool.stream_error_and_exit(
                 f"Error while retrieving tool metadata "

src/unstract/sdk/tool/stream.py

@@ -1,12 +1,14 @@
 import datetime
 import json
+import logging
 import os
 from typing import Any
 
 from deprecated import deprecated
 
 from unstract.sdk.constants import Command, LogLevel, LogStage, ToolEnv
 from unstract.sdk.utils import ToolUtils
+from unstract.sdk.utils.common_utils import UNSTRACT_TO_PY_LOG_LEVEL
 
 
 class StreamMixin:
@@ -30,8 +32,35 @@ def __init__(self, log_level: LogLevel = LogLevel.INFO, **kwargs) -> None:
         self._exec_by_tool = ToolUtils.str_to_bool(
             os.environ.get(ToolEnv.EXECUTION_BY_TOOL, "False")
         )
+        if self.is_exec_by_tool:
+            self._configure_logger()
         super().__init__(**kwargs)
 
+    @property
+    def is_exec_by_tool(self):
+        """Flag to determine if SDK library is used in a tool's context.
+
+        Returns:
+            bool: True if SDK is used by a tool else False
+        """
+        return self._exec_by_tool
+
+    def _configure_logger(self) -> None:
+        """Helps configure the logger for the tool run."""
+        rootlogger = logging.getLogger("")
+        # Avoids adding multiple handlers
+        if rootlogger.hasHandlers():
+            return
+        handler = logging.StreamHandler()
+        handler.setLevel(level=UNSTRACT_TO_PY_LOG_LEVEL[self.log_level])
+        handler.setFormatter(
+            logging.Formatter(
+                "[%(asctime)s] %(levelname)s in %(module)s: %(message)s",
+            )
+        )
+        rootlogger.addHandler(handler)
+        rootlogger.setLevel(level=UNSTRACT_TO_PY_LOG_LEVEL[self.log_level])
+
     def stream_log(
         self,
         log: str,

src/unstract/sdk/utils/common_utils.py

@@ -23,6 +23,14 @@ def generate_uuid() -> str:
     logging.ERROR: LogLevel.ERROR,
 }
 
+# Mapping from Unstract log level to python counterpart
+UNSTRACT_TO_PY_LOG_LEVEL = {
+    LogLevel.DEBUG: logging.DEBUG,
+    LogLevel.INFO: logging.INFO,
+    LogLevel.WARN: logging.WARNING,
+    LogLevel.ERROR: logging.ERROR,
+}
+
 
 def log_elapsed(operation):
     """Adds an elapsed time log.