nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/ingest/orm/broker_message.py

@@ -0,0 +1,181 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+
+from typing import cast
+
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.fields.conversation import Conversation
+from nucliadb.ingest.fields.file import File
+from nucliadb.ingest.fields.link import Link
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb_protos.resources_pb2 import (
+    ExtractedTextWrapper,
+    ExtractedVectorsWrapper,
+    FieldComputedMetadataWrapper,
+    FieldType,
+    LargeComputedMetadataWrapper,
+)
+from nucliadb_protos.writer_pb2 import BrokerMessage
+
+
+async def generate_broker_message(resource: Resource) -> BrokerMessage:
+    """Generate a full broker message from a resource. This means downloading
+    all the pointers minus the ones to external files that are not PB. Iterate
+    all resource fields and create a BrokerMessage
+    """
+    builder = _BrokerMessageBuilder()
+    bm = await builder.build_from(resource)
+    return bm
+
+
+class _BrokerMessageBuilder:
+    def __init__(self):
+        self.bm = BrokerMessage()
+
+    async def build_from(self, resource: Resource):
+        # clear the state and generate a new broker message
+        self.bm.Clear()
+
+        self.bm.kbid = resource.kb.kbid
+        self.bm.uuid = resource.uuid
+        basic = await resource.get_basic()
+        if basic is not None:
+            self.bm.basic.CopyFrom(basic)
+
+        self.bm.slug = self.bm.basic.slug
+        origin = await resource.get_origin()
+        if origin is not None:
+            self.bm.origin.CopyFrom(origin)
+        relations = await resource.get_relations()
+        if relations is not None:
+            for relation in relations.relations:
+                self.bm.relations.append(relation)
+
+        fields = await resource.get_fields(force=True)
+        for (type_id, field_id), field in fields.items():
+            # Value
+            await self.generate_field(type_id, field_id, field)
+
+            # Extracted text
+            await self.generate_extracted_text(type_id, field_id, field)
+
+            # Field Computed Metadata
+            await self.generate_field_computed_metadata(type_id, field_id, field)
+
+            if type_id == FieldType.FILE and isinstance(field, File):
+                field_extracted_data = await field.get_file_extracted_data()
+                if field_extracted_data is not None:
+                    self.bm.file_extracted_data.append(field_extracted_data)
+
+            elif type_id == FieldType.LINK and isinstance(field, Link):
+                link_extracted_data = await field.get_link_extracted_data()
+                if link_extracted_data is not None:
+                    self.bm.link_extracted_data.append(link_extracted_data)
+
+            # Field vectors
+            await self.generate_field_vectors(type_id, field_id, field)
+
+            # Large metadata
+            await self.generate_field_large_computed_metadata(type_id, field_id, field)
+
+        return self.bm
+
+    async def generate_field(
+        self,
+        type_id: FieldType.ValueType,
+        field_id: str,
+        field: Field,
+    ):
+        # Used for exporting a field
+        if type_id == FieldType.TEXT:
+            value = await field.get_value()
+            self.bm.texts[field_id].CopyFrom(value)
+        elif type_id == FieldType.LINK:
+            value = await field.get_value()
+            self.bm.links[field_id].CopyFrom(value)
+        elif type_id == FieldType.FILE:
+            value = await field.get_value()
+            self.bm.files[field_id].CopyFrom(value)
+        elif type_id == FieldType.CONVERSATION:
+            field = cast(Conversation, field)
+            value = await field.get_full_conversation()
+            self.bm.conversations[field_id].CopyFrom(value)
+
+    async def generate_extracted_text(
+        self,
+        type_id: FieldType.ValueType,
+        field_id: str,
+        field: Field,
+    ):
+        etw = ExtractedTextWrapper()
+        etw.field.field = field_id
+        etw.field.field_type = type_id
+        extracted_text = await field.get_extracted_text()
+        if extracted_text is not None:
+            etw.body.CopyFrom(extracted_text)
+            self.bm.extracted_text.append(etw)
+
+    async def generate_field_computed_metadata(
+        self,
+        type_id: FieldType.ValueType,
+        field_id: str,
+        field: Field,
+    ):
+        fcmw = FieldComputedMetadataWrapper()
+        fcmw.field.field = field_id
+        fcmw.field.field_type = type_id
+
+        field_metadata = await field.get_field_metadata()
+        if field_metadata is not None:
+            fcmw.metadata.CopyFrom(field_metadata)
+            fcmw.field.field = field_id
+            fcmw.field.field_type = type_id
+            self.bm.field_metadata.append(fcmw)
+            # Make sure cloud files are removed for exporting
+
+    async def generate_field_vectors(
+        self,
+        type_id: FieldType.ValueType,
+        field_id: str,
+        field: Field,
+    ):
+        vo = await field.get_vectors()
+        if vo is None:
+            return
+        evw = ExtractedVectorsWrapper()
+        evw.field.field = field_id
+        evw.field.field_type = type_id
+        evw.vectors.CopyFrom(vo)
+        self.bm.field_vectors.append(evw)
+
+    async def generate_field_large_computed_metadata(
+        self,
+        type_id: FieldType.ValueType,
+        field_id: str,
+        field: Field,
+    ):
+        lcm = await field.get_large_field_metadata()
+        if lcm is None:
+            return
+        lcmw = LargeComputedMetadataWrapper()
+        lcmw.field.field = field_id
+        lcmw.field.field_type = type_id
+        lcmw.real.CopyFrom(lcm)
+        self.bm.field_large_metadata.append(lcmw)

