nucliadb 6.2.0.post2675__py3-none-any.whl → 6.2.1__py3-none-any.whl

nucliadb/ingest/orm/broker_message.py

@@ -20,11 +20,13 @@
 
 from typing import cast
 
+from nucliadb.common import datamanagers
 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.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import (
     ExtractedTextWrapper,
     ExtractedVectorsWrapper,
@@ -90,7 +92,12 @@ class _BrokerMessageBuilder:
                     self.bm.link_extracted_data.append(link_extracted_data)
 
             # Field vectors
-            await self.generate_field_vectors(type_id, field_id, field)
+            async for vectorset_id, vs in datamanagers.vectorsets.iter(
+                resource.txn, kbid=resource.kb.kbid
+            ):
+                await self.generate_field_vectors(
+                    type_id, field_id, field, vectorset_id, vs.storage_key_kind
+                )
 
             # Large metadata
             await self.generate_field_large_computed_metadata(type_id, field_id, field)
@@ -155,13 +162,16 @@ class _BrokerMessageBuilder:
         type_id: FieldType.ValueType,
         field_id: str,
         field: Field,
+        vectorset: str,
+        storage_key_kind: VectorSetConfig.StorageKeyKind.ValueType,
     ):
-        vo = await field.get_vectors()
+        vo = await field.get_vectors(vectorset, storage_key_kind)
         if vo is None:
             return
         evw = ExtractedVectorsWrapper()
         evw.field.field = field_id
         evw.field.field_type = type_id
+        evw.vectorset_id = vectorset
         evw.vectors.CopyFrom(vo)
         self.bm.field_vectors.append(evw)
 

nucliadb/ingest/orm/entities.py

@@ -26,7 +26,6 @@ from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.exceptions import (
     AlreadyExists,
     EntitiesGroupNotFound,
-    NodeError,
 )
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.datamanagers.entities import (
@@ -37,6 +36,7 @@ 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.search.search.shards import query_shard
 from nucliadb_protos.knowledgebox_pb2 import (
     DeletedEntitiesGroups,
     EntitiesGroup,
@@ -53,9 +53,6 @@ from nucliadb_protos.nodereader_pb2 import (
 )
 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
 
@@ -218,20 +215,15 @@ class EntitiesManager:
                     ],
                 ),
             )
-            response = await node.reader.Search(request)  # type: ignore
+            response = await query_shard(node, shard_id, request)
             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:
-            if isinstance(result, Exception):
-                errors.capture_exception(result)
-                raise NodeError("Error while querying relation index")
 
         entities = {}
         for result in results:
@@ -307,6 +299,7 @@ class EntitiesManager:
         shard_manager = get_shard_manager()
 
         async def query_indexed_entities_group_names(node: AbstractIndexNode, shard_id: str) -> set[str]:
+            """Search all relation types"""
             request = SearchRequest(
                 shard=shard_id,
                 result_per_page=0,
@@ -315,25 +308,21 @@ class EntitiesManager:
                 paragraph=False,
                 faceted=Faceted(labels=["/e"]),
             )
-            response: SearchResponse = await node.reader.Search(request)  # type: ignore
+            response: SearchResponse = await query_shard(node, shard_id, request)
             try:
                 facetresults = response.document.facets["/e"].facetresults
-                return {facet.tag.split("/")[-1] for facet in facetresults}
             except KeyError:
                 # No entities found
                 return set()
