nucliadb 6.3.5.post3985__py3-none-any.whl → 6.3.5.post3995__py3-none-any.whl

nucliadb/ingest/orm/index_message.py

@@ -0,0 +1,409 @@
+# 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/>.
+#
+
+
+import asyncio
+from typing import Optional
+
+from nucliadb.common import datamanagers
+from nucliadb.ingest.fields.exceptions import FieldAuthorNotFound
+from nucliadb.ingest.fields.file import File
+from nucliadb.ingest.orm.brain_v2 import ResourceBrainV2 as ResourceBrain
+from nucliadb.ingest.orm.resource import Resource, get_file_page_positions
+from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
+from nucliadb_protos.noderesources_pb2 import Resource as IndexMessage
+from nucliadb_protos.resources_pb2 import Basic, FieldID, FieldType
+from nucliadb_protos.writer_pb2 import BrokerMessage
+from nucliadb_utils import const
+from nucliadb_utils.utilities import has_feature
+
+
+class IndexMessageBuilder:
+    def __init__(self, resource: Resource):
+        self.resource = resource
+        self.brain = ResourceBrain(resource.uuid)
+
+    async def _apply_resource_index_data(self, brain: ResourceBrain) -> None:
+        # Set the metadata at the resource level
+        basic = await self.resource.get_basic()
+        assert basic is not None
+        user_relations = await self.resource.get_user_relations()
+        origin = await self.resource.get_origin()
+        security = await self.resource.get_security()
+        await asyncio.to_thread(
+            brain.generate_resource_indexing_metadata,
+            basic,
+            user_relations,
+            origin,
+            self.resource._previous_status,
+            security,
+        )
+
+    async def _apply_field_index_data(
+        self,
+        brain: ResourceBrain,
+        fieldid: FieldID,
+        basic: Basic,
+        texts: bool = True,
+        paragraphs: bool = True,
+        vectors: bool = True,
+        relations: bool = True,
+        replace: bool = True,
+        vectorset_configs: Optional[list[VectorSetConfig]] = None,
+    ):
+        field = await self.resource.get_field(fieldid.field, fieldid.field_type)
+        extracted_text = await field.get_extracted_text()
+        field_computed_metadata = await field.get_field_metadata()
+        user_field_metadata = next(
+            (fm for fm in basic.fieldmetadata if fm.field == fieldid),
+            None,
+        )
+        if texts or paragraphs:
+            # We need to compute the texts when we're going to generate the paragraphs too, but we may not
+            # want to index them always.
+            skip_index_texts = not texts
+            replace_texts = replace and not skip_index_texts
+
+            if extracted_text is not None:
+                try:
+                    field_author = await field.generated_by()
+                except FieldAuthorNotFound:
+                    field_author = None
+                await asyncio.to_thread(
+                    brain.generate_texts_index_message,
+                    self.resource.generate_field_id(fieldid),
+                    extracted_text,
+                    field_computed_metadata,
+                    basic.usermetadata,
+                    field_author,
+                    replace_field=replace_texts,
+                    skip_index=skip_index_texts,
+                )
+        if paragraphs or vectors:
+            # The paragraphs are needed to generate the vectors. However, we don't need to index them
+            # in all cases.
+            skip_index_paragraphs = not paragraphs
+            replace_paragraphs = replace and not skip_index_paragraphs
+
+            # We need to compute the paragraphs when we're going to generate the vectors too.
+            if extracted_text is not None and field_computed_metadata is not None:
+                page_positions = (
+                    await get_file_page_positions(field) if isinstance(field, File) else None
+                )
+                await asyncio.to_thread(
+                    brain.generate_paragraphs_index_message,
+                    self.resource.generate_field_id(fieldid),
+                    field_computed_metadata,
+                    extracted_text,
+                    page_positions,
+                    user_field_metadata,
+                    replace_field=replace_paragraphs,
+                    skip_index=skip_index_paragraphs,
+                )
+        if vectors:
+            assert vectorset_configs is not None
+            for vectorset_config in vectorset_configs:
+                vo = await field.get_vectors(
+                    vectorset=vectorset_config.vectorset_id,
+                    storage_key_kind=vectorset_config.storage_key_kind,
+                )
+                if vo is not None:
+                    dimension = vectorset_config.vectorset_index_config.vector_dimension
+                    await asyncio.to_thread(
+                        brain.generate_vectors_index_message,
+                        self.resource.generate_field_id(fieldid),
+                        vo,
+                        vectorset=vectorset_config.vectorset_id,
+                        replace_field=replace,
+                        vector_dimension=dimension,
+                    )
+        if relations:
+            await asyncio.to_thread(
+                brain.generate_relations_index_message,
+                self.resource.generate_field_id(fieldid),
+                field_computed_metadata,
+                basic.usermetadata,
+                replace_field=replace,
+            )
+
+    def _apply_field_deletions(
+        self,
+        brain: ResourceBrain,
+        field_ids: list[FieldID],
+    ) -> None:
+        for field_id in field_ids:
+            brain.delete_field(self.resource.generate_field_id(field_id))
+
+    async def for_writer_bm(
+        self,
+        messages: list[BrokerMessage],
+        resource_created: bool,
+    ) -> IndexMessage:
+        """
+        Builds the index message for the broker messages coming from the writer.
+        The writer messages are not adding new vectors to the index.
+        """
+        assert all(message.source == BrokerMessage.MessageSource.WRITER for message in messages)
+
+        deleted_fields = get_bm_deleted_fields(messages)
+        self._apply_field_deletions(self.brain, deleted_fields)
+        await self._apply_resource_index_data(self.brain)
+        basic = await self.get_basic()
+        prefilter_update = needs_prefilter_update(messages)
+        if prefilter_update:
+            # Changes on some metadata at the resource level that is used for filtering require that we reindex all the fields
+            # in the texts index (as it is the one used for prefiltering).
+            fields_to_index = [
+                FieldID(field=field_id, field_type=field_type)
+                for field_type, field_id in await self.resource.get_fields(force=True)
+            ]
+        else:
+            # Simply process the fields that are in the message
+            fields_to_index = get_bm_modified_fields(messages)
+        for fieldid in fields_to_index:
+            if fieldid in deleted_fields:
+                continue
+            await self._apply_field_index_data(
+                self.brain,
+                fieldid,
+                basic,
+                texts=prefilter_update or needs_texts_update(fieldid, messages),
+                paragraphs=needs_paragraphs_update(fieldid, messages),
+                relations=False,  # Relations at the field level are not modified by the writer
+                vectors=False,  # Vectors are never added by the writer
+                replace=not resource_created,
+            )
+        return self.brain.brain
+
+    async def for_processor_bm(
+        self,
+        messages: list[BrokerMessage],
+    ) -> IndexMessage:
+        """
+        Builds the index message for the broker messages coming from the processor.
+        The processor can index new data to any index.
+        """
+        assert all(message.source == BrokerMessage.MessageSource.PROCESSOR for message in messages)
+        deleted_fields = get_bm_deleted_fields(messages)
+        self._apply_field_deletions(self.brain, deleted_fields)
+        await self._apply_resource_index_data(self.brain)
+        basic = await self.get_basic()
+        fields_to_index = get_bm_modified_fields(messages)
+        vectorsets_configs = await self.get_vectorsets_configs()
+        for fieldid in fields_to_index:
+            if fieldid in deleted_fields:
+                continue
+            await self._apply_field_index_data(
+                self.brain,
+                fieldid,
+                basic,
+                texts=needs_texts_update(fieldid, messages),
+                paragraphs=needs_paragraphs_update(fieldid, messages),
+                relations=needs_relations_update(fieldid, messages),
+                vectors=needs_vectors_update(fieldid, messages),
+                replace=True,
+                vectorset_configs=vectorsets_configs,
+            )
+        return self.brain.brain
+
+    async def full(self, reindex: bool) -> IndexMessage:
+        await self._apply_resource_index_data(self.brain)
+        basic = await self.get_basic()
+        fields_to_index = [
+            FieldID(field=field_id, field_type=field_type)
+            for field_type, field_id in await self.resource.get_fields(force=True)
+        ]
+        vectorsets_configs = await self.get_vectorsets_configs()
+        for fieldid in fields_to_index:
+            await self._apply_field_index_data(
+                self.brain,
+                fieldid,
+                basic,
+                texts=True,
+                paragraphs=True,
+                relations=True,
+                vectors=True,
+                replace=reindex,
+                vectorset_configs=vectorsets_configs,
+            )
+        return self.brain.brain
+
+    async def get_basic(self) -> Basic:
+        basic = await self.resource.get_basic()
+        assert basic is not None
+        return basic
+
+    async def get_vectorsets_configs(self) -> list[VectorSetConfig]:
+        """
+        Get the vectorsets config for the resource.
+        """
+        vectorset_configs = [
+            vectorset_config
+            async for _, vectorset_config in datamanagers.vectorsets.iter(
+                self.resource.txn, kbid=self.resource.kb.kbid
+            )
+        ]
+        return vectorset_configs
+
+
+def get_bm_deleted_fields(
+    messages: list[BrokerMessage],
+) -> list[FieldID]:
+    deleted = []
+    for message in messages:
+        for field in message.delete_fields:
+            if field not in deleted:
+                deleted.append(field)
+    return deleted
+
+
+def get_bm_modified_fields(messages: list[BrokerMessage]) -> list[FieldID]:
+    message_source = get_messages_source(messages)
+    modified = set()
+    for message in messages:
+        # Added or modified fields need indexing
+        for link in message.links:
+            modified.add((link, FieldType.LINK))
+        for file in message.files:
+            modified.add((file, FieldType.FILE))
+        for conv in message.conversations:
+            modified.add((conv, FieldType.CONVERSATION))
+        for text in message.texts:
+            modified.add((text, FieldType.TEXT))
+        if message.HasField("basic"):
+            # Add title and summary only if they have changed
+            if message.basic.title != "":
+                modified.add(("title", FieldType.GENERIC))
+            if message.basic.summary != "":
+                modified.add(("summary", FieldType.GENERIC))
+
+        if message_source == BrokerMessage.MessageSource.PROCESSOR:
+            # Messages with field metadata, extracted text or field vectors need indexing
+            for fm in message.field_metadata:
+                modified.add((fm.field.field, fm.field.field_type))
+            for et in message.extracted_text:
+                modified.add((et.field.field, et.field.field_type))
+            for fv in message.field_vectors:
+                modified.add((fv.field.field, fv.field.field_type))
+
+        if message_source == BrokerMessage.MessageSource.WRITER:
+            # Any field that has fieldmetadata annotations should be considered as modified
+            # and needs to be reindexed
+            if message.HasField("basic"):
+                for ufm in message.basic.fieldmetadata:
+                    modified.add((ufm.field.field, ufm.field.field_type))
+    return [FieldID(field=field, field_type=field_type) for field, field_type in modified]
+
+
+def get_messages_source(messages: list[BrokerMessage]) -> BrokerMessage.MessageSource.ValueType:
+    assert len(set(message.source for message in messages)) == 1
+    return messages[0].source
+
+
+def needs_prefilter_update(messages: list[BrokerMessage]) -> bool:
+    return any(message.reindex for message in messages)
+
+
+def needs_paragraphs_update(field_id: FieldID, messages: list[BrokerMessage]) -> bool:
+    return (
+        has_paragraph_annotations(field_id, messages)
+        or has_new_extracted_text(field_id, messages)
+        or has_new_field_metadata(field_id, messages)
+    )
+
+
+def has_paragraph_annotations(field_id: FieldID, messages: list[BrokerMessage]) -> bool:
+    for message in messages:
+        ufm = next(
+            (fm for fm in message.basic.fieldmetadata if fm.field == field_id),
+            None,
+        )
+        if ufm is None:
+            continue
+        if len(ufm.paragraphs) > 0:
+            return True
+    return False
+
+
+def has_new_field_metadata(
+    field_id: FieldID,
+    messages: list[BrokerMessage],
+) -> bool:
+    for message in messages:
+        for field_metadata in message.field_metadata:
+            if field_metadata.field == field_id:
+                return True
+    return False
+
+
+def has_new_extracted_text(
+    field_id: FieldID,
+    messages: list[BrokerMessage],
+) -> bool:
+    for message in messages:
+        for extracted_text in message.extracted_text:
+            if extracted_text.field == field_id:
+                return True
+    return False
+
+
+def needs_texts_update(
+    field_id: FieldID,
+    messages: list[BrokerMessage],
+) -> bool:
+    return has_new_extracted_text(field_id, messages) or has_new_field_metadata(field_id, messages)
+
+
+def needs_vectors_update(
+    field_id: FieldID,
+    messages: list[BrokerMessage],
+) -> bool:
+    for message in messages:
+        for field_vectors in message.field_vectors:
+            if field_vectors.field == field_id:
+                return True
+    return False
+
+
+def needs_relations_update(
+    field_id: FieldID,
+    messages: list[BrokerMessage],
+) -> bool:
+    return has_new_field_metadata(field_id, messages) or has_new_extracted_text(field_id, messages)
+
+
+async def get_resource_index_message(
+    resource: Resource,
+    reindex: bool = False,
+) -> IndexMessage:
+    """
+    Get the full index message for a resource.
+    """
+    if has_feature(
+        const.Features.INDEX_MESSAGE_GENERATION_V2,
+        context={
+            "kbid": resource.kb.kbid,
+        },
+    ):
+        im_builder = IndexMessageBuilder(resource)
+        return await im_builder.full(reindex=reindex)
+    else:
+        # TODO: remove this code when we remove the old index message generation
+        return (await resource.generate_index_message(reindex=reindex)).brain

