retab 0.0.37__py3-none-any.whl → 0.0.39__py3-none-any.whl

retab/resources/documents/client.py

@@ -1,22 +1,39 @@
 from io import IOBase
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import PIL.Image
 from pydantic import HttpUrl
 from pydantic_core import PydanticUndefined
+from openai.types.chat.chat_completion_reasoning_effort import ChatCompletionReasoningEffort
 
 from ..._resource import AsyncAPIResource, SyncAPIResource
-from ..._utils.json_schema import load_json_schema
+from ..._utils.json_schema import load_json_schema, filter_auxiliary_fields_json
 from ..._utils.mime import convert_mime_data_to_pil_image, prepare_mime_document
+from ..._utils.ai_models import assert_valid_model_extraction
 from ...types.documents.create_messages import DocumentCreateInputRequest, DocumentCreateMessageRequest, DocumentMessage
+from ...types.documents.extractions import DocumentExtractRequest, RetabParsedChatCompletion
+from ...types.documents.parse import ParseRequest, ParseResult, TableParsingFormat
 from ...types.browser_canvas import BrowserCanvas
 from ...types.mime import MIMEData
 from ...types.modalities import Modality
+from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
 from .extractions import AsyncExtractions, Extractions
 
 
+def maybe_parse_to_pydantic(schema: Schema, response: RetabParsedChatCompletion, allow_partial: bool = False) -> RetabParsedChatCompletion:
+    if response.choices[0].message.content:
+        try:
+            if allow_partial:
+                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:
+            pass
+    return response
+
+
 class BaseDocumentsMixin:
     def _prepare_create_messages(
         self,
@@ -69,17 +86,37 @@ class BaseDocumentsMixin:
             data={"document": mime_document.model_dump()},
         )
 
+    def _prepare_parse(
+        self,
+        document: Path | str | IOBase | MIMEData | PIL.Image.Image | HttpUrl,
+        fast_mode: bool = False,
+        table_parsing_format: TableParsingFormat = "html",
+        image_resolution_dpi: int = 72,
+        browser_canvas: BrowserCanvas = "A4",
+        idempotency_key: str | None = None,
+    ) -> PreparedRequest:
+        mime_document = prepare_mime_document(document)
+
+        parse_request = ParseRequest(
+            document=mime_document,
+            fast_mode=fast_mode,
+            table_parsing_format=table_parsing_format,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+        )
+        return PreparedRequest(method="POST", url="/v1/documents/parse", data=parse_request.model_dump(), idempotency_key=idempotency_key)
+
 
 class Documents(SyncAPIResource, BaseDocumentsMixin):
     """Documents API wrapper"""
 
     def __init__(self, client: Any) -> None:
         super().__init__(client=client)
-        self.extractions = Extractions(client=client)
+        # self.extractions_api = Extractions(client=client)
         # self.batch = Batch(client=client)
 
     def correct_image_orientation(self, document: Path | str | IOBase | MIMEData | PIL.Image.Image) -> PIL.Image.Image:
-        """Corrects the orientation of an image using the UiForm API.
+        """Corrects the orientation of an image using the Retab API.
 
         This method takes an image in various formats and returns a PIL Image with corrected orientation.
         Useful for handling images from mobile devices or cameras that may have incorrect EXIF orientation.
@@ -96,7 +133,7 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
 
         Raises:
             ValueError: If the input is not a valid image
-            UiformAPIError: If the API request fails
+            RetabAPIError: If the API request fails
         """
         request = self._prepare_correct_image_orientation(document)
         response = self._client._prepared_request(request)
@@ -112,7 +149,7 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
         idempotency_key: str | None = None,
     ) -> DocumentMessage:
         """
-        Create document messages from a file using the UiForm API.
+        Create document messages from a file using the Retab API.
 
         Args:
             document: The document to process. Can be a file path (Path or str) or a file-like object.
@@ -124,7 +161,7 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
             DocumentMessage: The processed document message containing extracted content.
 
         Raises:
-            UiformAPIError: If the API request fails.
+            RetabAPIError: If the API request fails.
         """
         request = self._prepare_create_messages(
             document=document, modality=modality, image_resolution_dpi=image_resolution_dpi, browser_canvas=browser_canvas, idempotency_key=idempotency_key
@@ -142,7 +179,7 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
         idempotency_key: str | None = None,
     ) -> DocumentMessage:
         """
-        Create document inputs (messages with schema) from a file using the UiForm API.
+        Create document inputs (messages with schema) from a file using the Retab API.
 
         Args:
             document: The document to process. Can be a file path (Path or str), file-like object, MIMEData, PIL Image, or URL.
