retab 0.0.36__py3-none-any.whl → 0.0.38__py3-none-any.whl

{uiform → retab}/resources/documents/extractions.py

@@ -2,7 +2,7 @@ import base64
 import json
 from io import IOBase
 from pathlib import Path
-from typing import Any, AsyncGenerator, Generator, Literal, Optional
+from typing import Any, AsyncGenerator, Generator
 
 from anthropic.types.message_param import MessageParam
 from openai.types.chat import ChatCompletionMessageParam
@@ -10,6 +10,7 @@ from openai.types.chat.chat_completion_reasoning_effort import ChatCompletionRea
 from openai.types.chat.parsed_chat_completion import ParsedChatCompletionMessage
 from openai.types.responses.response import Response
 from openai.types.responses.response_input_param import ResponseInputItemParam
+from pydantic_core import PydanticUndefined
 from pydantic import HttpUrl
 
 from ..._resource import AsyncAPIResource, SyncAPIResource
@@ -17,8 +18,9 @@ from ..._utils.ai_models import assert_valid_model_extraction
 from ..._utils.json_schema import filter_auxiliary_fields_json, load_json_schema, unflatten_dict
 from ..._utils.mime import MIMEData, prepare_mime_document
 from ..._utils.stream_context_managers import as_async_context_manager, as_context_manager
-from ...types.chat import ChatCompletionUiformMessage
+from ...types.chat import ChatCompletionRetabMessage
 from ...types.documents.extractions import DocumentExtractRequest, LogExtractionRequest, UiParsedChatCompletion, UiParsedChatCompletionChunk, UiParsedChoice
+from ...types.browser_canvas import BrowserCanvas
 from ...types.modalities import Modality
 from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