nucliadb/ingest/orm/metrics.py

@@ -22,6 +22,6 @@ from nucliadb_telemetry import metrics
 
 processor_observer = metrics.Observer(
     "nucliadb_ingest_processor",
-    labels={"type": ""},
+    labels={"type": "", "source": ""},
     error_mappings={"kb_conflict": KnowledgeBoxConflict},
 )

nucliadb/ingest/orm/processor/data_augmentation.py

@@ -23,8 +23,8 @@ from dataclasses import dataclass, field
 from typing import Optional
 
 from nucliadb.ingest.orm.resource import Resource
-from nucliadb.ingest.processing import ProcessingEngine, PushPayload, Source
-from nucliadb_models.text import PushTextFormat, Text
+from nucliadb.ingest.processing import ProcessingEngine
+from nucliadb.models.internal.processing import PushPayload, PushTextFormat, Source, Text
 from nucliadb_protos import resources_pb2, writer_pb2
 from nucliadb_protos.resources_pb2 import FieldType
 from nucliadb_utils.utilities import Utility, get_partitioning, get_utility

nucliadb/ingest/orm/processor/pgcatalog.py

@@ -23,6 +23,7 @@ from typing import cast
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.common.maindb.pg import PGDriver, PGTransaction
 from nucliadb.common.maindb.utils import get_driver
