add logger for retries

Author: Coniferish

_test_unstructured_client/test_utils_retries.py

@@ -1,5 +1,5 @@
-import os
 import pytest
+import logging
 
 import requests_mock
 
@@ -8,15 +8,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,7 +24,7 @@ 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(
@@ -45,3 +41,32 @@ 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)    
+
+    
+    assert "seconds before retry" in 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
+        
         try:
             return func(*args, **kwargs)
         except errors.SDKError as error:
@@ -86,5 +89,11 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> operations.PartitionResponse
                     )
 
             return func(*args, **kwargs)
-
+        
     return wrapper
+
+
+def log_retries(retry_count, sleep):
+    """Function for logging retries to give users visibility into requests."""
+    logging.basicConfig(level=logging.INFO, stream=sys.stdout)
+    logging.info(f"Retry attempt {retry_count}. Sleeping {round(sleep, 1)} seconds before retry.")

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, sleep=sleep)  # human code
             time.sleep(sleep)
             retries += 1