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

retab/resources/documents/client.py

@@ -1,6 +1,6 @@
 from io import IOBase
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import PIL.Image
 from pydantic import HttpUrl
@@ -12,7 +12,8 @@ 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, UiParsedChatCompletion
+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
@@ -21,7 +22,7 @@ from ...types.standards import PreparedRequest
 from .extractions import AsyncExtractions, Extractions
 
 
-def maybe_parse_to_pydantic(schema: Schema, response: UiParsedChatCompletion, allow_partial: bool = False) -> UiParsedChatCompletion:
+def maybe_parse_to_pydantic(schema: Schema, response: RetabParsedChatCompletion, allow_partial: bool = False) -> RetabParsedChatCompletion:
     if response.choices[0].message.content:
         try:
             if allow_partial:
@@ -85,13 +86,33 @@ 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_api = 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:
@@ -198,10 +219,10 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
         n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
-    ) -> UiParsedChatCompletion:
+    ) -> 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.
 
@@ -218,10 +239,10 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
             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:
-            UiParsedChatCompletion: Parsed response from the API
-            
+            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
@@ -258,16 +279,53 @@ class Documents(SyncAPIResource, BaseDocumentsMixin):
         )
 
         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
+            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, UiParsedChatCompletion.model_validate(response))
+        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):
@@ -275,7 +333,7 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
 
     def __init__(self, client: Any) -> None:
         super().__init__(client=client)
-        #self.extractions_api = AsyncExtractions(client=client)
+        # self.extractions_api = AsyncExtractions(client=client)
 
     async def create_messages(
         self,
@@ -383,10 +441,10 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
         n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
-    ) -> UiParsedChatCompletion:
+    ) -> 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.
 
@@ -403,10 +461,10 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
             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:
-            UiParsedChatCompletion: Parsed response from the API
-            
+            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
@@ -443,13 +501,50 @@ class AsyncDocuments(AsyncAPIResource, BaseDocumentsMixin):
         )
 
         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
+            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, UiParsedChatCompletion.model_validate(response))
+        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)

retab/resources/documents/extractions.py

@@ -19,14 +19,14 @@ from ..._utils.json_schema import filter_auxiliary_fields_json, load_json_schema
 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 ChatCompletionRetabMessage
-from ...types.documents.extractions import DocumentExtractRequest, LogExtractionRequest, UiParsedChatCompletion, UiParsedChatCompletionChunk, UiParsedChoice
+from ...types.documents.extractions import DocumentExtractRequest, LogExtractionRequest, RetabParsedChatCompletion, RetabParsedChatCompletionChunk, RetabParsedChoice
 from ...types.browser_canvas import BrowserCanvas
 from ...types.modalities import Modality
 from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
 
 
-def maybe_parse_to_pydantic(schema: Schema, response: UiParsedChatCompletion, allow_partial: bool = False) -> UiParsedChatCompletion:
+def maybe_parse_to_pydantic(schema: Schema, response: RetabParsedChatCompletion, allow_partial: bool = False) -> RetabParsedChatCompletion:
     if response.choices[0].message.content:
         try:
             if allow_partial:
@@ -152,7 +152,7 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """
         Process one or more documents using the Retab API.
 
@@ -170,7 +170,7 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
             idempotency_key: Idempotency key for request
             store: Whether to store the document in the Retab database
         Returns:
-            UiParsedChatCompletion: Parsed response from the API
+            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
@@ -195,7 +195,7 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         response = self._client._prepared_request(request)
 
         schema = Schema(json_schema=load_json_schema(json_schema))
-        return maybe_parse_to_pydantic(schema, UiParsedChatCompletion.model_validate(response))
+        return maybe_parse_to_pydantic(schema, RetabParsedChatCompletion.model_validate(response))
 
     @as_context_manager
     def stream(
@@ -212,7 +212,7 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
-    ) -> Generator[UiParsedChatCompletion, None, None]:
+    ) -> Generator[RetabParsedChatCompletion, None, None]:
         """
         Process one or more documents using the Retab API with streaming enabled.
 
@@ -231,7 +231,7 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
             store: Whether to store the document in the Retab database
 
         Returns:
-            Generator[UiParsedChatCompletion]: Stream of parsed responses
+            Generator[RetabParsedChatCompletion]: Stream of parsed responses
         Raises:
             ValueError: If neither document nor documents is provided, or if both are provided
             HTTPException: If the request fails
