diff --git a/mindee/v2/client_options/base_annotation_parameters.py b/mindee/v2/client_options/base_annotation_parameters.py new file mode 100644 index 00000000..50cdfab8 --- /dev/null +++ b/mindee/v2/client_options/base_annotation_parameters.py @@ -0,0 +1,2 @@ +class BaseAnnotationParameters: + """Base parameters for document annotations.""" diff --git a/mindee/v2/client_options/base_rag_document_upload_parameters.py b/mindee/v2/client_options/base_rag_document_upload_parameters.py new file mode 100644 index 00000000..7a56d43b --- /dev/null +++ b/mindee/v2/client_options/base_rag_document_upload_parameters.py @@ -0,0 +1,2 @@ +class BaseRagDocumentUploadParameters: + """Base parameters for document upload operations.""" diff --git a/mindee/v2/parsing/base_rag_annotation_response.py b/mindee/v2/parsing/base_rag_annotation_response.py new file mode 100644 index 00000000..ad001fd2 --- /dev/null +++ b/mindee/v2/parsing/base_rag_annotation_response.py @@ -0,0 +1,26 @@ +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"] diff --git a/mindee/v2/parsing/inference/field/base_field.py b/mindee/v2/parsing/inference/field/base_field.py index b8a12450..276e4480 100644 --- a/mindee/v2/parsing/inference/field/base_field.py +++ b/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) diff --git a/mindee/v2/parsing/inference/field/factory.py b/mindee/v2/parsing/inference/field/factory.py deleted file mode 100644 index f980fa5e..00000000 --- a/mindee/v2/parsing/inference/field/factory.py +++ /dev/null @@ -1,17 +0,0 @@ -from mindee.parsing.common.string_dict import StringDict -from mindee.v2.error.mindee_api_v2_error import MindeeAPIV2Error -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 - - -def create_field(raw_response: StringDict, indent_level: int = 0): - """The central parser function to be injected down the tree.""" - if "value" in raw_response: - return SimpleField(raw_response, indent_level) - if "items" in raw_response: - return ListField(raw_response, create_field, indent_level) - if "fields" in raw_response: - return ObjectField(raw_response, create_field, indent_level) - - raise MindeeAPIV2Error(f"Unrecognized field type in {raw_response}.") diff --git a/mindee/v2/parsing/inference/field/inference_fields.py b/mindee/v2/parsing/inference/field/inference_fields.py index 9b73043c..51636819 100644 --- a/mindee/v2/parsing/inference/field/inference_fields.py +++ b/mindee/v2/parsing/inference/field/inference_fields.py @@ -1,8 +1,11 @@ -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 @@ -10,18 +13,17 @@ 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: diff --git a/mindee/v2/parsing/inference/field/list_field.py b/mindee/v2/parsing/inference/field/list_field.py index af844aae..79969f0e 100644 --- a/mindee/v2/parsing/inference/field/list_field.py +++ b/mindee/v2/parsing/inference/field/list_field.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]: diff --git a/mindee/v2/parsing/inference/field/object_field.py b/mindee/v2/parsing/inference/field/object_field.py index fb0d4514..e14e088a 100644 --- a/mindee/v2/parsing/inference/field/object_field.py +++ b/mindee/v2/parsing/inference/field/object_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.""" diff --git a/mindee/v2/parsing/inference/field/simple_field.py b/mindee/v2/parsing/inference/field/simple_field.py index f6e6c809..f2b5fd97 100644 --- a/mindee/v2/parsing/inference/field/simple_field.py +++ b/mindee/v2/parsing/inference/field/simple_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.""" diff --git a/mindee/v2/product/extraction/extraction_result.py b/mindee/v2/product/extraction/extraction_result.py index 2f0c46ea..821aae00 100644 --- a/mindee/v2/product/extraction/extraction_result.py +++ b/mindee/v2/product/extraction/extraction_result.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"): diff --git a/mindee/v2/product/extraction/rag_documents/__init__.py b/mindee/v2/product/extraction/rag_documents/__init__.py new file mode 100644 index 00000000..aee756ad --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/__init__.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", +] diff --git a/mindee/v2/product/extraction/rag_documents/annotated_base_field.py b/mindee/v2/product/extraction/rag_documents/annotated_base_field.py new file mode 100644 index 00000000..c2fbfec3 --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/annotated_base_field.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") diff --git a/mindee/v2/product/extraction/rag_documents/annotated_fields.py b/mindee/v2/product/extraction/rag_documents/annotated_fields.py new file mode 100644 index 00000000..2e933ca2 --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/annotated_fields.py @@ -0,0 +1,49 @@ +from typing import TYPE_CHECKING, cast + +from mindee.parsing.common.string_dict import StringDict +from mindee.v2.parsing.inference.field.base_field import FieldType +from mindee.v2.product.extraction.rag_documents.annotated_base_field import ( + AnnotatedBaseField, + AnnotatedFieldsType, +) + +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, + ) + + +class AnnotatedFields(dict[str, AnnotatedFieldsType]): + """A dictionary of field names and their corresponding annotation.""" + + def __init__(self, raw_response: StringDict): + super().__init__() + for key, value in raw_response.items(): + self[key] = AnnotatedBaseField.build(value) + + def get_simple_field(self, field_name: str) -> "AnnotatedSimpleField": + """Retrieve a simple field by its name.""" + field = self.get(field_name) + if field and field.field_type == FieldType.SIMPLE: + return cast("AnnotatedSimpleField", field) + raise ValueError(f"Field {field_name} is not an AnnotatedSimpleField.") + + def get_list_field(self, field_name: str) -> "AnnotatedListField": + """Retrieve a list field by its name.""" + field = self.get(field_name) + if field and field.field_type == FieldType.LIST: + return cast("AnnotatedListField", field) + raise ValueError(f"Field {field_name} is not an AnnotatedListField.") + + def get_object_field(self, field_name: str) -> "AnnotatedObjectField": + """Retrieve an object field by its name.""" + field = self.get(field_name) + if field and field.field_type == FieldType.OBJECT: + return cast("AnnotatedObjectField", field) + raise ValueError(f"Field {field_name} is not an AnnotatedObjectField.") diff --git a/mindee/v2/product/extraction/rag_documents/annotated_list_field.py b/mindee/v2/product/extraction/rag_documents/annotated_list_field.py new file mode 100644 index 00000000..0ccb0bd7 --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/annotated_list_field.py @@ -0,0 +1,44 @@ +from mindee.parsing.common.string_dict import StringDict +from mindee.v2.parsing.inference.field.base_field import FieldType +from mindee.v2.product.extraction.rag_documents.annotated_base_field import ( + AnnotatedBaseField, +) +from mindee.v2.product.extraction.rag_documents.annotated_object_field import ( + AnnotatedObjectField, +) +from mindee.v2.product.extraction.rag_documents.annotated_simple_field import ( + AnnotatedSimpleField, +) + + +@AnnotatedBaseField.register("items") +class AnnotatedListField(AnnotatedBaseField): + """A ListField with additional configuration for annotation.""" + + items: list[AnnotatedBaseField] + + def __init__(self, raw_response: StringDict): + super().__init__(FieldType.LIST, raw_response) + self.items = [AnnotatedBaseField.build(item) for item in raw_response["items"]] + + @property + def simple_items(self) -> list[AnnotatedSimpleField]: + """List of items as ``AnnotatedSimpleField``.""" + simple_items = [] + for item in self.items: + if isinstance(item, AnnotatedSimpleField): + simple_items.append(item) + else: + raise ValueError("List item is not an AnnotatedSimpleField field.") + return simple_items + + @property + def object_items(self) -> list[AnnotatedObjectField]: + """List of items as ``AnnotatedObjectField``.""" + object_items = [] + for item in self.items: + if isinstance(item, AnnotatedObjectField): + object_items.append(item) + else: + raise ValueError("List item is not an AnnotatedObjectField field.") + return object_items diff --git a/mindee/v2/product/extraction/rag_documents/annotated_object_field.py b/mindee/v2/product/extraction/rag_documents/annotated_object_field.py new file mode 100644 index 00000000..9fbb91f0 --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/annotated_object_field.py @@ -0,0 +1,28 @@ +from mindee.parsing.common.string_dict import StringDict +from mindee.v2.parsing.inference.field.base_field import FieldType +from mindee.v2.product.extraction.rag_documents.annotated_base_field import ( + AnnotatedBaseField, +) +from mindee.v2.product.extraction.rag_documents.annotated_fields import AnnotatedFields +from mindee.v2.product.extraction.rag_documents.annotated_simple_field import ( + AnnotatedSimpleField, +) + + +@AnnotatedBaseField.register("fields") +class AnnotatedObjectField(AnnotatedBaseField): + """An ObjectField with additional configuration for annotation.""" + + fields: AnnotatedFields + + def __init__(self, raw_response: StringDict): + super().__init__(FieldType.OBJECT, raw_response) + self.fields = AnnotatedFields(raw_response["fields"]) + + def get_simple_field(self, field_name: str) -> AnnotatedSimpleField: + """Retrieve a Simple field by its name.""" + return self.fields.get_simple_field(field_name) + + def get_object_field(self, field_name: str) -> "AnnotatedObjectField": + """Retrieve an Object field by its name.""" + return self.fields.get_object_field(field_name) diff --git a/mindee/v2/product/extraction/rag_documents/annotated_simple_field.py b/mindee/v2/product/extraction/rag_documents/annotated_simple_field.py new file mode 100644 index 00000000..6712b0d5 --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/annotated_simple_field.py @@ -0,0 +1,16 @@ +from mindee.parsing.common.string_dict import StringDict +from mindee.v2.parsing.inference.field.base_field import FieldType +from mindee.v2.product.extraction.rag_documents.annotated_base_field import ( + AnnotatedBaseField, +) + + +@AnnotatedBaseField.register("value") +class AnnotatedSimpleField(AnnotatedBaseField): + """A SimpleField with additional configuration for annotation.""" + + value: str | float | bool | None + + def __init__(self, raw_response: StringDict): + super().__init__(FieldType.SIMPLE, raw_response) + self.value = raw_response["value"] diff --git a/mindee/v2/product/extraction/rag_documents/extraction_rag_annotation_response.py b/mindee/v2/product/extraction/rag_documents/extraction_rag_annotation_response.py new file mode 100644 index 00000000..1f1e841b --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/extraction_rag_annotation_response.py @@ -0,0 +1,29 @@ +from datetime import datetime + +from mindee.parsing.common.string_dict import StringDict +from mindee.v2.parsing.base_rag_annotation_response import BaseRagAnnotationResponse +from mindee.v2.product.extraction.rag_documents.rag_annotation import RagAnnotation + + +class ExtractionRagAnnotationResponse(BaseRagAnnotationResponse): + """Response for a RAG document.""" + + model_id: str + """Model identifier linked to the RAG document.""" + total_matches: int + """Number of times this document was used in an inference.""" + last_match_at: datetime | None = None + """Date and time of the latest matching inference, if any.""" + annotation: RagAnnotation | None = None + """Annotation metadata associated with the document.""" + + def __init__(self, raw_response: StringDict) -> None: + super().__init__(raw_response) + self.model_id = raw_response["model_id"] + self.total_matches = raw_response["total_matches"] + if raw_response.get("last_match_at"): + self.last_match_at = datetime.fromisoformat( + raw_response["last_match_at"].replace("Z", "+00:00") + ) + if raw_response.get("annotation"): + self.annotation = RagAnnotation(raw_response["annotation"]) diff --git a/mindee/v2/product/extraction/rag_documents/params/__init__.py b/mindee/v2/product/extraction/rag_documents/params/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/mindee/v2/product/extraction/rag_documents/rag_annotation.py b/mindee/v2/product/extraction/rag_documents/rag_annotation.py new file mode 100644 index 00000000..6fc08f00 --- /dev/null +++ b/mindee/v2/product/extraction/rag_documents/rag_annotation.py @@ -0,0 +1,12 @@ +from mindee.parsing.common.string_dict import StringDict +from mindee.v2.product.extraction.rag_documents.annotated_fields import AnnotatedFields + + +class RagAnnotation: + """A RAG annotation enriched with field-level configuration.""" + + fields: AnnotatedFields + """Annotated fields.""" + + def __init__(self, raw_response: StringDict): + self.fields = AnnotatedFields(raw_response["fields"]) diff --git a/tests/v2/product/extraction/test_extraction_response.py b/tests/v2/product/extraction/test_extraction_response.py index 8da8aec2..33d05223 100644 --- a/tests/v2/product/extraction/test_extraction_response.py +++ b/tests/v2/product/extraction/test_extraction_response.py @@ -7,14 +7,14 @@ FieldConfidence, InferenceFields, ListField, + ObjectField, + SimpleField, ) -from mindee.v2.parsing.inference.field.object_field import ObjectField -from mindee.v2.parsing.inference.field.simple_field import SimpleField from mindee.v2.parsing.inference.inference_active_options import InferenceActiveOptions from mindee.v2.parsing.inference.inference_file import InferenceFile from mindee.v2.parsing.inference.inference_model import InferenceModel from mindee.v2.parsing.inference.rag_metadata import RAGMetadata -from mindee.v2.product.extraction.extraction_inference import ExtractionInference +from mindee.v2.product.extraction import ExtractionInference from tests.utils import V2_PRODUCT_PATH from tests.v2.product.utils import get_product_samples diff --git a/tests/v2/product/extraction/test_rag_documents.py b/tests/v2/product/extraction/test_rag_documents.py new file mode 100644 index 00000000..4e33e812 --- /dev/null +++ b/tests/v2/product/extraction/test_rag_documents.py @@ -0,0 +1,124 @@ +import json + +import pytest + +from mindee.v2.product.extraction.rag_documents import ExtractionRagAnnotationResponse +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 tests.utils import V2_PRODUCT_PATH + + +@pytest.mark.v2 +def test_rag_documents_post_must_have_valid_properties(): + """Should load a POST response from a JSON string.""" + + response = _get_response("extraction/rag_documents/post_response.json") + assert response is not None + assert response.id == "cc831599-c545-48b7-aa27-6d7ccd5b8d32" + assert response.status == "Processing" + assert response.annotation is None + + +@pytest.mark.v2 +def test_rag_documents_get_draft_must_have_valid_properties(): + """Should load a GET response from a JSON string.""" + + response = _get_response("extraction/rag_documents/get_response_draft.json") + assert response is not None + assert response.id == "cc831599-c545-48b7-aa27-6d7ccd5b8d32" + assert response.status == "Draft" + assert response.annotation is not None + fields = response.annotation.fields + assert fields is not None + + # null simple field + tip_field = fields.get_simple_field("tip") + assert isinstance(tip_field, AnnotatedSimpleField) + assert tip_field.selected is False + assert tip_field.guidelines is None + assert tip_field.value is None + + # filled simple field + date_field = fields.get_simple_field("date") + assert isinstance(date_field, AnnotatedSimpleField) + assert date_field.selected is False + assert date_field.guidelines is None + assert date_field.value == "2019-11-02" + + # filled object field + locale_field = fields.get_object_field("locale") + assert isinstance(locale_field, AnnotatedObjectField) + assert locale_field.selected is False + assert locale_field.guidelines is None + assert locale_field.fields is not None + assert len(locale_field.fields) == 3 + assert locale_field.get_simple_field("country").value == "US" + assert locale_field.get_simple_field("currency").value == "USD" + assert locale_field.get_simple_field("language").value is None + + # list of simple fields + reference_numbers_field = fields.get_list_field("reference_numbers") + assert isinstance(reference_numbers_field, AnnotatedListField) + assert reference_numbers_field.selected is False + assert reference_numbers_field.guidelines is None + assert reference_numbers_field.items is not None + assert len(reference_numbers_field.items) == 1 + assert reference_numbers_field.simple_items[0].value == "2412/2019" + + # list of object fields + line_items_field = fields.get_list_field("line_items") + assert isinstance(line_items_field, AnnotatedListField) + assert line_items_field.selected is False + assert line_items_field.guidelines is None + assert line_items_field.items is not None + assert len(line_items_field.items) == 3 + + line_item_0 = line_items_field.object_items[0] + assert line_item_0.fields is not None + assert len(line_item_0.fields) == 8 + assert line_item_0.fields["description"].value == "Front and rear brake cables" + assert line_item_0.fields["quantity"].value == 1 + assert line_item_0.fields["unit_price"].value == 100 + assert line_item_0.fields["total_price"].value == 100 + assert line_item_0.fields["tax_rate"].value is None + assert line_item_0.fields["tax_amount"].value is None + assert line_item_0.fields["product_code"].value is None + assert line_item_0.fields["unit_measure"].value is None + + line_item_1 = line_items_field.object_items[1] + assert line_item_1.fields is not None + assert len(line_item_1.fields) == 8 + assert line_item_1.fields["description"].value == "New set of pedal arms" + assert line_item_1.fields["quantity"].value == 2 + assert line_item_1.fields["unit_price"].value == 25 + assert line_item_1.fields["total_price"].value == 50 + assert line_item_1.fields["tax_rate"].value is None + assert line_item_1.fields["tax_amount"].value is None + assert line_item_1.fields["product_code"].value is None + assert line_item_1.fields["unit_measure"].value is None + + line_item_2 = line_items_field.object_items[2] + assert line_item_2.fields is not None + assert len(line_item_2.fields) == 8 + assert line_item_2.fields["description"].value == "Labor 3hrs" + assert line_item_2.fields["quantity"].value == 3 + assert line_item_2.fields["unit_price"].value == 15 + assert line_item_2.fields["total_price"].value == 45 + assert line_item_2.fields["tax_rate"].value is None + assert line_item_2.fields["tax_amount"].value is None + assert line_item_2.fields["product_code"].value is None + assert line_item_2.fields["unit_measure"].value is None + + +def _get_response(path: str) -> ExtractionRagAnnotationResponse: + file_path = V2_PRODUCT_PATH / path + with open(file_path, encoding="utf-8") as f: + response_dict = json.load(f) + return ExtractionRagAnnotationResponse(response_dict)