nucliadb/ingest/orm/entities.py

@@ -21,24 +21,6 @@
 import asyncio
 from typing import AsyncGenerator, Optional
 
-from nucliadb_protos.knowledgebox_pb2 import (
-    DeletedEntitiesGroups,
-    EntitiesGroup,
-    EntitiesGroupSummary,
-    Entity,
-)
-from nucliadb_protos.nodereader_pb2 import (
-    Faceted,
-    RelationNodeFilter,
-    RelationPrefixSearchRequest,
-    RelationSearchRequest,
-    RelationSearchResponse,
-    SearchRequest,
-    SearchResponse,
-)
-from nucliadb_protos.utils_pb2 import RelationNode
-from nucliadb_protos.writer_pb2 import GetEntitiesResponse
-
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.exceptions import (
@@ -55,7 +37,25 @@ from nucliadb.common.datamanagers.entities import (
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.settings import settings
+from nucliadb_protos.knowledgebox_pb2 import (
+    DeletedEntitiesGroups,
+    EntitiesGroup,
+    EntitiesGroupSummary,
+    Entity,
+)
+from nucliadb_protos.nodereader_pb2 import (
+    Faceted,
+    RelationNodeFilter,
+    RelationPrefixSearchRequest,
+    RelationSearchResponse,
+    SearchRequest,
+    SearchResponse,
+)
+from nucliadb_protos.utils_pb2 import RelationNode
+from nucliadb_protos.writer_pb2 import GetEntitiesResponse
 from nucliadb_telemetry import errors
+from nucliadb_utils import const
+from nucliadb_utils.utilities import has_feature
 
 from .exceptions import EntityManagementException
 
@@ -199,37 +199,33 @@ class EntitiesManager:
         elif stored is not None and indexed is not None:
             entities_group = self.merge_entities_groups(indexed, stored)
         else:
-            entities_group = stored or indexed  # type: ignore
+            entities_group = stored or indexed
         return entities_group
 
     async def get_stored_entities_group(self, group: str) -> Optional[EntitiesGroup]:
-        return await datamanagers.entities.get_entities_group(
-            self.txn, kbid=self.kbid, group=group
-        )
+        return await datamanagers.entities.get_entities_group(self.txn, kbid=self.kbid, group=group)
 
     async def get_indexed_entities_group(self, group: str) -> Optional[EntitiesGroup]:
         shard_manager = get_shard_manager()
 
-        async def do_entities_search(
-            node: AbstractIndexNode, shard_id: str
-        ) -> RelationSearchResponse:
-            request = RelationSearchRequest(
-                shard_id=shard_id,
-                prefix=RelationPrefixSearchRequest(
+        async def do_entities_search(node: AbstractIndexNode, shard_id: str) -> RelationSearchResponse:
+            request = SearchRequest(
+                shard=shard_id,
+                relation_prefix=RelationPrefixSearchRequest(
                     prefix="",
                     node_filters=[
-                        RelationNodeFilter(
-                            node_type=RelationNode.NodeType.ENTITY, node_subtype=group
-                        )
+                        RelationNodeFilter(node_type=RelationNode.NodeType.ENTITY, node_subtype=group)
                     ],
                 ),
             )
-            return await node.reader.RelationSearch(request)  # type: ignore
+            response = await node.reader.Search(request)  # type: ignore
+            return response.relation
 
         results = await shard_manager.apply_for_all_shards(
             self.kbid,
             do_entities_search,
             settings.relation_search_timeout,
+            use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": self.kbid}),
             use_read_replica_nodes=self.use_read_replica_nodes,
         )
         for result in results:
@@ -239,9 +235,7 @@ class EntitiesManager:
 
         entities = {}
         for result in results:
-            entities.update(
-                {node.value: Entity(value=node.value) for node in result.prefix.nodes}
-            )
+            entities.update({node.value: Entity(value=node.value) for node in result.prefix.nodes})
 
         if not entities:
             return None
@@ -292,7 +286,7 @@ class EntitiesManager:
 
         # stored groups
         entities_key = KB_ENTITIES.format(kbid=self.kbid)
-        async for key in self.txn.keys(entities_key, count=-1):
+        async for key in self.txn.keys(entities_key):
             group = key.split("/")[-1]
             if exclude_deleted and group in deleted_groups:
                 continue
@@ -312,9 +306,7 @@ class EntitiesManager:
     ) -> set[str]:
         shard_manager = get_shard_manager()
 
-        async def query_indexed_entities_group_names(
-            node: AbstractIndexNode, shard_id: str
-        ) -> set[str]:
+        async def query_indexed_entities_group_names(node: AbstractIndexNode, shard_id: str) -> set[str]:
             request = SearchRequest(
                 shard=shard_id,
                 result_per_page=0,
@@ -335,6 +327,7 @@ class EntitiesManager:
             self.kbid,
             query_indexed_entities_group_names,
             settings.relation_types_timeout,
+            use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": self.kbid}),
             use_read_replica_nodes=self.use_read_replica_nodes,
         )
         for result in results:
@@ -347,9 +340,7 @@ class EntitiesManager:
         return set.union(*results)
 
     async def store_entities_group(self, group: str, eg: EntitiesGroup):
-        meta_cache = await datamanagers.entities.get_entities_meta_cache(
-            self.txn, kbid=self.kbid
-        )
+        meta_cache = await datamanagers.entities.get_entities_meta_cache(self.txn, kbid=self.kbid)
         duplicates = {}
         deleted = []
         duplicate_count = 0
@@ -373,9 +364,7 @@ class EntitiesManager:
 
         meta_cache.set_duplicates(group, duplicates)
         meta_cache.set_deleted(group, deleted)
-        await datamanagers.entities.set_entities_meta_cache(
-            self.txn, kbid=self.kbid, cache=meta_cache
-        )
+        await datamanagers.entities.set_entities_meta_cache(self.txn, kbid=self.kbid, cache=meta_cache)
 
         await datamanagers.entities.set_entities_group(
             self.txn, kbid=self.kbid, group_id=group, entities=eg
@@ -392,14 +381,10 @@ class EntitiesManager:
         await self.txn.delete(entities_key)
 
     async def mark_entities_group_as_deleted(self, group: str):
-        await datamanagers.entities.mark_group_as_deleted(
-            self.txn, kbid=self.kbid, group=group
-        )
+        await datamanagers.entities.mark_group_as_deleted(self.txn, kbid=self.kbid, group=group)
 
     async def unmark_entities_group_as_deleted(self, group: str):
-        await datamanagers.entities.unmark_group_as_deleted(
-            self.txn, kbid=self.kbid, group=group
-        )
+        await datamanagers.entities.unmark_group_as_deleted(self.txn, kbid=self.kbid, group=group)
 
     @staticmethod
     def merge_entities_groups(indexed: EntitiesGroup, stored: EntitiesGroup):

nucliadb/ingest/orm/exceptions.py

@@ -23,6 +23,10 @@ class NotFound(Exception):
     pass
 
 
+class KnowledgeBoxCreationError(Exception):
+    pass
+
+
 class KnowledgeBoxConflict(Exception):
     pass
 
@@ -48,3 +52,11 @@ class ResourceNotIndexable(Exception):
 
 class EntityManagementException(Exception):
     pass
+
+
+class VectorSetConflict(Exception):
+    pass
+
+
+class InvalidBrokerMessage(ValueError):
+    pass