nucliadb 6.7.2.post4862__py3-none-any.whl → 6.9.2.post5282__py3-none-any.whl

nucliadb/ingest/orm/index_message.py

@@ -20,11 +20,12 @@
 
 
 import asyncio
-from typing import Optional
+from typing import Optional, Sequence
 
 from nidx_protos.noderesources_pb2 import Resource as IndexMessage
 
 from nucliadb.common import datamanagers
+from nucliadb.ingest.fields.conversation import Conversation
 from nucliadb.ingest.fields.exceptions import FieldAuthorNotFound
 from nucliadb.ingest.fields.file import File
 from nucliadb.ingest.orm.brain_v2 import ResourceBrain
@@ -32,6 +33,7 @@ from nucliadb.ingest.orm.metrics import index_message_observer as observer
 from nucliadb.ingest.orm.resource import Resource, get_file_page_positions
 from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import Basic, FieldID, FieldType
+from nucliadb_protos.utils_pb2 import ExtractedText
 from nucliadb_protos.writer_pb2 import BrokerMessage
 
 
@@ -69,6 +71,7 @@ class IndexMessageBuilder:
         relations: bool = True,
         replace: bool = True,
         vectorset_configs: Optional[list[VectorSetConfig]] = None,
+        append_splits: Optional[set[str]] = None,
     ):
         field = await self.resource.get_field(fieldid.field, fieldid.field_type)
         extracted_text = await field.get_extracted_text()
@@ -120,6 +123,7 @@ class IndexMessageBuilder:
                     replace_field=replace_paragraphs,
                     skip_paragraphs_index=skip_paragraphs_index,
                     skip_texts_index=skip_texts_index,
+                    append_splits=append_splits,
                 )
         if vectors:
             assert vectorset_configs is not None
@@ -137,6 +141,7 @@ class IndexMessageBuilder:
                         vectorset=vectorset_config.vectorset_id,
                         replace_field=replace,
                         vector_dimension=dimension,
