:sprakles: add API for RAG operations

Author: ianardee

mindee/v2/parsing/base_rag_annotation_response.py

@@ -0,0 +1,27 @@
+from datetime import datetime
+
+from mindee.parsing.common import CommonResponse
+from mindee.parsing.common.string_dict import StringDict
+
+
+class BaseRagAnnotationResponse(CommonResponse):
+    """Base class for all RAG document responses from the V2 API."""
+
+    id: str
+    """Unique identifier of the RAG document."""
+    filename: str
+    """Original filename of the uploaded document."""
+    created_at: datetime
+    """Date and time of the document creation."""
+    status: str
+    """Current status of the RAG document."""
+
+    def __init__(self, raw_response: StringDict) -> None:
+        super().__init__(raw_response)
+        self.id = raw_response["id"]
+        self.filename = raw_response["filename"]
+        self.created_at = datetime.fromisoformat(
+            raw_response["created_at"].replace("Z", "+00:00")
+        )
+        self.status = raw_response["status"]
+

mindee/v2/parsing/inference/field/base_field.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import TYPE_CHECKING, ClassVar, TypeAlias, Union
 
 from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.inference.field.field_confidence import FieldConfidence
@@ -13,17 +14,33 @@ class FieldType(str, Enum):
     SIMPLE = "SimpleField"
 
 
+if TYPE_CHECKING:
+    from mindee.v2.parsing.inference.field.list_field import ListField
+    from mindee.v2.parsing.inference.field.object_field import ObjectField
+    from mindee.v2.parsing.inference.field.simple_field import SimpleField
+
+
+ResultFieldsType: TypeAlias = Union["SimpleField", "ObjectField", "ListField"]
+
+
 class BaseField:
-    """Field with base information."""
+    """Base class for V2 fields."""
 
     field_type: FieldType
-    _indent_level: int
+    """The type of field."""
     locations: list[FieldLocation]
+    """List of the location candidates for the value."""
     confidence: FieldConfidence | None
+    """Confidence associated with the field."""
+    _indent_level: int
+    """For pretty printing."""
+
+    _registry: ClassVar[dict[str, type[ResultFieldsType]]] = {}
 
     def __init__(
         self, field_type: FieldType, raw_response: StringDict, indent_level: int = 0
     ) -> None:
+
         self.field_type = field_type
         self._indent_level = indent_level
 