@@ -31,7 +33,7 @@ def maybe_parse_to_pydantic(schema: Schema, response: UiParsedChatCompletion, al
                 response.choices[0].message.parsed = schema._partial_pydantic_model.model_validate(filter_auxiliary_fields_json(response.choices[0].message.content))
             else:
                 response.choices[0].message.parsed = schema.pydantic_model.model_validate(filter_auxiliary_fields_json(response.choices[0].message.content))
-        except Exception as e:
+        except Exception:
             pass
     return response
 
@@ -40,15 +42,16 @@ class BaseExtractionsMixin:
     def prepare_extraction(
         self,
         json_schema: dict[str, Any] | Path | str,
-        document: Path | str | IOBase | HttpUrl | None,
-        image_resolution_dpi: int | None,
-        browser_canvas: Literal['A3', 'A4', 'A5'] | None,
-        model: str,
-        temperature: float,
-        modality: Modality,
-        reasoning_effort: ChatCompletionReasoningEffort,
-        stream: bool,
-        n_consensus: int = 1,
+        document: Path | str | IOBase | HttpUrl | None = None,
+        documents: list[Path | str | IOBase | HttpUrl] | None = None,
+        image_resolution_dpi: int = PydanticUndefined,  # type: ignore[assignment]
+        browser_canvas: BrowserCanvas = PydanticUndefined,  # type: ignore[assignment]
+        model: str = PydanticUndefined,  # type: ignore[assignment]
+        temperature: float = PydanticUndefined,  # type: ignore[assignment]
+        modality: Modality = PydanticUndefined,  # type: ignore[assignment]
+        reasoning_effort: ChatCompletionReasoningEffort = PydanticUndefined,  # type: ignore[assignment]
+        stream: bool = False,
+        n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         store: bool = False,
         idempotency_key: str | None = None,
     ) -> PreparedRequest:
@@ -56,26 +59,36 @@ class BaseExtractionsMixin:
 
         json_schema = load_json_schema(json_schema)
 
-        data = {
-            "json_schema": json_schema,
-            "document": prepare_mime_document(document).model_dump() if document is not None else None,
-            "model": model,
-            "temperature": temperature,
-            "stream": stream,
-            "modality": modality,
-            "store": store,
-            "reasoning_effort": reasoning_effort,
-            "n_consensus": n_consensus,
-        }
-        if image_resolution_dpi:
-            data["image_resolution_dpi"] = image_resolution_dpi
-        if browser_canvas:
-            data["browser_canvas"] = browser_canvas
+        # Handle both single document and multiple documents
+        if document is not None and documents is not None:
+            raise ValueError("Cannot provide both 'document' and 'documents' parameters. Use either one.")
+
+        # Convert single document to documents list for consistency
+        if document is not None:
+            processed_documents = [prepare_mime_document(document)]
+        elif documents is not None:
+            processed_documents = [prepare_mime_document(doc) for doc in documents]
+        else:
+            raise ValueError("Must provide either 'document' or 'documents' parameter.")
 
         # Validate DocumentAPIRequest data (raises exception if invalid)
-        document_extract_request = DocumentExtractRequest.model_validate(data)
+        request = DocumentExtractRequest(
+            json_schema=json_schema,
+            documents=processed_documents,
+            model=model,
+            temperature=temperature,
+            stream=stream,
+            modality=modality,
+            store=store,
+            reasoning_effort=reasoning_effort,
+            n_consensus=n_consensus,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+        )
 
-        return PreparedRequest(method="POST", url="/v1/documents/extractions", data=document_extract_request.model_dump(), idempotency_key=idempotency_key)
+        return PreparedRequest(
+            method="POST", url="/v1/documents/extractions", data=request.model_dump(mode="json", exclude_unset=True, exclude_defaults=True), idempotency_key=idempotency_key
+        )
 
     def prepare_log_extraction(
         self,
@@ -84,8 +97,8 @@ class BaseExtractionsMixin:
         model: str,
         temperature: float,
         completion: Any | None = None,
-        # The messages can be provided in different formats, we will convert them to the UiForm-compatible format
-        messages: list[ChatCompletionUiformMessage] | None = None,
+        # The messages can be provided in different formats, we will convert them to the Retab-compatible format
+        messages: list[ChatCompletionRetabMessage] | None = None,
         openai_messages: list[ChatCompletionMessageParam] | None = None,
         anthropic_messages: list[MessageParam] | None = None,
         anthropic_system_prompt: str | None = None,
@@ -99,12 +112,14 @@ class BaseExtractionsMixin:
                 # url is a base64 encoded string with the mime type and the content. For the dummy one we will send a .txt file with the text "No document provided"
                 url="data:text/plain;base64," + base64.b64encode(b"No document provided").decode("utf-8"),
             )
+        else:
+            mime_document = prepare_mime_document(document)
 
         return PreparedRequest(
             method="POST",
             url="/v1/documents/log_extraction",
             data=LogExtractionRequest(
-                document=prepare_mime_document(document) if document else mime_document,
+                document=mime_document,
                 messages=messages,
                 openai_messages=openai_messages,
                 anthropic_messages=anthropic_messages,
@@ -115,7 +130,7 @@ class BaseExtractionsMixin:
                 json_schema=json_schema,
                 model=model,
                 temperature=temperature,
-            ).model_dump(mode="json", by_alias=True),  # by_alias is necessary to enable serialization/deserialization ('schema' was being converted to 'schema_')
+            ).model_dump(mode="json"),
             raise_for_status=True,
         )
 
@@ -127,40 +142,55 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         self,
         json_schema: dict[str, Any] | Path | str,
         model: str,
-        document: Path | str | IOBase | HttpUrl | None,
-        image_resolution_dpi: int | None = None,
-        browser_canvas: Literal['A3', 'A4', 'A5'] | None = None,
-        temperature: float = 0,
-        modality: Modality = "native",
-        reasoning_effort: ChatCompletionReasoningEffort = "medium",
-        n_consensus: int = 1,
+        document: Path | str | IOBase | HttpUrl | None = None,
+        documents: list[Path | str | IOBase | HttpUrl] | None = None,
+        image_resolution_dpi: int = PydanticUndefined,  # type: ignore[assignment]
+        browser_canvas: BrowserCanvas = PydanticUndefined,  # type: ignore[assignment]
+        temperature: float = PydanticUndefined,  # type: ignore[assignment]
+        modality: Modality = PydanticUndefined,  # type: ignore[assignment]
+        reasoning_effort: ChatCompletionReasoningEffort = PydanticUndefined,  # type: ignore[assignment]
+        n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
     ) -> UiParsedChatCompletion:
         """
-        Process a document using the UiForm API.
+        Process one or more documents using the Retab API.
 
         Args:
             json_schema: JSON schema defining the expected data structure
-            document: Single document (as MIMEData) to process
             model: The AI model to use for processing
+            document: Single document to process (use either this or documents, not both)
+            documents: List of documents to process (use either this or document, not both)
+            image_resolution_dpi: Optional image resolution DPI
+            browser_canvas: Optional browser canvas size
             temperature: Model temperature setting (0-1)
             modality: Modality of the document (e.g., native)
             reasoning_effort: The effort level for the model to reason about the input data.
             n_consensus: Number of consensus extractions to perform (default: 1 which computes a single extraction and the likelihoods comes from the model logprobs)
             idempotency_key: Idempotency key for request
-            store: Whether to store the document in the UiForm database
+            store: Whether to store the document in the Retab database
         Returns:
-            DocumentAPIResponse
+            UiParsedChatCompletion: Parsed response from the API
         Raises:
-            HTTPException if the request fails
+            ValueError: If neither document nor documents is provided, or if both are provided
+            HTTPException: If the request fails
         """
 
