Logic fixes

Author: PastelStorm

_test_unstructured_client/integration/test_decorators.py

@@ -24,6 +24,7 @@
 from unstructured_client._hooks.custom import split_pdf_hook
 
 FAKE_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+TEST_TIMEOUT_MS = 300_000
 
 _HI_RES_STRATEGIES = ("hi_res", Strategy.HI_RES)
 
@@ -143,7 +144,7 @@ def test_integration_split_pdf_has_same_output_as_non_split(
     except requests.exceptions.ConnectionError:
         assert False, "The unstructured-api is not running on localhost:8000"
 
-    client = UnstructuredClient(api_key_auth=FAKE_KEY)
+    client = UnstructuredClient(api_key_auth=FAKE_KEY, timeout_ms=TEST_TIMEOUT_MS)
 
     with open(filename, "rb") as f:
         files = shared.Files(
@@ -215,7 +216,7 @@ def test_integration_split_pdf_with_caching(
     except requests.exceptions.ConnectionError:
         assert False, "The unstructured-api is not running on localhost:8000"
 
-    client = UnstructuredClient(api_key_auth=FAKE_KEY)
+    client = UnstructuredClient(api_key_auth=FAKE_KEY, timeout_ms=TEST_TIMEOUT_MS)
 
     with open(filename, "rb") as f:
         files = shared.Files(
@@ -282,7 +283,7 @@ def test_long_pages_hi_res(filename):
         split_pdf_concurrency_level=15
     ), )
 
-    client = UnstructuredClient(api_key_auth=FAKE_KEY)
+    client = UnstructuredClient(api_key_auth=FAKE_KEY, timeout_ms=TEST_TIMEOUT_MS)
 
     response = client.general.partition(
         request=req,
@@ -301,7 +302,7 @@ def test_integration_split_pdf_for_file_with_no_name():
     except requests.exceptions.ConnectionError:
         assert False, "The unstructured-api is not running on localhost:8000"
 
-    client = UnstructuredClient(api_key_auth=FAKE_KEY)
+    client = UnstructuredClient(api_key_auth=FAKE_KEY, timeout_ms=TEST_TIMEOUT_MS)
 
     with open("_sample_docs/layout-parser-paper-fast.pdf", "rb") as f:
         files = shared.Files(
@@ -357,7 +358,7 @@ def test_integration_split_pdf_with_page_range(
     except requests.exceptions.ConnectionError:
         assert False, "The unstructured-api is not running on localhost:8000"
 
-    client = UnstructuredClient(api_key_auth=FAKE_KEY)
+    client = UnstructuredClient(api_key_auth=FAKE_KEY, timeout_ms=TEST_TIMEOUT_MS)
 
     filename = "_sample_docs/layout-parser-paper.pdf"
     with open(filename, "rb") as f:
@@ -421,7 +422,7 @@ def test_integration_split_pdf_strict_mode(
     except requests.exceptions.ConnectionError:
         assert False, "The unstructured-api is not running on localhost:8000"
 
-    client = UnstructuredClient(api_key_auth=FAKE_KEY)
+    client = UnstructuredClient(api_key_auth=FAKE_KEY, timeout_ms=TEST_TIMEOUT_MS)
 
     with open(filename, "rb") as f:
         files = shared.Files(

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import io
 from asyncio import Task
 from collections import Counter
 from functools import partial
@@ -28,6 +29,7 @@
     MAX_PAGES_PER_SPLIT,
     MIN_PAGES_PER_SPLIT,
     SplitPdfHook,
+    _get_request_timeout_seconds,
     get_optimal_split_size, run_tasks,
 )
 from unstructured_client._hooks.types import BeforeRequestContext
@@ -47,6 +49,8 @@ async def example():
         requests.Response(),
         requests.Response(),
     ]
+    hook.api_failed_responses[operation_id] = [requests.Response()]
+    hook.operation_timeouts[operation_id] = 30.0
 
     assert len(hook.coroutines_to_execute[operation_id]) == 2
     assert len(hook.api_successful_responses[operation_id]) == 2
@@ -55,6 +59,18 @@ async def example():
 
     assert hook.coroutines_to_execute.get(operation_id) is None
     assert hook.api_successful_responses.get(operation_id) is None
+    assert hook.api_failed_responses.get(operation_id) is None
+    assert hook.operation_timeouts.get(operation_id) is None
+
+
+def test_unit_get_request_timeout_seconds_uses_request_timeout_extension():
+    request = httpx.Request(
+        "POST",
+        "http://localhost",
+        extensions={"timeout": {"connect": 10.0, "read": 30.0, "write": 20.0, "pool": 5.0}},
+    )
+
+    assert _get_request_timeout_seconds(request) == 30.0
 
 
 def test_unit_prepare_request_headers():
@@ -525,4 +541,52 @@ def test_before_request_raises_pdf_validation_error_when_pdf_check_fails():
         # Verify that the mocked functions were called as expected
         mock_get_fields.assert_called_once_with(mock_request)
         mock_read_pdf.assert_called_once_with(mock_pdf_file)
-        mock_check_pdf.assert_called_once_with(mock_pdf_reader)
+        mock_check_pdf.assert_called_once_with(mock_pdf_reader)
+
+
+def test_before_request_uses_in_memory_noop_request_for_split_pdf():
+    hook = SplitPdfHook()
+    mock_client = httpx.Client()
+    hook.sdk_init(base_url="http://localhost:8888", client=mock_client)
+
+    mock_hook_ctx = MagicMock()
+    mock_hook_ctx.operation_id = "partition"
+    mock_hook_ctx.config.timeout_ms = 12_000
+
+    mock_request = MagicMock(spec=httpx.Request)
+    mock_request.headers = {"Content-Type": "multipart/form-data"}
+    mock_request.url = httpx.URL("http://localhost:8888/general/v0/general")
+    mock_request.extensions = {"timeout": {"connect": 12.0, "read": 12.0, "write": 12.0, "pool": 12.0}}
+
+    mock_pdf_file = MagicMock()
+    mock_form_data = {
+        "split_pdf_page": "true",
+        "strategy": "fast",
+        "files": {
+            "filename": "test.pdf",
+            "content_type": "application/pdf",
+            "file": mock_pdf_file,
+        },
+    }
+    mock_pdf_reader = MagicMock()
+    mock_pdf_reader.get_num_pages.return_value = 100
+    mock_pdf_reader.pages = [MagicMock()] * 100
+    mock_pdf_reader.stream = io.BytesIO(b"fake-pdf-bytes")
+
+    with patch("unstructured_client._hooks.custom.request_utils.get_multipart_stream_fields") as mock_get_fields, \
+         patch("unstructured_client._hooks.custom.pdf_utils.read_pdf") as mock_read_pdf, \
+         patch("unstructured_client._hooks.custom.pdf_utils.check_pdf") as mock_check_pdf, \
+         patch("unstructured_client._hooks.custom.request_utils.get_base_url") as mock_get_base_url, \
+         patch.object(hook, "_trim_large_pages", side_effect=lambda pdf, fd: pdf), \
+         patch.object(hook, "_get_pdf_chunks_in_memory", return_value=[]):
+        mock_get_fields.return_value = mock_form_data
+        mock_read_pdf.return_value = mock_pdf_reader
+        mock_check_pdf.return_value = mock_pdf_reader
+        mock_get_base_url.return_value = "http://localhost:8888"
+
+        result = hook.before_request(mock_hook_ctx, mock_request)
+
+    assert isinstance(result, httpx.Request)
+    assert str(result.url) == "http://no-op"
+    assert result.headers["operation_id"]
+    assert result.extensions["timeout"]["read"] == 12.0

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -56,6 +56,27 @@
 MAX_PAGES_PER_SPLIT = 20
 HI_RES_STRATEGY = 'hi_res'
 MAX_PAGE_LENGTH = 4000
+TIMEOUT_BUFFER_SECONDS = 5
+
+
+def _get_request_timeout_seconds(request: httpx.Request) -> Optional[float]:
+    timeout = request.extensions.get("timeout")
+    if timeout is None:
+        return None
+
+    if isinstance(timeout, (int, float)):
+        return float(timeout)
+
+    if isinstance(timeout, dict):
+        timeout_values = [
+            float(value)
+            for value in timeout.values()
+            if isinstance(value, (int, float))
+        ]
+        if timeout_values:
+            return max(timeout_values)
+
+    return None
 
 def _run_coroutines_in_separate_thread(
         coroutines_task: Coroutine[Any, Any, list[tuple[int, httpx.Response]]],
@@ -72,6 +93,7 @@ async def run_tasks(
     coroutines: list[partial[Coroutine[Any, Any, httpx.Response]]],
     allow_failed: bool = False,
     concurrency_level: int = 10,
+    client_timeout: Optional[httpx.Timeout] = None,
 ) -> list[tuple[int, httpx.Response]]:
     """Run a list of coroutines in parallel and return the results in order.
 
@@ -84,14 +106,15 @@ async def run_tasks(
     """
 
 
-    # Use a variable to adjust the httpx client timeout, or default to 30 minutes
-    # When we're able to reuse the SDK to make these calls, we can remove this var
-    # The SDK timeout will be controlled by parameter
     limiter = asyncio.Semaphore(concurrency_level)
-    client_timeout_minutes = 60
-    if timeout_var := os.getenv("UNSTRUCTURED_CLIENT_TIMEOUT_MINUTES"):
-        client_timeout_minutes = int(timeout_var)
-    client_timeout = httpx.Timeout(60 * client_timeout_minutes)
+    if client_timeout is None:
+        # Use a variable to adjust the httpx client timeout, or default to 60 minutes.
+        # When we're able to reuse the SDK to make these calls, we can remove this var
+        # and let the SDK timeout flow through directly.
+        client_timeout_minutes = 60
+        if timeout_var := os.getenv("UNSTRUCTURED_CLIENT_TIMEOUT_MINUTES"):
+            client_timeout_minutes = int(timeout_var)
+        client_timeout = httpx.Timeout(60 * client_timeout_minutes)
 
     async with httpx.AsyncClient(timeout=client_timeout) as client:
         armed_coroutines = [coro(async_client=client, limiter=limiter) for coro in coroutines] # type: ignore
@@ -166,6 +189,7 @@ def __init__(self) -> None:
         self.api_failed_responses: dict[str, list[httpx.Response]] = {}
         self.executors: dict[str, futures.ThreadPoolExecutor] = {}
         self.tempdirs: dict[str, tempfile.TemporaryDirectory] = {}
+        self.operation_timeouts: dict[str, Optional[float]] = {}
         self.allow_failed: bool = DEFAULT_ALLOW_FAILED
         self.cache_tmp_data_feature: bool = DEFAULT_CACHE_TMP_DATA
         self.cache_tmp_data_dir: str = DEFAULT_CACHE_TMP_DATA_DIR
@@ -268,6 +292,7 @@ def before_request(
         # We need to pass it on to after_success so
         # we know which results are ours
         operation_id = str(uuid.uuid4())
+        self.operation_timeouts[operation_id] = _get_request_timeout_seconds(request)
 
         content_type = request.headers.get("Content-Type")
         if content_type is None:
@@ -397,14 +422,11 @@ def before_request(
         # This allows us to skip right to the AfterRequestHook and await all the calls
         # Also, pass the operation_id so after_success can await the right results
 
-        # Note: We need access to the async_client from the sdk_init hook in order to set
-        # up a mock request like this.
-        # For now, just make an extra request against our api, which should return 200.
-        # dummy_request = httpx.Request("GET",  "http://no-op")
         return httpx.Request(
             "GET",
-            f"{self.partition_base_url}/general/docs",
+            "http://no-op",
             headers={"operation_id": operation_id},
+            extensions=request.extensions.copy(),
         )
 
     async def call_api_partial(
@@ -620,15 +642,25 @@ def _await_elements(self, operation_id: str) -> Optional[list]:
             return None
 
         concurrency_level = self.concurrency_level.get(operation_id, DEFAULT_CONCURRENCY_LEVEL)
-        coroutines = run_tasks(tasks, allow_failed=self.allow_failed, concurrency_level=concurrency_level)
+        timeout_seconds = self.operation_timeouts.get(operation_id)
+        client_timeout = httpx.Timeout(timeout_seconds) if timeout_seconds is not None else None
+        coroutines = run_tasks(
+            tasks,
+            allow_failed=self.allow_failed,
+            concurrency_level=concurrency_level,
+            client_timeout=client_timeout,
+        )
 
         # sending the coroutines to a separate thread to avoid blocking the current event loop
         # this operation should be removed when the SDK is updated to support async hooks
         executor = self.executors.get(operation_id)
         if executor is None:
             raise RuntimeError("Executor not found for operation_id")
         task_responses_future = executor.submit(_run_coroutines_in_separate_thread, coroutines)
-        task_responses = task_responses_future.result()
+        if timeout_seconds is None:
+            task_responses = task_responses_future.result()
+        else:
+            task_responses = task_responses_future.result(timeout=timeout_seconds + TIMEOUT_BUFFER_SECONDS)
 
         if task_responses is None:
             return None
@@ -683,23 +715,20 @@ def after_success(
 
         # Grab the correct id out of the dummy request
         operation_id = response.request.headers.get("operation_id")
+        try:
+            elements = self._await_elements(operation_id)
 
-        elements = self._await_elements(operation_id)
-
-        # if fails are disallowed, return the first failed response
-        if not self.allow_failed and self.api_failed_responses.get(operation_id):
-            failure_response = self.api_failed_responses[operation_id][0]
-
-            self._clear_operation(operation_id)
-            return failure_response
-
-        if elements is None:
-            return response
+            # if fails are disallowed, return the first failed response
+            if not self.allow_failed and self.api_failed_responses.get(operation_id):
+                return self.api_failed_responses[operation_id][0]
 
-        new_response = request_utils.create_response(elements)
-        self._clear_operation(operation_id)
+            if elements is None:
+                return response
 
-        return new_response
+            return request_utils.create_response(elements)
+        finally:
+            if operation_id is not None:
+                self._clear_operation(operation_id)
 
     def after_error(
             self,
@@ -732,7 +761,9 @@ def _clear_operation(self, operation_id: str) -> None:
         """
         self.coroutines_to_execute.pop(operation_id, None)
         self.api_successful_responses.pop(operation_id, None)
+        self.api_failed_responses.pop(operation_id, None)
         self.concurrency_level.pop(operation_id, None)
+        self.operation_timeouts.pop(operation_id, None)
         executor = self.executors.pop(operation_id, None)
         if executor is not None:
             executor.shutdown(wait=True)