fix/ensure that split pdf requests are retried (#160)

We discovered that when a pdf is split into smaller chunks, those
requests are not being retried. Now that we have `allow_failed=False`,
this results in the whole document failing as soon as any of the child
requests hit a transient error. The fix is to reuse the `utils.Retry`
logic that the main code path uses. Copying the retry config in the hook
logic is not great, and we can work with Speakeasy to make the internal
logic more modular so we can reuse more. But for now, this will address
the current failures while we work on a better implementation.

Testing:
See the added unit test. The existing retry logic works for the final
split page, everything else needs to use the new logic. To test this, I
mocked a response from the server to return 502 for a low
`starting_page_number`, which we know will have to be handled by the
hooks.

Other changes:
Remove the "Not splitting" log. When the final split page is retried, it
triggers all the hooks again. We need to force `split_pdf_page=False` in
this request, and we don't need additional logging when this code is hit
again.

Author: awalker4

_test_unstructured_client/integration/test_decorators.py

@@ -1,9 +1,18 @@
+import httpx
+import json
 import pytest
 import requests
 from deepdiff import DeepDiff
+from httpx import Response
+
+from requests_toolbelt.multipart.decoder import MultipartDecoder  # type: ignore
+
 from unstructured_client import UnstructuredClient
 from unstructured_client.models import shared, operations
 from unstructured_client.models.errors import HTTPValidationError
+from unstructured_client.utils.retries import BackoffStrategy, RetryConfig
+from unstructured_client._hooks.custom import form_utils
+from unstructured_client._hooks.custom import split_pdf_hook
 
 FAKE_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
 
@@ -275,3 +284,80 @@ def test_integration_split_pdf_strict_mode(
         ],
     )
     assert len(diff) == 0
+
+
+@pytest.mark.asyncio
+async def test_split_pdf_requests_do_retry(monkeypatch):
+    """
+    Test that when we split a pdf, the split requests will honor retryable errors.
+    """
+    number_of_split_502s = 2
+    number_of_last_page_502s = 2
+
+    async def mock_send(_, request):
+        """
+        Return a predefined number of 502s for requests with certain starting_page_number values.
+
+        This is because N-1 splits are sent off in the hook logic. These need explicit retry handling.
+        The final split is returned to the SDK and gets the built in retry code.
+
+        We want to make sure both code paths are retried.
+        """
+        request_body = request.read()
+        decoded_body = MultipartDecoder(request_body, request.headers.get("Content-Type"))
+        form_data = form_utils.parse_form_data(decoded_body)
+
+        nonlocal number_of_split_502s
+        nonlocal number_of_last_page_502s
+
+        if number_of_split_502s > 0:
+            if "starting_page_number" in form_data and int(form_data["starting_page_number"]) < 3:
+                number_of_split_502s -= 1
+                return Response(502, request=request)
+
+        if number_of_last_page_502s > 0:
+            if "starting_page_number" in form_data and int(form_data["starting_page_number"]) > 12:
+                number_of_last_page_502s -= 1
+                return Response(502, request=request)
+
+        mock_return_data = [{
+            "type": "Title",
+            "text": "Hello",
+        }]
+
+        return Response(
+            200,
+            request=request,
+            content=json.dumps(mock_return_data),
+            headers={"Content-Type": "application/json"},
+        )
+
+    monkeypatch.setattr(split_pdf_hook.httpx.AsyncClient, "send", mock_send)
+
+    sdk = UnstructuredClient(
+        api_key_auth=FAKE_KEY,
+        server_url="localhost:8000",
+        retry_config=RetryConfig("backoff", BackoffStrategy(200, 1000, 1.5, 1000), False),
+    )
+
+    filename = "_sample_docs/layout-parser-paper.pdf"
+    with open(filename, "rb") as f:
+        files = shared.Files(
+            content=f.read(),
+            file_name=filename,
+        )
+
+    req = operations.PartitionRequest(
+        shared.PartitionParameters(
+            files=files,
+            split_pdf_page=True,
+            split_pdf_allow_failed=False,
+            strategy="fast",
+        )
+    )
+
+    res = await sdk.general.partition_async(request=req)
+
+    assert number_of_split_502s == 0
+    assert number_of_last_page_502s == 0
+    assert res.status_code == 200

src/unstructured_client/_hooks/custom/request_utils.py

@@ -19,6 +19,7 @@
     PARTITION_FORM_STARTING_PAGE_NUMBER_KEY,
     FormData,
 )
+import unstructured_client.utils as utils
 
 logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
 
@@ -69,10 +70,42 @@ async def call_api_async(
     )
 
     async with limiter:
-        response = await client.send(new_request)
+        response = await send_request_async_with_retries(client, new_request)
         return response
 
 
+async def send_request_async_with_retries(client: httpx.AsyncClient, request: httpx.Request):
+    # Hardcode the retry config until we can
+    # properly reuse the SDK logic
+    # (Values are in ms)
+    retry_config = utils.RetryConfig(
+        "backoff",
+        utils.BackoffStrategy(
+            initial_interval=2000,
+            max_interval=60000,
+            exponent=1.5,
+            max_elapsed_time=1000 * 60 * 5  # 5 minutes
+        ),
+        retry_connection_errors=True
+    )
+
+    retryable_codes = [
+        "502",
+        "503",
+        "504"
+    ]
+
+    async def do_request():
+        return await client.send(request)
+
+    response = await utils.retry_async(
+        do_request,
+        utils.Retries(retry_config, retryable_codes)
+    )
+
+    return response
+
+
 def prepare_request_headers(
     headers: httpx.Headers,
 ) -> httpx.Headers:

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -168,7 +168,6 @@ def before_request(
         form_data = form_utils.parse_form_data(decoded_body)
         split_pdf_page = form_data.get(PARTITION_FORM_SPLIT_PDF_PAGE_KEY)
         if split_pdf_page is None or split_pdf_page == "false":
-            logger.info("Partitioning without split.")
             return request
 
         logger.info("Preparing to split document for partition.")
@@ -287,6 +286,8 @@ async def call_api_partial(page):
         # `before_request` method needs to return a request so we skip sending the last page in parallel
         # and return that last page at the end of this method
 
+        # Need to make sure the final page does not trigger splitting again
+        form_data[PARTITION_FORM_SPLIT_PDF_PAGE_KEY] = "false"
         body = request_utils.create_request_body(
             form_data, last_page_content, file.file_name, last_page_number
         )
@@ -413,12 +414,13 @@ def after_error(
             If requests were run in parallel, and at least one was successful, a combined
             response object; otherwise, the original response and exception.
         """
+        operation_id = hook_ctx.operation_id
 
         # if fails are disallowed - return response and error objects immediately
         if not self.allow_failed:
+            self._clear_operation(operation_id)
             return (response, error)
 
-        operation_id = hook_ctx.operation_id
         # We know that this request failed so we pass a failed or empty response to `_await_elements` method
         # where it checks if at least on of the other requests succeeded
         elements = self._await_elements(operation_id, response or httpx.Response(status_code=200))