@@ -41,6 +58,27 @@ def __init__(
             for location in raw_response["locations"]:
                 self.locations.append(FieldLocation(location))
 
+    @classmethod
+    def register(cls, discriminator_key: str):
+        """Class decorator: subclasses declare which JSON key identifies them."""
+
+        def decorator(subclass):
+            cls._registry[discriminator_key] = subclass
+            return subclass
+
+        return decorator
+
+    @classmethod
+    def build(cls, raw_response: dict, indent_level: int) -> ResultFieldsType:
+        """Build an instance of the appropriate subclass."""
+
+        if not isinstance(raw_response, dict):
+            raise ValueError("Field must be a dict")
+        for key, subclass in cls._registry.items():
+            if key in raw_response:
+                return subclass(raw_response, indent_level)
+        raise ValueError("Invalid structure for field")
+
     def multi_str(self) -> str:
         """String representation of the field in a list."""
         return str(self)

mindee/v2/parsing/inference/field/factory.py

@@ -1,27 +1,29 @@
-from collections.abc import Callable
 from typing import TYPE_CHECKING, cast
 
 from mindee.parsing.common.string_dict import StringDict
-from mindee.v2.parsing.inference.field.base_field import BaseField, FieldType
+from mindee.v2.parsing.inference.field.base_field import (
+    BaseField,
+    FieldType,
+    ResultFieldsType,
+)
 
 if TYPE_CHECKING:
     from mindee.v2.parsing.inference.field.list_field import ListField
     from mindee.v2.parsing.inference.field.object_field import ObjectField
     from mindee.v2.parsing.inference.field.simple_field import SimpleField
 
 
-class InferenceFields(dict[str, BaseField]):
+class InferenceFields(dict[str, ResultFieldsType]):
     """Inference fields dict."""
 
     def __init__(
         self,
         raw_response: StringDict,
-        parser_func: Callable[[StringDict, int], BaseField],
         indent_level: int = 0,
     ) -> None:
         super().__init__()
         for key, value in raw_response.items():
-            self[key] = parser_func(value, indent_level)
+            self[key] = BaseField.build(value, indent_level)
 
     def __getattr__(self, item):
         try:

mindee/v2/parsing/inference/field/inference_fields.py

@@ -1,11 +1,10 @@
-from collections.abc import Callable
-
 from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.inference.field.base_field import BaseField, FieldType
 from mindee.v2.parsing.inference.field.object_field import ObjectField
 from mindee.v2.parsing.inference.field.simple_field import SimpleField
 
 
+@BaseField.register("items")
 class ListField(BaseField):
     """List field containing multiple fields."""
 
@@ -15,14 +14,13 @@ class ListField(BaseField):
     def __init__(
         self,
         raw_response: StringDict,
-        parser_func: Callable[[StringDict, int], BaseField],
         indent_level: int = 0,
     ):
         super().__init__(FieldType.LIST, raw_response, indent_level)
 
         self.items = []
         for item in raw_response["items"]:
-            self.items.append(parser_func(item, indent_level))
+            self.items.append(BaseField.build(item, indent_level))
 
     @property
     def simple_items(self) -> list[SimpleField]:

mindee/v2/parsing/inference/field/list_field.py

@@ -1,4 +1,3 @@
-from collections.abc import Callable
 from typing import TYPE_CHECKING, cast
 
 from mindee.parsing.common.string_dict import StringDict
@@ -10,19 +9,18 @@
     from mindee.v2.parsing.inference.field.simple_field import SimpleField
 
 
+@BaseField.register("fields")
 class ObjectField(BaseField):
     """Object field containing multiple fields."""
 
     fields: InferenceFields
     """Fields contained in the object."""
 
-    def __init__(
-        self, raw_response: StringDict, parser_func: Callable, indent_level: int = 0
-    ):
+    def __init__(self, raw_response: StringDict, indent_level: int = 0):
         super().__init__(FieldType.OBJECT, raw_response, indent_level)
         inner_fields = raw_response.get("fields", raw_response)
 
-        self.fields = InferenceFields(inner_fields, parser_func, self._indent_level + 1)
+        self.fields = InferenceFields(inner_fields, self._indent_level + 1)
 
     def single_str(self) -> str:
         """String representation of a single object field."""

mindee/v2/parsing/inference/field/object_field.py

@@ -2,6 +2,7 @@
 from mindee.v2.parsing.inference.field.base_field import BaseField, FieldType
 
 
+@BaseField.register("value")
 class SimpleField(BaseField):
     """Simple field containing a single value."""
 

mindee/v2/parsing/inference/field/simple_field.py

@@ -1,6 +1,5 @@
 from mindee.parsing.common.string_dict import StringDict
 from mindee.v2.parsing.inference.field import InferenceFields
-from mindee.v2.parsing.inference.field.factory import create_field
 from mindee.v2.parsing.inference.rag_metadata import RAGMetadata
 from mindee.v2.parsing.inference.raw_text import RawText
 
@@ -16,7 +15,7 @@ class ExtractionResult:
     """RAG metadata."""
 
     def __init__(self, raw_response: StringDict) -> None:
-        self.fields = InferenceFields(raw_response["fields"], create_field)
+        self.fields = InferenceFields(raw_response["fields"])
         if raw_response.get("raw_text"):
             self.raw_text = RawText(raw_response["raw_text"])
         if raw_response.get("rag"):

mindee/v2/product/extraction/extraction_result.py

@@ -0,0 +1,9 @@
+from mindee.v2.product.extraction.rag_documents.extraction_rag_annotation_response import (
+    ExtractionRagAnnotationResponse,
+)
+from mindee.v2.product.extraction.rag_documents.rag_annotation import RagAnnotation
+
+__all__ = [
+    "ExtractionRagAnnotationResponse",
+    "RagAnnotation",
+]

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

@@ -0,0 +1,63 @@
+from typing import TYPE_CHECKING, ClassVar, TypeAlias, Union
+
+from mindee.parsing.common.string_dict import StringDict
+from mindee.v2.parsing.inference.field.base_field import FieldType
+
+if TYPE_CHECKING:
+    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,
+    )
+
+AnnotatedFieldsType: TypeAlias = Union[
+    "AnnotatedSimpleField", "AnnotatedObjectField", "AnnotatedListField"
+]
+
+
+class AnnotatedBaseField:
+    """Base class for annotated fields."""
+
+    selected: bool = False
+    """When true, use the RAG information for the final result. When false, use the Data Schema information."""
+
+    guidelines: str | None = None
+    """Guidelines or instructions for processing this field."""
+
+    field_type: FieldType
+
+    _registry: ClassVar[dict[str, type[AnnotatedFieldsType]]] = {}
+
+    def __init__(self, field_type: FieldType, raw_response: StringDict):
+        if "selected" in raw_response and raw_response["selected"] is not None:
+            self.selected = raw_response["selected"]
+
+        if "guidelines" in raw_response and raw_response["guidelines"] is not None:
+            self.guidelines = raw_response.get("guidelines")
+
+        self.field_type = field_type
+
+    @classmethod
+    def register(cls, discriminator_key: str):
+        """Class decorator: subclasses declare which JSON key identifies them."""
+
+        def decorator(subclass):
+            cls._registry[discriminator_key] = subclass
+            return subclass
+
+        return decorator
+
+    @classmethod
+    def build(cls, raw_response: dict) -> AnnotatedFieldsType:
+        """Build an instance of the appropriate subclass."""
+
+        if not isinstance(raw_response, dict):
+            raise ValueError("Field must be a dict")
+        for key, subclass in cls._registry.items():
+            if key in raw_response:
+                return subclass(raw_response)
+        raise ValueError("Invalid structure for field")