@@ -155,7 +192,7 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
             DocumentMessage: The processed document message containing extracted content with schema context.
 
         Raises:
-            UiformAPIError: If the API request fails.
+            RetabAPIError: If the API request fails.
         """
         request = self._prepare_create_inputs(
             document=document,
@@ -168,13 +205,135 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
         response = self._client._prepared_request(request)
         return DocumentMessage.model_validate(response)
 
+    def extract(
+        self,
+        json_schema: dict[str, Any] | Path | str,
+        model: str,
+        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,
+    ) -> RetabParsedChatCompletion:
+        """
+        Process one or more documents using the Retab API for structured data extraction.
+
+        This method provides a direct interface to document extraction functionality,
+        intended to replace the current `.extractions.parse()` pattern.
+
+        Args:
+            json_schema: JSON schema defining the expected data structure
+            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
+            idempotency_key: Idempotency key for request
+            store: Whether to store the document in the Retab database
+
+        Returns:
+            RetabParsedChatCompletion: Parsed response from the API
+
+        Raises:
+            ValueError: If neither document nor documents is provided, or if both are provided
+            HTTPException: If the request fails
+        """
+        assert_valid_model_extraction(model)
+
+        json_schema = load_json_schema(json_schema)
+
+        # 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)
+        request = DocumentExtractRequest(
+            json_schema=json_schema,
+            documents=processed_documents,
+            model=model,
+            temperature=temperature,
+            stream=False,
+            modality=modality,
+            store=store,
+            reasoning_effort=reasoning_effort,
+            n_consensus=n_consensus,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+        )
+
+        prepared_request = PreparedRequest(
+            method="POST", url="/v1/documents/extract", data=request.model_dump(mode="json", exclude_unset=True, exclude_defaults=True), idempotency_key=idempotency_key
+        )
+
+        response = self._client._prepared_request(prepared_request)
+
+        schema = Schema(json_schema=load_json_schema(json_schema))
+        return maybe_parse_to_pydantic(schema, RetabParsedChatCompletion.model_validate(response))
+
+    def parse(
+        self,
+        document: Path | str | IOBase | MIMEData | PIL.Image.Image | HttpUrl,
+        fast_mode: bool = False,
+        table_parsing_format: TableParsingFormat = "html",
+        image_resolution_dpi: int = 72,
+        browser_canvas: BrowserCanvas = "A4",
+        idempotency_key: str | None = None,
+    ) -> ParseResult:
+        """
+        Parse a document and extract text content from each page.
+
+        This method processes various document types and returns structured text content
+        along with usage information. Supports different parsing modes and formats.
+
+        Args:
+            document: The document to parse. Can be a file path (Path or str), file-like object, MIMEData, PIL Image, or URL.
+            fast_mode: Use fast mode for parsing (may reduce quality). Defaults to False.
+            table_parsing_format: Format for parsing tables. Options: "html", "json", "yaml", "markdown". Defaults to "html".
+            image_resolution_dpi: DPI for image processing. Defaults to 72.
+            browser_canvas: Canvas size for document rendering. Defaults to "A4".
+            idempotency_key: Optional idempotency key for the request.
+
+        Returns:
+            ParseResult: Parsed response containing document metadata, usage information, and page text content.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        request = self._prepare_parse(
+            document=document,
+            fast_mode=fast_mode,
+            table_parsing_format=table_parsing_format,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+            idempotency_key=idempotency_key,
+        )
+        response = self._client._prepared_request(request)
+        return ParseResult.model_validate(response)
+
 
 class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
     """Documents API wrapper for asynchronous usage."""
 
     def __init__(self, client: Any) -> None:
         super().__init__(client=client)
