diff --git a/README.md b/README.md index 7710dce..00e995f 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,77 @@ edit_run = client.edit( > **Note:** The synchronous methods above have a 5-minute timeout and are best suited for onboarding and testing. For production workloads, use [polling helpers](#polling-helpers) or [webhooks](#webhook-verification) instead. +## Typed extraction with Pydantic + +The SDK supports [pydantic](https://docs.pydantic.dev/) models for fully typed extraction -- define your schema once and get end-to-end type safety from request to response: + +```python +from typing import List, Optional + +from pydantic import BaseModel, Field + +from extend_ai import Extend, ExtendCurrency, ExtendDate + +class LineItem(BaseModel): + description: Optional[str] = None + amount: Optional[ExtendCurrency] = None + +class Invoice(BaseModel): + invoice_number: Optional[str] = Field(None, description="The invoice number") + invoice_date: ExtendDate = Field(None, description="The invoice date") + line_items: List[LineItem] = Field(default_factory=list, description="Line items on the invoice") + total: Optional[ExtendCurrency] = Field(None, description="Total amount due") + +client = Extend(token="YOUR_API_KEY") + +result = client.extract( + file={"url": "https://example.com/invoice.pdf"}, + config={"schema": Invoice}, +) + +# output.value is a validated Invoice instance +if result.status == "PROCESSED" and result.output is not None: + invoice = result.output.value + print(invoice.invoice_number) # str | None + print(invoice.invoice_date) # datetime.date | None + if invoice.total is not None: + print(invoice.total.amount) # float | None + print(invoice.total.iso_4217_currency_code) # str | None +``` + +The model is converted to [Extend's JSON Schema format](https://docs.extend.ai/2026-02-09/extraction/schema) for the request, and the extraction output is validated back into model instances. Use `Field(description=...)` to guide the extraction. + +Primitive, enum, and date fields must be declared `Optional` -- extraction can return `null` for any field, so a non-Optional field raises `SchemaConversionError` before any request is sent. In the unlikely event that a completed run's output fails model validation, the SDK raises `ExtractOutputValidationError`, which preserves the completed run (including its raw output) on the error's `run` attribute. + +Pydantic model schemas are accepted everywhere an extraction schema can be provided: + +```python +# Polling (see below), including extractor config overrides +result = client.extract_runs.create_and_poll( + file={"url": "https://example.com/invoice.pdf"}, + config={"schema": Invoice}, +) + +# Creating and updating extractors +extractor = client.extractors.create(name="Invoice Extractor", config={"schema": Invoice}) +client.extractors.update(extractor.id, config={"schema": Invoice}) + +# Publishing extractor versions +client.extractor_versions.create(extractor.id, release_type="major", config={"schema": Invoice}) +``` + +### Custom field types + +The SDK provides field types for Extend-specific extraction behavior: + +| Type | Output type | Description | +|---|---|---| +| `ExtendDate` | `datetime.date \| None` | ISO date (plain `datetime.date` annotations work too) | +| `ExtendCurrency` | `ExtendCurrency(amount, iso_4217_currency_code)` | Currency with amount and code | +| `ExtendSignature` | `ExtendSignature(printed_name, signature_date, is_signed, title_or_role)` | Signature detection | + +Supported field types: `Optional[str]`, `Optional[float]`, `Optional[int]`, `Optional[bool]`, `Optional[datetime.date]`, `Optional[Literal[...]]` / string enums (converted to nullable enums), nested models, and lists of these (list items are non-Optional, e.g. `List[str]`). Unsupported constructs (non-Optional unions, dicts, recursive models, field aliases, etc.) raise `SchemaConversionError`. + ## Polling helpers Every run resource exposes a `create_and_poll()` method that creates the run and automatically polls until it reaches a terminal state (`PROCESSED`, `FAILED`, or `CANCELLED`): diff --git a/poetry.lock b/poetry.lock index 4bc6ae7..9992734 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.5.1 and should not be changed by hand. [[package]] name = "annotated-types" @@ -596,4 +596,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "bcf31a142c86d9e556553c8c260a93b563ac64a043076dbd48b26111d422c26e" +content-hash = "f1e2ab5eda730c574ef9098656ce135835405a944b8ad9dcd4c65ef061581e5f" diff --git a/pyproject.toml b/pyproject.toml index 3b3b12a..58b91bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,8 @@ python = "^3.8" httpx = ">=0.21.2" pydantic = ">= 1.9.2" pydantic-core = ">=2.18.2" -typing_extensions = ">= 4.0.0" +# Generic TypedDicts (used by the typed extraction wrappers) require >= 4.3.0 +typing_extensions = ">= 4.3.0" [tool.poetry.group.dev.dependencies] mypy = "==1.13.0" diff --git a/src/extend_ai/__init__.py b/src/extend_ai/__init__.py index 1fb7b15..71ae875 100644 --- a/src/extend_ai/__init__.py +++ b/src/extend_ai/__init__.py @@ -438,7 +438,22 @@ ClassifyRunsListResponse, ClassifyRunsListResponseParams, ) - from .wrapper import AsyncExtend, Extend, Webhooks, PollingOptions, PollingTimeoutError + from .wrapper import ( + AsyncExtend, + Extend, + ExtendCurrency, + ExtendDate, + ExtendSignature, + ExtractOutputValidationError, + PollingOptions, + PollingTimeoutError, + SchemaConversionError, + TypedExtractOutput, + TypedExtractRun, + Webhooks, + parse_extract_run, + pydantic_to_extend_schema, + ) from .edit_runs import ( EditRunsCreateRequestFile, EditRunsCreateRequestFileParams, @@ -1102,6 +1117,15 @@ "Webhooks": ".wrapper", "PollingOptions": ".wrapper", "PollingTimeoutError": ".wrapper", + "ExtendCurrency": ".wrapper", + "ExtendDate": ".wrapper", + "ExtendSignature": ".wrapper", + "ExtractOutputValidationError": ".wrapper", + "SchemaConversionError": ".wrapper", + "TypedExtractOutput": ".wrapper", + "TypedExtractRun": ".wrapper", + "parse_extract_run": ".wrapper", + "pydantic_to_extend_schema": ".wrapper", "ExtendEnvironment": ".environment", "ExternalDataValidationResult": ".types", "ExternalDataValidationResultParams": ".requests", @@ -1985,10 +2009,19 @@ def __dir__(): "ExcelSheetRange", "ExcelSheetRangeParams", "Extend", + "ExtendCurrency", + "ExtendDate", "ExtendEnvironment", + "ExtendSignature", + "ExtractOutputValidationError", "Webhooks", "PollingOptions", "PollingTimeoutError", + "SchemaConversionError", + "TypedExtractOutput", + "TypedExtractRun", + "parse_extract_run", + "pydantic_to_extend_schema", "ExternalDataValidationResult", "ExternalDataValidationResultParams", "ExternalDataValidationResultResponse", diff --git a/src/extend_ai/wrapper/__init__.py b/src/extend_ai/wrapper/__init__.py index d5edef6..1d0ab23 100644 --- a/src/extend_ai/wrapper/__init__.py +++ b/src/extend_ai/wrapper/__init__.py @@ -30,12 +30,33 @@ WebhookSignatureVerificationError, ) from .polling import PollingOptions, calculate_backoff_delay, poll_until_done, poll_until_done_async +from .schema import ( + ExtendCurrency, + ExtendDate, + ExtendSignature, + ExtractOutputValidationError, + SchemaConversionError, + TypedExtractOutput, + TypedExtractRun, + parse_extract_run, + pydantic_to_extend_schema, +) from .webhooks import RawWebhookEvent, SignedDataUrlPayload, WebhookEventWithSignedUrl, Webhooks __all__ = [ # Client "Extend", "AsyncExtend", + # Typed extraction schemas + "ExtendCurrency", + "ExtendDate", + "ExtendSignature", + "ExtractOutputValidationError", + "SchemaConversionError", + "TypedExtractOutput", + "TypedExtractRun", + "parse_extract_run", + "pydantic_to_extend_schema", # Webhooks "Webhooks", "RawWebhookEvent", diff --git a/src/extend_ai/wrapper/client.py b/src/extend_ai/wrapper/client.py index 55dff90..1f43473 100644 --- a/src/extend_ai/wrapper/client.py +++ b/src/extend_ai/wrapper/client.py @@ -32,44 +32,89 @@ import typing import httpx - +from ..batch_processor_run.client import AsyncBatchProcessorRunClient, BatchProcessorRunClient +from ..classifier_versions.client import AsyncClassifierVersionsClient, ClassifierVersionsClient +from ..classifiers.client import AsyncClassifiersClient, ClassifiersClient from ..client import AsyncExtend as GeneratedAsyncExtend from ..client import Extend as GeneratedExtend -from ..environment import ExtendEnvironment # Import all client types for proper type annotations -from ..files.client import FilesClient, AsyncFilesClient -from ..extractors.client import ExtractorsClient, AsyncExtractorsClient -from ..extractor_versions.client import ExtractorVersionsClient, AsyncExtractorVersionsClient -from ..classifiers.client import ClassifiersClient, AsyncClassifiersClient -from ..classifier_versions.client import ClassifierVersionsClient, AsyncClassifierVersionsClient -from ..splitters.client import SplittersClient, AsyncSplittersClient -from ..splitter_versions.client import SplitterVersionsClient, AsyncSplitterVersionsClient -from ..workflows.client import WorkflowsClient, AsyncWorkflowsClient -from ..evaluation_sets.client import EvaluationSetsClient, AsyncEvaluationSetsClient -from ..evaluation_set_items.client import EvaluationSetItemsClient, AsyncEvaluationSetItemsClient -from ..evaluation_set_runs.client import EvaluationSetRunsClient, AsyncEvaluationSetRunsClient -from ..processor.client import ProcessorClient, AsyncProcessorClient -from ..processor_run.client import ProcessorRunClient, AsyncProcessorRunClient -from ..processor_version.client import ProcessorVersionClient, AsyncProcessorVersionClient -from ..batch_processor_run.client import BatchProcessorRunClient, AsyncBatchProcessorRunClient - +from ..core.request_options import RequestOptions +from ..environment import ExtendEnvironment +from ..evaluation_set_items.client import AsyncEvaluationSetItemsClient, EvaluationSetItemsClient +from ..evaluation_set_runs.client import AsyncEvaluationSetRunsClient, EvaluationSetRunsClient +from ..evaluation_sets.client import AsyncEvaluationSetsClient, EvaluationSetsClient +from ..files.client import AsyncFilesClient, FilesClient +from ..processor.client import AsyncProcessorClient, ProcessorClient +from ..processor_run.client import AsyncProcessorRunClient, ProcessorRunClient +from ..processor_version.client import AsyncProcessorVersionClient, ProcessorVersionClient +from ..requests.extract_config_json import ExtractConfigJsonParams +from ..requests.extract_request_extractor import ExtractRequestExtractorParams +from ..requests.extract_request_file import ExtractRequestFileParams +from ..requests.multi_file_run_package import MultiFileRunPackageParams +from ..splitter_versions.client import AsyncSplitterVersionsClient, SplitterVersionsClient +from ..splitters.client import AsyncSplittersClient, SplittersClient +from ..types.extract_run import ExtractRun +from ..types.run_metadata import RunMetadata +from ..workflows.client import AsyncWorkflowsClient, WorkflowsClient from .resources import ( AsyncClassifyRunsClient, AsyncEditRunsClient, + AsyncExtractorsClient, + AsyncExtractorVersionsClient, AsyncExtractRunsClient, AsyncParseRunsClient, AsyncSplitRunsClient, AsyncWorkflowRunsClient, ClassifyRunsClient, EditRunsClient, + ExtractorsClient, + ExtractorVersionsClient, ExtractRunsClient, ParseRunsClient, SplitRunsClient, WorkflowRunsClient, ) +from .schema import ( + TypedExtractConfigParams, + TypedExtractorParams, + TypedExtractRun, + convert_typed_config, + convert_typed_extractor, + get_extractor_schema_model, + get_schema_model, + parse_extract_run, +) +from .schema.typed_run import ModelT from .webhooks import Webhooks +# this is used as the default value for optional parameters +OMIT = typing.cast(typing.Any, ...) + + +def _convert_extract_request( + extractor: typing.Any, config: typing.Any +) -> typing.Tuple[typing.Any, typing.Any, typing.Optional[type]]: + """ + Convert a pydantic model schema in an extract request's `config` or + `extractor["override_config"]` to Extend JSON Schema. Returns the + (possibly converted) extractor and config, and the schema model if one + was supplied. + """ + schema_model: typing.Optional[type] = None + + model = get_schema_model(config) + if model is not None: + schema_model = model + config = convert_typed_config(config) + + model = get_extractor_schema_model(extractor) + if model is not None: + schema_model = model + extractor = convert_typed_extractor(extractor) + + return extractor, config, schema_model + class Extend(GeneratedExtend): """ @@ -161,12 +206,96 @@ def __init__( self._workflow_runs_client: typing.Optional[WorkflowRunsClient] = None self._edit_runs_client: typing.Optional[EditRunsClient] = None self._parse_runs_client: typing.Optional[ParseRunsClient] = None + self._extractors_client: typing.Optional[ExtractorsClient] = None + self._extractor_versions_client: typing.Optional[ExtractorVersionsClient] = None @property def webhooks(self) -> Webhooks: """Webhook utilities for signature verification and event parsing.""" return self._webhooks + @typing.overload + def extract( + self, + *, + config: TypedExtractConfigParams[ModelT], + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + def extract( + self, + *, + extractor: TypedExtractorParams[ModelT], + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + def extract( + self, + *, + extractor: typing.Optional[ExtractRequestExtractorParams] = OMIT, + config: typing.Optional[ExtractConfigJsonParams] = OMIT, + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> ExtractRun: ... + + def extract( + self, + *, + extractor: typing.Optional[ + typing.Union[ExtractRequestExtractorParams, TypedExtractorParams[ModelT]] + ] = OMIT, + config: typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams[ModelT]]] = OMIT, + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> typing.Union[ExtractRun, TypedExtractRun[ModelT]]: + """ + Extract structured data from a file synchronously, waiting for the result. + + In addition to the generated `extract()` behavior, `config["schema"]` + (or `extractor["override_config"]["schema"]`) may be a pydantic model + class. The model is converted to Extend's JSON Schema format for the + request, and the extraction output is validated into instances of the + model, returned as a TypedExtractRun. + + Example: + from typing import Optional + from pydantic import BaseModel + + class Invoice(BaseModel): + invoice_number: Optional[str] = None + + result = client.extract( + file={"url": "https://example.com/invoice.pdf"}, + config={"schema": Invoice}, + ) + if result.output is not None: + print(result.output.value.invoice_number) # typed! + """ + converted_extractor, converted_config, schema_model = _convert_extract_request(extractor, config) + result = super().extract( + extractor=converted_extractor, + config=converted_config, + file=file, + package=package, + metadata=metadata, + request_options=request_options, + ) + if schema_model is not None: + return parse_extract_run(result, typing.cast(typing.Type[ModelT], schema_model)) + return result + # Run resources with create_and_poll support @property def extract_runs(self) -> ExtractRunsClient: @@ -218,13 +347,17 @@ def files(self) -> FilesClient: @property def extractors(self) -> ExtractorsClient: - """Extractors client.""" - return super().extractors # type: ignore[return-value] + """Extractors client with typed (pydantic) schema support.""" + if self._extractors_client is None: + self._extractors_client = ExtractorsClient(client_wrapper=self._client_wrapper) + return self._extractors_client @property def extractor_versions(self) -> ExtractorVersionsClient: - """Extractor versions client.""" - return super().extractor_versions # type: ignore[return-value] + """Extractor versions client with typed (pydantic) schema support.""" + if self._extractor_versions_client is None: + self._extractor_versions_client = ExtractorVersionsClient(client_wrapper=self._client_wrapper) + return self._extractor_versions_client @property def classifiers(self) -> ClassifiersClient: @@ -349,12 +482,80 @@ def __init__( self._workflow_runs_client: typing.Optional[AsyncWorkflowRunsClient] = None self._edit_runs_client: typing.Optional[AsyncEditRunsClient] = None self._parse_runs_client: typing.Optional[AsyncParseRunsClient] = None + self._extractors_client: typing.Optional[AsyncExtractorsClient] = None + self._extractor_versions_client: typing.Optional[AsyncExtractorVersionsClient] = None @property def webhooks(self) -> Webhooks: """Webhook utilities for signature verification and event parsing.""" return self._webhooks + @typing.overload + async def extract( + self, + *, + config: TypedExtractConfigParams[ModelT], + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + async def extract( + self, + *, + extractor: TypedExtractorParams[ModelT], + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + async def extract( + self, + *, + extractor: typing.Optional[ExtractRequestExtractorParams] = OMIT, + config: typing.Optional[ExtractConfigJsonParams] = OMIT, + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> ExtractRun: ... + + async def extract( + self, + *, + extractor: typing.Optional[ + typing.Union[ExtractRequestExtractorParams, TypedExtractorParams[ModelT]] + ] = OMIT, + config: typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams[ModelT]]] = OMIT, + file: typing.Optional[ExtractRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> typing.Union[ExtractRun, TypedExtractRun[ModelT]]: + """ + Extract structured data from a file synchronously, waiting for the result (async version). + + `config["schema"]` (or `extractor["override_config"]["schema"]`) may be + a pydantic model class, in which case the extraction output is validated + into instances of the model and returned as a TypedExtractRun. + """ + converted_extractor, converted_config, schema_model = _convert_extract_request(extractor, config) + result = await super().extract( + extractor=converted_extractor, + config=converted_config, + file=file, + package=package, + metadata=metadata, + request_options=request_options, + ) + if schema_model is not None: + return parse_extract_run(result, typing.cast(typing.Type[ModelT], schema_model)) + return result + # Run resources with create_and_poll support @property def extract_runs(self) -> AsyncExtractRunsClient: @@ -406,13 +607,17 @@ def files(self) -> AsyncFilesClient: @property def extractors(self) -> AsyncExtractorsClient: - """Extractors client.""" - return super().extractors # type: ignore[return-value] + """Extractors client with typed (pydantic) schema support.""" + if self._extractors_client is None: + self._extractors_client = AsyncExtractorsClient(client_wrapper=self._client_wrapper) + return self._extractors_client @property def extractor_versions(self) -> AsyncExtractorVersionsClient: - """Extractor versions client.""" - return super().extractor_versions # type: ignore[return-value] + """Extractor versions client with typed (pydantic) schema support.""" + if self._extractor_versions_client is None: + self._extractor_versions_client = AsyncExtractorVersionsClient(client_wrapper=self._client_wrapper) + return self._extractor_versions_client @property def classifiers(self) -> AsyncClassifiersClient: diff --git a/src/extend_ai/wrapper/resources/__init__.py b/src/extend_ai/wrapper/resources/__init__.py index b0c2345..ec7be97 100644 --- a/src/extend_ai/wrapper/resources/__init__.py +++ b/src/extend_ai/wrapper/resources/__init__.py @@ -3,6 +3,8 @@ from .classify_runs import AsyncClassifyRunsClient, ClassifyRunsClient from .edit_runs import AsyncEditRunsClient, EditRunsClient from .extract_runs import AsyncExtractRunsClient, ExtractRunsClient +from .extractor_versions import AsyncExtractorVersionsClient, ExtractorVersionsClient +from .extractors import AsyncExtractorsClient, ExtractorsClient from .parse_runs import AsyncParseRunsClient, ParseRunsClient from .split_runs import AsyncSplitRunsClient, SplitRunsClient from .workflow_runs import AsyncWorkflowRunsClient, WorkflowRunsClient @@ -10,6 +12,10 @@ __all__ = [ "ExtractRunsClient", "AsyncExtractRunsClient", + "ExtractorsClient", + "AsyncExtractorsClient", + "ExtractorVersionsClient", + "AsyncExtractorVersionsClient", "ClassifyRunsClient", "AsyncClassifyRunsClient", "SplitRunsClient", diff --git a/src/extend_ai/wrapper/resources/extract_runs.py b/src/extend_ai/wrapper/resources/extract_runs.py index 0733efd..932794e 100644 --- a/src/extend_ai/wrapper/resources/extract_runs.py +++ b/src/extend_ai/wrapper/resources/extract_runs.py @@ -1,5 +1,5 @@ """ -Extended ExtractRuns client with polling utilities. +Extended ExtractRuns client with polling utilities and typed schemas. Example: from extend_ai import Extend @@ -14,11 +14,26 @@ if result.status == "PROCESSED": print(result.output) + + # Or pass a pydantic model as the schema for typed output + from typing import Optional + from pydantic import BaseModel + + class Invoice(BaseModel): + invoice_number: Optional[str] = None + + result = client.extract_runs.create_and_poll( + file={"id": "file_xxx"}, + config={"schema": Invoice}, + ) + if result.output is not None: + print(result.output.value.invoice_number) # typed! """ -from typing import Any, Dict, Optional +import typing from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper +from ...core.request_options import RequestOptions from ...extract_runs.client import AsyncExtractRunsClient as GeneratedAsyncExtractRunsClient from ...extract_runs.client import ExtractRunsClient as GeneratedExtractRunsClient from ...extract_runs.requests.extract_runs_create_request_extractor import ExtractRunsCreateRequestExtractorParams @@ -28,13 +43,35 @@ from ...types.extract_run import ExtractRun from ...types.run_metadata import RunMetadata from ...types.run_priority import RunPriority -from ..polling import PollingOptions, poll_until_done, poll_until_done_async # Re-export for convenience -from ..polling import PollingTimeoutError +from ..polling import PollingOptions, PollingTimeoutError, poll_until_done, poll_until_done_async +from ..schema import ( + TypedExtractConfigParams, + TypedExtractorParams, + TypedExtractRun, + convert_typed_config, + convert_typed_extractor, + get_extractor_schema_model, + get_schema_model, + parse_extract_run, +) +from ..schema.typed_run import ModelT __all__ = ["ExtractRunsClient", "AsyncExtractRunsClient", "PollingTimeoutError"] +# this is used as the default value for optional parameters +OMIT = typing.cast(typing.Any, ...) + + +def _convert_create_args(extractor: typing.Any, config: typing.Any) -> typing.Tuple[typing.Any, typing.Any]: + """Convert a pydantic model schema in create() arguments, passing other values through.""" + if get_schema_model(config) is not None: + config = convert_typed_config(config) + if get_extractor_schema_model(extractor) is not None: + extractor = convert_typed_extractor(extractor) + return extractor, config + def _is_terminal_status(status: str) -> bool: """ @@ -45,6 +82,51 @@ def _is_terminal_status(status: str) -> bool: return status not in ("PROCESSING", "PENDING", "CANCELLING") +def _build_create_kwargs( + *, + file: typing.Optional[ExtractRunsCreateRequestFileParams], + package: typing.Optional[MultiFileRunPackageParams], + extractor: typing.Any, + config: typing.Any, + priority: typing.Optional[RunPriority], + metadata: typing.Optional[RunMetadata], +) -> typing.Tuple[typing.Dict[str, typing.Any], typing.Optional[type]]: + """ + Build create() kwargs (omitting None values), converting any pydantic model + schema to Extend JSON Schema. Returns the kwargs and the schema model, if + one was supplied. + """ + schema_model: typing.Optional[type] = None + + if config is not None: + model = get_schema_model(config) + if model is not None: + schema_model = model + config = convert_typed_config(config) + + if extractor is not None: + model = get_extractor_schema_model(extractor) + if model is not None: + schema_model = model + extractor = convert_typed_extractor(extractor) + + kwargs: typing.Dict[str, typing.Any] = {} + if file is not None: + kwargs["file"] = file + if package is not None: + kwargs["package"] = package + if extractor is not None: + kwargs["extractor"] = extractor + if config is not None: + kwargs["config"] = config + if priority is not None: + kwargs["priority"] = priority + if metadata is not None: + kwargs["metadata"] = metadata + + return kwargs, schema_model + + class ExtractRunsClient(GeneratedExtractRunsClient): """ Extended ExtractRuns client with create_and_poll method. @@ -56,17 +138,89 @@ class ExtractRunsClient(GeneratedExtractRunsClient): def __init__(self, *, client_wrapper: SyncClientWrapper): super().__init__(client_wrapper=client_wrapper) - def create_and_poll( + def create( self, *, - file: Optional[ExtractRunsCreateRequestFileParams] = None, - package: Optional[MultiFileRunPackageParams] = None, - extractor: Optional[ExtractRunsCreateRequestExtractorParams] = None, - config: Optional[ExtractConfigJsonParams] = None, - priority: Optional[RunPriority] = None, - metadata: Optional[RunMetadata] = None, - polling_options: Optional[PollingOptions] = None, + extractor: typing.Optional[ + typing.Union[ExtractRunsCreateRequestExtractorParams, TypedExtractorParams[ModelT]] + ] = OMIT, + config: typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams[ModelT]]] = OMIT, + file: typing.Optional[ExtractRunsCreateRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + priority: typing.Optional[RunPriority] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, ) -> ExtractRun: + """ + Create an extract run. See the generated client for full documentation. + + `config["schema"]` (or `extractor["override_config"]["schema"]`) may be a + pydantic model class; it is converted to Extend's JSON Schema format before + the request is sent. Note that `create()` returns immediately without + output — for validated, typed output use `create_and_poll()` or + `client.extract()` instead. + """ + converted_extractor, converted_config = _convert_create_args(extractor, config) + return super().create( + extractor=converted_extractor, + config=converted_config, + file=file, + package=package, + priority=priority, + metadata=metadata, + request_options=request_options, + ) + + @typing.overload + def create_and_poll( + self, + *, + config: TypedExtractConfigParams[ModelT], + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + def create_and_poll( + self, + *, + extractor: TypedExtractorParams[ModelT], + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + def create_and_poll( + self, + *, + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + extractor: typing.Optional[ExtractRunsCreateRequestExtractorParams] = None, + config: typing.Optional[ExtractConfigJsonParams] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> ExtractRun: ... + + def create_and_poll( + self, + *, + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + extractor: typing.Optional[ + typing.Union[ExtractRunsCreateRequestExtractorParams, TypedExtractorParams[ModelT]] + ] = None, + config: typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams[ModelT]]] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> typing.Union[ExtractRun, TypedExtractRun[ModelT]]: """ Creates an extract run and polls until it reaches a terminal state. @@ -81,13 +235,17 @@ def create_and_poll( package: A package of files for multi-file extraction. Mutually exclusive with `file` — provide one or the other. extractor: Reference to an existing extractor. - config: Inline extract configuration. + config: Inline extract configuration. `config["schema"]` may be a + pydantic model class, in which case the extraction output is + validated into instances of that model. priority: Priority of the run. metadata: Additional metadata for the run. polling_options: Options for polling behavior. Returns: - The final extract run when processing is complete. + The final extract run when processing is complete. If a pydantic + model was supplied as the schema, a TypedExtractRun whose output + values are instances of the model. Raises: PollingTimeoutError: If the run doesn't complete within max_wait_ms. @@ -101,32 +259,25 @@ def create_and_poll( if result.status == "PROCESSED": print(result.output) """ - # Build kwargs, only including non-None values to avoid passing null - kwargs: Dict[str, Any] = {} - if file is not None: - kwargs["file"] = file - if package is not None: - kwargs["package"] = package - if extractor is not None: - kwargs["extractor"] = extractor - if config is not None: - kwargs["config"] = config - if priority is not None: - kwargs["priority"] = priority - if metadata is not None: - kwargs["metadata"] = metadata + kwargs, schema_model = _build_create_kwargs( + file=file, package=package, extractor=extractor, config=config, priority=priority, metadata=metadata + ) # Create the extract run create_response = self.create(**kwargs) run_id = create_response.id # Poll until terminal state - return poll_until_done( + result = poll_until_done( retrieve=lambda: self.retrieve(run_id), is_terminal=lambda response: _is_terminal_status(response.status), options=polling_options, ) + if schema_model is not None: + return parse_extract_run(result, typing.cast(typing.Type[ModelT], schema_model)) + return result + class AsyncExtractRunsClient(GeneratedAsyncExtractRunsClient): """ @@ -136,44 +287,112 @@ class AsyncExtractRunsClient(GeneratedAsyncExtractRunsClient): def __init__(self, *, client_wrapper: AsyncClientWrapper): super().__init__(client_wrapper=client_wrapper) - async def create_and_poll( + async def create( self, *, - file: Optional[ExtractRunsCreateRequestFileParams] = None, - package: Optional[MultiFileRunPackageParams] = None, - extractor: Optional[ExtractRunsCreateRequestExtractorParams] = None, - config: Optional[ExtractConfigJsonParams] = None, - priority: Optional[RunPriority] = None, - metadata: Optional[RunMetadata] = None, - polling_options: Optional[PollingOptions] = None, + extractor: typing.Optional[ + typing.Union[ExtractRunsCreateRequestExtractorParams, TypedExtractorParams[ModelT]] + ] = OMIT, + config: typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams[ModelT]]] = OMIT, + file: typing.Optional[ExtractRunsCreateRequestFileParams] = OMIT, + package: typing.Optional[MultiFileRunPackageParams] = OMIT, + priority: typing.Optional[RunPriority] = OMIT, + metadata: typing.Optional[RunMetadata] = OMIT, + request_options: typing.Optional[RequestOptions] = None, ) -> ExtractRun: + """ + Create an extract run (async version). See the generated client for full + documentation. + + `config["schema"]` (or `extractor["override_config"]["schema"]`) may be a + pydantic model class; it is converted to Extend's JSON Schema format before + the request is sent. Note that `create()` returns immediately without + output — for validated, typed output use `create_and_poll()` or + `client.extract()` instead. + """ + converted_extractor, converted_config = _convert_create_args(extractor, config) + return await super().create( + extractor=converted_extractor, + config=converted_config, + file=file, + package=package, + priority=priority, + metadata=metadata, + request_options=request_options, + ) + + @typing.overload + async def create_and_poll( + self, + *, + config: TypedExtractConfigParams[ModelT], + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + async def create_and_poll( + self, + *, + extractor: TypedExtractorParams[ModelT], + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> TypedExtractRun[ModelT]: ... + + @typing.overload + async def create_and_poll( + self, + *, + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + extractor: typing.Optional[ExtractRunsCreateRequestExtractorParams] = None, + config: typing.Optional[ExtractConfigJsonParams] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> ExtractRun: ... + + async def create_and_poll( + self, + *, + file: typing.Optional[ExtractRunsCreateRequestFileParams] = None, + package: typing.Optional[MultiFileRunPackageParams] = None, + extractor: typing.Optional[ + typing.Union[ExtractRunsCreateRequestExtractorParams, TypedExtractorParams[ModelT]] + ] = None, + config: typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams[ModelT]]] = None, + priority: typing.Optional[RunPriority] = None, + metadata: typing.Optional[RunMetadata] = None, + polling_options: typing.Optional[PollingOptions] = None, + ) -> typing.Union[ExtractRun, TypedExtractRun[ModelT]]: """ Creates an extract run and polls until it reaches a terminal state (async version). `file` and `package` are mutually exclusive — provide one or the other. + `config["schema"]` may be a pydantic model class, in which case the + extraction output is validated into instances of that model. """ - # Build kwargs, only including non-None values to avoid passing null - kwargs: Dict[str, Any] = {} - if file is not None: - kwargs["file"] = file - if package is not None: - kwargs["package"] = package - if extractor is not None: - kwargs["extractor"] = extractor - if config is not None: - kwargs["config"] = config - if priority is not None: - kwargs["priority"] = priority - if metadata is not None: - kwargs["metadata"] = metadata + kwargs, schema_model = _build_create_kwargs( + file=file, package=package, extractor=extractor, config=config, priority=priority, metadata=metadata + ) # Create the extract run create_response = await self.create(**kwargs) run_id = create_response.id # Poll until terminal state - return await poll_until_done_async( + result = await poll_until_done_async( retrieve=lambda: self.retrieve(run_id), is_terminal=lambda response: _is_terminal_status(response.status), options=polling_options, ) + + if schema_model is not None: + return parse_extract_run(result, typing.cast(typing.Type[ModelT], schema_model)) + return result diff --git a/src/extend_ai/wrapper/resources/extractor_versions.py b/src/extend_ai/wrapper/resources/extractor_versions.py new file mode 100644 index 0000000..48211b4 --- /dev/null +++ b/src/extend_ai/wrapper/resources/extractor_versions.py @@ -0,0 +1,84 @@ +""" +Extended ExtractorVersions client with typed (pydantic) schema support. + +`config["schema"]` may be a pydantic model class; it is converted to Extend's +JSON Schema format before the request is sent. +""" + +import typing + +from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper +from ...core.request_options import RequestOptions +from ...extractor_versions.client import AsyncExtractorVersionsClient as GeneratedAsyncExtractorVersionsClient +from ...extractor_versions.client import ExtractorVersionsClient as GeneratedExtractorVersionsClient +from ...requests.extract_config_json import ExtractConfigJsonParams +from ...types.extractor_version import ExtractorVersion +from ...types.release_type import ReleaseType +from ...types.version_description import VersionDescription +from ..schema import TypedExtractConfigParams +from .extractors import convert_config_arg + +__all__ = ["ExtractorVersionsClient", "AsyncExtractorVersionsClient"] + +# this is used as the default value for optional parameters +OMIT = typing.cast(typing.Any, ...) + +_ConfigParam = typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams]] + + +class ExtractorVersionsClient(GeneratedExtractorVersionsClient): + """ + Extended ExtractorVersions client that accepts a pydantic model class as + `config["schema"]` in create(). + """ + + def __init__(self, *, client_wrapper: SyncClientWrapper): + super().__init__(client_wrapper=client_wrapper) + + def create( + self, + extractor_id: str, + *, + release_type: ReleaseType, + extend_workspace_id: typing.Optional[str] = None, + description: typing.Optional[VersionDescription] = OMIT, + config: _ConfigParam = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> ExtractorVersion: + return super().create( + extractor_id, + release_type=release_type, + extend_workspace_id=extend_workspace_id, + description=description, + config=convert_config_arg(config), + request_options=request_options, + ) + + +class AsyncExtractorVersionsClient(GeneratedAsyncExtractorVersionsClient): + """ + Extended AsyncExtractorVersions client that accepts a pydantic model class + as `config["schema"]` in create(). + """ + + def __init__(self, *, client_wrapper: AsyncClientWrapper): + super().__init__(client_wrapper=client_wrapper) + + async def create( + self, + extractor_id: str, + *, + release_type: ReleaseType, + extend_workspace_id: typing.Optional[str] = None, + description: typing.Optional[VersionDescription] = OMIT, + config: _ConfigParam = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> ExtractorVersion: + return await super().create( + extractor_id, + release_type=release_type, + extend_workspace_id=extend_workspace_id, + description=description, + config=convert_config_arg(config), + request_options=request_options, + ) diff --git a/src/extend_ai/wrapper/resources/extractors.py b/src/extend_ai/wrapper/resources/extractors.py new file mode 100644 index 0000000..c6f1a8c --- /dev/null +++ b/src/extend_ai/wrapper/resources/extractors.py @@ -0,0 +1,133 @@ +""" +Extended Extractors client with typed (pydantic) schema support. + +`config["schema"]` may be a pydantic model class; it is converted to Extend's +JSON Schema format before the request is sent. + +Example: + from typing import Optional + from pydantic import BaseModel, Field + from extend_ai import Extend + + class Invoice(BaseModel): + invoice_number: Optional[str] = Field(None, description="The invoice number") + + client = Extend(token="...") + extractor = client.extractors.create( + name="Invoice Extractor", + config={"schema": Invoice}, + ) +""" + +import typing + +from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper +from ...core.request_options import RequestOptions +from ...extractors.client import AsyncExtractorsClient as GeneratedAsyncExtractorsClient +from ...extractors.client import ExtractorsClient as GeneratedExtractorsClient +from ...extractors.requests.extractors_create_request_generate import ExtractorsCreateRequestGenerateParams +from ...requests.extract_config_json import ExtractConfigJsonParams +from ...types.extractor import Extractor +from ..schema import TypedExtractConfigParams, convert_typed_config, get_schema_model + +__all__ = ["ExtractorsClient", "AsyncExtractorsClient"] + +# this is used as the default value for optional parameters +OMIT = typing.cast(typing.Any, ...) + +_ConfigParam = typing.Optional[typing.Union[ExtractConfigJsonParams, TypedExtractConfigParams]] + + +def convert_config_arg(config: typing.Any) -> typing.Any: + """Convert a pydantic model schema in a config argument, passing other values through.""" + if get_schema_model(config) is not None: + return convert_typed_config(config) + return config + + +class ExtractorsClient(GeneratedExtractorsClient): + """ + Extended Extractors client that accepts a pydantic model class as + `config["schema"]` in create() and update(). + """ + + def __init__(self, *, client_wrapper: SyncClientWrapper): + super().__init__(client_wrapper=client_wrapper) + + def create( + self, + *, + name: str, + clone_extractor_id: typing.Optional[str] = OMIT, + config: _ConfigParam = OMIT, + generate: typing.Optional[ExtractorsCreateRequestGenerateParams] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> Extractor: + return super().create( + name=name, + clone_extractor_id=clone_extractor_id, + config=convert_config_arg(config), + generate=generate, + request_options=request_options, + ) + + def update( + self, + id: str, + *, + extend_workspace_id: typing.Optional[str] = None, + name: typing.Optional[str] = OMIT, + config: _ConfigParam = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> Extractor: + return super().update( + id, + extend_workspace_id=extend_workspace_id, + name=name, + config=convert_config_arg(config), + request_options=request_options, + ) + + +class AsyncExtractorsClient(GeneratedAsyncExtractorsClient): + """ + Extended AsyncExtractors client that accepts a pydantic model class as + `config["schema"]` in create() and update(). + """ + + def __init__(self, *, client_wrapper: AsyncClientWrapper): + super().__init__(client_wrapper=client_wrapper) + + async def create( + self, + *, + name: str, + clone_extractor_id: typing.Optional[str] = OMIT, + config: _ConfigParam = OMIT, + generate: typing.Optional[ExtractorsCreateRequestGenerateParams] = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> Extractor: + return await super().create( + name=name, + clone_extractor_id=clone_extractor_id, + config=convert_config_arg(config), + generate=generate, + request_options=request_options, + ) + + async def update( + self, + id: str, + *, + extend_workspace_id: typing.Optional[str] = None, + name: typing.Optional[str] = OMIT, + config: _ConfigParam = OMIT, + request_options: typing.Optional[RequestOptions] = None, + ) -> Extractor: + return await super().update( + id, + extend_workspace_id=extend_workspace_id, + name=name, + config=convert_config_arg(config), + request_options=request_options, + ) diff --git a/src/extend_ai/wrapper/resources/parse_runs.py b/src/extend_ai/wrapper/resources/parse_runs.py index 277ea51..03b5d55 100644 --- a/src/extend_ai/wrapper/resources/parse_runs.py +++ b/src/extend_ai/wrapper/resources/parse_runs.py @@ -21,12 +21,13 @@ from ...parse_runs.client import AsyncParseRunsClient as GeneratedAsyncParseRunsClient from ...parse_runs.client import ParseRunsClient as GeneratedParseRunsClient from ...parse_runs.requests.parse_runs_create_request_file import ParseRunsCreateRequestFileParams +from ...requests.data_retention import DataRetentionParams from ...requests.parse_config import ParseConfigParams from ...types.parse_run import ParseRun -from ..polling import PollingOptions, poll_until_done, poll_until_done_async +from ...types.run_metadata import RunMetadata # Re-export for convenience -from ..polling import PollingTimeoutError +from ..polling import PollingOptions, PollingTimeoutError, poll_until_done, poll_until_done_async __all__ = ["ParseRunsClient", "AsyncParseRunsClient", "PollingTimeoutError"] @@ -58,6 +59,8 @@ def create_and_poll( *, file: ParseRunsCreateRequestFileParams, config: Optional[ParseConfigParams] = None, + metadata: Optional[RunMetadata] = None, + data_retention: Optional[DataRetentionParams] = None, polling_options: Optional[PollingOptions] = None, ) -> ParseRun: """ @@ -68,6 +71,8 @@ def create_and_poll( Args: file: The file to parse (FileFromId or FileFromUrl). config: Parse configuration options. + metadata: Additional metadata for the run. + data_retention: Data retention policy override for the run. polling_options: Options for polling behavior. Returns: @@ -88,6 +93,10 @@ def create_and_poll( kwargs: Dict[str, Any] = {"file": file} if config is not None: kwargs["config"] = config + if metadata is not None: + kwargs["metadata"] = metadata + if data_retention is not None: + kwargs["data_retention"] = data_retention # Create the parse run create_response = self.create(**kwargs) @@ -114,6 +123,8 @@ async def create_and_poll( *, file: ParseRunsCreateRequestFileParams, config: Optional[ParseConfigParams] = None, + metadata: Optional[RunMetadata] = None, + data_retention: Optional[DataRetentionParams] = None, polling_options: Optional[PollingOptions] = None, ) -> ParseRun: """ @@ -123,6 +134,10 @@ async def create_and_poll( kwargs: Dict[str, Any] = {"file": file} if config is not None: kwargs["config"] = config + if metadata is not None: + kwargs["metadata"] = metadata + if data_retention is not None: + kwargs["data_retention"] = data_retention # Create the parse run create_response = await self.create(**kwargs) diff --git a/src/extend_ai/wrapper/schema/__init__.py b/src/extend_ai/wrapper/schema/__init__.py new file mode 100644 index 0000000..21c9199 --- /dev/null +++ b/src/extend_ai/wrapper/schema/__init__.py @@ -0,0 +1,71 @@ +""" +Schema utilities for typed extraction with pydantic models. + +Define your extraction schema as a pydantic model and pass it as +``config["schema"]`` to get end-to-end typing: the SDK converts the model to +Extend's JSON Schema format for the request, and validates the extraction +output back into model instances. + +Example: + from typing import List, Optional + from pydantic import BaseModel, Field + from extend_ai import Extend, ExtendCurrency, ExtendDate + + class LineItem(BaseModel): + description: Optional[str] = None + quantity: Optional[float] = None + price: Optional[ExtendCurrency] = None + + class Invoice(BaseModel): + invoice_number: Optional[str] = Field(None, description="The invoice number") + invoice_date: ExtendDate = Field(None, description="The invoice date") + line_items: List[LineItem] = Field(default_factory=list) + 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}, + ) + + # output.value is a validated Invoice instance + if result.output is not None: + print(result.output.value.invoice_number) + print(result.output.value.total.amount if result.output.value.total else None) +""" + +from .config_conversion import ( + TypedExtractConfigParams, + TypedExtractorParams, + convert_typed_config, + convert_typed_extractor, + get_extractor_schema_model, + get_schema_model, +) +from .conversion import SchemaConversionError, pydantic_to_extend_schema +from .custom_types import ExtendCurrency, ExtendDate, ExtendSignature +from .typed_run import ExtractOutputValidationError, TypedExtractOutput, TypedExtractRun, parse_extract_run + +__all__ = [ + # Custom field types + "ExtendCurrency", + "ExtendDate", + "ExtendSignature", + # Conversion + "SchemaConversionError", + "pydantic_to_extend_schema", + # Errors + "ExtractOutputValidationError", + # Typed configs (for annotations / advanced usage) + "TypedExtractConfigParams", + "TypedExtractorParams", + # Typed runs + "TypedExtractOutput", + "TypedExtractRun", + "parse_extract_run", + # Internal conversion helpers + "convert_typed_config", + "convert_typed_extractor", + "get_extractor_schema_model", + "get_schema_model", +] diff --git a/src/extend_ai/wrapper/schema/config_conversion.py b/src/extend_ai/wrapper/schema/config_conversion.py new file mode 100644 index 0000000..6a8601a --- /dev/null +++ b/src/extend_ai/wrapper/schema/config_conversion.py @@ -0,0 +1,93 @@ +""" +Detection and conversion of typed (pydantic) extract configs to API format. + +Used by ``Extend.extract()``, ``ExtractRunsClient.create_and_poll()``, +``ExtractorsClient.create()/update()``, and ``ExtractorVersionsClient.create()``. +""" + +import typing + +import pydantic +import typing_extensions +from ...requests.extract_advanced_options import ExtractAdvancedOptionsParams +from ...requests.parse_config import ParseConfigParams +from ...types.extract_base_processor import ExtractBaseProcessor +from ...types.processor_version_string import ProcessorVersionString +from .conversion import pydantic_to_extend_schema +from .typed_run import ModelT + +__all__ = [ + "TypedExtractConfigParams", + "TypedExtractorParams", + "convert_typed_config", + "convert_typed_extractor", + "get_extractor_schema_model", + "get_schema_model", +] + +_OVERRIDE_CONFIG_KEYS = ("override_config", "overrideConfig") + + +class TypedExtractConfigParams(typing_extensions.TypedDict, typing.Generic[ModelT], total=False): + """ + Extract configuration whose ``schema`` is a pydantic model class. + Extraction output will be validated against the model. + """ + + schema: typing_extensions.Required[typing.Type[ModelT]] + base_processor: ExtractBaseProcessor + base_version: str + extraction_rules: str + advanced_options: ExtractAdvancedOptionsParams + parse_config: ParseConfigParams + + +class TypedExtractorParams(typing_extensions.TypedDict, typing.Generic[ModelT]): + """ + Reference to an existing extractor whose ``override_config.schema`` is a + pydantic model class. Extraction output will be validated against the model. + """ + + id: str + version: typing_extensions.NotRequired[ProcessorVersionString] + override_config: TypedExtractConfigParams[ModelT] + + +def _as_schema_model(schema: typing.Any) -> typing.Optional[typing.Type[pydantic.BaseModel]]: + if isinstance(schema, type) and issubclass(schema, pydantic.BaseModel): + return schema + return None + + +def get_schema_model(config: typing.Any) -> typing.Optional[typing.Type[pydantic.BaseModel]]: + """Return the pydantic model used as ``config["schema"]``, if there is one.""" + if isinstance(config, typing.Mapping): + return _as_schema_model(config.get("schema")) + return None + + +def convert_typed_config(config: typing.Mapping[str, typing.Any]) -> typing.Dict[str, typing.Any]: + """Return a copy of ``config`` with its pydantic model schema converted to JSON Schema.""" + converted = dict(config) + converted["schema"] = pydantic_to_extend_schema(converted["schema"]) + return converted + + +def get_extractor_schema_model(extractor: typing.Any) -> typing.Optional[typing.Type[pydantic.BaseModel]]: + """Return the pydantic model used as ``extractor["override_config"]["schema"]``, if there is one.""" + if not isinstance(extractor, typing.Mapping): + return None + for key in _OVERRIDE_CONFIG_KEYS: + model = get_schema_model(extractor.get(key)) + if model is not None: + return model + return None + + +def convert_typed_extractor(extractor: typing.Mapping[str, typing.Any]) -> typing.Dict[str, typing.Any]: + """Return a copy of ``extractor`` with its override config's schema converted to JSON Schema.""" + converted = dict(extractor) + for key in _OVERRIDE_CONFIG_KEYS: + if get_schema_model(converted.get(key)) is not None: + converted[key] = convert_typed_config(converted[key]) + return converted diff --git a/src/extend_ai/wrapper/schema/conversion.py b/src/extend_ai/wrapper/schema/conversion.py new file mode 100644 index 0000000..cc2c4a3 --- /dev/null +++ b/src/extend_ai/wrapper/schema/conversion.py @@ -0,0 +1,353 @@ +""" +Converts pydantic models to Extend's JSON Schema format. + +The converter is strict: mistakes that would otherwise surface as a 400 from +the API, or worse as a validation failure after a completed extraction run, +are raised as SchemaConversionError before any request is sent. In +particular, fields whose emitted schema is nullable (primitives, enums, +dates) must be declared Optional, because extraction can return null for any +field and the output is validated back into the model. + +Structural limits (nesting depth, property counts, property key format) are +validated server-side. +""" + +import datetime as dt +import enum +import types +import typing + +import pydantic +import typing_extensions +from ...core.pydantic_utilities import IS_PYDANTIC_V2 +from .custom_types import get_extend_type + +__all__ = ["SchemaConversionError", "pydantic_to_extend_schema"] + +_NoneType = type(None) + +# typing.Literal and typing_extensions.Literal are distinct objects on some +# Python versions (e.g. 3.8), so origins must be checked against both. +_LITERAL_ORIGINS = {typing_extensions.Literal, getattr(typing, "Literal", typing_extensions.Literal)} + + +class SchemaConversionError(Exception): + """Raised when a pydantic model cannot be converted to Extend JSON Schema.""" + + def __init__(self, message: str, path: typing.Optional[typing.List[str]] = None): + self.path: typing.List[str] = list(path or []) + if self.path: + message = f"{message} at path: {'.'.join(self.path)}" + super().__init__(message) + + +def _iter_model_fields( + model: typing.Type[pydantic.BaseModel], +) -> typing.Iterator[typing.Tuple[str, typing.Any, typing.Optional[str], typing.Any]]: + """ + Yield (field_name, annotation, description, alias) for each field of a + pydantic model, working under both pydantic v1 and v2. + """ + # Raw class annotations (via get_type_hints) preserve Optional wrappers, + # which pydantic v1's `outer_type_` strips. + try: + hints = typing_extensions.get_type_hints(model, include_extras=True) + except Exception: + hints = {} + + if IS_PYDANTIC_V2: + for name, field in model.model_fields.items(): # type: ignore[attr-defined] + alias = field.alias or getattr(field, "validation_alias", None) + yield name, hints.get(name, field.annotation), field.description, alias + else: + for name, field in model.__fields__.items(): # type: ignore[attr-defined] + info = field.field_info # type: ignore[attr-defined] + annotation = hints.get(name) + if annotation is None: + annotation = field.outer_type_ # type: ignore[attr-defined] + if field.allow_none: # type: ignore[attr-defined] + annotation = typing.Optional[annotation] + yield name, annotation, getattr(info, "description", None), getattr(info, "alias", None) + + +def _is_union_origin(origin: typing.Any) -> bool: + if origin is typing.Union: + return True + union_type = getattr(types, "UnionType", None) # X | Y syntax on Python 3.10+ + return union_type is not None and origin is union_type + + +def _unwrap_annotation(annotation: typing.Any, path: typing.List[str]) -> typing.Any: + """ + Strip Annotated metadata and Optional/None-unions from an annotation, + returning the inner type. Unions of multiple non-None types are rejected. + """ + while True: + origin = typing_extensions.get_origin(annotation) + if origin is typing_extensions.Annotated: + annotation = typing_extensions.get_args(annotation)[0] + elif _is_union_origin(origin): + non_none = [arg for arg in typing_extensions.get_args(annotation) if arg is not _NoneType] + if len(non_none) != 1: + raise SchemaConversionError( + "Union types are not supported (only Optional[...] is allowed)", path + ) + annotation = non_none[0] + else: + return annotation + + +def _accepts_none(annotation: typing.Any) -> bool: + """Whether a value of None validates against the annotation.""" + origin = typing_extensions.get_origin(annotation) + if origin is typing_extensions.Annotated: + return _accepts_none(typing_extensions.get_args(annotation)[0]) + if _is_union_origin(origin): + return any(arg is _NoneType or _accepts_none(arg) for arg in typing_extensions.get_args(annotation)) + if origin in _LITERAL_ORIGINS: + return None in typing_extensions.get_args(annotation) + return annotation is _NoneType + + +def _require_nullable(annotation: typing.Any, kind: str, path: typing.List[str]) -> None: + """ + Fields whose emitted schema is nullable must accept None, otherwise + extraction output containing null would fail model validation after the + run has already completed. + """ + if not _accepts_none(annotation): + raise SchemaConversionError( + f"Field must be Optional: extraction can return null for any field, " + f"so declare it as Optional[{kind}]", + path, + ) + + +def _with_description(schema: typing.Dict[str, typing.Any], description: typing.Optional[str]) -> typing.Dict[str, typing.Any]: + if description: + schema["description"] = description + return schema + + +def _date_schema() -> typing.Dict[str, typing.Any]: + return {"type": ["string", "null"], "extend:type": "date"} + + +def _currency_schema() -> typing.Dict[str, typing.Any]: + return { + "type": "object", + "extend:type": "currency", + "properties": { + "amount": {"type": ["number", "null"]}, + "iso_4217_currency_code": {"type": ["string", "null"]}, + }, + "required": ["amount", "iso_4217_currency_code"], + "additionalProperties": False, + } + + +def _signature_schema() -> typing.Dict[str, typing.Any]: + return { + "type": "object", + "extend:type": "signature", + "properties": { + "printed_name": {"type": ["string", "null"]}, + "signature_date": {"type": ["string", "null"], "extend:type": "date"}, + "is_signed": {"type": ["boolean", "null"]}, + "title_or_role": {"type": ["string", "null"]}, + }, + "required": ["printed_name", "signature_date", "is_signed", "title_or_role"], + "additionalProperties": False, + } + + +def _enum_values(annotation: typing.Any, path: typing.List[str]) -> typing.List[typing.Optional[str]]: + """Extract string enum values from a Literal[...] or string Enum class.""" + if typing_extensions.get_origin(annotation) in _LITERAL_ORIGINS: + raw_values: typing.List[typing.Any] = [v for v in typing_extensions.get_args(annotation) if v is not None] + else: # enum.Enum subclass + raw_values = [member.value for member in annotation] + + values: typing.List[typing.Optional[str]] = [] + for value in raw_values: + if not isinstance(value, str): + raise SchemaConversionError( + f"Enums must only contain strings, got {type(value).__name__}: {value!r}", path + ) + values.append(value) + values.append(None) + return values + + +def _is_enum_annotation(annotation: typing.Any) -> bool: + if typing_extensions.get_origin(annotation) in _LITERAL_ORIGINS: + return True + return isinstance(annotation, type) and issubclass(annotation, enum.Enum) + + +def pydantic_to_extend_schema(model: typing.Type[pydantic.BaseModel]) -> typing.Dict[str, typing.Any]: + """ + Convert a pydantic model class to Extend's JSON Schema format. + + Every property is listed as required, and field descriptions come from + ``Field(description=...)``. Primitive, enum, and date fields must be + declared ``Optional`` — extraction can return ``null`` for any field, and + the emitted schema marks them nullable per Extend's schema requirements. + + Args: + model: A ``pydantic.BaseModel`` subclass describing the data to extract. + + Returns: + The Extend JSON Schema as a plain dict. + + Raises: + SchemaConversionError: If the model uses unsupported types, recursive + references, field aliases, or non-Optional nullable fields. + """ + if not (isinstance(model, type) and issubclass(model, pydantic.BaseModel)): + raise SchemaConversionError(f"Schema must be a pydantic BaseModel subclass, got {model!r}") + return _convert_object(model, [], frozenset()) + + +def _convert_object( + model: typing.Type[pydantic.BaseModel], + path: typing.List[str], + seen: typing.FrozenSet[type], +) -> typing.Dict[str, typing.Any]: + # Extend's schema format cannot express recursion, and recursive models + # would otherwise overflow the stack (fatally on some Python versions). + if model in seen: + raise SchemaConversionError( + f"Recursive model references are not supported: {model.__name__} refers back to itself", + path, + ) + seen = seen | {model} + + properties: typing.Dict[str, typing.Any] = {} + required: typing.List[str] = [] + + for name, annotation, description, alias in _iter_model_fields(model): + if alias: + raise SchemaConversionError( + f"Field aliases are not supported for extraction schemas " + f"(field {name!r} has alias {alias!r}): the extraction output uses field names, " + f"so aliased fields would silently validate to None. Remove the alias.", + path + [name], + ) + properties[name] = _convert_annotation(annotation, description, path + [name], seen) + required.append(name) + + return { + "type": "object", + "properties": properties, + "required": required, + "additionalProperties": False, + } + + +def _convert_annotation( + annotation: typing.Any, + description: typing.Optional[str], + path: typing.List[str], + seen: typing.FrozenSet[type], +) -> typing.Dict[str, typing.Any]: + inner = _unwrap_annotation(annotation, path) + + if typing_extensions.get_origin(inner) is list or inner is list: + args = typing_extensions.get_args(inner) + if not args: + raise SchemaConversionError("Arrays must declare an item type (use List[...])", path) + return _with_description({"type": "array", "items": _convert_array_item(args[0], path, seen)}, description) + + if _is_enum_annotation(inner): + kind = inner.__name__ if isinstance(inner, type) else "Literal[...]" + _require_nullable(annotation, kind, path) + return _with_description({"enum": _enum_values(inner, path)}, description) + + if isinstance(inner, type): + if issubclass(inner, pydantic.BaseModel): + extend_type = get_extend_type(inner) + if extend_type == "currency": + return _with_description(_currency_schema(), description) + if extend_type == "signature": + return _with_description(_signature_schema(), description) + return _with_description(_convert_object(inner, path, seen), description) + if issubclass(inner, bool): + _require_nullable(annotation, "bool", path) + return _with_description({"type": ["boolean", "null"]}, description) + if issubclass(inner, int): + _require_nullable(annotation, "int", path) + return _with_description({"type": ["integer", "null"]}, description) + if issubclass(inner, float): + _require_nullable(annotation, "float", path) + return _with_description({"type": ["number", "null"]}, description) + if issubclass(inner, dt.datetime): + raise SchemaConversionError( + "datetime.datetime is not supported; use datetime.date (or ExtendDate) for date fields", path + ) + if issubclass(inner, dt.date): + _require_nullable(annotation, "datetime.date", path) + return _with_description(_date_schema(), description) + if issubclass(inner, str): + _require_nullable(annotation, "str", path) + return _with_description({"type": ["string", "null"]}, description) + + raise SchemaConversionError(f"Unsupported type: {inner!r}", path) + + +def _convert_array_item( + annotation: typing.Any, + path: typing.List[str], + seen: typing.FrozenSet[type], +) -> typing.Dict[str, typing.Any]: + """ + Convert array item types, which have different rules than top-level types: + items can be objects or primitives, and primitive items are NOT nullable. + """ + inner = _unwrap_annotation(annotation, path) + + # Array items are never null in extraction output, so an Optional item + # annotation would misleadingly suggest otherwise. + if _accepts_none(annotation) and not (isinstance(inner, type) and issubclass(inner, pydantic.BaseModel)): + raise SchemaConversionError( + "Array items must not be Optional: extraction never returns null array items " + "(use e.g. List[str] instead of List[Optional[str]])", + path, + ) + + if _is_enum_annotation(inner): + raise SchemaConversionError( + "Enums are not supported as array items. " + "Array items must be objects or primitives (string, number, integer, boolean).", + path, + ) + + if isinstance(inner, type) and issubclass(inner, pydantic.BaseModel): + extend_type = get_extend_type(inner) + if extend_type == "currency": + return _currency_schema() + if extend_type == "signature": + return _signature_schema() + return _convert_object(inner, path, seen) + + if isinstance(inner, type): + if issubclass(inner, bool): + return {"type": "boolean"} + if issubclass(inner, int): + return {"type": "integer"} + if issubclass(inner, float): + return {"type": "number"} + if issubclass(inner, dt.datetime): + raise SchemaConversionError( + "datetime.datetime is not supported; use datetime.date (or ExtendDate) for date fields", path + ) + if issubclass(inner, dt.date): + return {"type": "string", "extend:type": "date"} + if issubclass(inner, str): + return {"type": "string"} + + raise SchemaConversionError( + f"Unsupported array item type: {inner!r}. " + "Array items must be objects or primitives (string, number, integer, boolean).", + path, + ) diff --git a/src/extend_ai/wrapper/schema/custom_types.py b/src/extend_ai/wrapper/schema/custom_types.py new file mode 100644 index 0000000..fab07dc --- /dev/null +++ b/src/extend_ai/wrapper/schema/custom_types.py @@ -0,0 +1,92 @@ +""" +Custom pydantic field types for Extend-specific extraction fields. + +These map to the `extend:type` custom field types in Extend's JSON Schema +format (see https://docs.extend.ai/2026-02-09/extraction/schema). Use them as +field annotations in the pydantic model you pass as an extraction schema. + +Example: + import datetime + from typing import Optional + from pydantic import BaseModel, Field + from extend_ai import ExtendCurrency, ExtendDate, ExtendSignature + + 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") + signature: Optional[ExtendSignature] = None +""" + +import datetime as dt +import typing + +import pydantic + +__all__ = ["ExtendCurrency", "ExtendDate", "ExtendSignature"] + + +# Fields annotated with `ExtendDate` (or plain `datetime.date`) convert to: +# +# {"type": ["string", "null"], "extend:type": "date"} +# +# The API returns an ISO date string (yyyy-mm-dd) or null; pydantic parses it +# back into a `datetime.date` when the output is validated. +ExtendDate = typing.Optional[dt.date] + + +class ExtendCurrency(pydantic.BaseModel): + """ + Currency field type. + + Converts to Extend's currency schema: + + { + "type": "object", + "extend:type": "currency", + "properties": { + "amount": {"type": ["number", "null"]}, + "iso_4217_currency_code": {"type": ["string", "null"]} + }, + "required": ["amount", "iso_4217_currency_code"] + } + """ + + __extend_type__: typing.ClassVar[str] = "currency" + + amount: typing.Optional[float] = None + iso_4217_currency_code: typing.Optional[str] = None + + +class ExtendSignature(pydantic.BaseModel): + """ + Signature field type. + + Converts to Extend's signature schema, which enables advanced signature + detection during parsing and post-processing heuristics that reduce false + positives on unsigned signature blocks: + + { + "type": "object", + "extend:type": "signature", + "properties": { + "printed_name": {"type": ["string", "null"]}, + "signature_date": {"type": ["string", "null"], "extend:type": "date"}, + "is_signed": {"type": ["boolean", "null"]}, + "title_or_role": {"type": ["string", "null"]} + }, + "required": ["printed_name", "signature_date", "is_signed", "title_or_role"] + } + """ + + __extend_type__: typing.ClassVar[str] = "signature" + + printed_name: typing.Optional[str] = None + signature_date: typing.Optional[dt.date] = None + is_signed: typing.Optional[bool] = None + title_or_role: typing.Optional[str] = None + + +def get_extend_type(model: type) -> typing.Optional[str]: + """Return the extend:type marker for a model class, if it has one.""" + return getattr(model, "__extend_type__", None) diff --git a/src/extend_ai/wrapper/schema/typed_run.py b/src/extend_ai/wrapper/schema/typed_run.py new file mode 100644 index 0000000..e6fec4b --- /dev/null +++ b/src/extend_ai/wrapper/schema/typed_run.py @@ -0,0 +1,181 @@ +""" +Typed extract run wrappers. + +When an extraction is created with a pydantic model as its schema, the SDK +returns a :class:`TypedExtractRun` whose ``output.value`` (and +``initial_output.value`` / ``reviewed_output.value``) are validated instances +of that model instead of plain dicts. +""" + +import typing + +import pydantic +from ...core.pydantic_utilities import IS_PYDANTIC_V2 +from ...types.created_at import CreatedAt +from ...types.extract_config import ExtractConfig +from ...types.extract_output import ExtractOutput +from ...types.extract_output_edits import ExtractOutputEdits +from ...types.extract_output_metadata import ExtractOutputMetadata +from ...types.extract_run import ExtractRun +from ...types.extractor_summary import ExtractorSummary +from ...types.extractor_version_summary import ExtractorVersionSummary +from ...types.file_summary import FileSummary +from ...types.processor_run_status import ProcessorRunStatus +from ...types.run_metadata import RunMetadata +from ...types.run_usage import RunUsage +from ...types.updated_at import UpdatedAt + +__all__ = [ + "ExtractOutputValidationError", + "ModelT", + "TypedExtractOutput", + "TypedExtractRun", + "parse_extract_run", +] + +ModelT = typing.TypeVar("ModelT", bound=pydantic.BaseModel) + + +class ExtractOutputValidationError(Exception): + """ + Raised when a completed extract run's output does not validate against the + pydantic model that was used as the extraction schema. + + The run itself completed successfully — only the client-side validation + failed — so the full run is preserved on the ``run`` attribute (including + ``run.id``, ``run.dashboard_url``, and the raw ``run.output``). + """ + + def __init__(self, message: str, run: ExtractRun): + self.run = run + super().__init__(message) + + +def _validate_model(model: typing.Type[ModelT], value: typing.Any) -> ModelT: + if IS_PYDANTIC_V2: + return typing.cast(ModelT, model.model_validate(value)) # type: ignore[attr-defined] + return typing.cast(ModelT, model.parse_obj(value)) + + +class TypedExtractOutput(typing.Generic[ModelT]): + """Extract output whose value is a validated pydantic model instance.""" + + value: ModelT + metadata: typing.Optional[ExtractOutputMetadata] + + def __init__(self, *, value: ModelT, metadata: typing.Optional[ExtractOutputMetadata]) -> None: + self.value = value + self.metadata = metadata + + def __repr__(self) -> str: + return f"TypedExtractOutput(value={self.value!r})" + + +class TypedExtractRun(typing.Generic[ModelT]): + """ + An extract run whose outputs are validated instances of the pydantic model + that was used as the extraction schema. + + Mirrors :class:`~extend_ai.types.extract_run.ExtractRun`; the original + response is available as ``raw``. + """ + + object: str + id: str + status: ProcessorRunStatus + output: typing.Optional[TypedExtractOutput[ModelT]] + initial_output: typing.Optional[TypedExtractOutput[ModelT]] + reviewed_output: typing.Optional[TypedExtractOutput[ModelT]] + failure_reason: typing.Optional[str] + failure_message: typing.Optional[str] + metadata: typing.Optional[RunMetadata] + reviewed: bool + edited: bool + edits: typing.Optional[typing.Dict[str, typing.Optional[ExtractOutputEdits]]] + config: ExtractConfig + extractor: typing.Optional[ExtractorSummary] + extractor_version: typing.Optional[ExtractorVersionSummary] + file: typing.Optional[FileSummary] + files: typing.Optional[typing.List[FileSummary]] + parse_run_id: typing.Optional[str] + dashboard_url: str + usage: typing.Optional[RunUsage] + created_at: CreatedAt + updated_at: UpdatedAt + raw: ExtractRun + """The original, untyped extract run response.""" + + def __init__(self, run: ExtractRun, model: typing.Type[ModelT]) -> None: + self.raw = run + self.object = run.object + self.id = run.id + self.status = run.status + self.output = _parse_output(run.output, model, run) + self.initial_output = _parse_output(run.initial_output, model, run) + self.reviewed_output = _parse_output(run.reviewed_output, model, run) + self.failure_reason = run.failure_reason + self.failure_message = run.failure_message + self.metadata = run.metadata + self.reviewed = run.reviewed + self.edited = run.edited + self.edits = run.edits + self.config = run.config + self.extractor = run.extractor + self.extractor_version = run.extractor_version + self.file = run.file + self.files = run.files + self.parse_run_id = run.parse_run_id + self.dashboard_url = run.dashboard_url + self.usage = run.usage + self.created_at = run.created_at + self.updated_at = run.updated_at + + def __repr__(self) -> str: + return f"TypedExtractRun(id={self.id!r}, status={self.status!r}, output={self.output!r})" + + +def _parse_output( + output: typing.Optional[ExtractOutput], model: typing.Type[ModelT], run: ExtractRun +) -> typing.Optional[TypedExtractOutput[ModelT]]: + if output is None: + return None + value = getattr(output, "value", None) + if value is None: + raise ExtractOutputValidationError( + f"Extract run {getattr(run, 'id', None)!r} has no 'value' on its output; typed schemas are " + "only supported for runs created with a JSON Schema config. " + "The full run is available on this error's `run` attribute.", + run=run, + ) + try: + validated = _validate_model(model, value) + except pydantic.ValidationError as exc: + raise ExtractOutputValidationError( + f"Output of extract run {getattr(run, 'id', None)!r} did not validate against " + f"{model.__name__}: {exc}\n" + "The run completed successfully; the full run (including its raw output) is " + "available on this error's `run` attribute.", + run=run, + ) from exc + return TypedExtractOutput( + value=validated, + metadata=getattr(output, "metadata", None), + ) + + +def parse_extract_run(run: ExtractRun, model: typing.Type[ModelT]) -> TypedExtractRun[ModelT]: + """ + Validate an extract run's outputs against a pydantic model. + + Args: + run: A completed extract run. + model: The pydantic model class that was used as the extraction schema. + + Returns: + A :class:`TypedExtractRun` whose output values are instances of ``model``. + + Raises: + ExtractOutputValidationError: If an output value does not conform to the + model. The completed run is preserved on the error's ``run`` attribute. + """ + return TypedExtractRun(run, model) diff --git a/tests/wrapper/test_schema_conversion.py b/tests/wrapper/test_schema_conversion.py new file mode 100644 index 0000000..6fc42e5 --- /dev/null +++ b/tests/wrapper/test_schema_conversion.py @@ -0,0 +1,645 @@ +"""Tests for pydantic-to-Extend JSON Schema conversion.""" + +import datetime as dt +import enum +import typing +from typing import Dict, List, Optional, Union + +import pydantic +import pytest +from typing_extensions import Literal + +from extend_ai.wrapper.schema import ( + ExtendCurrency, + ExtendDate, + ExtendSignature, + SchemaConversionError, + pydantic_to_extend_schema, +) + +CURRENCY_SCHEMA = { + "type": "object", + "extend:type": "currency", + "properties": { + "amount": {"type": ["number", "null"]}, + "iso_4217_currency_code": {"type": ["string", "null"]}, + }, + "required": ["amount", "iso_4217_currency_code"], + "additionalProperties": False, +} + +SIGNATURE_SCHEMA = { + "type": "object", + "extend:type": "signature", + "properties": { + "printed_name": {"type": ["string", "null"]}, + "signature_date": {"type": ["string", "null"], "extend:type": "date"}, + "is_signed": {"type": ["boolean", "null"]}, + "title_or_role": {"type": ["string", "null"]}, + }, + "required": ["printed_name", "signature_date", "is_signed", "title_or_role"], + "additionalProperties": False, +} + + +class TestBasicSchemaCreation: + def test_generates_valid_json_schema_from_model(self): + class Schema(pydantic.BaseModel): + name: Optional[str] = None + + assert pydantic_to_extend_schema(Schema) == { + "type": "object", + "properties": {"name": {"type": ["string", "null"]}}, + "required": ["name"], + "additionalProperties": False, + } + + def test_sets_additional_properties_false_at_root(self): + class Schema(pydantic.BaseModel): + name: Optional[str] = None + + assert pydantic_to_extend_schema(Schema)["additionalProperties"] is False + + def test_adds_all_properties_to_required(self): + class Schema(pydantic.BaseModel): + field1: Optional[str] = None + field2: Optional[float] = None + field3: Optional[bool] = None + + assert pydantic_to_extend_schema(Schema)["required"] == ["field1", "field2", "field3"] + + def test_rejects_non_model_schema(self): + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(dict) # type: ignore[arg-type] + + +class TestPrimitiveTypes: + def test_converts_nullable_string(self): + class Schema(pydantic.BaseModel): + field: Optional[str] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["field"] == {"type": ["string", "null"]} + + def test_converts_nullable_number(self): + class Schema(pydantic.BaseModel): + field: Optional[float] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["field"] == {"type": ["number", "null"]} + + def test_converts_nullable_integer(self): + class Schema(pydantic.BaseModel): + field: Optional[int] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["field"] == {"type": ["integer", "null"]} + + def test_converts_nullable_boolean(self): + class Schema(pydantic.BaseModel): + field: Optional[bool] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["field"] == {"type": ["boolean", "null"]} + + def test_non_optional_primitives_are_rejected(self): + class Schema(pydantic.BaseModel): + name: str + + with pytest.raises(SchemaConversionError, match=r"Optional\[str\].*name"): + pydantic_to_extend_schema(Schema) + + @pytest.mark.parametrize( + "annotation,expected_hint", + [(str, "str"), (int, "int"), (float, "float"), (bool, "bool"), (dt.date, "datetime.date")], + ids=["str", "int", "float", "bool", "date"], + ) + def test_each_non_optional_nullable_kind_is_rejected(self, annotation, expected_hint): + Schema = pydantic.create_model("Schema", field=(annotation, ...)) + + with pytest.raises(SchemaConversionError, match=f"Optional\\[{expected_hint}\\]"): + pydantic_to_extend_schema(Schema) + + def test_non_optional_primitive_with_default_is_still_rejected(self): + # A default only covers a *missing* field; extraction returns explicit + # nulls, which the model would still reject. + class Schema(pydantic.BaseModel): + name: str = "unknown" + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_includes_descriptions(self): + class Schema(pydantic.BaseModel): + name: Optional[str] = pydantic.Field(None, description="The customer name") + age: Optional[float] = pydantic.Field(None, description="Customer age in years") + + properties = pydantic_to_extend_schema(Schema)["properties"] + assert properties["name"] == {"type": ["string", "null"], "description": "The customer name"} + assert properties["age"] == {"type": ["number", "null"], "description": "Customer age in years"} + + +class TestEnumTypes: + def test_converts_literal_with_null_added(self): + class Schema(pydantic.BaseModel): + status: Optional[Literal["active", "inactive"]] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["status"] == {"enum": ["active", "inactive", None]} + + def test_converts_string_enum_with_null_added(self): + class Status(str, enum.Enum): + ACTIVE = "active" + INACTIVE = "inactive" + + class Schema(pydantic.BaseModel): + status: Optional[Status] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["status"] == {"enum": ["active", "inactive", None]} + + def test_preserves_description_on_enums(self): + class Schema(pydantic.BaseModel): + status: Optional[Literal["active", "inactive"]] = pydantic.Field(None, description="Account status") + + assert pydantic_to_extend_schema(Schema)["properties"]["status"] == { + "enum": ["active", "inactive", None], + "description": "Account status", + } + + def test_does_not_duplicate_null_in_literal(self): + class Schema(pydantic.BaseModel): + status: Optional[Literal["active", "inactive", None]] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["status"] == {"enum": ["active", "inactive", None]} + + def test_converts_single_string_literal(self): + class Schema(pydantic.BaseModel): + type: Optional[Literal["invoice"]] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["type"] == {"enum": ["invoice", None]} + + def test_rejects_non_string_literals(self): + class Schema(pydantic.BaseModel): + value: Optional[Literal[42]] = None + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_rejects_non_string_enums(self): + class Number(enum.Enum): + ONE = 1 + TWO = 2 + + class Schema(pydantic.BaseModel): + value: Optional[Number] = None + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_typing_literal_is_recognized(self): + # typing.Literal and typing_extensions.Literal are distinct objects on + # some Python versions; both must be treated as enums. + class Schema(pydantic.BaseModel): + status: Optional[typing.Literal["active", "inactive"]] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["status"] == {"enum": ["active", "inactive", None]} + + def test_literal_including_none_is_nullable_without_optional(self): + class Schema(pydantic.BaseModel): + status: Literal["active", "inactive", None] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["status"] == {"enum": ["active", "inactive", None]} + + def test_rejects_non_optional_literal(self): + class Schema(pydantic.BaseModel): + status: Literal["active", "inactive"] + + with pytest.raises(SchemaConversionError, match="Optional"): + pydantic_to_extend_schema(Schema) + + def test_rejects_non_optional_string_enum(self): + class Status(str, enum.Enum): + ACTIVE = "active" + + class Schema(pydantic.BaseModel): + status: Status + + with pytest.raises(SchemaConversionError, match="Optional"): + pydantic_to_extend_schema(Schema) + + +class TestArrayTypes: + def test_converts_array_of_objects(self): + class Item(pydantic.BaseModel): + name: Optional[str] = None + price: Optional[float] = None + + class Schema(pydantic.BaseModel): + items: List[Item] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["items"] == { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": {"type": ["string", "null"]}, + "price": {"type": ["number", "null"]}, + }, + "required": ["name", "price"], + "additionalProperties": False, + }, + } + + def test_converts_array_of_strings_non_nullable_items(self): + class Schema(pydantic.BaseModel): + tags: List[str] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["tags"] == { + "type": "array", + "items": {"type": "string"}, + } + + def test_converts_array_of_numbers(self): + class Schema(pydantic.BaseModel): + values: List[float] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["values"] == { + "type": "array", + "items": {"type": "number"}, + } + + def test_converts_array_of_integers(self): + class Schema(pydantic.BaseModel): + counts: List[int] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["counts"] == { + "type": "array", + "items": {"type": "integer"}, + } + + def test_converts_array_of_booleans(self): + class Schema(pydantic.BaseModel): + flags: List[bool] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["flags"] == { + "type": "array", + "items": {"type": "boolean"}, + } + + def test_converts_optional_array(self): + class Schema(pydantic.BaseModel): + tags: Optional[List[str]] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["tags"] == { + "type": "array", + "items": {"type": "string"}, + } + + def test_includes_description_on_arrays(self): + class Schema(pydantic.BaseModel): + items: List[str] = pydantic.Field(default_factory=list, description="List of items") + + assert pydantic_to_extend_schema(Schema)["properties"]["items"] == { + "type": "array", + "items": {"type": "string"}, + "description": "List of items", + } + + +class TestNestedObjects: + def test_converts_nested_objects(self): + class Address(pydantic.BaseModel): + street: Optional[str] = None + city: Optional[str] = None + zip: Optional[str] = None + + class Schema(pydantic.BaseModel): + address: Optional[Address] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["address"] == { + "type": "object", + "properties": { + "street": {"type": ["string", "null"]}, + "city": {"type": ["string", "null"]}, + "zip": {"type": ["string", "null"]}, + }, + "required": ["street", "city", "zip"], + "additionalProperties": False, + } + + def test_handles_deeply_nested_objects(self): + class Level3(pydantic.BaseModel): + value: Optional[str] = None + + class Level2(pydantic.BaseModel): + level3: Optional[Level3] = None + + class Level1(pydantic.BaseModel): + level2: Optional[Level2] = None + + class Schema(pydantic.BaseModel): + level1: Optional[Level1] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["level1"] == { + "type": "object", + "properties": { + "level2": { + "type": "object", + "properties": { + "level3": { + "type": "object", + "properties": {"value": {"type": ["string", "null"]}}, + "required": ["value"], + "additionalProperties": False, + }, + }, + "required": ["level3"], + "additionalProperties": False, + }, + }, + "required": ["level2"], + "additionalProperties": False, + } + + def test_preserves_description_on_nested_objects(self): + class Address(pydantic.BaseModel): + street: Optional[str] = None + + class Schema(pydantic.BaseModel): + address: Optional[Address] = pydantic.Field(None, description="Mailing address") + + assert pydantic_to_extend_schema(Schema)["properties"]["address"]["description"] == "Mailing address" + + +class TestExtendDate: + def test_converts_to_extend_type_date(self): + class Schema(pydantic.BaseModel): + invoice_date: ExtendDate = None + + assert pydantic_to_extend_schema(Schema)["properties"]["invoice_date"] == { + "type": ["string", "null"], + "extend:type": "date", + } + + def test_plain_date_annotation_converts_to_extend_type_date(self): + class Schema(pydantic.BaseModel): + invoice_date: Optional[dt.date] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["invoice_date"] == { + "type": ["string", "null"], + "extend:type": "date", + } + + def test_preserves_description(self): + class Schema(pydantic.BaseModel): + invoice_date: ExtendDate = pydantic.Field(None, description="The invoice date") + + assert pydantic_to_extend_schema(Schema)["properties"]["invoice_date"] == { + "type": ["string", "null"], + "extend:type": "date", + "description": "The invoice date", + } + + def test_works_in_arrays_with_non_nullable_format(self): + class Schema(pydantic.BaseModel): + dates: List[dt.date] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["dates"] == { + "type": "array", + "items": {"type": "string", "extend:type": "date"}, + } + + def test_rejects_datetime(self): + class Schema(pydantic.BaseModel): + timestamp: Optional[dt.datetime] = None + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + +class TestExtendCurrency: + def test_converts_to_extend_type_currency(self): + class Schema(pydantic.BaseModel): + total: Optional[ExtendCurrency] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["total"] == CURRENCY_SCHEMA + + def test_preserves_description(self): + class Schema(pydantic.BaseModel): + total: Optional[ExtendCurrency] = pydantic.Field(None, description="Total invoice amount") + + expected = dict(CURRENCY_SCHEMA, description="Total invoice amount") + assert pydantic_to_extend_schema(Schema)["properties"]["total"] == expected + + def test_works_in_arrays(self): + class Schema(pydantic.BaseModel): + amounts: List[ExtendCurrency] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["amounts"] == { + "type": "array", + "items": CURRENCY_SCHEMA, + } + + +class TestExtendSignature: + def test_converts_to_extend_type_signature(self): + class Schema(pydantic.BaseModel): + customer_signature: Optional[ExtendSignature] = None + + assert pydantic_to_extend_schema(Schema)["properties"]["customer_signature"] == SIGNATURE_SCHEMA + + def test_preserves_description(self): + class Schema(pydantic.BaseModel): + signature: Optional[ExtendSignature] = pydantic.Field(None, description="Customer signature") + + assert pydantic_to_extend_schema(Schema)["properties"]["signature"]["description"] == "Customer signature" + + def test_works_in_arrays(self): + class Schema(pydantic.BaseModel): + signatures: List[ExtendSignature] = [] + + assert pydantic_to_extend_schema(Schema)["properties"]["signatures"] == { + "type": "array", + "items": SIGNATURE_SCHEMA, + } + + +class TestSchemaConversionError: + def test_includes_path_in_error_message(self): + error = SchemaConversionError("Unsupported type", ["items", "nested", "field"]) + + assert str(error) == "Unsupported type at path: items.nested.field" + assert error.path == ["items", "nested", "field"] + + def test_works_without_path(self): + error = SchemaConversionError("General error") + + assert str(error) == "General error" + assert error.path == [] + + def test_conversion_errors_carry_field_path(self): + class Inner(pydantic.BaseModel): + mapping: Dict[str, str] = {} + + class Schema(pydantic.BaseModel): + inner: Optional[Inner] = None + + with pytest.raises(SchemaConversionError) as exc_info: + pydantic_to_extend_schema(Schema) + assert exc_info.value.path == ["inner", "mapping"] + + +class TestUnsupportedTypes: + def test_rejects_dict_fields(self): + class Schema(pydantic.BaseModel): + mapping: Dict[str, str] = {} + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_rejects_non_optional_unions(self): + class Schema(pydantic.BaseModel): + value: Union[str, int] = "" + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_rejects_array_of_enums(self): + class Schema(pydantic.BaseModel): + statuses: List[Literal["a", "b"]] = [] + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_rejects_array_of_string_enums(self): + class Status(str, enum.Enum): + A = "a" + + class Schema(pydantic.BaseModel): + statuses: List[Status] = [] + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_rejects_nested_arrays(self): + class Schema(pydantic.BaseModel): + matrix: List[List[str]] = [] + + with pytest.raises(SchemaConversionError): + pydantic_to_extend_schema(Schema) + + def test_rejects_optional_array_items(self): + class Schema(pydantic.BaseModel): + tags: List[Optional[str]] = [] + + with pytest.raises(SchemaConversionError, match="Array items must not be Optional"): + pydantic_to_extend_schema(Schema) + + +class TestRecursiveModels: + def test_rejects_directly_recursive_model(self): + class Node(pydantic.BaseModel): + name: Optional[str] = None + children: List["Node"] = [] + + if hasattr(Node, "model_rebuild"): + Node.model_rebuild() + else: + Node.update_forward_refs(Node=Node) + + with pytest.raises(SchemaConversionError, match="Recursive"): + pydantic_to_extend_schema(Node) + + def test_rejects_mutually_recursive_models(self): + class A(pydantic.BaseModel): + b: Optional["B"] = None + + class B(pydantic.BaseModel): + a: Optional[A] = None + + if hasattr(A, "model_rebuild"): + A.model_rebuild() + else: + A.update_forward_refs(B=B) + + with pytest.raises(SchemaConversionError, match="Recursive"): + pydantic_to_extend_schema(A) + + def test_allows_same_model_in_sibling_fields(self): + class Address(pydantic.BaseModel): + street: Optional[str] = None + + class Schema(pydantic.BaseModel): + billing_address: Optional[Address] = None + shipping_address: Optional[Address] = None + + json_schema = pydantic_to_extend_schema(Schema) + assert json_schema["properties"]["billing_address"] == json_schema["properties"]["shipping_address"] + + +class TestFieldAliases: + def test_rejects_aliased_fields(self): + class Schema(pydantic.BaseModel): + invoice_number: Optional[str] = pydantic.Field(None, alias="invoiceNumber") + + with pytest.raises(SchemaConversionError, match="alias"): + pydantic_to_extend_schema(Schema) + + def test_rejects_aliased_fields_in_nested_models(self): + class Inner(pydantic.BaseModel): + value: Optional[str] = pydantic.Field(None, alias="theValue") + + class Schema(pydantic.BaseModel): + inner: Optional[Inner] = None + + with pytest.raises(SchemaConversionError) as exc_info: + pydantic_to_extend_schema(Schema) + assert exc_info.value.path == ["inner", "value"] + + +class TestComplexSchemas: + def test_converts_realistic_invoice_schema(self): + class Vendor(pydantic.BaseModel): + name: Optional[str] = pydantic.Field(None, description="Vendor company name") + address: Optional[str] = pydantic.Field(None, description="Vendor address") + + class LineItem(pydantic.BaseModel): + description: Optional[str] = None + quantity: Optional[float] = None + unit_price: Optional[ExtendCurrency] = None + line_total: Optional[ExtendCurrency] = None + + class Invoice(pydantic.BaseModel): + invoice_number: Optional[str] = pydantic.Field(None, description="The invoice number") + invoice_date: ExtendDate = pydantic.Field(None, description="The invoice date") + due_date: ExtendDate = pydantic.Field(None, description="Payment due date") + vendor: Optional[Vendor] = pydantic.Field(None, description="Vendor information") + total_amount: Optional[ExtendCurrency] = pydantic.Field(None, description="Total invoice amount") + line_items: List[LineItem] = pydantic.Field(default_factory=list, description="Invoice line items") + status: Optional[Literal["draft", "sent", "paid", "overdue"]] = pydantic.Field( + None, description="Invoice status" + ) + + json_schema = pydantic_to_extend_schema(Invoice) + + assert json_schema["type"] == "object" + assert "invoice_number" in json_schema["required"] + assert "invoice_date" in json_schema["required"] + assert "line_items" in json_schema["required"] + + assert json_schema["properties"]["invoice_date"]["extend:type"] == "date" + assert json_schema["properties"]["total_amount"]["extend:type"] == "currency" + assert json_schema["properties"]["line_items"]["items"]["properties"]["unit_price"]["extend:type"] == "currency" + assert json_schema["properties"]["status"]["enum"] == ["draft", "sent", "paid", "overdue", None] + + def test_converts_contract_schema_with_signatures(self): + class Term(pydantic.BaseModel): + section: Optional[str] = None + content: Optional[str] = None + + class Contract(pydantic.BaseModel): + contract_id: Optional[str] = None + effective_date: ExtendDate = None + party_a_signature: Optional[ExtendSignature] = pydantic.Field(None, description="Party A signature") + party_b_signature: Optional[ExtendSignature] = pydantic.Field(None, description="Party B signature") + terms: List[Term] = [] + + json_schema = pydantic_to_extend_schema(Contract) + + assert json_schema["properties"]["party_a_signature"]["extend:type"] == "signature" + assert json_schema["properties"]["party_b_signature"]["extend:type"] == "signature" diff --git a/tests/wrapper/test_signature_parity.py b/tests/wrapper/test_signature_parity.py new file mode 100644 index 0000000..1de6c6b --- /dev/null +++ b/tests/wrapper/test_signature_parity.py @@ -0,0 +1,169 @@ +""" +Guards against wrapper code drifting from the Fern-generated SDK. + +The wrapper layer re-declares parts of the generated API surface: + +- `create_and_poll()` mirrors each generated `create()` signature +- `TypedExtractConfigParams` / `TypedExtractorParams` mirror the generated + request TypedDicts (with `schema` retyped to a pydantic model class) +- `TypedExtractRun` mirrors the fields of the generated `ExtractRun` + +When SDK regeneration adds a parameter, key, or field, these tests fail so the +wrapper gets updated in the same change. (Methods that *override* a generated +method, like `Extend.extract()` and `ExtractorsClient.create()`, are already +covered: mypy rejects overrides whose signatures are incompatible with the +generated superclass.) +""" + +import inspect +import typing +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from extend_ai.core.pydantic_utilities import IS_PYDANTIC_V2 +from extend_ai.wrapper.resources import ( + classify_runs, + edit_runs, + extract_runs, + parse_runs, + split_runs, + workflow_runs, +) + + +def _param_names(func: typing.Any) -> typing.Set[str]: + return {name for name in inspect.signature(func).parameters if name != "self"} + + +def _typed_dict_keys(td: typing.Any) -> typing.Set[str]: + return set(td.__annotations__) + + +RUN_CLIENTS = [ + classify_runs.ClassifyRunsClient, + classify_runs.AsyncClassifyRunsClient, + edit_runs.EditRunsClient, + edit_runs.AsyncEditRunsClient, + extract_runs.ExtractRunsClient, + extract_runs.AsyncExtractRunsClient, + parse_runs.ParseRunsClient, + parse_runs.AsyncParseRunsClient, + split_runs.SplitRunsClient, + split_runs.AsyncSplitRunsClient, + workflow_runs.WorkflowRunsClient, + workflow_runs.AsyncWorkflowRunsClient, +] + + +@pytest.mark.parametrize("wrapper_client", RUN_CLIENTS, ids=lambda cls: cls.__name__) +def test_create_and_poll_accepts_all_create_params(wrapper_client): + """Every parameter of the generated create() must be exposed by create_and_poll().""" + generated_client = wrapper_client.__mro__[1] + create_params = _param_names(generated_client.create) - {"request_options"} + create_and_poll_params = _param_names(wrapper_client.create_and_poll) + + missing = create_params - create_and_poll_params + assert not missing, ( + f"{wrapper_client.__name__}.create_and_poll() is missing parameters that " + f"{generated_client.__name__}.create() accepts: {sorted(missing)}. " + "Add them to create_and_poll() and forward them to create()." + ) + + +SYNC_RUN_CLIENTS = [cls for cls in RUN_CLIENTS if not cls.__name__.startswith("Async")] +ASYNC_RUN_CLIENTS = [cls for cls in RUN_CLIENTS if cls.__name__.startswith("Async")] + + +def _create_sentinel_kwargs(wrapper_client) -> typing.Dict[str, object]: + """One unique sentinel per generated create() parameter.""" + generated_client = wrapper_client.__mro__[1] + return {name: object() for name in _param_names(generated_client.create) - {"request_options"}} + + +def _bind_create_and_poll(wrapper_client, mock_cls): + client = MagicMock(spec=wrapper_client) + client.create = mock_cls() + client.retrieve = mock_cls() + client.retrieve.return_value = MagicMock(status="PROCESSED") + client.create_and_poll = wrapper_client.create_and_poll.__get__(client, wrapper_client) + return client + + +def _assert_all_forwarded(client, sentinels: typing.Dict[str, object], wrapper_client) -> None: + create_kwargs = client.create.call_args.kwargs + not_forwarded = [name for name, sentinel in sentinels.items() if create_kwargs.get(name) is not sentinel] + assert not not_forwarded, ( + f"{wrapper_client.__name__}.create_and_poll() accepted these parameters but did not " + f"forward them to create(): {sorted(not_forwarded)}" + ) + + +@pytest.mark.parametrize("wrapper_client", SYNC_RUN_CLIENTS, ids=lambda cls: cls.__name__) +def test_create_and_poll_forwards_all_create_params(wrapper_client): + """Accepting a parameter is not enough — it must reach create().""" + client = _bind_create_and_poll(wrapper_client, MagicMock) + sentinels = _create_sentinel_kwargs(wrapper_client) + + client.create_and_poll(**sentinels) + + _assert_all_forwarded(client, sentinels, wrapper_client) + + +@pytest.mark.parametrize("wrapper_client", ASYNC_RUN_CLIENTS, ids=lambda cls: cls.__name__) +async def test_async_create_and_poll_forwards_all_create_params(wrapper_client): + client = _bind_create_and_poll(wrapper_client, AsyncMock) + sentinels = _create_sentinel_kwargs(wrapper_client) + + await client.create_and_poll(**sentinels) + + _assert_all_forwarded(client, sentinels, wrapper_client) + + +def test_typed_extract_config_matches_generated_config_keys(): + from extend_ai.requests.extract_config_json import ExtractConfigJsonParams + from extend_ai.wrapper.schema import TypedExtractConfigParams + + generated_keys = _typed_dict_keys(ExtractConfigJsonParams) + typed_keys = _typed_dict_keys(TypedExtractConfigParams) + + missing = generated_keys - typed_keys + assert not missing, ( + f"TypedExtractConfigParams is missing keys that ExtractConfigJsonParams has: {sorted(missing)}. " + "Add them so typed configs accept the same options as untyped configs." + ) + extra = typed_keys - generated_keys + assert not extra, f"TypedExtractConfigParams has keys the generated config does not: {sorted(extra)}" + + +def test_typed_extractor_matches_generated_extractor_keys(): + from extend_ai.extract_runs.requests.extract_runs_create_request_extractor import ( + ExtractRunsCreateRequestExtractorParams, + ) + from extend_ai.wrapper.schema import TypedExtractorParams + + generated_keys = _typed_dict_keys(ExtractRunsCreateRequestExtractorParams) + typed_keys = _typed_dict_keys(TypedExtractorParams) + + assert generated_keys == typed_keys, ( + f"TypedExtractorParams keys {sorted(typed_keys)} differ from generated " + f"ExtractRunsCreateRequestExtractorParams keys {sorted(generated_keys)}." + ) + + +def test_typed_extract_run_mirrors_extract_run_fields(): + from extend_ai.types.extract_run import ExtractRun + from extend_ai.wrapper.schema import TypedExtractRun + + if IS_PYDANTIC_V2: + run_fields = set(ExtractRun.model_fields) + else: + run_fields = set(ExtractRun.__fields__) + + typed_fields = set(TypedExtractRun.__annotations__) + + missing = run_fields - typed_fields + assert not missing, ( + f"TypedExtractRun is missing fields that ExtractRun has: {sorted(missing)}. " + "Add the attributes to TypedExtractRun and copy them in its constructor." + ) diff --git a/tests/wrapper/test_typed_extraction.py b/tests/wrapper/test_typed_extraction.py new file mode 100644 index 0000000..22c33a0 --- /dev/null +++ b/tests/wrapper/test_typed_extraction.py @@ -0,0 +1,356 @@ +"""Tests for typed (pydantic schema) extraction across the wrapper clients.""" + +import datetime as dt +from typing import List, Optional +from unittest.mock import MagicMock + +import pydantic +import pytest + +from extend_ai.wrapper.schema import ( + ExtendCurrency, + ExtendDate, + ExtractOutputValidationError, + TypedExtractOutput, + TypedExtractRun, + parse_extract_run, +) + + +class LineItem(pydantic.BaseModel): + description: Optional[str] = None + amount: Optional[ExtendCurrency] = None + + +class Invoice(pydantic.BaseModel): + invoice_number: Optional[str] = None + invoice_date: ExtendDate = None + total: Optional[ExtendCurrency] = None + line_items: List[LineItem] = [] + + +INVOICE_OUTPUT_VALUE = { + "invoice_number": "INV-123", + "invoice_date": "2026-01-15", + "total": {"amount": 99.5, "iso_4217_currency_code": "USD"}, + "line_items": [ + {"description": "Widget", "amount": {"amount": 99.5, "iso_4217_currency_code": "USD"}}, + ], +} + + +def create_mock_run(status: str = "PROCESSED", value=None): + """Create a mock extract run with a JSON-schema output.""" + run = MagicMock() + run.id = "extract_run_test123" + run.status = status + run.object = "extract_run" + if value is None: + run.output = None + else: + run.output = MagicMock() + run.output.value = value + run.initial_output = None + run.reviewed_output = None + return run + + +# ============================================================================ +# parse_extract_run +# ============================================================================ + + +class TestParseExtractRun: + def test_parses_output_value_into_model_instance(self): + run = create_mock_run(value=INVOICE_OUTPUT_VALUE) + + typed = parse_extract_run(run, Invoice) + + assert isinstance(typed, TypedExtractRun) + assert isinstance(typed.output, TypedExtractOutput) + assert isinstance(typed.output.value, Invoice) + assert typed.output.value.invoice_number == "INV-123" + assert typed.output.value.invoice_date == dt.date(2026, 1, 15) + assert isinstance(typed.output.value.total, ExtendCurrency) + assert typed.output.value.total.amount == 99.5 + assert typed.output.value.line_items[0].description == "Widget" + + def test_copies_run_fields_and_keeps_raw(self): + run = create_mock_run(value=INVOICE_OUTPUT_VALUE) + + typed = parse_extract_run(run, Invoice) + + assert typed.id == run.id + assert typed.status == run.status + assert typed.raw is run + + def test_handles_none_outputs(self): + run = create_mock_run(status="FAILED", value=None) + + typed = parse_extract_run(run, Invoice) + + assert typed.output is None + assert typed.initial_output is None + assert typed.reviewed_output is None + + def test_null_field_values_validate_into_none(self): + run = create_mock_run( + value={"invoice_number": None, "invoice_date": None, "total": None, "line_items": []} + ) + + typed = parse_extract_run(run, Invoice) + + assert typed.output.value.invoice_number is None + assert typed.output.value.total is None + + def test_raises_for_output_without_value(self): + run = create_mock_run(value=INVOICE_OUTPUT_VALUE) + run.output = MagicMock(spec=[]) # legacy output shape: no `value` attribute + + with pytest.raises(ExtractOutputValidationError) as exc_info: + parse_extract_run(run, Invoice) + assert exc_info.value.run is run + + def test_validation_failure_preserves_completed_run_on_error(self): + run = create_mock_run(value={"invoice_number": "INV-1", "line_items": "not-a-list"}) + + with pytest.raises(ExtractOutputValidationError) as exc_info: + parse_extract_run(run, Invoice) + + error = exc_info.value + assert error.run is run + assert error.run.output.value["line_items"] == "not-a-list" + assert "Invoice" in str(error) + assert isinstance(error.__cause__, pydantic.ValidationError) + + +# ============================================================================ +# ExtractRunsClient.create_and_poll with typed schemas +# ============================================================================ + + +class TestCreateAndPollTypedSchema: + def setup_method(self): + from extend_ai.wrapper.resources.extract_runs import ExtractRunsClient + + self.wrapper = MagicMock(spec=ExtractRunsClient) + self.wrapper.create = MagicMock() + self.wrapper.retrieve = MagicMock() + self.wrapper.create_and_poll = ExtractRunsClient.create_and_poll.__get__(self.wrapper, ExtractRunsClient) + + def test_converts_config_schema_and_returns_typed_run(self): + self.wrapper.create.return_value = create_mock_run("PROCESSING") + self.wrapper.retrieve.return_value = create_mock_run("PROCESSED", value=INVOICE_OUTPUT_VALUE) + + result = self.wrapper.create_and_poll( + file={"id": "file_1"}, + config={"schema": Invoice, "base_processor": "extraction_performance"}, + ) + + create_kwargs = self.wrapper.create.call_args.kwargs + sent_schema = create_kwargs["config"]["schema"] + assert isinstance(sent_schema, dict) + assert sent_schema["type"] == "object" + assert sent_schema["properties"]["invoice_number"] == {"type": ["string", "null"]} + # Other config keys pass through untouched + assert create_kwargs["config"]["base_processor"] == "extraction_performance" + + assert isinstance(result, TypedExtractRun) + assert isinstance(result.output.value, Invoice) + assert result.output.value.invoice_number == "INV-123" + + def test_converts_extractor_override_config_schema(self): + self.wrapper.create.return_value = create_mock_run("PROCESSING") + self.wrapper.retrieve.return_value = create_mock_run("PROCESSED", value=INVOICE_OUTPUT_VALUE) + + result = self.wrapper.create_and_poll( + file={"id": "file_1"}, + extractor={"id": "extractor_abc", "override_config": {"schema": Invoice}}, + ) + + create_kwargs = self.wrapper.create.call_args.kwargs + sent_extractor = create_kwargs["extractor"] + assert sent_extractor["id"] == "extractor_abc" + assert isinstance(sent_extractor["override_config"]["schema"], dict) + assert sent_extractor["override_config"]["schema"]["type"] == "object" + + assert isinstance(result, TypedExtractRun) + + def test_untyped_config_passes_through_and_returns_plain_run(self): + self.wrapper.create.return_value = create_mock_run("PROCESSING") + processed = create_mock_run("PROCESSED", value=INVOICE_OUTPUT_VALUE) + self.wrapper.retrieve.return_value = processed + + json_config = {"schema": {"type": "object", "properties": {}}} + result = self.wrapper.create_and_poll(file={"id": "file_1"}, config=json_config) + + assert self.wrapper.create.call_args.kwargs["config"] is json_config + assert result is processed + assert not isinstance(result, TypedExtractRun) + + def test_typed_failed_run_has_no_output(self): + self.wrapper.create.return_value = create_mock_run("PROCESSING") + self.wrapper.retrieve.return_value = create_mock_run("FAILED", value=None) + + result = self.wrapper.create_and_poll(file={"id": "file_1"}, config={"schema": Invoice}) + + assert isinstance(result, TypedExtractRun) + assert result.status == "FAILED" + assert result.output is None + + +class TestAsyncCreateAndPollTypedSchema: + def setup_method(self): + from unittest.mock import AsyncMock + + from extend_ai.wrapper.resources.extract_runs import AsyncExtractRunsClient + + self.wrapper = MagicMock(spec=AsyncExtractRunsClient) + self.wrapper.create = AsyncMock() + self.wrapper.retrieve = AsyncMock() + self.wrapper.create_and_poll = AsyncExtractRunsClient.create_and_poll.__get__( + self.wrapper, AsyncExtractRunsClient + ) + + async def test_converts_config_schema_and_returns_typed_run(self): + self.wrapper.create.return_value = create_mock_run("PROCESSING") + self.wrapper.retrieve.return_value = create_mock_run("PROCESSED", value=INVOICE_OUTPUT_VALUE) + + result = await self.wrapper.create_and_poll(file={"id": "file_1"}, config={"schema": Invoice}) + + sent_schema = self.wrapper.create.call_args.kwargs["config"]["schema"] + assert isinstance(sent_schema, dict) + assert isinstance(result, TypedExtractRun) + assert result.output.value.invoice_number == "INV-123" + + +# ============================================================================ +# ExtractRunsClient.create (plain, non-polling) with typed schemas +# ============================================================================ + + +class TestPlainCreateTypedSchema: + def setup_method(self): + from extend_ai.wrapper.resources.extract_runs import ExtractRunsClient + + self.client = ExtractRunsClient(client_wrapper=MagicMock()) + self.raw_client = MagicMock() + self.client._raw_client = self.raw_client + + def test_create_converts_model_schema(self): + self.client.create(file={"id": "file_1"}, config={"schema": Invoice}) + + sent_config = self.raw_client.create.call_args.kwargs["config"] + assert isinstance(sent_config["schema"], dict) + assert sent_config["schema"]["type"] == "object" + + def test_create_converts_extractor_override_config_schema(self): + self.client.create( + file={"id": "file_1"}, + extractor={"id": "extractor_abc", "override_config": {"schema": Invoice}}, + ) + + sent_extractor = self.raw_client.create.call_args.kwargs["extractor"] + assert isinstance(sent_extractor["override_config"]["schema"], dict) + + def test_create_passes_through_json_schema(self): + json_config = {"schema": {"type": "object", "properties": {}}} + self.client.create(file={"id": "file_1"}, config=json_config) + + assert self.raw_client.create.call_args.kwargs["config"] is json_config + + +# ============================================================================ +# Extend.extract with typed schemas +# ============================================================================ + + +class TestClientExtractTypedSchema: + def setup_method(self): + from extend_ai.wrapper.client import Extend + + self.client = Extend(token="test-token") + self.raw_client = MagicMock() + self.client._raw_client = self.raw_client + + def test_converts_config_schema_and_returns_typed_run(self): + response = MagicMock() + response.data = create_mock_run("PROCESSED", value=INVOICE_OUTPUT_VALUE) + self.raw_client.extract.return_value = response + + result = self.client.extract( + file={"url": "https://example.com/invoice.pdf"}, + config={"schema": Invoice}, + ) + + sent_config = self.raw_client.extract.call_args.kwargs["config"] + assert isinstance(sent_config["schema"], dict) + assert sent_config["schema"]["type"] == "object" + + assert isinstance(result, TypedExtractRun) + assert isinstance(result.output.value, Invoice) + + def test_untyped_extract_returns_plain_run(self): + response = MagicMock() + run = create_mock_run("PROCESSED", value=INVOICE_OUTPUT_VALUE) + response.data = run + self.raw_client.extract.return_value = response + + result = self.client.extract(file={"url": "https://example.com/invoice.pdf"}) + + assert result is run + + +# ============================================================================ +# Extractors / ExtractorVersions with typed schemas +# ============================================================================ + + +class TestExtractorsTypedSchema: + def setup_method(self): + from extend_ai.wrapper.resources.extractors import ExtractorsClient + + self.client = ExtractorsClient(client_wrapper=MagicMock()) + self.raw_client = MagicMock() + self.client._raw_client = self.raw_client + + def test_create_converts_model_schema(self): + self.client.create(name="Invoice Extractor", config={"schema": Invoice}) + + sent_config = self.raw_client.create.call_args.kwargs["config"] + assert isinstance(sent_config["schema"], dict) + assert sent_config["schema"]["type"] == "object" + + def test_create_passes_through_json_schema(self): + json_config = {"schema": {"type": "object", "properties": {}}} + self.client.create(name="Invoice Extractor", config=json_config) + + assert self.raw_client.create.call_args.kwargs["config"] is json_config + + def test_update_converts_model_schema(self): + self.client.update("extractor_abc", config={"schema": Invoice}) + + sent_config = self.raw_client.update.call_args.kwargs["config"] + assert isinstance(sent_config["schema"], dict) + + def test_create_without_config_passes_omit(self): + self.client.create(name="Invoice Extractor") + + sent_config = self.raw_client.create.call_args.kwargs["config"] + assert sent_config is ... + + +class TestExtractorVersionsTypedSchema: + def setup_method(self): + from extend_ai.wrapper.resources.extractor_versions import ExtractorVersionsClient + + self.client = ExtractorVersionsClient(client_wrapper=MagicMock()) + self.raw_client = MagicMock() + self.client._raw_client = self.raw_client + + def test_create_converts_model_schema(self): + self.client.create("extractor_abc", release_type="minor", config={"schema": Invoice}) + + sent_config = self.raw_client.create.call_args.kwargs["config"] + assert isinstance(sent_config["schema"], dict) + assert sent_config["schema"]["properties"]["invoice_number"] == {"type": ["string", "null"]} diff --git a/uv.lock b/uv.lock index ce265ca..5bdf938 100644 --- a/uv.lock +++ b/uv.lock @@ -93,7 +93,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.21.2" }, { name = "pydantic", specifier = ">=1.9.2" }, { name = "pydantic-core", specifier = ">=2.18.2" }, - { name = "typing-extensions", specifier = ">=4.0.0" }, + { name = "typing-extensions", specifier = ">=4.3.0" }, ] [[package]]