feat: add document completion polling helper

README.md

@@ -197,6 +197,17 @@ doc_id = client.submit_document("report.pdf", wait=True)["doc_id"]
 print(client.chat("What was the 2023 operating margin?", doc_id=doc_id))
 ```
 
+To submit several documents first and wait for them later, use the public polling helper:
+
+```python
+submission = client.submit_document("report.pdf")
+document = client.wait_until_completed(
+    submission["doc_id"], timeout=300, poll_interval=5
+)
+```
+
+`poll_interval` is the initial delay between status requests. Repeated requests back off automatically to avoid excessive polling.
+
 | Capability | **Local** (this repo) | **Cloud** ([get an API key](https://developer.pageindex.ai/)) |
 |---|---|---|
 | Best for | text-heavy PDFs and local workflows | scanned, image-heavy, and large document collections |

pageindex/agent_tools.py

@@ -1637,8 +1637,9 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]:
         try:
             details.append(client.get_document(one_id))
         except PageIndexAPIError as exc:
-            # Batch only a definite not-found/denied (local raises carry no
-            # status); a cloud transport failure (429/5xx) propagates raw.
+            # Batch only a definite not-found/denied (including legacy
+            # client-side errors without a status); cloud transport failures
+            # (429/5xx) propagate raw.
             if exc.status_code not in (None, 403, 404):
                 raise
             missing.append(str(one_id))

pageindex/client.py

@@ -1,11 +1,13 @@
 """PageIndex SDK client: the 0.2.x cloud surface, now with a local mode."""
 from __future__ import annotations
 
+import math
 import os
 import re
 import threading
 import time
 import warnings
+from numbers import Real
 from typing import (TYPE_CHECKING, Any, Callable, Iterator, Literal, Mapping,
                     Optional, Union, cast, overload)
 
@@ -618,19 +620,70 @@ def submit_document(
                 stacklevel=2,
             )
         if wait:
-            self._wait_until_ready(result["doc_id"])
+            self.wait_until_completed(result["doc_id"])
         return result
 
-    def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None:
+    def wait_until_completed(
+        self,
+        doc_id: str,
+        timeout: float = 1800.0,
+        poll_interval: float = 2.0,
+    ) -> dict[str, Any]:
+        """Wait until an existing document is ready and return its metadata.
+
+        ``poll_interval`` is the initial delay. Repeated polls back off by 1.5x,
+        capped at 15 seconds (or the initial interval when it is larger).
+        """
         import requests
-        interval = 2.0
-        deadline = time.monotonic() + timeout
+
+        timings = {"timeout": timeout, "poll_interval": poll_interval}
+        validated_timings: dict[str, float] = {}
+        for name, value in timings.items():
+            if isinstance(value, bool) or not isinstance(value, Real):
+                raise ValueError(f"{name} must be a finite number greater than 0")
+            try:
+                seconds = float(value)
+            except (OverflowError, ValueError) as exc:
+                raise ValueError(
+                    f"{name} must be a finite number greater than 0"
+                ) from exc
+            if not math.isfinite(seconds) or seconds <= 0:
+                raise ValueError(f"{name} must be a finite number greater than 0")
+            validated_timings[name] = seconds
+
+        timeout_seconds = validated_timings["timeout"]
+        interval = validated_timings["poll_interval"]
+        max_interval = max(15.0, interval)
+        deadline = time.monotonic() + timeout_seconds
         poll_failures = 0
+        last_status = None
+
+        def timeout_error() -> PageIndexAPIError:
+            return PageIndexAPIError(
+                f"Timed out after {timeout_seconds:g}s waiting for document "
+                f"processing (doc_id: {doc_id}, last status: {last_status}). "
+                "Processing continues in the cloud — poll "
+                "get_document(doc_id) for status."
+            )
+
         while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise timeout_error()
+            status = None
             try:
-                status = self.get_document(doc_id).get("status")
-                poll_failures = 0
+                bounded_get_document = getattr(
+                    self._api, "_get_document_with_timeout", None
+                )
+                if bounded_get_document is None:
+                    document = self.get_document(doc_id)
+                else:
+                    document = bounded_get_document(
+                        doc_id, timeout=min(30.0, remaining)
+                    )
             except (PageIndexAPIError, requests.RequestException) as exc:
+                if time.monotonic() >= deadline:
+                    raise timeout_error() from exc
                 if getattr(exc, "status_code", None) in (401, 403, 404):
                     raise  # a definite answer, not a poll failure
                 # Tolerate transient poll failures; a 30-minute wait should
@@ -642,22 +695,25 @@ def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None:
                         f"{exc}. Processing continues in the cloud — poll "
                         "get_document(doc_id) for status."
                     ) from exc
-                status = None
+            else:
+                if time.monotonic() >= deadline:
+                    raise timeout_error()
+                status = document.get("status")
+                last_status = status
+                poll_failures = 0
+
             if status == "completed":
-                return
+                return document
             if status == "failed":
                 raise PageIndexAPIError(
                     f"Document processing failed (doc_id: {doc_id})."
                 )
-            if time.monotonic() >= deadline:
-                raise PageIndexAPIError(
-                    f"Timed out after {int(timeout)}s waiting for document "
-                    f"processing (doc_id: {doc_id}, last status: {status}). "
-                    "Processing continues in the cloud — poll "
-                    "get_document(doc_id) for status."
-                )
-            time.sleep(interval)
-            interval = min(interval * 1.5, 15.0)
+
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise timeout_error()
+            time.sleep(min(interval, remaining))
+            interval = min(interval * 1.5, max_interval)
 
     # ---------- OCR FUNCTIONALITY ----------
 

pageindex/cloud_api.py

@@ -326,10 +326,16 @@ def get_document(self, doc_id: str) -> Dict[str, Any]:
                 - createdAt (str): Creation timestamp in ISO format
                 - pageNum (int): Number of pages in the document
         """
+        return self._get_document_with_timeout(doc_id, timeout=30)
+
+    def _get_document_with_timeout(
+        self, doc_id: str, *, timeout: float
+    ) -> Dict[str, Any]:
+        """Fetch metadata with a caller-supplied request timeout."""
         response = requests.get(
             f"{self.BASE_URL}/doc/{_enc(doc_id)}/metadata/",
             headers=self._headers(),
-            timeout=30
+            timeout=timeout,
         )
         if response.status_code != 200:
             raise PageIndexAPIError(

pageindex/local_api.py

@@ -333,7 +333,10 @@ def _require_doc(self, doc_id: str, error_prefix: str) -> dict:
     def get_document(self, doc_id: str) -> dict[str, Any]:
         meta = self._store.get_meta(doc_id)
         if meta is None:
-            raise PageIndexAPIError("Failed to get document metadata: Document not found")
+            raise PageIndexAPIError(
+                "Failed to get document metadata: Document not found",
+                status_code=404,
+            )
         return {key: meta.get(key) for key in
                 ("id", "name", "description", "status", "createdAt", "pageNum", "folderId")}