add client methods

Author: ianardee

.github/workflows/_test-units.yml

@@ -59,9 +59,6 @@ jobs:
           - "windows-2022"
         python-version:
           - "3.10"
-          - "3.11"
-          - "3.12"
-          - "3.13"
           - "3.14"
     runs-on: ${{ matrix.os }}
     steps:

mindee/v2/client.py

@@ -10,12 +10,19 @@
 from mindee.logger import logger
 from mindee.mindee_http.cancellation_token import CancellationToken
 from mindee.parsing.common.common_response import CommonStatus
+from mindee.v2.client_options.base_annotation_parameters import BaseAnnotationParameters
 from mindee.v2.client_options.base_product_parameters import BaseProductParameters
+from mindee.v2.client_options.base_rag_document_upload_parameters import (
+    BaseRagDocumentUploadParameters,
+)
 from mindee.v2.client_options.base_search_parameters import (
     BaseSearchParameters,
     TypeSearchResponse,
 )
 from mindee.v2.mindee_http.mindee_api_v2 import MindeeAPIV2
+from mindee.v2.parsing.base_rag_annotation_response import (
+    TypeRagAnnotationResponse,
+)
 from mindee.v2.parsing.inference.base_inference_response import (
     TypeBaseInferenceResponse,
 )
@@ -166,6 +173,143 @@ def enqueue_and_get_result(
 
         raise MindeeError(f"Couldn't retrieve document after {try_counter + 1} tries.")
 
+    def upload_rag_document(
+        self,
+        input_source: LocalInputSource,
+        parameters: BaseRagDocumentUploadParameters[TypeRagAnnotationResponse],
+    ) -> TypeRagAnnotationResponse:
+        """
+        Not recommended for general use, prefer ``upload_and_get_rag_document``.
+        You will need to poll until the document is ready for use.
+        Add a document to the RAG database.
+        """
+        return self.mindee_api.req_post_rag_document(input_source, parameters)
+
+    def upload_and_get_rag_document(
+        self,
+        input_source: LocalInputSource,
+        parameters: BaseRagDocumentUploadParameters[TypeRagAnnotationResponse],
+        polling_options: PollingOptions | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ) -> TypeRagAnnotationResponse:
+        """
+        Add a document to the RAG database and return the initial annotation.
+        """
+        initial_response = self.upload_rag_document(input_source, parameters)
+        if initial_response.status != "Processing":
+            return initial_response
+        if polling_options is None:
+            polling_options = PollingOptions()
+        return self._poll_for_rag_document(
+            initial_response, polling_options, cancellation_token
+        )
+
+    def get_rag_document(
+        self, response_type: type[TypeRagAnnotationResponse], document_id: str
+    ) -> TypeRagAnnotationResponse:
+        """
+        Not recommended for general use, prefer ``get_ready_rag_document``.
+        You will need to poll until the document is ready for use.
+        Get a document's info and annotations from the RAG database.
+        """
+        return self.mindee_api.req_get_rag_annotation(response_type, document_id)
+
+    def get_ready_rag_document(
+        self,
+        response_type: type[TypeRagAnnotationResponse],
+        document_id: str,
+        polling_options: PollingOptions | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ):
+        """
+        Get a document's info and annotations from the RAG database.
+        """
+        initial_response = self.get_rag_document(response_type, document_id)
+        if initial_response.status != "Processing":
+            return initial_response
+        if polling_options is None:
+            polling_options = PollingOptions()
+        return self._poll_for_rag_document(
+            initial_response, polling_options, cancellation_token
+        )
+
+    def update_rag_annotations(
+        self, parameters: BaseAnnotationParameters[TypeRagAnnotationResponse]
+    ) -> TypeRagAnnotationResponse:
+        """
+        Not recommended for general use, prefer ``update_and_get_rag_annotations``.
+        You will need to poll until the document is ready for use.
+        Update a document's annotations in the RAG database.
+        """
+        return self.mindee_api.req_patch_rag_annotation(parameters)
+
+    def update_and_get_rag_annotations(
+        self,
+        parameters: BaseAnnotationParameters[TypeRagAnnotationResponse],
+        polling_options: PollingOptions | None = None,
+        cancellation_token: CancellationToken | None = None,
+    ) -> TypeRagAnnotationResponse:
+        """
+        Update a document's annotations in the RAG database.
+        """
+        initial_response = self.update_rag_annotations(parameters)
+        if initial_response.status != "Processing":
+            return initial_response
+        if polling_options is None:
+            polling_options = PollingOptions()
+        return self._poll_for_rag_document(
+            initial_response, polling_options, cancellation_token
+        )
+
+    def delete_extraction_rag_document(self, document_id: str) -> bool:
+        """
+        Delete a document from the RAG database.
+        For extraction models only.
+        """
+        return self.mindee_api.req_delete_extraction_rag_document(document_id)
+
+    def _poll_for_rag_document(
+        self,
+        initial_response: TypeRagAnnotationResponse,
+        polling_options: PollingOptions,
+        cancellation_token: CancellationToken | None = None,
+    ) -> TypeRagAnnotationResponse:
+        """
+        Poll until the document is finished processing or the max number of attempts is reached.
+        """
+        logger.info("Polling for RAG document ID: %s", initial_response.id)
+        max_retries = polling_options.max_retries + 1
+
+        logger.debug(
+            "Waiting %s seconds before attempting to retrieve the result...",
+            polling_options.initial_delay_sec,
+        )
+
+        if cancellation_token and cancellation_token.is_canceled:
+            raise MindeeError("Request canceled through cancellation token.")
+
+        sleep(polling_options.initial_delay_sec)
+        document_id = initial_response.id
+        retry_count = 1
+
+        while retry_count < max_retries:
+            if cancellation_token and cancellation_token.is_canceled:
+                raise MindeeError("Request canceled through cancellation token.")
+
+            sleep(polling_options.delay_sec)
+            logger.info("Poll attempt %s of %s", retry_count, max_retries)
+
+            response = self.get_rag_document(type(initial_response), document_id)
+            retry_count += 1
+
+            if response.status == "Processing":
+                continue
+            if response.status == "Failed":
+                raise MindeeError("Job failed without an error payload.")
+            return response
+
+        raise MindeeError(f"RAG polling not complete after {retry_count - 1} attempts.")
+
     def search(
         self, params: BaseSearchParameters[TypeSearchResponse]
     ) -> TypeSearchResponse:

mindee/v2/error/mindee_http_error_v2.py

@@ -1,5 +1,6 @@
 import json
 
+from mindee.logger import logger
 from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.error import ErrorItem, ErrorResponse, IErrorResponse
 
@@ -48,4 +49,6 @@ def handle_error_v2(raw_response: StringDict) -> None:
     """
     if "status" not in raw_response or "detail" not in raw_response:
         raise MindeeHTTPUnknownErrorV2(json.dumps(raw_response, indent=2))
-    raise MindeeHTTPErrorV2(ErrorResponse(raw_response))
+    error_response = ErrorResponse(raw_response)
+    logger.error("\n%s", error_response)
+    raise MindeeHTTPErrorV2(error_response)

mindee/v2/mindee_http/mindee_api_v2.py

@@ -11,7 +11,11 @@
 from mindee.mindee_http.settings_mixin import SettingsMixin
 from mindee.parsing.common.string_dict import StringDict
 from mindee.v1.mindee_http.base_settings import USER_AGENT
+from mindee.v2.client_options.base_annotation_parameters import BaseAnnotationParameters
 from mindee.v2.client_options.base_product_parameters import BaseProductParameters
+from mindee.v2.client_options.base_rag_document_upload_parameters import (
+    BaseRagDocumentUploadParameters,
+)
 from mindee.v2.client_options.base_search_parameters import (
     BaseSearchParameters,
     TypeSearchResponse,
@@ -21,6 +25,7 @@
     MindeeHTTPUnknownErrorV2,
     handle_error_v2,
 )
+from mindee.v2.parsing.base_rag_annotation_response import TypeRagAnnotationResponse
 from mindee.v2.parsing.inference.base_inference_response import (
     TypeBaseInferenceResponse,
 )
@@ -224,6 +229,74 @@ def req_get_search_models(
             handle_error_v2(dict_response)
         return SearchResponse(dict_response)
 
+    def req_post_rag_document(
+        self,
+        input_source: LocalInputSource,
+        params: BaseRagDocumentUploadParameters[TypeRagAnnotationResponse],
+    ) -> TypeRagAnnotationResponse:
+        """Add a document to the RAG database."""
+        response_class = params.get_response_class()
+        slug = response_class.get_product_slug()
+        response = self.post_caller(
+            url=f"{self.url_root}/v2/products/{slug}/rag-documents",
+            headers=self.base_headers,
+            files={"file": input_source.read_contents(params.close_file)},
+            data=params.get_request_parameters(),
+            follow_redirects=False,
+            timeout=self.request_timeout,
+        )
+        dict_response = self._response_json(response)
+        if not is_valid_sync_response(response):
+            handle_error_v2(dict_response)
+        return response_class(dict_response)
+
+    def req_get_rag_annotation(
+        self, response_class: type[TypeRagAnnotationResponse], document_id: str
+    ) -> TypeRagAnnotationResponse:
+        """Get a document's info and annotations from the RAG database."""
+        slug = response_class.get_product_slug()
+        response = self.get_caller(
+            url=f"{self.url_root}/v2/products/{slug}/rag-documents/{document_id}",
+            headers=self.base_headers,
+            follow_redirects=False,
+            timeout=self.request_timeout,
+        )
+        dict_response = self._response_json(response)
+        if not is_valid_sync_response(response):
+            handle_error_v2(dict_response)
+        return response_class(dict_response)
+
+    def req_patch_rag_annotation(
+        self, params: BaseAnnotationParameters[TypeRagAnnotationResponse]
+    ) -> TypeRagAnnotationResponse:
+        """Update a document's annotations in the RAG database."""
+        response_class = params.get_response_class()
+        slug = response_class.get_product_slug()
+        response = self.patch_caller(
+            url=f"{self.url_root}/v2/products/{slug}/rag-documents/{params.document_id}",
+            headers=self.base_headers,
+            timeout=self.request_timeout,
+            follow_redirects=False,
+            json=params.get_request_parameters(),
+        )
+        dict_response = self._response_json(response)
+        if not is_valid_sync_response(response):
+            handle_error_v2(dict_response)
+        return response_class(dict_response)
+
+    def req_delete_extraction_rag_document(self, document_id: str) -> bool:
+        """
+        Deletes a document from the RAG database.
+        For extraction models only.
+        """
+        response = self.delete_caller(
+            url=f"{self.url_root}/v2/products/extraction/rag-documents/{document_id}",
+            headers=self.base_headers,
+            timeout=self.request_timeout,
+            follow_redirects=False,
+        )
+        return response.is_success()
+
     @property
     def get_caller(self) -> Callable:
         """Caller for GET requests."""

mindee/v2/product/extraction/rag_documents/__init__.py

@@ -1,3 +1,12 @@
+from mindee.v2.product.extraction.rag_documents.annotated_list_field import (
+    AnnotatedListField,
+)
+from mindee.v2.product.extraction.rag_documents.annotated_object_field import (
+    AnnotatedObjectField,
+)
+from mindee.v2.product.extraction.rag_documents.annotated_simple_field import (
+    AnnotatedSimpleField,
+)
 from mindee.v2.product.extraction.rag_documents.extraction_rag_annotation_response import (
     ExtractionRagAnnotationResponse,
 )
@@ -10,6 +19,9 @@
 from mindee.v2.product.extraction.rag_documents.rag_annotation import RagAnnotation
 
 __all__ = [
+    "AnnotatedListField",
+    "AnnotatedObjectField",
+    "AnnotatedSimpleField",
     "ExtractionRagAnnotationResponse",
     "RagAnnotation",
     "RagDocumentAnnotationParameters",

mindee/v2/product/extraction/rag_documents/annotated_base_field.py

@@ -67,4 +67,4 @@ def build(cls, raw_response: dict) -> AnnotatedFieldsType:
         for key, subclass in cls._registry.items():
             if key in raw_response:
                 return subclass(raw_response)
-        raise ValueError("Invalid structure for field")
+        raise ValueError(f"Invalid structure for field: '{raw_response}'")