diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9cb7b7e..dbde7a2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,8 +27,8 @@ repos: name: conventional branch name entry: >- bash -c 'b=$(git branch --show-current); - [ -z "$b" ] || echo "$b" | grep -qE "^(main|master|develop|(feature|bugfix|hotfix|release|chore)/[a-z0-9._/-]+)$" - || { echo "branch \"$b\" must be /, e.g. feature/add-retries (types: feature, bugfix, hotfix, release, chore)"; exit 1; }' -- + [ -z "$b" ] || echo "$b" | grep -qE "^(main|master|develop|(feature|fix|bugfix|hotfix|release|chore)/[a-z0-9._/-]+)$" + || { echo "branch \"$b\" must be /, e.g. feature/add-retries (types: feature, fix, bugfix, hotfix, release, chore)"; exit 1; }' -- language: system pass_filenames: false always_run: true # branch name is not tied to any changed file diff --git a/AGENTS.md b/AGENTS.md index 0ec021e..75680d6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -60,7 +60,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. +- **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`): **both inputs end up normalized by `normalize_response_format_json`** — `$defs`/`$ref` inlined (`_inline_refs`), nullable `anyOf` collapsed to `type: [X, "null"]` (`_collapse_nullable`), draft-2020-12 `$schema` marker added (an existing one is kept). A pydantic model goes `model_json_schema()` → normalize (`convert_pydantic_to_response_format_json`); a dict is first validated against the draft-2020-12 meta-schema via `jsonschema` (`validate_response_format_json`, raises `SchemaError`) and then normalized too — **not** passed through as it used to be, because the endpoint 422s on `$ref` and a dict is usually just someone's own `model_json_schema()` call, which carries them (the original bug: nested models only worked via the model-class path). A `#/$defs/` ref with no target raises `SchemaError` rather than a bare `KeyError` from inside the recursion. `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. `_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. diff --git a/README.md b/README.md index 38e59e8..ea920cf 100644 --- a/README.md +++ b/README.md @@ -403,7 +403,10 @@ with LightOn() as client: ``` Or pass the schema dict directly, it's validated against the JSON-Schema -meta-schema (raises `jsonschema.SchemaError` if malformed) and sent as-is: +meta-schema (raises `jsonschema.SchemaError` if malformed), then normalized the +same way a model is — the endpoint rejects `$ref`, so `$defs`/`$ref` are inlined +whether the schema came from a model class or from your own +`Model.model_json_schema()` call: ```python with LightOn() as client: diff --git a/lighton/utils.py b/lighton/utils.py index cb2ae05..cf6eede 100644 --- a/lighton/utils.py +++ b/lighton/utils.py @@ -5,6 +5,7 @@ from typing import Any from jsonschema import Draft202012Validator +from jsonschema.exceptions import SchemaError from pydantic import BaseModel _DRAFT = "https://json-schema.org/draft/2020-12/schema" @@ -36,7 +37,10 @@ def _inline_refs(node: Any, defs: dict[str, Any]) -> Any: if isinstance(node, dict): ref = node.get("$ref") if isinstance(ref, str) and ref.startswith("#/$defs/"): - target = defs[ref.rsplit("/", 1)[-1]] + name = ref.rsplit("/", 1)[-1] + if name not in defs: + raise SchemaError(f"unresolved $ref {ref!r}: no such entry in $defs") + target = defs[name] siblings = { k: _inline_refs(v, defs) for k, v in node.items() if k != "$ref" } @@ -71,8 +75,8 @@ def _collapse_nullable(node: Any) -> Any: def validate_response_format_json(schema: dict[str, Any]) -> dict[str, Any]: """Validate a raw response_format schema against the draft-2020-12 meta-schema. - For dict schemas passed straight through to vLLM (no pydantic model to vouch - for them), this catches a malformed schema client-side instead of at the API. + For dict schemas handed to vLLM without a pydantic model to vouch for them, + this catches a malformed schema client-side instead of at the API. Args: schema: A dict holding a JSON Schema. @@ -87,12 +91,34 @@ def validate_response_format_json(schema: dict[str, Any]) -> dict[str, Any]: return schema +def normalize_response_format_json(schema: dict[str, Any]) -> dict[str, Any]: + """Normalize a JSON Schema into the self-contained shape vLLM wants. + + `$defs`/`$ref` inlined, nullable `anyOf` collapsed to `type: [X, "null"]`, + draft-2020-12 `$schema` marker added (an existing one is kept). The endpoint + rejects `$ref`, so every schema goes through here, whether it came from a + pydantic model or was passed in as a dict. + + Args: + schema: A dict holding a JSON Schema, possibly with `$defs`/`$ref`. + + Returns: + An equivalent self-contained schema, free of `$defs`/`$ref`. + + Raises: + jsonschema.exceptions.SchemaError: If a `#/$defs/` ref has no target. + """ + defs = schema.get("$defs", {}) + inlined = _inline_refs({k: v for k, v in schema.items() if k != "$defs"}, defs) + return {"$schema": _DRAFT, **_collapse_nullable(inlined)} + + def convert_pydantic_to_response_format_json(model: type[BaseModel]) -> dict[str, Any]: """Convert a pydantic model class to a vLLM guided-generation `response_format` schema. - Runs `model_json_schema()`, then normalizes: `$defs`/`$ref` inlined into a - self-contained schema, nullable `anyOf` collapsed to `type: [X, "null"]`, and - the draft-2020-12 `$schema` marker added. + Runs `model_json_schema()` through `normalize_response_format_json`, which a + nested model needs: pydantic emits `$defs`/`$ref` for every sub-model and the + endpoint rejects those. Args: model: The pydantic model class describing the extraction target. @@ -100,7 +126,4 @@ def convert_pydantic_to_response_format_json(model: type[BaseModel]) -> dict[str Returns: A self-contained JSON Schema dict suitable for vLLM guided generation. """ - raw = model.model_json_schema() - defs = raw.get("$defs", {}) - inlined = _inline_refs({k: v for k, v in raw.items() if k != "$defs"}, defs) - return {"$schema": _DRAFT, **_collapse_nullable(inlined)} + return normalize_response_format_json(model.model_json_schema()) diff --git a/lighton/verbs/extract.py b/lighton/verbs/extract.py index c13f70f..aa15c0a 100644 --- a/lighton/verbs/extract.py +++ b/lighton/verbs/extract.py @@ -13,20 +13,23 @@ from lighton.types.api import ExtractJobResponse from lighton.utils import ( convert_pydantic_to_response_format_json, + normalize_response_format_json, validate_response_format_json, ) from lighton.verbs._base import _VerbClient def _as_json_schema(schema: type[BaseModel] | dict[str, Any]) -> dict[str, Any]: - """A pydantic model class → a vLLM guided-generation schema; a dict is validated. + """Either input → the self-contained vLLM guided-generation schema. - A dict is validated against the JSON-Schema meta-schema (raises on malformed) - and otherwise returned untouched. A pydantic model is converted via - `convert_pydantic_to_response_format_json`. + A dict is validated against the JSON-Schema meta-schema (raises on malformed), + then normalized; a pydantic model class is converted, which normalizes too. + Both go through `normalize_response_format_json` because the endpoint rejects + `$ref`, and a dict hand-built from `model_json_schema()` carries them just as + a model class does. """ if isinstance(schema, dict): - return validate_response_format_json(schema) + return normalize_response_format_json(validate_response_format_json(schema)) if isinstance(schema, type) and issubclass(schema, BaseModel): return convert_pydantic_to_response_format_json(schema) raise TypeError("schema must be a pydantic BaseModel subclass or a dict") diff --git a/tests/e2e/cli.py b/tests/e2e/cli.py index bd2501c..e99cb44 100644 --- a/tests/e2e/cli.py +++ b/tests/e2e/cli.py @@ -52,13 +52,32 @@ class DocumentSummary(BaseModel): - """Doc-agnostic extraction schema (see --extract-schema in the module docs).""" + """Doc-agnostic extraction schema, flat: no sub-models, so no `$defs`/`$ref`.""" title: str = Field(description="Document title.") summary: str = Field(description="One-sentence summary of the document.") language: str = Field(description="Primary language, as an ISO 639-1 code.") +class Section(BaseModel): + """A section heading; sub-model of `DocumentOutline`.""" + + heading: str = Field(description="Section heading, verbatim as written.") + page: int | None = Field(None, description="Page it starts on; null if unclear.") + + +class DocumentOutline(BaseModel): + """Nested schema: `model_json_schema()` emits `$defs`/`$ref` for both sub-models. + + The API 422s on `$ref`, so this only reaches it because the SDK inlines them. + Reuses `DocumentSummary` as a sub-model on purpose: the same model then appears + both nested and standalone. + """ + + overview: DocumentSummary = Field(description="Summary of the whole document.") + sections: list[Section] = Field(description="Every top-level section heading.") + + @dataclass class Ctx: client: LightOn @@ -286,7 +305,7 @@ def parse(c: Ctx) -> None: @step def extract(c: Ctx) -> None: - """sync extract → async extract job (schema: DocumentSummary).""" + """sync → async job (flat schema) → nested schema, as a model and as a raw dict.""" doc = c.docs[0] r = c.client.extract(DocumentSummary, path=doc) assert r.result and r.result.data, "sync extract returned no data" @@ -298,6 +317,19 @@ def extract(c: Ctx) -> None: 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") + # A 422 on either call means $ref reached the API: the SDK stopped inlining. + nested = c.client.extract(DocumentOutline, path=doc) + assert nested.result and nested.result.data, "nested extract returned no data" + _say(f"nested (model class): {nested.result.data}") + + raw = DocumentOutline.model_json_schema() # carries $defs/$ref verbatim + assert "$defs" in raw, "pydantic stopped emitting $defs — this case is now moot" + as_dict = c.client.extract(raw, path=doc) + assert as_dict.result and as_dict.result.data, ( + "nested dict extract returned no data" + ) + _say(f"nested (raw dict): {as_dict.result.data}") + @step def batch(c: Ctx) -> None: diff --git a/tests/test_extract.py b/tests/test_extract.py index 8867068..3d4f660 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -1,4 +1,4 @@ -"""extract verb: pydantic → vLLM schema, raw-dict passthrough, url vs path upload.""" +"""extract verb: pydantic/dict → vLLM schema, url vs path upload.""" import json @@ -9,7 +9,12 @@ from lighton import ExecMode, LightOn, LightOnConfiguration from lighton.exceptions import LightOnError -from lighton.utils import validate_response_format_json +from lighton.utils import ( + convert_pydantic_to_response_format_json, + validate_response_format_json, +) + +_DRAFT = "https://json-schema.org/draft/2020-12/schema" def make_client(handler) -> LightOn: @@ -68,10 +73,44 @@ def handler(req: httpx.Request) -> httpx.Response: make_client(handler).extract(raw, url="https://x/i.pdf", options={"async": False}) assert seen["body"]["document"] == "https://x/i.pdf" - assert seen["body"]["schema"] == raw # dict passed through untouched + assert seen["body"]["schema"] == {"$schema": _DRAFT, **raw} assert seen["body"]["options"] == {"async": False} +def test_extract_dict_schema_refs_are_inlined(): + """A dict built from model_json_schema() carries $defs/$ref; the API rejects them.""" + + class Address(BaseModel): + city: str + + class Company(BaseModel): + name: str + addr: Address + sites: list[Address] = [] + + seen = {} + + def handler(req: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(req.content) + return httpx.Response(200, json=_OK) + + make_client(handler).extract(Company.model_json_schema(), url="https://x/i.pdf") + schema = seen["body"]["schema"] + assert "$defs" not in schema and "$ref" not in json.dumps(schema) + city = {"city": {"title": "City", "type": "string"}} + assert schema["properties"]["addr"]["properties"] == city + assert schema["properties"]["sites"]["items"]["properties"] == city + # same result as handing the model class over directly + assert schema == convert_pydantic_to_response_format_json(Company) + + +def test_extract_dict_schema_with_dangling_ref_raises(): + client = make_client(lambda req: httpx.Response(200, json=_OK)) + raw = {"type": "object", "properties": {"addr": {"$ref": "#/$defs/Nope"}}} + with pytest.raises(SchemaError, match="unresolved"): + client.extract(raw, url="https://x/i.pdf") + + def test_extract_path_sends_multipart(tmp_path): raw = {"type": "object", "properties": {"total": {"type": "number"}}} f = tmp_path / "doc.png" @@ -87,7 +126,7 @@ def handler(req: httpx.Request) -> httpx.Response: assert seen["ctype"].startswith("multipart/form-data") # file part + schema/options as JSON-encoded form fields assert b'name="file"' in seen["body"] and b"doc.png" in seen["body"] - assert json.dumps(raw).encode() in seen["body"] + assert json.dumps({"$schema": _DRAFT, **raw}).encode() in seen["body"] assert b'{"async": false}' in seen["body"]