✨ add support for additional utilities

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.parsing.v2.base_inference import BaseInference, BaseInferenceResponse, TypeInferenceV2
 from mindee.parsing.v2.inference_response import InferenceResponse
 from mindee.parsing.v2.job_response import JobResponse
 
@@ -41,7 +42,8 @@ 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.
@@ -52,16 +54,18 @@ def enqueue_inference(
         :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()
 
         if not is_valid_post_response(response):
             handle_error_v2(dict_response)
         return JobResponse(dict_response)
 
+
     def get_job(self, job_id: str) -> JobResponse:
         """
         Get the status of an inference that was previously enqueued.
@@ -79,13 +83,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_type: Type[BaseInference] = InferenceResponse
+    ) -> InferenceResponse:
         """
         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_type: Class of the product to instantiate.
         :return: An inference response.
         """
         logger.debug("Fetching inference: %s", inference_id)
@@ -94,19 +103,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_type(dict_response)
 
-    def enqueue_and_get_inference(
+    def _enqueue_and_get(
         self,
         input_source: Union[LocalInputSource, UrlInputSource],
-        params: InferenceParameters,
-    ) -> InferenceResponse:
+        params: Union[InferenceParameters, UtilityParameters],
+        inference_type: Optional[Type[BaseInference]] = BaseInference
+    ) -> Union[InferenceResponse, BaseInferenceResponse]:
         """
         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_type: The product class to use for the response object.
 
         :return: A valid inference response.
         """
@@ -117,9 +127,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_type if inference_type.get_slug() 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 +149,48 @@ 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)
+                return self.get_inference(job_response.job.id, inference_type)
             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_type: Type[TypeInferenceV2],
+        input_source: Union[LocalInputSource, UrlInputSource],
+        params: UtilityParameters
+    ) -> TypeInferenceV2:
+        """
+        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_type: The product class to use for the response object.
+
+        :return: A valid inference response.
+        """
+        response = self._enqueue_and_get(input_source, params, inference_type)
+        assert isinstance(response, inference_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,17 @@
 from mindee.input.workflow_options import WorkflowOptions
 
 __all__ = [
+    "Base64Input",
+    "BaseParameters",
+    "BytesInput",
+    "FileInput",
     "InputType",
     "LocalInputSource",
-    "UrlInputSource",
+    "LocalResponse",
+    "PageOptions",
     "PathInput",
-    "FileInput",
-    "Base64Input",
-    "BytesInput",
-    "WorkflowOptions",
     "PollingOptions",
-    "PageOptions",
-    "LocalResponse",
+    "UrlInputSource",
+    "UtilityParameters",
+    "WorkflowOptions",
 ]

mindee/input/base_parameters.py

@@ -0,0 +1,19 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Optional, List
+
+from mindee import PollingOptions
+
+
+@dataclass
+class BaseParameters(ABC):
+    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.utility_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
@@ -76,13 +76,10 @@ def __post_init__(self) -> None:
         elif isinstance(self.replace, str):
             self.replace = DataSchemaReplace(**json.loads(self.replace))
 
-
 @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 +91,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/base_inference/__init__.py

@@ -0,0 +1,4 @@
+from mindee.parsing.v2.base_inference.base_inference import BaseInference
+from mindee.parsing.v2.base_inference.base_inference_response import BaseInferenceResponse, TypeInferenceV2
+from mindee.parsing.v2.base_inference.split.split_inference import SplitInference
+from mindee.parsing.v2.base_inference.split.split_response import SplitResponse

mindee/parsing/v2/base_inference/base_inference.py

@@ -0,0 +1,25 @@
+from abc import ABC
+
+from mindee.parsing.common.string_dict import StringDict
+from mindee.parsing.v2 import InferenceModel, InferenceFile
+
+
+class BaseInference(ABC):
+    model: InferenceModel
+    """Model info for the inference."""
+    file: InferenceFile
+    """File info for the inference."""
+    id: str
+    """ID of the inference."""
+    _slug: str
+    """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"])
+
+    @classmethod
+    def get_slug(cls):
+        """Getter for the inference slug."""
+        return cls._slug

mindee/parsing/v2/base_inference/base_inference_response.py

@@ -0,0 +1,24 @@
+from typing import TypeVar, Generic
+
+from mindee.parsing.common.string_dict import StringDict
+from mindee.parsing.v2.base_inference.base_inference import BaseInference
+
+from mindee.parsing.v2.common_response import CommonResponse
+
+TypeInferenceV2 = TypeVar("TypeInferenceV2", bound=BaseInference)
+
+class BaseInferenceResponse(CommonResponse, Generic[TypeInferenceV2]):
+    inference: TypeInferenceV2
+    """The inference result for a split utility request"""
+
+    def __init__(self, raw_response: StringDict) -> None:
+        super().__init__(raw_response)
+        self.inference = self._set_inference_type(raw_response["inference"])
+
+    def _set_inference_type(self, inference_response: StringDict):
+        """
+        Sets the inference type.
+
+        :param inference_response: Server response.
+        """
+        raise NotImplementedError()