-        self.extractions = AsyncExtractions(client=client)
+        # self.extractions_api = AsyncExtractions(client=client)
 
     async def create_messages(
         self,
@@ -185,7 +344,7 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
         idempotency_key: str | None = None,
     ) -> DocumentMessage:
         """
-        Create document messages from a file using the UiForm API asynchronously.
+        Create document messages from a file using the Retab API asynchronously.
 
         Args:
             document: The document to process. Can be a file path (Path or str) or a file-like object.
@@ -195,7 +354,7 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
             DocumentMessage: The processed document message containing extracted content.
 
         Raises:
-            UiformAPIError: If the API request fails.
+            RetabAPIError: If the API request fails.
         """
         request = self._prepare_create_messages(
             document=document,
@@ -217,7 +376,7 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
         idempotency_key: str | None = None,
     ) -> DocumentMessage:
         """
-        Create document inputs (messages with schema) from a file using the UiForm API asynchronously.
+        Create document inputs (messages with schema) from a file using the Retab API asynchronously.
 
         Args:
             document: The document to process. Can be a file path (Path or str), file-like object, MIMEData, PIL Image, or URL.
@@ -230,7 +389,7 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
             DocumentMessage: The processed document message containing extracted content with schema context.
 
         Raises:
-            UiformAPIError: If the API request fails.
+            RetabAPIError: If the API request fails.
         """
         request = self._prepare_create_inputs(
             document=document,
@@ -244,7 +403,7 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
         return DocumentMessage.model_validate(response)
 
     async def correct_image_orientation(self, document: Path | str | IOBase | MIMEData | PIL.Image.Image) -> PIL.Image.Image:
-        """Corrects the orientation of an image using the UiForm API asynchronously.
+        """Corrects the orientation of an image using the Retab API asynchronously.
 
         This method takes an image in various formats and returns a PIL Image with corrected orientation.
         Useful for handling images from mobile devices or cameras that may have incorrect EXIF orientation.
@@ -261,9 +420,131 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
 
         Raises:
             ValueError: If the input is not a valid image
-            UiformAPIError: If the API request fails
+            RetabAPIError: If the API request fails
         """
         request = self._prepare_correct_image_orientation(document)
         response = await self._client._prepared_request(request)
         mime_response = MIMEData.model_validate(response["document"])
         return convert_mime_data_to_pil_image(mime_response)
+
+    async def extract(
+        self,
+        json_schema: dict[str, Any] | Path | str,
+        model: str,
+        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,
+    ) -> RetabParsedChatCompletion:
+        """
+        Process one or more documents using the Retab API for structured data extraction asynchronously.
+
+        This method provides a direct interface to document extraction functionality,
+        intended to replace the current `.extractions.parse()` pattern.
+
+        Args:
+            json_schema: JSON schema defining the expected data structure
+            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
+            idempotency_key: Idempotency key for request
+            store: Whether to store the document in the Retab database
+
+        Returns:
+            RetabParsedChatCompletion: Parsed response from the API
+
+        Raises:
+            ValueError: If neither document nor documents is provided, or if both are provided
+            HTTPException: If the request fails
+        """
+        assert_valid_model_extraction(model)
+
+        json_schema = load_json_schema(json_schema)
+
+        # 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)
+        request = DocumentExtractRequest(
+            json_schema=json_schema,
+            documents=processed_documents,
+            model=model,
+            temperature=temperature,
+            stream=False,
+            modality=modality,
+            store=store,
+            reasoning_effort=reasoning_effort,
+            n_consensus=n_consensus,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+        )
+
+        prepared_request = PreparedRequest(
+            method="POST", url="/v1/documents/extract", data=request.model_dump(mode="json", exclude_unset=True, exclude_defaults=True), idempotency_key=idempotency_key
+        )
+
+        response = await self._client._prepared_request(prepared_request)
+
+        schema = Schema(json_schema=load_json_schema(json_schema))
+        return maybe_parse_to_pydantic(schema, RetabParsedChatCompletion.model_validate(response))
+
+    async def parse(
+        self,
+        document: Path | str | IOBase | MIMEData | PIL.Image.Image | HttpUrl,
+        fast_mode: bool = False,
+        table_parsing_format: TableParsingFormat = "html",
+        image_resolution_dpi: int = 72,
+        browser_canvas: BrowserCanvas = "A4",
+        idempotency_key: str | None = None,
+    ) -> ParseResult:
+        """
+        Parse a document and extract text content from each page asynchronously.
+
+        This method processes various document types and returns structured text content
+        along with usage information. Supports different parsing modes and formats.
+
+        Args:
+            document: The document to parse. Can be a file path (Path or str), file-like object, MIMEData, PIL Image, or URL.
+            fast_mode: Use fast mode for parsing (may reduce quality). Defaults to False.
+            table_parsing_format: Format for parsing tables. Options: "html", "json", "yaml", "markdown". Defaults to "html".
+            image_resolution_dpi: DPI for image processing. Defaults to 72.
+            browser_canvas: Canvas size for document rendering. Defaults to "A4".
+            idempotency_key: Optional idempotency key for the request.
+
+        Returns:
+            ParseResult: Parsed response containing document metadata, usage information, and page text content.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        request = self._prepare_parse(
+            document=document,
+            fast_mode=fast_mode,
+            table_parsing_format=table_parsing_format,
+            image_resolution_dpi=image_resolution_dpi,
+            browser_canvas=browser_canvas,
+            idempotency_key=idempotency_key,
+        )
+        response = await self._client._prepared_request(request)
+        return ParseResult.model_validate(response)