@@ -266,16 +266,16 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         schema = Schema(json_schema=load_json_schema(json_schema))
 
         # Request the stream and return a context manager
-        ui_parsed_chat_completion_cum_chunk: UiParsedChatCompletionChunk | None = None
-        # Initialize the UiParsedChatCompletion object
-        ui_parsed_completion: UiParsedChatCompletion = UiParsedChatCompletion(
+        ui_parsed_chat_completion_cum_chunk: RetabParsedChatCompletionChunk | None = None
+        # Initialize the RetabParsedChatCompletion object
+        ui_parsed_completion: RetabParsedChatCompletion = RetabParsedChatCompletion(
             id="",
             created=0,
             model="",
             object="chat.completion",
             likelihoods={},
             choices=[
-                UiParsedChoice(
+                RetabParsedChoice(
                     index=0,
                     message=ParsedChatCompletionMessage(content="", role="assistant"),
                     finish_reason=None,
@@ -286,7 +286,7 @@ class Extractions(SyncAPIResource, BaseExtractionsMixin):
         for chunk_json in self._client._prepared_request_stream(request):
             if not chunk_json:
                 continue
-            ui_parsed_chat_completion_cum_chunk = UiParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
+            ui_parsed_chat_completion_cum_chunk = RetabParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
             # Basic stuff
             ui_parsed_completion.id = ui_parsed_chat_completion_cum_chunk.id
             ui_parsed_completion.created = ui_parsed_chat_completion_cum_chunk.created
@@ -353,7 +353,7 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """
         Extract structured data from one or more documents asynchronously.
 
@@ -371,7 +371,7 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
             idempotency_key: Idempotency key for request
             store: Whether to store the document in the Retab database
         Returns:
-            UiParsedChatCompletion: Parsed response from the API.
+            RetabParsedChatCompletion: Parsed response from the API.
         Raises:
             ValueError: If neither document nor documents is provided, or if both are provided
         """
@@ -392,7 +392,7 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         )
         response = await self._client._prepared_request(request)
         schema = Schema(json_schema=load_json_schema(json_schema))
-        return maybe_parse_to_pydantic(schema, UiParsedChatCompletion.model_validate(response))
+        return maybe_parse_to_pydantic(schema, RetabParsedChatCompletion.model_validate(response))
 
     @as_async_context_manager
     async def stream(
@@ -409,7 +409,7 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         n_consensus: int = PydanticUndefined,  # type: ignore[assignment]
         idempotency_key: str | None = None,
         store: bool = False,
-    ) -> AsyncGenerator[UiParsedChatCompletion, None]:
+    ) -> AsyncGenerator[RetabParsedChatCompletion, None]:
         """
         Extract structured data from one or more documents asynchronously with streaming.
 
@@ -427,7 +427,7 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
             idempotency_key: Idempotency key for request
             store: Whether to store the document in the Retab database
         Returns:
-            AsyncGenerator[UiParsedChatCompletion, None]: Stream of parsed responses.
+            AsyncGenerator[RetabParsedChatCompletion, None]: Stream of parsed responses.
         Raises:
             ValueError: If neither document nor documents is provided, or if both are provided
 
@@ -460,16 +460,16 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
             idempotency_key=idempotency_key,
         )
         schema = Schema(json_schema=load_json_schema(json_schema))
-        ui_parsed_chat_completion_cum_chunk: UiParsedChatCompletionChunk | None = None
-        # Initialize the UiParsedChatCompletion object
-        ui_parsed_completion: UiParsedChatCompletion = UiParsedChatCompletion(
+        ui_parsed_chat_completion_cum_chunk: RetabParsedChatCompletionChunk | None = None
+        # Initialize the RetabParsedChatCompletion object
+        ui_parsed_completion: RetabParsedChatCompletion = RetabParsedChatCompletion(
             id="",
             created=0,
             model="",
             object="chat.completion",
             likelihoods={},
             choices=[
-                UiParsedChoice(
+                RetabParsedChoice(
                     index=0,
                     message=ParsedChatCompletionMessage(content="", role="assistant"),
                     finish_reason=None,
@@ -481,7 +481,7 @@ class AsyncExtractions(AsyncAPIResource, BaseExtractionsMixin):
         async for chunk_json in self._client._prepared_request_stream(request):
             if not chunk_json:
                 continue
-            ui_parsed_chat_completion_cum_chunk = UiParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
+            ui_parsed_chat_completion_cum_chunk = RetabParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
             # Basic stuff
             ui_parsed_completion.id = ui_parsed_chat_completion_cum_chunk.id
             ui_parsed_completion.created = ui_parsed_chat_completion_cum_chunk.created

retab/resources/evaluations/documents.py

@@ -10,7 +10,7 @@ from ..._utils.mime import prepare_mime_document
 from ...types.evaluations import DocumentItem, EvaluationDocument, PatchEvaluationDocumentRequest
 from ...types.mime import MIMEData
 from ...types.standards import PreparedRequest, DeleteResponse, FieldUnset
-from ...types.documents.extractions import UiParsedChatCompletion
+from ...types.documents.extractions import RetabParsedChatCompletion
 
 
 class DocumentsMixin:
@@ -134,13 +134,13 @@ class Documents(SyncAPIResource, DocumentsMixin):
         request = self.prepare_delete(evaluation_id, document_id)
         return self._client._prepared_request(request)
 
-    def llm_annotate(self, evaluation_id: str, document_id: str) -> UiParsedChatCompletion:
+    def llm_annotate(self, evaluation_id: str, document_id: str) -> RetabParsedChatCompletion:
         """
         Annotate a document with an LLM. This method updates the document (within the evaluation) with the latest extraction.
         """
         request = self.prepare_llm_annotate(evaluation_id, document_id)
         response = self._client._prepared_request(request)
-        return UiParsedChatCompletion(**response)
+        return RetabParsedChatCompletion(**response)
 
 
 class AsyncDocuments(AsyncAPIResource, DocumentsMixin):
@@ -223,11 +223,11 @@ class AsyncDocuments(AsyncAPIResource, DocumentsMixin):
         request = self.prepare_delete(evaluation_id, document_id)
         return await self._client._prepared_request(request)
 
-    async def llm_annotate(self, evaluation_id: str, document_id: str) -> UiParsedChatCompletion:
+    async def llm_annotate(self, evaluation_id: str, document_id: str) -> RetabParsedChatCompletion:
         """
         Annotate a document with an LLM.
         This method updates the document (within the evaluation) with the latest extraction.
         """
         request = self.prepare_llm_annotate(evaluation_id, document_id)
         response = await self._client._prepared_request(request)
-        return UiParsedChatCompletion(**response)
+        return RetabParsedChatCompletion(**response)

retab/resources/evaluations/iterations.py

@@ -9,7 +9,7 @@ from ...types.inference_settings import InferenceSettings
 from ...types.metrics import DistancesResult
 from ...types.modalities import Modality
 from ...types.standards import DeleteResponse, PreparedRequest, FieldUnset
-from ...types.documents.extractions import UiParsedChatCompletion
+from ...types.documents.extractions import RetabParsedChatCompletion
 
 
 class IterationsMixin:
@@ -238,7 +238,7 @@ class Iterations(SyncAPIResource, IterationsMixin):
         response = self._client._prepared_request(request)
         return Iteration(**response)
 
-    def process_document(self, evaluation_id: str, iteration_id: str, document_id: str) -> UiParsedChatCompletion:
+    def process_document(self, evaluation_id: str, iteration_id: str, document_id: str) -> RetabParsedChatCompletion:
         """
         Process a single document within an iteration.
         This method updates the iteration document with the latest extraction.
@@ -248,13 +248,13 @@ class Iterations(SyncAPIResource, IterationsMixin):
             document_id: The ID of the document
 
         Returns:
-            UiParsedChatCompletion: The parsed chat completion
+            RetabParsedChatCompletion: The parsed chat completion
         Raises:
             HTTPException if the request fails
         """
         request = self.prepare_process_document(evaluation_id, iteration_id, document_id)
         response = self._client._prepared_request(request)
-        return UiParsedChatCompletion(**response)
+        return RetabParsedChatCompletion(**response)
 
     def status(self, evaluation_id: str, iteration_id: str) -> IterationDocumentStatusResponse:
         """
@@ -417,7 +417,7 @@ class AsyncIterations(AsyncAPIResource, IterationsMixin):
         response = await self._client._prepared_request(request)
         return Iteration(**response)
 
-    async def process_document(self, evaluation_id: str, iteration_id: str, document_id: str) -> UiParsedChatCompletion:
+    async def process_document(self, evaluation_id: str, iteration_id: str, document_id: str) -> RetabParsedChatCompletion:
         """
         Process a single document within an iteration.
         This method updates the iteration document with the latest extraction.
@@ -427,13 +427,13 @@ class AsyncIterations(AsyncAPIResource, IterationsMixin):
             document_id: The ID of the document
 
         Returns:
-            UiParsedChatCompletion: The parsed chat completion
+            RetabParsedChatCompletion: The parsed chat completion
         Raises:
             HTTPException if the request fails
         """
         request = self.prepare_process_document(evaluation_id, iteration_id, document_id)
         response = await self._client._prepared_request(request)
-        return UiParsedChatCompletion(**response)
+        return RetabParsedChatCompletion(**response)
 
     async def status(self, evaluation_id: str, iteration_id: str) -> IterationDocumentStatusResponse:
         """

retab/resources/jsonlUtils.py

@@ -18,7 +18,7 @@ from pydantic_core import PydanticUndefined
 from tqdm import tqdm
 
 from .._resource import AsyncAPIResource, SyncAPIResource
-from .._utils.ai_models import assert_valid_model_extraction, find_provider_from_model
+from .._utils.ai_models import assert_valid_model_extraction, get_provider_for_model
 from .._utils.chat import convert_to_anthropic_format, convert_to_openai_format, separate_messages
 from .._utils.display import Metrics, display_metrics, process_dataset_and_compute_metrics
 from .._utils.json_schema import load_json_schema
@@ -278,7 +278,7 @@ class Datasets(SyncAPIResource, BaseDatasetsMixin):
         Returns:
             A tuple of (client instance, provider type string)
         """
-        provider = find_provider_from_model(model)
+        provider = get_provider_for_model(model)
 
         if provider == "OpenAI":
             return OpenAI(api_key=self._client.headers["OpenAI-Api-Key"]), provider

retab/resources/processors/client.py

@@ -12,7 +12,7 @@ from ..._resource import AsyncAPIResource, SyncAPIResource
 from ..._utils.ai_models import assert_valid_model_extraction
 from ..._utils.mime import MIMEData, prepare_mime_document
 from ...types.browser_canvas import BrowserCanvas
-from ...types.documents.extractions import UiParsedChatCompletion
+from ...types.documents.extractions import RetabParsedChatCompletion
 from ...types.logs import ProcessorConfig, UpdateProcessorRequest
 from ...types.modalities import Modality
 from ...types.pagination import ListMetadata
@@ -349,7 +349,7 @@ class Processors(SyncAPIResource, ProcessorsMixin):
         temperature: float | None = None,
         seed: int | None = None,
         store: bool = True,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """Submit documents to a processor for processing.
 
         Args:
@@ -361,11 +361,11 @@ class Processors(SyncAPIResource, ProcessorsMixin):
             store: Whether to store the results
 
         Returns:
-            UiParsedChatCompletion: The processing result
+            RetabParsedChatCompletion: The processing result
         """
         request = self.prepare_submit(processor_id=processor_id, document=document, documents=documents, temperature=temperature, seed=seed, store=store)
         response = self._client._prepared_request(request)
-        return UiParsedChatCompletion.model_validate(response)
+        return RetabParsedChatCompletion.model_validate(response)
 
 
 class AsyncProcessors(AsyncAPIResource, ProcessorsMixin):
@@ -470,7 +470,7 @@ class AsyncProcessors(AsyncAPIResource, ProcessorsMixin):
         temperature: float | None = None,
         seed: int | None = None,
         store: bool = True,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """Submit documents to a processor for processing.
 
         Args:
@@ -482,8 +482,8 @@ class AsyncProcessors(AsyncAPIResource, ProcessorsMixin):
             store: Whether to store the results
 
         Returns:
-            UiParsedChatCompletion: The processing result
+            RetabParsedChatCompletion: The processing result
         """
         request = self.prepare_submit(processor_id=processor_id, document=document, documents=documents, temperature=temperature, seed=seed, store=store)
         response = await self._client._prepared_request(request)
-        return UiParsedChatCompletion.model_validate(response)
+        return RetabParsedChatCompletion.model_validate(response)