diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dbde7a2..b8686f7 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|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; }' -- + [ -z "$b" ] || echo "$b" | grep -qE "^(main|master|develop|(feat|feature|fix|bugfix|hotfix|release|chore)/[a-z0-9._/-]+)$" + || { echo "branch \"$b\" must be /, e.g. feat/add-retries (types: feat, 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 75680d6..f2fb0ea 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -60,9 +60,9 @@ 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`): **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. +- **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` XOR `file` (raises `ValueError` unless exactly one; `path` → multipart, the other two → JSON body) plus a `schema` that is **either a pydantic model class** or a **raw JSON-Schema dict**; returns `ExtractJobResponse`. `file` is an **already-ingested** file (`File` or bare id → the API's `file_id`, coerced by `_id()` in `utils.py`, the scalar sibling of `_ids()`), no re-upload; `File` is imported under `TYPE_CHECKING` only, so the annotation costs no import cycle. The multipart `file` **part** still isn't in the OpenAPI schema (`ExtractRequest` models `document`/`file_id`/`schema`/`options`) but the endpoint accepts it, verified by curl; on multipart, `schema`/`options` ride as JSON-encoded form fields alongside it. Note the name collision: the `file=` **param** means file_id, the multipart part is what `path=` sends. `ask` takes the same `schema` (pydantic class or dict) and sends it as the API's `response_format` for structured output; the answer then comes back as JSON **text** in `AskResponse.answer` (the response model is unchanged), and the SDK deliberately does **not** parse it back, callers do `Model.model_validate_json(resp.answer)`, since a model class is only one of the two accepted inputs and re-validating would make the return type depend on which one was passed. Schema handling (in `utils.py`, `as_json_schema()` is the shared entry point for both verbs): **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. +- Deferred: tag/content_type/attribute filters, streaming, add the params when needed. Also unwrapped from the current OpenAPI schema: `POST /api/v3/content-types/scope` (`FacetScopeRequest`/`FacetScopeResponse`, LLM scope inference), `WorkspaceTaxonomy` on the workspace list response, and the `ServiceMaintenance503` body (a 503 maps to the generic error today, no dedicated exception). - **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 *connection* errors (exp. backoff); `_request` itself handles **HTTP 429**, retries up to diff --git a/README.md b/README.md index ea920cf..78bd402 100644 --- a/README.md +++ b/README.md @@ -134,7 +134,8 @@ More on file management (list, fetch, tags, delete) and polling in Four actions live directly on the client. `ask` and `search` query your **indexed** documents, scope them with `workspaces=`, `tags=`, or `files=` (objects or bare ids). -`parse` and `extract` process a document **on the fly**, no indexing required. Full +`parse` and `extract` process a document **on the fly**, no indexing required, +`extract` can also target a file you already ingested, with `file=`. Full reference at [developers.lighton.ai](https://developers.lighton.ai). The per-verb snippets below assume a `client` opened with `with LightOn() as client:`. @@ -157,6 +158,26 @@ for r in resp.results: # the chunks used as grounding print(r.source.filename, r.score) ``` +Pass `schema=` to constrain the answer to **structured output**, same inputs as +`extract` (a pydantic model or a JSON-Schema dict, describing an object). The +answer comes back as JSON *text* in `.answer`, so parse it yourself: + +```python +from pydantic import BaseModel, Field + + +class Revenue(BaseModel): + amount: float = Field(description="Revenue figure, in millions.") + currency: str = Field(description="ISO 4217 code, e.g. 'EUR'.") + quarter: str | None = Field(None, description="Fiscal quarter, or null.") + + +resp = client.ask("What were Q4 revenues?", workspaces=[42], schema=Revenue) +revenue = Revenue.model_validate_json(resp.answer) +print(revenue.amount, revenue.currency) +print(resp.results) # sources still come back alongside +``` + ### `search`: retrieval only, no generation Hybrid semantic + lexical retrieval that returns ranked chunks with scores, source @@ -206,6 +227,7 @@ per page. See [Extract](#extract) below for the full schema guide. ```python resp = client.extract(schema=InvoiceModel, path="invoice.pdf") +# or, on a file already in your index: client.extract(schema=InvoiceModel, file=f) print(resp.result.data) ``` @@ -285,7 +307,9 @@ with LightOn() as client: doc.delete() ``` -Once a file reaches `embedded`, it's retrievable by `ask`/`search`. +Once a file reaches `embedded`, it's retrievable by `ask`/`search`. You can also +run `extract` straight on it, `client.extract(schema=Invoice, file=doc)`, instead +of uploading the document a second time (see [Extract](#extract)). ## Async jobs & polling @@ -359,10 +383,16 @@ for page in job.result.pages: ## Extract -`extract(schema, *, path | url)` pulls structured data from a document, pass a -local `path` to upload (multipart) or a public `url` to fetch, exactly one (same -as `parse`). The `schema` drives guided generation and can be **a pydantic -model** or a **raw JSON-Schema dict**, use whichever you have. +`extract(schema, *, path | url | file)` pulls structured data from a document. +Pass exactly one source: + +- `path=`: a local file, uploaded multipart +- `url=`: a publicly accessible URL the server fetches +- `file=`: a file **already ingested** into your index (a `File` or a bare id), + no re-upload, the cheap option when the document is already there + +The `schema` drives guided generation and can be **a pydantic model** or a **raw +JSON-Schema dict**, use whichever you have. A pydantic model is the easy path: nested models, `list[...]`, and `X | None` fields all convert to a valid vLLM `response_format` schema for you. @@ -404,7 +434,7 @@ with LightOn() as client: Or pass the schema dict directly, it's validated against the JSON-Schema meta-schema (raises `jsonschema.SchemaError` if malformed), then normalized the -same way a model is — the endpoint rejects `$ref`, so `$defs`/`$ref` are inlined +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: diff --git a/lighton/types/api/__init__.py b/lighton/types/api/__init__.py index 785b40c..caa3911 100644 --- a/lighton/types/api/__init__.py +++ b/lighton/types/api/__init__.py @@ -445,7 +445,7 @@ class FileCreateRequestSerializerV3(BaseModel): - name: Custom filename (optional, defaults to uploaded filename) - title: Custom title for the document (optional) - workspace_id: Workspace ID where the document will be stored (required) - - parser: Deprecated — ignored, the platform always uses its default pipeline + - parser: Deprecated: ignored, the platform always uses its default pipeline """ file: Annotated[AnyUrl, Field(description="The file to upload (binary data)")] diff --git a/lighton/utils.py b/lighton/utils.py index cf6eede..8585f2d 100644 --- a/lighton/utils.py +++ b/lighton/utils.py @@ -16,11 +16,16 @@ def _compact(**kw: Any) -> dict[str, Any]: return {k: v for k, v in kw.items() if v is not None} +def _id(item: int | Any) -> int: + """Coerce a resource or an int to its id (duck-typed on `.id`).""" + return item if isinstance(item, int) else item.id + + def _ids(items: list[int] | list[Any] | None) -> list[int] | None: """Coerce a list of resources or ints to a list of ids (duck-typed on `.id`).""" if items is None: return None - return [x if isinstance(x, int) else x.id for x in items] + return [_id(x) for x in items] def _inline_refs(node: Any, defs: dict[str, Any]) -> Any: @@ -127,3 +132,30 @@ def convert_pydantic_to_response_format_json(model: type[BaseModel]) -> dict[str A self-contained JSON Schema dict suitable for vLLM guided generation. """ return normalize_response_format_json(model.model_json_schema()) + + +def as_json_schema(schema: type[BaseModel] | dict[str, Any]) -> dict[str, Any]: + """Either guided-generation input → the self-contained schema to send. + + Shared by `extract` (`schema`, the extraction target) and `ask` (`schema` → + `response_format`, constraining the answer). 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 endpoints reject `$ref`, and a + dict hand-built from `model_json_schema()` carries them just as a class does. + + Args: + schema: A pydantic model class or a dict holding a JSON Schema. + + Returns: + A self-contained JSON Schema dict, free of `$defs`/`$ref`. + + Raises: + jsonschema.exceptions.SchemaError: If a dict schema is malformed. + TypeError: If `schema` is neither a dict nor a pydantic model class. + """ + if isinstance(schema, dict): + 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/lighton/verbs/ask.py b/lighton/verbs/ask.py index 53b92e9..cd9c383 100644 --- a/lighton/verbs/ask.py +++ b/lighton/verbs/ask.py @@ -2,12 +2,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast + +from pydantic import BaseModel from lighton.enums import RelevanceScoring from lighton.tag import resolve_ids from lighton.types.api import AskResponse -from lighton.utils import _compact, _ids +from lighton.utils import _compact, _ids, as_json_schema from lighton.verbs._base import _VerbClient if TYPE_CHECKING: @@ -30,6 +32,7 @@ def ask( max_results: int | None = None, relevance_scoring: RelevanceScoring | None = None, model: str | None = None, + schema: type[BaseModel] | dict[str, Any] | None = None, ) -> AskResponse: """POST /api/v3/ask, ask a grounded question over indexed documents. @@ -46,6 +49,11 @@ def ask( relevance_scoring: RelevanceScoring, .scoring_and_filtering (default), .scoring_only, or .none. model: LLM for answer generation; platform default if omitted. + schema: Constrain the answer to structured output, a pydantic model + class or a JSON-Schema dict (same inputs as `extract`, sent as + the API's `response_format`; must describe an object). The answer + then comes back as JSON *text* in `.answer`, parse it with + `YourModel.model_validate_json(resp.answer)`. Returns: The answer plus the ranked results used as context. @@ -59,6 +67,7 @@ def ask( max_results=max_results, relevance_scoring=relevance_scoring, model=model, + response_format=as_json_schema(schema) if schema is not None else None, ) return AskResponse.model_validate( self._request("POST", "/api/v3/ask", json=body) diff --git a/lighton/verbs/extract.py b/lighton/verbs/extract.py index aa15c0a..6cf0c12 100644 --- a/lighton/verbs/extract.py +++ b/lighton/verbs/extract.py @@ -4,35 +4,18 @@ import json from pathlib import Path -from typing import Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, overload from pydantic import BaseModel from lighton.enums import ExecMode from lighton.job import ExtractJob 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.utils import _id, as_json_schema from lighton.verbs._base import _VerbClient - -def _as_json_schema(schema: type[BaseModel] | dict[str, Any]) -> dict[str, Any]: - """Either input → the self-contained vLLM guided-generation schema. - - 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 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") +if TYPE_CHECKING: + from lighton.file import File class ExtractMixin(_VerbClient): @@ -43,6 +26,7 @@ def extract( *, path: str | Path | None = ..., url: str | None = ..., + file: File | int | None = ..., options: dict[str, Any] | None = ..., mode: Literal[ExecMode.SYNC] = ..., ) -> ExtractJobResponse: ... @@ -53,6 +37,7 @@ def extract( *, path: str | Path | None = ..., url: str | None = ..., + file: File | int | None = ..., options: dict[str, Any] | None = ..., mode: Literal[ExecMode.ASYNC], wait: bool = ..., @@ -64,6 +49,7 @@ def extract( *, path: str | Path | None = None, url: str | None = None, + file: File | int | None = None, options: dict[str, Any] | None = None, mode: ExecMode = ExecMode.SYNC, wait: bool = False, @@ -74,6 +60,8 @@ def extract( Pass exactly one of: path: A local file to upload (multipart). url: A publicly accessible URL to fetch. + file: An already-ingested file (File object or id), no re-upload, + the server reads the document it already has. Args: schema: The guided-generation schema driving extraction, either a @@ -92,18 +80,20 @@ def extract( async, pollable (wait=False) or already finished (wait=True). Raises: - ValueError: If not exactly one of path/url is given, or wait=True + ValueError: If not exactly one of path/url/file 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 sum(x is not None for x in (path, url, file)) != 1: + raise ValueError( + "extract() requires exactly one of 'path', 'url' or 'file'" + ) 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) + json_schema = as_json_schema(schema) if path is not None: path = Path(path) # multipart: schema/options ride as JSON-encoded form fields. @@ -118,7 +108,8 @@ def extract( data=data, ) else: - body: dict[str, Any] = {"document": url, "schema": json_schema} + source = {"document": url} if url is not None else {"file_id": _id(file)} + body: dict[str, Any] = {**source, "schema": json_schema} if options is not None: body["options"] = options resp = self._request("POST", "/api/v3/extract", json=body) diff --git a/tests/e2e/cli.py b/tests/e2e/cli.py index e99cb44..228d475 100644 --- a/tests/e2e/cli.py +++ b/tests/e2e/cli.py @@ -78,6 +78,13 @@ class DocumentOutline(BaseModel): sections: list[Section] = Field(description="Every top-level section heading.") +class GroundedAnswer(BaseModel): + """`ask(schema=...)` structured output: the LLM answer is constrained to this.""" + + answer: str = Field(description="The answer, in one or two sentences.") + confident: bool = Field(description="True if the sources fully support it.") + + @dataclass class Ctx: client: LightOn @@ -278,7 +285,7 @@ def search(c: Ctx) -> None: @step def ask(c: Ctx) -> None: - """grounded answer over the workspace.""" + """grounded answer over the workspace → structured answer via schema=.""" query = c.ask_query or f"What does the document say about {_topic(c)}?" r = c.client.ask( query, @@ -289,6 +296,11 @@ def ask(c: Ctx) -> None: assert r.answer, f"ask({query!r}) returned an empty answer" _say(f"answer ({len(r.results)} source chunk(s)): {r.answer[:160]}") + # structured output: the answer comes back as JSON text matching the schema. + s = c.client.ask(query, workspaces=[c.workspace()], schema=GroundedAnswer) + parsed = GroundedAnswer.model_validate_json(s.answer) # raises if off-schema + _say(f"structured: confident={parsed.confident} {parsed.answer[:120]}") + @step def parse(c: Ctx) -> None: @@ -305,7 +317,7 @@ def parse(c: Ctx) -> None: @step def extract(c: Ctx) -> None: - """sync → async job (flat schema) → nested schema, as a model and as a raw dict.""" + """sync → async job → nested schema (model + raw dict) → an ingested file=.""" doc = c.docs[0] r = c.client.extract(DocumentSummary, path=doc) assert r.result and r.result.data, "sync extract returned no data" @@ -323,13 +335,18 @@ def extract(c: Ctx) -> None: _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" + 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}") + # file=: extract from the already-ingested file, no re-upload + by_id = c.client.extract(DocumentSummary, file=c.uploaded()) + assert by_id.result and by_id.result.data, "extract by file_id returned no data" + _say(f"by file_id {c.uploaded().id}: {by_id.result.data}") + @step def batch(c: Ctx) -> None: diff --git a/tests/test_ask.py b/tests/test_ask.py index 22b391e..25ece74 100644 --- a/tests/test_ask.py +++ b/tests/test_ask.py @@ -3,11 +3,18 @@ import json import httpx +from pydantic import BaseModel from lighton import LightOn, LightOnConfiguration, Tag, Workspace from lighton.enums import RelevanceScoring +class Clause(BaseModel): + """A sub-model, so the generated schema has a reference to inline.""" + + text: str + + def make_client(handler) -> LightOn: return LightOn( "k", @@ -77,3 +84,37 @@ def handler(req: httpx.Request) -> httpx.Response: # names are resolved via Tag.list, mixed with a bare id make_client(handler).ask("q", tags=["legal", 4]) assert seen["body"] == {"query": "q", "tag_id": [4, 3]} + + +def test_ask_structured_output_sends_response_format(): + seen = {} + + class Verdict(BaseModel): + outcome: str + clauses: list[Clause] + + def handler(req: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(req.content) + return httpx.Response(200, json={"results": [], "answer": '{"outcome": "ok"}'}) + + resp = make_client(handler).ask("did it pass?", schema=Verdict) + fmt = seen["body"]["response_format"] + # normalized like extract's: nested sub-model inlined, no $ref for the API to reject + assert "$defs" not in fmt and "$ref" not in json.dumps(fmt) + assert fmt["properties"]["clauses"]["items"]["properties"]["text"] == { + "title": "Text", + "type": "string", + } + # the answer is JSON *text*, the caller parses it + assert json.loads(resp.answer) == {"outcome": "ok"} + + +def test_ask_without_schema_omits_response_format(): + seen = {} + + def handler(req: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(req.content) + return httpx.Response(200, json={"results": [], "answer": ""}) + + make_client(handler).ask("q") + assert "response_format" not in seen["body"] diff --git a/tests/test_extract.py b/tests/test_extract.py index 3d4f660..945d26e 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -7,7 +7,7 @@ from jsonschema.exceptions import SchemaError from pydantic import BaseModel -from lighton import ExecMode, LightOn, LightOnConfiguration +from lighton import ExecMode, File, LightOn, LightOnConfiguration from lighton.exceptions import LightOnError from lighton.utils import ( convert_pydantic_to_response_format_json, @@ -213,3 +213,28 @@ def test_extract_wait_requires_async_mode(): 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] + + +def test_extract_by_file_id_sends_json(): + seen = {} + raw = {"type": "object", "properties": {"total": {"type": "number"}}} + + def handler(req: httpx.Request) -> httpx.Response: + seen["body"] = json.loads(req.content) + return httpx.Response(200, json=_OK) + + client = make_client(handler) + # an already-ingested file: File object or bare id, no re-upload + client.extract(raw, file=File(id=7, filename="doc.pdf")) + assert seen["body"]["file_id"] == 7 and "document" not in seen["body"] + client.extract(raw, file=7) + assert seen["body"]["file_id"] == 7 + + +def test_extract_requires_exactly_one_source_of_three(): + client = make_client(lambda req: httpx.Response(200, json=_OK)) + raw = {"type": "object"} + with pytest.raises(ValueError, match="exactly one"): + client.extract(raw, url="https://x/i.pdf", file=7) + with pytest.raises(ValueError, match="exactly one"): + client.extract(raw, path="d.png", file=7)