✨ add support for split

Author: sebastianMindee

mindee/client_v2.py

@@ -1,10 +1,10 @@
 from time import sleep
-from typing import Optional, Union
+from typing import Optional, Union, Type
 
 from mindee.client_mixin import ClientMixin
 from mindee.error.mindee_error import MindeeError
 from mindee.error.mindee_http_error_v2 import handle_error_v2
-from mindee.input import UrlInputSource
+from mindee.input import UrlInputSource, UtilityParameters
 from mindee.input.inference_parameters import InferenceParameters
 from mindee.input.polling_options import PollingOptions
 from mindee.input.sources.local_input_source import LocalInputSource
@@ -15,6 +15,7 @@
     is_valid_post_response,
 )
 from mindee.parsing.v2.common_response import CommonStatus
+from mindee.v2 import BaseInferenceResponse
 from mindee.parsing.v2.inference_response import InferenceResponse
 from mindee.parsing.v2.job_response import JobResponse
 
@@ -41,20 +42,21 @@ def __init__(self, api_key: Optional[str] = None) -> None:
     def enqueue_inference(
         self,
         input_source: Union[LocalInputSource, UrlInputSource],
-        params: InferenceParameters,
+        params: Union[InferenceParameters, UtilityParameters],
+        slug: Optional[str] = None,
     ) -> JobResponse:
         """
         Enqueues a document to a given model.
 
         :param input_source: The document/source file to use. Can be local or remote.
-
         :param params: Parameters to set when sending a file.
+        :param slug: Slug for the endpoint.
+
         :return: A valid inference response.
         """
         logger.debug("Enqueuing inference using model: %s", params.model_id)
-
         response = self.mindee_api.req_post_inference_enqueue(
-            input_source=input_source, params=params
+            input_source=input_source, params=params, slug=slug
         )
         dict_response = response.json()
 
@@ -79,13 +81,18 @@ def get_job(self, job_id: str) -> JobResponse:
         dict_response = response.json()
         return JobResponse(dict_response)
 
-    def get_inference(self, inference_id: str) -> InferenceResponse:
+    def get_inference(
+        self,
+        inference_id: str,
+        inference_response_type: Type[InferenceResponse] = InferenceResponse,
+    ) -> BaseInferenceResponse:
         """
         Get the result of an inference that was previously enqueued.
 
         The inference will only be available after it has finished processing.
 
         :param inference_id: UUID of the inference to retrieve.
+        :param inference_response_type: Class of the product to instantiate.
         :return: An inference response.
         """
         logger.debug("Fetching inference: %s", inference_id)
@@ -94,19 +101,20 @@ def get_inference(self, inference_id: str) -> InferenceResponse:
         if not is_valid_get_response(response):
             handle_error_v2(response.json())
         dict_response = response.json()
-        return InferenceResponse(dict_response)
+        return inference_response_type(dict_response)
 