+            else:
+                return {facet.tag.split("/")[-1] for facet in facetresults}
 
         results = await shard_manager.apply_for_all_shards(
             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:
-            if isinstance(result, Exception):
-                errors.capture_exception(result)
-                raise NodeError("Error while looking for relations types")
 
         if not results:
             return set()

nucliadb/ingest/orm/knowledgebox.py

@@ -27,7 +27,6 @@ from grpc.aio import AioRpcError
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardNotFound
-from nucliadb.common.cluster.manager import get_index_node
 from nucliadb.common.cluster.utils import get_shard_manager
 
 # XXX: this keys shouldn't be exposed outside datamanagers
@@ -49,7 +48,6 @@ from nucliadb.ingest.orm.exceptions import (
 from nucliadb.ingest.orm.metrics import processor_observer
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb.ingest.orm.utils import choose_matryoshka_dimension, compute_paragraph_key
-from nucliadb.ingest.settings import settings
 from nucliadb.migrator.utils import get_latest_version
 from nucliadb_protos import knowledgebox_pb2, noderesources_pb2, nodewriter_pb2, writer_pb2
 from nucliadb_protos.knowledgebox_pb2 import (
@@ -58,8 +56,10 @@ from nucliadb_protos.knowledgebox_pb2 import (
     KnowledgeBoxConfig,
     SemanticModelMetadata,
     StoredExternalIndexProviderMetadata,
+    VectorSetPurge,
 )
 from nucliadb_protos.resources_pb2 import Basic
+from nucliadb_utils.settings import is_onprem_nucliadb
 from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.utilities import (
     get_audit,
@@ -74,6 +74,9 @@ KB_KEYS = "/kbs/{kbid}/"
 KB_TO_DELETE_BASE = "/kbtodelete/"
 KB_TO_DELETE_STORAGE_BASE = "/storagetodelete/"
 
+RESOURCE_TO_DELETE_STORAGE_BASE = "/resourcestoragetodelete"
+RESOURCE_TO_DELETE_STORAGE = f"{RESOURCE_TO_DELETE_STORAGE_BASE}/{{kbid}}/{{uuid}}"
+
 KB_TO_DELETE = f"{KB_TO_DELETE_BASE}{{kbid}}"
 KB_TO_DELETE_STORAGE = f"{KB_TO_DELETE_STORAGE_BASE}{{kbid}}"
 
@@ -100,9 +103,9 @@ class KnowledgeBox:
         *,
         kbid: str,
         slug: str,
+        semantic_models: dict[str, SemanticModelMetadata],
         title: str = "",
         description: str = "",
-        semantic_models: Optional[dict[str, SemanticModelMetadata]] = None,
         external_index_provider: CreateExternalIndexProviderMetadata = CreateExternalIndexProviderMetadata(),
         hidden_resources_enabled: bool = False,
         hidden_resources_hide_on_creation: bool = False,
@@ -117,7 +120,7 @@ class KnowledgeBox:
             raise KnowledgeBoxCreationError(
                 "Cannot hide new resources if the hidden resources feature is disabled"
             )
-        if semantic_models is None or len(semantic_models) == 0:
+        if len(semantic_models) == 0:
             raise KnowledgeBoxCreationError("KB must define at least one semantic model")
 
         rollback_ops: list[Callable[[], Coroutine[Any, Any, Any]]] = []
@@ -143,6 +146,7 @@ class KnowledgeBox:
                 kb_shards.actual = -1
 
                 vs_external_indexes = []
+
                 for vectorset_id, semantic_model in semantic_models.items():  # type: ignore
                     # if this KB uses a matryoshka model, we can choose a different
                     # dimension
@@ -169,6 +173,7 @@ class KnowledgeBox:
                             vector_dimension=dimension,
                         ),
                         matryoshka_dimensions=semantic_model.matryoshka_dimensions,
+                        storage_key_kind=knowledgebox_pb2.VectorSetConfig.StorageKeyKind.VECTORSET_PREFIX,
                     )
                     await datamanagers.vectorsets.set(txn, kbid=kbid, config=vectorset_config)
 
@@ -339,6 +344,8 @@ class KnowledgeBox:
         if exists is False:
             logger.error(f"{kbid} KB does not exists on Storage")
 
+        nidx_api = get_nidx_api_client()
+
         async with driver.transaction() as txn:
             storage_to_delete = KB_TO_DELETE_STORAGE.format(kbid=kbid)
             await txn.set(storage_to_delete, b"")
@@ -351,25 +358,17 @@ class KnowledgeBox:
                 logger.warning(f"Shards not found for KB while purging it", extra={"kbid": kbid})
             else:
                 for shard in shards_obj.shards:
-                    # Delete the shard on nodes
-                    for replica in shard.replicas:
-                        node = get_index_node(replica.node)
-                        if node is None:
-                            logger.error(
-                                f"No node {replica.node} found, let's continue. Some shards may stay orphaned",
-                                extra={"kbid": kbid},
-                            )
-                            continue
+                    if shard.nidx_shard_id:
                         try:
-                            await node.delete_shard(replica.shard.id)
+                            await nidx_api.DeleteShard(noderesources_pb2.ShardId(id=shard.nidx_shard_id))
                             logger.debug(
-                                f"Succeded deleting shard from nodeid={replica.node} at {node.address}",
-                                extra={"kbid": kbid, "node_id": replica.node},
+                                f"Succeded deleting shard",
+                                extra={"kbid": kbid, "shard_id": shard.nidx_shard_id},
                             )
                         except AioRpcError as exc:
                             if exc.code() == StatusCode.NOT_FOUND:
                                 continue
-                            raise ShardNotFound(f"{exc.details()} @ {node.address}")
+                            raise ShardNotFound(f"{exc.details()} @ shard {shard.nidx_shard_id}")
 
             await txn.commit()
         await cls.delete_all_kb_keys(driver, kbid)
@@ -415,9 +414,16 @@ class KnowledgeBox:
                 logger.exception("Error deleting slug")
 
     async def storage_delete_resource(self, uuid: str):
-        await self.storage.delete_resource(
-            self.kbid, uuid, max_parallel=settings.ingest_delete_resource_storage_max_parallel
-        )
+        if is_onprem_nucliadb():
+            await self.storage.delete_resource(self.kbid, uuid)
+        else:
+            # Deleting from storage can be slow, so we schedule its deletion and the purge cronjob
+            # will take care of it
+            await self.schedule_delete_resource(self.kbid, uuid)
+
+    async def schedule_delete_resource(self, kbid: str, uuid: str):
+        key = RESOURCE_TO_DELETE_STORAGE.format(kbid=kbid, uuid=uuid)
+        await self.txn.set(key, b"")
 
     async def delete_resource(self, uuid: str):
         with processor_observer({"type": "delete_resource_maindb"}):
@@ -479,6 +485,12 @@ class KnowledgeBox:
             self.txn, kbid=self.kbid, vectorset_id=config.vectorset_id
         ):
             raise VectorSetConflict(f"Vectorset {config.vectorset_id} already exists")
+
+        # To ensure we always set the storage key kind, we overwrite it with the
+        # correct value. This whole enum business is to maintain bw/c with KBs
+        # pre-vectorsets, so any new vectorset should use the vectorset prefix
+        # key kind
+        config.storage_key_kind = knowledgebox_pb2.VectorSetConfig.StorageKeyKind.VECTORSET_PREFIX
         await datamanagers.vectorsets.set(self.txn, kbid=self.kbid, config=config)
 
         # Remove the async deletion mark if it exists, just in case there was a previous deletion
@@ -491,11 +503,21 @@ class KnowledgeBox:
         await shard_manager.create_vectorset(self.kbid, config)
 
     async def delete_vectorset(self, vectorset_id: str):
-        await datamanagers.vectorsets.delete(self.txn, kbid=self.kbid, vectorset_id=vectorset_id)
+        vectorset_count = await datamanagers.vectorsets.count(self.txn, kbid=self.kbid)
+        if vectorset_count == 1:
+            raise VectorSetConflict("Deletion of your last vectorset is not allowed")
+
+        deleted = await datamanagers.vectorsets.delete(
+            self.txn, kbid=self.kbid, vectorset_id=vectorset_id
+        )
+        if deleted is None:
+            # already deleted
+            return
 
         # mark vectorset for async deletion
         deletion_mark_key = KB_VECTORSET_TO_DELETE.format(kbid=self.kbid, vectorset=vectorset_id)
-        await self.txn.set(deletion_mark_key, b"")
+        payload = VectorSetPurge(storage_key_kind=deleted.storage_key_kind)
+        await self.txn.set(deletion_mark_key, payload.SerializeToString())
 
         shard_manager = get_shard_manager()
         await shard_manager.delete_vectorset(self.kbid, vectorset_id)

nucliadb/ingest/orm/processor/data_augmentation.py

@@ -20,13 +20,13 @@
 
 import logging
 from dataclasses import dataclass, field
-from typing import Optional, Sequence
+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_protos import resources_pb2, writer_pb2
-from nucliadb_protos.resources_pb2 import FieldID, FieldType
+from nucliadb_protos.resources_pb2 import FieldType
 from nucliadb_utils.utilities import Utility, get_partitioning, get_utility
 
 logger = logging.getLogger("ingest-processor")
@@ -50,7 +50,7 @@ async def get_generated_fields(bm: writer_pb2.BrokerMessage, resource: Resource)
     ingest the processed thing later).
 
     Given a broker message and a resource, this function returns the list of