+from nucliadb_protos.noderesources_pb2 import Resource as IndexMessage
 from nucliadb_telemetry import metrics
 
 from ..resource import Resource
@@ -39,7 +40,7 @@ def pgcatalog_enabled(kbid):
 
 
 @observer.wrap({"type": "update"})
-async def pgcatalog_update(txn: Transaction, kbid: str, resource: Resource):
+async def pgcatalog_update(txn: Transaction, kbid: str, resource: Resource, index_message: IndexMessage):
     if not pgcatalog_enabled(kbid):
         return
 
@@ -69,7 +70,7 @@ async def pgcatalog_update(txn: Transaction, kbid: str, resource: Resource):
                 "title": resource.basic.title,
                 "created_at": created_at,
                 "modified_at": modified_at,
-                "labels": list(resource.indexer.brain.labels),
+                "labels": list(index_message.labels),
             },
         )
 

nucliadb/ingest/orm/processor/processor.py

@@ -38,6 +38,7 @@ from nucliadb.ingest.orm.exceptions import (
     ResourceNotIndexable,
     SequenceOrderViolation,
 )
+from nucliadb.ingest.orm.index_message import IndexMessageBuilder
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.orm.metrics import processor_observer
 from nucliadb.ingest.orm.processor import sequence_manager
@@ -312,25 +313,11 @@ class Processor:
                     await self.apply_resource(message, resource, update=(not created))
 
                 # index message
