fix: remove false success log & improve retry log (#111)

In the PR introducing more logging two things turned out lacking:
- there was a false success log in split-page after error hook
- the retry logs were only for 5XX responses where speakeasy also has
retry set for ConnectionError exceptions

This PR removes the faulty success log, moves success/failure logs to the `LoggerHook` and introduces logging for `ConnectionError` retries.

---------

Co-authored-by: Filip Knefel <[email protected]>

Author: ds-filipknefel

_test_unstructured_client/unit/test_custom_hooks.py

@@ -43,7 +43,7 @@ def test_unit_retry_with_backoff_does_retry(caplog):
 
 
 @pytest.mark.parametrize("status_code", [500, 503])
-def test_unit_backoff_strategy_logs_retries(status_code: int, caplog):
+def test_unit_backoff_strategy_logs_retries_5XX(status_code: int, caplog):
     caplog.set_level(logging.INFO)
     filename = "README.md"
     backoff_strategy = BackoffStrategy(
@@ -64,11 +64,38 @@ def test_unit_backoff_strategy_logs_retries(status_code: int, caplog):
         req = shared.PartitionParameters(files=files)
         with pytest.raises(Exception):
             session.general.partition(req, retries=retries)
+
     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))
 
 
+def test_unit_backoff_strategy_logs_retries_connection_error(caplog):
+    caplog.set_level(logging.INFO)
+    filename = "README.md"
+    backoff_strategy = BackoffStrategy(
+        initial_interval=10, max_interval=100, exponent=1.5, max_elapsed_time=300
+    )
+    retries = RetryConfig(
+        strategy="backoff", backoff=backoff_strategy, retry_connection_errors=True
+    )
+    with requests_mock.Mocker() as mock:
+        # mock a connection error response to POST request
+        mock.post("https://api.unstructured.io/general/v0/general", exc=requests.exceptions.ConnectionError)
+        session = UnstructuredClient(api_key_auth=FAKE_KEY)
+
+        with open(filename, "rb") as f:
+            files = shared.Files(content=f.read(), file_name=filename)
+
+        req = shared.PartitionParameters(files=files)
+        with pytest.raises(Exception):
+            session.general.partition(req, retries=retries)
+
+    pattern = re.compile(f"Failed to process a request due to connection error .*? "
+                         "Attempting retry number 1 after sleep.")
+    assert bool(pattern.search(caplog.text))
+    
+
 @pytest.mark.parametrize(
     "server_url",
     [

src/unstructured_client/_hooks/custom/logger_hook.py

@@ -9,19 +9,20 @@
     AfterErrorContext,
     AfterErrorHook,
     SDKInitHook,
+    AfterSuccessHook,
 )
 from collections import defaultdict
 
 logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
 
 
-class LoggerHook(AfterErrorHook, SDKInitHook):
+class LoggerHook(AfterErrorHook, AfterSuccessHook, 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):
+    def log_retries(self, response: Optional[requests.Response],  error: Optional[Exception], operation_id: str,):
         """Log retries to give users visibility into requests."""
 
         if response is not None and response.status_code // 100 == 5:
@@ -33,6 +34,15 @@ def log_retries(self, response: Optional[requests.Response], operation_id: str):
             )
             if response.text:
                 logger.info("Server message - %s", response.text)
+        
+        elif error is not None and isinstance(error, requests.exceptions.ConnectionError):
+            logger.info(
+                "Failed to process a request due to connection error - %s. "
+                "Attempting retry number %d after sleep.",
+                error,
+                self.retries_counter[operation_id],
+            )
+
 
     def sdk_init(
         self, base_url: str, client: requests.Session
@@ -43,7 +53,9 @@ def sdk_init(
     def after_success(
         self, hook_ctx: AfterSuccessContext, response: requests.Response
     ) -> Union[requests.Response, Exception]:
-        del self.retries_counter[hook_ctx.operation_id]
+        self.retries_counter.pop(hook_ctx.operation_id, None)
+        # NOTE: In case of split page partition this means - at least one of the splits was partitioned successfully
+        logger.info("Successfully partitioned the document.")
         return response
 
     def after_error(
@@ -54,5 +66,18 @@ def after_error(
     ) -> 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)
+        self.log_retries(response, error, hook_ctx.operation_id)
+
+        if response and response.status_code == 200:
+            # NOTE: Even though this is an after_error method, due to split_pdf_hook logic we may get
+            # a success here when one of the split requests was partitioned successfully
+            logger.info("Successfully partitioned the document.")
+        
+        else:
+            logger.error("Failed to partition the document.")
+            if response:
+                logger.error("Server responded with %d - %s", response.status_code, response.text)
+            if error is not None:
+                logger.error("Following error occurred - %s", error)
+        
         return response, error

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -302,7 +302,6 @@ def after_success(
 
         updated_response = request_utils.create_response(response, elements)
         self._clear_operation(operation_id)
-        logger.info("Successfully processed the request.")
         return updated_response
 
     def after_error(
@@ -335,21 +334,16 @@ def after_error(
         successful_responses = self.api_successful_responses.get(operation_id)
 
         if elements is None or successful_responses is None:
-            logger.info("Successfully processed the request.")
             return (response, error)
 
         if len(successful_responses) == 0:
-            logger.error("Failed to process the request.")
-            if error is not None:
-                logger.error(error)
             self._clear_operation(operation_id)
             return (response, error)
 
         updated_response = request_utils.create_response(
             successful_responses[0], elements
         )
         self._clear_operation(operation_id)
-        logger.info("Successfully processed the request.")
         return (updated_response, None)
 
     def _clear_operation(self, operation_id: str) -> None:

src/unstructured_client/_hooks/registration.py

@@ -27,6 +27,9 @@ def init_hooks(hooks: Hooks):
     logger_hook = LoggerHook()
     split_pdf_hook = SplitPdfHook()
 
+    # NOTE: logger_hook should stay registered last as logs the status of
+    # request and whether it will be retried which can be changed by e.g. split_pdf_hook
+
     # Register SDK Init hooks
     hooks.register_sdk_init_hook(clean_server_url_hook)
     hooks.register_sdk_init_hook(logger_hook)
@@ -37,8 +40,9 @@ def init_hooks(hooks: Hooks):
 
     # Register After Error hooks
     hooks.register_after_success_hook(split_pdf_hook)
+    hooks.register_after_success_hook(logger_hook)
 
     # Register After Error hooks
     hooks.register_after_error_hook(suggest_defining_url_hook)
     hooks.register_after_error_hook(split_pdf_hook)
-    hooks.register_after_error_hook(logger_hook)
+    hooks.register_after_error_hook(logger_hook)