-    generated fields, that can be empty.
+    generated fields, that can be empty. It skips fields with errors.
 
     """
     generated_fields = GeneratedFields()
@@ -60,34 +60,12 @@ async def get_generated_fields(bm: writer_pb2.BrokerMessage, resource: Resource)
         return generated_fields
 
     # search all fields
-
-    all_fields = await resource.get_all_field_ids(for_update=False)
-    fields: Sequence[FieldID]
-    if all_fields is None:
-        fields = []
-    else:
-        fields = all_fields.fields
-
-    for field_id in bm.texts:
-        field = FieldID(field_type=FieldType.TEXT, field=field_id)
-        if field not in fields:
+    for field_id, text in bm.texts.items():
+        errors = [e for e in bm.errors if e.field_type == FieldType.TEXT and e.field == field_id]
+        has_error = len(errors) > 0
+        if text.generated_by.WhichOneof("author") == "data_augmentation" and not has_error:
             generated_fields.texts.append(field_id)
 
-    for field_id in bm.links:
-        field = FieldID(field_type=FieldType.LINK, field=field_id)
-        if field not in fields:
-            generated_fields.links.append(field_id)
-
-    for field_id in bm.files:
-        field = FieldID(field_type=FieldType.FILE, field=field_id)
-        if field not in fields:
-            generated_fields.files.append(field_id)
-
-    for field_id in bm.conversations:
-        field = FieldID(field_type=FieldType.CONVERSATION, field=field_id)
-        if field not in fields:
-            generated_fields.conversations.append(field_id)
-
     return generated_fields
 
 

nucliadb/ingest/orm/processor/processor.py

@@ -275,7 +275,6 @@ class Processor:
 
                     if message.source == writer_pb2.BrokerMessage.MessageSource.WRITER:
                         resource = await kb.get(uuid)
-
                         if resource is None:
                             # It's a new resource
                             resource = await kb.add_resource(uuid, message.slug, message.basic)
@@ -737,7 +736,11 @@ 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:
+    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()])
+    ):
         return True
     for field_paragraphs in index_message.paragraphs.values():
         for paragraph in field_paragraphs.paragraphs.values():