-        assert document is not None, "Either document or messages must be provided"
-
         # Validate DocumentAPIRequest data (raises exception if invalid)
         request = self.prepare_extraction(
-            json_schema, document, image_resolution_dpi, browser_canvas, model, temperature, modality, reasoning_effort, False, n_consensus=n_consensus, store=store, idempotency_key=idempotency_key
+            json_schema=json_schema,
+            document=document,
+            documents=documents,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+            model=model,
+            temperature=temperature,
+            modality=modality,
+            reasoning_effort=reasoning_effort,
+            stream=False,
+            n_consensus=n_consensus,
+            store=store,
+            idempotency_key=idempotency_key,
         )
         response = self._client._prepared_request(request)
 
@@ -172,45 +202,66 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         self,
         json_schema: dict[str, Any] | Path | str,
         model: str,
-        document: Path | str | IOBase | HttpUrl | None,
-        image_resolution_dpi: int | None = None,
-        browser_canvas: Literal['A3', 'A4', 'A5'] | None = None,
-        temperature: float = 0,
-        modality: Modality = "native",
-        reasoning_effort: ChatCompletionReasoningEffort = "medium",
-        n_consensus: int = 1,
+        document: Path | str | IOBase | HttpUrl | None = None,
+        documents: list[Path | str | IOBase | HttpUrl] | None = None,
+        image_resolution_dpi: int = PydanticUndefined,  # type: ignore[assignment]
+        browser_canvas: BrowserCanvas = PydanticUndefined,  # type: ignore[assignment]
+        temperature: float = PydanticUndefined,  # type: ignore[assignment]
+        modality: Modality = PydanticUndefined,  # type: ignore[assignment]
+        reasoning_effort: ChatCompletionReasoningEffort = PydanticUndefined,  # type: ignore[assignment]
+        n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
     ) -> Generator[UiParsedChatCompletion, None, None]:
         """
-        Process a document using the UiForm API with streaming enabled.
+        Process one or more documents using the Retab API with streaming enabled.
 
         Args:
             json_schema: JSON schema defining the expected data structure
-            document: Single document (as MIMEData) to process
+            model: The AI model to use for processing
+            document: Single document to process (use either this or documents, not both)
+            documents: List of documents to process (use either this or document, not both)
             image_resolution_dpi: Optional image resolution DPI.
             browser_canvas: Optional browser canvas size.
-            model: The AI model to use for processing
             temperature: Model temperature setting (0-1)
             modality: Modality of the document (e.g., native)
             reasoning_effort: The effort level for the model to reason about the input data.
             n_consensus: Number of consensus extractions to perform (default: 1 which computes a single extraction and the likelihoods comes from the model logprobs)
             idempotency_key: Idempotency key for request
-            store: Whether to store the document in the UiForm database
+            store: Whether to store the document in the Retab database
 
         Returns:
-            Generator[DocumentExtractResponse]: Stream of parsed responses
+            Generator[UiParsedChatCompletion]: Stream of parsed responses
         Raises:
-            HTTPException if the request fails
+            ValueError: If neither document nor documents is provided, or if both are provided
+            HTTPException: If the request fails
         Usage:
         ```python
-        with uiform.documents.extractions.stream(json_schema, document, model, temperature, reasoning_effort, modality) as stream:
+        # Single document
+        with retab.documents.extractions.stream(json_schema, model, document=document) as stream:
+            for response in stream:
+                print(response)
+
+        # Multiple documents
+        with retab.documents.extractions.stream(json_schema, model, documents=[doc1, doc2]) as stream:
             for response in stream:
                 print(response)
         ```
         """
         request = self.prepare_extraction(
-            json_schema, document, image_resolution_dpi, browser_canvas, model, temperature, modality, reasoning_effort, True, n_consensus=n_consensus, store=store, idempotency_key=idempotency_key
+            json_schema=json_schema,
+            document=document,
+            documents=documents,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+            model=model,
+            temperature=temperature,
+            modality=modality,
+            reasoning_effort=reasoning_effort,
+            stream=True,
+            n_consensus=n_consensus,
+            store=store,
+            idempotency_key=idempotency_key,
         )
         schema = Schema(json_schema=load_json_schema(json_schema))
 