+                        append_splits=append_splits,
                     )
         if relations:
             await asyncio.to_thread(
@@ -150,7 +155,7 @@ class IndexMessageBuilder:
     def _apply_field_deletions(
         self,
         brain: ResourceBrain,
-        field_ids: list[FieldID],
+        field_ids: Sequence[FieldID],
     ) -> None:
         for field_id in field_ids:
             brain.delete_field(self.resource.generate_field_id(field_id))
@@ -158,20 +163,19 @@ class IndexMessageBuilder:
     @observer.wrap({"type": "writer_bm"})
     async def for_writer_bm(
         self,
-        messages: list[BrokerMessage],
+        message: BrokerMessage,
         resource_created: bool,
     ) -> IndexMessage:
         """
-        Builds the index message for the broker messages coming from the writer.
+        Builds the index message for the broker message 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)
+        assert message.source == BrokerMessage.MessageSource.WRITER
 
-        deleted_fields = get_bm_deleted_fields(messages)
-        self._apply_field_deletions(self.brain, deleted_fields)
+        self._apply_field_deletions(self.brain, message.delete_fields)
         await self._apply_resource_index_data(self.brain)
         basic = await self.get_basic()
-        prefilter_update = needs_prefilter_update(messages)
+        prefilter_update = needs_prefilter_update(message)
         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).
@@ -181,16 +185,16 @@ class IndexMessageBuilder:
             ]
         else:
             # Simply process the fields that are in the message
-            fields_to_index = get_bm_modified_fields(messages)
+            fields_to_index = get_bm_modified_fields(message)
         for fieldid in fields_to_index:
-            if fieldid in deleted_fields:
+            if fieldid in message.delete_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),
+                texts=prefilter_update or needs_texts_update(fieldid, message),
+                paragraphs=needs_paragraphs_update(fieldid, message),
                 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,
@@ -200,32 +204,45 @@ class IndexMessageBuilder:
     @observer.wrap({"type": "processor_bm"})
     async def for_processor_bm(
         self,
-        messages: list[BrokerMessage],
+        message: 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)
+        assert message.source == BrokerMessage.MessageSource.PROCESSOR
+        self._apply_field_deletions(self.brain, message.delete_fields)
         await self._apply_resource_index_data(self.brain)
         basic = await self.get_basic()
-        fields_to_index = get_bm_modified_fields(messages)
+        fields_to_index = get_bm_modified_fields(message)
         vectorsets_configs = await self.get_vectorsets_configs()
         for fieldid in fields_to_index:
-            if fieldid in deleted_fields:
+            if fieldid in message.delete_fields:
                 continue
+
+            # For conversation fields, we only replace the full field if it is not an append messages operation.
+            # All other fields are always replaced upon modification.
+            replace_field = True
+            modified_splits = None
+            if fieldid.field_type == FieldType.CONVERSATION:
+                modified_splits = await get_bm_modified_split_ids(fieldid, message, self.resource)
+                stored_splits = await get_stored_split_ids(fieldid, self.resource)
+                is_append_messages_op = modified_splits.issubset(stored_splits) and 0 < len(
+                    modified_splits
+                ) < len(stored_splits)
+                replace_field = not is_append_messages_op
+
             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,
+                texts=needs_texts_update(fieldid, message),
+                paragraphs=needs_paragraphs_update(fieldid, message),
+                relations=needs_relations_update(fieldid, message),
+                vectors=needs_vectors_update(fieldid, message),
+                replace=replace_field,
                 vectorset_configs=vectorsets_configs,
+                append_splits=modified_splits,
             )
         return self.brain.brain
 
@@ -270,130 +287,128 @@ class IndexMessageBuilder:
         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)
+def get_bm_modified_fields(message: BrokerMessage) -> list[FieldID]:
     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))
+    # 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"):
-            # 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))
+            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(message: BrokerMessage) -> bool:
+    return message.reindex
 
 
-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:
+def needs_paragraphs_update(field_id: FieldID, message: 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)
+        has_paragraph_annotations(field_id, message)
+        or has_new_extracted_text(field_id, message)
+        or has_new_field_metadata(field_id, message)
     )
 
 
-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_paragraph_annotations(field_id: FieldID, message: BrokerMessage) -> bool:
+    ufm = next(
+        (fm for fm in message.basic.fieldmetadata if fm.field == field_id),
+        None,
+    )
+    if ufm is None:
+        return False
+    return len(ufm.paragraphs) > 0
 
 
 def has_new_field_metadata(
     field_id: FieldID,
-    messages: list[BrokerMessage],
+    message: BrokerMessage,
 ) -> bool:
-    for message in messages:
-        for field_metadata in message.field_metadata:
-            if field_metadata.field == field_id:
-                return True
-    return False
+    return any(field_metadata.field == field_id for field_metadata in message.field_metadata)
 
 
 def has_new_extracted_text(
     field_id: FieldID,
-    messages: list[BrokerMessage],
+    message: BrokerMessage,
 ) -> bool:
-    for message in messages:
-        for extracted_text in message.extracted_text:
-            if extracted_text.field == field_id:
-                return True
-    return False
+    return any(extracted_text.field == field_id for extracted_text in message.extracted_text)
 
 
 def needs_texts_update(
     field_id: FieldID,
-    messages: list[BrokerMessage],
+    message: BrokerMessage,
 ) -> bool:
-    return has_new_extracted_text(field_id, messages) or has_new_field_metadata(field_id, messages)
+    return has_new_extracted_text(field_id, message) or has_new_field_metadata(field_id, message)
 
 
 def needs_vectors_update(
     field_id: FieldID,
-    messages: list[BrokerMessage],
+    message: BrokerMessage,
 ) -> bool:
-    for message in messages:
-        for field_vectors in message.field_vectors:
-            if field_vectors.field == field_id:
-                return True
-    return False
+    return any(field_vectors.field == field_id for field_vectors in message.field_vectors)
+
+
+async def get_bm_modified_split_ids(
+    conversation_field_id: FieldID,
+    message: BrokerMessage,
+    resource: Resource,
+) -> set[str]:
+    message_etw = next(
+        (etw for etw in message.extracted_text if etw.field == conversation_field_id), None
+    )
+    if message_etw is None:
+        return set()
+    storage = resource.storage
+    if message_etw.HasField("file"):
+        raw_payload = await storage.downloadbytescf(message_etw.file)
+        message_extracted_text = ExtractedText()
+        message_extracted_text.ParseFromString(raw_payload.read())
+        raw_payload.flush()
+    else:
+        message_extracted_text = message_etw.body
+    return set(message_extracted_text.split_text.keys())
+
+
+async def get_stored_split_ids(
+    conversation_field_id: FieldID,
+    resource: Resource,
+) -> set[str]:
+    fid = conversation_field_id
+    conv: Conversation = await resource.get_field(fid.field, fid.field_type, load=False)
+    splits_metadata = await conv.get_splits_metadata()
+    return set(splits_metadata.metadata)
 
 
 def needs_relations_update(
     field_id: FieldID,
-    messages: list[BrokerMessage],
+    message: BrokerMessage,
 ) -> bool:
-    return has_new_field_metadata(field_id, messages) or has_new_extracted_text(field_id, messages)
+    return has_new_field_metadata(field_id, message) or has_new_extracted_text(field_id, message)
 
 
 async def get_resource_index_message(

nucliadb/ingest/orm/knowledgebox.py

@@ -24,7 +24,7 @@ from uuid import uuid4
 
 from grpc import StatusCode
 from grpc.aio import AioRpcError
-from nidx_protos import noderesources_pb2
+from nidx_protos import nidx_pb2, noderesources_pb2
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardNotFound
@@ -36,7 +36,6 @@ from nucliadb.common.datamanagers.resources import (
     KB_RESOURCE_SLUG_BASE,
 )
 from nucliadb.common.external_index_providers.base import VectorsetExternalIndex
-from nucliadb.common.external_index_providers.pinecone import PineconeIndexManager
 from nucliadb.common.maindb.driver import Driver, Transaction
 from nucliadb.common.maindb.pg import PGTransaction
 from nucliadb.common.nidx import get_nidx_api_client
@@ -53,7 +52,6 @@ from nucliadb.migrator.utils import get_latest_version
 from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_protos.knowledgebox_pb2 import (
     CreateExternalIndexProviderMetadata,
-    ExternalIndexProviderType,
     KnowledgeBoxConfig,
     SemanticModelMetadata,
     StoredExternalIndexProviderMetadata,
@@ -110,6 +108,7 @@ class KnowledgeBox:
         external_index_provider: CreateExternalIndexProviderMetadata = CreateExternalIndexProviderMetadata(),
         hidden_resources_enabled: bool = False,
         hidden_resources_hide_on_creation: bool = False,
+        prewarm_enabled: bool = False,
     ) -> tuple[str, str]:
         """Creates a new knowledge box and return its id and slug."""
 
@@ -127,7 +126,7 @@ class KnowledgeBox:
         rollback_ops: list[Callable[[], Coroutine[Any, Any, Any]]] = []
 
         try:
-            async with driver.transaction() as txn:
+            async with driver.rw_transaction() as txn:
                 exists = await datamanagers.kb.get_kb_uuid(
                     txn, slug=slug
                 ) or await datamanagers.kb.exists_kb(txn, kbid=kbid)
@@ -196,6 +195,7 @@ class KnowledgeBox:
                     migration_version=get_latest_version(),
                     hidden_resources_enabled=hidden_resources_enabled,
                     hidden_resources_hide_on_creation=hidden_resources_hide_on_creation,
+                    prewarm_enabled=prewarm_enabled,
                 )
                 config.external_index_provider.CopyFrom(stored_external_index_provider)
                 await datamanagers.kb.set_config(txn, kbid=kbid, config=config)
@@ -222,7 +222,7 @@ class KnowledgeBox:
                 shard_manager = get_shard_manager()
                 # XXX creating a shard is a slow IO operation that requires a write
                 # txn to be open!
-                await shard_manager.create_shard_by_kbid(txn, kbid)
+                await shard_manager.create_shard_by_kbid(txn, kbid, prewarm_enabled=prewarm_enabled)
                 # shards don't need a rollback as they will be eventually purged
 
                 await txn.commit()
@@ -245,43 +245,90 @@ class KnowledgeBox:
     @classmethod
     async def update(
         cls,
-        txn: Transaction,
-        uuid: str,
+        driver: Driver,
+        kbid: str,
+        *,
         slug: Optional[str] = None,
-        config: Optional[KnowledgeBoxConfig] = None,
+        title: Optional[str] = None,
+        description: Optional[str] = None,
+        migration_version: Optional[int] = None,
+        external_index_provider: Optional[StoredExternalIndexProviderMetadata] = None,
+        hidden_resources_enabled: Optional[bool] = None,
+        hidden_resources_hide_on_creation: Optional[bool] = None,
+        prewarm_enabled: Optional[bool] = None,
     ) -> str:
-        exist = await datamanagers.kb.get_config(txn, kbid=uuid, for_update=True)
-        if not exist:
-            raise datamanagers.exceptions.KnowledgeBoxNotFound()
-
-        if slug:
-            await txn.delete(datamanagers.kb.KB_SLUGS.format(slug=exist.slug))
-            await txn.set(
-                datamanagers.kb.KB_SLUGS.format(slug=slug),
-                uuid.encode(),
-            )
-            if config:
-                config.slug = slug
-            else:
-                exist.slug = slug
+        async with driver.rw_transaction() as txn:
+            stored = await datamanagers.kb.get_config(txn, kbid=kbid, for_update=True)
+            if not stored:
+                raise datamanagers.exceptions.KnowledgeBoxNotFound()
+
+            if slug:
+                await txn.delete(datamanagers.kb.KB_SLUGS.format(slug=stored.slug))
+                await txn.set(
+                    datamanagers.kb.KB_SLUGS.format(slug=slug),
+                    kbid.encode(),
+                )
+                stored.slug = slug
 
-        if config and exist != config:
-            exist.MergeFrom(config)
-            exist.hidden_resources_enabled = config.hidden_resources_enabled
-            exist.hidden_resources_hide_on_creation = config.hidden_resources_hide_on_creation
+            if title is not None:
+                stored.title = title
+            if description is not None:
+                stored.description = description
 
-        if exist.hidden_resources_hide_on_creation and not exist.hidden_resources_enabled:
-            raise KnowledgeBoxCreationError(
-                "Cannot hide new resources if the hidden resources feature is disabled"
-            )
+            if migration_version is not None:
+                stored.migration_version = migration_version
+
+            if external_index_provider is not None:
+                stored.external_index_provider.MergeFrom(external_index_provider)
+
+            if hidden_resources_enabled is not None:
+                stored.hidden_resources_enabled = hidden_resources_enabled
+            if hidden_resources_hide_on_creation is not None:
+                stored.hidden_resources_hide_on_creation = hidden_resources_hide_on_creation
+
+            update_nidx_prewarm = None
+            if prewarm_enabled is not None:
+                if stored.prewarm_enabled != prewarm_enabled:
+                    update_nidx_prewarm = prewarm_enabled
+                stored.prewarm_enabled = prewarm_enabled
+
+            if stored.hidden_resources_hide_on_creation and not stored.hidden_resources_enabled:
+                raise KnowledgeBoxCreationError(
+                    "Cannot hide new resources if the hidden resources feature is disabled"
+                )
 
-        await datamanagers.kb.set_config(txn, kbid=uuid, config=exist)
+            await datamanagers.kb.set_config(txn, kbid=kbid, config=stored)
 
-        return uuid
+            await txn.commit()
+
+        if update_nidx_prewarm is not None:
+            await cls.configure_shards(driver, kbid, prewarm=update_nidx_prewarm)
+
+        return kbid
+
+    @classmethod
+    async def configure_shards(cls, driver: Driver, kbid: str, *, prewarm: bool):
+        shards_obj = await datamanagers.atomic.cluster.get_kb_shards(kbid=kbid)
+        if shards_obj is None:
+            logger.warning(f"Shards not found for KB while updating pre-warm flag", extra={"kbid": kbid})
+            return
+
+        nidx_shard_ids = [shard.nidx_shard_id for shard in shards_obj.shards]
+
+        nidx_api = get_nidx_api_client()
+        if nidx_api is not None and len(nidx_shard_ids) > 0:
+            configs = [
+                nidx_pb2.ShardConfig(
+                    shard_id=shard_id,
+                    prewarm_enabled=prewarm,
+                )
+                for shard_id in nidx_shard_ids
+            ]
+            await nidx_api.ConfigureShards(nidx_pb2.ShardsConfig(configs=configs))
 
     @classmethod
     async def delete(cls, driver: Driver, kbid: str):
-        async with driver.transaction() as txn:
+        async with driver.rw_transaction() as txn:
             exists = await datamanagers.kb.exists_kb(txn, kbid=kbid)
             if not exists:
                 return
@@ -347,7 +394,7 @@ class KnowledgeBox:
 
         nidx_api = get_nidx_api_client()
 
-        async with driver.transaction() as txn:
+        async with driver.rw_transaction() as txn:
             storage_to_delete = KB_TO_DELETE_STORAGE.format(kbid=kbid)
             await txn.set(storage_to_delete, b"")
 
@@ -377,7 +424,7 @@ class KnowledgeBox:
     @classmethod
     async def delete_all_kb_keys(cls, driver: Driver, kbid: str, chunk_size: int = 1_000):
         prefix = KB_KEYS.format(kbid=kbid)
-        async with driver.transaction() as txn:
+        async with driver.rw_transaction() as txn:
             await txn.delete_by_prefix(prefix)
             await txn.commit()
 
@@ -535,10 +582,7 @@ class KnowledgeBox:
         request: CreateExternalIndexProviderMetadata,
         indexes: list[VectorsetExternalIndex],
     ) -> StoredExternalIndexProviderMetadata:
-        if request.type != ExternalIndexProviderType.PINECONE:
-            return StoredExternalIndexProviderMetadata(type=request.type)
-        # Only pinecone is supported for now
-        return await PineconeIndexManager.create_indexes(kbid, request, indexes)
+        return StoredExternalIndexProviderMetadata(type=request.type)
 
     @classmethod
     async def _maybe_delete_external_indexes(
@@ -546,10 +590,7 @@ class KnowledgeBox:
         kbid: str,
         stored: StoredExternalIndexProviderMetadata,
     ) -> None:
-        if stored.type != ExternalIndexProviderType.PINECONE:
-            return
-        # Only pinecone is supported for now
-        await PineconeIndexManager.delete_indexes(kbid, stored)
+        return
 
 
 def chunker(seq: Sequence, size: int):

nucliadb/ingest/orm/processor/auditing.py

@@ -34,7 +34,7 @@ async def collect_audit_fields(
         return []
 
     audit_storage_fields: list[audit_pb2.AuditField] = []
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         kb = KnowledgeBox(txn, storage, message.kbid)
         resource = Resource(txn, storage, kb, message.uuid)
         field_keys = await resource.get_fields_ids()