nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/ingest/processing.py

@@ -25,7 +25,7 @@ import uuid
 from collections import defaultdict
 from contextlib import AsyncExitStack
 from enum import Enum
-from typing import Any, Optional
+from typing import Any
 
 import aiohttp
 import backoff
@@ -132,19 +132,19 @@ def to_processing_driver_type(file_backend_driver: FileBackendConfig) -> Process
 class ProcessingEngine:
     def __init__(
         self,
-        nuclia_service_account: Optional[str] = None,
-        nuclia_zone: Optional[str] = None,
-        nuclia_public_url: Optional[str] = None,
-        nuclia_processing_cluster_url: Optional[str] = None,
-        onprem: Optional[bool] = False,
-        nuclia_jwt_key: Optional[str] = None,
+        nuclia_service_account: str | None = None,
+        nuclia_zone: str | None = None,
+        nuclia_public_url: str | None = None,
+        nuclia_processing_cluster_url: str | None = None,
+        onprem: bool | None = False,
+        nuclia_jwt_key: str | None = None,
         days_to_keep: int = 3,
         driver: FileBackendConfig = FileBackendConfig.GCS,
     ):
         self.nuclia_service_account = nuclia_service_account
         self.nuclia_zone = nuclia_zone
         if nuclia_public_url is not None:
-            self.nuclia_public_url: Optional[str] = nuclia_public_url.format(zone=nuclia_zone)
+            self.nuclia_public_url: str | None = nuclia_public_url.format(zone=nuclia_zone)
         else:
             self.nuclia_public_url = None
 
@@ -196,7 +196,7 @@ class ProcessingEngine:
         return jwt.encode(payload, self.nuclia_jwt_key, algorithm="HS256")
 
     def generate_file_token_from_fieldfile(
-        self, file: FieldFilePB, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file: FieldFilePB, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         if self.nuclia_jwt_key is None:
             raise AttributeError("Nuclia JWT key not set")
@@ -235,7 +235,7 @@ class ProcessingEngine:
     )
     @processing_observer.wrap({"type": "file_field_upload"})
     async def convert_filefield_to_str(
-        self, file: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         # Upload file without storing on Nuclia DB
         headers = {}
@@ -273,7 +273,7 @@ class ProcessingEngine:
         ).decode()
 
     def convert_external_filefield_to_str(
-        self, file_field: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file_field: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         if self.nuclia_jwt_key is None:
             raise AttributeError("Nuclia JWT key not set")
@@ -313,7 +313,7 @@ class ProcessingEngine:
         self,
         file: FieldFilePB,
         storage: Storage,
-        classif_labels: Optional[list[ClassificationLabel]] = None,
+        classif_labels: list[ClassificationLabel] | None = None,
     ) -> str:
         """It's already an internal file that needs to be uploaded"""
         if self.onprem is False:
@@ -438,7 +438,7 @@ class ProcessingEngine:
             queue=QueueType(queue_type) if queue_type is not None else None,
         )
 