@@ -260,8 +311,8 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         model: str,
         temperature: float,
         completion: Any | None = None,
-        # The messages can be provided in different formats, we will convert them to the UiForm-compatible format
-        messages: list[ChatCompletionUiformMessage] | None = None,
+        # The messages can be provided in different formats, we will convert them to the Retab-compatible format
+        messages: list[ChatCompletionRetabMessage] | None = None,
         openai_messages: list[ChatCompletionMessageParam] | None = None,
         anthropic_messages: list[MessageParam] | None = None,
         anthropic_system_prompt: str | None = None,
@@ -270,10 +321,10 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         openai_responses_output: Response | None = None,
     ) -> None:
         request = self.prepare_log_extraction(
-            document,
-            json_schema,
-            model,
-            temperature,
+            document=document,
+            json_schema=json_schema,
+            model=model,
+            temperature=temperature,
             completion=completion,
             messages=messages,
             openai_messages=openai_messages,
@@ -292,36 +343,52 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         self,
         json_schema: dict[str, Any] | Path | str,
         model: str,
-        document: Path | str | IOBase | HttpUrl | None,
-        image_resolution_dpi: int | None = None,
-        browser_canvas: Literal['A3', 'A4', 'A5'] | None = None,
-        temperature: float = 0,
-        modality: Modality = "native",
-        reasoning_effort: ChatCompletionReasoningEffort = "medium",
-        n_consensus: int = 1,
+        document: Path | str | IOBase | HttpUrl | None = None,
+        documents: list[Path | str | IOBase | HttpUrl] | None = None,
+        image_resolution_dpi: int = PydanticUndefined,  # type: ignore[assignment]
+        browser_canvas: BrowserCanvas = PydanticUndefined,  # type: ignore[assignment]
+        temperature: float = PydanticUndefined,  # type: ignore[assignment]
+        modality: Modality = PydanticUndefined,  # type: ignore[assignment]
+        reasoning_effort: ChatCompletionReasoningEffort = PydanticUndefined,  # type: ignore[assignment]
+        n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
     ) -> UiParsedChatCompletion:
         """
-        Extract structured data from a document asynchronously.
+        Extract structured data from one or more documents asynchronously.
 
         Args:
             json_schema: JSON schema defining the expected data structure.
-            document: Path, string, or file-like object representing the document.
+            model: The AI model to use.
+            document: Single document to process (use either this or documents, not both)
+            documents: List of documents to process (use either this or document, not both)
             image_resolution_dpi: Optional image resolution DPI.
             browser_canvas: Optional browser canvas size.
-            model: The AI model to use.
             temperature: Model temperature setting (0-1).
             modality: Modality of the document (e.g., native).
             reasoning_effort: The effort level for the model to reason about the input data.
             n_consensus: Number of consensus extractions to perform (default: 1 which computes a single extraction and the likelihoods comes from the model logprobs)
             idempotency_key: Idempotency key for request
-            store: Whether to store the document in the UiForm database
+            store: Whether to store the document in the Retab database
         Returns:
-            DocumentExtractResponse: Parsed response from the API.
+            UiParsedChatCompletion: Parsed response from the API.
+        Raises:
+            ValueError: If neither document nor documents is provided, or if both are provided
         """
         request = self.prepare_extraction(
-            json_schema, document, image_resolution_dpi, browser_canvas, model, temperature, modality, reasoning_effort, False, n_consensus=n_consensus, store=store, idempotency_key=idempotency_key
+            json_schema=json_schema,
+            document=document,
+            documents=documents,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+            model=model,
+            temperature=temperature,
+            modality=modality,
+            reasoning_effort=reasoning_effort,
+            stream=False,
+            n_consensus=n_consensus,
+            store=store,
+            idempotency_key=idempotency_key,
         )
         response = await self._client._prepared_request(request)
         schema = Schema(json_schema=load_json_schema(json_schema))
@@ -332,41 +399,65 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         self,
         json_schema: dict[str, Any] | Path | str,
         model: str,
