final fixes

Author: ianardee

mindee/v2/client.py

@@ -295,14 +295,13 @@ def _poll_for_rag_document(
         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":
+                sleep(polling_options.delay_sec)
                 continue
             if response.status == "Failed":
                 raise MindeeError("Job failed without an error payload.")

mindee/v2/client_options/base_annotation_parameters.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 from typing import Generic
 
+from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.base_rag_annotation_response import TypeRagAnnotationResponse
 
 
@@ -17,7 +18,7 @@ class BaseAnnotationParameters(ABC, Generic[TypeRagAnnotationResponse]):
     """Response class for the annotation."""
 
     @abstractmethod
-    def get_request_parameters(self) -> dict[str, str]:
+    def get_request_parameters(self) -> dict[str, str | StringDict]:
         """Gets the request parameters for the upload request."""
 
     def get_response_class(self) -> type[TypeRagAnnotationResponse]:

mindee/v2/mindee_http/mindee_api_v2.py

@@ -295,7 +295,7 @@ def req_delete_extraction_rag_document(self, document_id: str) -> bool:
             timeout=self.request_timeout,
             follow_redirects=False,
         )
-        return response.is_success()
+        return response.is_success
 
     @property
     def get_caller(self) -> Callable:

mindee/v2/product/extraction/rag_documents/params/rag_document_annotation_parameters.py

@@ -1,5 +1,5 @@
 import json
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from typing import Any, ClassVar
 
 from mindee.parsing.common.string_dict import StringDict
@@ -33,14 +33,14 @@ def __post_init__(self) -> None:
         elif isinstance(self.annotation, dict):
             self.annotation = RagAnnotation(self.annotation)
 
-    def get_request_parameters(self) -> dict[str, str]:
+    def get_request_parameters(self) -> dict[str, str | StringDict]:
         """Gets the request parameters for the upload request."""
         parameters: dict[str, Any] = {}
 
         if self.status:
             parameters["status"] = self.status
 
         if self.annotation is not None and isinstance(self.annotation, RagAnnotation):
-            parameters["annotation"] = str(self.annotation)
+            parameters["annotation"] = asdict(self.annotation)
 
         return parameters

tests/v2/product/extraction/test_rag_documents.py

@@ -4,6 +4,7 @@
 
 from mindee.v2.product.extraction.rag_documents import (
     ExtractionRagAnnotationResponse,
+    RagAnnotation,
     RagDocumentAnnotationParameters,
     RagDocumentUploadParameters,
 )
@@ -18,6 +19,14 @@
 )
 from tests.utils import V2_PRODUCT_PATH
 
+expected_annotation = {
+    "fields": {
+        "simple": {"selected": False, "guidelines": None, "value": True},
+        "list": {"selected": False, "guidelines": None, "items": []},
+        "object": {"selected": False, "guidelines": None, "fields": {}},
+    }
+}
+
 
 @pytest.mark.v2
 def test_post_parameters_must_init():
@@ -30,17 +39,41 @@ def test_post_parameters_must_init():
 @pytest.mark.v2
 def test_patch_parameters_must_init_from_dict():
     """Should init PATCH parameters from an annotation dict."""
-    annotation = {"fields": {"is_good": {"value": True}}}
+    annotation = {
+        "fields": {
+            "simple": {"value": True},
+            "list": {"items": []},
+            "object": {"fields": {}},
+        }
+    }
     parameters = RagDocumentAnnotationParameters(
         document_id="invalid-document-id", status="Active", annotation=annotation
     )
     req_params = parameters.get_request_parameters()
     assert parameters.document_id == "invalid-document-id"
     assert req_params["status"] == "Active"
-    assert (
-        req_params["annotation"]
-        == '{"fields":{"is_good":{"selected":false,"value":true}}}'
+    assert req_params["annotation"] == expected_annotation
+
+
+@pytest.mark.v2
+def test_patch_parameters_must_init_from_object():
+    """Should init PATCH parameters from an annotation instance."""
+    annotation = RagAnnotation(
+        {
+            "fields": {
+                "simple": {"value": True},
+                "list": {"items": []},
+                "object": {"fields": {}},
+            }
+        }
     )
+    parameters = RagDocumentAnnotationParameters(
+        document_id="invalid-document-id", status="Active", annotation=annotation
+    )
+    req_params = parameters.get_request_parameters()
+    assert parameters.document_id == "invalid-document-id"
+    assert req_params["status"] == "Active"
+    assert req_params["annotation"] == expected_annotation
 
 
 @pytest.mark.v2
@@ -49,15 +82,12 @@ def test_patch_parameters_must_init_from_string():
     parameters = RagDocumentAnnotationParameters(
         document_id="invalid-document-id",
         status="Active",
-        annotation='{"fields": {"is_good": {"value": true}}}',
+        annotation='{"fields": {"simple": {"value": true},"list": {"items": []},"object": {"fields": {}}}}',
     )
     req_params = parameters.get_request_parameters()
     assert parameters.document_id == "invalid-document-id"
     assert req_params["status"] == "Active"
-    assert (
-        req_params["annotation"]
-        == '{"fields":{"is_good":{"selected":false,"value":true}}}'
-    )
+    assert req_params["annotation"] == expected_annotation
 
 
 @pytest.mark.v2

tests/v2/product/extraction/test_rag_documents_integration.py

@@ -3,8 +3,8 @@
 import pytest
 
 from mindee import PathInput
-from mindee.error.mindee_error import MindeeError
 from mindee.v2 import Client
+from mindee.v2.error import MindeeHTTPErrorV2
 from mindee.v2.product.extraction.rag_documents import (
     ExtractionRagAnnotationResponse,
     RagDocumentAnnotationParameters,
@@ -28,9 +28,8 @@ def test_rag_document_lifecycle_must_succeed():
     input_source = PathInput(
         V2_PRODUCT_PATH / "extraction" / "financial_document" / "default_sample.jpg"
     )
-    parameters = RagDocumentUploadParameters(model_id=extraction_model_id)
-
-    post_response = client.upload_and_get_rag_document(input_source, parameters)
+    post_parameters = RagDocumentUploadParameters(model_id=extraction_model_id)
+    post_response = client.upload_and_get_rag_document(input_source, post_parameters)
     assert post_response is not None
 
     post_annotation = post_response.annotation
@@ -101,5 +100,5 @@ def test_rag_document_lifecycle_must_succeed():
     delete_response = client.delete_extraction_rag_document(document_id)
     assert delete_response is True
 
-    with pytest.raises(MindeeError):
+    with pytest.raises(MindeeHTTPErrorV2):
         client.get_rag_document(ExtractionRagAnnotationResponse, document_id)