-    def enqueue_and_get_inference(
+    def _enqueue_and_get(
         self,
         input_source: Union[LocalInputSource, UrlInputSource],
-        params: InferenceParameters,
+        params: Union[InferenceParameters, UtilityParameters],
+        inference_response_type: Optional[Type[InferenceResponse]] = InferenceResponse,
     ) -> InferenceResponse:
         """
         Enqueues to an asynchronous endpoint and automatically polls for a response.
 
         :param input_source: The document/source file to use. Can be local or remote.
-
         :param params: Parameters to set when sending a file.
+        :param inference_response_type: The product class to use for the response object.
 
         :return: A valid inference response.
         """
@@ -117,9 +125,14 @@ def enqueue_and_get_inference(
             params.polling_options.delay_sec,
             params.polling_options.max_retries,
         )
-        enqueue_response = self.enqueue_inference(input_source, params)
+        slug = (
+            inference_response_type.inference.get_slug()
+            if inference_response_type
+            else None
+        )
+        enqueue_response = self.enqueue_inference(input_source, params, slug)
         logger.debug(
-            "Successfully enqueued inference with job id: %s", enqueue_response.job.id
+            "Successfully enqueued document with job id: %s", enqueue_response.job.id
         )
         sleep(params.polling_options.initial_delay_sec)
         try_counter = 0
@@ -134,8 +147,57 @@ def enqueue_and_get_inference(
                     f"Parsing failed for job {job_response.job.id}: {detail}"
                 )
             if job_response.job.status == CommonStatus.PROCESSED.value:
-                return self.get_inference(job_response.job.id)
+                result = self.get_inference(
+                    job_response.job.id, inference_response_type or InferenceResponse
+                )
+                assert isinstance(result, InferenceResponse), (
+                    f'Invalid response type "{type(result)}"'
+                )
+                return result
             try_counter += 1
             sleep(params.polling_options.delay_sec)
 
         raise MindeeError(f"Couldn't retrieve document after {try_counter + 1} tries.")
+
+    def enqueue_and_get_inference(
+        self,
+        input_source: Union[LocalInputSource, UrlInputSource],
+        params: InferenceParameters,
+    ) -> InferenceResponse:
+        """
+        Enqueues to an asynchronous endpoint and automatically polls for a response.
+
+        :param input_source: The document/source file to use. Can be local or remote.
+
+        :param params: Parameters to set when sending a file.
+
+        :return: A valid inference response.
+        """
+        response = self._enqueue_and_get(input_source, params)
+        assert isinstance(response, InferenceResponse), (
+            f'Invalid response type "{type(response)}"'
+        )
+        return response
+
+    def enqueue_and_get_utility(
+        self,
+        inference_response_type: Type[InferenceResponse],
+        input_source: Union[LocalInputSource, UrlInputSource],
+        params: UtilityParameters,
+    ) -> InferenceResponse:
+        """
+        Enqueues to an asynchronous endpoint and automatically polls for a response.
+
+        :param input_source: The document/source file to use. Can be local or remote.
+
+        :param params: Parameters to set when sending a file.
+
+        :param inference_response_type: The product class to use for the response object.
+
+        :return: A valid inference response.
+        """
+        response = self._enqueue_and_get(input_source, params, inference_response_type)
+        assert isinstance(response, inference_response_type), (
+            f'Invalid response type "{type(response)}"'
+        )
+        return response

mindee/input/__init__.py

@@ -1,4 +1,7 @@
 from mindee.input.local_response import LocalResponse
+from mindee.input.base_parameters import BaseParameters
+from mindee.input.inference_parameters import InferenceParameters
+from mindee.input.utility_parameters import UtilityParameters
 from mindee.input.page_options import PageOptions
 from mindee.input.polling_options import PollingOptions
 from mindee.input.sources.base_64_input import Base64Input
@@ -11,15 +14,18 @@
 from mindee.input.workflow_options import WorkflowOptions
 
 __all__ = [
+    "Base64Input",
+    "BaseParameters",
+    "BytesInput",
+    "FileInput",
     "InputType",
+    "InferenceParameters",
     "LocalInputSource",
-    "UrlInputSource",
+    "LocalResponse",
+    "PageOptions",
     "PathInput",
-    "FileInput",
-    "Base64Input",
-    "BytesInput",
-    "WorkflowOptions",
     "PollingOptions",
-    "PageOptions",
-    "LocalResponse",
+    "UrlInputSource",
+    "UtilityParameters",
+    "WorkflowOptions",
 ]

mindee/input/base_parameters.py

@@ -0,0 +1,21 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Optional, List
+
+from mindee.input.polling_options import PollingOptions
+
+
+@dataclass
+class BaseParameters(ABC):
+    """Base class for parameters accepted by all V2 endpoints."""
+
+    model_id: str
+    """ID of the model, required."""
+    alias: Optional[str] = None
+    """Use an alias to link the file to your own DB. If empty, no alias will be used."""
+    webhook_ids: Optional[List[str]] = None
+    """IDs of webhooks to propagate the API response to."""
+    polling_options: Optional[PollingOptions] = None
+    """Options for polling. Set only if having timeout issues."""
+    close_file: bool = True
+    """Whether to close the file after parsing."""

mindee/input/inference_parameters.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass, asdict
 from typing import List, Optional, Union
 
-from mindee.input.polling_options import PollingOptions
+from mindee.input.base_parameters import BaseParameters
 
 
 @dataclass
@@ -44,7 +44,7 @@ class DataSchemaField(StringDataClass):
     guidelines: Optional[str] = None
     """Optional extraction guidelines."""
     nested_fields: Optional[dict] = None
-    """Subfields when type is `nested_object`. Leave empty for other types"""
+    """Subfields when type is `nested_object`. Leave empty for other types."""
 
 
 @dataclass
@@ -78,11 +78,9 @@ def __post_init__(self) -> None:
 
 
 @dataclass
-class InferenceParameters:
+class InferenceParameters(BaseParameters):
     """Inference parameters to set when sending a file."""
 
-    model_id: str
-    """ID of the model, required."""
     rag: Optional[bool] = None
     """Enhance extraction accuracy with Retrieval-Augmented Generation."""
     raw_text: Optional[bool] = None
@@ -94,14 +92,6 @@ class InferenceParameters:
     Boost the precision and accuracy of all extractions.
     Calculate confidence scores for all fields, and fill their ``confidence`` attribute.
     """
-    alias: Optional[str] = None
-    """Use an alias to link the file to your own DB. If empty, no alias will be used."""
-    webhook_ids: Optional[List[str]] = None
-    """IDs of webhooks to propagate the API response to."""
-    polling_options: Optional[PollingOptions] = None
-    """Options for polling. Set only if having timeout issues."""
-    close_file: bool = True
-    """Whether to close the file after parsing."""
     text_context: Optional[str] = None
     """
     Additional text context used by the model during inference.

mindee/input/utility_parameters.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+
+from mindee.input.base_parameters import BaseParameters
+
+
+@dataclass
+class UtilityParameters(BaseParameters):
+    """
+    Parameters accepted by any of the asynchronous **inference** utility v2 endpoints.
+    """

mindee/mindee_http/mindee_api_v2.py

@@ -4,7 +4,7 @@
 import requests
 
 from mindee.error.mindee_error import MindeeApiV2Error
-from mindee.input import LocalInputSource, UrlInputSource
+from mindee.input import LocalInputSource, UrlInputSource, UtilityParameters
 from mindee.input.inference_parameters import InferenceParameters
 from mindee.logger import logger
 from mindee.mindee_http.base_settings import USER_AGENT
@@ -74,34 +74,37 @@ def set_from_env(self) -> None:
     def req_post_inference_enqueue(
         self,
         input_source: Union[LocalInputSource, UrlInputSource],
-        params: InferenceParameters,
+        params: Union[InferenceParameters, UtilityParameters],
+        slug: Optional[str] = None,
     ) -> requests.Response:
         """
         Make an asynchronous request to POST a document for prediction on the V2 API.
 
         :param input_source: Input object.
         :param params: Options for the enqueueing of the document.
+        :param slug: Slug to use for the enqueueing, defaults to 'inferences'.
         :return: requests response.
         """
+        slug = slug if slug else "inferences"
         data: Dict[str, Union[str, list]] = {"model_id": params.model_id}
-        url = f"{self.url_root}/inferences/enqueue"
-
-        if params.rag is not None:
-            data["rag"] = str(params.rag).lower()
-        if params.raw_text is not None:
-            data["raw_text"] = str(params.raw_text).lower()
-        if params.confidence is not None:
-            data["confidence"] = str(params.confidence).lower()
-        if params.polygon is not None:
-            data["polygon"] = str(params.polygon).lower()
+        url = f"{self.url_root}/{slug}/enqueue"
+        if isinstance(params, InferenceParameters):
+            if params.rag is not None:
+                data["rag"] = str(params.rag).lower()
+            if params.raw_text is not None:
+                data["raw_text"] = str(params.raw_text).lower()
+            if params.confidence is not None:
+                data["confidence"] = str(params.confidence).lower()
+            if params.polygon is not None:
+                data["polygon"] = str(params.polygon).lower()
+            if params.text_context and len(params.text_context):
+                data["text_context"] = params.text_context
+            if params.data_schema is not None:
+                data["data_schema"] = str(params.data_schema)
         if params.webhook_ids and len(params.webhook_ids) > 0:
             data["webhook_ids"] = params.webhook_ids
         if params.alias and len(params.alias):
             data["alias"] = params.alias
-        if params.text_context and len(params.text_context):
-            data["text_context"] = params.text_context
-        if params.data_schema is not None:
-            data["data_schema"] = str(params.data_schema)
 
         if isinstance(input_source, LocalInputSource):
             files = {"file": input_source.read_contents(params.close_file)}

mindee/parsing/v2/inference.py

@@ -1,36 +1,20 @@
 from mindee.parsing.common.string_dict import StringDict
+from mindee.v2.parsing.inference import BaseInference
 from mindee.parsing.v2.inference_active_options import InferenceActiveOptions
-from mindee.parsing.v2.inference_file import InferenceFile
-from mindee.parsing.v2.inference_model import InferenceModel
 from mindee.parsing.v2.inference_result import InferenceResult
 
 
-class Inference:
+class Inference(BaseInference):
     """Inference object for a V2 API return."""
 
-    id: str
-    """ID of the inference."""
-    model: InferenceModel
-    """Model info for the inference."""
-    file: InferenceFile
-    """File info for the inference."""
     result: InferenceResult
     """Result of the inference."""
     active_options: InferenceActiveOptions
     """Active options for the inference."""
+    _slug: str = "inferences"
+    """Slug of the inference."""
 
     def __init__(self, raw_response: StringDict):
-        self.id = raw_response["id"]
-        self.model = InferenceModel(raw_response["model"])
-        self.file = InferenceFile(raw_response["file"])
+        super().__init__(raw_response)
         self.result = InferenceResult(raw_response["result"])
         self.active_options = InferenceActiveOptions(raw_response["active_options"])
-
-    def __str__(self) -> str:
-        return (
-            f"Inference\n#########"
-            f"\n{self.model}"
-            f"\n\n{self.file}"
-            f"\n\n{self.active_options}"
-            f"\n\n{self.result}\n"
-        )