fix(extract): resolve refs into plain json schema (#8)

Author: MightyGoldenOctopus

.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 <type>/<description>, 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 <type>/<description>, 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

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 `<path>/<id>`, 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.

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:

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,20 +91,39 @@ 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.
 
     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())

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")

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:

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"]