-
-                if resource:
-                    if any(needs_reindex(m) for m in messages):
-                        # when reindexing, let's just generate full new index message
-                        # TODO - This should be improved in the future as it's not optimal for very large resources:
-                        # As of now, there are some API operations that require fully reindexing all the fields of a resource.
-                        # An example of this is classification label changes - we need to reindex all the fields of a resource to
-                        # propagate the label changes to the index.
-                        resource.replace_indexer(await resource.generate_index_message(reindex=True))
-                    else:
-                        # TODO - Ideally we should only update the fields that have been changed in the current transaction.
-                        await resource.compute_global_text()
-                        await resource.compute_global_tags(resource.indexer)
-                        await resource.compute_security(resource.indexer)
-
                 if resource and resource.modified:
-                    await pgcatalog_update(txn, kbid, resource)
+                    index_message = await self.generate_index_message(resource, messages, created)
+                    await pgcatalog_update(txn, kbid, resource, index_message)
                     await self.index_resource(  # noqa
-                        resource=resource,
+                        index_message=index_message,
                         txn=txn,
                         uuid=uuid,
                         kbid=kbid,
@@ -451,7 +438,7 @@ class Processor:
     @processor_observer.wrap({"type": "index_resource"})
     async def index_resource(
         self,
-        resource: Resource,
+        index_message: PBBrainResource,
         txn: Transaction,
         uuid: str,
         kbid: str,
@@ -460,9 +447,8 @@ class Processor:
         kb: KnowledgeBox,
         source: nodewriter_pb2.IndexMessageSource.ValueType,
     ) -> None:
-        validate_indexable_resource(resource.indexer.brain)
+        validate_indexable_resource(index_message)
         shard = await self.get_or_assign_resource_shard(txn, kb, uuid)
-        index_message = resource.indexer.brain
         external_index_manager = await get_external_index_manager(kbid=kbid)
         if external_index_manager is not None:
             await self.external_index_add_resource(external_index_manager, uuid, index_message)
@@ -476,6 +462,56 @@ class Processor:
                 source=source,
             )
 
