Skip to content

Add pydantic model support for extraction schemas - #77

Closed
jordanalexmeyer wants to merge 3 commits into
masterfrom
pydantic-extraction-schemas
Closed

Add pydantic model support for extraction schemas#77
jordanalexmeyer wants to merge 3 commits into
masterfrom
pydantic-extraction-schemas

Conversation

@jordanalexmeyer

Copy link
Copy Markdown
Contributor

Summary

Adds typed extraction schemas via pydantic, mirroring the TypeScript SDK's Zod support. Pass a pydantic.BaseModel subclass as config["schema"] and the SDK:

  1. Converts the model to Extend's JSON Schema subset for the request (primitives forced nullable, all properties required, enums get a null option, additionalProperties: false)
  2. Validates the extraction output back into model instances, returned as a TypedExtractRun[Model] — a step beyond the TS SDK, which only casts at compile time
from typing import List, Optional
from pydantic import BaseModel, Field
from extend_ai import Extend, ExtendCurrency, ExtendDate

class Invoice(BaseModel):
    invoice_number: Optional[str] = Field(None, description="The invoice number")
    invoice_date: ExtendDate = Field(None, description="The invoice date")
    total: Optional[ExtendCurrency] = Field(None, description="Total amount due")

client = Extend(token="...")
result = client.extract(file={"url": "https://example.com/invoice.pdf"}, config={"schema": Invoice})

result.output.value.invoice_number   # str | None — a validated Invoice instance
result.output.value.invoice_date     # datetime.date | None
result.output.value.total.amount     # float | None

Integration points (full TS parity)

  • client.extract() (sync + async)
  • client.extract_runs.create_and_poll() (sync + async), including extractor["override_config"]["schema"]
  • client.extractors.create() / update() (new wrapper clients)
  • client.extractor_versions.create() (new wrapper client)

@overload signatures give full static inference — mypy resolves result.output.value to the schema model, while untyped (dict) configs keep returning plain ExtractRun.

Details

  • New wrapper/schema/ module: ExtendDate / ExtendCurrency / ExtendSignature field types (plain datetime.date annotations also map to the date type), pydantic_to_extend_schema() conversion, SchemaConversionError carrying the field path, typed run wrappers
  • Supported field types: str, float, int, bool, datetime.date, Literal[...] / string enums, nested models, and lists of these; unsupported constructs (unions, dicts, datetime.datetime) raise SchemaConversionError
  • Works with both pydantic v1 and v2 (via the existing IS_PYDANTIC_V2 compat pattern)
  • All code lives under the .fernignore-protected wrapper/ directory
  • README gains a "Typed extraction with Pydantic" section

Out of scope (matching TS): extend:name, per-enum-value extend:descriptions, typed schemas on plain extract_runs.create().

Testing

  • Ported the TS schema.test.ts conversion suite (~45 cases) plus typed-extraction integration tests for all four endpoints
  • pytest tests/wrapper tests/custom: 212 passed (1 pre-existing skip)
  • Schema tests also verified against pydantic 1.10.13 in a separate environment
  • mypy .: clean (1,134 files); ruff clean on all touched files

Pass a pydantic BaseModel subclass as config["schema"] (or
extractor["override_config"]["schema"]) and the SDK converts it to
Extend's JSON Schema subset for the request, then validates the
extraction output back into model instances (TypedExtractRun).

Mirrors the TypeScript SDK's Zod support across the same four
integration points: extract(), extract_runs.create_and_poll(),
extractors.create()/update(), and extractor_versions.create().

- New wrapper/schema module: ExtendDate/ExtendCurrency/ExtendSignature
  field types, pydantic_to_extend_schema conversion with
  SchemaConversionError, typed run wrappers, config detection helpers
- Works with pydantic v1 and v2
- Overload signatures give full static inference:
  result.output.value is typed as the schema model
The wrapper layer re-declares parts of the generated API surface, which
can silently go stale when Fern regenerates the SDK with new parameters.
True overrides (extract(), extractors.create()/update(), etc.) are
already protected because mypy rejects overrides incompatible with the
generated superclass, but create_and_poll(), the typed config
TypedDicts, and TypedExtractRun had no guard.

These tests fail CI whenever a generated create() gains a parameter
that create_and_poll() doesn't forward, a config key is missing from
the typed TypedDicts, or ExtractRun gains a field TypedExtractRun
doesn't mirror.

Also fixes drift the new tests caught: parse_runs.create_and_poll()
was missing the metadata and data_retention parameters that the
generated parse_runs.create() accepts.
Stop silently force-nullabling primitives in the converter — that was
defeating the API's 2026-02-09 strict schema validation and could defer
user mistakes until after a paid run completed. Non-Optional primitives,
enums, and dates now raise SchemaConversionError before any request is
sent. Also:

- Detect recursive models and raise SchemaConversionError instead of a
  fatal stack overflow
- Recognize typing.Literal on Python 3.8 (distinct from typing_extensions)
- Reject field aliases (they caused silent None validation) and
  Optional array items
- Wrap residual output validation failures in ExtractOutputValidationError
  that preserves the completed ExtractRun (id, dashboard URL, raw output)
- Convert pydantic schemas in plain extract_runs.create() so users don't
  hit a cryptic encoder error
- Bump typing_extensions floor to >=4.3.0 for generic TypedDict support
@jordanalexmeyer

Copy link
Copy Markdown
Contributor Author

Superseded by #78, which stacks these pydantic commits on top of a fresh 1.17.0 Fern regeneration from documentation main.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant