✨ add RAG search API

Author: ianardee

.github/workflows/_test-units.yml

@@ -47,7 +47,7 @@ jobs:
 
     - name: Unit testing with pytest
       run: |
-        pytest --cov mindee --cov-fail-under 87
+        pytest --cov mindee --cov-fail-under 86
 
   pytest-lite:
     name: Run Unit Tests (Lite)

mindee/v2/client.py

@@ -57,7 +57,7 @@ def enqueue(
         :return: A valid inference response.
         """
         logger.debug("Enqueuing inference using model: %s", params.model_id)
-        return self.mindee_api.enqueue(input_source, params)
+        return self.mindee_api.req_post_product_enqueue(input_source, params)
 
     def get_job(self, job_id: str) -> JobResponse:
         """
@@ -70,7 +70,7 @@ def get_job(self, job_id: str) -> JobResponse:
         """
         logger.debug("Fetching job: %s", job_id)
 
-        return self.mindee_api.get_job(job_id)
+        return self.mindee_api.req_get_job_by_id(job_id)
 
     def get_result(
         self,
@@ -88,7 +88,7 @@ def get_result(
         """
         logger.debug("Fetching result: %s", inference_id)
 
-        return self.mindee_api.get_result(response_type, inference_id)
+        return self.mindee_api.req_get_product_result_by_id(response_type, inference_id)
 
     def get_result_from_url(
         self, response_type: type[TypeBaseInferenceResponse], url: str
@@ -100,7 +100,7 @@ def get_result_from_url(
         :param url: URL of the inference to retrieve.
         :return: The result of the inference.
         """
-        return self.mindee_api.get_result_by_url(response_type, url)
+        return self.mindee_api.req_get_product_result_by_url(response_type, url)
 
     def enqueue_and_get_result(
         self,
@@ -179,7 +179,7 @@ def search_models(
         :param model_type: Type of the model to filter by.
         :return: A list of models matching the provided criteria.
         """
-        return self.mindee_api.get_models(name, model_type)
+        return self.mindee_api.req_get_search_models(name, model_type)
 
     def close(self) -> None:
         """Closes the underlying HTTP client."""

mindee/v2/client_options/base_product_parameters.py

@@ -7,7 +7,7 @@
 
 @dataclass
 class BaseProductParameters(ABC):
-    """Base parameters for sending a document to a product."""
+    """Base parameters for sending a file to a Mindee V2 product."""
 
     model_id: str
     """Model ID to use for the inference. Required."""
@@ -32,7 +32,7 @@ class BaseProductParameters(ABC):
     """Whether to close the file after product."""
 
     _slug: ClassVar[str]
-    """Slug of the endpoint."""
+    """Slug of the product."""
 
     def get_request_parameters(self) -> dict[str, str | list[str]]:
         """

mindee/v2/client_options/base_search_parameters.py

@@ -0,0 +1,32 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import ClassVar
+
+
+@dataclass
+class BaseSearchParameters(ABC):
+    """Base parameters for searches."""
+
+    page: int | None
+    """1-based page index."""
+
+    per_page: int | None
+    """Number of items per page."""
+
+    _slug: ClassVar[str]
+    """Slug of the endpoint."""
+
+    def get_request_parameters(self) -> dict[str, str | list[str]]:
+        """
+        Gets the request parameters for the search request.
+
+        :return: A dict of parameters.
+        """
+        data: dict[str, str | list[str]] = {}
+
+        if self.page is not None:
+            data["page"] = str(self.page)
+        if self.per_page is not None:
+            data["per_page"] = str(self.per_page)
+
+        return data

mindee/v2/mindee_http/mindee_api_v2.py

@@ -55,7 +55,7 @@ def __init__(self, api_key: str | None, http_client: httpx.Client | None = None)
             else os.environ.get(API_KEY_V2_ENV_NAME, API_KEY_V2_DEFAULT)
         )
         self.set_base_url(BASE_URL_DEFAULT)
-        self.set_from_env()
+        self._set_from_env()
         if not self.api_key:
             raise MindeeAPIV2Error(
                 f"Missing API key,"
@@ -77,7 +77,7 @@ def base_headers(self) -> dict[str, str]:
             "User-Agent": USER_AGENT,
         }
 
-    def set_from_env(self) -> None:
+    def _set_from_env(self) -> None:
         """Set various parameters from environment variables, if present."""
         env_vars = {
             BASE_URL_ENV_NAME: self.set_base_url,
@@ -89,22 +89,21 @@ def set_from_env(self) -> None:
                 func(env_val)
                 logger.debug("Value was set from env: %s", name)
 
-    def req_post_inference_enqueue(
+    def req_post_product_enqueue(
         self,
         input_source: LocalInputSource | URLInputSource,
         params: BaseProductParameters,
-        slug: str,
-    ) -> httpx.Response:
+    ) -> JobResponse:
         """
-        Make a request to POST a document for enqueue on the V2 API.
+        Send a file to the asynchronous processing queue for inference processing.
 
         :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: httpx response.
         """
         data = params.get_request_parameters()
-        url = f"{self.url_root}/v2/{slug}/enqueue"
+        slug = params.get_enqueue_slug()
+        url = f"{self.url_root}/v2/products/{slug}/enqueue"
         post_kwargs: StringDict = {}
         if isinstance(input_source, LocalInputSource):
             post_kwargs["files"] = {
@@ -116,182 +115,115 @@ def req_post_inference_enqueue(
         post_caller: Callable
         if self.http_client is None or self.http_client.is_closed:
             post_caller = httpx.post
-            post_kwargs["timeout"] = self.request_timeout
         else:
             post_caller = self.http_client.post
-        return post_caller(
+
+        response = post_caller(
             url,
             headers=self.base_headers,
             data=data,
+            timeout=self.request_timeout,
             **post_kwargs,
         )
+        dict_response = self._response_json(response)
+
+        if not is_valid_post_response(response):
+            handle_error_v2(dict_response)
+        return JobResponse(dict_response)
 
-    def req_get_job(self, job_id: str) -> httpx.Response:
+    def req_get_job_by_id(self, job_id: str) -> JobResponse:
         """
-        Sends a request matching a given queue_id. Returns either a Job or a Document.
+        Get the result of an inference that was previously enqueued.
 
         :param job_id: Job ID, returned by the enqueue request.
         """
         get_caller: Callable
-        get_kwargs: StringDict = {}
         if self.http_client is None or self.http_client.is_closed:
             get_caller = httpx.get
-            get_kwargs["timeout"] = self.request_timeout
         else:
             get_caller = self.http_client.get
-        return get_caller(
+
+        response = get_caller(
             url=f"{self.url_root}/v2/jobs/{job_id}",
             headers=self.base_headers,
             follow_redirects=False,
-            **get_kwargs,
+            timeout=self.request_timeout,
         )
+        dict_response = self._response_json(response)
+        if not is_valid_get_response(response):
+            handle_error_v2(dict_response)
+        return JobResponse(dict_response)
 
-    def req_get_inference_by_url(self, url: str) -> httpx.Response:
+    def req_get_product_result_by_url(
+        self, response_type: type[ResponseT], url: str
+    ) -> ResponseT:
         """
-        Sends a request matching a given inference_id. Returns either a Job or a
-        Document.
+        Get the result of an inference that was previously enqueued.
 
         :param url: URL to use for the request.
+        :param response_type: Type of the response to return.
         :return: Response object from the request.
         """
         get_caller: Callable
-        get_kwargs: StringDict = {}
         if self.http_client is None or self.http_client.is_closed:
             get_caller = httpx.get
-            get_kwargs["timeout"] = self.request_timeout
         else:
             get_caller = self.http_client.get
-        return get_caller(
+
+        response = get_caller(
             url=url,
             headers=self.base_headers,
             follow_redirects=False,
-            **get_kwargs,
+            timeout=self.request_timeout,
         )
+        dict_response = self._response_json(response)
+        if not is_valid_get_response(response):
+            handle_error_v2(dict_response)
+        return response_type(dict_response)
 
-    def req_get_inference(self, inference_id: str, slug: str) -> httpx.Response:
+    def req_get_product_result_by_id(
+        self, response_type: type[ResponseT], inference_id: str
+    ) -> ResponseT:
         """
         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.
+        :param response_type: Type of the response to return.
         """
-        get_caller: Callable
-        get_kwargs: StringDict = {}
-        if self.http_client is None or self.http_client.is_closed:
-            get_caller = httpx.get
-            get_kwargs["timeout"] = self.request_timeout
-        else:
-            get_caller = self.http_client.get
-        return get_caller(
-            url=f"{self.url_root}/v2/{slug}/{inference_id}",
-            headers=self.base_headers,
-            follow_redirects=False,
-            **get_kwargs,
+        slug = response_type.get_result_slug()
+        return self.req_get_product_result_by_url(
+            response_type=response_type,
+            url=f"{self.url_root}/v2/products/{slug}/results/{inference_id}",
         )
 
     def req_get_search_models(
         self, name: str | None, model_type: str | None
-    ) -> httpx.Response:
+    ) -> SearchResponse:
         """
         Searches for a list of models matching criteria.
         :param name: Name pattern to search for.
         :param model_type: Type of model to search for (exact match).
         :return: Response object containing search results.
         """
         get_caller: Callable
-        get_kwargs: StringDict = {}
         if self.http_client is None or self.http_client.is_closed:
             get_caller = httpx.get
-            get_kwargs["timeout"] = self.request_timeout
         else:
             get_caller = self.http_client.get
         params = {}
         if name:
             params["name"] = name
         if model_type:
             params["model_type"] = model_type
-        return get_caller(
+
+        response = get_caller(
             url=f"{self.url_root}/v2/search/models",
             headers=self.base_headers,
             params=params,
             follow_redirects=False,
-            **get_kwargs,
-        )
-
-    def enqueue(
-        self,
-        input_source: LocalInputSource | URLInputSource,
-        params: BaseProductParameters,
-    ) -> JobResponse:
-        """
-        Enqueues a document to a given model.
-        :param input_source: Input object.
-        :param params: Parameters
-        :return: A valid inference Response.
-        """
-        response = self.req_post_inference_enqueue(
-            input_source=input_source, params=params, slug=params.get_enqueue_slug()
+            timeout=self.request_timeout,
         )
         dict_response = self._response_json(response)
-
-        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.
-
-        Can be used for polling.
-
-        :param job_id: UUID of the job to retrieve.
-        :return: A job response.
-        """
-        response = self.req_get_job(job_id)
-        dict_response = self._response_json(response)
-        if not is_valid_get_response(response):
-            handle_error_v2(dict_response)
-        return JobResponse(dict_response)
-
-    def get_result(self, response_type: type[ResponseT], inference_id: str):
-        """
-        Get the result of an inference that was previously enqueued.
-
-        :param response_type: Type of the response to return.
-        :param inference_id: UUID of the inference to retrieve.
-        :return: The result of the inference.
-        """
-        response = self.req_get_inference(inference_id, response_type.get_result_slug())
-        dict_response = self._response_json(response)
-        if not is_valid_get_response(response):
-            handle_error_v2(dict_response)
-        return response_type(dict_response)
-
-    def get_result_by_url(self, response_type: type[ResponseT], url: str):
-        """
-        Get the result of an inference that was previously enqueued by its URL.
-
-        :param response_type: Type of the response to return.
-        :param url: URL of the inference to retrieve.
-        :return: The result of the inference.
-        """
-        response = self.req_get_inference_by_url(url)
-        dict_response = self._response_json(response)
-        if not is_valid_get_response(response):
-            handle_error_v2(dict_response)
-        return response_type(dict_response)
-
-    def get_models(self, name: str | None, model_type: str | None):
-        """
-        Get a list of models matching the provided name and type.
-
-        :param name: Name of the model to filter by.
-        :param model_type: Type of the model to filter by.
-        :return: A list of models matching the provided criteria.
-        """
-        logger.debug("Fetching models matching: name=%s and type=%s", name, model_type)
-        response = self.req_get_search_models(name, model_type)
-        dict_response = self._response_json(response)
         if not is_valid_get_response(response):
             handle_error_v2(dict_response)
         return SearchResponse(dict_response)

mindee/v2/parsing/inference/base_inference_response.py

@@ -12,7 +12,7 @@ class BaseInferenceResponse(ABC, CommonResponse):
     """The inference result for a split utility request"""
 
     _slug: ClassVar[str]
-    """Slug of the inference."""
+    """Slug of the product."""
 
     def __str__(self) -> str:
         return str(self.inference)