-    async def delete_from_processing(self, *, kbid: str, resource_id: Optional[str] = None) -> None:
+    async def delete_from_processing(self, *, kbid: str, resource_id: str | None = None) -> None:
         """
         Delete a resource from processing. This prevents inflight resources from being processed
         and wasting resources.
@@ -479,7 +479,7 @@ class DummyProcessingEngine(ProcessingEngine):
         pass
 
     async def convert_filefield_to_str(
-        self, file: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         self.calls.append([file])
         index = len(self.values["convert_filefield_to_str"])
@@ -487,7 +487,7 @@ class DummyProcessingEngine(ProcessingEngine):
         return f"convert_filefield_to_str,{index}"
 
     def convert_external_filefield_to_str(
-        self, file_field: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file_field: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         self.calls.append([file_field])
         index = len(self.values["convert_external_filefield_to_str"])
@@ -498,7 +498,7 @@ class DummyProcessingEngine(ProcessingEngine):
         self,
         file: FieldFilePB,
         storage: Storage,
-        classif_labels: Optional[list[ClassificationLabel]] = None,
+        classif_labels: list[ClassificationLabel] | None = None,
     ) -> str:
         self.calls.append([file, storage])
         index = len(self.values["convert_internal_filefield_to_str"])
@@ -516,5 +516,5 @@ class DummyProcessingEngine(ProcessingEngine):
         self.values["send_to_process"].append([item, partition])
         return ProcessingInfo(seqid=len(self.calls), account_seq=0, queue=QueueType.SHARED)
 
-    async def delete_from_processing(self, *, kbid: str, resource_id: Optional[str] = None) -> None:
+    async def delete_from_processing(self, *, kbid: str, resource_id: str | None = None) -> None:
         self.calls.append([kbid, resource_id])

nucliadb/ingest/serialize.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional, Union
 
 import nucliadb_models as models
 from nucliadb.common import datamanagers
@@ -32,15 +31,22 @@ from nucliadb.ingest.fields.link import Link
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.orm.resource import Resource as ORMResource
 from nucliadb_models.common import FieldTypeName
+from nucliadb_models.metadata import Extra, Origin, Relation
 from nucliadb_models.resource import (
     ConversationFieldData,
     ConversationFieldExtractedData,
     Error,
     ExtractedDataType,
     ExtractedDataTypeName,
+    ExtractedText,
+    FieldComputedMetadata,
+    FieldQuestionAnswers,
+    FileExtractedData,
     FileFieldData,
     FileFieldExtractedData,
     GenericFieldData,
+    LargeComputedMetadata,
+    LinkExtractedData,
     LinkFieldData,
     LinkFieldExtractedData,
     QueueType,
@@ -48,6 +54,7 @@ from nucliadb_models.resource import (
     ResourceData,
     TextFieldData,
     TextFieldExtractedData,
+    VectorObject,
 )
 from nucliadb_models.search import ResourceProperties
 from nucliadb_models.security import ResourceSecurity
@@ -56,84 +63,15 @@ from nucliadb_protos.writer_pb2 import FieldStatus
 from nucliadb_utils.utilities import get_storage
 
 
-async def set_resource_field_extracted_data(
-    field: Field,
-    field_data: ExtractedDataType,
-    field_type_name: FieldTypeName,
-    wanted_extracted_data: list[ExtractedDataTypeName],
-) -> None:
-    if field_data is None:
-        return
-
-    if ExtractedDataTypeName.TEXT in wanted_extracted_data:
-        data_et = await field.get_extracted_text()
-        if data_et is not None:
-            field_data.text = from_proto.extracted_text(data_et)
-
-    metadata_wanted = ExtractedDataTypeName.METADATA in wanted_extracted_data
-    shortened_metadata_wanted = ExtractedDataTypeName.SHORTENED_METADATA in wanted_extracted_data
-    if metadata_wanted or shortened_metadata_wanted:
-        data_fcm = await field.get_field_metadata()
-
-        if data_fcm is not None:
-            field_data.metadata = from_proto.field_computed_metadata(
-                data_fcm, shortened=shortened_metadata_wanted and not metadata_wanted
-            )
-
-    if ExtractedDataTypeName.LARGE_METADATA in wanted_extracted_data:
-        data_lcm = await field.get_large_field_metadata()
-        if data_lcm is not None:
-            field_data.large_metadata = from_proto.large_computed_metadata(data_lcm)
-
-    if ExtractedDataTypeName.VECTOR in wanted_extracted_data:
-        # XXX: our extracted API is not vectorset-compatible, so we'll get the
-        # first vectorset and return the values. Ideally, we should provide a
-        # way to select a vectorset
-        vectorset_id = None
-        async with datamanagers.with_ro_transaction() as txn:
-            async for vectorset_id, vs in datamanagers.vectorsets.iter(
-                txn=txn,
-                kbid=field.resource.kb.kbid,
-            ):
-                break
-        assert vectorset_id is not None, "All KBs must have at least a vectorset"
-        data_vec = await field.get_vectors(vectorset_id, vs.storage_key_kind)
-        if data_vec is not None:
-            field_data.vectors = from_proto.vector_object(data_vec)
-
-    if ExtractedDataTypeName.QA in wanted_extracted_data:
-        qa = await field.get_question_answers()
-        if qa is not None:
-            field_data.question_answers = from_proto.field_question_answers(qa)
-
-    if (
-        isinstance(field, File)
-        and isinstance(field_data, FileFieldExtractedData)
-        and ExtractedDataTypeName.FILE in wanted_extracted_data
-    ):
-        data_fed = await field.get_file_extracted_data()
-        if data_fed is not None:
-            field_data.file = from_proto.file_extracted_data(data_fed)
-
-    if (
-        isinstance(field, Link)
-        and isinstance(field_data, LinkFieldExtractedData)
-        and ExtractedDataTypeName.LINK in wanted_extracted_data
-    ):
-        data_led = await field.get_link_extracted_data()
-        if data_led is not None:
-            field_data.link = from_proto.link_extracted_data(data_led)
-
-
 async def serialize(
     kbid: str,
-    rid: Optional[str],
+    rid: str | None,
     show: list[ResourceProperties],
     field_type_filter: list[FieldTypeName],
     extracted: list[ExtractedDataTypeName],
-    service_name: Optional[str] = None,
-    slug: Optional[str] = None,
-) -> Optional[Resource]:
+    service_name: str | None = None,
+    slug: str | None = None,
+) -> Resource | None:
     driver = get_driver()
     async with driver.ro_transaction() as txn:
         return await managed_serialize(
@@ -148,45 +86,56 @@ async def serialize(
         )
 
 
-async def serialize_field_errors(
-    field: Field,
-    serialized: Union[
-        TextFieldData, FileFieldData, LinkFieldData, ConversationFieldData, GenericFieldData
-    ],
-):
-    status = await field.get_status()
-    if status is None:
-        status = FieldStatus()
-    serialized.status = status.Status.Name(status.status)
-    if status.errors:
-        serialized.errors = []
-        for error in status.errors:
-            serialized.errors.append(
-                Error(
-                    body=error.source_error.error,
-                    code=error.source_error.code,
-                    code_str=writer_pb2.Error.ErrorCode.Name(error.source_error.code),
-                    created=error.created.ToDatetime(),
-                    severity=writer_pb2.Error.Severity.Name(error.source_error.severity),
-                )
-            )
-        serialized.error = serialized.errors[-1]
-
-
 async def managed_serialize(
     txn: Transaction,
     kbid: str,
-    rid: Optional[str],
+    rid: str | None,
     show: list[ResourceProperties],
     field_type_filter: list[FieldTypeName],
     extracted: list[ExtractedDataTypeName],
-    service_name: Optional[str] = None,
-    slug: Optional[str] = None,
-) -> Optional[Resource]:
+    service_name: str | None = None,
+    slug: str | None = None,
+) -> Resource | None:
     orm_resource = await get_orm_resource(txn, kbid, rid=rid, slug=slug, service_name=service_name)
     if orm_resource is None:
         return None
 
+    return await serialize_resource(orm_resource, show, field_type_filter, extracted)
+
+
+async def get_orm_resource(
+    txn: Transaction,
+    kbid: str,
+    rid: str | None,
+    slug: str | None = None,
+    service_name: str | None = None,
+) -> ORMResource | None:
+    storage = await get_storage(service_name=service_name)
+
+    kb = KnowledgeBox(txn, storage, kbid)
+
+    if rid is None:
+        if slug is None:
+            raise ValueError("Either rid or slug parameters should be used")
+
+        rid = await kb.get_resource_uuid_by_slug(slug)
+        if rid is None:
+            # Could not find resource uuid from slug
+            return None
+
+    orm_resource = await kb.get(rid)
+    if orm_resource is None:
+        return None
+
+    return orm_resource
+
+
+async def serialize_resource(
+    orm_resource: ORMResource,
+    show: list[ResourceProperties],
+    field_type_filter: list[FieldTypeName],
+    extracted: list[ExtractedDataTypeName],
+) -> Resource:
     resource = Resource(id=orm_resource.uuid)
 
     include_values = ResourceProperties.VALUES in show
@@ -230,29 +179,18 @@ async def managed_serialize(
             resource.queue = QueueType[orm_resource.basic.QueueType.Name(orm_resource.basic.queue)]
 
             if ResourceProperties.RELATIONS in show:
-                relations = await orm_resource.get_user_relations()
-                resource.usermetadata.relations = [
-                    from_proto.relation(rel) for rel in relations.relations
-                ]
+                resource.usermetadata.relations = await serialize_user_relations(orm_resource)
 
     if ResourceProperties.ORIGIN in show:
-        await orm_resource.get_origin()
-        if orm_resource.origin is not None:
-            resource.origin = from_proto.origin(orm_resource.origin)
+        resource.origin = await serialize_origin(orm_resource)
 
     if ResourceProperties.EXTRA in show:
-        await orm_resource.get_extra()
-        if orm_resource.extra is not None:
-            resource.extra = from_proto.extra(orm_resource.extra)
+        resource.extra = await serialize_extra(orm_resource)
 
     include_errors = ResourceProperties.ERRORS in show
 
     if ResourceProperties.SECURITY in show:
-        await orm_resource.get_security()
-        resource.security = ResourceSecurity(access_groups=[])
-        if orm_resource.security is not None:
-            for gid in orm_resource.security.access_groups:
-                resource.security.access_groups.append(gid)
+        resource.security = await serialize_security(orm_resource)
 
     if (field_type_filter and (include_values or include_extracted_data)) or include_errors:
         await orm_resource.get_fields()
@@ -360,38 +298,157 @@ async def managed_serialize(
     return resource
 
 
-async def get_orm_resource(
-    txn: Transaction,
-    kbid: str,
-    rid: Optional[str],
-    slug: Optional[str] = None,
-    service_name: Optional[str] = None,
-) -> Optional[ORMResource]:
-    storage = await get_storage(service_name=service_name)
+async def serialize_origin(resource: ORMResource) -> Origin | None:
+    origin = await resource.get_origin()
+    if origin is None:
+        return None
 
-    kb = KnowledgeBox(txn, storage, kbid)
+    return from_proto.origin(origin)
 
-    if rid is None:
-        if slug is None:
-            raise ValueError("Either rid or slug parameters should be used")
 
-        rid = await kb.get_resource_uuid_by_slug(slug)
-        if rid is None:
-            # Could not find resource uuid from slug
-            return None
+async def serialize_extra(resource: ORMResource) -> Extra | None:
+    extra = await resource.get_extra()
+    if extra is None:
+        return None
+    return from_proto.extra(extra)
 
-    orm_resource = await kb.get(rid)
-    if orm_resource is None:
+
+async def serialize_user_relations(resource: ORMResource) -> list[Relation]:
+    relations = await resource.get_user_relations()
+    return [from_proto.relation(rel) for rel in relations.relations]
+
+
+async def serialize_security(resource: ORMResource) -> ResourceSecurity:
+    security = ResourceSecurity(access_groups=[])
+
+    security_pb = await resource.get_security()
+    if security_pb is not None:
+        for gid in security_pb.access_groups:
+            security.access_groups.append(gid)
+
+    return security
+
+
+async def serialize_field_errors(
+    field: Field,
+    serialized: (
+        TextFieldData | FileFieldData | LinkFieldData | ConversationFieldData | GenericFieldData
+    ),
+):
+    status = await field.get_status()
+    if status is None:
+        status = FieldStatus()
+    serialized.status = status.Status.Name(status.status)
+    if status.errors:
+        serialized.errors = []
+        for error in status.errors:
+            serialized.errors.append(
+                Error(
+                    body=error.source_error.error,
+                    code=error.source_error.code,
+                    code_str=writer_pb2.Error.ErrorCode.Name(error.source_error.code),
+                    created=error.created.ToDatetime(),
+                    severity=writer_pb2.Error.Severity.Name(error.source_error.severity),
+                )
+            )
+        serialized.error = serialized.errors[-1]
+
+
+async def set_resource_field_extracted_data(
+    field: Field,
+    field_data: ExtractedDataType,
+    field_type_name: FieldTypeName,
+    wanted_extracted_data: list[ExtractedDataTypeName],
+) -> None:
+    if field_data is None:
+        return
+
+    if ExtractedDataTypeName.TEXT in wanted_extracted_data:
+        field_data.text = await serialize_extracted_text(field)
+
+    metadata_wanted = ExtractedDataTypeName.METADATA in wanted_extracted_data
+    shortened_metadata_wanted = ExtractedDataTypeName.SHORTENED_METADATA in wanted_extracted_data
+    if metadata_wanted or shortened_metadata_wanted:
+        field_data.metadata = await serialize_extracted_metadata(
+            field, shortened=shortened_metadata_wanted and not metadata_wanted
+        )
+
+    if ExtractedDataTypeName.LARGE_METADATA in wanted_extracted_data:
+        field_data.large_metadata = await serialize_extracted_large_metadata(field)
+
+    if ExtractedDataTypeName.VECTOR in wanted_extracted_data:
+        field_data.vectors = await serialize_extracted_vectors(field)
+
+    if ExtractedDataTypeName.QA in wanted_extracted_data:
+        field_data.question_answers = await serialize_extracted_question_answers(field)
+
+    if (
+        isinstance(field, File)
+        and isinstance(field_data, FileFieldExtractedData)
+        and ExtractedDataTypeName.FILE in wanted_extracted_data
+    ):
+        field_data.file = await serialize_file_extracted_data(field)
+
+    if (
+        isinstance(field, Link)
+        and isinstance(field_data, LinkFieldExtractedData)
+        and ExtractedDataTypeName.LINK in wanted_extracted_data
+    ):
+        field_data.link = await serialize_link_extracted_data(field)
+
+
+async def serialize_extracted_text(field: Field) -> ExtractedText | None:
+    data_et = await field.get_extracted_text()
+    if data_et is None:
         return None
+    return from_proto.extracted_text(data_et)
 
-    return orm_resource
 
+async def serialize_extracted_metadata(field: Field, *, shortened: bool) -> FieldComputedMetadata | None:
+    data_fcm = await field.get_field_metadata()
+    if data_fcm is None:
+        return None
+    return from_proto.field_computed_metadata(data_fcm, shortened)
 
-async def get_resource_uuid_by_slug(
-    kbid: str, slug: str, service_name: Optional[str] = None
-) -> Optional[str]:
-    storage = await get_storage(service_name=service_name)
-    driver = get_driver()
-    async with driver.ro_transaction() as txn:
-        kb = KnowledgeBox(txn, storage, kbid)
-        return await kb.get_resource_uuid_by_slug(slug)
+
+async def serialize_extracted_large_metadata(field: Field) -> LargeComputedMetadata | None:
+    data_lcm = await field.get_large_field_metadata()
+    if data_lcm is None:
+        return None
+    return from_proto.large_computed_metadata(data_lcm)
+
+
+async def serialize_extracted_vectors(field: Field) -> VectorObject | None:
+    # XXX: our extracted API is not vectorset-compatible, so we'll get the
+    # first vectorset and return the values. Ideally, we should provide a
+    # way to select a vectorset
+    vectorset_id = None
+    async with datamanagers.with_ro_transaction() as txn:
+        async for vectorset_id, vs in datamanagers.vectorsets.iter(txn=txn, kbid=field.kbid):
+            break
+    assert vectorset_id is not None, "All KBs must have at least a vectorset"
+    data_vec = await field.get_vectors(vectorset_id, vs.storage_key_kind)
+    if data_vec is None:
+        return None
+    return from_proto.vector_object(data_vec)
+
+
+async def serialize_extracted_question_answers(field: Field) -> FieldQuestionAnswers | None:
+    qa = await field.get_question_answers()
+    if qa is None:
+        return None
+    return from_proto.field_question_answers(qa)
+
+
+async def serialize_file_extracted_data(field: File) -> FileExtractedData | None:
+    data_fed = await field.get_file_extracted_data()
+    if data_fed is None:
+        return None
+    return from_proto.file_extracted_data(data_fed)
+
+
+async def serialize_link_extracted_data(field: Link) -> LinkExtractedData | None:
+    data_led = await field.get_link_extracted_data()
+    if data_led is None:
+        return None
+    return from_proto.link_extracted_data(data_led)