fix issues

Author: sebastianMindee

mindee/client_v2.py

@@ -1,5 +1,5 @@
 from time import sleep
-from typing import Optional, Union, Type
+from typing import Optional, Union, Type, TypeVar
 
 from mindee.client_mixin import ClientMixin
 from mindee.error.mindee_error import MindeeError
@@ -19,6 +19,10 @@
 from mindee.parsing.v2.inference_response import InferenceResponse
 from mindee.parsing.v2.job_response import JobResponse
 
+TypeBaseInferenceResponse = TypeVar(
+    "TypeBaseInferenceResponse", bound=BaseInferenceResponse
+)
+
 
 class ClientV2(ClientMixin):
     """
@@ -84,7 +88,7 @@ def get_job(self, job_id: str) -> JobResponse:
     def get_inference(
         self,
         inference_id: str,
-        inference_response_type: Type[InferenceResponse] = InferenceResponse,
+        inference_response_type: Type[BaseInferenceResponse] = InferenceResponse,
     ) -> BaseInferenceResponse:
         """
         Get the result of an inference that was previously enqueued.
@@ -96,8 +100,11 @@ def get_inference(
         :return: An inference response.
         """
         logger.debug("Fetching inference: %s", inference_id)
+        slug = None
+        if inference_response_type and inference_response_type is not InferenceResponse:
+            slug = "utilities/" + inference_response_type.get_inference_slug()
 
-        response = self.mindee_api.req_get_inference(inference_id)
+        response = self.mindee_api.req_get_inference(inference_id, slug)
         if not is_valid_get_response(response):
             handle_error_v2(response.json())
         dict_response = response.json()
@@ -107,8 +114,10 @@ def _enqueue_and_get(
         self,
         input_source: Union[LocalInputSource, UrlInputSource],
         params: Union[InferenceParameters, UtilityParameters],
-        inference_response_type: Optional[Type[InferenceResponse]] = InferenceResponse,
-    ) -> InferenceResponse:
+        inference_response_type: Optional[
+            Type[BaseInferenceResponse]
+        ] = InferenceResponse,
+    ) -> BaseInferenceResponse:
         """
         Enqueues to an asynchronous endpoint and automatically polls for a response.
 
@@ -125,11 +134,9 @@ def _enqueue_and_get(
             params.polling_options.delay_sec,
             params.polling_options.max_retries,
         )
-        slug = (
-            inference_response_type.inference.get_slug()
-            if inference_response_type
-            else None
-        )
+        slug = None
+        if inference_response_type and inference_response_type is not InferenceResponse:
+            slug = "utilities/" + inference_response_type.get_inference_slug()
         enqueue_response = self.enqueue_inference(input_source, params, slug)
         logger.debug(
             "Successfully enqueued document with job id: %s", enqueue_response.job.id
@@ -150,9 +157,6 @@ def _enqueue_and_get(
                 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)
@@ -181,10 +185,10 @@ def enqueue_and_get_inference(
 
     def enqueue_and_get_utility(
         self,
-        inference_response_type: Type[InferenceResponse],
+        inference_response_type: Type[TypeBaseInferenceResponse],
         input_source: Union[LocalInputSource, UrlInputSource],
         params: UtilityParameters,
-    ) -> InferenceResponse:
+    ) -> TypeBaseInferenceResponse:
         """
         Enqueues to an asynchronous endpoint and automatically polls for a response.
 

mindee/mindee_http/mindee_api_v2.py

@@ -85,22 +85,12 @@ def req_post_inference_enqueue(
         :param slug: Slug to use for the enqueueing, defaults to 'inferences'.
         :return: requests response.
         """
-        slug = slug if slug else "inferences"
+        if not slug:
+            slug = "inferences"
         data: Dict[str, Union[str, list]] = {"model_id": params.model_id}
         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)
+            self._set_inference_params(data, params)
         if params.webhook_ids and len(params.webhook_ids) > 0:
             data["webhook_ids"] = params.webhook_ids
         if params.alias and len(params.alias):
@@ -127,6 +117,28 @@ def req_post_inference_enqueue(
             raise MindeeApiV2Error("Invalid input source.")
         return response
 
+    def _set_inference_params(
+        self, data: dict[str, Union[str, list]], params: InferenceParameters
+    ) -> None:
+        """
+        Sets the inference-specific parameters.
+
+        :param data: Data dict to fill.
+        :param params: Parameters to add.
+        """
+        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)
+
     def req_get_job(self, job_id: str) -> requests.Response:
         """
         Sends a request matching a given queue_id. Returns either a Job or a Document.
@@ -140,14 +152,22 @@ def req_get_job(self, job_id: str) -> requests.Response:
             allow_redirects=False,
         )
 
-    def req_get_inference(self, inference_id: str) -> requests.Response:
+    def req_get_inference(
+        self, inference_id: str, slug: Optional[str]
+    ) -> requests.Response:
         """
         Sends a request matching a given queue_id. Returns either a Job or a Document.
 
         :param inference_id: Inference ID, returned by the job request.
+        :param slug: Slug of the inference, defaults to nothing.
         """
+
+        if not slug:
+            url = f"{self.url_root}/inferences/{inference_id}"
+        else:
+            url = f"{self.url_root}/{slug}/{inference_id}"
         return requests.get(
-            f"{self.url_root}/inferences/{inference_id}",
+            url,
             headers=self.base_headers,
             timeout=self.request_timeout,
             allow_redirects=False,

mindee/parsing/v2/inference_response.py

@@ -10,6 +10,7 @@ class InferenceResponse(BaseInferenceResponse[Inference]):
 
     inference: Inference
     """Inference result."""
+    inference_type = Inference
 
     def __init__(self, raw_response: StringDict) -> None:
         super().__init__(raw_response)

mindee/v2/parsing/inference/base_inference.py

@@ -9,14 +9,14 @@
 class BaseInference(ABC):
     """Base class for V2 inference objects."""
 
+    _slug: str
+    """Slug of the inference."""
     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"]

mindee/v2/parsing/inference/base_inference_response.py

@@ -1,16 +1,18 @@
-from typing import TypeVar, Generic
+from typing import ClassVar, Type, TypeVar, Generic
 
 from mindee.parsing.common.string_dict import StringDict
-from mindee.v2.parsing.inference.base_inference import TypeBaseInference
+from mindee.v2.parsing.inference.base_inference import BaseInference, TypeBaseInference
 
 from mindee.parsing.v2.common_response import CommonResponse
 
 
 class BaseInferenceResponse(CommonResponse, Generic[TypeBaseInference]):
     """Base class for V2 inference responses."""
 
-    inference: TypeBaseInference
+    inference: BaseInference
     """The inference result for a split utility request"""
+    inference_type: ClassVar[Type[BaseInference]]
+    """Inference class used for slug derivation."""
 
     def __init__(self, raw_response: StringDict) -> None:
         super().__init__(raw_response)
@@ -24,5 +26,10 @@ def _set_inference_type(self, inference_response: StringDict):
         """
         raise NotImplementedError()
 
+    @classmethod
+    def get_inference_slug(cls) -> str:
+        """Getter for the inference slug."""
+        return cls.inference_type.get_slug()
+
 
 TypeInferenceResponse = TypeVar("TypeInferenceResponse", bound=BaseInferenceResponse)

mindee/v2/parsing/inference/split/split_response.py

@@ -1,6 +1,3 @@
-from typing import override
-
-from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.inference.base_inference_response import (
     BaseInferenceResponse,
 )
@@ -10,8 +7,10 @@
 class SplitResponse(BaseInferenceResponse[SplitInference]):
     """Represent a split inference response from Mindee V2 API."""
 
-    @override
-    def _set_inference_type(self, inference_response: StringDict):
+    inference: SplitInference
+    inference_type = SplitInference
+
+    def _set_inference_type(self, inference_response):
         """
         Sets the inference type.