-        document: Path | str | IOBase | HttpUrl | None,
-        image_resolution_dpi: int | None = None,
-        browser_canvas: Literal['A3', 'A4', 'A5'] | None = None,
-        temperature: float = 0,
-        modality: Modality = "native",
-        reasoning_effort: ChatCompletionReasoningEffort = "medium",
-        n_consensus: int = 1,
+        document: Path | str | IOBase | HttpUrl | None = None,
+        documents: list[Path | str | IOBase | HttpUrl] | None = None,
+        image_resolution_dpi: int = PydanticUndefined,  # type: ignore[assignment]
+        browser_canvas: BrowserCanvas = PydanticUndefined,  # type: ignore[assignment]
+        temperature: float = PydanticUndefined,  # type: ignore[assignment]
+        modality: Modality = PydanticUndefined,  # type: ignore[assignment]
+        reasoning_effort: ChatCompletionReasoningEffort = PydanticUndefined,  # type: ignore[assignment]
+        n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
     ) -> AsyncGenerator[UiParsedChatCompletion, None]:
         """
-        Extract structured data from a document asynchronously with streaming.
+        Extract structured data from one or more documents asynchronously with streaming.
 
         Args:
             json_schema: JSON schema defining the expected data structure.
-            document: Path, string, or file-like object representing the document.
             model: The AI model to use.
+            document: Single document to process (use either this or documents, not both)
+            documents: List of documents to process (use either this or document, not both)
+            image_resolution_dpi: Optional image resolution DPI.
+            browser_canvas: Optional browser canvas size.
             temperature: Model temperature setting (0-1).
             modality: Modality of the document (e.g., native).
             reasoning_effort: The effort level for the model to reason about the input data.
             n_consensus: Number of consensus extractions to perform (default: 1 which computes a single extraction and the likelihoods comes from the model logprobs)
             idempotency_key: Idempotency key for request
-            store: Whether to store the document in the UiForm database
+            store: Whether to store the document in the Retab database
         Returns:
-            AsyncGenerator[DocumentExtractResponse, None]: Stream of parsed responses.
+            AsyncGenerator[UiParsedChatCompletion, None]: Stream of parsed responses.
+        Raises:
+            ValueError: If neither document nor documents is provided, or if both are provided
 
         Usage:
         ```python
-        async with uiform.documents.extractions.stream(json_schema, document, model, temperature, reasoning_effort, modality) as stream:
+        # Single document
+        async with retab.documents.extractions.stream(json_schema, model, document=document) as stream:
+            async for response in stream:
+                print(response)
+
+        # Multiple documents
+        async with retab.documents.extractions.stream(json_schema, model, documents=[doc1, doc2]) as stream:
             async for response in stream:
                 print(response)
         ```
         """
         request = self.prepare_extraction(
-            json_schema, document, image_resolution_dpi, browser_canvas, model, temperature, modality, reasoning_effort, True, n_consensus=n_consensus, store=store, idempotency_key=idempotency_key
+            json_schema=json_schema,
+            document=document,
+            documents=documents,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+            model=model,
+            temperature=temperature,
+            modality=modality,
+            reasoning_effort=reasoning_effort,
+            stream=True,
+            n_consensus=n_consensus,
+            store=store,
+            idempotency_key=idempotency_key,
         )
         schema = Schema(json_schema=load_json_schema(json_schema))
         ui_parsed_chat_completion_cum_chunk: UiParsedChatCompletionChunk | None = None
@@ -416,8 +507,8 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         model: str,
         temperature: float,
         completion: Any | None = None,
-        # The messages can be provided in different formats, we will convert them to the UiForm-compatible format
-        messages: list[ChatCompletionUiformMessage] | None = None,
+        # The messages can be provided in different formats, we will convert them to the Retab-compatible format
+        messages: list[ChatCompletionRetabMessage] | None = None,
         openai_messages: list[ChatCompletionMessageParam] | None = None,
         anthropic_messages: list[MessageParam] | None = None,
         anthropic_system_prompt: str | None = None,
@@ -426,10 +517,10 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         openai_responses_output: Response | None = None,
     ) -> None:
         request = self.prepare_log_extraction(
-            document,
-            json_schema,
-            model,
-            temperature,
+            document=document,
+            json_schema=json_schema,
+            model=model,
+            temperature=temperature,
             completion=completion,
             messages=messages,
             openai_messages=openai_messages,