+    async def generate_index_message_v2(
+        self,
+        resource: Resource,
+        messages: list[writer_pb2.BrokerMessage],
+        resource_created: bool,
+    ) -> PBBrainResource:
+        builder = IndexMessageBuilder(resource)
+        message_source = messages_source(messages)
+        if message_source == nodewriter_pb2.IndexMessageSource.WRITER:
+            with processor_observer({"type": "generate_index_message", "source": "writer"}):
+                return await builder.for_writer_bm(messages, resource_created)
+        elif message_source == nodewriter_pb2.IndexMessageSource.PROCESSOR:
+            with processor_observer({"type": "generate_index_message", "source": "processor"}):
+                return await builder.for_processor_bm(messages)
+        else:  # pragma: no cover
+            raise InvalidBrokerMessage(f"Unknown broker message source: {message_source}")
+
+    async def generate_index_message_v1(
+        self,
+        resource: Resource,
+        messages: list[writer_pb2.BrokerMessage],
+    ) -> PBBrainResource:
+        if any(needs_reindex(m) for m in messages):
+            # when reindexing, let's just generate full new index message
+            # TODO - This should be improved in the future as it's not optimal for very large resources:
+            # As of now, there are some API operations that require fully reindexing all the fields of a resource.
+            # An example of this is classification label changes - we need to reindex all the fields of a resource to
+            # propagate the label changes to the index.
+            resource.replace_indexer(await resource.generate_index_message(reindex=True))
+        else:
+            # TODO - Ideally we should only update the fields that have been changed in the current transaction.
+            await resource.compute_global_text()
+            await resource.compute_global_tags(resource.indexer)
+            await resource.compute_security(resource.indexer)
+        return resource.indexer.brain
+
+    async def generate_index_message(
+        self,
+        resource: Resource,
+        messages: list[writer_pb2.BrokerMessage],
+        resource_created: bool = False,
+    ) -> PBBrainResource:
+        if has_feature(
+            const.Features.INDEX_MESSAGE_GENERATION_V2,
+            context={"kbid": resource.kb.kbid},
+        ):
+            return await self.generate_index_message_v2(resource, messages, resource_created)
+        else:
+            return await self.generate_index_message_v1(resource, messages)
+
     async def external_index_delete_resource(
         self, external_index_manager: ExternalIndexManager, resource_uuid: str
     ):
@@ -564,7 +600,10 @@ class Processor:
         resource: Resource,
         update: bool = False,
     ):
-        """Apply broker message to resource object in the database"""
+        """
+        Apply broker message to resource object in the persistence layers (maindb and storage).
+        DO NOT add any indexing logic here.
+        """
         if update:
             await self.maybe_update_resource_basic(resource, message)
 
@@ -675,30 +714,9 @@ class Processor:
         try:
             async with self.driver.transaction() as txn:
                 kb.txn = resource.txn = txn
-
-                shard_id = await datamanagers.resources.get_resource_shard_id(
-                    txn, kbid=kb.kbid, rid=resource.uuid
-                )
-                shard = None
-                if shard_id is not None:
-                    shard = await kb.get_resource_shard(shard_id)
-                if shard is None:
-                    logger.warning(
-                        "Unable to mark resource as error, shard is None. "
-                        "This should not happen so you did something special to get here."
-                    )
-                    return
-
                 resource.basic.metadata.status = resources_pb2.Metadata.Status.ERROR
                 await resource.set_basic(resource.basic)
                 await txn.commit()
-
-            resource.indexer.set_processing_status(
-                basic=resource.basic, previous_status=resource._previous_status
-            )
-            await self.index_node_shard_manager.add_resource(
-                shard, resource.indexer.brain, seqid, partition=partition, kb=kb.kbid
-            )
         except Exception:
             logger.warning("Error while marking resource as error", exc_info=True)
 
@@ -745,11 +763,7 @@ def has_vectors_operation(index_message: PBBrainResource) -> bool:
     """
     Returns True if the index message has any vectors to index or to delete.
     """
-    if (
-        len(index_message.sentences_to_delete) > 0
-        or len(index_message.paragraphs_to_delete) > 0
-        or any([len(deletions.items) for deletions in index_message.vector_prefixes_to_delete.values()])
-    ):
+    if any([len(deletions.items) for deletions in index_message.vector_prefixes_to_delete.values()]):
         return True
     for field_paragraphs in index_message.paragraphs.values():
         for paragraph in field_paragraphs.paragraphs.values():