From a8676be7e856b4589f35e7281845894e8184d98b Mon Sep 17 00:00:00 2001 From: Jean-Adrien DUCASTAING Date: Mon, 3 Aug 2026 16:21:31 +0200 Subject: [PATCH] feat(job): allow to wait for async jobs --- AGENTS.md | 2 +- README.md | 19 ++++++++++++++++ lighton/job.py | 33 +++++++++++++++++++++++++++ lighton/verbs/extract.py | 22 +++++++++++++++--- lighton/verbs/parse.py | 21 +++++++++++++++-- tests/e2e/cli.py | 23 ++++--------------- tests/test_extract.py | 49 ++++++++++++++++++++++++++++++++++++++++ tests/test_parse.py | 18 +++++++++++++++ 8 files changed, 162 insertions(+), 25 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 121d8be..0ec021e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -61,7 +61,7 @@ too, no documented value set). - **Sync only.** `httpx.Client`. No async client until a real event-loop caller needs one, `_request` is the only logic to mirror. - **One `_request`** does auth header, error mapping (→ raises), and JSON parse. All calls route through it. A 2xx body that isn't JSON → `MalformedResponseError`. - **Primary verbs** live one-per-file in `verbs/` as mixins (`AskMixin`/`SearchMixin`/`ParseMixin`/`ExtractMixin`) composed onto `LightOn`. Each references `self._request`; the stub on `_VerbClient` (their shared base) makes them type-check in isolation, and `LightOn._request` overrides it at runtime. Keeps `_client.py` to just the transport core. They take explicit typed params and return the generated response models via `model_validate`. `ask`/`search` take `workspaces`/`tags`/`files` (lists of `Workspace`/`Tag`/`File` objects or bare ids; `_ids()` in `utils.py` coerces via duck-typed `.id` → the API's `workspace_id`/`tag_id`/`file_id`; server-side, `file_id` can't combine with `workspace_id`/`tag_id`, and `tag_id` is OR-matched). `tags` additionally accepts **name strings**, resolved through `tag.resolve_ids` (same helper as `File.tag`), so the verb `cast`s `self` to `LightOn` (the mixin `self` is typed `_VerbClient`) to call `Tag.list`; resolution only lists when a name is present. `parse` takes keyword-only `path` XOR `url` (multipart vs JSON body; raises `ValueError` unless exactly one). `extract` takes keyword-only `path` XOR `url` (multipart vs JSON body; raises `ValueError` unless exactly one, same as `parse`) plus a `schema` that is **either a pydantic model class** or a **raw JSON-Schema dict**; returns `ExtractJobResponse`. The multipart `file` upload isn't in the OpenAPI schema (`ExtractRequest` models only `document`/`schema`/`options`) but the endpoint accepts it, verified by curl; on multipart, `schema`/`options` ride as JSON-encoded form fields alongside the `file` part. Schema handling (in `utils.py`): a dict is validated against the draft-2020-12 meta-schema via `jsonschema` (`validate_response_format_json`, raises `SchemaError`) and otherwise passed through; a pydantic model is converted to a vLLM guided-generation `response_format` schema by `convert_pydantic_to_response_format_json`, `model_json_schema()` then normalized: `$defs`/`$ref` inlined (`_inline_refs`), nullable `anyOf` collapsed to `type: [X, "null"]` (`_collapse_nullable`), draft-2020-12 `$schema` marker added. `jsonschema` is a runtime dep (meta-schema validation is its job; hand-rolling would be flimsy). Ceiling: `_inline_refs` recurses through refs, so a self-referential model would overflow, fine, guided-gen grammars can't express unbounded recursion anyway. -- **Async jobs.** `parse`/`extract` take `mode: ExecMode` (default `ExecMode.SYNC`); `ExecMode.ASYNC` (uppercase members, value `"async"`, and lowercase `async` can't be a member name) sends `options={"async": true}`. `ExecMode` lives in `enums.py` (StrEnum, exported). Async returns a **pollable job handle** (`job.py`): `parse(mode=ASYNC)` → `ParseJob`, `extract(mode=ASYNC)` → `ExtractJob`; sync returns the full response model as before. Each verb has two `@overload`s keyed on `mode: Literal[ExecMode.SYNC|ASYNC]` so callers get the exact return type (`ParseResponse` vs `ParseJob`) instead of the union, the impl signature keeps the `ExecMode` default and the `... | ...Job` return. `Job.poll(page=None)` GETs `/`, absorbs the response onto itself in place (mirrors `_ActiveRecord._absorb`), returns self; `.done` (terminal, `completed_at` set) and `.succeeded` (`status == completed`) read state. `_Job` is a hand-written curated model (`extra="ignore"`) holding the shared plumbing + fields; `ParseJob`/`ExtractJob` subclass it ONLY because `result` differs (`ParseResult.pages` vs `ExtractResult.data`, whose optional fields make a union ambiguous), parse also has `error`. The job binds to the client via the `_VerbClient` transport surface (all it needs is `_request`), not a full `LightOn` (keeps the mixin's `self` assignable without a cast). `JobStatus` (enums.py) has only the documented `pending`/`completed`, the API doesn't publish the failure vocab, so it's for call-site comparison (StrEnum, unknown server values compare unequal, never validated onto the field), and the "poll until `.succeeded`, raise once `.done`" pattern keys off `completed_at`, not a failure string. No auto-wait helper, callers loop with `time.sleep` (see README); add one if asked. +- **Async jobs.** `parse`/`extract` take `mode: ExecMode` (default `ExecMode.SYNC`); `ExecMode.ASYNC` (uppercase members, value `"async"`, and lowercase `async` can't be a member name) sends `options={"async": true}`. `ExecMode` lives in `enums.py` (StrEnum, exported). Async returns a **pollable job handle** (`job.py`): `parse(mode=ASYNC)` → `ParseJob`, `extract(mode=ASYNC)` → `ExtractJob`; sync returns the full response model as before. Each verb has two `@overload`s keyed on `mode: Literal[ExecMode.SYNC|ASYNC]` so callers get the exact return type (`ParseResponse` vs `ParseJob`) instead of the union, the impl signature keeps the `ExecMode` default and the `... | ...Job` return. `Job.poll(page=None)` GETs `/`, absorbs the response onto itself in place (mirrors `_ActiveRecord._absorb`), returns self; `.done` (terminal, `completed_at` set) and `.succeeded` (`status == completed`) read state. `_Job` is a hand-written curated model (`extra="ignore"`) holding the shared plumbing + fields; `ParseJob`/`ExtractJob` subclass it ONLY because `result` differs (`ParseResult.pages` vs `ExtractResult.data`, whose optional fields make a union ambiguous), parse also has `error`. The job binds to the client via the `_VerbClient` transport surface (all it needs is `_request`), not a full `LightOn` (keeps the mixin's `self` assignable without a cast). `JobStatus` (enums.py) has only the documented `pending`/`completed`, the API doesn't publish the failure vocab, so it's for call-site comparison (StrEnum, unknown server values compare unequal, never validated onto the field), and the "poll until `.succeeded`, raise once `.done`" pattern keys off `completed_at`, not a failure string. `_Job.wait(timeout=300, poll=2)` is the auto-wait: a `File.wait`-style poll loop (no webhook exists) that returns self once terminal, raises `TimeoutError` past the deadline and `LightOnError` if `not .succeeded` (detail from `error` when the subclass has one, `getattr`, since only `ParseJob` does). The verbs expose it as `wait=False`/`timeout=300.0` (**same pair as `Workspace.ingest`**), declared **only on the ASYNC `@overload`** so `wait=True` without `mode=ASYNC` is a static error *and* a `ValueError` (sync already blocks); the two negative tests carry a `# ty: ignore[no-matching-overload]`. `wait=True` still returns the job (not the sync response model), so the return-type overloads stay two. No `poll` knob on the verbs, callers who need one use `job.wait(poll=...)`. - Deferred: tag/content_type/attribute filters, streaming, add the params when needed. - **Config object.** Non-essential knobs (`base_url`, `timeout`, `retries`, `transport`) live in `LightOnConfiguration` (pydantic, `arbitrary_types_allowed`). `api_key` stays a direct `LightOn()` arg; falls back to `LIGHTON_API_KEY` env. - **Retries / rate limiting.** Two layers: `httpx.HTTPTransport(retries=)` handles diff --git a/README.md b/README.md index 8bbaf33..38e59e8 100644 --- a/README.md +++ b/README.md @@ -321,6 +321,25 @@ for row in job.result.data: reads naturally; raising once `job.done` (terminal but not successful) means a stuck or failed job surfaces instead of looping forever. +If you don't need live progress, don't write the loop: pass `wait=True` to block +until the job is terminal (same `wait=` / `timeout=` pair as `ingest`), or call +`job.wait()` yourself. Both return the finished job, raise `TimeoutError` past +`timeout` (default 300s), and raise `LightOnError` if the job ends in failure, so +the `result` is there when the call returns. + +```python +# async endpoint (no sync timeout to hit), but blocking, no polling code +job = client.extract(schema=Letter, path="big-scan.pdf", mode=ExecMode.ASYNC, wait=True) +for row in job.result.data: + print(row) + +# equivalent, and how to tune the poll interval +job = client.parse(path="big.pdf", mode=ExecMode.ASYNC).wait(timeout=1800, poll=5) +``` + +`wait=True` only makes sense with `ExecMode.ASYNC` (sync already blocks); passing +it without is a `ValueError`. + `parse` is the same shape, on failure a `ParseJob` carries an `error` block you can raise with directly: diff --git a/lighton/job.py b/lighton/job.py index 778a41e..881648b 100644 --- a/lighton/job.py +++ b/lighton/job.py @@ -9,11 +9,13 @@ from __future__ import annotations +import time from typing import Self from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, PrivateAttr from lighton.enums import JobStatus +from lighton.exceptions import LightOnError from lighton.types.api import ( ExtractDocument, ExtractResult, @@ -83,6 +85,37 @@ def poll(self, *, page: int | None = None) -> Self: setattr(self, field, getattr(fresh, field)) return self + def wait(self, timeout: float = 300.0, poll: float = 2.0) -> Self: + """Block (polling) until the job is terminal, mirrors `File.wait`. + + ponytail: dumb poll loop, the API offers no webhook. + + Args: + timeout: Max seconds to wait before raising TimeoutError. + poll: Seconds to sleep between status checks. + + Returns: + `self`, once the job has completed successfully. + + Raises: + TimeoutError: If `timeout` elapses before the job is terminal. + LightOnError: If the job ends in a terminal-failure state. + """ + deadline = time.monotonic() + timeout + while not self.done: + if time.monotonic() > deadline: + raise TimeoutError( + f"job {self.id} still {self.status} after {timeout}s" + ) + time.sleep(poll) + self.poll() + if not self.succeeded: + # `error` only exists on ParseJob; extract reports failure via status. + raise LightOnError( + f"job {self.id} failed: {getattr(self, 'error', None) or self.status}" + ) + return self + @property def done(self) -> bool: """True once the job is terminal, finished, whether it succeeded or failed.""" diff --git a/lighton/verbs/extract.py b/lighton/verbs/extract.py index 4a8d7c9..c13f70f 100644 --- a/lighton/verbs/extract.py +++ b/lighton/verbs/extract.py @@ -52,6 +52,8 @@ def extract( url: str | None = ..., options: dict[str, Any] | None = ..., mode: Literal[ExecMode.ASYNC], + wait: bool = ..., + timeout: float = ..., ) -> ExtractJob: ... def extract( self, @@ -61,6 +63,8 @@ def extract( url: str | None = None, options: dict[str, Any] | None = None, mode: ExecMode = ExecMode.SYNC, + wait: bool = False, + timeout: float = 300.0, ) -> ExtractJobResponse | ExtractJob: """POST /api/v3/extract, extract structured data from a document. @@ -76,13 +80,24 @@ def extract( mode: ExecMode.SYNC (default) runs inline and returns the extracted data. ExecMode.ASYNC queues the job and returns an ``ExtractJob`` handle, call ``.poll()`` until ``.succeeded``. + wait: Async only. Block until the job is terminal, so the returned + ``ExtractJob`` already carries its ``result``. + timeout: Seconds to wait when wait=True before raising TimeoutError. Returns: - ``ExtractJobResponse`` (with data) when sync; a pollable ``ExtractJob`` - when async. + ``ExtractJobResponse`` (with data) when sync; an ``ExtractJob`` when + async, pollable (wait=False) or already finished (wait=True). + + Raises: + ValueError: If not exactly one of path/url is given, or wait=True + without ExecMode.ASYNC (sync already blocks). + TimeoutError: If wait=True and `timeout` elapses first. + LightOnError: If wait=True and the job ends in failure. """ if (path is None) == (url is None): raise ValueError("extract() requires exactly one of 'path' or 'url'") + if wait and mode != ExecMode.ASYNC: + raise ValueError("wait=True only applies to mode=ExecMode.ASYNC") if mode == ExecMode.ASYNC: options = {**(options or {}), "async": True} json_schema = _as_json_schema(schema) @@ -105,5 +120,6 @@ def extract( body["options"] = options resp = self._request("POST", "/api/v3/extract", json=body) if mode == ExecMode.ASYNC: - return ExtractJob._bind(self, "/api/v3/extract", resp) + job = ExtractJob._bind(self, "/api/v3/extract", resp) + return job.wait(timeout) if wait else job return ExtractJobResponse.model_validate(resp) diff --git a/lighton/verbs/parse.py b/lighton/verbs/parse.py index 2bb9ade..f95ac7e 100644 --- a/lighton/verbs/parse.py +++ b/lighton/verbs/parse.py @@ -28,6 +28,8 @@ def parse( path: str | Path | None = ..., url: str | None = ..., mode: Literal[ExecMode.ASYNC], + wait: bool = ..., + timeout: float = ..., ) -> ParseJob: ... def parse( self, @@ -35,6 +37,8 @@ def parse( path: str | Path | None = None, url: str | None = None, mode: ExecMode = ExecMode.SYNC, + wait: bool = False, + timeout: float = 300.0, ) -> ParseResponse | ParseJob: """POST /api/v3/parse, parse a document into per-page text. @@ -46,13 +50,25 @@ def parse( mode: ExecMode.SYNC (default) runs inline and returns the full ``ParseResponse``. ExecMode.ASYNC queues the job and returns a ``ParseJob`` handle, call ``.poll()`` until ``.succeeded``. + wait: Async only. Block until the job is terminal, so the returned + ``ParseJob`` already carries its ``result``. + timeout: Seconds to wait when wait=True before raising TimeoutError. Returns: - ``ParseResponse`` when sync; a pollable ``ParseJob`` when async. + ``ParseResponse`` when sync; a ``ParseJob`` when async, pollable + (wait=False) or already finished (wait=True). + + Raises: + ValueError: If not exactly one of path/url is given, or wait=True + without ExecMode.ASYNC (sync already blocks). + TimeoutError: If wait=True and `timeout` elapses first. + LightOnError: If wait=True and the job ends in failure. """ if (path is None) == (url is None): raise ValueError("parse() requires exactly one of 'path' or 'url'") is_async = mode == ExecMode.ASYNC + if wait and not is_async: + raise ValueError("wait=True only applies to mode=ExecMode.ASYNC") options = {"async": True} if is_async else None if path is not None: path = Path(path) @@ -68,5 +84,6 @@ def parse( body["options"] = options resp = self._request("POST", "/api/v3/parse", json=body) if is_async: - return ParseJob._bind(self, "/api/v3/parse", resp) + job = ParseJob._bind(self, "/api/v3/parse", resp) + return job.wait(timeout) if wait else job return ParseResponse.model_validate(resp) diff --git a/tests/e2e/cli.py b/tests/e2e/cli.py index a58d584..bd2501c 100644 --- a/tests/e2e/cli.py +++ b/tests/e2e/cli.py @@ -26,7 +26,6 @@ from datetime import datetime from pathlib import Path from types import FunctionType -from typing import TypeVar import typer from pydantic import BaseModel, Field @@ -37,10 +36,8 @@ Attribute, ContentType, ExecMode, - ExtractJob, File, LightOn, - ParseJob, RelevanceScoring, Role, SearchMode, @@ -52,7 +49,6 @@ DOCS_DIR = Path(__file__).parent / "documents" JOB_TIMEOUT = 300.0 PREREQS = ("workspace", "upload") # implied by --only: the rest build on them -_JobT = TypeVar("_JobT", ParseJob, ExtractJob) class DocumentSummary(BaseModel): @@ -119,19 +115,6 @@ def _topic(c: Ctx) -> str: return c.topic -def _poll(job: _JobT, timeout: float = JOB_TIMEOUT) -> _JobT: - """Poll an async parse/extract job until it is terminal.""" - deadline = time.monotonic() + timeout - while not job.done: - if time.monotonic() > deadline: - raise TimeoutError(f"job {job.id} still {job.status} after {timeout}s") - time.sleep(2.0) - job.poll() - if not job.succeeded: - raise RuntimeError(f"job {job.id} finished unsuccessfully: {job.status}") - return job - - # --- steps ------------------------------------------------------------------ @@ -296,7 +279,7 @@ def parse(c: Ctx) -> None: assert pages, "sync parse returned no pages" _say(f"sync: {len(pages)} page(s), page 1 is {len(pages[0].markdown)} chars") - job = _poll(c.client.parse(path=doc, mode=ExecMode.ASYNC)) + job = c.client.parse(path=doc, mode=ExecMode.ASYNC, wait=True, timeout=JOB_TIMEOUT) assert job.result and job.result.pages, "async parse returned no pages" _say(f"async: job {job.id} completed in {job.processing_time_ms}ms") @@ -309,7 +292,9 @@ def extract(c: Ctx) -> None: assert r.result and r.result.data, "sync extract returned no data" _say(f"sync: {r.result.data}") - job = _poll(c.client.extract(DocumentSummary, path=doc, mode=ExecMode.ASYNC)) + job = c.client.extract( + DocumentSummary, path=doc, mode=ExecMode.ASYNC, wait=True, timeout=JOB_TIMEOUT + ) assert job.result and job.result.data, "async extract returned no data" _say(f"async: job {job.id} completed in {job.processing_time_ms}ms") diff --git a/tests/test_extract.py b/tests/test_extract.py index 0b5f8bf..8867068 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -8,6 +8,7 @@ from pydantic import BaseModel from lighton import ExecMode, LightOn, LightOnConfiguration +from lighton.exceptions import LightOnError from lighton.utils import validate_response_format_json @@ -125,3 +126,51 @@ def test_extract_malformed_dict_schema_raises(): # "type" must be a string/array, not an int — invalid per the meta-schema with pytest.raises(SchemaError): client.extract({"type": 123}, url="https://x/i.pdf") + + +_DONE = {**_OK, "completed_at": "2026-01-01T00:00:01Z"} + + +def test_extract_async_wait_returns_finished_job(): + def handler(req: httpx.Request) -> httpx.Response: + assert req.method == "POST" # already terminal, no poll needed + return httpx.Response(200, json=_DONE) + + job = make_client(handler).extract( + {"type": "object"}, url="https://x/i.pdf", mode=ExecMode.ASYNC, wait=True + ) + assert job.done and job.succeeded + assert job.result is not None and job.result.data == [{"total": 42}] + + +def test_job_wait_polls_until_done(): + seq = iter([_OK, _OK, _DONE]) # not terminal until completed_at is set + + def handler(req: httpx.Request) -> httpx.Response: + if req.method == "POST": + return httpx.Response(202, json={"id": "e1", "status": "pending"}) + return httpx.Response(200, json=next(seq)) + + job = make_client(handler).extract( + {"type": "object"}, url="https://x/i.pdf", mode=ExecMode.ASYNC + ) + assert job.wait(timeout=5, poll=0) is job and job.done # poll=0 → no real sleep + + +def test_job_wait_raises_on_terminal_failure(): + failed = {"id": "e1", "status": "failed", "completed_at": "2026-01-01T00:00:01Z"} + + def handler(req: httpx.Request) -> httpx.Response: + return httpx.Response(200 if req.method == "GET" else 202, json=failed) + + with pytest.raises(LightOnError, match="failed"): + make_client(handler).extract( + {"type": "object"}, url="https://x/i.pdf", mode=ExecMode.ASYNC, wait=True + ) + + +def test_extract_wait_requires_async_mode(): + client = make_client(lambda req: httpx.Response(200, json=_OK)) + with pytest.raises(ValueError, match="wait=True"): + # wait without ASYNC is also a static error, hence the ignore + client.extract({"type": "object"}, url="https://x/i.pdf", wait=True) # ty: ignore[no-matching-overload] diff --git a/tests/test_parse.py b/tests/test_parse.py index 0854605..af54563 100644 --- a/tests/test_parse.py +++ b/tests/test_parse.py @@ -98,3 +98,21 @@ def handler(req: httpx.Request) -> httpx.Response: assert same is job assert job.succeeded and job.done assert job.result is not None and job.result.pages == [] + + +def test_parse_async_wait_returns_finished_job(): + def handler(req: httpx.Request) -> httpx.Response: + assert req.method == "POST" # already terminal, no poll needed + return httpx.Response(200, json=_parse_body("parse_Kg", 8)) + + job = make_client(handler).parse( + url="https://example.com/d.pdf", mode=ExecMode.ASYNC, wait=True + ) + assert job.done and job.succeeded and job.result is not None + + +def test_parse_wait_requires_async_mode(): + client = make_client(lambda req: httpx.Response(200, json=_parse_body("p1", 8))) + with pytest.raises(ValueError, match="wait=True"): + # wait without ASYNC is also a static error, hence the ignore + client.parse(url="https://example.com/d.pdf", wait=True) # ty: ignore[no-matching-overload]