Merge pull request #39 from Unstructured-IO/jj/log-retry

Author: Coniferish

_test_unstructured_client/test_utils_retries.py

@@ -1,5 +1,6 @@
-import os
 import pytest
+import logging
+import re
 
 import requests_mock
 
@@ -8,15 +9,11 @@
 from unstructured_client.utils.retries import BackoffStrategy, RetryConfig
 
 
-def get_api_key():
-    api_key = os.getenv("UNS_API_KEY")
-    if api_key is None:
-        raise ValueError("""UNS_API_KEY environment variable not set. 
-Set it in your current shell session with `export UNS_API_KEY=<api_key>`""")
-    return api_key
+FAKE_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
 
-# this test requires UNS_API_KEY be set in your shell session. Ex: `export UNS_API_KEY=<api_key>`
-def test_backoff_strategy():
+
+def test_retry_with_backoff_does_retry(caplog):
+    caplog.set_level(logging.INFO)
     filename = "README.md"
     backoff_strategy = BackoffStrategy(
         initial_interval=100, max_interval=1000, exponent=1.5, max_elapsed_time=3000
@@ -28,13 +25,10 @@ def test_backoff_strategy():
     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)
-        session = UnstructuredClient(api_key_auth=get_api_key())
+        session = UnstructuredClient(api_key_auth=FAKE_KEY)
 
         with open(filename, "rb") as f:
-            files=shared.Files(
-                content=f.read(),
-                file_name=filename,
-            )
+            files=shared.Files(content=f.read(), file_name=filename)
 
         req = shared.PartitionParameters(files=files)
 
@@ -45,3 +39,28 @@ def test_backoff_strategy():
 
         # the number of retries varies
         assert len(mock.request_history) > 1 
+
+
+def test_backoff_strategy_logs_retries(caplog):
+    caplog.set_level(logging.INFO)
+    filename = "README.md"
+    backoff_strategy = BackoffStrategy(
+        initial_interval=100, max_interval=1000, exponent=1.5, max_elapsed_time=3000
+    )
+    retries = RetryConfig(
+        strategy="backoff", backoff=backoff_strategy, retry_connection_errors=True
+    )
+    
+    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)
+        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"{re.escape('Retry attempt #1. Sleeping')}.*{'seconds before retry'}")
+    assert bool(pattern.search(caplog.text))

src/unstructured_client/general.py

@@ -4,7 +4,7 @@
 from typing import Any, List, Optional
 from unstructured_client import utils
 from unstructured_client.models import errors, operations, shared
-from unstructured_client.utils._decorators import suggest_defining_url_if_401  # human code
+from unstructured_client.utils._human_utils import suggest_defining_url_if_401  # human code
 
 class General:
     sdk_configuration: SDKConfiguration

src/unstructured_client/sdk.py

@@ -6,7 +6,7 @@
 from typing import Callable, Dict, Union
 from unstructured_client import utils
 from unstructured_client.models import shared
-from unstructured_client.utils._decorators import clean_server_url  # human code
+from unstructured_client.utils._human_utils import clean_server_url  # human code
 
 class UnstructuredClient:
     r"""Unstructured Pipeline API: Partition documents with the Unstructured library"""

src/unstructured_client/utils/_decorators.py renamed to src/unstructured_client/utils/_human_utils.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
 import functools
-from typing import cast, Callable, TYPE_CHECKING, Optional
+import logging
+import sys
+from typing import Callable, Optional, TYPE_CHECKING, cast
 from typing_extensions import ParamSpec
-from urllib.parse import urlparse, urlunparse, ParseResult
+from urllib.parse import ParseResult, urlparse, urlunparse
 import warnings
 
-from unstructured_client.models import errors, operations
-
 if TYPE_CHECKING:
     from unstructured_client.general import General
+    from unstructured_client.models import operations
 
 
 _P = ParamSpec("_P")
@@ -75,6 +76,8 @@ def suggest_defining_url_if_401(
 
     @functools.wraps(func)
     def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> operations.PartitionResponse:
+        from unstructured_client.models import errors  # pylint: disable=C0415
+
         try:
             return func(*args, **kwargs)
         except errors.SDKError as error:
@@ -88,3 +91,18 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> operations.PartitionResponse
             return func(*args, **kwargs)
 
     return wrapper
+
+
+def log_retries(retry_count: int, sleep: float, exception: Exception):
+    """Function for logging retries to give users visibility into requests."""
+    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s', stream=sys.stdout)
+    logger = logging.getLogger('unstructured-client')
+    logger.setLevel(logging.INFO)
+    logger.info(
+        "Response status code: %s Retry attempt #%s. Sleeping %s seconds before retry.", 
+        exception.response.status_code,
+        retry_count,
+        round(sleep, 1),
+    )
+    if bool(exception.response.text):
+        logger.info(exception.response.text)

src/unstructured_client/utils/retries.py

@@ -6,6 +6,8 @@
 
 import requests
 
+from unstructured_client.utils._human_utils import log_retries  # human code
+
 
 class BackoffStrategy:
     initial_interval: int
@@ -116,5 +118,6 @@ def retry_with_backoff(func, initial_interval=500, max_interval=60000, exponent=
                      exponent**retries + random.uniform(0, 1))
             if sleep > max_interval/1000:
                 sleep = max_interval/1000
+            log_retries(retry_count=retries+1, sleep=sleep, exception=exception)  # human code
             time.sleep(sleep)
             retries += 1