chore: extend logging (#96)

- Added logs to split page hook to improve transparency on the process.
- Changed log split hook member variable for clarity.
- Added retry counter to retry hook logs.
- Moved logging configuration to `__init__` from the `sdk_init` method
of a single hook and removed stdout as target for the logger. Logs now
go to stderr (the default).

Tested on `_sample_docs/layout-parser-paper-fast.pdf` and
`_sample_docs/layout-parser-paper.pdf`.
First one will log that it cannot be split if run with pdf split.

---------

Co-authored-by: Filip Knefel <[email protected]>
Co-authored-by: Marek Połom <[email protected]>

Author: 3 people

_test_unstructured_client/unit/test_custom_hooks.py

@@ -42,7 +42,8 @@ def test_unit_retry_with_backoff_does_retry(caplog):
         assert len(mock.request_history) > 1
 
 
-def test_unit_backoff_strategy_logs_retries(caplog):
+@pytest.mark.parametrize("status_code", [500, 503])
+def test_unit_backoff_strategy_logs_retries(status_code: int, caplog):
     caplog.set_level(logging.INFO)
     filename = "README.md"
     backoff_strategy = BackoffStrategy(
@@ -53,8 +54,8 @@ def test_unit_backoff_strategy_logs_retries(caplog):
     )
 
     with requests_mock.Mocker() as mock:
-        # mock a 500 status code for POST requests to the api
-        mock.post("https://api.unstructured.io/general/v0/general", status_code=500)
+        # mock a 500/503 status code for POST requests to the api
+        mock.post("https://api.unstructured.io/general/v0/general", status_code=status_code)
         session = UnstructuredClient(api_key_auth=FAKE_KEY)
 
         with open(filename, "rb") as f:
@@ -63,7 +64,8 @@ def test_unit_backoff_strategy_logs_retries(caplog):
         req = shared.PartitionParameters(files=files)
         with pytest.raises(Exception):
             session.general.partition(req, retries=retries)
-    pattern = re.compile("Response status code: 500. Sleeping before retry.")
+    pattern = re.compile(f"Failed to process a request due to API server error with status code {status_code}. "
+                        "Attempting retry number 1 after sleep.")
     assert bool(pattern.search(caplog.text))
 
 

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -43,18 +43,18 @@ async def example():
         pass
 
     hook.coroutines_to_execute[operation_id] = [example(), example()]
-    hook.api_responses[operation_id] = [
+    hook.api_successful_responses[operation_id] = [
         requests.Response(),
         requests.Response(),
     ]
 
     assert len(hook.coroutines_to_execute[operation_id]) == 2
-    assert len(hook.api_responses[operation_id]) == 2
+    assert len(hook.api_successful_responses[operation_id]) == 2
 
     hook._clear_operation(operation_id)
 
     assert hook.coroutines_to_execute.get(operation_id) is None
-    assert hook.api_responses.get(operation_id) is None
+    assert hook.api_successful_responses.get(operation_id) is None
 
 
 def test_unit_prepare_request_payload():

src/unstructured_client/_hooks/custom/__init__.py

@@ -1,4 +1,5 @@
 from .clean_server_url_hook import CleanServerUrlSDKInitHook
-from .log_retries_after_error_hook import LogRetriesAfterErrorHook
+from .logger_hook import LoggerHook
 from .suggest_defining_url import SuggestDefiningUrlIf401AfterErrorHook
 from .split_pdf_hook import SplitPdfHook
+import logging

src/unstructured_client/_hooks/custom/log_retries_after_error_hook.py

@@ -0,0 +1,58 @@
+import logging
+from typing import Optional, Tuple, Union, DefaultDict
+
+import requests
+
+from unstructured_client._hooks.custom.common import UNSTRUCTURED_CLIENT_LOGGER_NAME
+from unstructured_client._hooks.types import (
+    AfterSuccessContext,
+    AfterErrorContext,
+    AfterErrorHook,
+    SDKInitHook,
+)
+from collections import defaultdict
+
+logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
+
+
+class LoggerHook(AfterErrorHook, SDKInitHook):
+    """Hook providing custom logging"""
+
+    def __init__(self) -> None:
+        self.retries_counter: DefaultDict[str, int] = defaultdict(int)
+
+    def log_retries(self, response: Optional[requests.Response], operation_id: str):
+        """Log retries to give users visibility into requests."""
+
+        if response is not None and response.status_code // 100 == 5:
+            logger.info(
+                "Failed to process a request due to API server error with status code %d. "
+                "Attempting retry number %d after sleep.",
+                response.status_code,
+                self.retries_counter[operation_id],
+            )
+            if response.text:
+                logger.info("Server message - %s", response.text)
+
+    def sdk_init(
+        self, base_url: str, client: requests.Session
+    ) -> Tuple[str, requests.Session]:
+        logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+        return base_url, client
+
+    def after_success(
+        self, hook_ctx: AfterSuccessContext, response: requests.Response
+    ) -> Union[requests.Response, Exception]:
+        del self.retries_counter[hook_ctx.operation_id]
+        return response
+
+    def after_error(
+        self,
+        hook_ctx: AfterErrorContext,
+        response: Optional[requests.Response],
+        error: Optional[Exception],
+    ) -> Union[Tuple[Optional[requests.Response], Optional[Exception]], Exception]:
+        """Concrete implementation for AfterErrorHook."""
+        self.retries_counter[hook_ctx.operation_id] += 1
+        self.log_retries(response, hook_ctx.operation_id)
+        return response, error

src/unstructured_client/_hooks/custom/logger_hook.py

@@ -108,7 +108,9 @@ def call_api(
         return requests.Response()
 
 
-def prepare_request_headers(headers: CaseInsensitiveDict[str]) -> CaseInsensitiveDict[str]:
+def prepare_request_headers(
+    headers: CaseInsensitiveDict[str],
+) -> CaseInsensitiveDict[str]:
     """Prepare the request headers by removing the 'Content-Type' and 'Content-Length' headers.
 
     Args:
@@ -161,3 +163,16 @@ def create_response(response: requests.Response, elements: list) -> requests.Res
     response_copy.headers.update({"Content-Length": content_length})
     setattr(response_copy, "_content", content)
     return response_copy
+
+
+def log_after_split_response(status_code: int, split_number: int):
+    if status_code == 200:
+        logger.info(
+            "Successfully partitioned set #%d, elements added to the final result.",
+            split_number,
+        )
+    else:
+        logger.warning(
+            "Failed to partition set #%d, its elements will be omitted in the final result.",
+            split_number,
+        )