From 6d41ecc29a249dea0902e4c55e335a6c77dbad65 Mon Sep 17 00:00:00 2001 From: denis-samatov Date: Mon, 31 Aug 2026 15:33:56 +0700 Subject: [PATCH 1/4] feat: expose document completion polling --- pageindex/client.py | 71 ++++++++++++++++++----- tests/test_agent_tools.py | 119 +++++++++++++++++++++++++++++++++++++- 2 files changed, 173 insertions(+), 17 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 8d6aec6fd..e16420a27 100644 --- a/pageindex/client.py +++ b/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 Any, Callable, Iterator, Mapping, Optional, Union, cast from .errors import PageIndexAPIError @@ -590,17 +592,57 @@ 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} + for name, value in timings.items(): + if ( + isinstance(value, bool) + or not isinstance(value, Real) + or not math.isfinite(float(value)) + or value <= 0 + ): + raise ValueError(f"{name} must be a finite number greater than 0") + + timeout_seconds = float(timeout) + interval = float(poll_interval) + max_interval = max(15.0, interval) + deadline = time.monotonic() + timeout_seconds poll_failures = 0 + last_status = None + first_poll = True + + 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: + if not first_poll and time.monotonic() >= deadline: + raise timeout_error() + first_poll = False + status = None try: - status = self.get_document(doc_id).get("status") + document = self.get_document(doc_id) + status = document.get("status") + last_status = status poll_failures = 0 except (PageIndexAPIError, requests.RequestException) as exc: if getattr(exc, "status_code", None) in (401, 403, 404): @@ -614,22 +656,19 @@ 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 + 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 ---------- diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index db400273b..89b0d01d6 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2383,7 +2383,21 @@ def instructions(self): cloud.agent_instructions() -# ── submit_document(wait=True) ── +# ── document completion polling ── + + +class _FakeClock: + def __init__(self): + self.now = 0.0 + self.sleeps = [] + + def monotonic(self): + return self.now + + def sleep(self, seconds): + self.sleeps.append(seconds) + self.now += seconds + class _FakeCloudAPI: def __init__(self, statuses): @@ -2412,6 +2426,109 @@ def build(statuses): return build +def test_wait_until_completed_returns_final_document(fake_cloud_client): + cloud = fake_cloud_client(["processing", "completed"]) + + result = cloud.wait_until_completed("pi-fake") + + assert result == {"id": "pi-fake", "status": "completed"} + assert cloud._api.polls == 2 + + +@pytest.mark.parametrize("parameter", ["timeout", "poll_interval"]) +@pytest.mark.parametrize( + "bad_value", [0, -1, float("nan"), float("inf"), True, "2"] +) +def test_wait_until_completed_rejects_invalid_timing( + fake_cloud_client, parameter, bad_value +): + cloud = fake_cloud_client(["completed"]) + + with pytest.raises(ValueError, match=parameter): + cloud.wait_until_completed("pi-fake", **{parameter: bad_value}) + + assert cloud._api.polls == 0 + + +def test_wait_until_completed_uses_custom_initial_interval( + fake_cloud_client, monkeypatch +): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + cloud = fake_cloud_client(["processing", "processing", "completed"]) + + cloud.wait_until_completed("pi-fake", poll_interval=4) + + assert clock.sleeps == [4.0, 6.0] + + +def test_wait_until_completed_does_not_poll_after_deadline( + fake_cloud_client, monkeypatch +): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + cloud = fake_cloud_client(["processing"]) + + with pytest.raises(PageIndexAPIError, match="Timed out"): + cloud.wait_until_completed("pi-fake", timeout=3, poll_interval=2) + + assert clock.sleeps == [2.0, 1.0] + assert cloud._api.polls == 2 + + +@pytest.mark.parametrize("status_code", [401, 403, 404]) +def test_wait_reraises_definite_poll_answers( + fake_cloud_client, monkeypatch, status_code +): + cloud = fake_cloud_client(["processing"]) + polls = {"n": 0} + + def denied(doc_id): + polls["n"] += 1 + raise PageIndexAPIError( + f"Failed to get document metadata: {status_code}", + status_code=status_code, + ) + + monkeypatch.setattr(cloud, "get_document", denied) + with pytest.raises(PageIndexAPIError, match=str(status_code)) as err: + cloud.wait_until_completed("pi-fake") + + assert polls["n"] == 1 + assert "Processing continues" not in str(err.value) + + +def test_wait_timeout_keeps_last_successful_status( + fake_cloud_client, monkeypatch +): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + cloud = fake_cloud_client(["processing"]) + responses = iter([ + {"id": "pi-fake", "status": "processing"}, + PageIndexAPIError("temporary 502"), + ]) + + def poll(doc_id): + response = next(responses) + if isinstance(response, Exception): + raise response + return response + + monkeypatch.setattr(cloud, "get_document", poll) + with pytest.raises(PageIndexAPIError, match="last status: processing"): + cloud.wait_until_completed("pi-fake", timeout=1, poll_interval=0.5) + + +def test_wait_completed_document_does_not_sleep(fake_cloud_client, monkeypatch): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + cloud = fake_cloud_client(["completed"]) + + assert cloud.wait_until_completed("pi-fake")["status"] == "completed" + assert clock.sleeps == [] + + def test_submit_wait_polls_until_completed(fake_cloud_client): cloud = fake_cloud_client(["processing", "processing", "completed"]) result = cloud.submit_document("whatever.pdf", wait=True) From b93aa91e257a31cb325778fd4bf82c5853d9961e Mon Sep 17 00:00:00 2001 From: denis-samatov Date: Mon, 31 Aug 2026 15:35:26 +0700 Subject: [PATCH 2/4] docs: show delayed document polling --- README.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/README.md b/README.md index 52dbc8db1..8d71e2a6c 100644 --- a/README.md +++ b/README.md @@ -199,6 +199,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 | From 9c82434e155d059de0bcc7e6193fcd2e81d0ac70 Mon Sep 17 00:00:00 2001 From: denis-samatov Date: Mon, 31 Aug 2026 15:57:26 +0700 Subject: [PATCH 3/4] fix: handle polling edge cases --- pageindex/agent_tools.py | 5 ++- pageindex/client.py | 21 ++++++---- pageindex/local_api.py | 5 ++- tests/test_agent_tools.py | 87 ++++++++++++++++++++++++++++++++++++++- 4 files changed, 106 insertions(+), 12 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 6e5bd4c5a..57f265cd2 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1636,8 +1636,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)) diff --git a/pageindex/client.py b/pageindex/client.py index e16420a27..7614883e9 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -609,17 +609,22 @@ def wait_until_completed( import requests 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) - or not math.isfinite(float(value)) - or value <= 0 - ): + 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 = float(timeout) - interval = float(poll_interval) + timeout_seconds = validated_timings["timeout"] + interval = validated_timings["poll_interval"] max_interval = max(15.0, interval) deadline = time.monotonic() + timeout_seconds poll_failures = 0 diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 40aa3aedc..175d69527 100644 --- a/pageindex/local_api.py +++ b/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")} diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 89b0d01d6..6c2ddc9aa 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2437,7 +2437,17 @@ def test_wait_until_completed_returns_final_document(fake_cloud_client): @pytest.mark.parametrize("parameter", ["timeout", "poll_interval"]) @pytest.mark.parametrize( - "bad_value", [0, -1, float("nan"), float("inf"), True, "2"] + "bad_value", + [ + 0, + -1, + float("nan"), + float("inf"), + True, + "2", + pytest.param(10**1000, id="overflow-positive"), + pytest.param(-(10**1000), id="overflow-negative"), + ], ) def test_wait_until_completed_rejects_invalid_timing( fake_cloud_client, parameter, bad_value @@ -2462,6 +2472,30 @@ def test_wait_until_completed_uses_custom_initial_interval( assert clock.sleeps == [4.0, 6.0] +def test_wait_until_completed_caps_backoff(fake_cloud_client, monkeypatch): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + cloud = fake_cloud_client([ + "processing", "processing", "processing", "processing", "completed" + ]) + + cloud.wait_until_completed("pi-fake", poll_interval=8) + + assert clock.sleeps == [8.0, 12.0, 15.0, 15.0] + + +def test_wait_until_completed_preserves_large_initial_interval( + fake_cloud_client, monkeypatch +): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + cloud = fake_cloud_client(["processing", "completed"]) + + cloud.wait_until_completed("pi-fake", poll_interval=20) + + assert clock.sleeps == [20.0] + + def test_wait_until_completed_does_not_poll_after_deadline( fake_cloud_client, monkeypatch ): @@ -2529,6 +2563,29 @@ def test_wait_completed_document_does_not_sleep(fake_cloud_client, monkeypatch): assert clock.sleeps == [] +def test_wait_until_completed_returns_real_local_document(client, store_path): + seed_doc(store_path, "pi-local", "local.pdf") + + document = client.wait_until_completed("pi-local") + + assert document["id"] == "pi-local" + assert document["status"] == "completed" + + +def test_wait_until_completed_reraises_missing_local_document( + client, monkeypatch +): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + + with pytest.raises(PageIndexAPIError, match="Document not found") as err: + client.wait_until_completed("pi-missing") + + assert err.value.status_code == 404 + assert "Processing continues in the cloud" not in str(err.value) + assert clock.sleeps == [] + + def test_submit_wait_polls_until_completed(fake_cloud_client): cloud = fake_cloud_client(["processing", "processing", "completed"]) result = cloud.submit_document("whatever.pdf", wait=True) @@ -2576,12 +2633,40 @@ def test_submit_wait_poll_error_carries_doc_id(fake_cloud_client, monkeypatch): recoverable, like the timeout and failed branches do.""" cloud = fake_cloud_client(["processing"]) + polls = {"n": 0} + def boom(doc_id): + polls["n"] += 1 raise PageIndexAPIError("Failed to get document metadata: 502") monkeypatch.setattr(cloud, "get_document", boom) with pytest.raises(PageIndexAPIError, match="pi-fake"): cloud.submit_document("whatever.pdf", wait=True) + assert polls["n"] == 3 + + +def test_wait_recovers_after_two_consecutive_poll_failures( + fake_cloud_client, monkeypatch +): + cloud = fake_cloud_client(["processing"]) + responses = iter([ + PageIndexAPIError("temporary 502"), + PageIndexAPIError("temporary 503"), + {"id": "pi-fake", "status": "completed"}, + ]) + polls = {"n": 0} + + def poll(doc_id): + polls["n"] += 1 + response = next(responses) + if isinstance(response, Exception): + raise response + return response + + monkeypatch.setattr(cloud, "get_document", poll) + + assert cloud.wait_until_completed("pi-fake")["status"] == "completed" + assert polls["n"] == 3 def test_submit_wait_reraises_definite_poll_answers(fake_cloud_client, From ab9335e677ce00419b05bc686ff994f1fb8c480a Mon Sep 17 00:00:00 2001 From: denis-samatov Date: Mon, 31 Aug 2026 16:19:34 +0700 Subject: [PATCH 4/4] fix: enforce polling request deadline --- pageindex/client.py | 26 ++++++++++++++------ pageindex/cloud_api.py | 8 +++++- tests/test_agent_tools.py | 52 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 8 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 7614883e9..36f19f4ba 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -629,7 +629,6 @@ def wait_until_completed( deadline = time.monotonic() + timeout_seconds poll_failures = 0 last_status = None - first_poll = True def timeout_error() -> PageIndexAPIError: return PageIndexAPIError( @@ -640,16 +639,23 @@ def timeout_error() -> PageIndexAPIError: ) while True: - if not first_poll and time.monotonic() >= deadline: + remaining = deadline - time.monotonic() + if remaining <= 0: raise timeout_error() - first_poll = False status = None try: - document = self.get_document(doc_id) - status = document.get("status") - last_status = 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 @@ -661,6 +667,12 @@ def timeout_error() -> PageIndexAPIError: f"{exc}. Processing continues in the cloud — poll " "get_document(doc_id) for status." ) from exc + else: + if time.monotonic() >= deadline: + raise timeout_error() + status = document.get("status") + last_status = status + poll_failures = 0 if status == "completed": return document diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index df3c0088c..7f5604449 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -318,10 +318,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( diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 6c2ddc9aa..780fd607e 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2510,6 +2510,58 @@ def test_wait_until_completed_does_not_poll_after_deadline( assert cloud._api.polls == 2 +def test_wait_bounds_cloud_request_and_rejects_late_completion(monkeypatch): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + request_timeouts = [] + + class Response: + status_code = 200 + + @staticmethod + def json(): + return {"id": "pi-cloud", "status": "completed"} + + def slow_get(url, headers, timeout): + request_timeouts.append(timeout) + clock.now += 1.5 + return Response() + + monkeypatch.setattr("pageindex.cloud_api.requests.get", slow_get) + cloud = PageIndexCloudClient(api_key="pi-test-key") + + with pytest.raises(PageIndexAPIError, match="Timed out"): + cloud.wait_until_completed("pi-cloud", timeout=1) + + assert request_timeouts == [1.0] + + +def test_wait_cloud_request_keeps_existing_thirty_second_cap(monkeypatch): + clock = _FakeClock() + monkeypatch.setattr(client_module, "time", clock) + request_timeouts = [] + + class Response: + status_code = 200 + + @staticmethod + def json(): + return {"id": "pi-cloud", "status": "completed"} + + def quick_get(url, headers, timeout): + request_timeouts.append(timeout) + clock.now += 0.25 + return Response() + + monkeypatch.setattr("pageindex.cloud_api.requests.get", quick_get) + cloud = PageIndexCloudClient(api_key="pi-test-key") + + document = cloud.wait_until_completed("pi-cloud", timeout=60) + + assert document["status"] == "completed" + assert request_timeouts == [30.0] + + @pytest.mark.parametrize("status_code", [401, 403, 404]) def test_wait_reraises_definite_poll_answers( fake_cloud_client, monkeypatch, status_code